Source-linked AI summary
Partition the Support, Reconstruct the Residual: Training-Free Sparse Attention for Video Generation and World Models
Pardis Taghavi, Reza Langari, Gaurav Pandey
TL;DR
Training-free sparse attention must address both disagreement in shared query routes and errors from skipped interactions. SparsePR couples response-based partitioning with probe-fitted residual reconstruction, reducing reconstruction error while preserving generation quality at 22.0–26.0% executed-pair density and delivering 1.48×–2.61× speedups.
Problem
Per-query attention concentration and retained mass do not determine shared-route support or post-softmax error, limiting executable sparse-attention design.
Method
SparsePR couples response-based query/K/V partitioning with probe-fitted, call-specific affine reconstruction of skipped-interaction residuals.
Results
1.48×–2.61× end-to-end speedups at 22.0–26.0% realized executed-pair density accompany closely matched dense benchmark quality across four video models.
Takeaways & Limitations
Effective sparse execution depends on both partition geometry and residual reconstruction, with probe fitting providing the largest fidelity gain.
Takeaways & Limitations
Probe-fitted reconstruction restricts its feature-dependent correction to output directions observed in exact probe residuals and does not assume globally low-rank residuals.
Abstract
from arXiv · showhide
Training-free block-sparse attention can accelerate video transformers, but row-wise attention concentration does not by itself specify an executable sparse operator. Queries sharing a block route may have poorly overlapping supports, while retained attention mass alone does not determine the post-softmax error from skipped interactions. We show that partition geometry affects both pooled support and the predictability of the remaining residual from the sparse output. We introduce SparsePR, which combines Response-Coupled Partitioning with Probe-Fitted Residual Reconstruction. Sampled-query key responses form paired K/V groups, whose centroids induce query-response coordinates for shared routing. A small set of exact query rows then calibrates a call-specific affine correction from the sparse output within the output subspace observed in the probe residuals. Across four heterogeneous video generation and world models, SparsePR consistently reduces attention-reconstruction error. Ablations show that probe fitting accounts for most of this reduction, while response-coupled partitioning lowers hard-drop error and improves reconstruction under a finite probe budget. SparsePR preserves generation quality at 22.0-26.0% realized executed-pair density while achieving 1.48x-2.61x end-to-end speedups. Project page: https://pardistaghavi.github.io/SparsePR-website/
1 INTRODUCTION
Long spatiotemporal video tokens make quadratic self-attention a major inference bottleneck, while attention concentration alone does not define accurate executable block sparsity. SparsePR addresses this by coupling response-based partitioning with probe-fitted residual reconstruction.
- Motivation: Long, high-resolution video sequences make quadratic self-attention a major inference bottleneck, motivating training-free sparse attention.Video generation and world models process long spatiotemporal token sequences.
- Problem: Per-query attention concentration does not determine shared-route support, and retained attention mass alone does not determine post-softmax error.Under renormalized hard drop, error also depends on the difference between omitted- and retained-support outputs.
- Problem: Partition choice affects both executable block-sparse routing and how much post-softmax residual can be represented as an affine function of the sparse output.The analysis separates query partitioning, paired K/V partitioning, routing, and skipped-interaction rules.
- SparsePR: SparsePR combines Response-Coupled Partitioning with Probe-Fitted Residual Reconstruction in a training-free sparse-attention framework.Sampled queries define key-response coordinates for paired K/V grouping, whose centroids define query-response coordinates for shared routing.
2 RELATED WORK
Prior work accelerates video-model attention through structured sparsity, executable token partitioning, and approximations or reconstruction of skipped interactions. SparsePR focuses on full-sequence attention, where partitioning and skipped-interaction treatment determine executable density and post-softmax error.
- Sparse Attention for Video Generative Models and Video World Models: Training-free sparse attention for video DiTs reduces inference cost through structured sparse patterns, online block selection, or current-call support estimation.Representative methods exploit spatial, temporal, layer, or head structures.
- Sparse Attention for Video Generative Models and Video World Models: Autoregressive video models use sparse attention for causal K/V cache compression and approximate retrieval, whereas this work studies full-sequence attention.The setting emphasizes how partitioning and skipped-interaction treatment affect executable density and post-softmax error.
- Token Partitioning for Executable Block Sparsity: Earlier executable block-sparsity methods organize tokens with LSH, online k-means, query clustering, asymmetric clustering, or learned permutations.For video DiTs, SVG2 clusters query and key activations and permutes the resulting groups into contiguous layouts.
- Attention Approximation and Residual Reconstruction: Attention approximations use random features, landmarks, sparse approximations, or low-rank approximations, while visual-generation methods approximate or recover skipped-interaction contributions.Examples include reusing softmax statistics, pooled corrections, and multiresolution K/V representations.
3 SPARSE ATTENTION FORMULATION AND OBSERVATIONS
This section formalizes block-sparse attention as a combination of partition construction, route selection, and skipped-interaction treatment, then identifies three dependencies governing approximation error. Shared-route support can expand beyond per-query support, retained mass does not determine post-softmax error, and partition geometry changes residual representability.
- 3 SPARSE ATTENTION FORMULATION AND OBSERVATIONS: Block-sparse attention is determined by partition construction, route selection, and the treatment of skipped interactions, which create distinct approximation-error sources.The formulation distinguishes executable sparsity from the structural choices that determine its approximation.
- 3.1 STRUCTURAL OBSERVATIONS ON SPARSE ATTENTION: O1: Per-query sparsity does not determine shared-route support because pooled support depends on overlap among individual query supports.The pooled-support diagnostic isolates query-side support disagreement before further expansion from selecting K/V groups.
- 3.1 STRUCTURAL OBSERVATIONS ON SPARSE ATTENTION: At α = 0.9, Wan2.2 median support increases from 6.2% per query to 22.9% after pooling across eight queries.The vertical displacement above the diagonal measures support expansion caused by limited overlap among key supports.
- 3.1 STRUCTURAL OBSERVATIONS ON SPARSE ATTENTION: O2: Retained attention mass alone does not determine output error, because comparable retained masses can yield different post-softmax residual norms.Error generally decreases as retained mass increases, but model-dependent variation remains at comparable mass levels.
- 3.1 STRUCTURAL OBSERVATIONS ON SPARSE ATTENTION: O3: Partition geometry changes residual representability, so residual magnitude alone does not indicate affine predictability from the sparse output.The residual decomposes into an affine-feature-space component and an orthogonal component.
- 3.1 STRUCTURAL OBSERVATIONS ON SPARSE ATTENTION: The affine-versus-orthogonal decomposition is an oracle diagnostic, because computing it requires the full dense residual and is not part of sparse execution.The projected term is the best least-squares affine approximation of the residual from the sparse output.
- 3.1 STRUCTURAL OBSERVATIONS ON SPARSE ATTENTION: Across four models, response-coupled partitioning increases the affine-explainable fraction by 3.2 to 14.9 percentage points and reduces affine-orthogonal energy to 0.285×–0.653× the semantic-partition baseline.The comparison uses matched group counts, routing, renormalized hard drop, and realized pair density.
4 METHOD
SparsePR combines response-coupled partitioning with probe-fitted residual reconstruction in a single, training-free attention procedure. It builds routing groups from current-call response geometry, then calibrates a call-specific affine correction using exact probe rows and residual-derived output directions.
- 4 METHOD: SparsePR combines Response-Coupled Partitioning and Probe-Fitted Residual Reconstruction to address shared-support and residual-error effects.The two components operate online and without training, followed by sparse execution with exact-pair accounting.
- Response-Coupled Partitioning: Response-coupled partitioning groups paired K/V tokens by sampled-query response profiles, then clusters queries using responses over the resulting K/V groups.The coupled, asymmetric construction uses low-dimensional response coordinates and performs K/V and query clustering in a single pass without alternating refinement.
- Probe-Fitted Residual Reconstruction: Probe-Fitted Residual Reconstruction uses M ≪ Nq exact query rows to fit an affine residual estimator from sparse outputs.Probe indices are stratified across query groups, with coverage weighting and ridge regularization supporting calibration from a limited probe set.
- Probe-Fitted Residual Reconstruction: The reconstruction restricts the feature-dependent correction to output directions observed in the exact probe residuals.The basis comes from leading right singular vectors of the weighted, centered probe-residual matrix; this restriction does not assume the complete residual matrix is globally low rank.
- Sparse Execution: Sparse execution permutes tokens into group-major layouts, selects cells under a query–key pair budget using true cell costs, and evaluates retained interactions exactly with row renormalization.The implementation uses FlashInfer’s variable-block sparse-attention primitive and restores the output to the original query order.
5 EXPERIMENTS
Experiments evaluate SparsePR across four heterogeneous video-generation and physical-world models under matched dense and sparse protocols. They measure fidelity, generation and physical-world quality, efficiency, and ablation behavior.
- Implementation details: 22% executed-pair targets apply to HunyuanVideo, Wan2.2, and Cosmos-Predict2.5, while Cosmos3-Nano uses 26%.These targets include probe rows; experiments use BF16 on a single NVIDIA H100 GPU.
- Models and benchmarks: Experiments cover four models spanning text-to-video, image-to-video, and physical-world generation and prediction.The evaluated models are HunyuanVideo-13B, Wan2.2-I2V-A14B, Cosmos-Predict2.5-14B, and Cosmos3-Nano-16B.
- Evaluation protocol: Dense and sparse runs use identical conditioning, preprocessing, seeds, sampling, inference, guidance, resolution, and frame-count settings.Reproduced baselines use the same samples, while prior-work results are marked with †.
- Metrics: Four metric groups assess dense-reference fidelity, generation quality, physical-world performance, and efficiency.Reported metrics include PSNR, SSIM, LPIPS, ImgQual, subject consistency, PBench Quality, executed-pair density, end-to-end FLOPs, and generation speedup.
- Ablations: Table 2 compares partitioning and reconstruction at matched routed-pair density, while Tables 3 and 4 test density sensitivity and probe-selection or estimator choices.Table 2 reports mean and p99 normalized attention-output error, with lower values preferred; Table 4 uses the response partition with M = 64 and r = 16.
6 CONCLUSION … B.2 QUERY-RESPONSE COORDINATES FROM K/V GROUPS
SparsePR treats training-free block-sparse attention as an executable operator whose partitioning affects shared support and skipped-interaction residuals. It combines response-based grouping with call-specific probe-fitted residual correction, with probe fitting providing the largest fidelity improvement across four video generation and world models.
- 6 CONCLUSION: SparsePR combines Response-Coupled Partitioning with Probe-Fitted Residual Reconstruction to reduce attention-reconstruction error across four video generation and world models.The conclusion identifies probe fitting as providing the largest fidelity improvement.
- A PROBE-FITTED RESIDUAL RECONSTRUCTION DETAILS: The final estimator regresses the residual from the observable sparse output using a call-specific affine model for the active attention call.Query-response coordinates guide partitioning and probe selection, but are not residual-regression features.
- A.1 HARD-DROP RESIDUAL AND AFFINE MODEL: The exact post-softmax residual depends on omitted attention mass and omitted-support output, which are unavailable during sparse execution.This motivates estimating the residual from quantities observed during the sparse call.
- A.2 QUERY-GROUP-STRATIFIED EXACT PROBE SELECTION: M = 64 ≪Nq exact probe rows are evaluated against all keys, contributing M/Nq executed-pair density of 0.054% to 0.145%.Queries are selected using query partitions and query-response coordinates, without inspecting the dense residual.
- A.3 WEIGHTED AFFINE FITTING AND PROBE-RESIDUAL OUTPUT SUBSPACE: The weighted ridge fit projects predicted corrections into the output subspace spanned by leading right singular vectors of measured exact probe residuals.The basis is built from current-call probe residuals rather than fitted residuals and does not assume the complete residual is globally low rank.
- B RESPONSE-COUPLED PARTITIONING AND SPARSE EXECUTION: Response-Coupled Partitioning forms paired K/V groups by clustering key-response coordinates derived from sampled-query pre-softmax response profiles.Each value inherits its paired key’s assignment, while values do not enter the partition feature.
- B.1 KEY-RESPONSE COORDINATES FOR PAIRED K/V GROUPS: Key-response coordinates compare pre-softmax response profiles under deterministic sampled queries, with eigendirections weighted by normalized eigenvalues.Applying k-means to these coordinates produces the paired K/V groups.
- B.2 QUERY-RESPONSE COORDINATES FROM K/V GROUPS: Query-response coordinates are built from equally weighted centered nonempty K/V-group centroids, then clustered to form query partitions and stratify probe selection.The K/V and query stages form a coupled, single-pass construction without alternating refinement.
C IMPLEMENTATION AND EVALUATION DETAILS · C.1 GPU IMPLEMENTATION · C.2 SHARED HYPERPARAMETERS
The implementation combines GPU-optimized routing and correction kernels with batched linear algebra. Experiments use shared method settings in BF16 on one NVIDIA H100, with model-specific executed-pair targets.
- C.1 GPU IMPLEMENTATION: Warm-started GPU k-means and a fused query-response projection/RMS-normalization kernel implement the routing path.Residual fitting, probe-residual SVD, projection, and correction use small batched GPU linear-algebra operations.
- C.1 GPU IMPLEMENTATION: 0.5064 ms is the fused query-projection latency at the evaluated production shape, versus 1.6170 ms for the reference implementation.The fused projection has relative feature error 3.68×10−4 against the reference.
- C.2 SHARED HYPERPARAMETERS: All experiments use BF16 on a single NVIDIA H100 GPU.The method settings are shared across models except for the selected total executed-pair target.
- C.2 SHARED HYPERPARAMETERS: Table C.1 reports the method configuration used in the evaluated models.The supplied passage identifies the table but does not enumerate its configuration values.
- C.2 SHARED HYPERPARAMETERS: Table C.2 reports probe-row fractions for the evaluated sequence lengths with M = 64.The supplied passage identifies the table but does not enumerate the probe-row fractions.
C.3 TIMING PROTOCOL
The timing protocol matches dense and sparse runs across all listed experimental conditions and counts probe interactions in executed-pair density. Reported latency includes the full SparsePR pipeline, from response-coordinate construction through output restoration.
- Matched conditions: Dense and sparse runs match checkpoints, inputs, seeds, schedules, inference steps, guidance, resolutions, and frame counts.This controls timing and quality comparisons across the two execution modes.
- Density accounting: Executed-pair density counts exact query–key interactions, including probe rows.The density therefore reflects both sparse attention and exact probe computation.
- Density sensitivity: 22% is the reference density for comparing normalized attention-output error across semantic and response-coupled partitioning.Figure D.1 reports mean and p99 error versus total executed-pair density under the same residual-repair configuration.
D CROSS-MODEL DENSITY SENSITIVITY
Response-coupled partitioning consistently outperforms semantic partitioning across executed-pair densities from 12% to 35%, reducing mean error across all densities and models and lowering p99 error over nearly the full range.
- D CROSS-MODEL DENSITY SENSITIVITY: 12%–35% executed-pair densities: response-coupled partitioning reduces mean error across all densities and models versus semantic partitioning.The comparison uses the same residual-repair configuration.
- D CROSS-MODEL DENSITY SENSITIVITY: p99 error also decreases over nearly the full density range, showing partitioning gains persist across operating points rather than depending on one density.
E MORE QUALITATIVE RESULTS
The section presents qualitative comparisons of SparsePR across Cosmos3-Nano-16B, HunyuanVideo-13B, and Wan2.2-I2V-A14B. Two figures examine Cosmos3-Nano-16B, while additional figures cover HunyuanVideo-13B and Wan2.2-I2V-A14B.
- Figures E.1 and E.2 present qualitative comparisons on Cosmos3-Nano-16B.
- Figure E.3 presents a qualitative comparison on HunyuanVideo-13B.
- Figure E.4 presents a qualitative comparison on Wan2.2-I2V-A14B.