DilutionAttention v0.1.0 · results

Attention that keeps working past its training context

Dilution attention divides each query's softmax bid on a key by the demand that key has already absorbed from earlier queries. Trained at 16K tokens with no position encoding, 209-million-parameter models built on it retrieve planted keywords out to 128K tokens — eight times their training context — in both seeds, with no scaling rule and no fine-tuning. Softmax trained the same way retrieves 0.17–0.19 at its own context and nothing beyond it.

8×

training context, retrieval held

0.91–1.00 to 128K, both seeds

0.17–0.19

softmax at its own 16K context

same recipe, no position encoding

learned parameters in the operator

same O(T²) class as softmax

2.2×

softmax's training time

the cost, measured at 16K

The headline comparison is between 24-layer, 209-million-parameter decoders trained at a 16K context on 4.19 billion FineWeb-Edu tokens (a Chinchilla-matched budget), two seeds per arm. Every arm shares the architecture, data, schedule and learning rate; only the attention operator, its position encoding and the seed differ.

The operator

Standard attention lets every query spend a fixed budget across the keys behind it, and nothing records how much of a key other queries have already used. Dilution keeps the softmax bids and adds a price: before a bid is honoured it is divided by the demand its key has already received from earlier queries, so a key that many queries have claimed yields a smaller share. Each query's shares are then renormalised to sum to one.

p         = causal_softmax(q k^T / sqrt(d))     # bids
committed = exclusive_cumsum_cols(p)            # demand from earlier queries
share     = p / (committed + p + eps)           # each bid against cumulative demand
A         = share / rowsum(share);   out = A @ v

One softmax, one cumulative sum and one renormalisation: the same O(T²) class as standard attention, with nothing learned inside the operator. Cached decoding needs the ordinary key/value cache plus one running demand total per key and head.

The closest precedent is temporal attention (Sankaran et al., 2016), which divides each raw attention score by the scores earlier decoder steps gave the same source position, then renormalises. Dilution keeps that division but changes what is accumulated and what it is divided by. It accumulates softmax-normalised bids, so every query adds exactly one unit of demand in total, whatever the magnitude of its raw scores, and every share stays at or below one. And it counts the current bid in the demand, which keeps the operator defined where a key first becomes visible.

The second change is load-bearing in causal self-attention. Dividing by earlier demand alone gives zero at each key's first query — the key itself — so that entry wins its row outright; a model trained that way kept its first layer as an identity map and lost 0.072 nats (one seed each). Whether normalised bids beat raw exponentials in practice is untested: our raw-exponential runs overflowed. The contribution here is this form as a language model's self-attention, the controlled study around it, and exact fused kernels.

Retrieval past the trained context

The models were asked to retrieve a keyword planted at varying depths in haystacks from 3.8K to 256K tokens. Everything to the right of the vertical rule is longer than any sequence the model saw in training. Each row is one trained model; the seeds are shown separately because how far a model reaches varies between seeds of the same recipe.

Likelihood-scored retrieval: eight candidate keywords, one disqualified by a control prior, so chance is 1/7 ≈ 0.14. In-context cells use 20 depth placements and cells past the context use 10. Models with RoPE layers use dynamic NTK scaling past the trained context; models without position encoding run unmodified. – not evaluated. † scored with the earlier version of the protocol: the checkpoint was trained further afterwards and could not be re-scored.

Long documents

Keyword retrieval is synthetic, so the same models were also scored on real long text: 29 books from PG-19, out-of-domain for a model trained on educational web pages, read to 128K tokens. Loss is averaged over each range of positions.

Past the trained context, dilution's loss stays flat to eight times its training length, while softmax with RoPE and NTK scaling ends 1.7–1.9 nats worse. Inside the context the ordering reverses on these out-of-domain books: dilution trails softmax with RoPE by 0.05–0.08 nats. Dilution's other 1× seed was trained on to twice the tokens; at that budget both dilution seeds stay flat to 128K (3.33–3.41 nats) while both softmax-without-position-encoding seeds climb to 7.3–8.5.

Generated text

Loss is measured on text the model is given. To check that the models also write, each continued the same 256 validation prompts, and a larger model (SmolLM2-1.7B) scored how natural the continuations were.

Nucleus sampling (top-p 0.9), 256 prompts per cell; the human continuation scores 0.65–0.70. Inside the trained context the models are level. At twice the context, softmax with RoPE writes incoherent text while dilution stays within 0.03 of its in-context score. The judge reads only the last 4,096 tokens of each prompt, so this shows dilution remains a working language model past its window, not that it uses information from 32K tokens back. With no position encoding at twice the training budget, dilution scores 1.00 / 1.02 at 2× past the context, softmax 3.33 / 2.31.

What it costs

Inside the trained context the gain is small, and dilution is slower to train. The alternating hybrid, which gives each operator its own best position mode, has the best in-domain loss and recovers part of the speed.

Validation loss on one common slice of held-out FineWeb-Edu, seeds 1337 / 2357. *This softmax seed hit gradient spikes early in training and never recovered. Throughput is the logged training rate on one RTX PRO 6000 Blackwell with the released fused kernels: at 16K, dilution takes about 2.2× softmax's training time. The in-domain margin depends on context length: softmax with RoPE is ahead by 0.036 nats at a 1K context, and dilution ahead by 0.028 at 16K.

What this does not establish

These are controlled results at up to 209 million parameters. They do not establish that the behaviour holds at substantially larger scale, and the retrieval test is synthetic: likelihood-scored keyword retrieval, not a downstream long-context task.

How far retrieval extends varies between seeds, and loss does not show it. Both seeds above hold to eight times their context, but trained on to twice the tokens, one seed reaches only about twice its context while the other stays near-flat, at nearly identical loss (2.857 and 2.854). Across the settings we ran, the range runs from 2× to 8×, and we report it as a range rather than a single number.

Why dilution extrapolates is not established. A plausible reason is that the running demand total carries no distance scale, but that has not been tested causally. Inside the trained context, dilution is not a better language model than softmax with RoPE in general: in-domain it leads by 0.028 nats at 16K and trails by 0.036 at 1K, and on out-of-domain books it trails by 0.05–0.24 nats.

Everything behind these charts

Every number on this page is backed by a raw record. The release ships the reference implementation and fused Triton kernels with exact gradients for training, prefill and cached decoding; the training and evaluation code; the configuration, metrics and data provenance of the 101 runs behind the results; every result file behind these charts, including per-sample generation metrics from which a script recomputes every published statistic; an evidence ledger that keeps withdrawn claims; the weights of all 16 models; and a technical report. Code is Apache-2.0 licensed; documentation and results are CC BY 4.0.

Code & study on GitHub → Technical report (DOI) → Model weights → Release v0.1.0 → Archived software (DOI) → Back to news →