PaperlayerView exact revision on arXiv

Source-linked AI summary

PARK: Accurate Block Retrieval for Sparse Attention in Video Diffusion Transformers

Yun Dai, Jiarui Wen, Huiping Zhuang, Cen Chen, Ziqian Zeng

arXiv:2609.38978v1cs.CV

TL;DR

In video diffusion transformers, quadratic full attention makes generation costly, while inaccurate sparse block retrieval can hurt quality or waste computation. PARK preserves each query’s attention preferences and clusters keys using current-query information. Experiments report preserved generation quality with end-to-end speedups over dense attention.

  • Problem

    Centroid-based block retrieval can misestimate block importance because query averaging loses individual preferences and Euclidean key clustering ignores how current queries score keys.

    Source: §1 · p. 2, §4.2 · p. 5

  • Method

    PARK is a training-free method that averages separately normalized attention distributions from original queries and clusters transformed keys using information from current queries.

    Source: §1 · p. 3

  • Results

    Across four models, PARK preserves high generation quality while achieving 1.53×–1.76× end-to-end speedups over dense attention.

    Source: §7 · p. 9

  • Takeaways & Limitations

    PARK achieves the best quality-efficiency trade-off among the compared training-free sparse attention methods on HunyuanVideo and Wan.

    Source: §1 · p. 3

    The evaluation is limited to video generation with diffusion models and does not establish applicability to text-based large language models or other multimodal tasks.

    Source: K LIMITATIONS · p. 28

Abstract

from arXiv · show

Diffusion Transformers (DiTs) have become a dominant architecture for video generation, but their efficiency is limited by the quadratic complexity of full attention. Sparse attention reduces this cost by retrieving important blocks and computing attention only within them, but inaccurate retrieval can either degrade generation quality or yield unnecessary computation. We identify two retrieval mismatches in methods that retrieve blocks using the averaged representations of query and key blocks: (i) query-side aggregation mismatch, where averaging queries before Softmax fails to preserve their individual attention preferences, and (ii) key-side clustering metric mismatch, where standard Euclidean clustering in the original key space can group keys with dissimilar QK scores under the current query, so their average representation may not accurately represent how the current query scores individual keys. These mismatches can lead to inaccurate block retrieval. To address these mismatches, we propose PARK, a training-free sparse attention method for accurate block retrieval. PARK retains every original query, independently normalizes its attention over key blocks, and then averages these distributions within each query block. It also uses information from the current queries to transform keys before clustering, so that keys receiving similar QK scores are grouped together. A fused GPU kernel further reduces the overhead of block retrieval. Experiments on HunyuanVideo and Wan demonstrate that PARK improves block retrieval accuracy and preserves generation quality while accelerating inference, achieving the best quality-efficiency trade-off among the compared sparse attention methods.

1 INTRODUCTION

Full attention’s quadratic cost makes long, high-resolution video generation computationally demanding. PARK targets inaccurate block retrieval with two complementary changes: preserving query-specific attention and clustering keys according to their scores under current queries.

1 INTRODUCTION

  • Inaccurate retrieval can miss important blocks and degrade quality, or select unimportant blocks and waste computation, limiting attainable speedup.Retaining more blocks to compensate for retrieval errors also limits speedup.Source: §1 · p. 2
  • Centroid-based retrieval has two mismatches: query averaging can erase individual key preferences, while Euclidean key clustering may group keys scored differently by current queries.Both mismatches can make block-importance estimates inaccurate.Source: §1 · pp. 2–3
  • PARK is training-free and addresses the mismatches with Per-query Attention-Mass Aggregation (PAMA) and Query-Response Key Clustering (QRKC).PAMA preserves original queries; QRKC uses current-query information to cluster keys by QK-score similarity.Source: §1 · p. 3
  • PARK improves block retrieval accuracy, preserves generation quality, and achieves the best quality-efficiency trade-off among compared training-free sparse attention methods on HunyuanVideo and Wan.The method is presented as delivering high end-to-end speedup.Source: §1 · p. 3

2 RELATED WORK

Related work spans sparse-attention methods that require training and training-free approaches that exploit sparsity in pretrained video DiTs. Within the latter, token reordering groups similar tokens into contiguous blocks.

2 RELATED WORK

  • Training-based methods fine-tune video diffusion models using routing, changing block partitions, or learned predictors and assignments, generally requiring extra training or auxiliary modules.Examples include VSA, VMoBA, DSV, BSA, and MoGA.Source: §2 · p. 3
  • Training-free methods use offline masks or runtime block-importance estimates, while token-reordering methods cluster similar tokens into contiguous blocks.The described approaches include pattern-aware reordering, K-means clustering, and query–key co-clustering.Source: §2 · p. 3

3 PRELIMINARY

Full attention normalizes scores over keys independently for each query, whereas sparse attention skips selected query–key blocks. Block importance estimates guide retrieval, and attention recall measures how much dense-attention weight the selected mask covers.

3.1 ATTENTION AND SPARSE ATTENTION

  • Full attention computes query–key attention with Softmax applied independently along the key dimension for each query, at quadratic cost in sequence length.Sparse attention uses a block-level mask to skip selected interactions.Source: §3.1 · p. 4

3.2 BLOCK IMPORTANCE AND ATTENTION RECALL

  • Block importance is the average attention weight queries in a query block assign to a key block, normalized across key blocks for each query block.Blocks are retrieved in descending estimated importance until the cumulative estimate reaches the Top-p threshold.Source: §3.2 · p. 4
  • Attention recall measures the proportion of full-attention weights covered by the sparse block mask; higher recall indicates more accurate retrieval.The mask retains or skips each query–key block as a whole.Source: §3.2 · p. 4, §3.1 · p. 4

4 MOTIVATION

Centroid-based retrieval compresses query and key blocks, but these compressions introduce distinct approximations relative to exact block importance. The analysis identifies query aggregation and key clustering as separate sources of retrieval error.

4.1 QUERY-SIDE AGGREGATION MISMATCH

  • Query-centroid estimation applies Softmax after replacing a query block with its centroid, unlike exact block importance, which normalizes each query’s attention before averaging within the block.The empirical comparison finds that retaining original queries with key centroids achieves the highest recall among compressed representations across ordering schemes and models.Source: §4.1 · pp. 4–5
  • Query-centroid error depends on the interaction between attention curvature and within-block query variance, so reducing query variance alone does not ensure small error.The analysis identifies the Hessian of the block-attention function and query covariance as the leading-error factors.Source: §4.1 · p. 5

4.2 KEY-SIDE CLUSTERING METRIC MISMATCH

  • Key-centroid estimation is more accurate when keys in each block receive similar dot-product scores from current queries, but standard Euclidean clustering does not account for those queries.The key-side error is proportional to within-block key variance after projection onto the current query.Source: §4.2 · p. 5, Figure 2 · p. 6

5 METHODOLOGY

PARK combines PAMA, which preserves query-specific attention preferences, with QRKC, which clusters keys using a metric derived from current queries to improve block retrieval.

5.1 PAMA: PER-QUERY ATTENTION-MASS AGGREGATION

  • PAMA retains original queries, normalizes each query’s attention over key blocks separately, then averages weights within each query block to estimate block importance.A fused GPU kernel combines centroid scoring, row-wise Softmax, and averaging, keeping intermediate logits and probabilities on-chip.Source: §5 · p. 6, §5.1 · p. 6

5.2 QRKC: QUERY-RESPONSE KEY CLUSTERING

  • QRKC transforms keys using a query-dependent metric so Euclidean K-means groups keys with similar QK scores under current queries.When GQ is positive definite, the transformation uses its Cholesky factor; clustering results and sparse masks are reused across diffusion steps and recomputed every R steps.Source: §5.2 · p. 7
  • Unlike SVG2 and SVOO, PARK keeps original queries for importance estimation and applies query-informed clustering only to keys.SVG2 clusters queries and keys independently, while SVOO couples their partitions through bidirectional co-clustering; both estimate block importance from centroids.Source: §5.2 · p. 7

6 EXPERIMENTS

Across four video DiT models, PARK achieves the strongest reported quality-efficiency trade-off among the compared sparse-attention methods, with quality close to dense attention and faster inference.

6.1 EXPERIMENTAL SETUP

  • Experiments cover four open-source video DiTs at 720p, using VBench prompts and quality, dense-attention similarity, and end-to-end latency metrics.All models use 50 denoising steps; experiments run on one NVIDIA A800 GPU.Source: §6.1 · pp. 7–8

6.2 QUALITY AND EFFICIENCY EVALUATION

  • PARK’s generated videos more closely match dense attention than SVG2, with better-preserved fine-grained visual details.Image Quality improves by 2.91 and 2.70 percentage points on HunyuanVideo and Wan2.1-1.3B, respectively.Source: §6.2 · p. 9, Figure 3 · p. 9
  • PARK achieves the best quality-efficiency trade-off among compared sparse-attention methods, with 1.55×, 1.76×, 1.53×, and 1.66× speedups over dense attention on Wan2.2-14B, Wan2.1-1.3B, Wan2.1-14B, and HunyuanVideo, respectively.It also achieves the highest PSNR and SSIM and lowest LPIPS among compared methods across all four models.Source: §6.2 · pp. 8–9, Table 2 · p. 8

6.3 ABLATION STUDY

  • Removing PAMA degrades PSNR, SSIM, and LPIPS, while removing QRKC also reduces similarity to dense attention; variant latency differences remain within 0.3%.The ablation uses a fixed subset of VBench prompts and otherwise identical model configurations and inference settings.Source: §6.3 · p. 9, Table 3 · p. 9

7 CONCLUSION

PARK is a training-free sparse attention method for fast video diffusion. Across four models, it preserves high generation quality while accelerating inference.

  • Across four models, PARK preserves high generation quality while achieving 1.53×–1.76× end-to-end speedups over dense attention.Source: §7 · p. 9
  • PARK addresses query-side aggregation mismatch with PAMA, which preserves query-specific attention preferences during block retrieval.Source: §7 · p. 9
  • PARK addresses key-side clustering metric mismatch with QRKC, which clusters keys using a query-dependent metric.Source: §7 · p. 9

A DIRECT EVALUATION OF RETRIEVAL MISMATCHES

The evaluation isolates query-side aggregation and key-side clustering mismatches by comparing retrieval estimates with exact references. Retaining original queries and clustering keys by QK-score similarity reduce their respective approximation errors.

A DIRECT EVALUATION OF RETRIEVAL MISMATCHES

  • The query comparison measures each estimate against exact block importance, including key-centroid approximation in both estimates.The query-centroid, full-key configuration separately measures the query-side approximation error; Table 4 reports mean TV errors of 0.1335 and 0.1667.Source: A DIRECT EVALUATION OF RETRIEVAL… · p. 13, Table 4 · p. 14
  • 47.7% and 49.0%: retaining original queries reduces mean TV from 0.1774 to 0.0927 on Wan2.1-14B and from 0.2092 to 0.1066 on HunyuanVideo.The comparison uses the same key centroids; query-centroid estimation averages queries before normalization.Source: A DIRECT EVALUATION OF RETRIEVAL… · p. 13, Figure 4 · p. 13
  • For key clustering, the evaluation compares Euclidean clustering with QRKC using original queries and each partition’s own exact reference.Centered log-mass error subtracts the mean error across key blocks for each query because a common additive offset does not affect Softmax.Source: A DIRECT EVALUATION OF RETRIEVAL… · pp. 13–14
  • 17.2% and 16.9%: QRKC reduces mean centered log-mass RMSE from 0.4960 to 0.4109 on Wan2.1-14B and from 0.4949 to 0.4113 on HunyuanVideo.The reductions are relative to standard Euclidean clustering.Source: Figure 4 · p. 13, A DIRECT EVALUATION OF RETRIEVAL… · p. 14

B KEY-SPACE AND QK-SCORE APPROXIMATION ERRORS

The appendix compares key-space and QK-score approximation errors, tests QRKC’s effect on attention recall, and derives the query-centroid approximation error. It relates that query-side error to within-block query covariance and block-weight curvature.

B KEY-SPACE AND QK-SCORE APPROXIMATION ERRORS

  • QRKC increases raw-key MSE while reducing QK-score error on both models, showing that smaller Euclidean key-centroid distances do not necessarily yield more accurate QK scores.Raw-key MSE averages squared key-space distances; QK-score error averages squared scaled-score errors over query–key pairs.Source: B KEY-SPACE AND QK-SCORE APPROXIMATION… · pp. 14–15, Table 5 · p. 14

C EFFECT OF QUERY-RESPONSE KEY CLUSTERING

D DETAILED DERIVATIONS FOR CENTROID-BASED BLOCK RETRIEVAL

  • The derivations retain all original keys for query-centroid analysis and all original queries for key-centroid analysis, holding partitions fixed within comparisons.The appendix gives complete expansions underlying the analyses in Sections 4.1 and 4.2.Source: D DETAILED DERIVATIONS FOR… · p. 16

D.1 DERIVATION OF THE QUERY-SIDE AGGREGATION MISMATCH

  • The query-centroid error’s first-order effect cancels; its leading term couples within-block query covariance with the curvature of exact normalized block weight.A query centroid can either underestimate or overestimate a key block’s importance; identical queries within a block make the estimate exact.Source: §D.1 · pp. 17–18
  • The Hessian of a block’s exact attention weight depends on attention-weighted statistics over all original keys, including keys outside that block.Those statistics account for competition with all other keys under Softmax normalization.Source: §D.1 · p. 18

D.2 DERIVATION OF THE QUERY-RESPONSE KEY METRIC

The derivation links key-centroid approximation error to the variance of scaled QK scores within blocks, then uses the current queries to define a clustering metric that reflects that error. A factorized key transformation lets Euclidean K-means optimize this query-dependent objective without increasing the feature dimension.

D.2 DERIVATION OF THE QUERY-RESPONSE KEY METRIC

  • For bounded, sufficiently small score deviations, the leading error term is half the within-block variance of scaled QK scores, or 1/(2d) times unscaled-score variance.The leading term is not the full approximation error; higher-order terms remain.Source: §D.2 · p. 20
  • Jensen’s inequality makes the centroid estimate a lower bound on exact block log-mass, with higher-order error terms arising from exponential and logarithmic remainders.The leading error term is proportional to key variance after projection onto the current query.Source: §D.2 · p. 20
  • The query-induced metric weights key-centroid distortion by the current queries’ uncentered second moment, whereas ordinary Euclidean clustering uses the identity metric.The resulting objective measures squared error in scaled attention logits, summed over keys and averaged over queries.Source: §D.2 · p. 21
  • A factorized transformation realizes the query-dependent distances in d dimensions, so Euclidean K-means optimizes the same QRKC objective without changing its clustering complexity.The total measured transformation overhead remains below 2 ms across evaluated settings.Source: §D.2 · pp. 21–22
  • Stabilization adds an isotropic penalty, so the implemented objective differs from the unregularized one and can change assignments in directions weakly weighted by the original metric.With α = 0.05 and the described numerical stabilization, no Cholesky decomposition failures were observed in the experiments.Source: §D.2 · p. 23

E FUSED PAMA KERNEL

PARK’s fused PAMA kernels reduce block-importance estimation overhead, while QRKC clusters transformed keys and applies the resulting labels to original keys and values. Across the examined schedules and sparsity settings, PARK maintains generation quality as recomputation or latency changes.

E FUSED PAMA KERNEL

  • Triton Full-K estimates PAMA block importance across all layers in under 1 s on both models, and four recomputations account for less than 0.21% of end-to-end latency.The two-pass kernel recomputes logits to keep register use independent of the number of key centroids for fixed tile sizes.Source: E FUSED PAMA KERNEL · p. 24

F QRKC PSEUDOCODE

  • QRKC uses transformed keys only for cluster assignments, then applies labels to original keys and values and computes PAMA centroids in the original key space.The transformation retains the feature dimension and K-means computational complexity for fixed token, cluster, and iteration counts.Source: F QRKC PSEUDOCODE · p. 25

G SENSITIVITY TO QUERY AND KEY CENTROID COUNTS

H CLUSTERING AND SPARSE-MASK REUSE

  • Three and four recomputations yield nearly identical similarity to dense attention, and all four tested schedules have similar end-to-end latency.Increasing recomputations generally improves similarity; the study uses HunyuanVideo with 50 denoising steps.Source: H CLUSTERING AND SPARSE-MASK REUSE · p. 26, Table 9 · p. 26

I EFFICIENCY-QUALITY TRADE-OFF STUDY

  • As latency is reduced, similarity to dense attention decreases, but PARK maintains high generation quality across tested sparsity settings, with relatively stable VBench-Q scores.Figure 7 presents the efficiency-quality trade-off under different sparsity settings.Source: Figure 7 · p. 27, I EFFICIENCY-QUALITY TRADE-OFF STUDY · p. 27

J QUERY AND KEY COMPRESSION UNDER TOP-K RETRIEVAL

  • At a fixed retrieval ratio of ρ = 0.1, Table 10 compares top-k attention recall across query/key representations and token orderings, using full-query and full-key estimation as references.The table reports mean, fifth-percentile, and minimum recall across attention heads and layers; higher is better.Source: J QUERY AND KEY COMPRESSION UNDER… · p. 28, Table 10 · p. 28

K LIMITATIONS

The evaluation is limited to accelerating video generation with diffusion models, and does not establish PARK’s applicability to language models or other multimodal tasks. Visual comparisons report quality similar to dense attention across four video-generation models.

K LIMITATIONS

  • PARK is evaluated for video generation with diffusion models; applicability to text-based large language models and other multimodal tasks is not established and may require adaptation.Further research is needed to determine whether PARK can provide similar efficiency benefits while preserving output quality in those settings.Source: K LIMITATIONS · p. 28

L VISUALIZATION OF THE GENERATED VIDEOS

Loading 2609.38978v1…