微融合：学习浅层扩散变压器

摘要

扩散变压器在图像生成方面展示出卓越的能力，但往往伴随着过多的参数化，导致在实际应用中存在相当大的推理开销。在这项工作中，我们提出了TinyFusion，一种深度修剪方法，旨在通过端到端学习从扩散变压器中去除多余的层。我们方法的核心原则是创建一个具有高可恢复性的修剪模型，使其在微调后能够恢复强大的性能。为实现这一目标，我们引入了一种可微分采样技术，使修剪可学习化，并配以一个协同优化的参数来模拟未来的微调。虽然先前的研究侧重于在修剪后最小化损失或错误，但我们的方法明确地对修剪模型在微调后的性能进行建模和优化。实验结果表明，这种可学习的范式为扩散变压器的层修剪提供了实质性的好处，超越了现有的基于重要性和错误的方法。此外，TinyFusion在各种架构上都表现出强大的泛化能力，如DiTs、MARs和SiTs。对DiT-XL的实验表明，TinyFusion可以以不到预训练成本的7%构建一个浅层扩散变压器，在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.

微融合：学习浅层扩散变压器

TinyFusion: Diffusion Transformers Learned Shallow

摘要

Summary

Support

Support