Technical Summary: Heterogeneous Step Allocation (HSA) for Efficient Video Generation

1. Problem Statement

Diffusion Transformers (DiTs) have established state-of-the-art performance in video generation but suffer from immense computational costs. Standard inference protocols apply a uniform number of denoising steps ($T$) to every spatiotemporal token in the sequence. This monolithic approach ignores the inherent asymmetry of visual data: background regions, static patches, and coarse structures are perceptually simpler and require fewer denoising steps than detail-rich foreground objects or high-motion regions. Furthermore, human vision is less sensitive to quality degradation in highly dynamic components compared to static ones. Consequently, applying a full denoising budget uniformly to all tokens is computationally wasteful and fails to align with perceptual priorities.

Existing acceleration methods, such as step-level feature caching (e.g., TeaCache), typically make global decisions: at any given iteration, either all tokens are computed or all tokens reuse cached outputs. These methods do not exploit spatial or token-level heterogeneity, treating the sequence as a monolithic block rather than a collection of tokens with varying dynamic requirements.

2. Methodology: Heterogeneous Step Allocation (HSA)

The authors propose Heterogeneous Step Allocation (HSA), a training-free, model-agnostic inference algorithm that assigns token-specific denoising trajectories within a pretrained DiT. The core idea is to partition the token sequence into groups based on their velocity dynamics, allocating fewer steps to slower-evolving tokens while maintaining a full schedule for a baseline group.

2.1 Token Partitioning and Scheduling

The set of $N$ tokens is partitioned into $K$ disjoint groups ($G_1, \dots, G_K$). Each group $G_k$ is assigned a step budget $s_k$, which must be a divisor of the total steps $T$.

• Baseline Group ($G_1$): Assigned the full budget $s_1 = T$. These tokens are active at every iteration.
• Reduced Groups: Assigned smaller budgets (e.g., $T/2, T/4$). These tokens are active only when the global iteration $i$ satisfies $i \pmod{T/s_k} = 0$.
• Dynamic Selection: The paper evaluates four strategies for assigning tokens to groups. The Dynamic Token Selection strategy, which assigns tokens based on the per-token L1 relative change of velocity predictions (identifying slow-evolving regions like backgrounds), consistently yields the best performance.

2.2 KV-Cache Synchronization

A primary challenge in heterogeneous scheduling is that tokens with different step counts exist at different noise levels ($\sigma$) at any given wall-clock iteration. Self-attention requires all tokens to attend to one another, creating a synchronization mismatch. HSA resolves this via a KV-cache synchronization mechanism:

1. Active Tokens: Only the subset of active tokens ($A_i$) performs a forward pass through the QKV projections.
2. Cache Update: Fresh Key ($K$) and Value ($V$) projections for active tokens are written into a per-layer cache, overwriting stale entries.
3. Full-Context Attention: Active tokens compute queries against the full cached $K, V$ matrices (containing both fresh entries for active tokens and stale-but-valid entries for skipped tokens).
4. Complexity Reduction: This reduces the self-attention complexity from $O(N^2)$ to $O(|A_i| \cdot N)$, where $|A_i|$ is the number of active tokens at step $i$. Skipped tokens incur no query computation, cross-attention, or feed-forward operations during inactive iterations.

2.3 Cached Euler Update

To advance the latent states of skipped tokens without additional model evaluations, HSA employs a cached Euler update.

• Each token stores the velocity ($\hat{v}_n$) predicted at its most recent active step.
• At every global iteration, all tokens (active and skipped) are advanced by the same incremental Euler step $(\sigma_{i+1} - \sigma_i)$.
• Active tokens use their freshly computed velocity; skipped tokens reuse their cached velocity.
• This update is formulated as a single tensor operation over all $N$ tokens, avoiding branching and maintaining GPU efficiency.

2.4 Phase-Aware Caching Window

Recognizing that early and late denoising stages are critical for global structure and fine details respectively, HSA restricts caching to the middle portion of the trajectory. The first and last $m$ steps execute full, uncached forward passes for all tokens, ensuring structural integrity and detail refinement are not compromised by stale approximations.

3. Key Contributions

1. HSA Algorithm: A fully training-free, plug-and-play inference paradigm that assigns heterogeneous denoising trajectories to spatiotemporal token groups based on velocity dynamics.
2. KV-Cache Synchronization: A mechanism that preserves the full $N$-token attention context for active tokens while entirely bypassing inactive tokens, reducing per-iteration attention cost to $O(|A_i| \cdot N)$.
3. Cached Euler Update: A simple tensor operation that advances skipped tokens using cached velocities, maintaining generation quality without extra model evaluations.
4. Comprehensive Evaluation: An analysis of four token-grouping strategies and four budget presets, demonstrating robustness across different inference budgets.

4. Experimental Results

The authors evaluated HSA on Wan-2.1-1.3B and Wan-2.2-A14B (Text-to-Video and Image-to-Video) and LTX-2 (Audio-Video).

• Quality-Runtime Trade-off: HSA achieves a superior Pareto frontier compared to uniform Flow Matching (FM) and state-of-the-art caching methods like TeaCache.
• Aggressive Acceleration: The advantages are most pronounced at aggressive acceleration regimes (50% and 25% runtime). At 25% runtime, HSA significantly outperforms baselines on the VBench benchmark (Total Score: 79.87% vs. 75.68% for FM and 78.33% for TeaCache).
• Dimensional Robustness: While baseline methods suffer from "catastrophic dimension collapse" (severe degradation in specific VBench dimensions like motion smoothness or temporal consistency) under tight budgets, HSA robustly tracks the full-budget reference across all 16 evaluation dimensions.
• Image-to-Video (I2V): HSA successfully preserves image-conditioning alignment and visual fidelity in I2V tasks, whereas uniform FM baselines at reduced steps ($T=10$) exhibit severe structural collapse.
• No Offline Profiling: Unlike TeaCache, which requires expensive offline profiling to estimate model output fluctuations, HSA operates without any offline tuning or fine-tuning.

5. Significance and Claims

The paper claims that HSA addresses a fundamental inefficiency in DiT inference by aligning computational resources with the perceptual and dynamic complexity of visual data. By treating tokens individually rather than globally, HSA achieves significant speedups without the quality degradation typically associated with aggressive acceleration.

The authors emphasize that HSA is a flexible framework that can be instantiated with various token-grouping strategies. While dynamic selection currently yields the strongest results, the framework anticipates further improvements with more sophisticated grouping strategies. The work lays the groundwork for a new class of inference algorithms that leverage the inherent heterogeneity of spatiotemporal data to democratize high-fidelity video generation.