HyperAIHyperAI

Command Palette

Search for a command to run...

TriAttention: 삼각함수 기반 KV Compression을 통한 효율적인 Long Reasoning

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

초록

대규모 언어 모델(LLMs)에서의 확장된 추론(Extended reasoning)은 심각한 KV cache 메모리 병목 현상을 야기합니다. 주요 KV cache 압축 방법들은 최근 RoPE 적용 이후(post-RoPE)의 query attention score를 사용하여 KV의 중요도를 추정합니다. 그러나 RoPE 과정에서 query는 위치에 따라 회전하므로, 대표성을 갖는 query가 매우 적어지게 되며, 이는 부적절한 top-key 선택과 불안정한 추론으로 이어집니다.이러한 문제를 해결하기 위해 본 논문에서는 RoPE 적용 전(pre-RoPE) 공간을 주목합니다. 우리는 해당 공간에서 Q 및 K 벡터가 고정된 비영(non-zero) 중심점 주변에 고도로 집중되어 있으며, 위치 변화에 관계없이 안정적으로 유지된다는 점을 관찰하였습니다—이를 Q/K concentration이라 명명합니다. 우리는 이러한 concentration 현상으로 인해 query가 특정 거리(예: 가장 가까운 key)에 있는 key를 우선적으로 attend 하게 되며, 중심점이 삼각 급수(trigonometric series)를 통해 어떤 거리가 선호되는지를 결정한다는 것을 보여줍니다.이를 바탕으로, 우리는 이러한 중심점들을 활용하여 key의 중요도를 추정하는 TriAttention을 제안합니다. TriAttention은 삼각 급수를 통해 이러한 중심점에 의해 특징지어지는 거리 선호도를 활용하여 위치에 따른 key의 점수를 산출하며, 추가적인 중요도 추정 신호로 Q/K norm을 활용합니다.32K-token 생성 조건의 AIME25 benchmark 테스트 결과, TriAttention은 Full Attention과 동일한 추론 정확도를 유지하면서도 2.5배 높은 throughput 또는 10.7배의 KV memory 절감 효과를 달성했습니다. 반면, 기존의 주요 baseline 모델들은 동일한 효율성 수준에서 약 절반 정도의 정확도만을 기록했습니다. TriAttention을 통해 Full Attention 사용 시 Out-of-memory(OOM)가 발생할 수 있는 긴 context 상황에서도 단일 소비자용 GPU에서 OpenClaw 배포가 가능해졌습니다.

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.


AI로 AI 구축

아이디어에서 출시까지 — 무료 AI 코코딩, 즉시 사용 가능한 환경, 최적의 GPU 가격으로 AI 개발을 가속화하세요.

AI 협업 코딩
바로 사용 가능한 GPU
최적의 가격

HyperAI Newsletters

최신 정보 구독하기
한국 시간 매주 월요일 오전 9시 에 이번 주의 최신 업데이트를 메일로 발송합니다
이메일 서비스 제공: MailChimp