HyperAIHyperAI

Command Palette

Search for a command to run...

TriAttention: كفاءة الاستدلال الطويل باستخدام Trigonometric KV Compression

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

الملخص

تؤدي عمليات الاستدلال الممتد (Extended reasoning) في النماذج اللغوية الكبيرة (LLMs) إلى خلق اختناقات شديدة في ذاكرة KV cache. وتعتمد أبرز طرق ضغط KV cache الحالية على تقدير أهمية الـ KV باستخدام درجات الانتباه (attention scores) المستمدة من الـ queries الأخيرة بعد عملية RoPE. ومع ذلك، فإن الـ queries تدور مع تغير الموقع أثناء عملية RoPE، مما يجعل الـ queries الممثلة قليلة للغاية، ويؤدي بدوره إلى اختيار سيئ للـ top-keys وعدم استقرار في عملية الاستدلال.ولتجنب هذه المشكلة، توجهنا إلى مساحة ما قبل RoPE (pre-RoPE space)، حيث لاحظنا أن متجهات Q و K تتركز بشكل مكثف حول مراكز ثابتة غير صفرية وتظل مستقرة عبر المواقع المختلفة، وهو ما نسميه بتركيز Q/K (Q/K concentration). وقد أظهرنا أن هذا التركيز يجعل الـ queries تنتبه بشكل تفضيلي للـ keys عند مسافات محددة (على سبيل المثال، أقرب الـ keys)، حيث تحدد هذه المراكز المسافات المفضلة عبر سلسلة مثلثية (trigonometric series).وبناءً على ذلك، نقترح TriAttention لتقدير أهمية الـ keys من خلال الاستفادة من هذه المراكز. فباستخدام السلسلة المثلثية، نستخدم تفضيل المسافة الذي تميزه هذه المراكز لتسجيل درجات الـ keys وفقًا لمواقعها، كما نستفيد من قيم Q/K norms كإشارة إضافية لتقدير الأهمية.وفي اختبار AIME25 مع توليد 32K-token، يحقق TriAttention دقة استدلال تماثل Full Attention، بينما يحقق إنتاجية (throughput) أعلى بمقدار 2.5 ضعف أو تقليلاً في ذاكرة KV بمقدار 10.7 ضعف، في حين لا تحقق النماذج المرجعية (baselines) الرائدة سوى نصف الدقة تقريبًا عند نفس مستوى الكفاءة. يتيح TriAttention نشر OpenClaw على وحدة GPU واحدة مخصصة للمستهلكين، حيث كان السياق الطويل (long context) سيؤدي لولا ذلك إلى نفاذ الذاكرة (out-of-memory) عند استخدام 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.


بناء الذكاء الاصطناعي بالذكاء الاصطناعي

من الفكرة إلى الإطلاق — سرّع تطوير الذكاء الاصطناعي الخاص بك مع المساعدة البرمجية المجانية بالذكاء الاصطناعي، وبيئة جاهزة للاستخدام، وأفضل أسعار لوحدات معالجة الرسومات.

البرمجة التعاونية باستخدام الذكاء الاصطناعي
وحدات GPU جاهزة للعمل
أفضل الأسعار

HyperAI Newsletters

اشترك في آخر تحديثاتنا
سنرسل لك أحدث التحديثات الأسبوعية إلى بريدك الإلكتروني في الساعة التاسعة من صباح كل يوم اثنين
مدعوم بواسطة MailChimp