PaperlayerView exact revision on arXiv

Source-linked AI summary

ClusterAttention: A training-free speedup of bidirectional attention

Kasper Nordenram, Amelie Dittmann

arXiv:2608.26965v1cs.LGcs.CV

TL;DR

ClusterAttention addresses the quadratic cost of bidirectional attention and the limitations of sparse methods that depend on input structure or slow clustering. It uses interaction-aware recursive clustering with fixed-size clusters and centroid compensation. Across tabular and video-generation evaluations, it provides substantial speedups while preserving accuracy or improving closeness to dense attention, subject to clustering-latency and cluster-quality limitations.

  • Problem

    Bidirectional attention is quadratic in token count, while existing sparse methods may rely on input structure or slow clustering processes.

  • Method

    ClusterAttention uses interaction-aware transformations, fast recursive clustering, fixed-size clusters, and centroid compensation for excluded clusters.

  • Results

    ClusterAttention retains over 99% relative accuracy on TabPFN-3 while achieving speedups from around 2x to almost 6x as dataset size grows.

  • Takeaways & Limitations

    The method applies training-free sparse attention to unstructured inputs in a single forward pass and is evaluated across vision, tabular, and video-generation domains.

  • Takeaways & Limitations

    At lower token counts, clustering latency, especially eigenvalue decompositions, can dominate and reduce the speed advantage.

Abstract

from arXiv · show

This paper introduces ClusterAttention, a general training-free speedup of bidirectional attention layers. Existing sparse attention methods either rely on structure in the input, such as order in language or spatial proximity in images, or use slow clustering processes amortized over several forward passes. ClusterAttention instead uses a fast recursive clustering method that adapts to the geometry of the keys and queries in each attention head to produce useful clusters. This method allows setting the size of the clusters arbitrarily. We utilize this by setting all clusters to be a fixed size that is a power of two, allowing the block-sparse attention to run at the same latency per query-key interaction as dense attention on GPUs. We also derive an expression for the output error in sparse attention, that explains the counterintuitive experimental finding that tight clusters can lead to larger errors than random clusters. We then derive the error when excluded clusters are compensated through their centroids, and show that this error shrinks with tighter clusters. We integrate this compensation into the method. On large-scale tabular data ClusterAttention speeds up TabPFN-3 arXiv:2605.13986 by two to six times, while retaining at least 99% of the dense accuracy. To our knowledge, it is the first training-free method that can be successfully applied in the setting of unstructured input and a single forward pass. For video generation with Wan 2.1-14B T2V arXiv:2503.20314 , ClusterAttention achieves output closer to dense attention and a larger speedup (1.8x versus 1.4x) compared to SVOO arXiv:2603.18636 , a leading method developed specifically for this domain, both run without offline calibration.

1 Preliminaries

Bidirectional attention is widely used but scales quadratically with token count, motivating sparse approximations that cluster tokens, select cluster interactions, and optionally compensate for excluded clusters. The paper analyzes sparse-attention error through missed attention mass and output deviation, then introduces centroid compensation to make cluster tightness reduce error.

  • 1.1 Motivation for sparse attention: Bidirectional attention has quadratic computation in token count because each token attends to all tokens, although many interactions are weak.
  • 1.1 Motivation for sparse attention: Sparse attention routes each query to a subset of keys instead of computing all key-query interactions.
  • 1.1 Motivation for sparse attention: ClusterAttention clusters keys and queries, scores cluster pairs, selects key-clusters, performs sparse attention, and optionally compensates for unselected clusters.
  • 1.1 Motivation for sparse attention: Speedup requires clustering, scoring, selection, and compensation overhead to remain smaller than the saved attention computation.
  • 1.2 Error in sparse attention: Sparse-attention error depends on missed attention mass and the deviation between included and excluded attention outputs.
  • 1.2 Error in sparse attention: Random clusters can outperform clusters optimized for attention-mass recall because they may reduce deviation between included and excluded outputs when key-value covariance is high.
  • 1.2 Error in sparse attention: Centroid compensation includes unselected clusters through key-value centroid interactions, making tighter clusters reduce the resulting error.
  • 1.2 Error in sparse attention: The Jensen term and covariance term form the compensation error decomposition; empirically, the Jensen term was roughly an order of magnitude larger on DINOv2.

2 Related works

Related sparse-attention methods differ in how they cluster tokens and route interactions. ClusterAttention combines key-query interaction-aware transformations with recursive principal-component splitting into predetermined cluster sizes for fast, hardware-efficient sparse attention.

  • 2 Related works: SpargeAttn uses input structure such as ordering or space-filling curves to create clusters and is training-free.
  • 2 Related works: Clustered Attention represents each query cluster with its centroid and uses fast locality-sensitive hashing followed by K-Means in Hamming space.
  • 2 Related works: AdaCluster separately clusters keys and queries per head with different methods, using Euclidean key clustering and angle-based query clustering.
  • 2 Related works: SVOO performs offline layer-wise sparsity profiling, iteratively refines key and query clusters, and reclusters every N = 20 diffusion steps.
  • 2 Related works: ClusterAttention accounts for different key and query roles, their interactions, and tiling constraints when constructing clusters and routing sparse attention.
  • 2 Related works: ClusterAttention uses interaction-aware transformations and recursive principal-component splitting to produce predetermined cluster sizes with fast clustering and dense-like latency per interaction.

3 Method

ClusterAttention clusters keys and queries separately within each head, assigns key clusters to query clusters, and computes selected block-sparse attention. Its recursive and transformed clustering methods adapt similarity to attention geometry while supporting fixed-size clusters and optional compensation.

  • Overview: ClusterAttention separately clusters keys and queries within each attention head, then performs per-head assignment and attention.The method is organized into clustering, assignment, and attention stages.
  • Recursive splitting: Recursive splitting projects keys or queries onto an approximate first principal component and partitions them into clusters near a predetermined size c.The threshold is chosen so the number of tokens on one side is a multiple of c; at most one final cluster differs in size before padding.
  • Assignment: Assignment ranks key clusters for each query cluster using centroid dot products, optionally corrected for cluster variance, and selects top-k clusters or enough clusters to meet an attention-mass threshold.The number of selected clusters can therefore be fixed or adaptive to an estimated attention-mass recall.
  • Attention: Attention uses a block-sparse kernel, while the evaluated implementation uses cluster sizes of 128 for keys and 64 for queries.The kernel was modified to return the attention denominator.
  • Recursive splitting: Basic recursive splitting repeatedly estimates a principal direction, partitions vectors near a balanced threshold, and recurses until clusters contain c vectors.Power iteration is warm-started from the parent estimate, and splitting is applied separately to keys and queries.
  • Recursive splitting: Diagonalized recursive splitting computes principal components once, then splits along the highest-variance coordinate at each recursive step.This avoids repeated power iterations and is faster end-to-end than the basic version.
  • Attention-aware transforms: Key clustering can use a query-dependent transformed space so distances reflect differences in attention logits rather than irrelevant key dimensions.The transform is computed on-the-fly from the actual queries, and its square root is obtained from an eigenvalue decomposition; small negative eigenvalues are clamped to zero.
  • Attention-aware transforms: Query clustering uses different representations for adaptive attention-mass selection and top-k assignment because softmax makes the two cases asymmetric.The adaptive case groups queries with similar key-cluster attention distributions, whereas the top-k case emphasizes similar rankings.

4 Results

ClusterAttention is evaluated across vision, tabular data, and video generation, using Pareto comparisons of quality or distortion against latency. The experiments also examine clustering variants, compensation, and practical implementation effects.

  • Evaluation scope: ClusterAttention is tested in vision transformers, tabular transformers, and diffusion transformers for video generation.DINOv2-L is evaluated outside its native resolution range, while TabPFN-3 and Wan 2.1-14B T2V are evaluated in their normal operational ranges.
  • 4.1 DINOv2-L: DINOv2-L is evaluated on DIV8K at 21,904, 62,500, and 143,641 tokens across three image resolutions.The evaluated resolutions are 2072 × 2072, 3500 × 3500, and 5306 × 5306 pixels.
  • 4.1 DINOv2-L: 0.99 cosine similarity with dense attention is used as a suggested reference operating point for comparing latency in the representation-distortion Pareto front.The authors emphasize that the appropriate threshold depends on the downstream application and must be measured empirically.
  • Video generation: SVOO has substantial clustering overhead and is not competitive in the single-pass setting evaluated here.The authors also report that SVOO attention appears slower, potentially because of kernel quantization and ragged cluster sizes.
  • Ablations: Top-k with striped mean-compensation is generally Pareto dominant, while compensation provides a larger quality improvement for top-k than for adaptive selection.The adaptive improvement is smaller and concentrated mostly at higher sparsities, with a low-resolution exception where uncompensated results dominate.

5 Discussion

ClusterAttention’s performance depends on token count and latency-component balance: it is competitive at larger scales, while clustering and compensation can dominate in different regimes.

  • 5.1 Overall performance: At around 20,000 tokens, routing overhead makes ClusterAttention∗slower than SpargeAttn and dense attention; at medium and high resolutions, it is Pareto-dominant.The SpargeAttn and ClusterAttention∗ crossover appears between 25,000 and 30,000 tokens.
  • 5.1 Overall performance: ClusterAttention provides better overall video-generation quality retention and a higher speedup than SVOO, despite lower speedups than SVOO reports.The authors attribute the discrepancy partly to faster dense attention in their testing and do not fully diagnose the quality difference.
  • 5.2 Latency in detail: Clustering overhead dominates at low token counts, with eigenvalue decompositions appearing to be the main few-token cost.The overhead is approximately flat up to around 10,000 tokens, then curves upward.
  • 5.2 Latency in detail: With fixed low k, sparse attention overtakes clustering near 20,000 tokens, while striped-mean compensation overtakes attention near 250,000 tokens and then dominates.For highly sparse attention, the authors identify compensation latency as the most impactful improvement target.
  • 5.2 Latency in detail: At 10% token attention, attention overtakes clustering slightly later and then grows quickly, reaching about twice the SMC computation and data transfer.The cited comparison corresponds to 0.1/1⁄64 = 6.4 times the computation and data transfer in the excluded-cluster compensation.
  • 5.2 Latency in detail: Cluster assignment approaches quadratic scaling but remains essentially the smallest latency component because its constant is small.

6 Future work

Future work targets latency, compensation, clustering quality, routing, broader attention variants, and stronger evaluation. The authors identify concrete bottlenecks and several unimplemented or insufficiently tested improvements.

  • 6 Future work: Better compensation of excluded clusters is identified as a principal improvement direction.
  • 6.1 Improvements: At lower token counts, clustering latency is the main bottleneck, especially because of eigenvalue decompositions; a faster non-diagonalized alternative has not been implemented.
  • 6.1 Improvements: At higher token counts, SMC and actual attention are the main latency bottlenecks, with better implementations potentially making SMC several times faster.
  • 6.1 Improvements: SVOO achieves markedly lower error than adaptive ClusterAttention without SMC at the same attention budget on both DINOv2 and TabPFN-3.This indicates headroom for improving ClusterAttention’s cluster quality.
  • 6.1 Improvements: Possible clustering improvements include variable power-of-two cluster sizes for dense and sparse regions and nonlinear representation transformations.
  • 6.1 Improvements: MuSe’s exponential tilting could improve compensation, but its memory-transfer requirements make competitiveness in this setting unclear and the authors have not tested it extensively.
  • 6.1 Improvements: Current centroid-centroid selection is reasonable because error scales with excluded-cluster attention weight, but it is not necessarily optimal.
  • 6 Future work: The study uses only multi-head attention; extensions to GQA and MQA require ablation and analysis to determine suitable approaches.

A Error of mean-compensated sparse attention

The paper derives query- and key-aware representations for clustering attention interactions, analyzes sparse-attention error, and uses centroid compensation to represent excluded clusters. It also describes recursive splitting and practical complexity considerations for constructing these representations and clusters.

  • Error of mean-compensated sparse attention: Centroid key-and-value compensation incorporates excluded attention clusters using cluster-size weighting.
  • Error of mean-compensated sparse attention: The sparse-attention output error is decomposed into normalization error and the difference between dense and compensated outputs.
  • Error of mean-compensated sparse attention: The quantity δc = ¯wc − wc is nonnegative by Jensen’s inequality and exponential convexity, while the perturbations are zero-mean.
  • Query representation: Query clustering seeks a forward property, where nearby representations yield similar attention weights, and a backward property, where similar weights imply nearby representations.
  • Query representation: Mean-centering keys removes softmax shift invariance, so small representation distances produce small attention-weight differences.
  • Query representation: The key-aware query representation uses relative key-logit rankings, with covariance-based geometry and normalization recovering scale invariance.
  • Key representation: Full query-aware projections are avoided by using landmark queries, optional principal-component projection, and increased softmax temperature when needed.
  • Complexity: Recursive splitting is asymptotically dominated by an n log2 n term, while tested token counts keep selection overhead from dominating ClusterAttention latency.

E Evaluation images

The evaluation uses sampled, square-cropped high-resolution DIV8K images for DINOv2-L.

  • The evaluation image set consists of 12 sampled high-resolution DIV8K images after square-cropping.

F.1 Transforms

These figures compare latency and representation distortion for top-k and adaptive methods, with and without transform-related components, against a 0.99 cosine-similarity reference.

  • Top-k methods without SMC: Figure 5 compares top-k methods without SMC using latency-versus-representation-distortion Pareto fronts.
  • Adaptive methods without SMC: Figure 6 compares adaptive methods without SMC on the same latency-versus-representation-distortion axes.
  • Top-k methods with SMC: Figure 7 ablates transforms for top-k methods with SMC using latency-versus-representation-distortion Pareto fronts.
  • Adaptive methods with SMC: Figure 9-style adaptive comparisons use the same Pareto-front framework, while the dashed line marks 0.99 cosine similarity with dense attention.

F.2 Striped mean-compensation

These ablation figures evaluate compensation through latency-versus-representation-distortion Pareto fronts, using 0.99 cosine similarity with dense attention as the reference operating point.

  • Top-k methods: Figure 9 ablates compensation for top-k methods on DINOv2 using latency-versus-representation-distortion Pareto fronts.
  • Adaptive methods: Figure 10 performs the corresponding compensation ablation for adaptive methods.

F.3 Top-k versus adaptive

Figure 11 compares selection methods by plotting latency against representation distortion on DINOv2. The dashed reference marks 0.99 cosine similarity with dense attention.

  • Figure 11 presents Pareto fronts comparing selection methods on latency versus representation distortion.The evaluation uses DINOv2 representations.
  • Representation distortion is measured relative to dense attention using cosine similarity.
  • The dashed line denotes 0.99 cosine similarity with dense attention.

G Adaptive methods on DINOv2

Figure 12 compares adaptive ClusterAttention with SVOO on DINOv2 using Pareto fronts of latency and representation distortion. Distortion is measured by average cosine similarity with dense attention.

  • Figure 12 compares adaptive ClusterAttention and SVOO on DINOv2.
  • The figure plots latency against representation distortion for the two methods.
  • Representation distortion is measured as average cosine similarity with dense attention.
Loading 2608.26965v1…