Online attention: the rescaling that avoids the weight matrix

Abstract: a normalized sum can be computed in blocks

Producing one attention row does not require storing all its weights. A maximum, a normalizer, and an accumulated vector suffice, provided all change scale together when a larger score arrives. We derive the identity, verify it with logits that overflow naive exponentiation, and show a concrete bug: rescaling only the denominator gives a wrong output even when softmax appears numerically stable. This is a worked derivation with executed code, not a new algorithm or GPU benchmark.

1. The object we want to compute

A query q and N keys k_j of dimension d_k produce dimensionless scores s_j=qᵀk_j/√d_k. Any bias is included in the score. Each key has a value v_j in R^(d_v). Output o is a convex combination of values: coefficients are positive and sum to one. We consider one row, no dropout, and at least one unmasked key. Excluded keys contribute nothing. The question concerns evaluation order, not changing the mathematical function.

Subtracting the same maximum multiplies numerator and denominator by exp(−m), leaving the ratio unchanged. For finite scores, exponents become nonpositive: no term exceeds one and at least one equals one. Thus l lies between 1 and N in real arithmetic. This prevents exponential overflow, not unlimited numerical error: tiny terms can underflow and sums still round.

2. The invariant and change of scale

After t keys maintain m_t, the prefix maximum; l_t, the exponential sum relative to that maximum; and u_t, the corresponding value-weighted sum. A pair (s,v) arrives. The new maximum is m′=max(m_t,s). Multiply every old contribution exp(s_i−m_t) by exp(m_t−m′), turning it into exp(s_i−m′). Apply that factor to u_t too, since values used the same weights. Then add the new key’s contribution.

The proof is a loop invariant: substitute the definitions of l and u into the updates to obtain exactly the sums over the extended prefix. Initialize one key with m=s_1, l=1, u=v_1, avoiding −∞−(−∞). An empty state requires explicit handling. Distinguish unnormalized numerator u from output o: treating o as u loses the previous factor l.

3. An example exposing the bug

Use s=(1000,1001,999), v_1=(1,0), v_2=(0,2), and v_3=(3,−1). Python reports overflow for exp(1000), while stabilized weights are proportional to (e^−1,1,e^−2). After the first key l=1 and u=(1,0). The second raises the maximum from 1000 to 1001: the old numerator becomes (e^−1,0), not (1,0). Thus l=1+e^−1 and u=(e^−1,2). The third leaves the maximum unchanged and adds e^−2·(3,−1).

The final result is o=(0.514820191, 1.240451338). Rescaling l but not u at the second step instead gives first component 0.935332675. The second remains accidentally correct because its previous contribution was zero: checking only one component could miss the bug. Both outputs are finite. Absence of NaN and overflow does not establish algorithmic correctness.

4. Merging blocks without retaining weights

The same algebra merges disjoint sets A and B, each summarized by (m,l,u). Choose m=max(m_A,m_B), rescale both normalizers and numerators, then sum. The result represents the union, so in real arithmetic the operation is associative and commutative. Floating-point reduction trees can change the last bits. This permits parallel blocks without claiming bitwise identity with sequential order.

The archive compares block sizes 1, 7, 32, and 257 against stable materialization for 257 keys and four-component values. Python’s generator uses seed 20260924; scores are uniform in [−a,a] for a=1,10,1000 and values in [−2,2]. Maximum absolute error over twelve comparisons is about 6.66×10^−16. This is a synthetic Python-float check, not a universal bound or FP16/BF16 measurement. Code, seed, and results are archived.

5. Memory, complexity, and remaining measurements

For one row, persistent numerator state needs d_v+2 scalars, besides the query, key/value blocks, and temporary buffers. With d_v=64 in FP32 this is 264 bytes, not total kernel memory. A complete N×N FP32 score matrix takes 4N² bytes per head and batch element: 64 MiB at N=4096, 256 MiB at N=8192. A naive implementation may also materialize normalized weights. Avoiding these matrices differs from compressing K and V and does not remove a decoder’s required KV cache.

Dense computation of all query-key pairs remains quadratic: O(N²(d_k+d_v)). The recurrence does not turn exact attention into linear attention. A scalar Python loop can be slower than vectorized multiplication; hardware gains require suitable tiling, fusion, and memory access. Measuring them requires GPU, dtype, dimensions, mask, batch, versions, warmup, and synchronization details. We have not performed that benchmark and report no speedup.

Fully masked rows need an explicit convention: 0/0 defines no distribution. The code rejects a state without keys and requires finite scores; filter masked keys first. Mixed-sign values can cancel in the numerator, while reduced precision requires an accumulator choice. Gradients and dropout require additional logic too: checking one-row forward computation does not certify a complete training implementation.

6. Sources and conclusion

Milakov and Gimelshein describe online normalization (2018, arXiv v2); their Tesla V100, CUDA 9.1 experiments use different batch sizes, so timings do not transfer to our code. Dao and colleagues’ FlashAttention (2022, arXiv v2) links tiled computation to reduced HBM traffic and backward recomputation. We read algorithms, analysis, and experiments. This page explains one foundation without reproducing benchmarks or surveying subsequent developments.

The verifiable conclusion is the invariant shared by numerator and denominator. The maximum is not merely an overflow trick: it defines the numerical scale of accumulated contributions. When it changes, convert all previous contributions. That step removes the need for a weight matrix while preserving the attention result in real arithmetic.

Milakov M., Gimelshein N. (2018), Online normalizer calculation for softmax, arXiv:1805.02867v2.

Dao T., Fu D. Y., Ermon S., Rudra A., Ré C. (2022), FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, arXiv:2205.14135v2.

from math import exp
scores = [1000., 1001., 999.]
values = [[1., 0.], [0., 2.], [3., -1.]]
m, total, numerator = scores[0], 1., values[0][:]
for s, v in zip(scores[1:], values[1:]):
    new_m = max(m, s)
    old_scale, new_scale = exp(m-new_m), exp(s-new_m)
    total = total*old_scale + new_scale
    numerator = [u*old_scale + x*new_scale for u, x in zip(numerator, v)]
    m = new_m
print([u/total for u in numerator])

Code, data, and instructions · JSON. Educational calculations executed with Python 3.14.0; figures with Matplotlib 3.11.2. AI-assisted analysis, without claiming peer review or human review. Original illustrative ImageGen cover: it does not document EL-AI people, premises, or installations. Sources accessed 24 September 2026.