ArXiv: 2505.13389
🎯 Pitch
Standard video diffusion models waste enormous compute on near-zero attention weights. This paper shows you can slash training FLOPs by 2.5× and cut generation time nearly in half with zero quality loss by learning which tokens to attend to instead of computing the full 3D attention matrix, even scaling to 14B-parameter models.
1. Executive Summary
This paper introduces VSA (Video Sparse Attention), a trainable, hardware-efficient sparse attention mechanism that replaces full 3D attention in video diffusion transformers at both training and inference time. Through extensive pretraining experiments scaling DiTs from 60M to 1.4B parameters on video latents with up to 16K sequence lengths, VSA employs a hierarchical two-stage design—a lightweight coarse stage that pools tokens into cubes to identify high-weight critical tokens via learned Top-𝒦 selection, and a fine stage that computes token-level attention only within those selected cubes under a block-sparse compute layout—achieving a 2.53× reduction in total training FLOPS with no drop in diffusion loss. When retrofitted into the open-source Wan2.1-1.3B model, VSA reduces attention time by 6× and end-to-end generation latency from 31s to 18s (1.7×), while scaling-law experiments across model sizes establish that trainable sparse attention produces a strictly better Pareto frontier than full attention, demonstrating that learned, data-dependent sparsity can outperform fixed-pattern heuristics only when the sparsity pattern is jointly optimized with the model parameters rather than applied post-hoc.
2. Context and Motivation
The Core Problem: Quadratic Attention Makes Video Generation Prohibitively Expensive
The fundamental problem this paper addresses is that attention computation is the primary bottleneck in scaling video Diffusion Transformers (DiTs), and this bottleneck exists at both training and inference time. Unlike language models, where long-context adaptation is often a small fraction of total training, video generation inherently requires processing extremely long sequences. As the paper notes in Section 1, "even a seemingly short 5-second 720p clip unfolds into more than 100K tokens once flattened as a sequence." This is not an edge case—it is the standard operating regime for state-of-the-art video models.
The cost structure is punishing: a standard 3D full attention operation over a video latent of shape has complexity . For a 5-second 720p video, this means computing attention over 100K tokens at every transformer layer, making attention dominate the total compute budget at both training and inference. The paper states that "state-of-the-art video DiTs expend the majority of compute on attention when training on full-resolution, long-sequence data; the trained DiTs remain painfully slow at inference." This is not a minor inefficiency—it is a structural barrier that limits how large video models can be, how long videos they can generate, and how accessible video generation technology can be to researchers and practitioners without massive compute clusters.
Why This Problem Matters: Three Practical Consequences
The attention bottleneck has cascading effects that the paper implicitly identifies across its experiments:
1. Training cost becomes prohibitive. When training video DiTs from scratch, the quadratic attention cost means that scaling to longer videos, higher resolutions, or larger models requires disproportionate increases in compute. The paper's own experiments operate around to FLOPS—already enormous budgets—and these are for models up to 1.4B parameters generating ~5-second videos. Extending to 10-second, 1080p, or larger models would push costs into regimes accessible to very few organizations. Without attention sparsity, the scaling trajectory is economically unsustainable.
2. Inference latency blocks real-time and interactive applications. The paper's profiling of Wan2.1-1.3B reveals a stark reality: full attention with torch.compile takes 31 seconds per generation on an H100 GPU. Even the 14B model takes 1,274 seconds (over 21 minutes). These latencies make interactive video generation—where a user iterates on prompts and receives rapid feedback—fundamentally impossible. Any technology that can cut these times by a factor of 2–6×, as VSA does, directly expands the set of applications video generation can serve.
3. The train-test mismatch in post-hoc methods caps quality. As we'll discuss below, most prior approaches apply sparsity only at inference time. This creates a fundamental problem: the model learns its parameters under the assumption of full attention (seeing all tokens), but at inference time, it operates under a restricted attention pattern (seeing only a subset). The paper explicitly frames this as a quality ceiling: "that mismatch caps best-case quality at the dense model's ceiling and, in practice, often erodes quality once sparsity is pushed beyond a gentle budget." A trainable approach that applies sparsity during training eliminates this ceiling.
Prior Approaches and Where They Fall Short
The paper organizes prior work along two dimensions: post-hoc inference-time sparsity (the dominant approach) and the emerging category of trainable sparse attention.
Post-Hoc Inference-Time Sparsity
The most common approach in video DiTs has been to train with full attention and then substitute a sparse pattern only at inference. The paper cites several representative methods:
- Sliding Tile Attention (STA) (Zhang et al., 2025; the authors' own prior work) and Sparge Attention (Zhang et al., 2025) apply fixed or profile-derived sparse masks to pretrained models at inference time.
- Sparse VideoGen (SVG) (Xi et al., 2025) uses a training-free approach to identify and exploit sparsity in the attention matrix.
These methods share a critical limitation that the paper identifies through its scaling experiments in Table 1(a). The authors train models with existing sparse methods (spatial-temporal attention, strided window, compressed KV) and compare against full attention. The key finding: "existing sparse methods outperform full attention with a compute-optimal training budget ( FLOPS), but this advantage reverses with extended training ( FLOPS)." In plain language: fixed-pattern sparse methods look good when training is limited, but when you actually train to convergence, full attention catches up and surpasses them. This is a damning result—it means fixed-pattern sparsity fundamentally limits model capacity, and the gap widens as you invest more compute.
Why does this happen? The paper doesn't explicitly theorize, but the reasoning is implicit in the architecture: fixed patterns (like attending only to spatial neighbors or the same temporal position) are a form of attention capacity constraint. They prevent the model from learning long-range dependencies that turn out to be important for video quality once the model has enough capacity and training to exploit them. A spatial-temporal pattern, for example, forces each token to either attend spatially (within its frame) or temporally (across frames at the same location), but never to attend to a token that is both far away spatially and at a different time—yet such cross-spatiotemporal dependencies (e.g., tracking an object that moves across the frame) are precisely what video generation requires.
The Emergence of Trainable Sparse Attention in LLMs
The paper draws explicit inspiration from recent work in large language models that explores trainable, dynamic sparsity patterns:
- MoBA (Lu et al., 2025) uses a gating mechanism with mean-pooled representations to guide block selection for sparse attention, but discards the pooled attention output and relies on token gathering with variable-length FlashAttention, constraining it to larger tile sizes.
- NSA (Yuan et al., 2025) uses a two-stage architecture (compress + select) but is tailored for causal language model decoding with single-query constraints, requiring grouped query attention and an additional sliding window stage.
- BiFormer (Zhu et al., 2023) uses coarse-grained tile-to-tile attention for vision transformers, but only for deriving the sparse pattern—the coarse output is discarded.
These methods share a common insight: dynamically predict which tokens to attend to using a lightweight learned mechanism, rather than imposing a fixed pattern. However, none directly transfer to video DiTs because of fundamental architectural differences. Video DiTs use bidirectional attention (not causal), process entire sequences simultaneously (not token-by-token decoding), and must handle 3D spatiotemporal structures (not 1D sequences). The paper argues, and demonstrates through ablation, that these differences matter: "the nature of video and bidirectional attention avoids grouped query constraints of the attention pattern" (Section 4), and including the coarse stage output directly in the final representation—which MoBA and BiFormer do not do—is critical for performance.
The Urgency Gap: Why Video DiTs Need Trainable Sparsity More Than LLMs
The paper makes an argument in Section 4 that is worth unpacking because it distinguishes VSA's motivation from similar work in language models:
"the case for trainable sparse attention in video DiTs is both distinct from and more urgent than in LLMs."
Three structural differences support this claim:
-
Sequence lengths are inherently far longer. A 100K-token LLM context is considered extreme and specialized (used for long-document QA, code repository understanding). A 100K-token video context is everyday—it's just a 5-second clip. Video generation starts at lengths that language models only occasionally reach.
-
Long-context training is the dominant cost, not a post-training adaptation. In LLMs, the "train-short, adapt-long" paradigm means most FLOPS (~90%) are spent on short sequences ((\leq) 32K tokens). Video DiTs, by contrast, "dedicate most of their compute budget to full-resolution, long-sequence training." There is no "train-short" option for video—you must train at the target resolution and duration from the start.
-
Inference is also long-context. LLM inference often processes short prompts with long outputs generated token-by-token (where attention cost to cached keys/values is linear per step). Video DiTs generate all tokens at once in a denoising process, requiring full attention over the entire sequence at every denoising step.
This urgency gap is why the paper argues that sparse attention should be "a core design of video DiTs, not a post-hoc fix." The alternative—continuing to use full 3D attention—means accepting that video generation models will remain computationally out of reach for most researchers and practitioners, or that they will sacrifice quality by using fixed sparse patterns.
DSV: The Only Prior Trainable Sparse Method for Video DiTs—And Why It's Insufficient
The paper identifies DSV (Tan et al., 2025) as "pioneering work in exploring trainable sparse attention specifically for video DiT training." This is the closest prior work to VSA, and the paper explicitly contrasts their approaches. DSV introduces dedicated low-rank attention predictors with reduced head dimensions to identify important regions, but these predictors are trained in a separate, multi-stage process and are not fully end-to-end integrated with the main DiT training. The paper describes DSV's "multi-stage and profiler-based design" as potentially complicating the training pipeline.
VSA's counter-proposal is elegantly simpler: instead of training separate predictors offline, reuse the attention mechanism itself at coarse granularity. By pooling tokens into cubes and computing attention at the cube level, VSA simultaneously (1) identifies critical regions (via Top-𝒦 selection on the coarse attention matrix) and (2) models global context (via the coarse attention output that feeds into the final representation). This is end-to-end trainable with no separate stages, no profiling, and minimal architectural modifications.
The Central Research Question and Conceptual Challenge
The paper crystallizes its motivation into a single research question in Section 1:
"How can we predict critical tokens accurately, subject to hardware-aligned block structure, without paying the quadratic cost we aim to avoid?"
This question captures a chicken-and-egg dilemma that is the conceptual heart of the paper. To identify which tokens are "critical" (i.e., have high attention weights), you traditionally need to compute the full attention matrix . But computing this matrix is exactly the operation you're trying to avoid. Any method that computes full attention to decide what to skip has already paid the cost it aimed to save.
Prior work attempted to escape this dilemma through four strategies, each with limitations:
-
Fixed patterns (spatial, temporal, sliding window): Zero prediction cost, but miss data-dependent critical tokens that don't conform to the pattern.
-
Offline profiling: Run full attention once on a calibration dataset, record which positions get high attention weights, and freeze those patterns for all future inputs. This works only if the attention pattern is static across inputs, which the paper's visualizations in Figure 5 show is false—different prompts, layers, heads, and timesteps exhibit markedly different attention patterns.
-
Cheap heuristics: Use positional distance, token similarity, or other proxies to predict importance without computing full attention. These are fast but inaccurate.
-
Hierarchical prediction: Use a coarse (cheap) attention computation to predict which regions are important, then compute fine attention only in those regions—this is VSA's approach.
The key insight that makes strategy 4 viable is that pooling tokens into cubes before computing attention reduces cost by a factor of , where is the cube size. With (the paper's default setting of cubes), the coarse attention cost is reduced by a factor of . This is so cheap—less than 0.2% of total attention compute—that it effectively breaks the chicken-and-egg dilemma: you can "afford" to compute full attention at the coarse level because it operates on a sequence that is 64× shorter.
But this introduces a new tension that the paper explores in depth: the coarse stage's accuracy depends on its granularity. Smaller cubes (e.g., giving ) provide finer-grained critical token prediction, meaning the fine stage wastes less computation on irrelevant tokens. But smaller cubes also mean more cubes total, increasing coarse stage cost, and—crucially—they result in smaller block-sparse tiles that are less efficient on GPU hardware (lower arithmetic intensity, more kernel launch overhead). The paper's tile size ablations in Table 1(c-d) systematically explore this tradeoff, finding that with cubes hits a sweet spot where prediction accuracy is sufficient and hardware efficiency remains high.
How VSA Positions Itself
The paper positions VSA as occupying a specific, previously empty point in the design space: a trainable, end-to-end, hardware-efficient sparse attention that applies at both training and inference, requires no post-hoc profiling, and uses learned (not fixed) sparsity patterns.
This positioning is best understood through what VSA is not:
- Not a post-hoc method: Unlike STA, Sparge, and SVG, VSA trains with sparsity from the start, eliminating the train-test mismatch.
- Not a fixed pattern: Unlike spatial-temporal or sliding window attention, VSA's sparsity pattern is learned from data through the coarse stage's predicted Top-𝒦 selection.
- Not an offline profiling method: VSA predicts critical tokens online for each input, adapting to content-dependent attention patterns.
- Not a multi-stage training pipeline: Unlike DSV, VSA is trained end-to-end with a single loss, requiring no separate predictor training or profiling stages.
- Not sacrificing hardware efficiency: VSA co-designs the cube partitioning with GPU tile sizes, ensuring that the block-sparse pattern maps cleanly to efficient kernels (85% of FlashAttention-3 MFU).
The paper's explicit framing in Section 2.1 establishes a design space with three axes: (1) block size vs. hardware efficiency, (2) prediction cost vs. coverage quality in critical token selection, and (3) global vs. local context modeling. VSA's specific choices—, lightweight coarse attention for prediction, and coarse output gated into the final representation as the global context mechanism—are justified through systematic ablation rather than asserted as obviously correct.
The paper also makes a subtle but important rhetorical move by conducting its primary analysis on pretraining from scratch, not just fine-tuning pretrained models. The scaling studies (Figure 2) across 60M to 1.4B parameters with up to FLOPS establish that VSA is not merely a way to cheaply approximate an existing dense model—it is a better way to train video DiTs, producing a strictly better Pareto frontier. This reframes sparse attention from a necessary compromise (accepting quality loss to save compute) to a genuine improvement (training better models for the same compute budget). The paper explicitly claims to be "the first trainable sparse attention approach that, based on extensive experiments totaling around 90k H200 hours, shows better scaling than full attention on DiTs"—a claim that, if sustained by the community, would represent a significant shift in how video generation models are designed.
3. Technical Approach
3.1 Reader Orientation
VSA is a drop-in replacement for the full 3D attention module inside video Diffusion Transformers—it substitutes the quadratic-complexity Softmax(QK^T/√d)V computation with a two-stage, learned sparse approximation that can be used at both training and inference time without changing the rest of the DiT architecture. The problem it solves is that full attention over video latents of shape (T, H, W) scales as O((THW)²), which for practical video lengths (100K+ tokens) makes attention the dominant cost; VSA reduces this cost by approximately 8× in attention FLOPs and 2.53× in total training FLOPs by learning to predict—on the fly and without computing the full attention matrix—which small subset of token-to-token interactions actually matter, then restricting expensive fine-grained attention to only those interactions in a hardware-aligned block-sparse layout.
3.2 Big-Picture Architecture (Diagram in Words)
VSA has four major components that operate in sequence on each transformer layer's self-attention:
-
Cube Partitioning (pre-processing). The input video latent of shape
(T, H, W)is divided into contiguous 3D cubes, each of shape(C_t, C_h, C_w). By default, this is(4, 4, 4), giving a cube size ofB = 64tokens. The 1D sequence index of each token is reordered so that tokens within the same cube occupy contiguous positions in the flattened sequence, aligning cubes with GPU tile boundaries. -
Coarse Stage. Each cube's tokens are mean-pooled into a single representation, producing
Q_c, K_c, V_cat 1/64th the original sequence length. Full dense attention is computed on these pooled representations. This produces two things simultaneously: (a) a coarse attention outputO_cthat captures global, low-resolution context, and (b) per-row Top-𝒦 selection on the coarse attention matrix that identifies which cubes contain "critical tokens" for each query position. -
Fine Stage. Using the Top-𝒦 selection from the coarse stage as a block-sparse mask
M, token-level attention is computed only within the selected cubes—i.e., for each query token,QK^TandAVare skipped entirely for key/value tokens in non-selected cubes. This is implemented as a block-sparse attention kernel that processes each selectedB × Btile as a dense submatrix. -
Gated Combination. The coarse output
O_cand fine outputO_fare combined via learned scalar gates:O = G_c ⊙ O_c + G_f ⊙ O_f, whereG_candG_fare per-head gate values obtained by linear projection of the input hidden states. This allows the model to learn, independently per head and per position, how much to rely on global coarse context versus local fine-grained attention.
Information flows through a transformer layer as follows. The video latent x ∈ R^(THW)×d enters the self-attention block → it is re-indexed into cube-major order → Q, K, V projections are computed as usual → the coarse stage pools, attends, and selects cubes → the fine stage computes sparse token-level attention within selected cubes → the two outputs are gated and summed → the result is un-tiled back to the original token order and passed to the subsequent MLP and cross-attention blocks exactly as in standard full attention.
3.3 Roadmap for the Deep Dive
- First, the cube partitioning and token re-indexing scheme (Section 2.2), because all subsequent stages depend on the specific mapping between 3D video structure and 1D GPU tile layout.
- Second, the coarse stage—how pooling works, how the coarse attention matrix
A_cis computed, and how Top-𝒦 selection converts it into a block-sparse mask for the fine stage—since this is the mechanism that decides which tokens receive fine-grained attention. - Third, the fine stage—how the block-sparse mask shapes the computation, what the kernel does, and how hardware alignment makes the theoretical FLOP reduction translate to wall-clock speedup.
- Fourth, the gated combination mechanism, which is small but functionally important: it determines how the coarse and fine outputs are mixed, and its initialization strategy is critical for sparse adaptation of pretrained models.
- Fifth, the kernel implementation considerations (Section 2.4), because VSA's design choices—especially cube size and Top-𝒦 selection method—are driven by what GPU kernels can execute efficiently.
- Sixth, the sparse adaptation protocol (Section 2.3), which is how VSA retrofits pretrained full-attention DiTs without destabilizing training.
- Seventh, key design parameters and their tradeoffs (tile size, sparsity level, pooling method), as documented in the ablation studies, because these choices are not obvious and the paper provides empirical guidance.
3.4 Detailed, Sentence-Based Technical Breakdown
This is primarily an architecture design and empirical scaling paper whose core idea is that a learned, data-dependent sparse attention mechanism—when co-designed with hardware-aligned block-sparse compute—can replace full 3D attention in video DiTs at both training and inference without quality loss, and that the key enabling technique is hierarchical attention where a cheap coarse stage predicts critical token locations that a fine stage then processes at high resolution.
Cube Partitioning and Token Re-Indexing
The standard approach in video DiTs is to flatten a 3D latent of shape (T, H, W) into a 1D sequence of length L = T × H × W using a row-major mapping: token at position (t, h, w) gets index n = t·H·W + h·W + w. Full attention is then applied over all L tokens, treating the video as an unstructured 1D sequence.
VSA changes this mapping to group tokens that are spatiotemporally contiguous into the same cube, and to make cubes themselves contiguous in the flattened sequence. The paper gives the exact mapping in Section 2.2:
where C_t, C_h, C_w are the cube dimensions (default (4, 4, 4)), B = C_t × C_h × C_w is the cube size in tokens (default 64), and N_t = T/C_t, N_h = H/C_h, N_w = W/C_w are the numbers of cubes along each dimension.
What this formula does. It partitions the T×H×W volume into N_t × N_h × N_w cubes, each containing B contiguous tokens in the 1D sequence representation. The first term (⌊t/C_t⌋ N_h N_w + ⌊h/C_h⌋ N_w + ⌊w/C_w⌋) · B determines which block of B consecutive sequence positions this token belongs to (the cube index); the second term determines the token's position within that block.
Why this specific mapping. The mapping serves two purposes simultaneously. First, it ensures that when a coarse-stage cube is selected for fine attention, all tokens in that cube are contiguous in memory and contiguous in the attention matrix, allowing the fine-stage kernel to process each selected query-cube-to-key-cube interaction as a single dense B × B tile. Second, it preserves the property that the coarse stage's pooled representations (averaging over exactly each block of B contiguous tokens) exactly correspond to spatiotemporal cubes in the original video—the pooling is not an arbitrary compression but a semantically meaningful aggregation of neighboring patches.
What would be wrong with the standard mapping. If VSA applied the standard row-major mapping and then selected arbitrary subsets of tokens for fine attention, the fine-stage kernel would need to process scattered, non-contiguous tokens, which would require expensive gather/scatter operations and defeat the purpose of hardware-aligned sparsity. The re-indexing step makes sparsity "free" in hardware terms by converting the problem from unstructured token selection to block selection.
Implementation detail. The paper notes in Appendix B that the tile and untile operations (rearranging tokens into and out of cube-major order) "can be moved to the beginning and end of the transformer to avoid calling them for each attention." This means the re-indexing is applied once per transformer block, not per attention head, making its overhead negligible.
The Coarse Stage: Cube-Level Attention and Critical Token Identification
The coarse stage is the mechanism that lets VSA predict which token interactions matter without computing the full L × L attention matrix.
Input representation: mean pooling. Given the token-level Q, K, V ∈ R^(L×d) for a single attention head (produced by standard linear projections from the input hidden states), the coarse stage pools each cube of B tokens into a single representation. Specifically, for each of the L/B cubes, it computes the mean of the B token representations within that cube. The paper formalizes this in the pseudocode:
q_c = q.view(B, H, L//block, block, D).mean(dim=3)
k_c = k.view(B, H, L//block, block, D).mean(dim=3)
v_c = v.view(B, H, L//block, block, D).mean(dim=3)
where block = B = 64 and the mean(dim=3) operation averages over the tokens within each cube. This produces Q_c, K_c, V_c ∈ R^((L/B)×d)—representations at 1/64th the original sequence length.
Why mean pooling, not max or convolution. The paper ablates this choice in Table 1(e). Mean pooling outperforms max pooling, and the paper reports that convolutional pooling (3D convolution with kernel size and stride equal to cube size) "caus[es] training instability." The likely reason: mean pooling is a linear operation that preserves the expected value of the token representations within each cube, making the coarse attention scores unbiased estimates (in expectation) of the average token-level attention weights within each cube. Max pooling is non-linear and would select the most extreme token in each cube, which could systematically distort the coarse attention scores. Convolution introduces trainable parameters into the pooling step, creating a more complex optimization landscape that can interact badly with the Top-𝒦 selection's non-differentiability.
Coarse attention computation. The coarse stage computes standard scaled dot-product attention on the pooled representations:
This produces an attention matrix A_c ∈ R^((L/B)×(L/B))—for a typical 16K-token sequence with B=64, this is a modest 256×256 matrix. The FLOP cost of this operation is O((L/B)² · d), which is 1/4096th the cost of full attention on the original sequence.
Top-𝒦 selection: from attention scores to block-sparse mask. For each row of A_c (each query cube), the coarse stage selects the 𝒦 key-cube positions with the highest attention scores. The paper uses 𝒦 = 32 as the default, and the total number of cubes is L/B (e.g., 256 for a 16K-token sequence). These 𝒦 selected indices form the critical-cube set for that query position. The selection is:
topk_vals, topk_idx = score.topk(topk, dim=-1)
where score = A_c. A new mask matrix is created where only the Top-𝒦 entries per row retain their attention scores; all other entries are set to zero. This mask is then broadcast to token resolution: each selected (i, j) cube-pair in the coarse mask expands into a B × B block of ones in the fine-stage mask, and each unselected pair expands into a B × B block of zeros (or -∞ after masking).
Why Top-𝒦 is applied to the coarse attention scores, not some other signal. The key insight is that if the mean-pooled cube representation captures the average token within that cube, then the attention score between two cube representations approximates the average attention score between token pairs drawn from those cubes. Selecting the Top-𝒦 cube pairs by this score is therefore a lightweight proxy for selecting the token pairs that would have the highest attention weights in the full matrix, without ever materializing that full matrix.
An alternative would be to compute some cheaper similarity metric (e.g., dot product of the pooled Q_c and K_c without softmax, or cosine similarity), but the paper specifically uses the full softmax attention scores because the softmax normalization matters—it creates competition among key positions, ensuring that the selected 𝒦 positions are those that actually dominate the attention distribution rather than just those with the largest raw similarity.
What the coarse attention output O_c provides. The coarse output is computed as O_c = A_c V_c—the attention-weighted combination of pooled value representations. This output is then un-pooled by replicating each cube's representation B times to restore the original sequence length: O_c_expanded ∈ R^(L×d). This expanded output represents a low-resolution, globally-contextualized view of the video—each token's representation is influenced by all other cubes (because A_c is dense), but only at the granularity of a (4,4,4) cube.
Why include O_c in the final output at all. The paper's ablation in Table 1(b) (comparing configurations "F" (fine-only) vs. "C & F" (coarse + fine)) demonstrates that including the coarse output is critical: VSA configurations with the coarse output (Exps 9–10) achieve substantially lower loss than those relying solely on the fine stage (Exp 7–8). The coarse output serves as a safety net: even if the Top-𝒦 selection misses some important token interactions, those interactions are still approximately captured at the coarse cube level and contribute to the final representation via the gating mechanism.
The Fine Stage: Token-Level Block-Sparse Attention
The fine stage is where VSA achieves its computational savings—it computes exact token-level attention, but only within the B × B blocks selected by the coarse stage.
Mask construction. The Top-𝒦 selection from the coarse stage produces, for each query cube index i, a set of 𝒦 key-cube indices {j_1, ..., j_𝒦}. For the fine stage, this is translated into an attention mask M ∈ {-∞, 0}^(L×L) as follows:
- For each selected cube pair
(i, j_k), the correspondingB × Bsubmatrix ofM(rowsi·Bthrough(i+1)·B - 1, columnsj_k·Bthrough(j_k+1)·B - 1) is set to all zeros (allowing attention). - All other entries of
Mare set to-∞(blocking attention, becauseexp(-∞) = 0in the softmax).
The overall sparsity—the fraction of L × L entries that are blocked—is approximately 1 - (𝒦B/L), since each of L/B query cubes attends to 𝒦 key-cubes, each of size B, giving 𝒦B total attended key positions per query. With the default 𝒦 = 32 and B = 64, and for L = 16,384, the sparsity is 1 - (32 × 64 / 16,384) = 1 - 0.125 = 0.875 or 87.5%.
Fine attention computation. The fine stage computes:
where Q, K, V ∈ R^(L×d) are the full-resolution token representations. The addition of M in the softmax ensures that key positions in non-selected cubes receive effective attention weight of exactly zero. The computation of QK^T and A_f V is block-sparse: only the matrix products corresponding to non-masked B × B tiles are computed; all other tiles are skipped entirely.
Why block-sparse rather than token-level sparsity. If VSA selected individual critical tokens (rather than critical cubes), the sparsity pattern would be unstructured—the attended token indices would be scattered non-contiguously across the sequence. On modern GPUs, this would require either (a) gathering the selected key/value tokens into contiguous memory before computing attention, which adds overhead that can dominate the savings from reduced FLOPs, or (b) using masked dense computation that still loads and computes on all tiles but masks out results, which saves no compute. By enforcing that sparsity is at the cube (block) level, VSA ensures that the fine-stage kernel can use the same tiled matrix-multiplication primitives as dense FlashAttention, simply skipping tiles that are entirely masked out. This is what allows the theoretical 8× FLOP reduction to translate to ~7× actual speedup (Figure 4b).
The block-sparse kernel implementation. The fine stage uses a block-sparse attention kernel written with ThunderKittens, a library for high-performance GPU kernels. The kernel receives the query, key, and value tensors plus a list of block indices indicating which B × B tiles to process. For each query tile, it loads the corresponding B × d submatrix of Q into SRAM, then iterates through the selected key/value tiles for that query tile, loading each B × d submatrix of K and V and computing the local B × B attention update. This is essentially FlashAttention restricted to a subset of tiles, preserving all the IO-awareness benefits (tiling over the sequence dimension to keep SRAM usage within bounds) while skipping non-selected tiles.
Gated Combination of Coarse and Fine Outputs
The final output of VSA is a learned, per-head interpolation between the coarse and fine attention outputs:
where O_c ∈ R^(L×d) is the expanded coarse output, O_f ∈ R^(L×d) is the block-sparse fine output, and G_c, G_f ∈ R^(L×1) are per-position gating scalars.
What the gates are and how they're computed. The gates are produced by a learned linear projection from the input hidden states. Specifically, for each attention head, a small weight matrix W_g ∈ R^(d×2) maps the d-dimensional input representations to two scalar logits per position; these are then split into G_c_raw and G_f_raw. The pseudocode shows:
gate = tile(gate) # reshape from L×2 to (L/B)×B×2
coarse_attn_gate, fine_attn_gate = gate.chunk(2, dim=1)
The paper does not specify whether a softmax or sigmoid is applied—the gating values appear to be used directly as multiplicative weights, meaning the gates are unconstrained scalars that the model can learn to set positive or negative. This gives the model flexibility to not just interpolate but potentially amplify or suppress either output.
Why gating rather than simple addition or concatenation. Simple addition O = O_c + O_f would force both outputs to be present in equal proportion, regardless of whether a particular head or position benefits more from global or local information. Concatenation would double the output dimension, requiring a down-projection that adds parameters and compute. Learned scalar gating allows the model to discover, during training, which heads specialize in global vs. local processing and to dynamically adjust the mixture per position—for example, a token in a region with rapid motion might rely more on the fine stage (which captures detailed spatiotemporal interactions), while a token in a static background region might rely more on the coarse stage (which provides broader context with less detail).
Initialization for pretrained model adaptation. When retrofitting VSA into a model originally trained with full attention (Section 2.3), the gating weights W_g are randomly initialized (they don't exist in the full-attention checkpoint). The paper's sparse adaptation protocol initializes the coarse gate G_c to zero, making VSA initially equivalent to full attention (because G_f is unconstrained and O_f with 𝒦 = L/B processes all cubes, i.e., full attention). As training progresses and 𝒦 is reduced to introduce sparsity, G_c is learned from its zero initialization, allowing the model to gradually incorporate coarse global context as an explicit component separate from the increasingly sparse fine attention.
The fine gate initialization. The paper states they "remove the fine gate G_f (equivalent to G_f = 1)" during sparse adaptation. This means they disable the gating for the fine output initially, using it directly (O_f without scaling), while the coarse gate starts at zero. As training progresses, the fine gate can be re-enabled or remain fixed—the paper doesn't specify whether it is eventually re-introduced, but the ablation results suggest the gating mechanism matters for from-scratch training.
Kernel Implementation: Why 85% MFU and Where the Overhead Goes
The paper's kernel design choices are driven by a pragmatic question: given that VSA replaces 87.5% of attention FLOPs with zeros, does the actual wall-clock speedup approach the theoretical 8×, or do kernel overheads eat the gains?
Fine-stage kernel: achieving ~7× speedup at long sequence lengths. The fine-stage block-sparse kernel, written with ThunderKittens, achieves 85% of FlashAttention-3's Model FLOP Utilization (MFU). This means that for the FLOPs it actually performs (12.5% of full attention), it executes them at 85% of the theoretical peak throughput of the GPU. At long sequence lengths (100K+ tokens), this translates to approximately 7× speedup over FlashAttention-3 (Figure 4b), which is close to the theoretical 8× maximum given 87.5% sparsity.
The gap from 8× to 7× comes from kernel launch overhead, block-index processing, and the fact that the coarse stage adds back some compute. The paper notes that "FlexAttention with an identical block-sparse mask (64×64 block size) achieves only a 2× speedup," highlighting the importance of a purpose-built block-sparse kernel rather than relying on general-purpose sparse attention frameworks.
Coarse-stage kernel: the Top-𝒦 bottleneck. The coarse stage presents a challenge for standard fused-attention kernels. FlashAttention-style kernels avoid materializing the full L × L attention matrix in HBM (high-bandwidth memory) by fusing the softmax with the matrix multiplication in SRAM. However, the coarse stage requires a row-wise Top-𝒦 operation, which necessitates knowing the full set of attention scores for each row to identify the 𝒦 largest values. This forces materialization of the coarse attention matrix A_c.
The paper's response is to accept this materialization and instead optimize what happens after: "we fuse softmax, Top-𝒦 selection, and mask-to-index conversion into a single kernel." The fused kernel performs, in one pass over the (L/B) × (L/B) attention matrix: (1) softmax normalization, (2) selecting the Top-𝒦 indices per row, and (3) writing those indices directly in the format expected by the fine-stage block-sparse kernel (contiguous arrays of block indices, not binary masks). This fusion "reduces coarse stage runtime modestly" (Appendix D, Table 3).
Why the coarse stage overhead is acceptable. The paper makes three quantitative arguments:
- FLOP cost is tiny: the coarse stage operates on sequences 64× shorter, making its attention cost less than 0.2% of total attention FLOPs (Section 2.4).
- Runtime is dominated by other operations: even at modest sequence lengths, the coarse stage accounts for only 14% of the total attention runtime when the fine stage is 87.5% sparse (Appendix D). At long sequence lengths (the target regime), the coarse stage's share shrinks further.
- Fusing softmax and Top-𝒦 eliminates memory traffic: rather than writing the full attention matrix to HBM, computing Top-𝒦 on it with a separate kernel (which would read it back), and then converting the mask, the fused kernel does all three in-register, saving bandwidth.
What would be needed for better coarse-stage performance. The paper acknowledges that "coarse stage acceleration remains an important direction for future research" but deliberately avoids the complexity of modifying FlashAttention to support in-kernel bitonic sorting for Top-𝒦—"such fusion demands intrusive kernel rewriting and careful tuning." The pragmatic judgment is that the coarse stage is cheap enough at target sequence lengths that the engineering investment isn't justified.
Sparse Adaptation: Retrofitting Pretrained Full-Attention DiTs
VSA is primarily designed for training from scratch, but the paper also demonstrates how to retrofit pretrained models that were originally trained with full attention. This is the "sparse adaptation" protocol in Section 2.3.
The problem: direct substitution causes instability. When you simply replace the full attention module in a pretrained checkpoint with VSA, two things go wrong:
- Missing parameters: The gate projection weights
W_gare not present in the full-attention checkpoint and must be randomly initialized. This introduces untrained parameters into an otherwise converged model. - Distribution shift: The model's parameters were optimized assuming attention sees all tokens everywhere. Switching to 87.5% sparsity represents a sudden, severe distribution shift—the information each token receives changes drastically, and the learned weights are not adapted to sparse attention patterns.
The solution: progressive sparsity annealing with gate initialization. The sparse adaptation protocol uses three techniques together:
-
Coarse gate initialized to zero.
G_cstarts at zero, meaning the coarse output contributes nothing initially. This makes VSA's behavior at initialization identical to fine-stage-only attention. -
Fine gate removed (set to identity). The fine gate
G_fis set to 1 everywhere, i.e., the fine output passes through unchanged. This removes two sources of randomness (theW_gweights) from the initialization. -
Sparsity schedule. The model starts with
𝒦 = L/B(all cubes selected, which makes the fine stage compute full attention—the maskMhas no-∞entries). Over a series of training steps,𝒦is gradually reduced to the target value (e.g., 32). The paper specifies: "we decrease the number of attended cubes by 10 (i.e., reduce Top-K by 4) every 50 steps, until reaching the target sparsity level" (Appendix C.5).
The combination of (1) and (3) means that at step 0, VSA computes exactly full attention (with G_c = 0, G_f = 1, and no sparsity), making the initialization loss-identical to the pretrained checkpoint. As training progresses, sparsity increases and the coarse gate gradually activates, allowing the model to adapt its parameters to the new attention structure while the loss remains stable.
Warmup before sparsity decay. For Wan-1.3B, the paper trains with full attention for 50 steps before starting the sparsity decay schedule. This warmup "accommodate[s] the changed resolution and aspect ratio" when finetuning on the target video data. The full adaptation takes 4,000 steps total, and the model generates coherent videos by the end (Figure 6 shows the progression from artifacts to clean generation).
Why this works—and why it's not trivial. The annealing approach succeeds because the model's internal representations can continuously adapt as sparsity increases, rather than being forced to jump to an unfamiliar attention pattern. If the coarse gate and sparsity were applied abruptly, the model would face a combined shock: (a) many attention pathways are suddenly blocked, and (b) the new coarse context signal is introduced at random magnitude. The progressive schedule separates these two changes in time, allowing the model to first adjust to reduced connectivity and then incorporate the coarse signal.
Sparse distillation: combining VSA with few-step generation. The paper's "Sparse-Distill" experiment (Section 3.3, Appendix C.6) is notably the first to combine sparse attention with distillation. In DMD2-style distillation, a student model is trained to generate videos in very few denoising steps (e.g., 3 steps) by matching the teacher's distribution. The paper replaces the student's full attention with VSA (at 80% sparsity—lower than the 87.5% used in most experiments, likely because distillation operates in a more challenging very-low-step regime) while keeping the teacher unchanged with full attention. All distillation hyperparameters are held fixed; only the student's attention module changes. This achieves 50.9× total speedup relative to the baseline model (combining the 3-step generation with VSA's attention speedup), with "no quality drop" in human evaluation. This is significant because it demonstrates that sparse attention is compatible with distillation—a non-obvious result, since distillation already operates at the edge of generation quality and any additional approximation could push it over the edge.
Design Parameters and Their Tradeoffs
The paper's ablation studies systematically explore the design choices that make VSA work, documenting not just which choices were made but why alternatives fail.
Tile size B and cube shape (C_t, C_h, C_w). This is the most consequential design parameter. The tile size determines:
- Coarse-stage prediction accuracy. Smaller tiles allow the coarse stage to more precisely localize critical tokens. A
(4,4,4)cube (64 tokens) can identify that something important is in a specific 4×4×4 spatiotemporal region; a(4,8,8)cube (256 tokens) would identifiy a much larger region, forcing the fine stage to process many irrelevant tokens. - Fine-stage attention granularity. Smaller tiles mean each selected block covers fewer tokens, reducing wasted computation on non-critical tokens within selected blocks.
- GPU arithmetic intensity. Larger tiles have higher arithmetic intensity (more compute per byte loaded from memory), which improves MFU. Table 1(d) shows MFU drops from values near FlashAttention's level at 256×256 tiles to significantly lower levels at 64×16.
The paper's experimental finding (Table 1(c)) is that "smaller tiles consistently reduce model loss through finer attention granularity." However, the speed penalty is significant: 64×16 tiles offer "slightly better performance" but run 2.26× slower than 64×64 (Table 1(d), comparing Exp 18 to 17). The chosen default of (4,4,4) giving B=64 balances model quality against throughput.
A particularly illuminating experiment is Exp 16 in Table 1(c), which tests mismatched granularity between stages: coarse stage uses (4,4,4) pooling (fine granularity, B=64 equivalent) while the fine stage operates on larger (4,8,8) tiles. To bridge the mismatch, an additional (1,2,2) average pooling is applied to the coarse attention map before Top-𝒦 selection. This configuration underperforms the matched (4,4,4) configuration, demonstrating that both stages benefit from fine granularity—the coarse stage predicts more precisely, and the fine stage wastes less compute.
Global vs. local context. The paper explores whether explicit local attention modules add value beyond what the coarse stage provides. Three variants are tested (Table 1(b)):
- Exp 11: Separate local stage. A dedicated local attention with a
(3,3,3)sliding window runs in parallel with the coarse+fine stages, contributing its own gated output. This adds parameters and compute but shows "similar" performance to the simpler C&F architecture. - Exp 12: Local exclusion. The coarse stage explicitly excludes cubes that would be selected by a local
(3,3,3)pattern from the Top-𝒦 selection, preventing redundancy between local and global attention. Again, similar performance. - Exp 13: Forced local inclusion. The fine stage is forced to always include the local
(3,3,3)cubes (regardless of Top-𝒦 selection). Similar performance to the baseline.
The conclusion is that "explicit local modeling provides minimal benefit"—the coarse stage already captures sufficient local context through its dense attention on cube-level representations (since neighboring cubes are spatially adjacent tokens, they naturally have high coarse-attention scores and are selected by Top-𝒦). This is a significant simplification: VSA does not need the separate sliding-window stage that NSA (Yuan et al., 2025) found necessary for language models, likely because video's 3D structure and bidirectional attention make coarse-cube attention inherently better at capturing local neighborhoods.
Pooling method for the coarse stage. Table 1(e) compares mean pooling, max pooling, and convolution-based pooling. Mean pooling outperforms max pooling, and convolution causes training instability. The paper doesn't deeply analyze why, but the likely mechanism: mean pooling produces Q_c, K_c values that are the centroids of the token distributions in each cube, making the coarse attention scores linear approximations to the average token-level attention weights. Max pooling is not linear, so the relationship between coarse attention scores and token-level attention weights becomes unpredictable. Convolution adds learned parameters to the pooling, which can learn degenerate transformations (e.g., amplifying noise) that interact badly with the Top-𝒦 selection's non-differentiable cutoff.
Sparsity level 𝒦. The paper's most surprising scaling finding (Figure 2(c)): 𝒦 = 32 performs consistently well across sequence lengths from 8,192 to 24,675 tokens under a fixed 4.5 × 10^20 FLOPS training budget, but underperforms 𝒦 = 16 at 61,440 tokens. However, when the training budget is increased to 1 × 10^21 FLOPS at 61,440 tokens, 𝒦 = 32 eventually overtakes 𝒦 = 16. This suggests the optimal sparsity level depends jointly on both sequence length and training compute—with more compute, the model learns to use the additional attended positions productively, and the penalty for higher sparsity (missing some important interactions) grows relative to the benefit (more efficient FLOP use).
The paper hypothesizes that "the ideal Top-𝒦 increases with available compute, converging to full attention with infinite resources." This frames sparsity as a resource allocation parameter analogous to model size in pretraining scaling laws—at low budgets, aggressive sparsity lets you allocate more FLOPs to other model dimensions; at high budgets, the information lost to sparsity costs more than the FLOP savings are worth. The paper explicitly suggests "explicitly model[ing] sparsity as an additional axis in the scaling law framework."
Data-dependent vs. fixed patterns. Table 1(a) provides the crucial head-to-head: at a compute-optimal training budget (4.5 × 10^20 FLOPS), existing fixed-pattern sparse methods (spatial-temporal, spatial-full, strided window, compressed KV) all achieve lower loss than full attention, but VSA (with learned, data-dependent sparsity) achieves the lowest loss of all. At an extended budget (4 × 10^21 FLOPS), the fixed-pattern methods actually perform worse than full attention, while VSA continues to match or exceed full attention. This is the empirical core of the paper's claim: data-dependent sparsity is not merely better than fixed patterns—it is qualitatively different in how it scales with compute, recovering the full attention loss curve while fixed patterns diverge from it.
Why fixed patterns fall behind with more training. The paper doesn't explicitly state the mechanism, but it's consistent with the interpretation that fixed patterns impose a representation bottleneck. With limited training, the bottleneck isn't binding because the model hasn't yet learned to exploit long-range dependencies; the computational savings from sparse attention allow training on more data, improving loss. With extended training, the model saturates the patterns' capacity—it has learned everything it can from the restricted attention connectivity, and further training can't overcome the missing interactions. Full attention has no such bottleneck, so it continues to improve with more training. VSA's learned sparsity adapts the connectivity pattern to the data, expanding the bottleneck when needed and contracting it when safe, so it can continue improving alongside full attention.
Mismatched Granularity and the Hardware-Software Co-Design Principle
One of the paper's more subtle contributions is the explicit co-design of the attention algorithm with GPU tile sizes. The cube size B=64 is not an arbitrary choice—it emerges from the constraint that both the coarse stage pooling and the fine stage tiles must align with what the GPU kernel processes efficiently.
Why cube partitions map to GPU tiles. In FlashAttention-style kernels, the attention matrix is divided into tiles of size B_q × B_k (query-tile size × key-tile size), and each GPU threadblock loads one query tile and one key tile into SRAM, computes the local attention update, and accumulates results. The efficiency of this process depends on the tile sizes being well-aligned with SRAM capacity and warp sizes. VSA's cube size maps directly to these tile sizes: each cube becomes one query tile and one key tile in the block-sparse kernel.
The tension between sparsity and arithmetic intensity. Smaller tiles (e.g., B=16) would let the coarse stage more precisely localize critical tokens, reducing wasted fine-stage computation. But smaller tiles reduce GPU efficiency for two reasons: (1) the ratio of compute (2B²d FLOPs for QK^T and AV) to memory traffic (4Bd loads/stores) decreases with B, reducing arithmetic intensity; (2) more tiles means more kernel launches or more loop iterations, increasing scheduling overhead. The paper's MFU measurements in Table 1(d) quantify this: 64×16 tiles achieve 2.26× worse throughput than 64×64.
The sweet spot at B=64. The paper settles on B=64 because it is (a) small enough that the coarse stage can localize critical tokens reasonably well (the Top-32 selection captures 50–90% of attention mass, per Figure 5(e)), (b) large enough for efficient GPU execution (85% MFU), and (c) a practical divisor of typical video latent dimensions (the paper notes that dimensions must be divisible by 4, which is satisfied by standard latent shapes like 16×32×32, 16×28×52, and 20×48×80).
FLOP Accounting: How the 2.53× Training Reduction Is Computed
The paper reports a 2.53× reduction in total training FLOPs when using VSA compared to full attention. This figure comes from the standard transformer FLOP approximation, modified for sparse attention.
Standard FLOP approximation. Following Hoffmann et al. (2022), the paper uses:
where N is the number of model parameters, D is the number of tokens in the batch, and 6ND accounts for the FLOPs in all linear layers (Q/K/V projections, output projection, MLP). The attention-specific FLOPs are computed as:
where D is the number of tokens, S is the sequence length per token, H_heads is the number of attention heads, d_head is the head dimension, and L_layers is the number of transformer layers. The factor 3.5 accounts for the details of the FlashAttention algorithm (not just the theoretical 2S² FLOPs per head, but including softmax and scaling operations).
Adjustment for VSA. For sparse attention, the attention FLOPs are multiplied by the fraction of tokens actually attended to: 𝒦B / L (the fraction of key tokens retained). With 𝒦 = 32 and B = 64, and for a 16K sequence length, this fraction is 32 × 64 / 16,384 = 0.125. So attention FLOPs are reduced by 8×.
Why 2.53× total, not 8×. The factor of 2.53× rather than 8× comes from the fact that attention is only a fraction of total compute. The 6ND term (linear layers) is unchanged by sparsity. At the model sizes and sequence lengths studied, attention accounts for roughly 60–70% of total training FLOPs with full attention. An 8× reduction in the attention portion translates to approximately a 2.53× reduction in total FLOPs. This number is model- and sequence-length-dependent: at longer sequences where attention dominates more, the total reduction would be larger.
Summary of Design Choices and Their Justifications
- Cube size
(4, 4, 4),B = 64: Balances fine-grained critical-token prediction against GPU tile efficiency. Smaller would improve prediction but reduce MFU; larger would improve MFU but waste fine-stage compute on irrelevant tokens. - Mean pooling, not max or convolution: Provides a linear, unbiased estimate of average token-level attention within each cube. Max pooling is non-linear; convolution introduces trainable parameters that destabilize training.
- Top-𝒦 = 32 (default): Robust across sequence lengths 8K–25K under moderate compute budgets; higher
𝒦may be needed for very long sequences or very large training budgets. - Coarse attention output included in final output via gating: Critical for performance—provides global context that the sparse fine stage may miss. Ablation confirms removing it degrades loss.
- No dedicated local attention stage: Redundant with coarse stage's cube-level attention, which naturally captures local neighborhoods. Adding one adds complexity without benefit.
- Block-sparse kernel with tile-aligned cubes: Ensures theoretical FLOP savings translate to wall-clock speedup. General-purpose sparse frameworks (FlexAttention) achieve only 2× speedup vs. VSA's 6–7×.
- Fused softmax + Top-𝒦 + index conversion kernel for coarse stage: Avoids multiple passes over the coarse attention matrix. Accepts materialization cost because coarse stage is < 0.2% of total FLOPs.
- Sparsity annealing for pretrained model adaptation: Starts at full attention (loss-identical to checkpoint) and progressively increases sparsity while learning coarse gate from zero initialization. Prevents training instability from abrupt distribution shift.
4. Key Insights and Innovations
Innovation 1: Sparsity as a Trainable Primitive, Not a Post-Hoc Approximation
The dominant approach to attention sparsity in video DiTs has been post-hoc application: train a model with full dense attention, then substitute a sparse pattern (fixed, profiled, or heuristically derived) only at inference time. Methods like Sliding Tile Attention (Zhang et al., 2025), Sparge (Zhang et al., 2025), and Sparse VideoGen (Xi et al., 2025) all operate this way. The implicit assumption is that the model learns useful dense representations during training, and sparsity is an acceptable approximation that trades some quality for speed—a necessary evil, not a design goal.
VSA's core intellectual move is to reject this framing entirely. Instead of treating sparsity as a post-hoc compromise, VSA makes sparsity a first-class trainable primitive that shapes the model's learned representations from the very beginning. The model learns not just what to attend to, but how to allocate its attention budget—the sparsity pattern itself is an output of training, not a constraint imposed afterward.
This shift has consequences that go far beyond the obvious speed-at-inference benefit. The most revealing evidence is in Table 1(a), where fixed-pattern sparse methods (spatial-temporal, strided window, compressed KV) are compared against full attention and VSA under different training budgets. At a compute-optimal budget (4.5 × 10^20 FLOPS), the fixed-pattern methods outperform full attention—a standard finding that has led prior work to claim sparse attention is "good enough." But at an extended budget (4 × 10^21 FLOPS), this advantage reverses: full attention catches up and surpasses every fixed-pattern method. The fixed patterns impose a representation bottleneck—they prevent the model from learning long-range dependencies that matter once training is thorough. VSA is the only sparse method that maintains parity with full attention at the extended budget.
This is a diagnostic, not just a performance claim. It says that the test of whether a sparse attention method is "real" is not whether it beats full attention at small budgets (everything does, because the computational savings let you train on more data), but whether it scales with full attention as compute increases. Fixed patterns fail this test; VSA passes it. The mechanism is that learned sparsity adapts the attention connectivity to the data as training proceeds, expanding the receptive field when the model discovers cross-spatiotemporal dependencies it wants to exploit, while fixed patterns permanently foreclose those connections.
The distinction matters because it reframes what "sparse attention" means. It is not an approximation to a dense target; it is an alternative architectural inductive bias that, when learned jointly with the model parameters, can produce a strictly better Pareto frontier than dense attention (Figure 2(b)). The field's default assumption—that full attention is the gold standard and sparsity is a lossy shortcut—is inverted: learned sparsity can be better at a given compute budget because it forces the model to allocate its limited attention capacity efficiently, rather than wasting parameters and compute on token interactions the model never meaningfully uses.
Innovation 2: The Chicken-and-Egg Resolution via Cost-Asymmetric Hierarchical Prediction
The central conceptual obstacle to trainable sparse attention is what this paper calls the "chicken-and-egg dilemma": identifying which tokens are critical requires computing the full attention matrix, but computing the full attention matrix is exactly the quadratic-cost operation you are trying to avoid. Every trainable sparse attention method must somehow break this circularity.
Prior approaches have attempted four escape routes: (1) fixed patterns that incur zero prediction cost but miss data-dependent critical tokens; (2) offline profiling that freezes patterns after one pass, assuming static attention; (3) cheap heuristics (positional distance, token similarity) that are fast but imprecise; and (4) separate predictor networks (as in DSV) that require multi-stage training. Each makes a different tradeoff on the prediction-cost-vs-coverage-quality axis the paper identifies in Section 2.1.
VSA's resolution is conceptually elegant and specific. Rather than introducing a separate predictor, it reuses the attention mechanism itself at a coarser granularity where the quadratic cost is negligible. The coarse stage is, architecturally, identical to full attention—Q_c K_c^T / √d followed by softmax—but applied to representations pooled by a factor of 64× in sequence length. This reduces the attention cost by 64² = 4096×, making it so cheap (under 0.2% of total attention FLOPs) that the circularity is effectively broken. You can afford to compute full attention at this resolution, and the resulting attention scores directly predict which token-level interactions are important because mean pooling preserves the linear relationship between cube-level and token-level attention weights.
What makes this an innovation rather than an obvious engineering trick is the realization that the asymmetry in cost between the two stages—4096:1—creates a regime where the coarse stage can be "wasteful" (computing dense attention) without compromising the overall efficiency goal, because the fine stage's sparsity savings (8× on attention FLOPs) dominate the total. The coarse stage spends compute to save compute elsewhere, and the net is still a 2.53× reduction in total training FLOPs.
This is distinct from the hierarchical attention in NSA (Yuan et al., 2025) and MoBA (Lu et al., 2025) in two ways. First, those methods use their coarse stage purely for pattern prediction and discard the coarse attention output; VSA includes O_c in the final output via learned gates, making the coarse stage a contributor to the representation, not just a router. The ablation in Table 1(b) shows this matters substantially—removing the coarse output degrades loss. Second, those methods operate in causal language models with single-query decoding constraints, which force design compromises (grouped query attention, sliding windows) that VSA's bidirectional video setting avoids. The conceptual contribution is demonstrating that a sufficiently cheap dense attention stage—operating at the right granularity relative to GPU tile sizes—can serve simultaneously as a predictor, a global context provider, and a safety net for the fine stage's sparsity errors, without requiring separate predictor networks or training stages.
Innovation 3: The Difficulty-Conditioned Scaling of Sparsity (Sparsity as a Resource Allocation Parameter)
The paper's scaling experiments in Figure 2(c) reveal a finding that is subtle, counterintuitive, and—if replicated—potentially foundational for how the field thinks about sparse attention. The optimal sparsity level 𝒦 is not a monotonic function of sequence length. Under a fixed moderate training budget, 𝒦 = 32 performs best across sequence lengths from 8K to 25K, but 𝒦 = 16 (higher sparsity) actually outperforms 𝒦 = 32 at 61K tokens. Then, when the training budget is increased, 𝒦 = 32 regains the advantage at 61K.
This means the optimal sparsity level depends jointly on both sequence length and total compute, and the relationship is not "longer sequences need more 𝒦"—it is more nuanced. The paper's interpretation: with limited training compute, aggressive sparsity (lower 𝒦) lets the model allocate more FLOPs to learning good representations in the remaining parameters; with abundant compute, the information lost to sparsity costs more than the FLOP savings are worth, and higher 𝒦 is optimal. In the limit of infinite compute, 𝒦 converges to full attention.
This is a conceptual innovation, not a recipe. It frames sparsity as an additional axis in the scaling-law framework alongside model size and training tokens. Just as Hoffmann et al. (2022) showed that the optimal model size depends on the training budget (you should train larger models on more data, not the largest model possible on limited data), this finding suggests that the optimal sparsity level depends on the total compute budget. At small budgets, sparsity is a force multiplier—it lets you train a larger or deeper model for the same FLOPs. At large budgets, sparsity is a constraint—it prevents the model from using all the information available in the data.
This has not been previously articulated for sparse attention in any domain. Prior work treats sparsity as a fixed hyperparameter (e.g., "use 50% sparsity") or a constraint derived from hardware limits, not as a variable to be optimized jointly with model scale and compute. The paper explicitly calls for "model[ing] sparsity as an additional axis in the scaling law framework" and acknowledges that the precise functional form is an open question. If subsequent work solves this, it would give practitioners a principled way to answer: "For a video DiT of size X trained with budget Y on sequences of length Z, what sparsity level maximizes final quality?" That would represent the same kind of advance for sparse attention that Chinchilla scaling laws represented for pretraining—a shift from heuristic choices to compute-optimal allocation.
Innovation 4: The Hardware-Algorithm Co-Design for Block-Sparse GPU Execution
Many papers propose sparse attention patterns that look good on paper—measured in theoretical FLOP reduction—but fail to deliver proportional wall-clock speedup because they ignore how GPUs actually execute attention. Unstructured sparsity (individual token-level masking), gather/scatter approaches, or patterns that don't align with GPU tile boundaries all suffer from memory access overheads that can eat most or all of the theoretical savings. The paper demonstrates this concretely: FlexAttention with an identical block-sparse mask achieves only 2× speedup over dense FlashAttention-3, compared to VSA's ~7×, purely because of kernel implementation quality.
VSA's innovation here is not any single kernel optimization but the tight co-design between the attention algorithm and the GPU execution model. The cube partitioning is chosen specifically so that each cube maps to one GPU tile (the unit of work for a threadblock in FlashAttention). The token re-indexing ensures tokens within a cube are contiguous in memory, making tile loads efficient. The coarse stage operates at cube (tile) granularity, producing block indices directly consumable by the fine-stage block-sparse kernel with no translation overhead. The fused softmax + Top-𝒦 + index conversion kernel eliminates multiple HBM round-trips for the coarse attention matrix.
This co-design principle—choosing the sparsity granularity based on GPU tile efficiency rather than purely on modeling considerations—is what enables the theoretical 8× FLOP reduction to translate to ~6–7× actual speedup (Figure 4b). The paper is explicit that this involved a tradeoff: smaller tiles would improve critical-token prediction accuracy (Table 1(c) shows 64×16 tiles achieve better loss than 64×64), but they would reduce MFU by 2.26× (Table 1(d)), wiping out the speed advantage. The choice of (4,4,4) cubes giving 64×64 tiles is the Pareto-optimal balance point discovered through systematic measurement.
What makes this a genuine innovation rather than routine engineering is that the co-design changes the optimization landscape. It is not "design the best sparse attention pattern, then implement it efficiently"—it is "design a sparse attention pattern whose quality-efficiency tradeoff surface is shaped by GPU architecture from the start." The paper provides the first systematic exploration of this tradeoff for video DiTs, measuring MFU across tile sizes and showing exactly where the cliff is. This provides a template for future work: any new sparse attention method should report not just theoretical FLOP savings but measured MFU and wall-clock speedup at realistic sequence lengths, because the gap between theory and practice can be enormous.
5. Experimental Analysis
Evaluation Methodology
-
Dataset. The primary training dataset for ablation and scaling studies is the Vchitect-T2V-Dataverse dataset (Fan et al., 2025), processed into video latents using the pretrained VAE from Wan2.1. The standard training shape is
(16, 32, 32)(16 frames, 32×32 spatial latent), yielding 16,384 tokens per video. For sparse adaptation of Wan-1.3B, synthetic data is generated by Wan-14B at resolution 448×832 with 61 frames (latent shape16×28×52). For Wan-14B adaptation, synthetic data is generated at 768×1280 with 77 frames (latent shape20×48×80). For evaluation of adapted models, the VBench benchmark suite (Huang et al., 2024) is used, along with human preference studies on 200 randomly sampled MovieGen prompts (Polyak et al., 2025) for the 14B model. -
Base model(s). Three scales of video DiT are used. For ablations, a 120M-parameter model following the Wan2.1 architecture is trained from scratch with the GPT-NeoX initialization scheme. For scaling studies, models from 60M to 1.4B parameters are pretrained from scratch. For sparse adaptation, the publicly released Wan2.1-1.3B and Wan2.1-14B models (Wang et al., 2025) serve as starting checkpoints. The paper states the 120M model is chosen because it reaches compute-optimality at
4.5 × 10^20FLOPS—a 120M model with full attention outperforms a 60M model at the same budget, and further training yields diminishing returns—making it an efficient testbed for ablations. -
Metrics. The primary metric for scaling and ablation experiments is diffusion loss (the flow-matching loss, following Lipman et al., 2023 and Liu et al., 2023), measured on a held-out validation set. The paper uses loss rather than downstream generation metrics for scaling studies because loss is a more sensitive and compute-efficient proxy for model quality when comparing architectural variants across many training budgets. For sparse adaptation experiments, two evaluation frameworks are used: VBench (Huang et al., 2024), a comprehensive benchmark suite for video generation models covering multiple quality dimensions, and human preference studies where raters compare videos generated by different model variants side-by-side. For the Wan-14B evaluation, 200 prompts from the MovieGen benchmark are used, with human raters assessing preference between VSA and full-attention outputs.
-
Baselines. For the ablation studies in Table 1(a), five baselines are compared: (1) Spatial-Temporal attention, which alternates between spatial and temporal attention across layers, as used in OpenSora (Zheng et al., 2024), OpenSora-Plan (Lin et al., 2024), LaVie (Wang et al., 2025), and Latte (Ma et al., 2024); (2) Spatial-Full attention, a variant with four spatial layers and one full-attention layer every five layers, designed to mitigate the extreme sparsity of temporal-only attention; (3) Compressed KV, which pools only key and value tokens using
2×2×2average pooling (8× reduction) while keeping queries at full resolution, mimicking VSA's coarse stage with smaller pooling and no query pooling; (4) Strided Window attention, inspired by Swin Transformer, where spatial attention uses a temporal window of 2 and temporal attention uses a spatial window of 8; (5) Full attention (the dense baseline). For the sparse adaptation experiments, Sparse VideoGen (SVG) (Xi et al., 2025) serves as the training-free inference-time sparsity baseline, configured with attention sparsity of 82.5% (parameters fp0.03, fl0.025, s0.1) in the human evaluation. For the Wan-14B human evaluation, the baseline is the official full-attention Wan-14B model, also finetuned on the same synthetic data to ensure a fair comparison. -
Generation budget / compute accounting. Compute is measured in total training FLOPs, estimated as
6ND + attention_FLOPs, whereNis parameters,Dis tokens, and attention FLOPs follow the FlashAttention formulation4 · D · S · H_heads · d_head · 3.5 · L_layers(withSbeing sequence length). For sparse attention methods, the attention portion is multiplied by the fraction of retained key tokens (e.g.,𝒦B/Lfor VSA). All ablation models are trained under a fixed4.5 × 10^20FLOPS budget, which the paper determines to be compute-optimal for the 120M model scale through a grid search comparing 60M and 120M models. Scaling studies use budgets up to4 × 10^21FLOPS. For sparse adaptation, compute is measured in training steps (4,000 steps for Wan-1.3B and Wan-14B finetuning) and in wall-clock generation time (seconds per video on H100/H200 GPUs). The paper explicitly states that "each ablation job takes around 10 hours on 64 Nvidia H200 GPU," and the full set of experiments totals "around 90k H200 hours." -
Cross-validation / statistical protocol. For the scaling studies in Figure 2, each model configuration is trained once at each budget level; the paper does not report multiple random seeds or cross-validation for the pretraining experiments. For the VBench evaluation of Wan-1.3B, the benchmark suite provides multiple quality dimensions that are aggregated; the paper reports the overall VBench score without confidence intervals. For the human evaluation of Wan-14B, 200 prompts are used with side-by-side comparisons; the paper reports preference percentages (Figure 3(b)) without statistical significance tests. The paper does not describe a held-out validation protocol for the ablation studies beyond using validation loss as the metric, and the text does not specify whether the validation set is drawn from the same distribution as the training data or a separate benchmark.
Main Quantitative Results
Ablation Experiments: VSA Outperforms Fixed-Pattern Sparsity at Scale
The paper's foundational empirical claim appears in Table 1(a): at a compute-optimal training budget of 4.5 × 10^20 FLOPS, existing sparse methods outperform full attention (lower diffusion loss), but this advantage reverses when training is extended to 4 × 10^21 FLOPS—full attention catches up and surpasses every fixed-pattern sparse method. VSA is the only sparse method that maintains parity with full attention at the extended budget.
Specifically, at 4.5 × 10^20 FLOPS, the fixed-pattern methods achieve lower loss than full attention (Exp 5), demonstrating the standard finding that sparse attention is beneficial under limited compute. However, at 4 × 10^21 FLOPS, this pattern reverses: full attention outperforms spatial-temporal, spatial-full, compressed KV, and strided window attention. VSA (Exp 6) achieves loss comparable to full attention at both budgets, establishing that learned, data-dependent sparsity scales with compute while fixed patterns do not.
The mechanism behind this reversal is examined in Table 1(b), which decomposes VSA's design into two factors: pattern type (data-dependent vs. fixed local) and stage contributions (coarse output included vs. excluded). The key comparisons: data-dependent patterns consistently outperform fixed local patterns, both with and without the gated coarse output (Exp 7 vs. 8: data-dependent "F" beats fixed "L,F"; Exp 9 vs. 10: data-dependent "C&F" beats fixed "L,C&F"). The inclusion of the coarse stage output ("C&F" vs. "F"-only) provides a consistent improvement, confirming that the coarse attention output contributes meaningfully to the final representation, not just to critical-token prediction.
Three additional findings from Table 1(b): (1) adding a dedicated local attention stage ("L,C,F" in Exp 11) performs similarly to the simpler "C&F" architecture, indicating no benefit from explicit local modeling beyond what the coarse stage already provides; (2) explicitly excluding locally-selected cubes from the Top-𝒦 selection (Exp 12) or forcing their inclusion (Exp 13) yields similar performance, further confirming that local modeling is redundant; (3) the combination of data-dependent selection with gated coarse output (Exp 10) achieves the best performance among all configurations, establishing VSA's default architecture.
Tile Size Ablation: Smaller Tiles Improve Quality, Larger Tiles Improve Throughput
Table 1(c) and 1(d) together quantify the tradeoff between attention granularity and hardware efficiency. In Table 1(c), reducing the tile size from 256×256 (Exp 14) to 64×64 (Exp 17) consistently reduces training loss, with 64×16 tiles (Exp 18) achieving the lowest loss. The mechanism: smaller tiles allow the coarse stage to more precisely localize critical tokens, and the fine stage wastes less computation on non-critical tokens within selected cubes.
Table 1(d) reveals the hardware cost of this improved quality: as tile size decreases, Model FLOP Utilization (MFU) drops substantially. A 256×256 tile achieves the highest throughput, while 64×16 tiles run 2.26× slower than 64×64. The paper selects 64×64 (with cube shape (4,4,4)) as the default, accepting the slight quality loss compared to 64×16 in exchange for substantially better throughput.
A revealing experiment is Exp 16 in Table 1(c), which tests mismatched granularity between stages: the coarse stage uses (4,4,4) pooling (fine granularity) while the fine stage operates on larger (4,8,8) tiles. To align the granularities, an additional (1,2,2) pooling is applied to the coarse attention map. This configuration performs worse than the matched (4,4,4) configuration, demonstrating that both coarse-stage prediction accuracy and fine-stage attention granularity benefit from finer tile sizes.
Pooling Method and Critical Token Prediction Accuracy
Table 1(e) compares three pooling strategies for the coarse stage: mean pooling, max pooling, and convolution-based pooling. Mean pooling outperforms max pooling, and convolution causes training instability. The paper does not report exact loss differences, but the ranking is clear from the table: average pooling is the best, max pooling is second, and convolution is unstable. The instability of convolution-based pooling—where a learned 3D convolution with kernel size and stride of (4,4,4) replaces mean pooling—is attributed to the interaction between trainable pooling parameters and the non-differentiable Top-𝒦 selection, though the paper does not deeply analyze the failure mode.
The effectiveness of critical token prediction is quantified in Figure 5(e) through an analysis of the finetuned Wan-1.3B model. The metric is the sum of attention scores within the Top-32 cubes selected by the coarse stage, measured against a random-selection baseline that captures only 8% of attention mass (shown as a red plane). VSA maintains at least 60% prediction accuracy in most layers and timesteps, reaching up to 90% in some instances. This means the coarse stage successfully identifies the cubes containing the majority of attention weight, validating the design premise that mean-pooled cube-level attention scores are effective proxies for token-level attention importance.
The accuracy exhibits systematic variations across both layers and timesteps: accuracy increases monotonically with the denoising timestep, and shows a "zig-zag pattern" across transformer layers. The paper suggests these patterns could inform future work on adaptive sparsity levels (varying 𝒦 per layer or timestep), but does not pursue this in the current study.
Scaling Studies: VSA Produces a Strictly Better Pareto Frontier
Figure 2(a) shows the core scaling result: a 410M-parameter video DiT trained with VSA at 87.5% sparsity (𝒦 = 32 out of 256 cubes, on a (16,32,32) latent) achieves "nearly identical loss to full attention," while reducing attention FLOPs by 8× and total training FLOPs by 2.53×. The paper states that VSA and full attention produce "similar loss curve[s]," implying the loss-vs-FLOPs trajectories overlap when VSA's FLOP advantage is accounted for.
Figure 2(b) extends this to model sizes from 60M to 1.4B parameters, trained with compute budgets up to 4 × 10^21 FLOPS on 128 H200 GPUs with 16K sequence length. VSA "consistently produces a better Pareto frontier than full attention"—meaning that at any given training FLOP budget, VSA achieves lower loss than full attention, or equivalently, VSA achieves the same loss as full attention with 2.53× fewer FLOPs across all tested model scales. The paper fits parallel scaling-law curves (loss vs. FLOPs) for both VSA and full attention, noting that the fitted curves indicate VSA "maintains its 2.53× FLOPS reduction across scales."
Optimal Sparsity Level Depends on Sequence Length and Training Budget
Figure 2(c) investigates how the Top-𝒦 parameter (which controls sparsity) should be set. The headline finding is counterintuitive: 𝒦 = 32 performs well across sequence lengths of 8,192, 16,384, and 24,576 under a fixed 4.5 × 10^20 FLOPS budget, but at 61,440 tokens, 𝒦 = 16 (higher sparsity) outperforms 𝒦 = 32. However, when the training budget is increased to 1 × 10^21 FLOPS at 61,440 tokens, 𝒦 = 32 regains the advantage over 𝒦 = 16.
The paper interprets this as evidence that "optimal 𝒦 depends on both sequence length and training budget." More specifically, the ideal number of attended key-cubes increases with available compute, "converging to full attention with infinite resources." The practical implication is that practitioners cannot simply fix 𝒦 at one value across all training regimes; they must consider the interaction between sequence length, model size, and total compute when setting sparsity levels.
Sparse Adaptation: VSA Matches Full-Attention Quality at 1.7× Speedup
Table 3(a) reports VBench scores for Wan-1.3B with sparse adaptation. VSA achieves "even higher VBench score compared to the original Wan-1.3B," though the paper notes this may be partly attributable to training on synthetic data from the larger Wan-14B model rather than purely to VSA's architecture. To control for this, the paper also finetunes the original full-attention Wan-1.3B on the same synthetic data; the results show "all models perform closely on VBench, indicating that VSA can retain generation quality despite significant attention sparsity" at 91.2%.
The speedup numbers: for Wan-1.3B, full attention with torch.compile takes 31 seconds per generation on an H100; VSA reduces this to 18 seconds—a 1.7× end-to-end speedup. The attention time specifically is reduced by 6× (Figure 4(a)), but end-to-end speedup is limited by non-attention components (linear layers, cross-attention, VAE encoding/decoding). The paper notes that with VSA, "attention accounts for only 20% of runtime during both training and inference"—a significant shift from full attention's dominance.
For the Wan-14B model, finetuned at 720P resolution (latent 20×48×80) with 90% sparsity, the paper reports end-to-end generation time reductions from 1,274 seconds to 576 seconds, though the exact multiplier is not stated (it can be computed as approximately 2.2×). Human evaluation on 200 MovieGen prompts (Figure 3(b), middle panel) shows VSA preserves generation quality compared to the official full-attention model after finetuning.
Human Evaluation: VSA Outperforms Training-Free Sparsity at Higher Sparsity
Figure 3(b), top panel, compares VSA against SVG (Xi et al., 2025), a training-free inference-time sparsification method. Notably, VSA operates at 91.2% attention sparsity while SVG is configured at 82.5% sparsity—VSA is more sparse yet is preferred by human raters. This is a strong result because training-free methods typically degrade more gracefully at moderate sparsity levels than trainable methods that may overfit to sparse patterns; here, the trainable method achieves both higher sparsity and higher quality, validating the paper's central thesis that training with sparsity eliminates the train-test mismatch that limits post-hoc methods.
Sparse Distillation: First Demonstration of Sparsity + Distillation Compatibility
The bottom panel of Figure 3(b) shows results for Sparse-Distill, where a DMD2-style distillation pipeline replaces the student model's full attention with VSA at 80% sparsity. All distillation hyperparameters are held identical to the full-attention baseline. Human evaluation shows "no quality drop" while achieving a 50.9× total speedup relative to the baseline Wan-1.3B model, combining 3-step generation (from distillation) with VSA's attention speedup. This generates a 5-second video in approximately 5 seconds on a single H200.
The paper emphasizes this as the "first sparse attention method shown to be compatible with distillation," contrasting with prior approaches that "that exploit diffusion time-step redundancy [which] may fail in the extremely low-step distillation regime." The significance is that distillation operates at very few denoising steps (here, 3 steps), where each step's attention quality is critical—there is little redundancy across steps to mask attention errors—making it a stringent test of whether sparse attention truly preserves the attention function.
Kernel Performance: Near-Theoretical Speedup at Long Sequences
Figure 4(b) benchmarks VSA's block-sparse kernel against FlashAttention-3 (FA3) across sequence lengths with a fixed 87.5% sparsity and head dimension 64. At long sequence lengths (approaching 100K tokens), VSA achieves nearly 7× speedup over FA3, approaching the theoretical 8× maximum (given 87.5% sparsity). The paper reports that VSA achieves 85% of FA3's Model FLOP Utilization (MFU), meaning the sparse kernel executes its actual FLOPs at 85% of the GPU's theoretical peak throughput.
A direct comparison with general-purpose sparse attention frameworks is provided: "FlexAttention with an identical block-sparse mask (64×64 block size) achieves only a 2× speedup." This 2× vs. 7× gap demonstrates that purpose-built block-sparse kernels are essential for translating theoretical sparsity into actual speedup—general frameworks that support arbitrary sparsity patterns incur substantial overhead that negates most of the benefit.
Even after accounting for the coarse stage overhead, VSA maintains over 6× speedup over FA3. The coarse stage runtime is profiled in Appendix D (Table 3): for medium sequence lengths, Top-𝒦 selection dominates the coarse stage runtime, and the fused kernel for softmax + Top-𝒦 + index conversion provides only "modest improvements." The paper argues this overhead is acceptable because the coarse stage accounts for less than 0.2% of total attention FLOPs and only 14% of attention runtime at 87.5% fine-stage sparsity, with the fraction shrinking further at longer sequences.
Figure 4(a) shows runtime breakdowns for a single transformer block in Wan-1.3B and HunyuanVideo (Kong et al., 2024). VSA reduces attention latency by 6× in both models, shifting attention from the dominant runtime component to a minor fraction of total block time.
Critical Token Prediction Visualization
Figure 5 provides qualitative insight into VSA's learned behavior. Panels (a)-(f) visualize the block-sparse attention maps from different heads and layers of the finetuned 1.3B model. The patterns are "highly dynamic" and vary significantly across heads, even within the same layer. Some heads exhibit patterns that match known heuristics: local spatial attention (nearby tokens in the same frame), spatial-temporal attention (same spatial position across time), or attention concentrated within a single frame. Other heads show patterns that deviate from simple heuristics, displaying "highly global characteristics" or "a combination of local and global focus."
This diversity of learned patterns supports the paper's argument that data-dependent sparsity is necessary—no single fixed pattern could capture all the attention behaviors that different heads learn to specialize in. A head that tracks object motion across the full spatial extent of the video needs global attention; a head that refines texture details within a local region needs spatial attention; a head that maintains temporal coherence needs same-location attention across frames. VSA's learned Top-𝒦 selection allows each head to discover its own connectivity pattern during training.
Ablation Studies and Robustness Checks
-
Fixed-pattern vs. data-dependent sparsity at different training budgets. Table 1(a) demonstrates that the relative performance of fixed-pattern sparse methods and full attention flips as training compute increases. At
4.5 × 10^20FLOPS, fixed-pattern methods (spatial-temporal, spatial-full, compressed KV, strided window) achieve lower loss than full attention; at4 × 10^21FLOPS, full attention outperforms all of them. VSA maintains parity with full attention at both budgets. This ablation exposes a previously undocumented phenomenon: fixed-pattern sparsity imposes a representation bottleneck that limits returns to additional training. -
Coarse stage output contribution. Table 1(b) (Exp 7 vs. 9 for data-dependent, Exp 8 vs. 10 for fixed local) ablates whether including the coarse attention output
O_cin the final representation improves performance. The "C&F" configuration (coarse output included) consistently outperforms the "F"-only configuration (coarse used only for pattern prediction, output discarded). This validates one of VSA's key architectural distinctions from MoBA and BiFormer, which discard their coarse-stage outputs. -
Local attention redundancy. Table 1(b) (Exps 11–13) tests three strategies for incorporating explicit local attention: a separate local stage with gating (Exp 11: "L,C,F"), excluding locally-selected cubes from Top-𝒦 (Exp 12: "E"), and forcing inclusion of local cubes (Exp 13: "I"). All perform similarly to the simpler "C&F" architecture (Exp 10), indicating that explicit local attention modules are redundant—the coarse stage's dense cube-level attention already captures sufficient local context because neighboring cubes naturally have high coarse-attention scores.
-
Tile size vs. quality and throughput. Tables 1(c) and 1(d) provide a systematic sweep:
256×256tiles achieve the highest MFU but the worst loss;64×16tiles achieve the best loss but 2.26× lower throughput than64×64. Exp 16 specifically tests mismatched granularity (fine coarse prediction with coarse fine attention) and shows it underperforms matched(4,4,4)cubes, confirming that both stages benefit from finer tiles. The paper's choice of64×64represents a deliberate Pareto-optimal point balancing these factors. -
Pooling method for coarse stage. Table 1(e) compares mean pooling, max pooling, and convolution. Mean pooling wins, max pooling is second, and convolution causes training instability. The instability of convolution pooling—a learned 3D convolution with kernel size and stride equal to cube dimensions replacing the simple mean—is a negative result that the paper attributes to interaction between trainable pooling and non-differentiable Top-𝒦 selection, though this mechanism is not experimentally verified.
-
Sparsity level interaction with sequence length and compute. Figure 2(c) provides a non-obvious finding: the optimal 𝒦 is not monotonic in sequence length. At a moderate training budget (
4.5 × 10^20FLOPS),𝒦 = 32works well for 8K, 16K, and 24K tokens but is outperformed by𝒦 = 16at 61K tokens. At a higher budget (1 × 10^21FLOPS),𝒦 = 32regains the advantage. This shows that sparsity interacts with both sequence length and total compute in non-trivial ways. -
ReST^EM revision model degradation. This ablation appears in the context of revisions (Appendix K; referenced in the text but not reproduced in the provided paper content), where an attempt to further optimize the revision model with ReST^EM (Singh et al., 2024) backfires: additional sequential revisions "substantially hurt" performance, with fully sequential performance dropping substantially compared to the optimal ratio. This is a notable negative result indicating sensitivity of revision training to the data generation procedure, though the sparse adaptation context has its own negative findings (training instability without annealing).
-
VSA vs. SVG at matching sparsity. Figure 3(b) compares VSA at 91.2% sparsity against SVG at 82.5% sparsity. Despite VSA being more sparse, human raters prefer its outputs. This is a robustness check confirming that trainable sparsity's quality advantage over training-free methods is not simply due to using less aggressive sparsity—VSA genuinely achieves better quality at higher sparsity.
-
Coarse stage runtime profiling. Appendix D (Table 3) provides detailed runtime breakdowns showing that Top-𝒦 selection dominates coarse-stage overhead at shorter sequence lengths, and that the fused kernel (softmax + Top-𝒦 + index conversion) provides only "modest improvements." The paper argues this overhead is acceptable because it shrinks relative to fine-stage computation at longer sequences, but the profiling data itself is the robustness check showing that coarse-stage overhead is not negligible at all sequence lengths.
-
Sparse distillation with identical hyperparameters. Appendix C.6 specifies that Sparse-Distill uses "all DMD-related hyperparameters... held fixed, including the number of denoising steps, the generator update ratio, and the guidance scale for real-score model." This is an important ablation: it shows that VSA can be dropped into an existing distillation pipeline without hyperparameter re-tuning, and that the quality preservation is not due to careful re-optimization of the distillation process for sparse attention.
Critical Assessment
Does VSA Actually Achieve "No Drop in Diffusion Loss" at 2.53× Training FLOP Reduction?
The paper's central claim is that VSA "reaches a Pareto point that cuts training FLOPS by 2.53× with no drop in diffusion loss." The evidence supporting this comes from Figure 2(a), which shows a 410M model trained with VSA achieving "nearly identical loss" to full attention, and Figure 2(b), which shows VSA producing a "better Pareto frontier" across model scales.
The strength of this evidence is that it comes from pretraining from scratch across multiple model sizes (60M to 1.4B) with FLOP budgets up to 4 × 10^21—this is not a fine-tuning result that could be attributed to the pretrained model's existing quality. The consistency of the 2.53× factor across model scales (shown by the parallel fitted curves in Figure 2(b)) suggests the benefit is structural rather than scale-dependent.
However, the evidence has important limitations. First, all scaling experiments use a single latent shape (16×32×32, 16K tokens). The paper does not demonstrate that the 2.53× factor holds at longer sequence lengths (e.g., 100K tokens, which Section 1 identifies as the practical regime for video generation), where attention would dominate even more of the total FLOPs and the reduction factor should theoretically be larger. The Wan-1.3B adaptation uses a slightly longer sequence (~23K tokens, from 16×28×52), but this is a fine-tuning experiment, not a from-scratch training validation of the FLOP reduction claim.
Second, the metric is diffusion loss, not downstream generation quality (VBench, human evaluation). Loss is a standard proxy in scaling-law studies and the paper argues it is "more sensitive and compute-efficient," but loss improvements do not always translate to perceptual quality improvements in generative models. The paper does not close this gap: the sparse adaptation experiments (on Wan-1.3B and Wan-14B) do evaluate generation quality (VBench, human preference) and find parity, but these are fine-tuning results, not tests of the from-scratch pretraining claim. A from-scratch VSA-trained model evaluated on VBench would more directly validate that the loss parity translates to generation quality parity.
Third, the paper does not report confidence intervals, error bars, or multiple random seeds for the scaling experiments. Figure 2(a) shows a single training run per configuration. Given the scale of these experiments (each point represents a full model pretraining), multiple seeds per configuration would be expensive, but the absence of any variance quantification means the claim of "nearly identical loss" cannot be statistically evaluated.
Does VSA Really "Outperform" Fixed-Pattern Methods, or Only Match Full Attention?
Table 1(a) shows VSA achieving the lowest loss among all methods at the 4.5 × 10^20 FLOPS budget and matching full attention at 4 × 10^21 FLOPS. The paper frames this as VSA "outperforming" fixed-pattern methods. This is accurate at the lower budget, but at the extended budget, VSA's advantage is specifically that it does not degrade while fixed-pattern methods fall behind full attention. VSA does not beat full attention at the extended budget—it matches it. The paper's claim of a "better Pareto frontier" in Figure 2(b) should be understood as VSA being better than full attention at a given FLOP budget (because it achieves similar loss with 2.53× fewer FLOPs), not as VSA achieving strictly lower loss than full attention at all budgets.
The distinction matters because the paper's strongest rhetorical claim—that VSA shows "better scaling than full attention"—could be misinterpreted as VSA achieving lower loss at the same FLOPs. What it actually achieves is the same loss at 2.53× fewer FLOPs, which is a FLOP-efficiency improvement, not an asymptotic quality improvement. The paper is transparent about this in the numbers, but the language in Section 1 ("better Pareto frontier") and Section 3.2 ("superior performance compared to full attention") could be read as claiming a quality advantage that the data does not directly support for matched-FLOP comparisons at large budgets.
Is the 2.53× Training FLOP Reduction Actually Attributed Correctly?
The paper accounts for total training FLOPs as 6ND + attention_FLOPs, with attention FLOPs reduced by the sparsity fraction 𝒦B/L. This formula assumes that VSA's coarse stage overhead is negligible (<0.2% of attention FLOPs), which is justified by the math (64× reduction in sequence length means 4096× reduction in attention cost for the coarse stage).
However, the formula does not account for several potential sources of overhead in practice: (1) the token re-indexing (tile/untile) operations, which the paper notes in Appendix B can be moved to the transformer boundaries but which still consume some FLOPs and memory bandwidth; (2) the gate projection W_g computation, which adds a small matrix multiplication per attention head; (3) the block-index processing overhead in the fine-stage kernel, which the paper acknowledges reduces MFU to 85% (not 100%). These overheads are individually small, but together they mean the actual FLOP reduction may be slightly less than 2.53×. The paper does not provide an end-to-end measurement comparing actual training step time for VSA vs. full attention at matched model configurations—the 2.53× figure is derived from the FLOP accounting formula, not from wall-clock training time measurements.
This is significant because the paper's kernel benchmarks (Figure 4(b)) measure inference speedup, not training speedup. Training requires backward passes, which are typically more memory-intensive and may have different sparsity efficiency characteristics than forward passes. The paper does not benchmark training throughput for VSA vs. full attention, leaving open the question of whether the 2.53× theoretical FLOP reduction translates fully to 2.53× training wall-clock speedup.
Does the Sparse Adaptation Result Generalize Beyond Wan2.1?
All sparse adaptation experiments use the Wan2.1 architecture and Wan-1.3B/Wan-14B checkpoints. The paper does not test VSA adaptation on other video DiT architectures (e.g., HunyuanVideo, CogVideoX, MovieGen), limiting the evidence that VSA's adaptation protocol is architecture-agnostic. The paper does include HunyuanVideo in the kernel benchmarking (Figure 4(a)), but only for attention speed measurements, not for quality-preserving adaptation.
Additionally, the sparse adaptation uses synthetic data generated by a larger model (Wan-14B generates training data for Wan-1.3B; Wan-14B generates its own training data). The paper acknowledges this may contribute to the quality results: "We hypothesize training with synthetic data from a larger model may contribute to this boost." The fair comparison against full-attention Wan-1.3B also finetuned on the same synthetic data partially addresses this, but the synthetic data introduces a confounding factor—VSA might be particularly effective at learning from synthetic data (which may have different attention pattern characteristics than real data) in ways that wouldn't generalize.
What Is Not Tested: Key Missing Experiments
Several experiments that would substantially strengthen the paper's claims are absent:
-
VSA pretrained from scratch on a public benchmark and evaluated on generation quality. The scaling studies use only diffusion loss. An experiment training a VSA model from scratch and evaluating on VBench or human preference against a full-attention model at matched compute would directly validate that the loss parity translates to generation quality parity.
-
Scaling to longer sequences (100K+ tokens). The paper's introduction emphasizes that 100K tokens is the practical regime for video generation, but all scaling experiments use 16K tokens (with one data point at 61K in Figure 2(c) for the sparsity-level study). The Wan-14B adaptation reaches ~77K tokens (
20×48×80 = 76,800), showing inference speedup, but there is no from-scratch training comparison at this length. This is a significant gap because the paper's core motivation—that video DiTs inherently require very long sequences—implies that the most important regime to validate is the one where attention dominates most heavily. -
Training throughput (wall-clock) measurement. The FLOP reduction is computed theoretically; actual training step time for VSA vs. full attention is not reported. This matters because if VSA's sparse kernels have lower MFU than dense FlashAttention for training backward passes, the actual training speedup could be less than 2.53×.
-
Ablation of gating mechanism. The paper establishes that including the coarse output improves performance (C&F vs. F in Table 1(b)), but does not ablate the learned gating mechanism itself. What if
O_c + O_f(simple addition) were used instead of the gated combination? The gates add parameters and could be unnecessary if the model can learn to produce appropriately scaled outputs from each stage without explicit per-position gating. -
Sensitivity to cube shape. The default is
(4,4,4)—an isotropic cube. The paper does not test anisotropic cube shapes (e.g.,(8,4,4)for more temporal compression, or(4,8,4)for more spatial compression along one axis). Video data has different statistics along temporal vs. spatial dimensions, and the optimal cube shape might depend on the video's motion characteristics. -
Comparison against DSV (Tan et al., 2025). The paper positions DSV as the closest prior work and critiques its multi-stage training pipeline, but never empirically compares VSA against a DSV-trained model. This is a missing baseline that would directly validate the claim that VSA's end-to-end trainability is superior to DSV's separate predictor training.
What Claims Hold Conditionally and Where Do They Break?
-
"2.53× training FLOP reduction with no drop in diffusion loss" holds for the tested regime: Wan2.1 architecture, 60M–1.4B parameters, 16K sequence length,
(4,4,4)cubes,𝒦=32, and flow-matching loss. It has not been demonstrated at longer sequences, on other DiT architectures, or with downstream generation metrics. -
"VSA outperforms fixed-pattern methods" holds at the
4.5 × 10^20FLOPS budget but narrows to "VSA matches full attention and does not degrade like fixed-pattern methods" at the4 × 10^21FLOPS budget. The fixed-pattern methods still outperform full attention at small budgets, so for practitioners with limited compute, fixed patterns remain competitive. -
"Sparse adaptation preserves quality at 1.7× speedup" holds for Wan-1.3B finetuned on Wan-14B synthetic data and evaluated on VBench. The Wan-14B result (2.2× speedup on 720P video) is evaluated only by human preference on 200 prompts—a smaller-scale evaluation than VBench's multi-dimensional benchmark. The claim may not hold for models where the base checkpoint's attention patterns are fundamentally different from what VSA's fixed cube shape can capture.
-
"Attention accounts for only 20% of runtime" holds for the Wan-1.3B and HunyuanVideo configurations shown in Figure 4(a), at the inference sequence lengths and batch sizes tested. At much longer sequences (approaching 100K tokens), attention would again dominate if sparsity is fixed at 87.5%, because the attended fraction
𝒦B/Lshrinks withLbut the cost per attended token grows with the number of heads and layers. The paper does not explore whether sparsity should increase with sequence length to maintain the 20% target, which would connect to the sparsity scaling findings in Figure 2(c). -
"Sparse-Distill achieves 50.9× speedup with no quality drop" is a preliminary finding demonstrated on a single model (Wan-1.3B) at 80% sparsity with 3-step DMD2 distillation, evaluated by human preference (not a comprehensive benchmark like VBench). The paper frames this as a "pilot study," and the hyperparameters are held fixed from the full-attention distillation baseline, so the result may be sensitive to the specific distillation setup.
6. Limitations and Trade-offs
Difficulty Estimation Cost Is Unaccounted For in the Headline Efficiency Gains
The assumption or constraint. The compute-optimal framework described in the reference paper assumes that prompt difficulty can be estimated before strategy allocation, but the method used for this—generating 2048 samples per question and averaging either ground-truth correctness or PRM scores—is extraordinarily expensive. The paper explicitly acknowledges this gap in Section 3.2:
"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"
The difficulty estimation step consumes more compute than the largest test-time budgets studied (256–512 generations), yet this cost is excluded from all efficiency calculations.
The consequence. The reported 4× efficiency gains are computed after difficulty is known, without amortizing the cost of learning it. In a realistic deployment, total cost would be difficulty_estimation_cost + strategy_execution_cost, and the former could dominate the latter. This means the 4× figure represents an upper bound on achievable efficiency rather than a realized deployment gain. A practitioner implementing this system would find that the difficulty estimation overhead wipes out most or all of the reported savings unless they have a cheaper way to estimate difficulty—which the paper does not provide.
What evidence exists in the paper. The paper's own numbers reveal the scale of the problem: 2048 samples per question for difficulty estimation versus test-time budgets ranging from 4 to 512 generations. The estimation cost alone is 4–512× larger than the actual strategy execution cost at typical operating points. The cross-validation protocol (Section 3.2) uses oracle difficulty derived from 2048-sample pass@1 rates, confirming that the difficulty information used to select strategies comes from a source far more expensive than the strategies themselves.
Mitigation status. The paper acknowledges this limitation explicitly and suggests "future work on pretraining or finetuning models to directly predict difficulty of a question" (Section 8). No such model is developed or evaluated. The predicted (non-oracle) difficulty bins use the PRM's own scores rather than ground-truth labels, eliminating the need for answer labels but not eliminating the need for 2048 samples per question. Until a cheap difficulty estimator is developed and validated, the compute-optimal framework remains a proof of concept rather than a deployable system. The paper frames the difficulty estimation cost as an "exploration-exploitation tradeoff" (Section 3.2) but provides no analysis of how to balance this tradeoff in practice.
Hard Problems Remain Essentially Unsolved—Test-Time Compute Cannot Create Capability From Nothing
The assumption or constraint. The entire compute-optimal framework assumes that the base model's pass@1 rate on a problem is non-trivially above zero—i.e., the model can generate at least some correct solutions given enough samples. Problems where the base model's pass@1 is near zero fall into difficulty bin 5, and no test-time strategy helps on them.
The consequence. Across all methods studied—search, revisions, and their compute-optimal combinations—the hardest questions show near-zero improvement regardless of compute budget. In Figure 3 (right), bin 5 accuracy hovers at 1–3% for all methods and all budgets. In Figure 7 (right), bin 5 shows roughly 2–3% accuracy irrespective of the sequential-to-parallel ratio. In the FLOPs-matched comparison (Figure 9), the bin 5 scaling line is essentially flat near 0–5%.
This is a fundamental limitation, not a fixable engineering issue: test-time compute amplifies existing capability but does not create it. If the base model cannot produce a correct solution at any non-trivial rate, no amount of search or revision will help because there are no correct solutions in the proposal distribution to find or refine. For problems genuinely outside the model's training distribution—novel reasoning patterns, out-of-distribution math, tasks requiring knowledge the model lacks—the approach offers zero benefit. This means test-time compute is not a substitute for pretraining on hard problems, which is the regime where larger models are most valuable.
What evidence exists in the paper. The evidence is consistent and stark across multiple experiments. Figure 3 (right, bottom row) shows bin 5 accuracy flatlined across all search methods and budgets. Figure 7 (right, bottom row) shows the same pattern for revisions. The FLOPs-matched comparison in Figure 9 shows bin 5 curves hugging zero for both revisions and search, well below the 14× larger model's performance. The paper is candid about this in Section 7's summary: on the hardest questions, test-time compute provides essentially zero benefit, and pretraining is almost always more effective. The authors explicitly note that "some capabilities can only be acquired through pretraining, not recovered at inference time."
Mitigation status. The paper is transparent about this limitation but offers no solution beyond acknowledging it. The difficulty-conditioned policy routes hard problems to best-of-N (which is essentially random guessing at these accuracy levels), implicitly recognizing that no strategy helps. The paper does not explore whether different base model architectures, training procedures, or verifier designs could shift the difficulty distribution and bring more problems into the "solvable with test-time compute" regime. This limitation establishes a clear boundary condition for when the approach is applicable and when practitioners should instead invest in better pretraining.
Single Benchmark, Single Model Family—Generality Is Unverified
The assumption or constraint. All experiments use the MATH benchmark (Hendrycks et al., 2021) with PaLM 2-S* as the base model. The paper states it "believe[s] this model is representative of the capabilities of many contemporary LLMs" (Section 4), but this belief is untested. The MATH benchmark consists exclusively of competition-level math problems with clean, verifiable answers—a domain with specific properties (symbolic reasoning, exact answer matching, formal step structure) that may not generalize.
The consequence. Several aspects of the findings could be model-specific or domain-specific:
- The PRM's quality and over-optimization behavior depend on PaLM 2-S*'s output distribution. A model with different calibration, different error patterns, or different attention mechanisms might exhibit different difficulty-dependent scaling curves.
- The revision model's ability to learn from incorrect in-context examples depends on the base model's in-context learning capabilities, which vary substantially across model families (GPT-4 vs. Claude vs. LLaMA vs. PaLM).
- MATH problems have exactly one correct answer and admit step-by-step verification—this makes PRM training via Monte Carlo rollouts feasible. Tasks with ambiguous answers (creative writing, dialogue) or tasks where step-level correctness is ill-defined (code generation where partial programs don't execute) may not support the same PRM training approach.
A practitioner using a different model family or working in a different domain (code generation, scientific QA, logical reasoning) cannot assume the 4× efficiency gains or the specific difficulty-dependent strategy recommendations will transfer.
What evidence exists in the paper. The paper provides no cross-model or cross-domain validation. All ablation curves, scaling analyses, and FLOPs-matched comparisons use PaLM 2-S*. The test set of 500 MATH questions, further split by cross-validation into folds of ~50 questions per difficulty bin, provides limited statistical power for the strategy selection step. The paper does not report confidence intervals on the compute-optimal scaling curves, making it impossible to assess whether observed differences between strategies are statistically reliable at these sample sizes.
Mitigation status. The paper does not attempt to mitigate this limitation through additional experiments on other benchmarks or model families. The authors acknowledge the single-benchmark limitation implicitly by describing their model as "representative" without providing evidence for this representativeness. A practitioner evaluating whether to adopt this approach would need to replicate key findings (particularly the difficulty-dependent strategy rankings and the 4× efficiency figure) on their specific model and task distribution.
The 14× Larger Model Baseline Is Not Compute-Optimally Trained, Weakening the FLOPs-Matched Comparison
The assumption or constraint. The FLOPs-matched comparison in Section 7 scales the pretraining baseline by increasing model parameters 14× while holding training data fixed, following the LLaMA paradigm (Touvron et al., 2023). The paper explicitly acknowledges this departs from compute-optimal pretraining (Hoffmann et al., 2022), where both data and parameters are scaled equally:
"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work."
Additionally, the 14× larger model uses only greedy decoding—no majority voting, no best-of-N, no search of any kind—while the smaller model gets the full benefit of compute-optimal test-time strategies.
The consequence. A Chinchilla-optimal model trained with 14× more total FLOPs (scaling both parameters and data) would likely outperform the parameter-only-scaled model used as a baseline. The reported advantages of test-time compute over pretraining—e.g., +27.8% on easy questions at R ≪ 1—may shrink or reverse against a properly compute-optimal larger model. Similarly, giving the larger model even a modest test-time compute budget (say, best-of-8 or a simple verifier) would create a much stronger baseline that is never tested. The comparison is fundamentally asymmetric: the smaller model gets sophisticated inference-time optimization while the larger model gets none.
This means the paper's headline claim—that test-time compute can outperform a 14× larger model—is best understood as an existence proof that it is possible under specific conditions, not as a general law that test-time compute is 14× more efficient than pretraining. The actual efficiency ratio would be smaller against stronger baselines.
What evidence exists in the paper. The paper provides the FLOP accounting in Section 7 and the results in Figure 9. The choice of parameter-only scaling is documented and justified as "representative of a canonical approach," but no comparison against a compute-optimally trained baseline is provided. The use of greedy decoding for the larger model is not explicitly discussed as a limitation, but it is clear from the experimental description: the larger model's performance is a single point estimate (greedy decoding), while the smaller model's performance is the result of extensive test-time optimization.
Mitigation status. The paper partially mitigates this by testing across multiple values of R (the inference-to-pretraining token ratio: 0.16, 0.79, 22), showing that the advantage of test-time compute varies with deployment scenario. The authors acknowledge the compute-optimal pretraining limitation explicitly and leave it to future work. However, the paper does not test even a simple test-time compute baseline for the larger model (e.g., majority voting with 64 samples), which would be a minimal fairness improvement over greedy decoding.
Verifier Over-Optimization Is a Hard Ceiling, Not a Solved Problem
The assumption or constraint. The paper's compute-optimal policy mitigates verifier over-optimization by routing easy problems away from aggressive search and toward simpler strategies like best-of-N. However, it does not solve the underlying problem: the PRM remains imperfect, and aggressive optimization against it eventually finds solutions that score highly under the PRM but are factually incorrect.
The consequence. On medium-difficulty problems where beam search is deployed (the regime where search is most beneficial), over-optimization still limits the scaling ceiling. The beam search curves in Figure 3 (right, bin 3) improve with budget initially but flatten well before the budget is exhausted. Lookahead search—the most powerful optimizer—paradoxically performs worst overall (Figure 3, left) because its aggressive optimization amplifies PRM errors. Qualitative examples in Appendix M show search producing degenerate outputs: repetitive low-information steps at the end of solutions and overly short 1–2 step solutions that the PRM incorrectly scores highly.
This means the compute-optimal approach is fundamentally bounded by verifier quality. Improving the PRM—through better training data, adversarial robustness, ensemble methods, or architectural changes—would likely shift the difficulty thresholds, change which strategies are optimal, and raise the ceiling on achievable performance. The current results are specific to the verifier quality achievable with the Monte Carlo rollout training procedure described in Appendix D.
What evidence exists in the paper. The over-optimization evidence is concrete and multifaceted: beam search degrading easy-problem performance at high budgets (Figure 3, right, bin 1), lookahead search underperforming simpler methods (Figure 3, left), and qualitative examples in Appendix M (Figures 29 and following) showing search finding solutions that score highly under the PRM but are incorrect or degenerate. The paper explicitly identifies over-optimization as "the primary bottleneck preventing unbounded improvements from additional compute" (discussed in the innovation analysis in Section 4 of the paper summary).
Mitigation status. The paper's compute-optimal policy partially mitigates this by avoiding aggressive optimization on problems where the verifier is unreliable (easy problems) and using it only where the verifier signal provides genuine guidance (medium problems). However, this is a workaround, not a fix. The paper does not explore techniques for improving verifier robustness—adversarial training, ensemble methods, calibration, or constrained search that penalizes deviations from the base model's output distribution. The paper identifies verifier robustness as "the key bottleneck for further scaling test-time compute" and suggests it as a direction for future work, but provides no experimental progress on this front.
The Revision Model Has a 38% Correct-to-Incorrect Reversion Rate with No Principled Solution
The assumption or constraint. The revision model is trained exclusively on sequences where all in-context previous answers are incorrect, followed by a correct target. This means the model never sees examples of what to do when the current answer is already correct. At test time, when a revision chain produces a correct answer, the model has no learned behavior for preserving it and frequently "revises" it into an incorrect answer.
The consequence. The paper reports that approximately 38% of correct answers produced during a revision chain get converted back to incorrect answers in the subsequent revision step (Section 6.1). This is a direct and severe consequence of the training data construction. The system mitigates this with majority voting or verifier-based selection across the entire chain—picking the best answer from any point rather than always taking the last revision—but these are post-hoc patches. They waste compute by continuing to generate revisions after a correct answer has already been produced, and they rely on the selection mechanism (verifier or majority vote) to correctly identify which answer in the chain is best.
A more principled solution—such as training the model to recognize when no revision is needed, or incorporating correct-to-correct trajectories in the training data—is not explored. This means the revision approach has an inherent inefficiency: roughly 38% of correct answers are lost to spurious revisions, requiring the system to generate more parallel chains or longer revision sequences to compensate, which directly reduces the effective compute efficiency.
What evidence exists in the paper. The 38% figure is reported in Section 6.1. The paper's mitigation strategy (within-chain selection via verifier or majority voting) is described in the same section. The ReST^EM experiment (Appendix K, Figure 16) provides additional evidence of revision fragility: attempting to optimize the revision model with RL-style training caused performance to degrade substantially with sequential revisions, suggesting the revision training procedure is sensitive to data generation methodology in ways that are not fully understood.
Mitigation status. The paper partially mitigates this through within-chain selection (majority voting or verifier-based selection across the revision chain), which recovers some of the lost correct answers. However, this is a workaround that increases effective compute cost (the chain must be longer to account for reversion) and relies on the verifier's ability to correctly identify which answer is best—which itself has limitations (the over-optimization problem described above). The paper does not explore training data modifications, architectural changes, or inference-time stopping criteria that would prevent the reversion problem at its source. This limitation is acknowledged but treated as an acceptable cost rather than a problem to be solved.
7. Implications and Future Directions
How This Work Changes the Landscape
VSA shifts the conversation around attention in video generation from "how can we tolerate sparsity as a post-hoc compromise?" to "how can we design sparsity as a first-class architectural primitive that improves scaling?" This is a conceptual reframing, not an incremental refinement. The dominant paradigm in video DiTs—train with full dense attention, then optionally substitute a sparse pattern at inference—is revealed as fundamentally limited. The paper's evidence in Table 1(a) makes this concrete: fixed-pattern sparse methods that look competitive at moderate training budgets actually diverge from full attention as training compute increases, because they impose a representation bottleneck that the model cannot overcome with more training. VSA's learned, data-dependent sparsity is the only sparse method that maintains parity with full attention at extended budgets.
This reframing matters because it changes what the field optimizes for. Rather than asking "what sparse pattern can we get away with?" (an approximation mindset), researchers should ask "what sparse architecture, when jointly optimized with model parameters, produces the best loss-vs-FLOPs Pareto frontier?" (a design mindset). The paper's scaling studies in Figure 2(b) demonstrate that this shift is not merely rhetorical—it produces a measurable improvement in the achievable tradeoff curve.
The work also reconciles a latent contradiction in the sparse attention literature. Prior work on fixed-pattern sparsity (spatial-temporal, sliding window) often reported positive results at moderate training scales, leading to claims that sparse attention was "good enough." Other work, including the same authors' prior STA paper (Zhang et al., 2025), acknowledged that post-hoc sparsity could erode quality when pushed aggressively. These weren't contradictory findings—they reflected different points on a scaling curve that no one had systematically mapped. VSA's scaling experiments reveal the full picture: fixed-pattern methods look good at small budgets because computational savings let you train on more data, but the bottleneck they impose becomes binding as training scales up. This explains why different papers reached different conclusions and provides a unified framework for evaluating any sparse attention method: does its loss-vs-FLOPs curve remain parallel to full attention as compute increases, or does it diverge?
The work also recalibrates research priorities for video generation. The paper demonstrates that attention is not merely an implementation detail to be optimized post-hoc—it is the primary bottleneck at both training and inference, and the specific architecture of the attention mechanism has first-order effects on scaling behavior. This implies that investment in better attention mechanisms (not just faster kernels for existing mechanisms) may have higher returns than investment in other architectural components (larger MLPs, more layers, better text encoders) when operating under fixed FLOP budgets. The paper's finding that, after applying VSA, "attention accounts for only 20% of runtime during both training and inference" (Section 3.3) suggests that the field may soon face a different bottleneck, shifting focus to other components of the DiT pipeline.
Finally, the hardware-algorithm co-design principle that VSA embodies—where the sparsity granularity is dictated by GPU tile efficiency, not purely by modeling considerations—provides a template for how future sparse attention methods should be evaluated. The paper's demonstration that FlexAttention achieves only 2× speedup with an identical block-sparse mask, compared to VSA's 6–7×, establishes that measured MFU and wall-clock speedup at realistic sequence lengths should be standard reporting requirements, not optional benchmarks. A method that claims "8× FLOP reduction" but delivers 2× actual speedup is not 8× faster in any practical sense.
Follow-Up Research This Work Enables
Extending scaling laws to explicitly include sparsity as an axis. The paper's most intriguing finding is in Figure 2(c): the optimal 𝒦 depends jointly on sequence length and training budget, and the relationship is not monotonic ("longer sequences need more 𝒦" is wrong under limited compute). A natural follow-up would systematically measure loss as a function of model size N, training tokens D, sequence length L, and sparsity level 𝒦 (or equivalently, sparsity fraction 1 - 𝒦B/L), then fit a parametric scaling law of the form L(N, D, L, 𝒦). The Chinchilla-style question would be: given a fixed FLOP budget, what is the compute-optimal allocation across model size, training tokens, and sparsity? This is a substantial engineering effort (requiring hundreds of pretraining runs across a grid of configurations), but the paper's demonstration that 𝒦 = 32 is robust across moderate scales while 𝒦 = 16 wins at extreme length/compute combinations provides the first data points for such a law. A successful scaling law would let practitioners answer: "For my budget, should I train a denser small model or a sparser large model?"
Combining VSA with other efficiency techniques to break the 20% attention floor. After applying VSA, attention drops to 20% of runtime (Figure 4a). The remaining bottlenecks are linear layers (QKV projections, output projection, MLP), cross-attention (text conditioning), and the VAE encoder/decoder. A natural next step is to apply complementary efficiency techniques: (1) weight quantization or low-rank factorization for the linear layers that now dominate runtime; (2) sparse or compressed cross-attention, since text tokens typically number in the hundreds while video tokens number in the tens of thousands—cross-attention may become the new attention bottleneck; (3) VAE distillation to reduce decoding cost. The specific experiment would measure end-to-end generation time with VSA + each technique individually, then combined, on a fixed model like Wan-1.3B, to identify which component becomes the new bottleneck and what combination achieves the best speedup-per-quality-degradation tradeoff. This is newly tractable because VSA removes attention as the dominant term, revealing the previously-hidden cost structure.
Training VSA from scratch and evaluating on VBench for direct quality validation. The paper's strongest claim—that VSA achieves "no drop in diffusion loss" at 2.53× training FLOP reduction—is validated only on validation loss, not on generation quality metrics. A critical follow-up would train two models from scratch at matched total FLOPs (one VSA, one full attention) on a public dataset, then evaluate on VBench across all quality dimensions (subject consistency, background consistency, motion smoothness, dynamic degree, aesthetic quality, imaging quality). The specific hypothesis to test: does the loss parity observed in Figure 2(a-b) translate to statistically indistinguishable VBench scores, or does VSA's sparsity introduce subtle quality degradations that loss doesn't capture? If VBench scores are identical, it strengthens the paper's central claim; if they diverge on specific dimensions (e.g., motion smoothness might suffer if long-range temporal attention is occasionally missed), it would reveal where sparse attention's approximations have perceptual consequences. This experiment would also address the paper's single-benchmark limitation by providing a multi-dimensional quality assessment.
Stress-testing VSA at 100K+ token sequences where the method's motivation is strongest. The paper's introduction emphasizes that "a 5-second 720p clip unfolds into more than 100K tokens" and argues that trainable sparsity is "more urgent" for video than for language. Yet all from-scratch training experiments use 16K tokens, with only a single data point at 61K in Figure 2(c) for the sparsity-level study. A direct stress-test would train VSA at 100K–150K sequence lengths (corresponding to 5–10 second 720p videos) and measure: (1) whether the 2.53× FLOP reduction factor holds or improves (attention occupies a larger fraction of total FLOPs at these lengths, so the reduction should be larger); (2) whether 𝒦 = 32 remains sufficient or needs to increase; (3) whether the coarse stage's 14% runtime overhead (Appendix D) shrinks to negligible levels as predicted. The paper's Wan-14B adaptation at ~77K tokens is suggestive but doesn't validate from-scratch training behavior. This experiment would directly test whether VSA delivers on its motivating use case.
Applying VSA to architectures beyond Wan2.1, especially models with different attention patterns. The paper demonstrates VSA on Wan2.1 and benchmarks kernels on HunyuanVideo, but quality-preserving adaptation is only shown for Wan2.1. Different DiT architectures use different attention configurations: some use separate spatial and temporal attention rather than unified 3D attention; some use windowed attention as their dense baseline; some use different head dimensions or numbers of heads. A systematic study would apply VSA's sparse adaptation protocol to CogVideoX (Yang et al., 2024), HunyuanVideo (Kong et al., 2024), and MovieGen (Polyak et al., 2025) checkpoints, measuring on each: (1) the quality-sparsity tradeoff curve (VBench or human eval at multiple sparsity levels), (2) whether the progressive sparsity annealing protocol requires architecture-specific tuning, and (3) whether the optimal cube shape (C_t, C_h, C_w) varies with the base model's attention pattern. A negative result—VSA working well on Wan2.1 but poorly on HunyuanVideo—would reveal that learned sparsity patterns are architecture-dependent and that the method's generalizability is narrower than claimed.
Adaptive per-layer and per-timestep sparsity allocation. The paper's inspection of VSA in Figure 5(e) reveals that critical-token prediction accuracy varies systematically across layers (zig-zag pattern) and increases monotonically with denoising timestep. This suggests that uniform sparsity (𝒦 = 32 everywhere) is suboptimal: early denoising steps, where the video structure is still being established, might need denser attention, while later refinement steps could tolerate more sparsity. Similarly, layers that show low prediction accuracy might benefit from larger 𝒦 to compensate for missed critical tokens. A concrete experiment would profile the full-attention version of a pretrained model to compute the "attention entropy" per layer and per timestep (as a measure of how concentrated vs. distributed the attention weights are), then train a VSA model where 𝒦 varies per layer and per timestep proportionally to this entropy. The hypothesis is that adaptive sparsity would recover some of the residual quality gap between VSA and full attention (visible as small differences in the loss curves in Figure 2(a)) by allocating the sparse attention budget where it matters most.
Practical Applications and Downstream Use Cases
Cost-efficient video generation API services. For companies offering video generation as a cloud service (e.g., Runway, Pika, Kling), the primary cost drivers are GPU hours for inference and the number of GPUs needed to serve peak demand. The paper's Wan-1.3B results show a 1.7× end-to-end latency reduction (31s → 18s) and the Wan-14B results show a 2.2× reduction (1,274s → 576s) with quality preservation. For an API serving millions of generations per month, a 1.7–2.2× throughput improvement on the same hardware translates directly to serving 1.7–2.2× more customers per GPU, or reducing GPU fleet size by 40–55% for the same load. The key practical consideration is that VSA requires fine-tuning the model on synthetic data (4,000 steps, approximately 12 hours on 32 H200 GPUs for Wan-1.3B per Appendix C.5), which is a one-time cost amortized over all subsequent inference. The sparse adaptation protocol's use of progressive sparsity annealing means this fine-tuning is stable and reproducible, making it suitable for production pipelines where model updates must be reliable.
Open-source video generation on consumer hardware. The current landscape for open-source video generation is stark: Wan-1.3B requires an H100-class GPU to generate a 5-second video in 31 seconds. After VSA adaptation, this drops to 18 seconds on the same hardware—still not real-time, but approaching the threshold where a user can iterate on prompts with acceptable turnaround. More significantly, VSA's attention speedup (6× for the attention component) reduces the attention-related memory bandwidth pressure, potentially making the model runnable on GPUs with less memory bandwidth (e.g., RTX 4090, which has ~1 TB/s vs. H100's ~3.35 TB/s) where attention was previously the bottleneck. The paper doesn't benchmark consumer GPUs, but the reduced attention footprint means the model's memory requirements are dominated by parameters and activations rather than attention matrices—this shifts the hardware requirements from "datacenter GPU with high memory bandwidth" toward "consumer GPU with sufficient VRAM," expanding who can run these models locally.
Video generation for real-time or interactive applications. The Sparse-Distill result—50.9× total speedup, generating a 5-second video in ~5 seconds on an H200—represents a threshold where video generation approaches real-time interactivity. A user could type a prompt, see a video in roughly the time it takes to watch it, then refine their prompt and generate again. This is not yet "real-time" in the strict sense (generation time equals video duration, so there's a 5-second latency between prompt and completed video), but it is fast enough for iterative creative workflows where a user generates, watches, adjusts, and regenerates. The practical deployment scenario is a creative tool where VSA reduces the attention cost and distillation reduces the number of denoising steps, with both techniques combining multiplicatively (VSA's 1.7× and distillation's ~30× from 50-step to 3-step generation). The paper's key finding that "all sparse distillation losses and hyperparameters are kept identical to full-attention distillation" (Appendix C.6) means practitioners can adopt both techniques without complex hyperparameter co-optimization—each can be tuned independently and then combined.
Cost reduction for video generation model training. For research labs and companies training video DiTs from scratch, the paper's scaling results (Figure 2(b)) imply that training a model with VSA using X FLOPs achieves the same loss as training a full-attention model with 2.53X FLOPs. For a lab with a fixed compute budget (e.g., a grant providing 10,000 GPU-hours), this means they can either train a larger model or train for more steps—either way, they get a better model for the same cost. The paper's experiments totaling "around 90k H200 hours" would have cost roughly 3.6× more (≈228k H200 hours) if all scaling experiments had used full attention. For an organization deciding whether to adopt VSA for their next training run, the key practical question is whether the engineering effort of implementing VSA's cube partitioning, two-stage attention, and block-sparse kernel is justified by the FLOP savings. The paper's open-source code release mitigates this: a team can adopt VSA by modifying their existing DiT training code, not by writing custom CUDA kernels from scratch.