MIT·NVIDIA, 트라이어텐션 제안…KV 캐시 10배 절감

Editors Pick Agentic AI Artificial Intelligence AI Infrastructure Tech News AI Paper Summary Technology AI Shorts Applications Deep Learning Language Model Large Language Model Machine Learning New Releases Open Source Software Engineering Staff 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 — 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. A team of researchers from MIT, NVIDIA, and Zhejiang University proposed a method called TriAttention that directly addresses this problem. On the AIME25 mathematical reasoning benchmark with 32K-token generation, TriAttention matches Full Attention accuracy while achieving 2.5× higher throughput or 10.7× KV memory reduction. Leading baselines achieve only about half the accuracy at the same efficiency level. The Problem with Existing KV Cache Compression To understand why TriAttention is important, it helps to understand the standard approach to KV cache compression. Most existing methods — including SnapKV, H2O, and R-KV — 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. The catch is that these methods operate in what the research team calls post-RoPE space. RoPE, or Rotary Position Embedding , 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. This rotation means that only the most recently generated queries have orientations that are ‘up to date' for estimating which keys are important right now. Prior work has confirmed this empirically: increasing the observation window for importance estimation does not help — 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. This problem is especially acute for what the research team calls retrieval heads — 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. The Pre-RoPE Observation: Q/K Concentration The key insight in TriAttention comes from looking at Query and Key vectors before RoPE rotation is applied — 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 , and measures it using the Mean Resultant Length R — a standard directional statistics measure where R → 1 means tight clustering and R → 0 means dispersion in all directions. 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 — they are an intrinsic property of the model's 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–0.980. 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. From Concentration to a Trigonometric Series The research team then show mathematically that when Q and K vectors are concentrated around their centers, the attention logit — the raw score before softmax that determines how much a query attends to a key — 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 (the relative positional gap between query and key), expressed as a trigonometric series: logit ( Δ ) ≈ ∑ f ‖ q ‾ f ‖ ‖ k ‾ f ‖ ⏟ amplitude cos ⁡ ( ω f Δ + ϕ ‾ f ⏟ phase ) = ∑ f [ a f cos ⁡ ( ω f Δ ) + b f sin ⁡ ( ω f Δ ) ] \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)] Here, Δ is the positional distance, ω f are the RoPE rotation frequencies for each frequency band f, and the coefficients a f and b f 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. 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–0.9. The research team further validates this on GLM-4.7-Flash, which uses Multi-head Latent Attention (MLA) rather than standard Grouped-Query Attention — 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. How TriAttention Uses This 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: The Trigonometric Series Score (S trig ) 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. S trig ( k , Δ ) = ∑ f ‖ 𝔼 [ q f ] ‖ ⋅ ‖ k f ‖ ⋅ cos ⁡ ( ω f Δ + ϕ f ) S_{\text{trig}}(k, \Delta) = \sum_{f} \|\mathbb{E}[q_f]\| \cdot \|k_f\| \cdot \cos(\omega_f \Delta + \phi_f) The Norm-Based Score (S norm ) 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. S norm ( 0 ) ( k ) = ∑ f 𝔼 [ ‖ q f ‖ ] ⋅ ‖ k f ‖ S_{\text{norm}}^{(0)}(k) = \sum_{f} \mathbb{E}[\|q_f\|] \cdot \|k_f\| The two scores are combined using the Mean Resultant Length R as an adaptive weight: when concentration is high, S trig do