Not All Tokens Need 40 Steps: Heterogeneous Step Allocation in Diffusion Transformers for Efficient Video Generation

Abstract: Diffusion Transformers (DiTs) have achieved state-of-the-art video generation quality, but they incur immense computational cost because standard inference applies the same number of denoising steps uniformly to every token in the sequence. It is well known that human vision ignores vast amounts of redundant motion. Why, then, do our densest models treat every spatiotemporal token with equal priority? In this paper, we introduce Heterogeneous Step Allocation (HSA), a training-free inference algorithm that assigns varying step budgets to different spatiotemporal tokens based on their velocity dynamics. To resolve the resulting sequence-length mismatch without sacrificing global context, HSA introduces a KV-cache synchronization mechanism that allows active tokens to attend to the full sequence while entirely bypassing inactive tokens. Furthermore, we derive a cached Euler update that advances the latent states of skipped tokens in a single operation without additional model evaluations. We evaluate HSA on the Wan-2 and LTX-2 models for both text-to-video (T2V) and image-to-video (I2V) generation. Our results demonstrate that HSA significantly outperforms previous state-of-the-art caching methods and the vanilla Flow Matching baseline, especially at aggressive acceleration regimes (e.g., 50% and 25% runtimes). Crucially, HSA achieves a superior quality-runtime Pareto frontier without the need for expensive offline profiling, robustly preserving structural integrity and generation quality even under tight computational budgets. Project page: https://ernestchu.github.io/hsa

• The paper introduces a novel method that assigns heterogeneous denoising steps per token to efficiently generate high-fidelity videos.
• It leverages adaptive token grouping based on velocity dynamics and lightweight KV-cache synchronization to reduce compute costs.
• Empirical results on Wan-2.1 and Wan-2.2 DiTs demonstrate superior VBench scores while preserving semantic and structural integrity.

Heterogeneous Step Allocation for Efficient Diffusion Transformer-Based Video Generation

Motivation and Problem Statement

Diffusion Transformers (DiTs) have become foundational for high-fidelity video generation, leveraging self-attention for global context aggregation and iterative denoising for generative expressivity. However, the standard inference protocol uniformly applies the same number of denoising steps to every spatiotemporal token. This approach is computationally inefficient, as visual data inherently exhibits high spatial and temporal asymmetry: static backgrounds and coarse regions require significantly fewer updates, while dynamic, detail-rich foregrounds demand higher fidelity. Prior caching accelerations operate globally—skipping whole steps or blocks—failing to address spatial redundancy and the heterogeneity of per-token dynamics.

Heterogeneous Step Allocation: Methodological Advances

The paper introduces Heterogeneous Step Allocation (HSA), a training-free, model-agnostic inference algorithm for DiTs. HSA partitions the token sequence into groups based on velocity dynamics and assigns each group a distinct step budget (a divisor of the total number of steps, $T$). Tokens are denoised on heterogeneous schedules: frequent for salient, high-motion regions and infrequent for static, background regions. A subset of baseline tokens is always processed fully, anchoring trajectory and ensuring consistent global structure.

Synchronization across heterogeneous schedules poses nontrivial challenges due to sequence-length mismatch and varied noise levels. HSA resolves this via a lightweight KV-cache synchronization scheme: at each iteration, only active tokens (those scheduled for an update) perform query, key, and value projections, whereas inactive tokens have their cached representations reused. Active tokens attend against the full sequence—both fresh and cached—preserving global context with reduced attention complexity per iteration ($O(|A_i|\cdot N)$).

For latent updates, HSA employs a cached Euler step: inactive tokens are advanced using their most recently predicted velocity, avoiding unnecessary model calls. This results in a unified tensor operation, streamlining GPU compute and obviating irregular control flow.

Phase-aware caching restricts heterogeneously allocated steps to the central portion of the denoising trajectory, running full passes in critical early and late steps to avoid quality loss induced by stale KV features.

Experimental Evaluation

Quantitative evaluation is conducted primarily on Wan-2.1-1.3B and Wan-2.2-A14B DiTs for both text-to-video (T2V) and image-to-video (I2V) tasks. The VBench metric suite, aggregating quality and semantic alignment scores across sixteen dimensions, serves as the primary indicator.

HSA achieves a superior quality-runtime Pareto frontier when compared to uniform Flow Matching and the state-of-the-art TeaCache. The efficiency gains are most pronounced at aggressive runtime targets (i.e., 25% and 50% of the reference compute), where competing methods exhibit catastrophic dimension collapse while HSA robustly tracks all reference dimensions. Specifically, HSA outperforms baselines in VBench scores (e.g., at 50% runtime, HSA achieves 82.79 vs. 81.58 for FM and TC), preserves structural and semantic integrity, and maintains high image-conditioning fidelity in I2V tasks.

Ablation studies demonstrate that dynamic token selection (based on velocity L1 change) consistently yields higher quality than uniform, random, or spatially anchored allocations. The flexible grouping paradigm allows adaptation to diverse video modalities and model scales, with qualitative results supporting robustness even under tight budgets.

Relation to Prior Work

Step-level caching methods accelerate DiTs by globally skipping steps and reusing features, relying on offline profiling or adaptive sensitivity indicators. These include TeaCache, MagCache, EasyCache, and others, but all apply cache decisions uniformly across the token space. Attention-block and intra-step caching explore spatial redundancy but still do not assign explicit per-token step budgets. Recent works like X-Slim and HetCache incorporate spatial and block-level axes, but their token-level scheduling is reactive and lacks fixed per-token allocation.

HSA introduces an unexplored axis: explicit, deterministic, per-token heterogeneous step assignment. The combination of per-token schedules and full-context attention via KV-cache synchronization is not addressed by previous literature.

Implications and Future Directions

Practically, HSA significantly mitigates the computational bottleneck of large-scale DiTs, enabling efficient high-quality video generation on lower compute budgets. This democratizes content creation, reducing environmental impact and broadening accessibility.

Theoretically, HSA formalizes the intuition that generative processes benefit from spatially adaptive resource allocation. The framework is extensible: advanced token grouping strategies informed by perceptual or semantic priors could further improve quality and efficiency. The mechanisms developed—KV-cache synchronization and cached Euler updates—lay groundwork for broader exploration of heterogeneous scheduling in generative modeling.

Future work may consider joint optimization of token grouping with content-aware priors, integration with hierarchical diffusion schedules, or adaptation to diverse modalities beyond video (e.g., multimodal audio-visual synthesis). Distributional benchmarking (VBench, VBench++) should remain central, as per-sample metrics become less informative when alternative high-quality basins are reached.

Conclusion

Heterogeneous Step Allocation introduces a principled, training-free approach for inferring DiTs with spatially adaptive schedules, leveraging dynamic token partitioning, efficient synchronization, and cached updates to maximize the quality-runtime frontier. The empirical results substantiate its advantage, particularly under constrained compute. By embracing the intrinsic redundancy and heterogeneity in visual data, HSA sets a precedent for future generative modeling paradigms pursuing efficiency and perceptual alignment.