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%未満で浅い拡散トランスフォーマーを作成し、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.