WhiteMatter: All-to-All Cross-Layer Connections via KV Mixing

arXiv:2608.184862026-08-20

Letting every Transformer layer read every other layer's memory, like white-matter wiring in the brain, boosts language-model quality without adding depth

Normally each layer in a Transformer can only reuse the information it produced itself when looking back at past tokens, even though deeper layers hold richer information from the same token. WhiteMatter adds a small router network that mixes the hidden states of all layers at each token into a handful of shared channels, and lets each layer pick one channel to read, so shallow layers can access what deeper layers computed. Trained under matched conditions, WhiteMatter beat a vanilla Transformer with 50% more layers and, even after halving its cache size, kept most of that improvement.

METAL MEDIA explanatory visual

Letting every Transformer layer read every other layer's memory, like white-matter wiring in the brain, boosts language-model quality without adding depth

01Problem: when a Transformer looks back at a past token, each layer can only use the information it produced at its own depth for that token, wasting the deeper representations the model already computed.
02Method: at every token position, a router reads the hidden states from all layers and mixes them into k shared channels; each layer follows a fixed rule to read one of those channels for attention. Using fewer channels than layers (k less than L) shrinks the stored cache.
03Training challenge: building a token's channels requires that token's full layer computation to finish, but that computation depends on earlier tokens' channels, creating a circular dependency. The authors solve this with a cyclic, grouped iteration schedule that keeps computation parallel across tokens while converging to the right answer over a few passes.
04Results: trained on 8 billion tokens, the 16-layer WhiteMatter model lowered held-out perplexity (a measure of prediction quality, lower is better) by 8.2% compared with a same-depth vanilla Transformer, and even beat a 24-layer vanilla model. Halving the cache still kept a 6.3% improvement and beat an equal-cache baseline (LCKV) by 5.0%.
05Speed: the cyclic iteration schedule converged up to 13.9x faster than exact sequential (autoregressive) computation and 11.2x faster than the simplest parallel iteration approach.

An explanatory diagram made by METAL MEDIA, not a figure supplied by the paper's authors.

What they did

Problem: when a Transformer looks back at a past token, each layer can only use the information it produced at its own depth for that token, wasting the deeper representations the model already computed.
Method: at every token position, a router reads the hidden states from all layers and mixes them into k shared channels; each layer follows a fixed rule to read one of those channels for attention. Using fewer channels than layers (k less than L) shrinks the stored cache.
Training challenge: building a token's channels requires that token's full layer computation to finish, but that computation depends on earlier tokens' channels, creating a circular dependency. The authors solve this with a cyclic, grouped iteration schedule that keeps computation parallel across tokens while converging to the right answer over a few passes.
Results: trained on 8 billion tokens, the 16-layer WhiteMatter model lowered held-out perplexity (a measure of prediction quality, lower is better) by 8.2% compared with a same-depth vanilla Transformer, and even beat a 24-layer vanilla model. Halving the cache still kept a 6.3% improvement and beat an equal-cache baseline (LCKV) by 5.0%.
Speed: the cyclic iteration schedule converged up to 13.9x faster than exact sequential (autoregressive) computation and 11.2x faster than the simplest parallel iteration approach.

Figure 4: Held-out language-modeling quality versus non-embedding parameter count at an 8B-token budget. The connected vanilla points form the depth-scaling reference; point labels report per-token KV-cache size relative to the L=16 vanilla model. WhiteMatter is shown in half- and full-cache configurations; the LCKV configurations have four and seven warmup layers.

Table 1: Downstream evaluation. LAMBADA and WikiText report perplexity; the remaining columns report normalized accuracy in percent. Bold denotes the best result among the 16-layer models.
Model	LAMBADA ↓	WikiText ↓	PIQA ↑	HellaSwag ↑	ARC-E ↑	OBQA ↑
Vanilla 16L	127.47	49.34	60.88	31.67	47.39	29.00
LCKV w=4	107.52	48.81	62.57	32.52	45.66	31.20
LCKV w=7	102.97	49.02	62.24	32.40	46.21	30.00
WhiteMatter k=8	71.58	44.40	62.35	33.61	45.71	29.60
WhiteMatter k=16	60.73	43.28	63.55	33.80	46.21	29.40
Vanilla 24L	97.40	44.71	62.73	33.21	47.94	31.80
Vanilla 32L	79.39	41.44	63.82	34.35	47.90	32.20

Figure 5: Prefill convergence wall time versus group count g. Jacobi (g=1) and autoregressive evaluation (g=T) form the two endpoints. The 4-layer model was trained with exact autoregressive execution at length 1024; evaluation used T=2048 and the same channel-read policy for pass selection and timing.

Table 2: Measured per-token FLOPs for training, prefill, and decoding. Each pair of columns reports GFLOPs per token and the ratio to the 16-layer vanilla model.
Model	Training	Prefill	Decode
	GFLOP/tok	×	GFLOP/tok	×	GFLOP/tok	×
Vanilla 16L	0.444	1.00	0.142	1.00	0.179	1.00
Vanilla 24L	0.665	1.50	0.212	1.50	0.269	1.50
Vanilla 32L	0.887	2.00	0.283	2.00	0.359	2.00
LCKV w=4	1.421	3.20	0.935	6.61	0.173	0.97
LCKV w=7	1.174	2.65	0.738	5.21	0.175	0.97
WhiteMatter k=8	1.028	2.32	0.432	3.05	0.177	0.99
WhiteMatter k=16	1.111	2.50	0.467	3.30	0.184	1.03

Figure 6: Training-schedule and pool-rank ablations. (a) Model performance across training iteration schedules. Each cell reports one combination of iteration parameters, averaged over two seeds. The vertical axis nests the number of no-gradient passes nno-grad (inner labels) within the number of gradient-carrying passes ng (outer labels); the horizontal axis shows the training iteration schedule. (b) Test perplexity versus pool rank for WhiteMatter and two ablations.

Table 3: Complete zero-shot downstream results. LAMBADA and WikiText report perplexity; all other columns report accuracy in percent.
Model	LAMBADA	WikiText	PIQA	Wino.	BoolQ	Hella.	ARC-E	ARC-C	OBQA
Vanilla 16L	127.47	49.34	60.88	50.04	48.17	31.67	47.39	25.00	29.00
LCKV w=4	107.52	48.81	62.57	50.91	48.93	32.52	45.66	23.72	31.20
LCKV w=7	102.97	49.02	62.24	52.41	60.64	32.40	46.21	25.94	30.00
WhiteMatter k=8	71.58	44.40	62.35	51.38	60.40	33.61	45.71	25.26	29.60
WhiteMatter k=16	60.73	43.28	63.55	51.38	61.04	33.80	46.21	25.43	29.40
Vanilla 24L	97.40	44.71	62.73	50.99	56.67	33.21	47.94	26.37	31.80
Vanilla 32L	79.39	41.44	63.82	50.51	56.18	34.35	47.90	26.96	32.20

Figure 7: Whole-brain white-matter tractography. A population-averaged human structural connectome reconstructed from diffusion MRI, rendered as fiber tracts in sagittal (left) and coronal (right) views. The tracts span the brain and arc between distant regions in every direction; color encodes local fiber orientation (red: left–right, green: anterior–posterior, blue: superior–inferior). Rendered with DSI Studio (33) from its population-averaged human template (32), built from Human Connectome Project data (5).

Why it matters

The size of the cache a model keeps in memory and the compute it uses during generation directly determine serving cost and speed for real chatbots and assistants. This work shows a way to improve model quality without adding layers while also shrinking the cache footprint, offering a practical lever for running capable models more cheaply.

Figure 8: Convergence timing for the larger cyclic-trained model. This 8-layer, D=1024, k=8 model is evaluated at T=4096.

Terms in this paper

perplexity · a score measuring how well a language model predicts the next word; lower means better prediction
KV cache · the stored keys and values from past tokens that a Transformer reuses so it doesn't have to recompute them
router · a small neural network that decides how to combine information from multiple layers
autoregressive decoding · generating text one token at a time, each new token conditioned on all previously generated tokens
Jacobi / Gauss-Seidel iteration · numerical methods for solving circularly dependent equations by repeated approximation passes; the order and grouping of updates trades off parallelism against convergence speed

Original abstract (English)

In a Transformer, each layer attends to past tokens only through KV produced at its own depth, despite the presence of deeper representations during autoregressive decoding. Feedback architectures allow shallow consumer layers to attend to KV produced by deeper past-token representations, but give all consumer layers the same fixed connection patterns to source layers. We propose WhiteMatter, which connects every attention layer to the representations from all layers of each past token, with connection weights that can vary across consumer layers and adapt to the source token. For each token, a router implements these connections by mixing its $L$ layer states into $k$ KV channels that are cached for subsequent tokens; each consumer layer attends to one of the channels. The number of channels $k$ controls the KV-cache size. Setting $k<L$ reduces the cache's memory footprint. In our pretraining experiments, WhiteMatter outperforms a vanilla Transformer with 50% more layers and retains most of this gain with a 50% KV-cache compression.

Authors · Wenbo Zhang, Xiang Ren

Read on arXiv

Latest papers

All papers →

Latest from METAL MEDIA

Figures: Wenbo Zhang et al., arXiv:2608.18486, CC BY 4.0