TinyFusion: Diffusion Transformers Learned Shallow (2412.01199v1)

Abstract: 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 2$\times$ speedup with an FID score of 2.86, outperforming competitors with comparable efficiency. Code is available at https://github.com/VainF/TinyFusion.

• The paper proposes a learnable depth pruning method that intelligently removes redundant transformer layers to enhance inference speed.
• It co-optimizes differentiable layer masks with weight updates to explicitly model post-fine-tuning recoverability in diffusion transformers.
• Empirical results on architectures like DiT-XL show a 2× speedup and competitive FID scores, outperforming traditional compression techniques.

TinyFusion: Optimizing Diffusion Transformer Models through Learnable Depth Pruning

The architectural framework of diffusion transformers has solidified its position as an influential element in generative tasks, notably in image synthesis. Despite their capabilities, their significant parameter requirements often serve as roadblocks in practical, real-world applications because they induce substantial inference overheads. The research paper "TinyFusion: Diffusion Transformers Learned Shallow" presents an innovative method aimed at reducing this overhead through a novel depth pruning approach. TinyFusion proposes a systematic pruning strategy that intelligently removes redundant layers within diffusion transformer models while ensuring that the models retain high recoverability post-pruning.

Novel Contributions

TinyFusion introduces a learnable paradigm that involves differentiable sampling of layer masks. By co-optimizing these masks alongside weight updates aimed at predicting potential fine-tuning results, the approach shifts the pruning focus from immediate minimization of loss to explicitly modeling the post-fine-tuning performance of diffusion transformers. In contrast with prior techniques that predominantly rely on heuristic or empirical removal of seemingly redundant parameters, this method adopts a more prognostic standpoint.

Experimental Results

Empirical evaluations substantiate the effectiveness of TinyFusion across diverse architectures, namely DiTs, MARs, and SiTs. A standout experiment with the DiT-XL architecture indicates that TinyFusion managed to achieve a 2× speedup at less than 7% of the pre-training cost, coupled with an FID score of 2.86, outperforming competitive compression strategies. The pruned models demonstrate substantial merit in terms of computational efficiency while maintaining core performance parameters when subjected through extensive testing.

Implications and Claims

The research contends that traditional strategies focusing purely on obtaining models with low calibration loss after pruning may not sufficiently cater to diffusion transformers. Instead, TinyFusion builds on the insight that the structural recoverability of a model, in terms of maintaining performance upon supplementary fine-tuning, is an essential consideration. As such, this dynamic introduces a novel dimension into model compression—targeting and enhancing the recoverability trait in pruned networks.

Theoretical and Practical Implications

TinyFusion formulates a blueprint for a profound shift in how depth pruning might transition into a predictive, performance-oriented task. This paradigm can alleviate computational demands markedly, potentially leading to transformation in deployment scales of diffusion models in consumer-oriented applications that accept constraints on hardware capabilities. The practicality of layer pruning, demonstrated through linear acceleration strategics in practical setups, invites further exploration of enhanced granularity in pruning mechanisms, possibly extending into domain-specific adaptations within transformer applications.

Future Prospects

One avenue worth pursuing further could involve evaluating the potential of enhanced pruning mechanisms on other generative tasks, such as text-to-image synthesis, where the relevance of specific transformer layers may diverge significantly. In tandem, the exploration of different model architectures, potentially optimizing attention and MLP layers distinctly, may yield further insights that enhance the flexibility and efficacy of TinyFusion’s methodologies.

In summary, TinyFusion makes a substantial contribution to the dialogue on model efficiency, proposing a meaningful framework that challenges and extends current paradigms. This work lays a foundation for future endeavors aimed at propelling diffusion transformers toward greater applicability in a variety of computational environments.