HyperAIHyperAI

Command Palette

Search for a command to run...

TriAttention: Efficient Long Reasoning with Trigonometric KV Compression

Weian Mao Xi Lin Wei Huang Yuxin Xie Tianfu Fu Bohan Zhuang Song Han Yukang Chen

Abstract

Extended reasoning in large language models (LLMs) creates severe KV cache memory bottlenecks. Leading KV cache compression methods estimate KV importance using attention scores from recent post-RoPE queries. However, queries rotate with position during RoPE, making representative queries very few, leading to poor top-key selection and unstable reasoning. To avoid this issue, we turn to the pre-RoPE space, where we observe that Q and K vectors are highly concentrated around fixed non-zero centers and remain stable across positions -- Q/K concentration. We show that this concentration causes queries to preferentially attend to keys at specific distances (e.g., nearest keys), with the centers determining which distances are preferred via a trigonometric series. Based on this, we propose TriAttention to estimate key importance by leveraging these centers. Via the trigonometric series, we use the distance preference characterized by these centers to score keys according to their positions, and also leverage Q/K norms as an additional signal for importance estimation. On AIME25 with 32K-token generation, TriAttention matches Full Attention reasoning accuracy while achieving 2.5x higher throughput or 10.7x KV memory reduction, whereas leading baselines achieve only about half the accuracy at the same efficiency. TriAttention enables OpenClaw deployment on a single consumer GPU, where long context would otherwise cause out-of-memory with Full Attention.

One-sentence Summary

To address the instability of post-RoPE importance estimation, researchers propose TriAttention, a KV cache compression method that leverages Q/K concentration in the pre-RoPE space and uses a trigonometric series to model distance-based attention preferences, matching Full Attention accuracy on AIME25 with 32K-token generation while achieving 2.5x higher throughput or 10.7x KV memory reduction.

Key Contributions

  • The paper identifies a phenomenon called Q/K concentration in the pre-RoPE space, where query and key vectors cluster around stable, non-zero centers regardless of position.
  • This work introduces TriAttention, a method that estimates key importance by using a trigonometric series to characterize distance preferences based on these stable centers and incorporating Q/K norms as an additional signal.
  • Experimental results on benchmarks such as AIME25 demonstrate that TriAttention matches Full Attention reasoning accuracy while achieving 2.5x higher throughput or 10.7x KV memory reduction compared to existing baselines.

Introduction

As large language models engage in extended reasoning, the growing KV cache creates significant memory bottlenecks that hinder long-context performance. Existing compression methods typically estimate token importance in the post-RoPE space, but these approaches struggle because positional rotations cause query vectors to shift constantly. This rotation makes it difficult to identify representative queries and leads to unstable importance estimation or the loss of critical directional information. The authors leverage a discovered property called Q/K concentration in the pre-RoPE space, where vectors cluster around stable, fixed centers. By using a trigonometric series to model how these centers determine distance-based attention preferences, they propose TriAttention to estimate key importance more reliably and efficiently.

Dataset

Please provide the paper paragraphs you would like me to summarize. The text provided in your prompt only contains a title and a placeholder sentence regarding a benchmark, which does not contain the necessary technical details (composition, sources, sizes, or processing rules) required to draft the description.

Once you provide the full text, I will generate the concise dataset description following all your requirements.

Method

The authors leverage Rotary Position Embedding (RoPE) to model positional information through rotations in vector space, which is a foundational component of their method. RoPE operates by dividing a ddd-dimensional vector into d/2d/2d/2 two-dimensional subspaces, each associated with a frequency band fff. For each band, a rotation by angle ωfp\omega_f pωfp is applied at position ppp, where ωf=θ2f/d\omega_f = \theta^{-2f/d}ωf=θ2f/d and θ=10000\theta = 10000θ=10000 is a fixed constant. This rotation is expressed as a linear transformation on the vector components (x2f,x2f+1)(x_{2f}, x_{2f+1})(x2f,x2f+1), resulting in post-RoPE vectors. The authors observe that queries and keys in the pre-RoPE space are highly concentrated around non-zero centers, a phenomenon consistent across different positions and contexts, as illustrated in the distribution plots shown in the figure below.

Pre-RoPE distribution of query and key vectors
Pre-RoPE distribution of query and key vectors

This concentration is quantified using the Mean Resultant Length R=E[q]/E[q]R = \|\mathbb{E}[q]\| / \mathbb{E}[\|q\|]R=E[q]∥/E[q], where values approaching 1 indicate strong directional concentration. The authors show that this concentration enables the attention computation to be approximated by a trigonometric series. When query and key vectors are approximately constant, the attention logit simplifies to a sum of cosine and sine terms in the relative position Δ=pqpk\Delta = p_q - p_kΔ=pqpk, forming a trigonometric series with coefficients determined by the magnitudes and phases of the vectors.

Based on this analysis, the authors propose TriAttention, a KV cache compression method that scores key importance for pruning. The scoring function combines two components: a trigonometric series score StrigS_{\text{trig}}Strig and a norm-based score SnormS_{\text{norm}}Snorm. The trigonometric series score estimates attention based on distance preference, using the expected query center E[qf]\mathbb{E}[q_f]E[qf] as a proxy for future queries and computing a cosine similarity weighted by vector magnitudes and phase differences. The norm-based score accounts for variations around the center by using the expected query norm E[qf]\mathbb{E}[\|q_f\|]E[qf] and key magnitude kf\|k_f\|kf. The final combined score is S(k,Δ)=Strig(k,Δ)+Snorm(k)S(k, \Delta) = S_{\text{trig}}(k, \Delta) + S_{\text{norm}}(k)S(k,Δ)=Strig(k,Δ)+Snorm(k).

To adapt the scoring to varying levels of concentration, the authors introduce an adaptive weighting mechanism. The Mean Resultant Length RfR_fRf for each frequency band fff is used to scale the norm-based score. When RfR_fRf is high (strong concentration), the contribution of SnormS_{\text{norm}}Snorm is reduced, emphasizing the trigonometric series. When RfR_fRf is low (weak concentration), the full norm contribution is preserved. The final score Sfinal(k)S_{\text{final}}(k)Sfinal(k) is derived by averaging the score over multiple future query positions and applying a normalize-then-aggregate strategy for Grouped-Query Attention, where scores from different query heads are z-score normalized and combined via a maximum operation.

TriAttention scoring components
TriAttention scoring components

The method is implemented with window-based pruning, where key scoring and pruning are triggered every 128 generated tokens to reduce computational overhead. Keys are retained based on their final score, and the top-BBB keys are kept in the KV cache. The overall framework, as illustrated in the figure below, begins with offline calibration to compute query and key distribution centers, followed by the scoring process during inference, and concludes with the retention of top-scoring keys to produce a pruned attention map.

Method overview of TriAttention
Method overview of TriAttention

The effectiveness of this approach is demonstrated by its ability to maintain correct memory retention in recursive tasks, where losing intermediate states leads to error propagation and a corrupted final result, as shown in the figure below.

Memory retention in recursive simulation
Memory retention in recursive simulation

Experiment

The researchers evaluate TriAttention, a KV cache compression method, by testing its ability to reconstruct attention patterns through trigonometric series and its performance on mathematical reasoning, retrieval, and agentic tasks. Experiments across various architectures and benchmarks demonstrate that TriAttention effectively preserves essential information for long-chain reasoning and memory retention while significantly improving throughput and reducing memory footprints. The results show that the method maintains accuracy comparable to full attention even under aggressive compression, outperforming existing observation-based pruning baselines.

TriAttention achieves significant throughput improvements over Full Attention while maintaining comparable accuracy across multiple benchmarks. The method shows substantial speedup, particularly on MATH 500, and reduces KV cache memory requirements, enabling efficient long-context reasoning. TriAttention achieves up to 6.3× higher throughput than Full Attention on MATH 500. TriAttention matches Full Attention accuracy while reducing KV budget significantly on AIME24 and AIME25. TriAttention enables efficient long-context reasoning, allowing successful task completion within limited GPU memory.

TriAttention efficiency gains
TriAttention efficiency gains

TriAttention achieves the highest accuracy across all tested KV cache budgets compared to other compression methods. It matches or exceeds the performance of FullKV at lower memory usage, demonstrating superior efficiency and accuracy. TriAttention achieves the highest accuracy at all budget levels compared to other methods. TriAttention matches FullKV performance at lower memory usage, showing improved efficiency. TriAttention outperforms all baselines, including H2O, TOVA, and RaaS, across different KV cache budgets.

TriAttention outperforms baselines
TriAttention outperforms baselines

The experiment evaluates the impact of future offset range and spacing strategy on model accuracy. Results show that increasing the maximum distance improves performance, while geometric spacing outperforms linear spacing in maintaining accuracy. Increasing the maximum distance improves accuracy Geometric spacing outperforms linear spacing Accuracy varies with the number of offsets

Future offset ablation study
Future offset ablation study

{"caption": "AIME performance on reasoning and coding", "summary": "The the the table compares performance on AIME24 and AIME25 benchmarks between coding and reasoning tasks. Reasoning tasks show lower performance than coding tasks on both benchmarks, indicating a gap in model capabilities across domains.", "highlights": ["Reasoning tasks achieve lower scores than coding tasks on both AIME24 and AIME25.", "Performance is consistently higher on AIME24 compared to AIME25 for both coding and reasoning tasks.", "The gap between coding and reasoning performance is larger on AIME25 than on AIME24."]

[[IMG:http://api-rsrc.hyper.ai/2604.04921/41063cc6-280f-4700-a63f-3eba5c13a885/tex_resource/extracted_tables/table-4.png|]]

TriAttention achieves the highest average score on the RULER benchmark, surpassing other methods. The results demonstrate superior performance compared to SnapKV and PyramidKV, highlighting its effectiveness in retrieval tasks. TriAttention achieves the highest RULER average score among all methods TriAttention significantly outperforms SnapKV and PyramidKV on RULER TriAttention demonstrates strong retrieval capabilities on the RULER benchmark

TriAttention outperforms baselines on RULER
TriAttention outperforms baselines on RULER

TriAttention is evaluated across various benchmarks to validate its throughput, memory efficiency, and retrieval capabilities compared to full attention and existing compression methods. The results demonstrate that TriAttention significantly improves throughput and reduces KV cache requirements while maintaining or exceeding the accuracy of baseline models. Furthermore, ablation studies indicate that performance is optimized through specific offset ranges and geometric spacing strategies.


Build AI with AI

From idea to launch — accelerate your AI development with free AI co-coding, out-of-the-box environment and best price of GPUs.

AI Co-coding
Ready-to-use GPUs
Best Pricing

HyperAI Newsletters

Subscribe to our latest updates
We will deliver the latest updates of the week to your inbox at nine o'clock every Monday morning
Powered by MailChimp