PaperlayerView exact revision on arXiv

Source-linked AI summary

TriAttention: Efficient Long Reasoning with Trigonometric KV Compression

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

arXiv:2604.04921v1cs.CLcs.CV

TL;DR

Long reasoning makes KV caches a severe memory bottleneck, while post-RoPE importance estimation is unstable because queries rotate with position. TriAttention instead uses stable pre-RoPE Q/K centers, trigonometric distance preferences, and norms to score keys. On AIME25, it matches Full Attention accuracy while delivering 2.5× higher throughput or 10.7× lower KV memory.

  • Problem

    Long reasoning creates severe KV cache memory bottlenecks, while post-RoPE methods have limited representative queries because queries rotate with position.

  • Method

    TriAttention estimates key importance from stable pre-RoPE Q/K centers through trigonometric distance preferences, with Q/K norms as complementary signals.

  • Results

    2.5× higher throughput or 10.7× KV memory reduction is achieved on AIME25 while matching Full Attention accuracy.

  • Takeaways & Limitations

    TriAttention consistently outperforms baselines across mathematical reasoning and general tasks while matching Full Attention with lower resource use.

  • Takeaways & Limitations

    Post-RoPE attention methods are limited by tiny observation windows, while norm-based methods ignore direction and cannot readily exploit it after rotation.

Abstract

from arXiv · show

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.

1. Introduction

TriAttention addresses unstable post-RoPE KV importance estimation by using stable pre-RoPE Q/K centers, whose trigonometric distance preferences guide key scoring alongside Q/K norms. It matches Full Attention accuracy on AIME25 while substantially improving throughput or reducing KV memory.

  • Long reasoning produces tens-of-thousands-token sequences, making KV cache growth a severe memory bottleneck.
  • Post-RoPE methods use a tiny recent-query window because positional rotation makes older queries unrepresentative, causing important keys to go undetected.
  • Pre-RoPE Q/K vectors concentrate around fixed non-zero centers and remain stable across positions and contexts.
  • Q/K centers yield a trigonometric series whose attention-vs-distance curve predicts preferred key distances.
  • TriAttention scores and retains top keys using Q-center distance preferences, complementary Q/K norms, and concentration-based adaptive weighting.
  • 2.5× higher throughput or 10.7× KV memory reduction is achieved on AIME25 at Full Attention's equivalent 40.8% accuracy.

2. Background and Related Work

KV compression methods estimate token importance using post-RoPE representations, but positional rotations limit attention-based observations and obscure directional information for norm-based methods. RoPE itself applies frequency-dependent rotations, motivating analysis of stable pre-RoPE structure.

  • KV cache compression retains only a subset of KV pairs under a fixed memory budget, making future-key importance estimation central.
  • 2.1. Rotary Position Embedding: RoPE encodes position by rotating each two-dimensional frequency subspace at a frequency-dependent rate.
  • 2.1. Rotary Position Embedding: Pre-RoPE denotes Q/K vectors before positional rotation, whereas post-RoPE denotes vectors after rotation.
  • 2.2. KV Cache Compression: Existing compression methods include heuristic, attention-based, and norm-based approaches that estimate importance from post-RoPE representations.
  • 2.3. Limitations of Post-RoPE Methods: Post-RoPE attention methods are constrained by rotating queries, limiting useful observations and potentially missing keys that become important later.
  • 2.3. Limitations of Post-RoPE Methods: Post-RoPE norm-based methods use magnitudes but ignore directional information because positional rotations entangle direction with position.

3. Pre-RoPE and Trigonometric Series

Pre-RoPE Q/K vectors concentrate around stable non-zero centers, allowing attention patterns to be predicted as trigonometric functions of Q-K distance. Across heads and architectures, these center-based predictions correlate substantially with actual attention, supporting predictable distance preferences.

  • 3.1. The Pre-RoPE Concentration Phenomenon: Q and K vectors concentrate around non-zero centers across most attention heads and remain stable across positions and input contexts.The concentration is prevalent across models rather than limited to specific head types.
  • 3.1. The Pre-RoPE Concentration Phenomenon: Mean Resultant Length values approaching 1.0 across most Qwen3-8B heads confirm strong Q/K concentration.R = ∥E[q]∥/E[∥q∥], with R →1 indicating perfect concentration and R →0 indicating uniform dispersion.
  • 3.2. Predictable Distance Preferences: Approximating Q/K vectors by their centers transforms attention computation into a trigonometric series that depends only on Q-K distance.The series coefficients are determined by the Q/K centers, which shape the resulting attention-vs-distance curve.
  • 3.2. Predictable Distance Preferences: Different Q/K centers produce different distance-preference curves, including peaks at small distances or large distances.These curves encode local-attention preferences or attention-sink behavior through the centers’ trigonometric coefficients.
  • 3.3. Experimental Validation: Across Qwen3, Qwen2.5, and Llama3, predicted attention patterns achieve mean reconstruction correlations above 0.5, with distributions peaking around 0.6–0.9.A non-cherry-picked Qwen3 example reaches r̄ = 0.72; Qwen3 MRL values across Math, Coding, and Chat are nearly identical at 0.977–0.980.

4. TriAttention

TriAttention scores cached keys using distance preferences from a trigonometric series and norm information, balancing the signals according to Q/K concentration before pruning to a fixed budget.

  • Key importance scoring: TriAttention combines a trigonometric-series score for distance preferences with a complementary norm-based score.The method scores keys independently and retains only the top-scoring keys.
  • Key importance scoring: The trigonometric-series score uses Q centers and future Q-K distances to estimate how much attention each cached key will receive.Importance is averaged across multiple future offsets D = {1, 2, 4, . . . , 2^16}.
  • Key importance scoring: For norm scoring, expected query norms from calibration data weight each frequency band alongside the cached key representations.This accounts for variation around the Q/K centers that the trigonometric-series approximation does not capture.
  • Adaptive weighting via concentration: Q/K concentration controls the balance: strong concentration favors the trigonometric score, while weaker concentration increases the norm contribution.Concentration is quantified using the Mean Resultant Length R_f.
  • KV cache pruning: Pruning retains the top-B independently scored keys, while scoring is batched every 128 generated tokens to reduce computational overhead.When the cache exceeds budget B, the method scores all keys and evicts the rest.
  • KV cache pruning: In GQA, TriAttention z-score normalizes scores within each query head, aggregates them by maximum, and retains keys important to any query head.The final top-B keys are selected after aggregation.

5. Experiments

Experiments across reasoning benchmarks and memory-retention tests show that TriAttention consistently approaches or matches Full Attention while outperforming KV-compression baselines across accuracy, efficiency, and memory budgets.

  • Reasoning Task Performance: TriAttention consistently achieves the best performance among KV cache compression methods across evaluated models and mathematical reasoning datasets, closely approaching Full Attention.The evaluation covers Qwen3-8B, DeepSeek-R1-Distill-Llama-8B, DeepSeek-R1-Distill-Qwen-7B, and GPT-OSS-20B on AIME24, AIME25, and MATH 500.
  • Reasoning Task Performance: 43.3% on AIME25 at a 4,096-token budget surpasses Full Attention, while TriAttention matches Full Attention on MATH 500 at a 1,024-token budget.At a 2,048-token AIME25 budget, TriAttention scores 32.9% versus R-KV’s 17.5%.
  • Memory Retention Benchmark: Under memory pressure through recursion depth 16, TriAttention performs comparably to Full Attention, whereas R-KV’s accuracy drops from approximately 61% at depth 14 to 31% at depth 16.TriAttention begins to lag behind Full Attention only beyond depth 18.
  • Ablation Study: Removing the trigonometric-series score or concentration-based weighting degrades performance, while calibration remains stable across data quantities and qualities.The norm-only ablation reduces AIME24 accuracy from 45.8% to 40.4%, and calibration performance stays between 45.4% and 45.8% across 50k–960k tokens.
  • Throughput and Efficiency: At identical AIME25 accuracy of 40.8%, TriAttention provides 2.5× higher throughput or 10.7× KV memory reduction than Full Attention.On MATH 500, it reaches 1,405 versus 223 tokens/s, a 6.3× throughput gain, with comparable accuracy of 68.4% versus 69.6%.
  • Throughput and Efficiency: Compared with R-KV at comparable accuracy, TriAttention uses half the KV budget and achieves 85% higher throughput; at equal memory, it achieves higher accuracy.The comparable-accuracy budgets are 1,024 versus 2,048 tokens, with throughput of 1,405 versus 760 tokens/s.

6. Conclusion

The paper identifies stable pre-RoPE Q/K concentration and uses its trigonometric distance preferences to improve KV importance estimation, achieving strong accuracy-efficiency trade-offs.

  • Conclusion: Q/K concentration makes pre-RoPE query and key centers stable across domains and architectures, enabling predictable distance preferences through a trigonometric series.TriAttention uses these stable centers to avoid the instability of post-RoPE importance estimation.
  • Conclusion: TriAttention consistently outperforms baselines on mathematical reasoning and general tasks while matching Full Attention with 2.5× throughput and 10.7× KV memory reduction.The reported evaluations include AIME, MATH 500, LongBench, and RULER.

Impact Statement

TriAttention targets inference-time reasoning efficiency by using stable pre-RoPE Q/K statistics to estimate key importance through trigonometric distance preferences. The paper reports efficiency gains without compromising accuracy and identifies broader deployment and engineering implications.

  • Impact: The paper introduces no new model capabilities, data sources, or application domains beyond standard ML efficiency improvements.The stated implications include sustainability and usability benefits without novel ethical concerns beyond those associated with efficiency improvements.
  • Impact: 2.5× higher throughput and 10.7× lower KV memory are reported while matching Full Attention accuracy.These gains are presented as enabling more efficient and accessible deployment of reasoning-capable LLMs.
  • Technical basis: RoPE links attention logits to relative position through frequency-specific trigonometric terms and phase differences.The coefficients and phases depend on Q/K vectors, while concentrated vectors permit center-based approximation.
  • Scientific implication: The trigonometric series can synthesize distance-dependent attention patterns, including preferences for local positions or distant attention sinks.Different learned Q/K centers shape different attention-versus-distance curves.
  • Method: TriAttention exploits stable Q/K concentration in pre-RoPE space to estimate key importance through trigonometric series.When concentration is high, Q/K centers determine predictable attention-versus-distance curves.
  • Technical basis: Mean Resultant Length quantifies directional concentration, with R = 1 indicating identical directions and R = 0 indicating no concentration.The method uses concentration to determine when trigonometric scoring is reliable and when norm-based scoring should contribute more.

C. Recursive State Query Benchmark Details

The Recursive State Query benchmark evaluates whether KV compression preserves information needed for multi-step graph traversal. It emphasizes exact historical state retention, deterministic ground truth, and sensitivity to memory loss.

  • Benchmark task: The benchmark asks models to simulate k-step depth-first search on an undirected graph and report traversal state.Reported state includes the current node, complete stack path, and visited nodes.
  • Method visualization: Figure B visualizes TriAttention as a four-stage pipeline: compute centers, score distance preferences, add norm scores, and prune the cache.The visualization includes real attention maps corresponding to the schematic pipeline.
  • Benchmark rationale: The task distributes information across the sequence, scales difficulty with step count, and provides unique ground truth through deterministic traversal.Neighbors are selected in ascending order, making the expected result deterministic.
  • Evaluation metric: Stack exact match is the primary metric because it requires the complete traversal path to be correct in order.The paper identifies it as the strictest metric and most sensitive to information loss from KV pruning.
  • Benchmark rationale: Recursive simulation exposes memory failures because forgotten intermediate states corrupt subsequent return values during backtracking.The task requires recalling states from earlier in the sequence after nested descent.

F. Evaluation on Diverse Benchmarks

TriAttention generalizes beyond mathematical reasoning across LongBench and RULER, achieving the strongest reported aggregate or task-level results against the evaluated baselines.

  • LongBench: 48.1 average on LongBench is the highest across 16 subtasks, with TriAttention winning 11 of 16.It surpasses Ada-KV+SnapKV by +2.5, reporting 48.1 versus 45.6.
  • RULER: 66.1 on RULER is reported as outperforming all baselines, with a +10.5 gap over SnapKV.RULER evaluates retrieval tasks at 4K context.
  • Memory-constrained comparison: Against H2O on the 12 LongBench subtasks fitting in 48GB GPU memory, TriAttention wins 10 of 12.The comparison notes that H2O requires O(n^2) memory and cannot use FlashAttention.

G. Future Offset Design Ablation

The future-offset ablation shows that broader and geometrically spaced sampling improves scoring, while performance remains stable across calibration sizes and data sources.

  • Offset range: +7.1% accuracy results when the maximum future distance increases from 128 to 4096, reaching 48.8% from 41.7%.The result supports including longer future offsets in importance estimation.
  • Spacing strategy: 45.8% versus 28.7% accuracy shows geometric spacing outperforming linear spacing by 17.1 percentage points.The paper attributes this to denser sampling being needed at near distances.
  • Calibration robustness: 45.4–45.8% performance remains stable across calibration sizes from 50k to 960k tokens.The reported calibration-size range does not produce a substantial accuracy change.
  • Calibration robustness: 46.2% with Google homepage HTML is comparable to 46.7% with ShareGPT chat data.The paper reports no clear correlation between calibration-data quality and accuracy.

I. MLA Architecture Validation

The MLA evaluation tests whether Q/K concentration extends beyond standard attention architectures. MLA shows comparable or stronger concentration and reconstruction quality than GQA, with more heads exceeding the reported reconstruction threshold.

  • I. MLA Architecture Validation: MLA evaluation compares reconstruction quality (Pearson r) and directional concentration (MRL) against GQA.The comparison uses GLM-4.7-Flash with Multi-head Latent Attention and 940 heads.
  • I. MLA Architecture Validation: 96.6% of MLA heads achieve R > 0.95, compared with 84.7% for GQA.The result indicates that Q/K concentration is architecture-general and potentially more pronounced in MLA.
  • I. MLA Architecture Validation: LongBench evaluates 16 subtasks spanning QA, summarization, few-shot classification, retrieval, counting, and code at a 50% KV budget.The table marks the best result among compression methods in bold.
  • I. MLA Architecture Validation: RULER reports retrieval results for Qwen3-8B at a 50% KV budget and 4K context.The table identifies the best result among compression methods in bold.

J. Real-World Deployment: Agentic Task on a Single GPU

The OpenClaw deployment evaluates TriAttention for a long-context, multi-turn agent on a single RTX 4090 with limited KV-cache memory. Full attention runs out of memory, whereas TriAttention stays within budget and completes the document-processing task.

  • J. Real-World Deployment: Agentic Task on a Single GPU: The deployment serves Qwen3-32B with AWQ INT4 quantization on a single RTX 4090 with 24GB memory.Model weights consume most GPU memory, leaving limited space for the KV cache; the first prompt already exceeds 15k tokens and each interaction expands the context.
  • J. Real-World Deployment: Agentic Task on a Single GPU: The deployment scenario is a multi-turn OpenClaw interaction in which the agent reads and processes six documents.Each interaction round further expands the context, increasing pressure on the limited KV-cache budget.
  • J. Real-World Deployment: Agentic Task on a Single GPU: Full attention’s KV cache grows unboundedly and causes an out-of-memory error before the agent completes its task.The setup cannot use FlashAttention because it materializes the full O(n^2) attention matrix.
  • J. Real-World Deployment: Agentic Task on a Single GPU: TriAttention keeps KV-cache memory within budget throughout the session, allowing the agent to read all documents and generate the report.A screenshot of the completed session is shown in Figure C.
  • J. Real-World Deployment: Agentic Task on a Single GPU: The section also references future-offset and calibration-data sensitivity ablations on Qwen3-8B using AIME24.The supplied table descriptions identify these as ablation settings but do not provide their numerical outcomes.
Loading 2604.04921v1…