1 Introduction↩︎

In long-context LLMs the KV cache grows large, and to stay within a memory budget it must periodically be pruned. Which entries to evict is a consequential choice: drop a key the model still needs and its accuracy suffers.

TriAttention [1] is a recent and elegant answer to that choice. When the cache must be trimmed, it scores each cached key by how much attention it is likely to receive, and keeps the highest-scoring ones.

This paper contributes a small improvement to one step of that method. We first describe the part of TriAttention our result touches, then the improvement itself.

2 The TriAttention score↩︎

TriAttention’s score rests on a striking regularity in how trained models represent queries and keys.

2.0.0.1 Their empirical observation.

Across several leading-edge models (Qwen3, Qwen2.5, Llama3), and measured before the RoPE position-dependent rotation [2] is applied, a head’s query and key vectors do not spread out. They cluster tightly around a fixed, non-zero center, and that center barely moves across positions or inputs. Roughly \(90\%\) of heads show this, and it is stable across task domains. The consequence matters asymmetrically for what comes next. At pruning time a cached key is already known exactly; the only thing that must be approximated is the future query that will read it—and because queries concentrate, their center is a faithful stand-in for it. With the query pinned to that center, the only thing still varying in the attention logit is the rotation, i.e.the query–key distance \(\Delta\). Attention as a function of distance becomes a fixed, predictable curve you can compute once, offline—no need to watch live attention scores. This is why TriAttention scores keys in the pre-rotation space: it is where the query center is stable enough to be a reliable summary.

2.0.0.2 TriAttention Score.

At pruning time the keys are concrete (they sit in the cache), but the queries that will read them are in the future and do not exist yet. So TriAttention substitutes the calibrated query center \(\mathbb{E}[q_f]\)—the average of band-\(f\) queries over a calibration set—for the unknown future query and predicts the attention a key \(k\) would receive at distance \(\Delta\) ([1], Eq. 10): \[S(k,\Delta)=S_{\mathrm{trig}}(k,\Delta)+S_{\mathrm{norm}}(k), \label{eq:score}\tag{1}\] whose trigonometric term ([1], Eq. 6) collects the per-band contributions, \[S_{\mathrm{trig}}(k,\Delta)=\sum_f \lVert \mathbb{E}[q_f] \rVert\,\lVert k_f \rVert\cos(\omega_f\Delta+\varphi_f), \label{eq:strig}\tag{2}\] where:

• \(f\in\{0,\dots,d/2-1\}\) indexes the rotation’s two-dimensional bands, which TriAttention calls frequency bands;

• \(\omega_f=\theta^{-2f/d}\) is band \(f\)’s angular frequency—the rate at which the rotation turns that band;

• \(\varphi_f=\arg\mathbb{E}[q_f]-\arg k_f\) is the angle between the query center and the key’s band-\(f\) component \(k_f\);

• \(d\) is the head dimension; and

• \(\theta=10000\) is the rotary base TriAttention uses—the value RoPE adopted from the sinusoidal positional encoding of the original Transformer [3].

The second term \(S_{\mathrm{norm}}(k)\) is a norm-based correction that does not depend on \(\Delta\); it is mentioned for completeness, but is not relevant to our discussion in this paper.

2.0.0.3 The score as a picture in the plane.

The rotation acts on two coordinates of the query and key at a time, turning that pair by an angle. A pair of coordinates is a point in the plane, and a point in the plane is a complex number: an arrow with a length and an angle. So read band \(f\)’s pre-rotation query and key as complex numbers \(q_f,k_f\in\mathbb{C}\) (1A). Three small facts are all we need.

1. The bar flips the angle. The conjugate \(\overline{k_f}\) is the mirror image of \(k_f\) across the real axis: same length, opposite angle (1B).

2. Multiplying adds angles. So the product \(q_f\,\overline{k_f}\) is a single arrow whose angle is the query–key gap \(\arg q_f-\arg k_f\), and whose length is \(\lVert q_f \rVert\lVert k_f \rVert\) (1C); with the center in the query’s place this gap is \(\varphi_f\) (2 ).

3. The position-dependent rotation is itself just a multiplication: turning band \(f\) by the distance angle \(\omega_f\Delta\) means multiplying by \(e^{\,i\omega_f\Delta}\).

Put the three together, now with the score’s center \(\mathbb{E}[q_f]\) in place of the unknown query. The arrow \(\mathbb{E}[q_f]\,\overline{k_f}\,e^{\,i\omega_f\Delta}\) has length \(\lVert \mathbb{E}[q_f] \rVert\lVert k_f \rVert\) and angle \(\varphi_f+\omega_f\Delta\), and its real part is exactly \(\lVert \mathbb{E}[q_f] \rVert\lVert k_f \rVert\cos(\omega_f\Delta+\varphi_f)\)—one term of \(S_{\mathrm{trig}}\). Summing over bands, the trigonometric part of the score, 2 , is the real part of a sum of arrows—the same complex form RoFormer gives the rotary dot product ([2] Eq. 12): \[S_{\mathrm{trig}}(k,\Delta)=\operatorname{Re}\!\Big\{\sum_f \mathbb{E}[q_f]\,\overline{k_f}\,e^{\,i\omega_f\Delta}\Big\}. \label{eq:phasor}\tag{3}\]

3 The observation↩︎

3.0.0.1 The averaging step—the one we improve.

A cached key might be queried soon or much later. Its real importance is its importance across all those future distances, so TriAttention averages the score over a set of future offsets \(D=\{1,2,4,\dots,2^{16}\}\) ([1], Eq. 11): \[\tilde{S}(k)=\frac{1}{|D|}\sum_{\delta\in D} S(k,\Delta+\delta), \qquad |D|=17. \label{eq:avg}\tag{4}\] The offsets are spaced geometrically (powers of two) because attention changes fast at short distances and slowly at long ones. Done as written, 4 evaluates the whole band sum \(17\) times per key and averages—and that is exactly what the reference implementation does, materializing a [keys \(\times\) offsets \(\times\) bands] table of cosines and averaging over the offset axis.²

Now, consider where the offset \(\delta\) goes. In the phasor score, 3 , it appears only inside the rotation, and—because adding angles multiplies rotations—a turn by the shifted distance factors into the part we keep and the part we average: \[e^{\,i\omega_f(\Delta+\delta)} =\underbrace{e^{\,i\omega_f\Delta}}_{\text{distance, kept}}\; \underbrace{e^{\,i\omega_f\delta}}_{\text{offset, averaged}}. \label{eq:split}\tag{5}\] Only the offset factor depends on \(\delta\), and it touches neither the key nor the center. So averaging the phasor score over the \(17\) offsets reaches only that factor, where it collapses to one complex weight per band: \[\tilde{S}(k)=\operatorname{Re}\!\Big\{\sum_f \mathbb{E}[q_f]\,\overline{k_f}\,e^{\,i\omega_f\Delta}\, \underbrace{\tfrac{1}{|D|}\sum_{\delta\in D}e^{\,i\omega_f\delta}}_{\textstyle W_f}\Big\}+S_{\mathrm{norm}}(k). \label{eq:collapse}\tag{6}\] The averaged offset rotation \(W_f=\frac{1}{|D|}\sum_{\delta\in D}e^{\,i\omega_f\delta}\) (2) depends only on the frequencies \(\{\omega_f\}\) and the fixed offset set, so it is computed once, offline, alongside the centers.

With \(W_f\) in hand, the averaged score is a single evaluation of the original score, 1 , with the query center \(\mathbb{E}[q_f]\) nudged to a precomputed “offset-averaged center”³ \(\mu_f'=\mathbb{E}[q_f]\,W_f\), with \(\varphi_f'=\arg\mu_f'-\arg k_f\): \[\boxed{\;\tilde{S}(k)=\sum_f \lVert \mu_f' \rVert\,\lVert k_f \rVert\cos(\omega_f\Delta+\varphi_f')+S_{\mathrm{norm}}(k), \qquad \mu_f'=\mathbb{E}[q_f]\,W_f.\;} \label{eq:folded}\tag{7}\] After that, scoring a key costs one band-sum evaluation instead of \(17\), and—because 7 is an exact rewrite of 4 , not an approximation—every key receives precisely the score it had before. The pruned set is unchanged. The complex weights \(W_f\) need not appear in the scoring loop: folded offline into the centers (or the real coefficients), they leave the computation on the real cosine/sine form TriAttention already uses; 7 makes that concrete as a fixed \(2\times2\) per-band map and reports a numerical check.

4 What it saves, and what it does not↩︎

For one pruning event over \(N\) cached keys and \(F\) bands, the trigonometric term drops from \(O(N F\,|D|)\) band-cosine evaluations to \(O(N F)\), plus a one-time offline cost of \(O(F\,|D|)\) to build the weights \(\{W_f\}\). The reduction on the first term is an exact factor of \(|D|=17\).

It is worth being clear about the size and place of this. The saving is in the scoring arithmetic of the pruning step, which TriAttention already runs only once per window of \(\beta=128\) generated tokens. It is not a saving in the attention kernel that runs every token, nor in the cache budget, nor in how often pruning happens. It composes with, and does not replace, the method’s other engineering: the window batching that decides how often scoring runs, and the grouped-query normalize-then-max step that consumes the per-head \(\tilde{S}(k)\) unchanged. This rewrite is exact for the average in 4 , which is the published method and the default in the code; it relies on the average being linear, so it does not apply to a non-default variant that takes a maximum over offsets instead of a mean. In short: a key step that looked like it cost \(17\times\) costs \(1\times\), paid for once, offline, without changing any output—a small, exact tightening of one bolt in their machine.

5 Conclusion↩︎

TriAttention’s future-offset average admits an offline computation: because the offset lives entirely in the position-dependent rotation, it folds exactly into a per-band weight that is precomputed offline, turning a \(17\)-fold loop into a single evaluation, with no change to which keys are kept. We offer it as a small, exact complement to a method we expect to see more of, and as a reminder that “average over a geometric ladder of rotations” is often a precompute in disguise.

6 Derivation↩︎

The rewrite used in the body is exact; here is the one line behind it.

7 Real-coefficient form↩︎

No complex arithmetic is required at inference. TriAttention writes the single-distance score with constant real coefficients ([1], Eqs. 17–19), \[S_{\mathrm{trig}}(k,\Delta)=\sum_f\big[a_f\cos(\omega_f\Delta)+b_f\sin(\omega_f\Delta)\big],\] with \(a_f=\lVert \mathbb{E}[q_f] \rVert\lVert k_f \rVert\cos\varphi_f\) and \(b_f=-\lVert \mathbb{E}[q_f] \rVert\lVert k_f \rVert\sin\varphi_f\). In these coordinates, 6 is a fixed \(2\times2\) per-band linear map, with \(A_f=\operatorname{Re}W_f\) and \(B_f=\Im W_f\): \[\begin{align} \begin{pmatrix}a_f'\\ b_f'\end{pmatrix} &=\begin{pmatrix}A_f & B_f\\ -B_f & A_f\end{pmatrix} \begin{pmatrix}a_f\\ b_f\end{pmatrix},\\[4pt] \tilde{S}_{\mathrm{trig}}(k)&=\sum_f\big[a_f'\cos(\omega_f\Delta)+b_f'\sin(\omega_f\Delta)\big]. \end{align} \label{eq:realmap}\tag{8}\]
The existing real-valued scoring form runs unchanged on the pre-mapped coefficients; complex numbers appear only in this derivation. A per-band, offset-independent rescaling that the code already applies (a high-frequency mask) multiplies \(\lVert \mu_f' \rVert\) and is unaffected.

7.0.0.1 Numerical confirmation.

We checked 7 8 by comparing the literal \(17\)-fold offset loop (a faithful reimplementation of the reference scorer) against the complex, polar, and real-coefficient forms on synthetic centers and keys. All three match the baseline to machine precision (\(\max|\Delta|\!\sim\!10^{-12}\)) across \(\delta_{\max}\in\{2^7,2^{12},2^{16}\}\).

References↩︎