WhiteMatter: All-to-All Cross-Layer Connections via KV Source Mixing
Wenbo Zhang, Xiang Ren
cs.CL, cs.LG
2026-08-19
WhiteMatter mixes all-depth past-token states into shared KV channels. Matched-token full cache matches a 50% deeper Transformer; half cache beats matched models up to 1.3B.
A Transformer already builds a representation of every past token at every layer. Attention still reads keys and values from the same depth only. That choice keeps training and prompt prefill easy to parallelize. It also leaves unused the states the model has already paid for.
Feedback Transformer collapses those states into one weighted sum and shares the resulting KV with every layer. LCKV gives condensed layers the top-layer KV and lets warmup layers keep same-depth KV. Both designs squeeze every source into a single layer-width vector that has to serve every target. LCKV's own plots show that applying feedback KV at every layer hurts perplexity. Enlarging the cache does not close the gap: different layers appear to want different source depths.
Decoding often dominates inference for agents and long reasoning traces. The bet here is to let every layer reuse past-token states from any depth, without giving up a practical training and prefill path.
WhiteMatter replaces per-layer KV projections with a cross-layer pool. For each past token:
Attention is strictly causal: token i attends only to positions s < i. After the layer stack finishes, the pool writes that token's channels for later tokens. Reading the current token's own channels would cycle shallow and deep layers.
Prefill and training treat the cache as a fixed point. Jacobi iteration, used by LCKV, evaluates all tokens in parallel against the previous pass's cache. Updates cannot move right until the next full pass, so many passes are common. Cyclic Gauss-Seidel iteration splits positions into g strided groups (i mod g), runs groups in order, and refreshes KV after each group. Later groups in the same pass can read earlier groups' new KV, which helps because attention often lands on neighbors. Gradients flow only through the last ng passes. The kernel adapts FlashAttention-style tiling to a strided causal mask.
The main runs use g=8, router stride p=2, one no-gradient pass and two gradient-carrying passes.
Decoders follow the Qwen3 layout and train from scratch on FineWeb-Edu. The small suite is width 512 with 8B tokens. The larger suite is width 1792, about 1.3B parameters, with 10B tokens.
Held-out perplexity at the small scale (Figure 4):
| Method | Test PPL | KV vs 16L |
| Vanilla 16L | 21.75 | 1.0× |
| LCKV w=7 | 21.46 | 0.5× |
| FusedKV | 21.59 | 0.5× |
| WhiteMatter k=8 | 20.47 | 0.5× |
| WhiteMatter k=16 | 20.08 | 1.0× |
| Vanilla 24L | 20.18 | 1.5× |
Full-cache WhiteMatter is 7.7% below same-depth vanilla and roughly matches a 24-layer vanilla that has 50% more layers. Half-cache is 5.9% below same-depth vanilla. Mean zero-shot score over 11 tasks is 49.89 for k=16, against 47.21 for 16L vanilla and 48.52 for 24L. LAMBADA perplexity drops from 127.47 to 60.73.
At 1.3B, half-cache WhiteMatter (k=14, 1.326B) records held-out PPL 12.93 versus 13.51 for 28-layer vanilla (1.351B), a 4.3% relative cut. Mean zero-shot is 59.24 vs 57.48. This scale has no full-cache run and no LCKV or FusedKV baseline.
On a 4-layer model trained with exact autoregressive execution, cyclic g=16 hits within 1% of reference PPL in 4 passes at 7.32 ms/sequence. Jacobi needs 53 passes and 91.20 ms, a 12.5× wall-clock gap. At 1.3B, batch 64, A6000, all four architectures decode around 2,500 tokens/s. WhiteMatter peaks at 10.05 GiB versus vanilla's 16.59 GiB (39.4% lower); prefill peaks at 13.44 GiB versus 21.21 GiB. Prefill throughput is 31% of vanilla, 1.78× LCKV, and 2.92× Feedback Transformer.
Training FLOPs at width 512 are 2.32-2.50× vanilla. Three-pass prefill is 3.05-3.30×. Decode FLOPs stay within 0.99-1.03×.
Ablations on 327.7M tokens: a single shared mixture plus 16 independent KV projections still loses to WhiteMatter with k=4, despite 4× the cache. Static mixing weights raise PPL by 3.0% at k=1 and 1.9% at k=16 versus the dynamic router. Blocking deep-to-shallow feedback raises full-cache PPL by 4.1%. Training farther from the fixed point makes later iteration less stable; the strongest schedule is 32% lower PPL than the weakest.
For decode-heavy, cache-heavy workloads, half the KV and better quality is a real trade.
Decode compute and throughput stay in line with a standard model, while peak memory drops by about 40%.
None of this is free. Training and prefill cost more, and the schedule needs a strided causal attention kernel. All quality numbers come from training from scratch up to 1.3B on at most 10B tokens. There is no continued-training result on a pretrained large model. If prefill is the bottleneck (short generations, RAG), 31% of vanilla throughput can kill the design.
The architectural claim is narrower than share KV across layers. Distinct per-target mixtures and content-dependent routing both move the number. Cache size alone does not.
The paper lists three limits: iterative training and prefill remain more expensive than a standard Transformer; main quality results stop at 1.3B and 10B tokens; full-cache and the other KV-sharing baselines were only compared at the small scale.
A few more gaps sit in the numbers. The 1.3B comparison is half-cache only, so the claim of matching 50% more layers is untested at that size. Likelihood eval uses three cyclic passes for WhiteMatter, nine Jacobi passes for LCKV, and one pass for vanilla, so compute at test time is not matched. OpenBookQA barely moves at the small scale, and PIQA falls from 70.46 to 69.53 at 1.3B; the average is carried by LAMBADA, SciQ, and SQuAD. The white-matter analogy lives in an appendix and does not constrain the architecture.