Winzige Fusion: Flach gelernte Diffusions-Transformer

TinyFusion: Diffusion Transformers Learned Shallow

December 2, 2024

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

cs.AI

Zusammenfassung

Diffusions-Transformer haben bemerkenswerte Fähigkeiten in der Bildgenerierung gezeigt, kommen jedoch oft mit übermäßiger Parametrisierung, was zu erheblichem Inferenzaufwand in realen Anwendungen führt. In dieser Arbeit präsentieren wir TinyFusion, eine Tiefenpruningsmethode, die darauf abzielt, überflüssige Schichten aus Diffusions-Transformern durch End-to-End-Lernen zu entfernen. Das Kernprinzip unseres Ansatzes besteht darin, ein beschnittenes Modell mit hoher Wiederherstellbarkeit zu erstellen, das es ermöglicht, nach Feinabstimmung starke Leistungen wieder zu erlangen. Um dies zu erreichen, führen wir eine differenzierbare Abtasttechnik ein, um das Pruning erlernbar zu machen, gepaart mit einem ko-optimierten Parameter zur Simulation zukünftiger Feinabstimmung. Während frühere Arbeiten darauf abzielen, den Verlust oder Fehler nach dem Pruning zu minimieren, modelliert unsere Methode explizit und optimiert die Leistung von beschnittenen Modellen nach der Feinabstimmung. Experimentelle Ergebnisse deuten darauf hin, dass dieses erlernbare Paradigma erhebliche Vorteile für das Schichtenpruning von Diffusions-Transformern bietet und bestehende bedeutungsbasierte und fehlerbasierte Methoden übertrifft. Darüber hinaus zeigt TinyFusion eine starke Verallgemeinerung über verschiedene Architekturen wie DiTs, MARs und SiTs. Experimente mit DiT-XL zeigen, dass TinyFusion einen flachen Diffusions-Transformer zu weniger als 7% der Vor-Trainingskosten erstellen kann, was zu einer 2-fachen Beschleunigung mit einem FID-Score von 2.86 führt und Wettbewerber mit vergleichbarer Effizienz übertrifft. Der Code ist verfügbar unter 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.