Mise à l’échelle efficace des transformateurs de diffusion via μP

Les Transformers de diffusion sont apparus comme la base des modèles génératifs d'images, mais leur évolutivité est limitée par le coût élevé de l'ajustement des hyperparamètres (HP) à grande échelle. Récemment, la paramétrisation avec mise à jour maximale (muP) a été proposée pour les Transformers classiques, permettant un transfert stable des HP des petits aux grands modèles de langage et réduisant considérablement les coûts d'ajustement. Cependant, il n'est pas clair si cette méthode muP s'étend aux Transformers de diffusion, qui diffèrent tant sur le plan architectural qu'objectif. Dans ce travail, nous généralisons la muP standard aux Transformers de diffusion et validons son efficacité par des expériences à grande échelle. Tout d'abord, nous démontrons rigoureusement que la muP des Transformers de diffusion courants, tels que DiT, U-ViT, PixArt-alpha et MMDiT, est conforme à celle du Transformer classique, ce qui permet une application directe des méthodologies muP existantes. En exploitant ce résultat, nous montrons systématiquement que DiT-muP bénéficie d'une transférabilité robuste des HP. Notamment, DiT-XL-2-muP avec un taux d'apprentissage transféré converge 2,9 fois plus rapidement que le DiT-XL-2 original. Enfin, nous validons l'efficacité de la muP pour la génération d'images à partir de texte en faisant évoluer PixArt-alpha de 0,04 milliard à 0,61 milliard de paramètres et MMDiT de 0,18 milliard à 18 milliards de paramètres. Dans les deux cas, les modèles sous muP surpassent leurs baselines respectives tout en nécessitant un faible coût d'ajustement : seulement 5,5 % d'une exécution d'entraînement pour PixArt-alpha et 3 % du temps consommé par des experts humains pour MMDiT-18B. Ces résultats établissent la muP comme un cadre principiel et efficace pour l'évolutivité des Transformers de diffusion.