Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

DiffSparse: Accelerating Diffusion Transformers with Learned Token Sparsity

Haowei Zhu, Ji Liu, Ziqiong Liu, Dong Li, Junhai Yong, Bin Wang, Emad Barsoum

Abstract

Diffusion models demonstrate outstanding performance in image generation, but their multi-step inference mechanism requires immense computational cost. Previous works accelerate inference by leveraging layer or token cache techniques to reduce computational cost. However, these methods fail to achieve superior acceleration performance in few-step diffusion transformer models due to inefficient feature caching strategies, manually designed sparsity allocation, and the practice of retaining complete forward computations in several steps in these token cache methods. To tackle these challenges, we propose a differentiable layer-wise sparsity optimization framework for diffusion transformer models, leveraging token caching to reduce token computation costs and enhance acceleration. Our method optimizes layer-wise sparsity allocation in an end-to-end manner through a learnable network combined with a dynamic programming solver. Additionally, our proposed two-stage training strategy eliminates the need for full-step processing in existing methods, further improving efficiency. We conducted extensive experiments on a range of diffusion-transformer models, including DiT-XL/2, PixArt-$α$, FLUX, and Wan2.1. Across these architectures, our method consistently improves efficiency without degrading sample quality. For example, on PixArt-$α$ with 20 sampling steps, we reduce computational cost by $54\%$ while achieving generation metrics that surpass those of the original model, substantially outperforming prior approaches. These results demonstrate that our method delivers large efficiency gains while often improving generation quality.

DiffSparse: Accelerating Diffusion Transformers with Learned Token Sparsity

Abstract

Diffusion models demonstrate outstanding performance in image generation, but their multi-step inference mechanism requires immense computational cost. Previous works accelerate inference by leveraging layer or token cache techniques to reduce computational cost. However, these methods fail to achieve superior acceleration performance in few-step diffusion transformer models due to inefficient feature caching strategies, manually designed sparsity allocation, and the practice of retaining complete forward computations in several steps in these token cache methods. To tackle these challenges, we propose a differentiable layer-wise sparsity optimization framework for diffusion transformer models, leveraging token caching to reduce token computation costs and enhance acceleration. Our method optimizes layer-wise sparsity allocation in an end-to-end manner through a learnable network combined with a dynamic programming solver. Additionally, our proposed two-stage training strategy eliminates the need for full-step processing in existing methods, further improving efficiency. We conducted extensive experiments on a range of diffusion-transformer models, including DiT-XL/2, PixArt-, FLUX, and Wan2.1. Across these architectures, our method consistently improves efficiency without degrading sample quality. For example, on PixArt- with 20 sampling steps, we reduce computational cost by while achieving generation metrics that surpass those of the original model, substantially outperforming prior approaches. These results demonstrate that our method delivers large efficiency gains while often improving generation quality.

Paper Structure

This paper contains 55 sections, 10 equations, 5 figures, 10 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Acceleration of Diffusion Models.
  4. Method
  5. Preliminary
  6. Diffusion Transformer.
  7. Token-Wise Feature Caching Approach.
  8. Challenges in Existing Token Caching Approaches.
  9. DiffSparse Approach
  10. Token Selector.
  11. Learnable Sparsity Cost Predictor.
  12. Dynamic Programming Solver.
  13. Training Loss.
  14. Two-Stage Training Strategy.
  15. Experiments
  16. ...and 40 more sections

Figures (5)

  • Figure 1: DiffSparse uses a learnable sparsity-cost predictor and dynamic programming to learn per-layer sparsity under target ratio $R$. We generate binary masks from the chosen sparsity maps and candidate masks. A token selector reuses features from previous diffusion steps to skip unimportant tokens and speed sampling. To enable gradient flow through the binary masks, we apply Straight-Through Estimation (STE) and train our model using full-step sampling targets with LPIPS loss.
  • Figure 2: Comparison of our method with the baseline (PixArt-$\alpha$ with DPM-Solver++ using 20 steps) and existing methods under different acceleration rates.
  • Figure 3: Visualization of predicted layer sparsity of PixArt-$\alpha$ with 20 steps. In the figure, the x-axis denotes different network layers, the y-axis denotes sampling time steps, and the color gradient from blue to yellow indicates increasing sparsity.
  • Figure 4: Comparison of our method with the baseline (PixArt-$\alpha$ with DPM-Solver++ using 20 steps) under different acceleration rates.
  • Figure 5: Comparison between our DiffSparse, and ToCa with the baseline (PixArt-$\alpha$ with DPM-Solver++ using 20 steps under 512$\times$512 resolution).