{"id":706,"date":"2026-04-12T04:10:41","date_gmt":"2026-04-11T20:10:41","guid":{"rendered":"https:\/\/connectword.dpdns.org\/?p=706"},"modified":"2026-04-12T04:10:41","modified_gmt":"2026-04-11T20:10:41","slug":"researchers-from-mit-nvidia-and-zhejiang-university-propose-triattention-a-kv-cache-compression-method-that-matches-full-attention-at-2-5x-higher-throughput","status":"publish","type":"post","link":"https:\/\/connectword.dpdns.org\/?p=706","title":{"rendered":"Researchers from MIT, NVIDIA, and Zhejiang University Propose TriAttention: A KV Cache Compression Method That Matches Full Attention at 2.5\u00d7 Higher Throughput"},"content":{"rendered":"

Long-chain reasoning is one of the most compute-intensive tasks in modern large language models. When a model like DeepSeek-R1 or Qwen3 works through a complex math problem, it can generate tens of thousands of tokens before arriving at an answer. Every one of those tokens must be stored in what is called the KV cache \u2014 a memory structure that holds the Key and Value vectors the model needs to attend back to during generation. The longer the reasoning chain, the larger the KV cache grows, and for many deployment scenarios, especially on consumer hardware, this growth eventually exhausts GPU memory entirely.<\/p>\n

A team of researchers from MIT, NVIDIA, and Zhejiang University proposed a method called TriAttention<\/strong> that directly addresses this problem. On the AIME25 mathematical reasoning benchmark with 32K-token generation, TriAttention matches Full Attention accuracy while achieving 2.5\u00d7 higher throughput or 10.7\u00d7 KV memory reduction. Leading baselines achieve only about half the accuracy at the same efficiency level.<\/p>\n

\n
\"\"
https:\/\/arxiv.org\/pdf\/2604.04921<\/figcaption><\/figure>\n<\/div>\n

The Problem with Existing KV Cache Compression<\/strong><\/h3>\n

To understand why TriAttention is important, it helps to understand the standard approach to KV cache compression. Most existing methods \u2014 including SnapKV, H2O, and R-KV \u2014 work by estimating which tokens in the KV cache are important and evicting the rest. Importance is typically estimated by looking at attention scores: if a key receives high attention from recent queries, it is considered important and kept.<\/p>\n

The catch is that these methods operate in what the research team calls post-RoPE<\/em> space. RoPE, or Rotary Position Embedding<\/strong>, is the positional encoding scheme used by most modern LLMs including Llama, Qwen, and Mistral. RoPE encodes position by rotating the Query and Key vectors in a frequency-dependent way. As a result, a query vector at position 10,000 looks very different from the same semantic query at position 100, because its direction has been rotated by the position encoding.<\/p>\n

This rotation means that only the most recently generated queries have orientations that are \u2018up to date\u2019 for estimating which keys are important right now. Prior work has confirmed this empirically: increasing the observation window for importance estimation does not help \u2014 performance peaks at around 25 queries and declines after that. With such a tiny window, some keys that will become important later get permanently evicted.<\/p>\n

This problem is especially acute for what the research team calls retrieval heads<\/em> \u2014 attention heads whose function is to retrieve specific factual tokens from long contexts. The relevant tokens for a retrieval head can remain dormant for thousands of tokens before suddenly becoming essential to the reasoning chain. Post-RoPE methods, operating over a narrow observation window, see low attention on those tokens during the dormant period and permanently evict them. When the model later needs to recall that information, it is already gone, and the chain of thought breaks.<\/p>\n

The Pre-RoPE Observation: Q\/K Concentration<\/strong><\/h3>\n

The key insight in TriAttention comes from looking at Query and Key vectors before<\/em> RoPE rotation is applied \u2014 the pre-RoPE space. When the research team visualized Q and K vectors in this space, they found something consistent and striking: across the vast majority of attention heads and across multiple model architectures, both Q and K vectors cluster tightly around fixed, non-zero center points. The research team terms this property Q\/K concentration<\/strong>, and measures it using the Mean Resultant Length<\/strong> R \u2014 a standard directional statistics measure where R \u2192 1 means tight clustering and R \u2192 0 means dispersion in all directions.<\/p>\n

On Qwen3-8B, approximately 90% of attention heads exhibit R > 0.95, meaning their pre-RoPE Q\/K vectors are nearly perfectly concentrated around their respective centers. Critically, these centers are stable across different token positions and across different input sequences \u2014 they are an intrinsic property of the model\u2019s learned weights, not a property of any particular input. The research team further confirm that Q\/K concentration is domain-agnostic: measuring Mean Resultant Length across Math, Coding, and Chat domains on Qwen3-8B yields nearly identical values of 0.977\u20130.980.<\/p>\n

This stability is what post-RoPE methods cannot exploit. RoPE rotation disperses these concentrated vectors into arc patterns that vary with position. But in pre-RoPE space, the centers remain fixed.<\/p>\n

From Concentration to a Trigonometric Series<\/strong><\/h3>\n

The research team then show mathematically that when Q and K vectors are concentrated around their centers, the attention logit \u2014 the raw score before softmax that determines how much a query attends to a key \u2014 simplifies dramatically. Substituting the Q\/K centers into the RoPE attention formula, the logit reduces to a function that depends only on the Q-K distance<\/strong> (the relative positional gap between query and key), expressed as a trigonometric series:<\/p>\n

logit<\/mtext>(<\/mo><\/mrow>\u0394<\/mi><\/mrow>)<\/mo>\u2248<\/mo>\u2211<\/mo>f<\/mi><\/msub><\/munder>\u2016<\/mi>q<\/mi>\u203e<\/mo><\/mover>f<\/mi><\/msub>\u2016<\/mi>\u2016<\/mi>k<\/mi>\u203e<\/mo><\/mover>f<\/mi><\/msub>\u2016<\/mi><\/mrow>\u23df<\/mo><\/munder>amplitude<\/mtext><\/mrow>cos<\/mi>\u2061<\/mo><\/mrow>(<\/mo>\u03c9<\/mi>f<\/mi><\/msub>\u0394<\/mi><\/mrow>+<\/mo><\/munder>\u03d5<\/mi>\u203e<\/mo><\/mover>f<\/mi><\/msub>\u23df<\/mo><\/munder>phase<\/mtext><\/mrow>)<\/mo>=<\/mo>\u2211<\/mo>f<\/mi><\/msub>[<\/mo>a<\/mi>f<\/mi><\/msub>cos<\/mi>\u2061<\/mo><\/mrow>(<\/mo>\u03c9<\/mi>f<\/mi><\/msub>\u0394<\/mi><\/mrow>)<\/mo>+<\/mo>b<\/mi>f<\/mi><\/msub>sin<\/mi>\u2061<\/mo><\/mrow>(<\/mo>\u03c9<\/mi>f<\/mi><\/msub>\u0394<\/mi><\/mrow>)<\/mo>]<\/mo> text{logit}(Delta) approx sum_{f} underbrace{|bar{q}_f| |bar{k}_f|}_{text{amplitude}} cos(omega_f Delta + underbrace{bar{phi}_f}_{text{phase}}) = sum_{f} [a_f cos(omega_f Delta) + b_f sin(omega_f Delta)] <\/annotation><\/semantics><\/math><\/p>\n

Here, \u0394 is the positional distance, \u03c9f<\/sub> are the RoPE rotation frequencies for each frequency band f, and the coefficients af<\/sub> and bf<\/sub> are determined by the Q\/K centers. This series produces a characteristic attention-vs-distance curve for each head. Some heads prefer nearby keys (local attention), others prefer very distant keys (attention sinks). The centers, computed offline from calibration data, fully determine which distances are preferred.<\/p>\n

The research team validated this experimentally across 1,152 attention heads in Qwen3-8B and across Qwen2.5 and Llama3 architectures. The Pearson correlation between the predicted trigonometric curve and the actual attention logits has a mean above 0.5 across all heads, with many heads achieving correlations of 0.6\u20130.9. The research team further validates this on GLM-4.7-Flash, which uses Multi-head Latent Attention (MLA)<\/strong> rather than standard Grouped-Query Attention \u2014 a meaningfully different attention architecture. On MLA, 96.6% of heads exhibit R > 0.95, compared to 84.7% for GQA, confirming that Q\/K concentration is not specific to one attention design but is a general property of modern LLMs.<\/p>\n

How TriAttention Uses This<\/strong><\/h3>\n

TriAttention is a KV cache compression method that uses these findings to score keys without needing any live query observations. The scoring function has two components:<\/strong><\/p>\n

The Trigonometric Series Score<\/strong> (Strig<\/sub>) uses the Q center computed offline and the actual cached key representation to estimate how much attention the key will receive, based on its positional distance from future queries. Because a key may be attended to by queries at many future positions, TriAttention averages this score over a set of future offsets using geometric spacing.<\/p>\n

S<\/mi>trig<\/mtext><\/msub>(<\/mo>k<\/mi>,<\/mo><\/mrow>\u0394<\/mi><\/mrow>)<\/mo>=<\/mo>\u2211<\/mo>f<\/mi><\/msub>\u2016<\/mi>\ud835\udd3c<\/mi>[<\/mo>q<\/mi>f<\/mi><\/msub>]<\/mo>\u2016<\/mi>\u22c5<\/mo>\u2016<\/mi>k<\/mi>f<\/mi><\/msub>\u2016<\/mi>\u22c5<\/mo>cos<\/mi>\u2061<\/mo><\/mrow>(<\/mo>\u03c9<\/mi>f<\/mi><\/msub>\u0394<\/mi><\/mrow>+<\/mo>\u03d5<\/mi>f<\/mi><\/msub>)<\/mo>S_{text{trig}}(k, Delta) = sum_{f} |mathbb{E}[q_f]| cdot |k_f| cdot cos(omega_f Delta + phi_f)<\/annotation><\/semantics><\/math><\/p>\n

The Norm-Based Score<\/strong> (Snorm<\/sub>) handles the minority of attention heads where Q\/K concentration is lower. It weights each frequency band by the expected query norm contribution, providing complementary information about token salience beyond distance preference alone.<\/p>\n

S<\/mi>norm<\/mtext>(<\/mo>0<\/mn>)<\/mo><\/mrow><\/msubsup>(<\/mo>k<\/mi>)<\/mo>=<\/mo>\u2211<\/mo>f<\/mi><\/msub>\ud835\udd3c<\/mi>[<\/mo>\u2016<\/mi>q<\/mi>f<\/mi><\/msub>\u2016<\/mi>]<\/mo>\u22c5<\/mo>\u2016<\/mi>k<\/mi>f<\/mi><\/msub>\u2016<\/mi><\/mrow>S_{text{norm}}^{(0)}(k) = sum_{f} mathbb{E}[|q_f|] cdot |k_f|<\/annotation><\/semantics><\/math><\/p>\n

The two scores are combined using the Mean Resultant Length R as an adaptive weight: when concentration is high, Strig<\/sub> dominates; when concentration is lower, Snorm<\/sub> contributes more. Every 128 generated tokens, TriAttention scores all keys in the cache and retains only the top-B, evicting the rest.<\/p>\n

Results on Mathematical Reasoning<\/strong><\/h3>\n

On AIME24 with Qwen3-8B, TriAttention achieves 42.1% accuracy against Full Attention\u2019s 57.1%, while R-KV achieves only 25.4% at the same KV budget of 2,048 tokens. On AIME25, TriAttention achieves 32.9% versus R-KV\u2019s 17.5% \u2014 a 15.4 percentage point gap. On MATH 500 with only 1,024 tokens in the KV cache out of a possible 32,768, TriAttention achieves 68.4% accuracy against Full Attention\u2019s 69.6%.<\/p>\n

\n
\"\"
https:\/\/arxiv.org\/pdf\/2604.04921<\/figcaption><\/figure>\n<\/div>\n

The research team also introduces a Recursive State Query benchmark<\/strong> based on recursive simulation using depth-first search. Recursive tasks stress memory retention because the model must maintain intermediate states across long chains and backtrack to them later \u2014 if any intermediate state is evicted, the error propagates through all subsequent return values, corrupting the final result. Under moderate memory pressure up to depth 16, TriAttention performs comparably to Full Attention, while R-KV shows catastrophic accuracy degradation \u2014 dropping from approximately 61% at depth 14 to 31% at depth 16. This indicates R-KV incorrectly evicts critical intermediate reasoning states.<\/p>\n

On throughput, TriAttention achieves 1,405 tokens per second on MATH 500 against Full Attention\u2019s 223 tokens per second, a 6.3\u00d7 speedup. On AIME25, it achieves 563.5 tokens per second against 222.8, a 2.5\u00d7 speedup at matched accuracy.<\/p>\n

\n
\"\"
https:\/\/arxiv.org\/pdf\/2604.04921<\/figcaption><\/figure>\n<\/div>\n

Generalization Beyond Mathematical Reasoning<\/strong><\/h3>\n

The results extend well beyond math benchmarks. On LongBench<\/strong> \u2014 a 16-subtask benchmark covering question answering, summarization, few-shot classification, retrieval, counting, and code tasks \u2014 TriAttention achieves the highest average score of 48.1 among all compression methods at a 50% KV budget on Qwen3-8B, winning 11 out of 16 subtasks and surpassing the next best baseline, Ada-KV+SnapKV, by 2.5 points. On the RULER<\/strong> retrieval benchmark at a 4K context length, TriAttention achieves 66.1, a 10.5-point gap over SnapKV. These results confirm that the method is not tuned to mathematical reasoning alone \u2014 the underlying Q\/K concentration phenomenon transfers to general language tasks.<\/p>\n

Key Takeaways<\/strong><\/h3>\n