Can sparse attention close the quality gap in diffusion transformers?

The gap has three separable sources

Sparse attention is the standard way to cut the quadratic cost of diffusion transformers on long video and high-resolution 3D sequences, but aggressive sparsity usually degrades generation quality. The paper asks where that quality gap actually comes from, rather than assuming it is only a selector-accuracy problem.

Using controlled oracle comparisons at matched attention densities, the authors compare vanilla block-sparse attention (BSA) with three oracles that maximize retained attention mass using exact dense-attention probabilities: a block oracle that selects KV blocks per query block, a token oracle that selects individual KV tokens but shares one selection within each query block, and a per-query oracle that lets every query choose its own KV tokens.

Comparing these oracles isolates three sources of the gap: structural binding error, caused by binding KV tokens into blocks and forcing queries in the same group to share one selection; selection error, caused by approximate scoring such as mean pooling; and discarded tail error, caused because omitted tokens still carry nonzero attention mass. This decomposition, not a single accuracy number, is the paper's diagnostic contribution.

Mean-pooled block scoring misses high-mass tokens

Vanilla BSA scores each query-block/KV-block pair by the dot product of mean-pooled query and key representations, then keeps the highest-scoring KV blocks at a given attention density. The paper shows this pooled score does not preserve average attention probabilities: the softmax of the pooled query against keys is not equal to the mean over queries in the group of the true attention weights.

That mismatch matters because pooled scoring can miss KV tokens that carry high cumulative attention mass for the group. The oracle comparison confirms both constraints cause losses: moving from the block oracle to the token oracle with shared query blocks removes KV binding, and moving further to the per-query oracle removes query binding, with accuracy gains at each step. At matched realized densities, finer KV granularity and grouping more similar queries both improve fidelity.

Token-level selection needs tile-aligned query groups

MC-Sparse removes KV binding by selecting individual KV tokens for each query group across block boundaries, gathering selected keys and values online from their indices without repacking them into blocks. Because multi-head attention uses head-specific keys and values, efficient tiles must batch queries within each head, so the method forms groups of C similar queries that share one KV selection, with C matched to the kernel's query tile size.

Minimizing within-group variance of queries is equivalent to maximizing between-group variance, but standard clustering such as k-means does not guarantee equal-size groups, and hard balance constraints add assignment overhead and complicate GPU parallelization. The authors adopt a median-split version of principal direction divisive partitioning that recursively projects each group onto its principal direction and splits ordered queries at the median until every group contains exactly C queries, preserving tile alignment while promoting within-group similarity.

Exact token-level selection requires attention scores, which FlashAttention does not expose, so the authors implement a two-pass kernel for exact KV selection and a token-sparse attention kernel in CuTeDSL. This is the engineering cost that makes the finer granularity practical rather than merely accurate.

Reused exact selection beats recomputed mean pooling

Recomputing exact attention scores at every denoising step would be too costly without a trained indexer, so MC-Sparse exploits temporal stability: it performs exact selection and query grouping only at anchor steps and reuses the resulting query groups, KV indices, and dense-sparse output residual at subsequent reuse steps. Current queries, keys, and values are still recomputed at every step.

The paper reports that exact selections from earlier steps outperform mean-pooled selection recomputed at every step, and that attention patterns remain stable across nearby denoising steps. This is the empirical basis for amortizing both grouping and selection rather than approximating them at each step.

Anchor steps are the last step of a dense warm-up or refresh segment. The reported configurations use anchor steps {10, 26} for the video models and {1} for the 3D model, with warm-up layers and steps specified per model, such as 1/52 warm-up layers and 10/49 warm-up steps for Minimax-H3-Base.

Residual correction recovers the discarded tail

Discarded tokens typically have small attention weights but still make nonzero contributions, leaving an output gap to dense attention even with per-query oracle selection. Block-sparse methods approximate those contributions using block-mean representations or block-level statistics, but token-level selection cuts irregularly across blocks, preventing direct reuse of those summaries.

MC-Sparse instead exploits the temporal dimension: the dense-sparse residual varies less across denoising steps than the attention output itself, and exact selection further reduces that variation compared with mean-pooled selection. At an anchor step the method records the residual between dense and sparse attention outputs, and at a reuse step it adds that cached residual to the sparse attention output. The residual follows the same refresh schedule as the cached groups and KV indices, and forming it adds one sparse-attention evaluation plus an elementwise subtraction at each anchor.

Speedups come with negligible measured quality loss

Across video and 3D generation models, MC-Sparse achieves higher fidelity to dense-attention outputs and larger denoising speedups than the evaluated sparse-attention baselines, without visible quality degradation. Relative to dense attention, it delivers a 1.80x denoising speedup on Minimax-H3-Base and a 2.32x speedup on 3D asset generation, both with negligible quality loss.

These are denoising speedups relative to dense attention, not end-to-end generation speedups, and they are tied to the specific model configurations and warm-up schedules listed in the paper. The fidelity claim is stated as fidelity to dense-attention outputs, which is a different measure from perceptual quality or task accuracy, so the two should not be substituted for one another.

The paper also reports that when GPU memory is insufficient, inactive cache entries can be offloaded to CPU memory and prefetched ahead of use, overlapping transfers with computation to hide latency. This is a deployment detail for the caching scheme rather than a change to the accuracy argument.

The oracle analysis bounds what sparsity can recover

The three-error decomposition is derived from oracle comparisons that maximize retained attention mass using exact dense-attention probabilities at matched densities. That makes it a diagnostic upper-bound analysis: it shows how much of the gap is attributable to structural binding, selection error, and discarded tail mass, but the oracles themselves are not deployable selectors.

MC-Sparse is training-free and reuses cached metadata across denoising steps, so its behavior depends on attention stability between anchor and reuse steps. The reported anchor and warm-up configurations are model-specific, and the speedup figures are tied to those settings rather than to arbitrary sparsity levels or sequence lengths.

The evaluation covers video and 3D asset generation models. The paper does not establish that the same error decomposition or the same reuse schedule transfers to other diffusion transformer tasks, other sparsity targets, or training-time use.

Compare sparse methods at matched realized density

The most useful takeaway is a comparison rule rather than a single number: sparse-attention methods should be compared at matched realized attention density, and their quality gap should be attributed separately to selection granularity, scoring accuracy, and discarded tail mass. MC-Sparse's contribution is showing that all three matter, and that token-level selection plus residual compensation addresses the parts that block-level selection cannot.

For readers deciding whether this approach applies to their setting, the paper's own variables are the decision inputs: whether the workload is long-sequence video or high-resolution 3D generation, whether the model can tolerate a dense warm-up segment and periodic anchor steps, and whether the target sparsity is high enough that discarded tail mass becomes the binding constraint. The reported 1.80x and 2.32x denoising speedups are the reference points for those specific configurations, not a general guarantee.

Sources used in this answer