FlexiDiT : Votre Transformer de Diffusion Peut Facilement Générer des Échantillons de Haute Qualité avec Moins de Calcul

FlexiDiT: Your Diffusion Transformer Can Easily Generate High-Quality Samples with Less Compute

February 27, 2025

Auteurs: Sotiris Anagnostidis, Gregor Bachmann, Yeongmin Kim, Jonas Kohler, Markos Georgopoulos, Artsiom Sanakoyeu, Yuming Du, Albert Pumarola, Ali Thabet, Edgar Schönfeld

cs.AI

Résumé

Malgré leurs performances remarquables, les Transformers de Diffusion modernes sont freinés par des exigences substantielles en ressources lors de l'inférence, découlant de la quantité fixe et importante de calcul nécessaire pour chaque étape de débrui tage. Dans ce travail, nous revisitons le paradigme statique conventionnel qui alloue un budget de calcul fixe par itération de débrui tage et proposons plutôt une stratégie dynamique. Notre cadre simple et efficace en termes d'échantillonnage permet de convertir des modèles DiT pré-entraînés en modèles flexibles -- appelés FlexiDiT -- leur permettant de traiter des entrées avec des budgets de calcul variables. Nous démontrons comment un seul modèle flexible peut générer des images sans aucune perte de qualité, tout en réduisant les FLOPs requis de plus de 40\% par rapport à leurs homologues statiques, que ce soit pour la génération d'images conditionnées par classe ou par texte. Notre méthode est générale et agnostique aux modalités d'entrée et de conditionnement. Nous montrons comment notre approche peut être facilement étendue à la génération de vidéos, où les modèles FlexiDiT génèrent des échantillons avec jusqu'à 75\% de calcul en moins sans compromettre les performances.

English

Despite their remarkable performance, modern Diffusion Transformers are hindered by substantial resource requirements during inference, stemming from the fixed and large amount of compute needed for each denoising step. In this work, we revisit the conventional static paradigm that allocates a fixed compute budget per denoising iteration and propose a dynamic strategy instead. Our simple and sample-efficient framework enables pre-trained DiT models to be converted into flexible ones -- dubbed FlexiDiT -- allowing them to process inputs at varying compute budgets. We demonstrate how a single flexible model can generate images without any drop in quality, while reducing the required FLOPs by more than 40\% compared to their static counterparts, for both class-conditioned and text-conditioned image generation. Our method is general and agnostic to input and conditioning modalities. We show how our approach can be readily extended for video generation, where FlexiDiT models generate samples with up to 75\% less compute without compromising performance.