TinyFusion: Поверхностные трансформеры с диффузией, выученные

TinyFusion: Diffusion Transformers Learned Shallow

December 2, 2024

Авторы: Gongfan Fang, Kunjun Li, Xinyin Ma, Xinchao Wang

cs.AI

Аннотация

Диффузионные трансформеры продемонстрировали выдающиеся возможности в генерации изображений, но часто имеют избыточное параметризацию, что приводит к значительным накладным расходам при выводе в реальных приложениях. В данной работе мы представляем TinyFusion, метод обрезки глубины, разработанный для удаления избыточных слоев из диффузионных трансформеров с помощью обучения с применением конечных точек. Основным принципом нашего подхода является создание обрезанной модели с высокой восстанавливаемостью, позволяющей ей достичь высокой производительности после тонкой настройки. Для достижения этой цели мы вводим дифференцируемую технику выборки для обучения обрезки, совмещенную с кооптимизированным параметром для имитации будущей тонкой настройки. В то время как предыдущие работы сосредотачивались на минимизации потерь или ошибок после обрезки, наш метод явно моделирует и оптимизирует производительность обрезанных моделей после тонкой настройки. Экспериментальные результаты показывают, что этот обучаемый парадигма приносит существенные преимущества для обрезки слоев диффузионных трансформеров, превосходящих существующие методы на основе важности и ошибок. Кроме того, TinyFusion проявляет сильное обобщение на различных архитектурах, таких как DiTs, MARs и SiTs. Эксперименты с DiT-XL показывают, что TinyFusion может создать неглубокий диффузионный трансформер менее чем за 7% от стоимости предварительного обучения, достигая ускорения в 2 раза с оценкой FID 2,86, превосходя конкурентов с сопоставимой эффективностью. Код доступен по адресу 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.