Partition the Support, Reconstruct the Residual: Training-Free Sparse Attention for Video Generation and World Models

The Problem: Sparse Attention Doesn't Know What It's Missing

Video transformers use attention mechanisms that are computationally expensive—every query token must compare against every key-value pair. A natural fix is sparse attention: only compute a subset of those comparisons. But existing training-free sparse methods have a hidden flaw. Even if you identify which attention weights are large (row-wise concentration), that doesn't tell you how to handle the weights you skip. Two queries routed to the same sparse block may need very different information from the key-value pairs you're ignoring. Worse, the error introduced by dropping those interactions is unpredictable from the sparse output alone.

The authors show that how you partition queries into sparse blocks fundamentally changes both what information you retain and how well you can reconstruct what you lost.

How SparsePR Works

The method has two core ideas working together.

Response-Coupled Partitioning: Instead of partitioning queries arbitrarily, the authors sample a small set of queries and compute their exact attention responses (their full K/V interactions). These sampled responses are clustered into groups. The centroids of these groups define a coordinate system. All queries—both sampled and unsampled—are then routed to the nearest centroid, grouping them by similarity in what information they actually need. This ensures queries in the same sparse block have overlapping support (similar attention patterns), reducing the harm of using a single sparse block for both.

Probe-Fitted Residual Reconstruction: After computing sparse attention on a query, the output is incomplete—it's missing contributions from the skipped key-value pairs (the residual). Rather than assume this residual is negligible, SparsePR learns to predict it. A small calibration set of queries (the "probe") is used to fit a simple linear correction: an affine transformation that maps the sparse output to a better estimate of the full-attention output. This correction is specific to the current attention call, so it adapts to the actual sparse output observed.

Together: partition by actual attention needs, then correct for what you skipped.

Results and Impact

The authors test SparsePR on four models: video generation (Latte, Cosmos), and world models (World Models, Genie). Across all four, SparsePR reduces attention-reconstruction error—the gap between sparse and full attention—compared to prior training-free baselines.

In practice, SparsePR achieves 1.48x to 2.61x end-to-end speedups while keeping only 22–26% of attention pairs. Ablations show that probe-fitted reconstruction accounts for most of the error reduction; response-coupled partitioning adds robustness and helps under tight memory budgets.

Generation quality is preserved: visual outputs remain competitive with full attention.

Why this matters: video transformers are expensive at scale. This method accelerates them without retraining, making it immediately useful for practitioners. The insight—that partition geometry and residual predictability are intertwined—is also conceptually useful for future sparse-attention designs. By showing that you can calibrate away residual error with a tiny probe set, the work opens a practical path to training-free sparse inference that actually works.