Récemment, des chercheurs de l'Université de Californie à Irvine et d'autres institutions ont réduit le coût de formation des modèles de diffusion à 1890 dollars en utilisant des stratégies telles que le masquage différé, le MoE et l'expansion hiérarchique. C'est un ordre de grandeur inférieur aux 28400 dollars de la méthode la moins chère précédente (Wuerstchen), et encore moins que des modèles comme Stable Diffusion.
Les chercheurs ont essayé diverses approches pour réduire ces coûts énormes. Par exemple, le modèle de diffusion original nécessitait environ 1000 étapes du bruit à l'image, ce qui a été réduit à environ 20 étapes, voire moins. Lorsque les modules de base des modèles de diffusion ont été progressivement remplacés par DiT (Transformer) au lieu de Unet (CNN), certaines optimisations basées sur les caractéristiques des Transformers ont suivi, comme la quantification, le saut de certains calculs redondants dans l'attention, le pipeline, etc.
Cette étude a formé un modèle de diffusion de 11,6 milliards de paramètres à partir de zéro pour seulement 1890 dollars. C'est une amélioration d'un ordre de grandeur par rapport à l'état de l'art, donnant espoir aux gens ordinaires de pouvoir accéder à la pré-formation. Plus important encore, les techniques de réduction des coûts n'ont pas affecté les performances du modèle, les 11,6 milliards de paramètres donnant de très bons résultats. Outre l'aspect visuel, les indicateurs de données du modèle sont également excellents, avec un score FID très proche de Stable Diffusion 1.5 et DALL·E 2.
Les secrets pour économiser de l'argent comprennent principalement :
-
Stratégie de masquage différé : Utilisation d'un mélangeur de patchs pour prétraiter avant le masquage, intégrant les informations des patchs supprimés dans les patchs survivants, réduisant ainsi considérablement la baisse de performance due au masquage élevé.
-
Ajustement fin : Effectuer un léger ajustement fin (sans masque) après la pré-formation (avec masque) pour atténuer les artefacts de génération indésirables causés par l'utilisation du masque.
-
MoE et expansion hiérarchique : Utilisation de couches MoE simplifiées basées sur le routage de sélection d'experts, augmentant les paramètres et la capacité d'expression du modèle sans augmenter significativement les coûts de formation. Considération également de la méthode de mise à l'échelle hiérarchique, augmentant linéairement la largeur des blocs Transformer.
En termes de configuration expérimentale, les auteurs ont utilisé deux variantes de DiT : DiT-Tiny/2 et DiT-Xl/2, avec une taille de patch de 2. Tous les modèles ont été formés avec l'optimiseur AdamW avec décroissance du taux d'apprentissage en cosinus et forte décroissance de poids. Le front-end du modèle utilise l'auto-encodeur variationnel (VAE) à quatre canaux du modèle Stable-Diffusion-XL pour extraire les caractéristiques de l'image, et les performances du VAE à 16 canaux le plus récent ont également été testées dans la formation à grande échelle (version économique).
Les auteurs ont utilisé le cadre EDM comme configuration de formation unifiée pour tous les modèles de diffusion, utilisant le score FID et le score CLIP pour mesurer les performances des modèles de génération d'images. Le modèle CLIP le plus couramment utilisé a été choisi comme encodeur de texte.
L'ensemble de données d'entraînement comprenait trois ensembles de données d'images réelles (Conceptual Captions, Segment Anything, TextCaps), contenant 22 millions de paires image-texte. Comme SA1B ne fournit pas de vraies légendes, des légendes synthétiques générées par le modèle LLaVA ont été utilisées ici.