TinyFusion: Trasformatori di Diffusione Appresi Superficialmente

TinyFusion: Diffusion Transformers Learned Shallow

December 2, 2024

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

cs.AI

Abstract

I Transformer di Diffusione hanno dimostrato notevoli capacità nella generazione di immagini ma spesso presentano un'eccessiva parametrizzazione, con conseguente considerevole overhead di inferenza nelle applicazioni reali. In questo lavoro, presentiamo TinyFusion, un metodo di potatura della profondità progettato per rimuovere strati ridondanti dai transformer di diffusione tramite apprendimento end-to-end. Il principio fondamentale del nostro approccio è quello di creare un modello potato con un'elevata recuperabilità, consentendogli di ripristinare elevate prestazioni dopo il fine-tuning. Per raggiungere questo obiettivo, introduciamo una tecnica di campionamento differenziabile per rendere l'apprendimento della potatura possibile, abbinata a un parametro co-ottimizzato per simulare il futuro fine-tuning. Mentre i lavori precedenti si concentrano sulla minimizzazione della perdita o dell'errore dopo la potatura, il nostro metodo modella ed ottimizza esplicitamente le prestazioni post-fine-tuning dei modelli potati. I risultati sperimentali indicano che questo paradigma apprendibile offre notevoli vantaggi per la potatura degli strati dei transformer di diffusione, superando i metodi esistenti basati sull'importanza e sull'errore. Inoltre, TinyFusion mostra una forte generalizzazione tra diverse architetture, come DiTs, MARs e SiTs. Gli esperimenti con DiT-XL mostrano che TinyFusion può creare un transformer di diffusione superficiale a meno del 7% del costo di pre-training, ottenendo un aumento di velocità del 2 volte con un punteggio FID di 2.86, superando i concorrenti con efficienza comparabile. Il codice è disponibile su 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.