The future-offset average in TriAttention can be precomputed

Amarnath Mukherjee · July 2026 · full note (PDF)

A small, exact note on one step of a recent KV-cache method: the average TriAttention takes over future distances need not be recomputed per key — it folds, once and offline, into a single weight per frequency band.

Long-context language models keep a growing cache of past keys and values. To stay inside a memory budget, that cache has to be pruned, and which entries you evict matters: drop a key the model still needs and accuracy suffers. TriAttention is a recent and rather 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 note is about one small step inside that scoring. To capture a key's importance across the many distances at which it might later be read, TriAttention averages its score over a schedule of future offsets — in their implementation, the powers of two from $2^0$ to $2^{16}$. That particular schedule is a design choice, and nothing below depends on it. The point I'd like to make is that whatever the schedule, this average need not be recomputed for each key: it collapses, by a one-line identity, into a single weight per frequency band that you compute once, offline. After that, scoring a key is one evaluation rather than one per offset, and not a single key changes rank. It is a small result. But TriAttention is the kind of method we will see more of, and it is a clean example of a pattern worth naming: an average over a fixed schedule of rotations is often a precompute in disguise.

The score TriAttention computes

TriAttention rests on a striking empirical regularity. Measured before RoPE's position-dependent rotation is applied, a head's query and key vectors don't spread out — across Qwen3, Qwen2.5, and Llama3, they cluster tightly around a fixed, non-zero center that barely moves across positions or inputs (roughly 90% of heads). That matters asymmetrically. At pruning time a cached key is known exactly; the only unknown is the future query that will read it — and because queries concentrate, their calibrated center $\mathbb{E}[q_f]$ is a faithful stand-in for it. Pin the query to that center, and the only thing still varying in the attention logit is the rotation, i.e. the query–key distance $\Delta$.

So TriAttention predicts the attention a key $k$ would receive at distance $\Delta$ as a score that splits into a distance-dependent trigonometric part and a norm correction:

$$ S(k,\Delta) = S_{\mathrm{trig}}(k,\Delta) + S_{\mathrm{norm}}(k). $$

The part we care about collects a contribution from each two-dimensional rotary band $f$:

$$ S_{\mathrm{trig}}(k,\Delta) = \sum_f \lVert \mathbb{E}[q_f]\rVert\,\lVert k_f\rVert\, \cos(\omega_f \Delta + \varphi_f), $$

where $f \in \{0, \dots, d/2-1\}$ indexes the bands, $\omega_f = \theta^{-2f/d}$ is band $f$'s angular frequency (the rate at which that band's rotation turns, with rotary base $\theta = 10000$), and $\varphi_f = \arg\mathbb{E}[q_f] - \arg k_f$ is the angle between the query center and the key's band-$f$ component. The norm term $S_{\mathrm{norm}}(k)$ doesn't depend on $\Delta$; it plays no role here, so set it aside.

The same score, read as arrows in the plane

The rotation acts on two coordinates 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. If you did phasors once — three-phase circuits in a college EE class, say — this is the same picture. Read band $f$'s pre-rotation query and key as complex numbers $q_f, k_f \in \mathbb{C}$. Three small facts are all we need.

Three panels: a complex number as an arrow; its conjugate mirrored across the real axis; the product carrying the angle between query and key. — Reading band $f$ as arrows. (A) $q_f$ and $k_f$ are arrows with a length and an angle. (B) The conjugate $\overline{k_f}$ is $k_f$ mirrored across the real axis: same length, opposite angle. (C) Because multiplying complex numbers adds their angles, $q_f\,\overline{k_f}$ is one arrow carrying the gap *between* query and key, $\arg q_f - \arg k_f$. Multiply further by $e^{\,i\omega_f\Delta}$ to turn it by the distance angle; its real part is the cosine term of the score.

The bar flips the angle. The conjugate $\overline{k_f}$ is the mirror of $k_f$ across the real axis: same length, opposite angle.
Multiplying adds angles. So $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$; with the center in the query's place, that gap is exactly $\varphi_f$.
The position rotation is itself a multiplication. Turning band $f$ by the distance angle $\omega_f\Delta$ means multiplying by $e^{\,i\omega_f\Delta}$.

Put the three together, with the score's center $\mathbb{E}[q_f]$ in the query's place. 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 one term of $S_{\mathrm{trig}}$. Summing over bands, the trigonometric score is the real part of a sum of arrows — the same complex form RoFormer gives the rotary dot product:

$$ S_{\mathrm{trig}}(k,\Delta) = \operatorname{Re}\Big\{ \sum_f \mathbb{E}[q_f]\, \overline{k_f}\, e^{\,i\omega_f\Delta} \Big\}. $$

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}\}$, spaced geometrically because attention changes fast at short distances and slowly at long ones:

$$ \tilde S(k) = \frac{1}{|D|} \sum_{\delta \in D} S(k, \Delta + \delta). $$

Done literally, this evaluates the whole band-sum once per offset, for every key, and averages — which is what the reference implementation does, materializing a [keys × offsets × bands] table of cosines and averaging over the offset axis.

Now look at where the offset $\delta$ actually goes. In the phasor score it appears only inside the rotation, and because adding angles multiplies rotations, a turn by the shifted distance factors cleanly:

$$ e^{\,i\omega_f(\Delta + \delta)} = \underbrace{e^{\,i\omega_f\Delta}}_{\text{distance, kept}}\; \underbrace{e^{\,i\omega_f\delta}}_{\text{offset, averaged}}. $$

Only the offset factor depends on $\delta$, and it touches neither the key nor the center. So averaging over the offsets reaches only that factor, where it collapses into 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}}_{W_f} \Big\} + S_{\mathrm{norm}}(k). $$

Two panels of unit arrows on a circle averaging to a single green weight: high-frequency band cancels to a short weight, low-frequency band reinforces to a long one. — The offset average is an average of rotations. Each offset $\delta$ contributes a unit arrow $e^{\,i\omega_f\delta}$ on the circle; their average is the single weight $W_f$ (green). In a high-frequency band (left) the offsets fan out and largely cancel, so $W_f$ is short; in a low-frequency band (right) they cluster and reinforce, so $W_f$ is long. Either way it is one number per band, computed once.

The weight $W_f = \frac{1}{|D|}\sum_{\delta\in D} e^{\,i\omega_f\delta}$ depends only on the fixed frequencies $\{\omega_f\}$ and the chosen offset set — whatever that set is; the powers of two are TriAttention's choice, not something the collapse relies on. It has nothing to do with any particular key. So it is computed once, offline, alongside the centers — and then the averaged score is a single evaluation of the original score, with the query center nudged to a precomputed “offset-averaged center” $\mu_f' = \mathbb{E}[q_f]\,W_f$:

$$ \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, \quad \varphi_f' = \arg\mu_f' - \arg k_f. $$

Because this is an exact rewrite of the average — not an approximation of it — every key receives precisely the score it had before, and the pruned set is unchanged. And nothing here forces complex arithmetic into the hot loop: the weight $W_f$ folds offline into the same real cosine/sine coefficients TriAttention already uses, as a fixed $2\times2$ per-band map. The existing scorer runs unchanged on the pre-mapped coefficients. (The PDF works this out and reports a numerical check — the complex, polar, and real-coefficient forms all match the literal offset loop to machine precision.)

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. The reduction on that first term is an exact factor of $|D|$, the number of offsets in the schedule.

I want to be plain about the size and place of this, because it is easy to oversell. The saving is in the scoring arithmetic of the pruning step, which TriAttention already runs only once per window of 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 the method's other engineering rather than replacing any of it. And it relies on the average being linear, so it applies to the published mean-aggregation default but not to a non-default variant that takes a maximum over offsets instead.

In short: a step that looked like it cost one evaluation per offset costs a single evaluation, paid for once, offline, without changing any output. A small, exact tightening of one bolt in their machine.

Full note · PDF Precomputing the Future-Offset Average in TriAttention Includes the lemma, proof, the real-coefficient $2\times2$ form, and the numerical check.