Formalisme DDPM — du bruit à l'image - DDPM 3

Série — Modèles génératifs appliqués aux images de galaxies Ce billet est le troisième d’une série consacrée à l’application des modèles de diffusion (DDPM) aux images de galaxies du catalogue Galaxy Zoo 2. Le code complet est disponible sur github.com/GregS1t/DDPM_GalaxyZoo2.

Notations. Ce billet suit strictement les notations de Ho et al. (2020). Les correspondances avec le code sont indiquées au fil du texte.

1. Vue d’ensemble

Un DDPM repose sur deux processus Markoviens opposés (Fig. 1) :

• Le processus forward $q$ (fixé, sans paramètres) corrompu progressivement une image $\mathbf{x}_0$ en y ajoutant du bruit gaussien jusqu’à obtenir $\mathbf{x}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$.

• Le processus reverse $p_\theta$ (appris) reconstruit l’image propre en débruitant pas à pas depuis $\mathbf{x}_T$.

2. Le processus forward $q$

2.1 Transition élémentaire

À chaque pas $t$, le processus forward ajoute un bruit gaussien contrôlé par un schedule de variance $\beta_1, \dots, \beta_T$ :

Ho et al. fixent $\beta_t$ constants (non appris), croissant linéairement de $\beta_1 = 10^{-4}$ à $\beta_T = 0.02$.

2.2 Marginal en un seul pas — la propriété clé

En posant $\alpha_t := 1 - \beta_t$ et $\bar{\alpha}t := \prod{s=1}^{t} \alpha_s$, le processus forward admet un échantillonnage en forme close à n’importe quel pas $t$ :

Ce qui se réécrit de manière équivalente — c’est la formule directement utilisée dans le code :

Cette équation est au cœur de la méthode diffuse dans notre code :

# ddpm_model.py — DDPMModel.diffuse()
# Eq. 4 de Ho et al. : x_t = sqrt(alpha_bar_t)*x0 + sqrt(1-alpha_bar_t)*eps
alpha_bar_t = self.alpha_bars[t].view(-1, 1, 1, 1)
x_noisy = alpha_bar_t.sqrt() * x0 + (1.0 - alpha_bar_t).sqrt() * noise

Pourquoi c’est utile ? Sans cette propriété, il faudrait simuler les $t$ transitions successives pour obtenir $\mathbf{x}_t$. Ici, on saute directement à n’importe quel niveau de bruit en $O(1)$ — ce qui rend l’entraînement par stochastic gradient descent efficace : on tire $t$ aléatoirement à chaque itération.

3. Le processus reverse $p_\theta$

Le processus reverse est une chaîne de Markov apprise, définie comme :

avec $p(\mathbf{x}_T) = \mathcal{N}(\mathbf{0}, \mathbf{I})$ et des transitions gaussiennes :

Ho et al. fixent la variance à $\boldsymbol{\Sigma}\theta = \sigma_t^2 \mathbf{I}$ avec $\sigma_t^2 = \beta_t$ (non appris). Seule la moyenne $\boldsymbol{\mu}\theta$ est paramétrée par le réseau.

4. L’objectif d’entraînement

4.1 La borne variationnelle (ELBO)

L’entraînement minimise une borne supérieure de la log-vraisemblance négative (ELBO) :

• $L_T$ : KL entre la distribution finale du forward et $\mathcal{N}(\mathbf{0},\mathbf{I})$ Constant car $\beta_t$ est fixé → ignoré pendant l’entraînement.
• $L_{t-1}$ : KL entre la posterior forward et la transition reverse apprise. C’est le terme principal.
• $L_0$ : terme de reconstruction au dernier pas.

Les détails de calculs sont disponibles dans l’annexe A (Extended derivations) du papier de (Ho et al., 2015)

4.2 La posterior forward $q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0)$

La posterior forward conditionnée sur $\mathbf{x}_0$ s’écrit comme une gaussienne explicite — sa moyenne $\tilde{\boldsymbol{\mu}}_t$ et sa variance $\tilde{\beta}_t$ ont des expressions exactes calculables directement à partir de $\mathbf{x}_0$, $\mathbf{x}_t$ et des paramètres du schedule $\bar{\alpha}_t$, sans approximation ni intégrale numérique :

avec :

4.3 Paramétrisation par la prédiction du bruit

Ho et al. montrent qu’il est plus efficace de paramétrer $\boldsymbol{\mu}_\theta$ non pas comme une prédiction directe de la moyenne, mais via la prédiction du bruit $\boldsymbol{\epsilon}$.

En substituant la reparamétrisation de l’Eq. (4) dans l’expression de $\tilde{\boldsymbol{\mu}}_t$, on obtient :

où $\boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)$ est le réseau U-Net décrit dans le billet 1, conditionné sur $t$ par les embeddings sinusoïdaux du billet 2.

4.4 LOSS simplifiée $L_\text{simple}$

En substituant l’Eq. (11) dans les termes $L_{t-1}$ de l’ELBO, et en simplifiant les poids, Ho et al. proposent la fonction de Loss suivante :

C’est la MSE entre le bruit réel et le bruit prédit — directement implémentée dans compute_loss :

# ddpm_model.py — DDPMModel.compute_loss()
# L_simple (Ho et al., Eq. 14)
x_noisy, noise = self.diffuse(x0, t, noise=noise)   # Eq. 4
noise_pred      = self.predict_noise(x_noisy, t)     # eps_theta(x_t, t)
return torch.mean((noise - noise_pred) ** 2)          # MSE

Pourquoi $L_\text{simple}$ et non l’ELBO complet ? L’ELBO pondère chaque terme $L_{t-1}$ par $\frac{\beta_t^2}{2\sigma_t^2 \alpha_t (1-\bar{\alpha}t)}$ (Ho, Eq. 12). Cette pondération défavorise les petits $t$ (bruit faible), ce qui nuit à la qualité des échantillons. $L\text{simple}$ supprime cette pondération, donnant un poids égal à tous les pas de temps — résultat : meilleure qualité d’image en pratique (Ho et al., Table 2).

5. Algorithmes

Algorithme 1 — Entraînement

Implémentation dans train_gz2_ddpm.py :

# train_gz2_ddpm.py — boucle d'entraînement
x0 = batch[0].to(device)
t  = torch.randint(0, ddpm.n_steps, (batch_size,), device=device)

with autocast():
    loss = ddpm.compute_loss(x0, t)   # L_simple, Eq. 14

scaler.scale(loss).backward()

Algorithme 2 — Échantillonnage

Implémentation dans ddpm_model.py — generate_samples() :

# ddpm_model.py — generate_samples()
for t in reversed(range(ddpm.n_steps)):
    t_batch    = torch.full((n_samples,), t, device=device, dtype=torch.long)
    noise_pred = ddpm.predict_noise(x, t_batch)        # eps_theta(x_t, t)
    alpha_t    = ddpm.alphas[t]
    alpha_bar_t = ddpm.alpha_bars[t]

    # Eq. 11 : mise à jour vers x_{t-1}
    x = (1.0 / alpha_t.sqrt()) * (
        x - (1.0 - alpha_t) / (1.0 - alpha_bar_t).sqrt() * noise_pred
    )
    if t > 0:
        sigma_t = ddpm.betas[t].sqrt()   # sigma_t^2 = beta_t (option 1)
        x = x + sigma_t * torch.randn_like(x)

6. Résumé des équations clés

Références

1. Ho, J., Jain, A., & Abbeel, P. (2020). Denoising Diffusion Probabilistic Models. NeurIPS 2020. arXiv:2006.11239

2. Nichol, A., & Dhariwal, P. (2021). Improved Denoising Diffusion Probabilistic Models. ICML 2021. arXiv:2102.09672

Prochain billet : Application à Galaxy Zoo 2 — résultats →