TinyFusion: Transformadores de Difusão Superficial Aprendidos

TinyFusion: Diffusion Transformers Learned Shallow

December 2, 2024

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

cs.AI

Resumo

Os Transformadores de Difusão têm demonstrado capacidades notáveis na geração de imagens, mas frequentemente vêm com uma parametrização excessiva, resultando em considerável sobrecarga de inferência em aplicações do mundo real. Neste trabalho, apresentamos o TinyFusion, um método de poda de profundidade projetado para remover camadas redundantes dos transformadores de difusão por meio de aprendizado de ponta a ponta. O princípio central de nossa abordagem é criar um modelo podado com alta capacidade de recuperação, permitindo que ele recupere um desempenho forte após o ajuste fino. Para realizar isso, introduzimos uma técnica de amostragem diferenciável para tornar a poda aprendível, combinada com um parâmetro co-otimizado para simular o ajuste fino futuro. Enquanto trabalhos anteriores se concentram em minimizar a perda ou erro após a poda, nosso método modela e otimiza explicitamente o desempenho pós-ajuste fino dos modelos podados. Resultados experimentais indicam que este paradigma aprendível oferece benefícios substanciais para a poda de camadas de transformadores de difusão, superando métodos existentes baseados em importância e erro. Além disso, o TinyFusion exibe uma forte generalização em diversas arquiteturas, como DiTs, MARs e SiTs. Experimentos com DiT-XL mostram que o TinyFusion pode criar um transformador de difusão raso a menos de 7% do custo de pré-treinamento, alcançando um aumento de velocidade de 2 vezes com um escore FID de 2.86, superando concorrentes com eficiência comparável. O código está disponível em 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.