Boyang Zheng Nanye Ma Shengbang Tong Saining Xie

Résumé

La modélisation générative latente, dans laquelle un autoencodeur préentraîné projette les pixels dans un espace latent destiné au processus de diffusion, est devenue la stratégie standard des Transformers de diffusion (DiT) ; toutefois, le composant autoencodeur a peu évolué. La plupart des DiT continuent de s’appuyer sur l’encodeur VAE original, ce qui introduit plusieurs limitations : des architectures de base obsolètes, nuisant à la simplicité architecturale ; des espaces latents à faible dimension, restreignant la capacité d’information ; et des représentations faibles, résultant d’un entraînement fondé uniquement sur la reconstruction, et qui limitent finalement la qualité de génération. Dans ce travail, nous explorons le remplacement de la VAE par des encodeurs de représentations préentraînés (par exemple DINO, SigLIP, MAE), couplés à des décodeurs entraînés, formant ce que nous appelons des Autoencodeurs de Représentations (RAE). Ces modèles offrent à la fois des reconstructions de haute qualité et des espaces latents riches sur le plan sémantique, tout en permettant une architecture scalable basée sur les Transformers. Étant donné que ces espaces latents sont généralement de haute dimension, un défi majeur réside dans la capacité des Transformers de diffusion à opérer efficacement dans ces espaces. Nous analysons les sources de cette difficulté, proposons des solutions théoriquement motivées, et les validons empiriquement. Notre approche permet une convergence plus rapide, sans recourir à des pertes d’alignement de représentation auxiliaires. En utilisant une variante de DiT munie d’une tête DDT légère et large, nous obtenons des résultats remarquables en génération d’images sur ImageNet : un score FID de 1,51 à 256×256 (sans guidance) et de 1,13 à la fois à 256×256 et à 512×512 (avec guidance). Les RAE offrent des avantages clairs et devraient devenir la nouvelle référence par défaut pour l’entraînement des Transformers de diffusion.