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

2026-05-07

AI Summary

Standard diffusion Transformers in video generation apply a uniform denoising schedule across all spatiotemporal tokens, leading to computational redundancy. This work proposes Heterogeneous Step Allocation (HSA), a training-free inference algorithm that adaptively assigns denoising steps to tokens based on their motion dynamics while preserving global contextual consistency through KV cache synchronization and cached Euler updates. As the first online, motion-aware adaptive stepping method, HSA significantly outperforms existing baselines on Wan-2 and LTX-2 models, maintaining high-quality generation even at 50% or 25% of the original runtime and achieving a superior quality-efficiency Pareto frontier.

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

Problem

• Diffusion Transformers
• video generation
• denoising steps
• computational efficiency
• spatiotemporal tokens

Innovation

• Heterogeneous Step Allocation
• Diffusion Transformers
• KV-cache synchronization
• Efficient Video Generation
• Training-free Inference