TinyFusion: Ondiepe Diffusie Transformers Aangeleerd

TinyFusion: Diffusion Transformers Learned Shallow

December 2, 2024

Auteurs: Gongfan Fang, Kunjun Li, Xinyin Ma, Xinchao Wang

cs.AI

Samenvatting

Diffusion Transformers hebben opmerkelijke capaciteiten aangetoond in beeldgeneratie, maar gaan vaak gepaard met overmatige parameterisatie, wat resulteert in aanzienlijke inferentie-overhead in real-world toepassingen. In dit werk presenteren we TinyFusion, een dieptesnoeimethode ontworpen om overbodige lagen uit diffusion transformers te verwijderen via end-to-end leren. Het kernprincipe van onze aanpak is het creëren van een gesnoeid model met een hoge herstelbaarheid, waardoor het sterke prestaties kan herwinnen na fijnafstemming. Om dit te bereiken, introduceren we een differentieerbare bemonsteringstechniek om snoeien leerbaar te maken, in combinatie met een co-geoptimaliseerde parameter om toekomstige fijnafstemming te simuleren. Terwijl eerdere werken zich richten op het minimaliseren van verlies of fout na snoeien, modelleert en optimaliseert onze methode expliciet de prestaties na fijnafstemming van gesnoeide modellen. Experimentele resultaten geven aan dat dit leerparadigma aanzienlijke voordelen biedt voor laagsnoei van diffusion transformers, waarbij bestaande belangrijkheidsgebaseerde en foutgebaseerde methoden worden overtroffen. Bovendien vertoont TinyFusion sterke generalisatie over diverse architecturen, zoals DiTs, MARs en SiTs. Experimenten met DiT-XL tonen aan dat TinyFusion een ondiepe diffusion transformer kan creëren voor minder dan 7% van de voorafgaande trainingskosten, met een 2 keer versnelling en een FID-score van 2.86, waarbij concurrenten met vergelijkbare efficiëntie worden overtroffen. De code is beschikbaar op 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.