TinyFusion：学习浅层的扩散变换器
计算机视觉与模式识别
2024-12-03 v1 人工智能
机器学习
摘要
扩散变换器在图像生成方面已展现出卓越的能力，但往往伴随大量参数，导致实际应用中推理开销较大。本文提出了 TinyFusion，一种通过端到端学习剔除扩散变换器中冗余层的深度剥离方法。其核心原则是创建具有高可恢复性的剥离模型，以便在微调后恢复强大的性能。为实现这一目标，我们引入了可微分抽样技术以使剥离可学习，并配合一个共优化参数来模拟未来的微调。虽然先前方法侧重于最小化剥离后的损失或误差，但我们的方法明确建模和优化剥离后模型的性能。实验结果表明，这种可学习的范式对扩散变换器的层剥离具有显著优势，超越了现有的基于重要性的和基于误差的方法。此外，TinyFusion 在 diverse architectures（包括 DiTs、MARs 和 SiTs）上表现出强大的泛化能力。以 DiT-XL 为例，TinyFusion 可在预训练成本不足7%的情况下构建浅层扩散变换器，实现约2倍的加速，FID 分数为 2.86，优于在效率相当的情况下的竞争方法。代码可在 https://github.com/VainF/TinyFusion 获取。
引用
@article{arxiv.2412.01199,
title = {TinyFusion: Diffusion Transformers Learned Shallow},
author = {Gongfan Fang and Kunjun Li and Xinyin Ma and Xinchao Wang},
journal= {arXiv preprint arXiv:2412.01199},
year = {2024}
}