Modèles de diffusion : un entraînement révolutionnaire en 24 heures

Dans les deux précédents articles de cette série, nous avons exploré une variété d'astuces architecturales et d'entraînement pour les modèles de diffusion. Chaque idée a été évaluée individuellement pour mesurer le débit, la vitesse de convergence et la qualité d'image finale, tout en cherchant à comprendre les éléments qui font réellement la différence.

Dans cet article, nous abordons une question plus pratique : que se passe-t-il lorsque toutes les astuces qui ont fonctionné sont combinées ? Plutôt que d'optimiser une dimension à la fois, nous avons décidé d'empiler les ingrédients les plus prometteurs pour voir jusqu'où nous pouvons pousser les performances sous un budget de calcul strict.

Pour rendre cela concret, nous avons réalisé un speedrun de 24 heures avec un budget total de calcul d'environ 1500 $, soit 2 $ par heure et par GPU. Cela contraste fortement avec les débuts des modèles de diffusion, où former des modèles compétitifs pouvait coûter des millions de dollars. L'objectif ici est de démontrer à quel point le domaine a évolué et jusqu'où une ingénierie soignée peut vous mener en seulement une journée d'entraînement.

Ce speedrun n'est pas seulement une expérience amusante. Il servira probablement de base pour notre recette d'entraînement à grande échelle à l'avenir. En plus des résultats, nous mettons également notre code en open source, qui contient le code d'entraînement utilisé pour ce speedrun ainsi que le cadre expérimental des articles de blog précédents. Ainsi, vous pouvez reproduire, modifier et étendre tout vous-même.

La recette d'entraînement

X-prédiction et entraînement dans l'espace pixel

Nous utilisons la formulation de x-prédiction de l'article "Back to Basics: Let Denoising Generative Models Denoise" par Li et He, 2025. Comme vu dans la Partie 2, cela permet d'entraîner directement dans l'espace pixel et élimine complètement le besoin d'un VAE. Nous utilisons une taille de patch de 32 et un goulot d'étranglement de 256 dimensions dans la couche de projection de tokens initiale. Ce design maintient la longueur de séquence sous contrôle, rendant l'entraînement dans l'espace pixel gérable même à des résolutions plus élevées.

À 512px, la longueur de séquence est : (512 / 32)^2 = 256. À 1024px, la longueur de séquence devient : (1024 / 32)^2 = 1024. Au lieu de suivre le programme habituel de 256px → 512px → 1024px, nous commençons directement à 512px puis affinons à 1024px. Avec des comptes de tokens contrôlés et du matériel moderne, l'entraînement dans l'espace pixel n'est plus prohibitif. C'est simplement une formulation plus propre et plus directe.

Pertes perceptuelles

Un effet secondaire très intéressant de la prédiction de x0 directement dans l'espace pixel est que nous pouvons réutiliser une boîte à outils entière de la vision par ordinateur classique. Lorsque votre modèle produit des latents, la supervision perceptuelle devient délicate. Vous devez soit décoder de nouveau vers les pixels, soit définir des pertes dans un espace latent appris qui peut ou non s'aligner avec la perception humaine. Une fois que vous prédisez des pixels directement, tout redevient simple. Vous pouvez intégrer les pertes perceptuelles exactement comme elles ont été conçues à l'origine.

Nous nous inspirons de l'article "PixelGen : Pixel Diffusion Beats Latent Diffusion with Perceptual Loss" par Ma et al., où les auteurs introduisent des objectifs perceptuels supplémentaires en plus de la perte de diffusion. Ils montrent que l'ajout de signaux perceptuels peut améliorer de manière significative la vitesse de convergence et la qualité visuelle finale.

Pour ce run de 24h, nous ajoutons deux pertes auxiliaires : LPIPS (Zhang et al.) et une perte perceptuelle basée sur DINO (nous utilisons DINOv2 [Oquab et al.]). L'idée est simple : en plus de l'objectif de correspondance de flux standard, nous encourageons l'image propre prédite à correspondre à l'image cible dans un espace de caractéristiques perceptuelles. LPIPS capture la similarité perceptuelle de bas niveau, tandis que les caractéristiques de DINO fournissent un signal sémantique plus fort.

Nous conservons la même idée générale que dans l'article, mais nous avons ajusté quelques détails. Dans nos expériences, nous avons constaté empiriquement qu'il était plus efficace d'appliquer les pertes perceptuelles sur des images complètes regroupées plutôt que sur des caractéristiques par patch et de les appliquer à tous les niveaux de bruit. Ce sont de petits détails d'implémentation, mais dans notre configuration, ils ont systématiquement donné de meilleurs résultats.

Nous avons utilisé un poids de 0.1 pour la perte LPIPS et 0.01 pour la perte perceptuelle DINO, correspondant aux valeurs recommandées dans l'article original. Ces pertes sont légères par rapport au passage avant principal du transformer, et dans notre configuration, elles ajoutent seulement un léger surcoût tout en fournissant un gain de qualité constant.

Routage de tokens avec TREAD

Pour rendre chaque étape moins coûteuse, nous utilisons le routage de tokens avec TREAD (Krause et al., 2025), qui sélectionne aléatoirement une fraction de tokens et les laisse contourner un bloc contigu de transformers, puis les réinjecte plus tard pour ne rien perdre. Nous avons choisi TREAD plutôt que SPRINT (Park et al., 2025) principalement pour sa simplicité, et parce que la complexité supplémentaire de SPRINT ne semblait pas valoir les économies de calcul relativement faibles dans notre configuration (longueur de séquence de 64 contre 128 avec TREAD à 512px).