DDPM

Préambule

Le code relié à ce post se trouve sur mon github.

Ce post est écrit par moi (humain), mais l’IA m’aide à corriger les fautes de langue et améliorer mes skills de coloriage en Markdown.

Montons le niveau de ce blog avec du generative modeling qui constitue une grosse partie de mes prochains travaux. Plus particulièrement, on va se concentrer sur de la diffusion en utilisant tout d’abord des notions probabilistes (Denoising Diffusion Probabilistic Models). Ce post sera probablement mis à jour ou d’autres parties seront faites.

Je ne compte pas faire un énième post en dérivant les équations et développer le principe, car plein d’articles géniaux ont été faits : lilan-weng ou encore Sanderson. Je compte surtout parler de ce que je n’ai pas trouvé dans ces articles, les points importants que j’ai remarqués et comment coder cela concrètement.

Ce post sera actualisé, car je compte approfondir en continuant les grandes lignes des modèles de diffusion.

Introduction

Les modèles de diffusion sont des modèles génératifs devenus très importants après l’apparition des DDPM en 2020, en particulier parce qu’ils ont progressivement surpassé les GANs sur de nombreuses tâches de génération d’images.

DDPM : l’idée de base

Commençons par introduire le modèle basique, DDPM.

Les modèles de diffusion sont des modèles génératifs qui cherchent à apprendre une distribution de données $p_{\text{data}}$, par exemple une distribution d’image.

L’idée de DDPM est de ne pas apprendre directement cette distribution. À la place, on définit d’abord un processus de bruitage progressif : on part d’une image réelle $x_0 \sim p_{\text{data}}$, puis on ajoute petit à petit du bruit gaussien jusqu’à obtenir une variable $x_T$ proche d’un bruit pur $\mathcal{N}(0,I)$.

Le modèle apprend ensuite le processus inverse : partir d’un bruit pur et retirer progressivement le bruit pour revenir vers une image réaliste. Autrement dit, au lieu d’apprendre directement une distribution d’images très complexe, on apprend une suite de petits pas de débruitage.

C’est cette idée qui rend le problème plus simple : chaque étape $x_t \rightarrow x_{t-1}$ est locale, beaucoup plus facile à apprendre que la transformation complète $\mathcal{N}(0,I) \rightarrow p_{\text{data}}$.

Le processus forward est défini par :

\[\color{orange}{ q(x_t \mid x_{t-1}) = \mathcal{N} \left( x_t ; \sqrt{1-\beta_t}x_{t-1}, \beta_t I \right).}\]

On pose :

\[\alpha_t = 1-\beta_t, \qquad \bar{\alpha}_t = \prod_{s=1}^t \alpha_s.\]

Une propriété très utile des gaussiennes (appelée “the nice property” par L.Weng) est que l’on peut écrire directement $x_t$ en fonction de $x_0$ :

\[x_t = \sqrt{\bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha}_t}\varepsilon, \qquad \varepsilon \sim \mathcal{N}(0,I).\]

Le réseau est alors entraîné à prédire le bruit ajouté :

\[\color{red}{\boldsymbol{ \mathcal{L}_{\text{simple}} = \mathbb{E}_{t,x_0,\varepsilon} \left[ \left\| \varepsilon - \varepsilon_\theta(x_t,t) \right\|^2 \right]. }}\]

Un point très important est que la solution pour $\color{red}{\varepsilon_\theta(x_t,t)}$ que doit donner le modèle correspond à $\color{red}{\mathbb{E} \left[ \varepsilon \mid x_t, t \right]}$.

À la génération, on part de $x_T \sim \mathcal{N}(0,I)$, puis on applique progressivement le modèle appris pour obtenir $x_0$.

Le setup

MNIST est le dataset classique que je pensais inutile, car trop simple, néanmoins si l’on s’en sert bien sa simplicité est très utile. Utiliser MNIST et entraîner un modèle sur peu de données permet de faire des expériences rapidement (quelques minutes maximum pour entraîner) quand l’on a accès uniquement au GPU de Colab et donc de comprendre rapidement les modèles et les moduler. Il est aussi intéressant de concevoir des modèles assez puissants pour que ça marche, mais pas assez puissants pour voir les limites du modèle et ses échecs.

Ce que j’ai appris en codant le DDPM

Coder le modèle de diffusion n’est pas difficile, mais comprendre les choix qui ont été faits, leur importance et les voir à l’œuvre est plus compliqué.

L’architecture U-net est puissante

Le Deep Learning nous sert dans notre cas à prédire le bruit ajouté à l’instant $t$, soit $\varepsilon_\theta(x_t, t)$, et j’ai commencé simplement en voulant utiliser un MLP (neuronet dans mon code). J’ai voulu optimiser pendant 1h30 mon MLP pour faire descendre la loss, mais malheureusement même sur un dataset comme MNIST un MLP basique n’a aucune chance :
il doit apprendre de zéro la structure des images (des pixels proches sont vraisemblablement corrélés) ce que font les convolutions (biais inductif spatial) et doit apprendre à générer des images à des niveaux de bruits différents où il doit donner des détails fins à petit $t$ et la structure globale à grand $t$…
Le MLP a donc trop de choses à apprendre tout seul et sans un gros modèle avec énormément de paramètres et d’entraînement, c’est très compliqué.

Je me suis donc mis à faire un U-net, ce qui était bien plus costaud à coder, mais aussi très formateur.
Le U-net était à 6 années-lumière du MLP et demande peu de paramètres par rapport à ce qu’il propose.
Mon U-net comprend quelques améliorations récentes telles que le rescaling par $1/\sqrt{2}$ ou encore l’adaGN, ce qui donne de meilleurs résultats.
La raison réside dans l’architecture du U-net. Il apprend les détails fins grâce aux skip connections et la structure globale en faisant du downsampling/bottleneck, au contraire d’un MLP.
De plus, un U-net comprend la spatialité des images car il est constitué de convolutions : biais inductif spatial (je n’ai pas ajouté de transformers dans ce 1er U-net) et est donc très efficace sur des images et le problème de diffusion tout en demandant peu de compute (par rapport à un MLP équivalent).

Voici quelques exemples de génération du modèle de diffusion à partir d’un bruit aléatoire :

1k5 images, 32-Unet, 500 epoch

Étant donné que l’on entraîne notre modèle sur seulement 1 500 images avec un modèle de faible capacité, je suis plutôt satisfait du résultat.

Le temps est intéressant

Le temps paraît anodin et on n’y pense pas vraiment au premier abord (si ce n’est « le modèle peut savoir grâce au temps s’il doit ajouter plus ou moins de bruit »), mais 2 questions me sont venues.
« Que se passe-t-il si je supprime le temps dans mon U-net ? »
Le temps apparaît toujours dans le sampling avec le beta-schedule, mais je peux le supprimer du contexte de mon U-net.
Notons que le temps est encodé avec un embedding sinusoïdal, il faut donner des tenseurs au modèle.
Si l’on supprime le temps comme contexte du modèle, le modèle prend une entrée et essaie de la débruiter sans savoir a priori à quel point elle est bruitée :
On peut supposer qu’un modèle avec suffisamment de données et de capacité (ce qui n’est pas vraiment l’idée de ce projet) pourrait saisir la distribution par rapport au temps, mais cela semble instable (imaginez une image légèrement bruitée par hasard et le modèle la prendrait comme une entrée bruitée plutôt que comme une image de la distribution).

Le DDPM n’est tout de même pas facile à entraîner sur un petit dataset quand on veut voir les limites du modèle avec peu de compute, mais c’est là qu’on apprend le plus, I guess.

Il faut garder à l’esprit que les biais s’accumulent pendant le sampling (même si le bruit ajouté régule cette accumulation)et qu’avec moins de données ou moins d’époques, les résultats changent beaucoup.
Voici la différence en changeant de 200 epoch :

1k5 images, U-netdim=32, sans temps, 500 époques

Le modèle apprend la moyenne de tous les chiffres, ce qui donne cette forme floue (underfit), ou parfois il ne trouve pas la bonne direction dans l’espace et ne génère que du bruit.

700 époques

Ma deuxième question était : « comment le modèle apprend-il au cours du temps ? »

C’est là qu’intervient l’aspect intéressant de la « simple loss » du papier DDPM : Si l’on garde la loss ELBO initiale, les petits bruits sont moins pondérés et les plus grands ont plus de poids, ainsi le modèle s’entraîne davantage sur les grands bruits et moins sur les petits où il doit être précis pour retrouver les détails.
Utiliser une loss uniforme contrebalance cela en équilibrant l’apprentissage à travers le temps.
Vous avez accès à la loss pour chaque pas de temps dans mon code.

DDIM

Avec les expériences précédentes, on voit à quel point le modèle est sensible à la gestion des pas de temps et du temps, néanmoins une limitation du DDPM, c’est le sampling où les pas de temps jouent le rôle principal.
Une hypothèse principale du modèle est que l’on a des gaussiennes comme posterior avec des petits pas de temps, donc plus on fait de pas de temps plus cette hypothèse est vérifiée.
Pour l’instant, le modèle apprend à débruiter petit à petit, mais ça veut dire que quand on sample, on doit faire 300 passes avec le modèle (j’ai choisi 300 pas de temps) ce qui demande du temps quand on veut générer plusieurs images.

À plus grande échelle et sur des images de haute résolution (pas 28×28 comme MNIST), la différence de temps de sampling est abyssale entre un DDPM et un GAN : « For example, it takes around 20 hours to sample 50k images of size 32 × 32 from a DDPM, but less than a minute to do so from a GAN on an Nvidia 2080 Ti GPU. » Song et al.

Avant de passer au DDIM, j’ai essayé de simplement sauter des étapes dans le sampling DDPM classique au lieu d’aller de 1 en 1, mais ce n’était pas concluant.

Le principal insight de DDIM, c’est que la loss ne dépend pas de la loi jointe $\color{green}{\boldsymbol{q(x_1,\ldots,x_T \mid x_0)}}$, mais uniquement des $\color{green}{\boldsymbol{q(x_t \mid x_0)}}$.
Dans le reverse process, on va de 1 en 1 pour suivre la chaîne de Markov inverse, car c’est ce que nous impose la loi jointe $q(x_1,\ldots,x_T \mid x_0) = q(x_T \mid x_0)\prod_t q(x_{t-1} \mid x_t,x_0)$, mais la loss est en fait bien moins contraignante. C’est là où la propriété des gaussiennes « the nice property » est magique, on a un lien direct entre $x_t$ et $x_0$. Sans ces factorisations de gaussiennes, on serait contraint de suivre la loi inverse $q(x_{t-1} \mid x_t)$.
L’idée de DDIM, c’est de voir que l’on peut choisir un processus différent avec une loi jointe différente, mais qui a les mêmes marginales $q(x_t \mid x_0)$, ce qui donne exactement la même loss. Ils utilisent ainsi un sampling déterministe (non gaussien comme dans le DDPM où on rajoute du bruit à chaque pas de temps) et font $T/S$ pas au lieu de $T$.
Cela n’était pas concluant en DDPM à cause du bruit ajouté (stochasticité).
On ne peut néanmoins toujours pas passer de $T$ à $0$ directement, notre modèle découpe l’information petit à petit même avec DDIM.

Ce qui est néanmoins intéressant de noter, c’est que DDIM nécessite un meilleur modèle pour sampler aussi bien que DDPM. Le bruit ajouté en DDPM lors du sampling peut corriger l’erreur du modèle, ce qui freine l’accumulation d’erreur tout au long du sampling.
L’erreur va se propager dans le DDIM car c’est devenu déterministe et que l’on n’a plus cette stochasticité.

Avec assez de données, j’arrive à diviser par 5 le nombre de steps sans trop de perte de qualité:

Unet=64, 4k images, 700 époques, S=5

L’output est néanmoins moins varié que le DDPM classique et c’est dû à la nature déterministe du sampling et au fait que le modèle reste entraîné sur peu de données :

même modèle mais DDPM

Classifier-free guidance

On sait maintenant apprendre une distribution et sampler, mais on sample une image aléatoire. On aimerait conditionner et choisir ce que l’on sample.
Pour ça, on fait de la classifier guidance ou ici classifier-free guidance (vous allez voir pourquoi dans quelques lignes).
Géométriquement, c’est assez simple si on adopte la vue du score ($\nabla \log p_{t}(x_t)$) : Le score nous permet de nous diriger dans l’espace jusqu’à la distribution. Notre loss DDPM revient en fait à apprendre le score, ou du moins c’est équivalent (cf Score matching Y.Song). Cela vient de la formule de Tweedie (ajouter formule de Tweedie !!) En utilisant Bayes, on a directement que :

\[\nabla \log p(x_t \mid y) = \nabla \log p(x_t) + \nabla \log p(y \mid x_t).\]

Et c’est le deuxième terme qui correspond à un classifier qui nous permet d’aller vers un endroit choisi de la distribution. Maintenant, on ne va pas entraîner un classifier externe ici mais entraîner la moitié du temps le modèle à avoir le conditionnement comme données durant l’entraînement et l’autre moitié sans conditionnement.
Le classifier est interne au modèle, d’où le nom classifier-free guidance (cfg). Ainsi, on va pouvoir ajouter les deux contributions et nous diriger correctement dans l’espace.

Conditionnement sur 1 dans l’espace

References

• Ho, J., Jain, A., & Abbeel, P. (2020). Denoising Diffusion Probabilistic Models. NeurIPS 2020. arxiv.org/abs/2006.11239
• Song, J., Meng, C., & Ermon, S. (2020). Denoising Diffusion Implicit Models. ICLR 2021. arxiv.org/abs/2010.02502
• Song, Y., & Ermon, S. (2019). Generative Modeling by Estimating Gradients of the Data Distribution. NeurIPS 2019. arxiv.org/abs/1907.05600
• Weng, L. (2021). What are Diffusion Models? Lilian Weng’s Blog. lilianweng.github.io/posts/2021-07-11-diffusion-models