TinyFusion : Transformateurs de Diffusion Peu Profonds Appris

TinyFusion: Diffusion Transformers Learned Shallow

December 2, 2024

Auteurs: Gongfan Fang, Kunjun Li, Xinyin Ma, Xinchao Wang

cs.AI

Résumé

Les transformateurs de diffusion ont démontré des capacités remarquables dans la génération d'images, mais sont souvent caractérisés par une paramétrisation excessive, entraînant un surcoût significatif lors de l'inférence dans des applications réelles. Dans ce travail, nous présentons TinyFusion, une méthode d'élagage en profondeur conçue pour éliminer les couches redondantes des transformateurs de diffusion via un apprentissage de bout en bout. Le principe fondamental de notre approche est de créer un modèle élagué avec une capacité de récupération élevée, lui permettant de retrouver de bonnes performances après un affinage. Pour ce faire, nous introduisons une technique d'échantillonnage différenciable pour rendre l'élagage apprenable, associée à un paramètre co-optimisé pour simuler un affinage futur. Alors que les travaux antérieurs se concentrent sur la minimisation de la perte ou de l'erreur après l'élagage, notre méthode modélise et optimise explicitement les performances post-affinage des modèles élagués. Les résultats expérimentaux indiquent que ce paradigme apprenable offre des avantages substantiels pour l'élagage des couches des transformateurs de diffusion, surpassant les méthodes existantes basées sur l'importance et l'erreur. De plus, TinyFusion présente une forte généralisation à travers diverses architectures, telles que DiTs, MARs et SiTs. Les expériences avec DiT-XL montrent que TinyFusion peut créer un transformateur de diffusion peu profond pour moins de 7% du coût de pré-entraînement, atteignant un gain de vitesse de 2 fois avec un score FID de 2,86, surpassant les concurrents avec une efficacité comparable. Le code est disponible sur https://github.com/VainF/TinyFusion.

English

Diffusion Transformers have demonstrated remarkable capabilities in image generation but often come with excessive parameterization, resulting in considerable inference overhead in real-world applications. In this work, we present TinyFusion, a depth pruning method designed to remove redundant layers from diffusion transformers via end-to-end learning. The core principle of our approach is to create a pruned model with high recoverability, allowing it to regain strong performance after fine-tuning. To accomplish this, we introduce a differentiable sampling technique to make pruning learnable, paired with a co-optimized parameter to simulate future fine-tuning. While prior works focus on minimizing loss or error after pruning, our method explicitly models and optimizes the post-fine-tuning performance of pruned models. Experimental results indicate that this learnable paradigm offers substantial benefits for layer pruning of diffusion transformers, surpassing existing importance-based and error-based methods. Additionally, TinyFusion exhibits strong generalization across diverse architectures, such as DiTs, MARs, and SiTs. Experiments with DiT-XL show that TinyFusion can craft a shallow diffusion transformer at less than 7% of the pre-training cost, achieving a 2times speedup with an FID score of 2.86, outperforming competitors with comparable efficiency. Code is available at https://github.com/VainF/TinyFusion.