💻 computer science

DiffSparse: Accelerating Diffusion Transformers with Learned Token Sparsity

DiffSparse introduces a differentiable, end-to-end framework that optimizes layer-wise token sparsity and employs a two-stage training strategy to significantly accelerate diffusion transformers while maintaining or even improving generation quality.

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

Published 2026-04-07

📖 4 min read☕ Coffee break read

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a master chef trying to bake a perfect, complex cake (an image) from a bowl of raw ingredients (random noise). To get the cake just right, you have to taste the batter, adjust the sugar, mix it again, taste it again, and repeat this process about 20 times before it's ready. This is how Diffusion Models work: they take a step-by-step approach to create images, but it's incredibly slow and uses a lot of computer power, like running a marathon just to make a sandwich.

The paper introduces a new method called DiffSparse. Think of it as a smart sous-chef that helps the master chef bake the cake much faster without ruining the taste.

Here is how DiffSparse works, broken down into simple concepts:

1. The Problem: "Re-doing Work You Already Did"

In the current way of baking these cakes, the chef tastes the batter at every single step. But here's the catch: between step 10 and step 11, the batter doesn't change that much. The chef is wasting time re-tasting and re-mixing parts of the batter that are already perfect.

Previous attempts to speed this up tried to "skip" some steps or save some notes (caching), but they were clumsy. They were like a chef who decides, "Okay, I'll skip tasting the sugar for the next 3 steps," but they do it randomly or based on a rigid rule. Sometimes they skip too much, and the cake tastes weird. Sometimes they don't skip enough, and it's still slow.

2. The Solution: The "Smart Skip" (Token Sparsity)

DiffSparse is like giving the chef a smart assistant who knows exactly which parts of the batter need attention and which parts are fine to leave alone.

Tokens are like ingredients: In the computer world, an image is broken down into tiny pieces called "tokens."
The Assistant's Job: Instead of the chef re-mixing every ingredient at every step, the assistant looks at the batter and says, "Hey, this specific patch of chocolate chips hasn't changed since step 10. Let's just use the old memory of it and skip re-mixing it."
The Result: The chef only mixes the parts that actually need changing. This saves a massive amount of energy (computing power).

3. How the Assistant Learns (The "Brain")

The magic of DiffSparse isn't just that it skips things; it's how it learns to skip them.

No More Guessing: Old methods required humans to write rules like "Skip 20% of the mixing in step 5." DiffSparse learns the rules itself. It's like a student who watches the chef bake 10,000 cakes and figures out, "Oh, I see a pattern! The flour needs constant attention, but the vanilla extract is stable after step 3."
The Dynamic Planner: The system uses a mathematical "planner" (called Dynamic Programming) to look at the whole baking process at once. It asks, "If I save energy here, can I afford to spend a little more energy there to make sure the cake is perfect?" It balances the budget perfectly to ensure the cake (image) looks amazing even though we did less work.

4. The Two-Stage Training (The "Practice Run")

To make sure the assistant doesn't get lazy and ruin the cake, the researchers use a special two-step training method:

Stage 1: The assistant practices on a "full effort" version of the cake to learn what a perfect cake looks like.
Stage 2: The assistant starts practicing the "skipping" technique, but it's constantly checked against the perfect cake to make sure it hasn't forgotten how to bake.

The Results: Faster, Better, and Smarter

The paper tested this on several famous image generators (like PixArt and FLUX). The results were impressive:

Speed: They cut the computing work by nearly 50% (almost half the time!).
Quality: Surprisingly, the images weren't just "okay"; they were often better than the original slow versions. By focusing the computer's power only on the parts that needed it, the final image ended up sharper and more detailed.

In a Nutshell

Imagine you are walking through a museum.

Old Way: You stop and stare at every single painting for 10 seconds, even the ones you've seen a hundred times.
DiffSparse: You have a guide who whispers, "You know that painting? You've seen it. Keep walking. But stop and look closely at this new one."

DiffSparse is that smart guide. It teaches AI models to stop wasting time on things that haven't changed, allowing them to create beautiful images in a fraction of the time, without sacrificing quality. It turns a slow, exhausting marathon into a fast, efficient sprint.

1. Problem Statement

Diffusion models, particularly those based on Transformer architectures (Diffusion Transformers or DiTs), have achieved state-of-the-art performance in image and video generation. However, their multi-step inference mechanism incurs immense computational costs, hindering real-world deployment.

While previous acceleration methods exist, they suffer from specific limitations:

Inefficient Caching Strategies: Existing token caching methods (e.g., ToCa, DuCa) often rely on manually designed sparsity allocations and fixed schedules.
Redundant Full-Step Computation: To maintain generation quality, current methods often require preserving several "full-step" forward passes (where no tokens are skipped) during the denoising process. This limits the potential speedup of token-level caching.
Lack of Adaptability: Manual heuristics fail to dynamically optimize sparsity across different layers and timesteps, leading to suboptimal trade-offs between speed and quality.

2. Methodology: DiffSparse

The authors propose DiffSparse, a differentiable, end-to-end framework that optimizes layer-wise token sparsity in diffusion transformers. The method replaces manual heuristics with a learned approach comprising three core components:

A. Token Selector

Instead of random or fixed token selection, DiffSparse employs a Token Selector that assigns an importance score to each token based on a composite metric:

Self-Attention Influence: How much a token affects others.
Cross-Attention Focus: How strongly the token is guided by conditioning signals.
Cache Reuse Frequency: How often a token has been reused (high frequency may indicate accumulated error).
Spatial Bonus: Ensures spatially uniform coverage to prevent clustering.
Tokens are sorted by score, and a specific proportion (determined by the sparsity ratio) is selected for recomputation, while the rest reuse cached features.

B. Learnable Sparsity Cost Predictor

To determine which layers and timesteps should be sparse, the authors introduce a Learnable Sparsity Cost Predictor.

Input: It takes the current timestep and layer index.
Output: A cost matrix $C \in \mathbb{R}^{(T \times L) \times |S|}$ , where $T$ is timesteps, $L$ is layers, and $|S|$ is a set of candidate sparsity rates (e.g., 0%, 25%, 50%, 75%, 100%).
Function: It predicts the "cost" (degradation in generation quality) of applying a specific sparsity rate at a specific layer and timestep.
Scalability: The predictor size depends only on $T, L,$ and $|S|$ , not the token sequence length, allowing it to generalize to higher resolutions without retraining.

C. Dynamic Programming (DP) Solver

The framework formulates sparsity allocation as a constrained optimization problem.

Objective: Minimize the total sparsity cost (predicted by the Cost Predictor) while satisfying a global target sparsity ratio $R$ .
Algorithm: A DP solver iteratively computes the optimal sparsity configuration for all layers and timesteps. It finds the path through the cost matrix that minimizes cumulative error while meeting the pruning constraint.
Differentiability: Since the DP solver involves discrete decisions, the authors use the Straight-Through Estimator (STE) to allow gradients to flow back to the Cost Predictor during training.

D. Two-Stage Training Strategy

To overcome the reliance on full-step computations found in prior works, DiffSparse uses a novel two-stage training protocol:

Stage 1 (Warm-up): The model is trained with a preset number of full-step positions. The system learns the temporal cost of these full steps and the layer-wise sparsity costs independently.
Stage 2 (Refinement): The learned costs from Stage 1 are integrated. The model is fine-tuned to dynamically redistribute computation, effectively learning to replace full steps with sparse steps where the cost is low, thereby eliminating the need for rigid full-step schedules.

Training Loss: The framework uses LPIPS (Learned Perceptual Image Patch Similarity) loss as a perceptual distillation loss. The original full-model acts as the "teacher," and the pruned model as the "student," ensuring the accelerated model mimics the perceptual quality of the original.

3. Key Contributions

Differentiable Sparsity Optimization: A novel framework that automates layer-wise and timestep-wise sparsity allocation without manual tuning, using a learnable cost predictor and DP solver.
Elimination of Full-Step Dependencies: The two-stage training strategy removes the need for predefined full-step computations, unlocking the full acceleration potential of token caching.
Generalization: The learned sparsity predictor trained at lower resolutions (e.g., 256x256) transfers effectively to higher resolutions (e.g., 512x512) without retraining.
Comprehensive Evaluation: Extensive experiments across diverse foundation models (PixArt-α, DiT-XL/2, FLUX, Wan2.1) and modalities (text-to-image, class-conditional, text-to-video).

4. Experimental Results

The authors evaluated DiffSparse on multiple benchmarks, demonstrating superior efficiency and quality compared to State-of-the-Art (SOTA) methods like ToCa, TaylorSeer, and DeepCache.

PixArt-α (Text-to-Image):
- At 1.91× speedup (54% token pruning), DiffSparse achieved an FID of 27.79, outperforming ToCa (28.35) and TaylorSeer (29.08).
- Notably, the accelerated model achieved better FID scores than the original full model (28.20), suggesting the learned sparsity schedule improves convergence.
DiT-XL/2 (Class-Conditional):
- At 2.07× speedup, DiffSparse achieved an FID of 2.81, significantly better than ToCa (3.05) and DuCa (3.04).
Wan2.1 (Text-to-Video):
- Achieved the highest VBench score (43.83) with a 2.05× speedup, outperforming DuCa (43.56 at 1.69× speedup).
FLUX.1-schnell (Distilled Model):
- Even on 4-step distilled models, DiffSparse provided a 1.81× speedup with no quality degradation (Image Reward increased from 1.063 to 1.184).

5. Significance

DiffSparse represents a significant shift in diffusion model acceleration:

From Heuristic to Learned: It moves away from manual, rule-based sparsity schedules to a data-driven, differentiable optimization process.
Quality-Speed Trade-off: It breaks the traditional trade-off where acceleration leads to quality loss; in many cases, it improves generation quality by optimizing the distribution of computational resources.
Practical Deployment: By reducing the need for full-step computations and enabling cross-resolution generalization, DiffSparse makes high-quality generative AI more accessible for edge devices and real-time applications.

In summary, DiffSparse provides a robust, learnable framework that maximizes the efficiency of Diffusion Transformers by intelligently pruning redundant token computations while preserving, and often enhancing, the generative quality of the output.