微融合：學習淺層擴散轉換器

TinyFusion: Diffusion Transformers Learned Shallow

December 2, 2024

作者: Gongfan Fang, Kunjun Li, Xinyin Ma, Xinchao Wang

cs.AI

摘要

擴散Transformer 在圖像生成方面展現出卓越的能力，但往往伴隨著過多的參數化，導致在實際應用中存在相當大的推理開銷。在這項工作中，我們提出了TinyFusion，一種深度修剪方法，旨在通過端對端學習從擴散Transformer 中刪除多餘的層。我們方法的核心原則是創建一個具有高可恢復性的修剪模型，使其在微調後能夠恢復出強大的性能。為了實現這一目標，我們引入了一種可微取樣技術，使修剪變得可學習，並配合一個共同優化的參數來模擬未來的微調。儘管先前的研究著重於在修剪後最小化損失或錯誤，我們的方法明確地對修剪後模型的微調後性能進行建模和優化。實驗結果表明，這種可學習的範式對於擴散Transformer 的層修剪提供了顯著的好處，超越了現有的基於重要性和錯誤的方法。此外，TinyFusion 在各種架構上都表現出強大的泛化能力，如DiTs、MARs 和SiTs。對於DiT-XL 的實驗表明，TinyFusion 可以以不到預訓練成本的 7% 打造一個淺層擴散Transformer，在 FID 分數為 2.86 的情況下實現 2 倍加速，勝過具有可比效率的競爭對手。代碼可在 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.