Source-linked AI summary
Where Compute Matters: Heterogeneous Attention for Efficient Video Diffusion
Olga Zatsarynna, Denis Korzhenkov, Juergen Gall, Amir Habibian, Mohsen Ghafoorian
TL;DR
Video diffusion attention has quadratic cost over long token sequences, while existing efficient methods often allocate compute uniformly despite changing token difficulty. HetA-DiT routes tokens using predicted uncertainty, and on Wan models matches base-model VBench and VBench-2.0 performance while using dense attention for only about 20% of tokens.
Problem
Existing efficient methods often treat video tokens uniformly, although denoising difficulty varies across regions and steps, against a backdrop of quadratic self-attention cost.
Source: §1 · pp. 1–2Method
HetA-DiT predicts token-wise uncertainty and uses an adaptive per-sample, per-timestep threshold to route reliable tokens to local attention and uncertain tokens to dense attention.
Source: §1 · p. 2Results
∼20% of tokens use dense attention while HetA-DiT matches the base models on VBench and VBench-2.0.
Source: §1 · p. 2Takeaways & Limitations
Ablations confirm the importance of uncertainty-based routing, and experiments show quality preservation alongside reduced attention computation.
Source: §6 · p. 10
Abstract
from arXiv · showhide
Efficient video generation requires reducing the quadratic cost of self-attention over long spatio-temporal token sequences. Existing efficient-attention methods typically apply the same computation pattern to every token, even though denoising difficulty varies substantially across video regions and evolves throughout the generation process. We introduce HetA-DiT, a heterogeneous attention mechanism that adaptively allocates computation according to token difficulty. A lightweight uncertainty branch predicts a token-wise estimate of denoising difficulty, which is used to route uncertain tokens through dense global attention while processing more reliable tokens with efficient local attention. The resulting routing is content- and timestep-adaptive, retains global context where it matters most, and provides a single parameter for controlling the quality-efficiency trade-off. HetA-DiT is compatible with few-step distribution-matching distillation and introduces no additional Transformer evaluation at inference time by reusing uncertainty estimates from the preceding denoising step. We evaluate the method on DMD-distilled Wan2.2-5B and Wan2.1-1.3B models. Across VBench, VBench-2.0, and human preference evaluation, HetA-DiT maintains competitive generation quality while routing only approximately 20% of tokens through dense attention.
1 Introduction
Video diffusion attention is costly, and uniform compute allocation does not reflect how token denoising difficulty varies across content and steps. HetA-DiT routes tokens by uncertainty, matching base-model benchmark quality while sending only about 20% of tokens through dense attention.
1 Introduction
- Video DiT self-attention scales quadratically with long spatio-temporal token sequences, while denoising difficulty varies across regions and changes over steps, prompts, and frames.Source: §1 · p. 1
- HetA-DiT uses soft, attention-level, DMD-compatible token routing, with tuning reported at about 12 GPU-days.Source: §1 · p. 2
- ∼20% of tokens use dense attention while HetA-DiT matches the base models on VBench and VBench-2.0 and outperforms random and fixed-pattern routing at matched sparsity.Source: §1 · p. 2
- Existing efficiency methods leave a gap through uniform compute allocation, hard cache-or-compute decisions, or incompatibility with step distillation.Source: §1 · p. 2
- A lightweight uncertainty branch guides an adaptive per-sample, per-timestep threshold that sends reliable tokens to local attention and uncertain tokens to dense global attention.Source: §1 · p. 2
2 Related Work
Related efficiency work reduces denoising steps or attention cost, but generally does not allocate attention according to each token’s denoising difficulty. Caching and FFN-level expert routing also leave limitations for low-step deployment or the attention bottleneck.
2 Related Work
- Step-reduction and deployment-efficiency methods still treat every token identically within a denoising step.Source: §2 · p. 3
- Efficient-attention approaches make attention globally cheaper but apply the same operator to tokens regardless of denoising difficulty.Source: §2 · p. 3
- Feature caching methods such as ToCa [6] and LiteAttention make binary cache-or-recompute decisions, and their reliance on many-step redundancy limits them in low-step settings.Source: §2 · p. 3
- Video MoE DiTs [8] [1] route at the FFN sub-layer, while Wan2.2 A14B at 720p spends 82% of compute on self-attention and 13% on FFN.The authors describe attention-level routing as complementary to FFN-level MoE.Source: §2 · p. 3
3 Method
HetA-DiT allocates attention according to token-level denoising difficulty, using uncertainty estimates to preserve global attention for harder tokens and local attention for more reliable ones. Its two-stage design adds uncertainty prediction to a DMD model, then uses those predictions to route attention during fine-tuning and generation.
3.1.1 Self-Attention.
- Self-attention has quadratic cost in video-token count, while local attention limits each token to a spatio-temporal neighborhood but may lose important global information.The method therefore applies dense attention only where needed.Source: §3.1.1 · p. 4
3.1.2 Distribution Matching Distillation.
- DMD trains a student to match a pretrained diffusion teacher with fewer denoising steps, using teacher-based supervision in video latent space rather than paired ground-truth videos.The paper applies DMD fine-tuning as a starting point for further efficiency improvements.Source: §3.1.2 · pp. 4–5
3.2 Proposed Approach.
- HetA-DiT adds an uncertainty branch to a few-step DMD model and uses its predictions to route tokens between sparse local and full attention.The approach allocates computation per token according to predicted complexity.Source: §3 · p. 3, §3.2 · p. 5
3.2.1 Stage 1: Uncertainty Prediction.
- In Stage 1, an added branch predicts a token-wise uncertainty map alongside the denoised latent, and uncertainty-aware weighting avoids the usual per-sample DMD loss weighting.The loss encourages higher uncertainty for larger prediction errors while preventing arbitrarily large uncertainty estimates.Source: §3.2.1 · p. 5
3.2.2 Stage 2: Uncertainty-Guided Token Densification.
- In Stage 2, the frozen branch’s thresholded predictions form a mask used to route tokens in the following denoising step, and the resulting model is fine-tuned with the standard DMD objective.The uncertainty mask is reused to determine heterogeneous attention in subsequent denoising steps.Source: §3.2.2 · pp. 5–6, Figure 2 · p. 6
- HetA-DiT sends uncertain tokens to full attention and certain tokens to local attention, with experiments using a dense-token ratio of approximately 0.2.The threshold is adjustable, and selected token swaps across transformer blocks broaden full-attention coverage while uncertain tokens still receive it most of the time.Source: §3.2.2 · pp. 5–6, Figure 2 · p. 6
4 Experimental Setup
The experimental setup trains HetA-DiT in two stages and evaluates it on DMD-based Wan models with three inference steps. Evaluation includes benchmark prompts, a method-blinded human preference study, and compute and latency measurements.
Training setup
- Training uses 15k iterations for the first stage and 6k for the second, with an extra full-attention student pass during training to estimate the preceding step’s uncertainty.At inference, uncertainty from the preceding denoising step is reused, so no additional Transformer pass is needed.Source: §4 · p. 7
- Both stages use four H100 GPUs and prompt-only training, with visual supervision from the pretrained generative model rather than ground-truth videos.Training prompts include captions generated for a subset of VIPE1M videos.Source: §4 · p. 7
Evaluation setting
- Evaluation fine-tunes HetA-DiT from Wan2.2 5B and Wan2.1 1.3B DMD models, using three denoising steps and the full extended VBench and VBench-2.0 prompt sets.Generated resolutions are 121×704×1280 for Wan2.2-based models and 81×480×832 for Wan2.1-based models.Source: §4 · p. 7
Human preference study
- A method-blinded human study collected 1,248 paired comparisons from 19 evaluators across 1,082 videos and 667 prompts.Evaluators chose a preferred video or a tie based on semantic coherence and visual quality; video placement was randomized.Source: §4 · pp. 7–8
Compute and latency assessment
- The study measures FLOPs at attention and Transformer-block levels and profiles wall-clock latency using FlexAttention with optimized block-sparse kernels.Latency is measured against dense attention under the same model and resolution settings.Source: §4 · p. 8
5 Results
HetA-DiT maintains strong generation quality against dense and efficient baselines while reducing computation, with results spanning benchmark scores, human preference, latency, and routing ablations.
5 Results
- HetA-DiT performs on par with full Wan models and 3-step DMD on VBench, while outperforming strong efficient video-generation methods.The comparison covers Wan2.2 and Wan2.1 base models.Source: §5 · p. 8, Table 2 · p. 9, Table 3 · p. 9
- On VBench-2.0, HetA-DiT outperforms the full Wan2.2 model, which requires 50 denoising steps, and the step-distilled DMD model, which uses 2 steps.Source: §5 · p. 8
- In the 1,248-pair human preference study, HetA-DiT is preferred almost as often as 3-step DMD and significantly more often than the other two baselines.The compared baselines are DMD with 2 and 3 steps and Attention Surgery.Source: §5 · p. 8, Table 4 · p. 9
- At c=0.9, HetA-DiT achieves 1.73× Transformer-block speedup and a VBench score of 83.28; routing 90% locally saves more compute but lowers the evaluation score.Routing 70% or 80% locally yields comparable performance, with less dense computation producing larger speedups.Source: §5 · p. 9, Figure 4 · p. 10
- With 20% of tokens processed by dense attention on average, HetA-DiT is 1.63× faster than the full DiT model.The reported speedups are particularly larger for high-resolution latents; Table 5 includes amortized BlockMask construction cost.Source: §5 · p. 9, Table 5 · p. 9
5.1 Ablation Studies
- Uncertainty-based routing achieves the best total and semantic scores, consistently outperforming random and fixed-uniform routing at matched sparsity.The local-only variant, which sends no tokens to dense attention, is presented as a lower-bound setting.Source: §5.1 · p. 10, Table 6 · p. 10
- An 11×11×11 local-attention neighborhood offers the best overall trade-off, preserving strong quality while maintaining the highest speedup.A 13×13×13 neighborhood gives the highest total and quality scores but lower speedup; 9×9×9 is more efficient but weaker on semantic performance.Source: §5.1 · p. 10, Table 7 · p. 10
6 Conclusion
HetA-DiT allocates attention computation according to token denoising difficulty, using uncertainty to route tokens between local and dense attention. Wan-model experiments show preserved quality with reduced attention computation, and ablations support uncertainty-based routing.
6 Conclusion
- HetA-DiT uses a lightweight uncertainty branch to route reliable tokens to local attention and uncertain tokens to dense global attention, adapting sparsification to content and timestep.Experiments on Wan models preserve dense-attention baseline quality while substantially reducing attention computation; ablations confirm the importance of uncertainty-based routing.Source: §6 · p. 10
A.1 Training of HetA-DiT.
HetA-DiT training uses two fine-tuning stages within the standard DMD pipeline.
A.1 Training of HetA-DiT
- The first stage uses 15k iterations with per-GPU batch size 2 across 4 GPUs; the second uses 6k iterations with per-GPU batch size 1 across 4 GPUs.Student and critic learning rates are 1 × 10^-5 and 5 × 10^-6; the second-stage timestep shift is s=1, while the first-stage shift is s=5.Source: §A.1 · p. 15
A.2 Training of baselines.
The baseline methods use different implementation and training recipes, including hybrid-block distillation, sparse attention, and limited-step caching.
A.2 Training of baselines
- Attention Surgery replaces 20 transformer blocks with hybrid blocks in a 20×R8 configuration, while ReHyAt replaces all 30 Wan2.2 transformer blocks.Both use three training stages; chunks contain 3 latent frames with one-frame overlap, and their linear-attention branches use degree-2 polynomials.Source: §A.2 · p. 15
- VSA uses sparsity 0.8 and 4 × 4 × 4 spatiotemporal tiles, while Jenga uses sparsity 0.8, threshold probability 0.9, and tile size 64.VSA follows the FastVideo recipe; Jenga is non-trainable and is applied to the baseline DMD checkpoint.Source: §A.2 · p. 15
- ToCa caching is applied only on the middle of three denoising steps, with 60% of FFN tokens recomputed on that cached step.The first and last steps are evaluated fully; the scheduler yields an effective FFN refresh ratio of approximately 89.7% and cache ratio of approximately 10.3%.Source: §A.2 · p. 15
B Additional Results
B Additional Results reports per-dimension comparisons on VBench and VBench-2.0 and an ablation of random token selection. C Additional Qualitative Comparisons provides further comparisons with efficient video-generation baselines.
B Additional Results
- Figures 5 and 6 provide per-dimension comparisons of HetA-DiT with efficient video-generation models on VBench and VBench-2.0.Figure 5 is the VBench comparison; Figure 6 is the VBench-2.0 comparison.Source: B Additional Results · p. 15, Figure 5 · p. 16, Figure 6 · p. 16
- With a 7% random swap, dense-attention exposure reaches approximately 0.89 per token across 30 blocks, and this setting is used as the best balance; fully random selection at 20% degrades performance.The reported coverage probability is 1 − (1 − 0.07)^30 ≈ 0.89; adding randomness to uncertainty-based routing improves results, while purely random selection performs worse.Source: B Additional Results · pp. 15, 17, Table 8 · p. 17
C Additional Qualitative Comparisons
- Figures 7–11 add qualitative comparisons with efficient video-generation baselines across five prompts.The prompts include a Victorian house, a coastal beach, a person playing flute, a giraffe, and a bioluminescent jellyfish.Source: C Additional Qualitative Comparisons · p. 17
D Uncertainty Branch
The uncertainty analysis interprets a token’s prediction uncertainty through signal-to-noise ratio. Under its stated Gaussian assumptions, greater uncertainty corresponds to a noisier effective timestep and motivates assigning difficult tokens more global computation.
D Uncertainty Branch
- The analysis models a token’s prediction as a Gaussian estimate of the clean latent, then re-noises that estimate to derive an effective noise variance and signal-to-noise ratio.It assumes the prediction-noise variables are independent standard Gaussians and matches the resulting ratio to an artificial timestep.Source: D Uncertainty Branch · p. 17
- Under this interpretation, larger uncertainty makes a token behave as if it were at a noisier timestep, supporting more global computation for high-uncertainty tokens.The passage relates this interpretation to SNR matching in pyramidal or multi-scale diffusion methods.Source: D Uncertainty Branch · p. 17
- Figures 7–11 provide qualitative comparisons of HetA-DiT with efficient video-generation models across five prompts.The examples include a Victorian house, a coastal beach, a person playing flute, a giraffe, and a bioluminescent jellyfish.Source: Figure 7 · p. 18, Figure 8 · p. 19, Figure 9 · p. 20
- Figures 12–16 show examples of predicted uncertainty maps at diffusion timestep σ = 0.6.The captions identify the timestep but do not state a qualitative outcome.Source: Figure 12 · p. 23, Figure 13 · p. 23, Figure 14 · p. 23