TinyFusion: Transformadores de Difusión Aprendidos Superficiales

TinyFusion: Diffusion Transformers Learned Shallow

December 2, 2024

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

cs.AI

Resumen

Los Transformadores de Difusión han demostrado capacidades notables en la generación de imágenes, pero a menudo vienen con una excesiva parametrización, lo que resulta en una considerable sobrecarga de inferencia en aplicaciones del mundo real. En este trabajo, presentamos TinyFusion, un método de poda de profundidad diseñado para eliminar capas redundantes de los transformadores de difusión a través de un aprendizaje de extremo a extremo. El principio central de nuestro enfoque es crear un modelo podado con alta capacidad de recuperación, permitiéndole recuperar un rendimiento sólido después del ajuste fino. Para lograr esto, introducimos una técnica de muestreo diferenciable para hacer que la poda sea aprendible, emparejada con un parámetro co-optimizado para simular el ajuste fino futuro. Mientras que trabajos anteriores se centran en minimizar la pérdida o el error después de la poda, nuestro método modela y optimiza explícitamente el rendimiento posterior al ajuste fino de los modelos podados. Los resultados experimentales indican que este paradigma aprendible ofrece beneficios sustanciales para la poda de capas de los transformadores de difusión, superando a los métodos existentes basados en importancia y error. Además, TinyFusion muestra una fuerte generalización en diversas arquitecturas, como DiTs, MARs y SiTs. Experimentos con DiT-XL muestran que TinyFusion puede diseñar un transformador de difusión superficial por menos del 7% del costo de pre-entrenamiento, logrando una aceleración de 2 veces con una puntuación FID de 2.86, superando a competidores con eficiencia comparable. El código está disponible en 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.