PRX Partie 3 — Former un modèle de texte à image en 24h !

PRX Partie 3 — Former un modèle de texte à image en 24h !

Dans les deux derniers articles (Partie 1 et Partie 2), nous avons exploré une large gamme d'astuces architecturales et d'entraînement pour les modèles de diffusion. Nous avons essayé d'évaluer chaque idée isolément, mesurant le débit, la vitesse de convergence et la qualité d'image finale, tout en cherchant à comprendre ce qui fait réellement la différence.

Dans cet article, nous voulons répondre à une question beaucoup plus pratique :

Que se passe-t-il lorsque nous combinons toutes les astuces qui ont fonctionné ?

Au lieu d'optimiser une dimension à la fois, nous allons empiler les ingrédients les plus prometteurs ensemble et voir jusqu'où nous pouvons pousser les performances sous un budget de calcul strict.

Pour rendre les choses concrètes, nous réalisons un speedrun de 24 heures :

• Budget total de calcul : environ 1500 $ (2 $/heure/GPU)

Cela est très éloigné des 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ù un 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 (lien Github), qui contient :

• Le code d'entraînement utilisé pour ce speedrun
• Le cadre expérimental de l'article de blog précédent

Ainsi, vous pouvez reproduire, modifier et étendre tout vous-même.

La recette d'entraînement

Maintenant, passons en revue ce qui a été intégré dans ce run de 24h.

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" [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 [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.])
• 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 de :

• Appliquer les pertes perceptuelles sur des images complètes regroupées plutôt que sur des caractéristiques par patch
• 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).