AdaCluster: Adaptive Query-Key Clustering for Sparse Attention in Video Generation

Abstract: Video diffusion transformers (DiTs) suffer from prohibitive inference latency due to quadratic attention complexity. Existing sparse attention methods either overlook semantic similarity or fail to adapt to heterogeneous token distributions across layers, leading to model performance degradation. We propose AdaCluster, a training-free adaptive clustering framework that accelerates the generation of DiTs while preserving accuracy. AdaCluster applies an angle-similarity-preserving clustering method to query vectors for higher compression, and designs a euclidean-similarity-preserving clustering method for keys, covering cluster number assignment, threshold-wise adaptive clustering, and efficient critical cluster selection. Experiments on CogVideoX-2B, HunyuanVideo, and Wan-2.1 on one A40 GPU demonstrate up to 1.67-4.31x speedup with negligible quality degradation.

• The paper introduces an adaptive query-key clustering framework that leverages unit sphere normalization and multi-stage K-means to reduce attention computational cost.
• It achieves up to 4.31× acceleration with maintained visual quality (PSNR 30.58dB, lower LPIPS) on high-resolution, long-duration video generation tasks.
• It employs TensorQuest for efficient cluster selection, dynamically adapting cluster counts per layer to handle heterogeneous token distributions.

AdaCluster: Adaptive Query-Key Clustering for Sparse Attention in Video Generation

Introduction and Motivation

Diffusion Transformer (DiT) models underpin recent progress in text-to-video generation, yet their deployment remains constrained by quadratic complexity in attention computation. Sequence lengths in this domain are driven by both spatial resolution and frame count, resulting in attention bottlenecks that can consume over 75% of inference time for input sequences exceeding $70K$ tokens. Prior sparse attention methods either resort to static heuristics, overlook semantic structure, or fail to adapt to heterogeneous distributional shifts across layers, ultimately undermining fidelity or generalization. AdaCluster (2604.18348) introduces a principled, training-free, adaptive clustering framework that addresses these limitations via role-aware query and key clustering, layer-wise compressibility analysis, and efficient cluster selection, all tuned for video diffusion transformers.

Technical Design

AdaCluster decomposes attention acceleration into query clustering—targeting angular similarity—and key clustering—targeting Euclidean similarity—underlining the fundamentally distinct roles of queries and keys in $QK^\top$ scoring.

Query Clustering: Angular Compactness

AdaCluster normalizes query vectors onto the unit sphere, converting Euclidean clustering into angular clustering. This transformation increases intra-cluster compactness and enhances the compression ratio, validated by tighter clusters and lower Davies-Bouldin indices compared to conventional unnormalized clustering.

Figure 1: Query vector distributions before and after normalization, showing improved clustering efficiency on the unit sphere.

Empirically, compressing to $65$ clusters after normalization achieves equivalent intra-cluster distances as $235$ clusters in the unnormalized space, yielding a $3.6\times$ compression advantage. Critically, normalization preserves Top-$K$ relative ordering: for any keys $k_a$, $k_b$, if $q^\top k_a > q^\top k_b$, then their normalized scores maintain this order.

Key Clustering: Layer-Wise Compressibility and Multi-Stage K-means

Key token distributions exhibit pronounced heterogeneity across model layers, motivating AdaCluster's layer-adaptive clustering strategy. Compactness is quantified via intra-cluster squared reconstruction error (MSE), with cluster counts dynamically tuned per layer to maximize efficiency without degrading accuracy.

Figure 2: Layer-wise compactness scores delineate variation across Wan-2.1 and Hunyuan models, guiding cluster count assignment.

AdaCluster employs multi-stage K-means, progressively reclustering outlier tokens based on distance thresholds. Layers that exceed a preset cluster limit are flagged as "hard-to-compress" and default to full attention, with negligible impact on overall speedup.

Figure 3: Progressive clustering of outliers via multi-stage K-means ensures intra-cluster compactness and adaptive cluster counts.

Token distributions change gradually across denoising steps; thus, cluster centers are initialized from the previous timestep, further reducing computational overhead.

Figure 4: Consistency of token distribution across timesteps for different layers enables efficient center initialization.

Efficient Cluster Selection: TensorQuest

AdaCluster integrates TensorQuest, a matrix-multiplication reformulation of the Quest upper-bound scoring method, optimized for GPU tensor core execution. This enables efficient identification of critical Key-Value clusters with negligible computational delay, unlike prior CUDA-only approaches.

Figure 5: TensorQuest delivers significant speedup for Top-K selection under increasing input tokens compared to vanilla selection.

Empirical Results

AdaCluster is benchmarked on CogVideoX-2B, HunyuanVideo, and Wan-2.1 using an NVIDIA A40, generating 81-frame videos at resolutions up to $1920\times1120$. Experimental evaluations span fidelity (PSNR, SSIM, LPIPS), semantic metrics (imaging quality, background/subject consistency), and efficiency (end-to-end speedup).

Figure 6: AdaCluster achieves substantial reduction in inference latency versus Full Attention while preserving high visual quality (PSNR 30.58dB, 1.68$QK^\top$0 speedup).

AdaCluster consistently outperforms SpargeAttn and SVG2 in similarity and perceptual quality across models, with up to $QK^\top$1 acceleration at peak sequence lengths. Video quality remains stable across all prompts and models, and adaptive cluster assignment yields higher PSNR, SSIM, and notably lower LPIPS.

Figure 7: Time cost trend shows superior scalability and acceleration of AdaCluster with respect to sequence length.

Ablation studies confirm the necessity of layer-wise adaptive clustering, the contribution of TensorQuest, and the benefits of query normalization.

Theoretical and Practical Implications

AdaCluster’s role-aware clustering strategies and adaptive cluster allocation offer a blueprint for scalable attention modules in generative diffusion models. The technical innovations—unit sphere normalization, compressibility-driven cluster assignment, progressive reclustering, and tensor core optimization—can be directly extended to long-context image and video transformers, as well as to other attention-intensive tasks in generative AI. By preserving high-fidelity output while significantly reducing computational demand, AdaCluster enables practical deployment of DiTs for high-resolution, long-duration generation, and sets a precedent for future system-level and algorithmic sparsification.

Future Directions

Extensions may explore end-to-end trainable cluster assignment, integration with quantization for further efficiency, and adaptation to multimodal diffusion transformers. System-level fusion with operator optimizations (e.g., FlashAttention kernels) could provide multiplicative performance gains. Additionally, role-aware clustering and adaptive sparsification are conceptually transferable across transformer architectures in NLP, vision, and RL settings.

Conclusion

AdaCluster introduces adaptive, role-aware query-key clustering for sparse attention in video diffusion transformers, yielding robust acceleration while preserving high-fidelity generation. By accommodating layer-wise distributional heterogeneity and exploiting efficient cluster selection on hardware primitives, AdaCluster demonstrates practical and theoretically informed advances for scalable, efficient video generation (2604.18348).