Selective Attention, trained with no position encoding at 16K

While reviewing related work for DilutionAttention we found that Selective Attention has the same shape: both subtract from each attention logit a term built from what earlier queries did with the same key. In small experiments at a 2K context it behaved much like dilution with no position signal at all, which made it worth testing at the scale where the dilution results were measured.

The paper's listing builds the full attention matrix, which does not fit a 24-layer model at 16K tokens. We wrote fused Triton kernels — forward and full backward, tested against the paper's listing — that compute the same operator without it. With those it trains at about dilution's speed on the unchanged DilutionAttention harness, so its results sit directly beside the published runs.

Dilution attention has the same form, softmax(L − F), with a different history F: the log of the softmax demand each key has accumulated, per head. The two differ in what is accumulated, how it is subtracted and whose history is used, and those differences show up in the long-document results below.

Past the training length natural text costs Selective Attention retrieval, with a wide spread between seeds, and it loses facts in the middle of the input first. With three distractors, every operator scores 0.24–0.46 inside 16K and most misses are a distractor's keyword: the models find the planted facts but often attach the wrong one to the key. The position-free operators hold that level to 128K; softmax with RoPE falls toward chance.

Selective Attention exempts position 0 from selection because the paper assumes it holds a beginning-of-sequence token. Windows cut from a token stream start mid-document, and without a BOS token two of our three models could not retrieve a fact placed in the very first sentence of their input (2/7 and 1/7 placements), at every length, while every other depth was retrieved. Putting a BOS token at the start of every training window brought those two seeds to 7/7 and 6/7 at no measurable cost in loss. The failure did not appear in our 2K-context experiments, so a short screen may not catch it.

A flat curve shows that a model tolerates long input, not that it uses it. Scoring the same tokens with the whole book and with only a 16K rolling window, Selective Attention gains a little from text one to two windows back and is hurt by text more than four windows back, by up to 0.29 nats at 16–32×; dilution is about neutral to 131K and hurt about a third as much. Its lower loss on these books is in-window quality. Trained on to twice the tokens, its long-book loss climbs well above dilution's, which stays flat.

Where Selective Attention is ahead: it retrieves further on the standard test (all three seeds at 0.99–1.00 to 128K), models book text better inside its window by about 0.1 nats, and its kernel is slightly faster with half the memory. Validation loss on the training distribution is a tie.

Where dilution is ahead: stability. It needs no BOS token; its long-book loss stays flat to 32× its training length, at 1× and after twice the training, where Selective Attention's rises; and distant text harms it about a third as much. Its natural-text retrieval declines more gently on the one seed we evaluated, but the selective seeds spread so widely that this is not settled.

A plausible reading, not yet tested causally: Selective Attention's history grows linearly with the number of earlier queries and is shared by every head, so its suppression keeps strengthening with length and with training; dilution's log of demand grows slowly and is kept per head.

One attention layer, forward and backward at 16K, takes 27.7 ms with the selective kernel against 30.5 ms for the dilution kernel and 9.4 ms for flash softmax attention: 2.9× softmax's time, 0.91× dilution's, with half dilution's memory. For the whole model a 16K training step runs at about 36K tokens/s against 72K for softmax with RoPE, so at equal wall clock softmax sees about twice the tokens; every comparison here is at equal tokens. Cached decoding is written in PyTorch, not as a kernel.

One model size, one corpus and one context length, with three seeds of Selective Attention and two of each reference. The retrieval tests are synthetic — one planted keyword, in repeated text or in natural text with distractors — and how far retrieval reaches past 16K varies by seed.

Retrieval scores overstate long-context ability: neither position-free operator uses distant natural text to predict better, and Selective Attention's long-document loss degrades further with more training. Why its extrapolation varies between seeds and erodes with training is not resolved; damping the history term at evaluation time made things worse. The Forgetting Transformer and stick-breaking attention also train without position encoding and were not run at this scale. This is an independent implementation, not affiliated with or endorsed by the paper's authors.