Royer Research Labs

Research note · small-scale screen

Dilution in a Kimi-style hybrid

Kimi Linear (Moonshot AI) does most of its work in linear-time Kimi Delta Attention (KDA) layers and keeps full attention in one layer in four. We kept that recipe — the same layers in the same order — and changed only the attention layers, to dilution attention. Loss is unchanged. Retrieval past the training length becomes high and consistent across seeds, where the reference's changes sharply from one training run to the next.

3.643

loss, both versions

mean of 2 seeds each

0.97–1.00

retrieval at 4× training length

reference: 0.31–0.84

0.79–0.94

retrieval at 8× training length

reference: 0.27–0.61

¼

the key-value cache of all-attention

6.1 vs 24 KB per token

Every model is a 12-layer, 64-million-parameter decoder trained on 600 million FineWeb-Edu tokens at a 2K context, parameter-matched, two seeds each. This is a screen, not a final result: the same designs are now training at 209 million parameters and a 16K context.

Retrieval past the training length

A keyword is planted at varying depths in haystacks from 2K to 16K tokens. Everything right of the vertical rule is longer than any sequence the models saw in training. Each row is one trained model.

Retrieval past the training lengthKeyword retrieval by haystack length for 64M models trained at 2K. Swapping the attention layers of a Kimi-style hybrid for dilution keeps retrieval at 0.97 to 1.00 at four times the training length and 0.79 to 0.94 at eight times, against 0.31 to 0.84 and 0.27 to 0.61 for the reference, at the same loss. Twelve layers of softmax with RoPE fall to chance past the training length.Same Kimi-style hybrid, attention layers swapped for dilutionkeyword retrieval by haystack length · 64M parameters, trained at 2K · chance is 0.14past the training length · no fine-tuning2Kin context4K2×8K4×16K8×losslower is betterKimi recipe, attention swapped for dilutionseed 13371.001.000.970.793.640seed 23571.001.001.000.943.645Kimi recipe, reference attention (no position encoding)seed 13370.990.770.310.273.639seed 23571.000.970.840.613.64712 layers of softmax attention + RoPEseed 13370.640.140.140.143.682seed 23570.680.160.140.143.684
Likelihood-scored retrieval: eight candidate keywords, one disqualified by a control prior, so chance is 1/7 ≈ 0.14; the 2K column is 2,040 tokens, leaving room for the beginning-of-sequence token. Loss is on a common 2.1M-token FineWeb-Edu validation slice. The softmax model uses RoPE unscaled.

What changes, and what it costs

Same layers, same order, same parameter count; only the attention in layers 4, 8 and 12 differs. Where two values are shown they are the two seeds.

Reference attention against dilution in the Kimi recipeSame layers, same order, two seeds each. Loss is tied at 3.643. Dilution is ahead on retrieval at four and eight times the training length, on all-keyword retrieval and on using distant book text. The reference is ahead on training speed by 4 percent and on 64K prompt processing. Cache size is equal.Kimi-style hybrid: reference attention vs dilution · 64M, trained at 2KReferenceDilutionEdgeValidation loss, mean of 2 seeds3.6433.643tieRetrieval at 4× training length0.31 · 0.840.97 · 1.00dilutionRetrieval at 8× training length0.27 · 0.610.79 · 0.94dilutionAll-keyword retrieval at 4×0.19 · 0.660.94 · 0.97dilutionBook text 8–16K back (− helps)+0.047 · −0.003−0.057 · −0.054dilutionTraining speed at 2K241K tok/s231K tok/sreference64K-token prompt processing428K tok/s236K tok/sreferenceKey-value cache per token6.0 KB6.1 KBtie
Book text: loss on PG-19 books with the full preceding text minus loss with only the last 2K tokens, for tokens 8–16K into the book; negative means the distant text helped. All-keyword retrieval asks for every planted keyword rather than one. Speeds: batch 1 prompt processing and training throughput on one GPU, dilution on its fused kernel.

The gain is less that dilution retrieves more than that it retrieves reliably: two otherwise identical training runs of the reference reach 0.31 and 0.84 at four times the training length, and both dilution runs reach 0.97 or better. On long books, distant text helps the dilution version at every range we measured and does nothing for the reference past about 8K tokens.

The cost is speed. Dilution trains about 4% slower at this size and, in our current implementation, processes a 64K-token prompt at roughly half the reference's rate. Both hybrids are about 0.04 lower in loss than twelve layers of softmax attention with RoPE, which keeps four times the cache and retrieves at chance past its training length.

What was tested

Twelve layers, width 512, eight heads, SwiGLU feed-forward layers and LayerNorm, trained with a beginning-of-sequence token on every window. KDA layers at positions 1–3, 5–7 and 9–11, attention at 4, 8 and 12, as in Kimi Linear's three-to-one pattern. The reference attention is causal softmax with no position encoding, as in Kimi Linear's attention layers; the swap uses the dilution operator from our DilutionAttention release with its fused kernel. KDA uses the open-source flash-linear-attention implementation. Parameter counts are matched to within 0.02% through the feed-forward width.

What this does not establish

This is a 64-million-parameter screen trained at a 2K context, with two seeds per model. Our reference follows Kimi Linear's published layer recipe but uses standard multi-head attention rather than its compressed (MLA) attention and has no mixture of experts, so it is not the shipped model and this is not a comparison with Kimi Linear itself. Retrieval here is the standard synthetic test; harder retrieval and the use of distant natural text are weak for every model at this scale. KDA and Kimi Linear are Moonshot AI's work (arXiv 2510.26692). This note reports unreleased experiments; the 209M, 16K results will be compared against our published dilution, Selective Attention and softmax runs.

DilutionAttention results → DilutionAttention on GitHub → Kimi Linear (arXiv) → Back to news →