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

Wenbo Zhang; Xiang Ren

cs.CLAug 19, 2026

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

Authors: Wenbo Zhang, Xiang Ren

Organizations: University of Southern California

Abstract

When generating text, a Transformer produces representations of past tokens at every layer, but each layer can normally use only representations from the same depth. This restriction prevents the model from fully reusing information it has already computed. We introduce WhiteMatter, which allows every layer to draw on past-token representations from any depth. A learned mixer selects the most useful depths for the current context and combines their representations into shared key-value (KV) cache channels. Sharing these channels across layers can reduce the cache size. Given the same number of training tokens, WhiteMatter with a full-size cache performs comparably to a standard Transformer with 50% more layers. With half the KV cache, WhiteMatter outperforms matched standard Transformers at two model scales, up to 1.3B parameters. Cross-layer connections, however, introduce dependencies that slow training and prompt processing. We address this problem with cyclic iteration, which updates interleaved groups of tokens in turn while processing the tokens within each group in parallel. On a reference model trained with exact autoregressive execution, cyclic iteration converges 12.5x faster than standard Jacobi iteration.

Figures & tables

Panel of Figure 1 — Figure 1: KV production and consumption across layers. (a) Vanilla: each layer reads KV produced at the same depth. (b) Feedback Transformer: all layers read KV produced from the same weighted sum of source depths. (c) LCKV: warmup layers retain same-depth KV, while condensed layers read KV produced from the top layer in the block. (d) WhiteMatter: a router forms $k$ channels from all source depths, and each layer reads one of the channels.

Panel of Figure 2 — Figure 2: The cross-layer KV pool for one token position $i$ . Dashed dividers separate the three steps of § 3.1 . In Step 1 a data-dependent router mixes the $L$ per-layer states into $k$ shared channels. In Step 2 the resulting channels undergo KV projection; K normalization and RoPE are then applied to the keys before cache storage. In Step 3 each query-side layer reads one channel; the dashed arrow marks the cache boundary, as the stored channels are read while processing a later token. The key and value branches are processed independently.

Panel of Figure 3 — Figure 3: Schedules for resolving the feedback connections. Rows are computation steps and columns are tokens; each cell is shaded according to when its KV source was last updated. (a) Autoregressive: Feedback Transformer processes each token sequentially. (b) Jacobi: LCKV processes all tokens in parallel in multiple passes. (c) Cyclic: strided groups run in order, so later groups read earlier groups’ updates within a pass.

Model	Wiki.	LMB.	LMB.	BLiMP	PIQA	Hella.	ARC-E	ARC-C	Wino.	OBQA	SciQ	ReC.	SQD	Avg.
	PPL $\downarrow$	PPL $\downarrow$	acc. $\uparrow$	acc. $\uparrow$	acc. $\uparrow$	acc. n $\uparrow$	acc. $\uparrow$	acc. n $\uparrow$	acc. $\uparrow$	acc. n $\uparrow$	acc. $\uparrow$	EM $\uparrow$	acc. $\uparrow$	acc. $\uparrow$
Vanilla 16L	49.34	127.47	24.43	77.23	62.73	31.67	52.15	25.00	50.04	29.00	78.50	58.50	30.03	47.21
FusedKV	48.35	92.33	27.13	76.64	63.71	32.25	50.76	25.00	52.01	30.20	80.30	58.96	34.82	48.34
LCKV $w{=}4$	48.81	107.52	26.33	76.17	62.73	32.52	51.43	23.72	50.91	31.20	76.60	58.04	45.48	48.65
LCKV $w{=}7$	49.02	102.97	25.87	75.73	61.97	32.40	52.44	25.94	52.41	30.00	76.50	58.34	44.00	48.69
WhiteMatter $k{=}8$	44.40	71.58	29.65	78.83	63.22	33.61	52.65	25.26	51.38	29.60	80.50	60.24	36.70	49.24

Table 1: Zero-shot evaluation at

D{=}512

with

8

B training tokens. Both full- and half-cache WhiteMatter models achieved lower perplexity and higher average accuracy than all other models, including the 24L vanilla model, which has

50\%

more layers.

Figure 4: Language-modeling quality versus model size after

8

B training tokens. Lower perplexity is better; labels indicate KV-cache size relative to vanilla 16L. Full-cache WhiteMatter performs comparably to the deeper vanilla model, while half-cache WhiteMatter improves over same-depth vanilla.

Figure 7: (a) Training schedules. Training with more refinement generally improves perplexity and reduces degradation under further iteration, although approaching the best quality can require more inference passes. Row groups vary gradient-carrying passes, inner rows vary no-gradient passes, and columns compare Jacobi with cyclic schedules (

C_{g}

g

groups). (b) Channel count and ablations. More KV channels generally improve quality, but increasing cache capacity alone does not recover the benefit of specialized source mixtures. Content-dependent routing and deep-to-shallow connections also improve performance.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 8 — Figure 8: Whole-brain white-matter tractography. A population-averaged human structural connectome reconstructed from diffusion MRI, rendered as fiber tracts with DSI Studio ( Yeh, 2025 ) . The population-averaged template ( Yeh et al., 2018 ) was built from Human Connectome Project data ( Essen et al., 2013 ) .

Setting	$D{=}512$	$D{=}1792$
Intermediate width	$1536$	$5376$
Query / KV heads	$6$ / $3$	$14$ / $7$
Head dimension	$96$	$128$
Optimizer steps	$30{,}518$	$152{,}588$

Appendix

Table 3: Additional pretraining settings.

Model	Decode	Prefill	Training
Vanilla 16L	0.179	0.142	0.444
FusedKV	0.203	0.071	0.429
LCKV $w{=}4$	0.173	0.935	1.421
LCKV $w{=}7$	0.175	0.738	1.174
WhiteMatter $k{=}8$	0.177	0.432	1.028
WhiteMatter $k{=}16$	0.184	0.467	1.111

Appendix

Table 4: Per-token FLOPs for training, prefill, and decoding at

D{=}512

, in GFLOPs per token.

Figure 9: FP32 prefill convergence on the autoregressively trained reference model. Both partitions use

g\in\{2,4,8,16,32,64\}

groups. Curves show all evaluated passes. Horizontal lines mark exact autoregressive perplexity and the

1\%

threshold used for pass selection.

Figure 10: Iteration-schedule runtime for a larger cyclic-trained model. Cyclic iteration achieves the lowest runtime at the selected pass counts.

Explore similar work

Depth-Attention: Cross-Layer Value Mixing for Language Models
Jun 3, 2026Boyi Zeng, Yiqin Hao, Zitong Wang +7Layer-Wise Cross-Layer Interactions
LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University · Sun Yat-sen University · Shanghai Innovation Institute +2
You Do Not Fully Utilize Transformer's Representation Capacity
Feb 13, 2025Gleb Gerasimov, Yaroslav Aksenov, Nikita Balagansky +2Transformer Architectures Language Modeling
1T-Tech · 2Moscow Institute of Physics and Technology · 3HSE University
The Recurrent Transformer: Greater Effective Depth and Efficient Decoding
Apr 23, 2026Costin-Andrei Oncescu, Depen Morwani, Samy Jelassi +3Transformer Architectures Kv-Cache Management
Harvard University

Jun 3, 2026, cs.CLJ/K move · Enter open · S save

Depth-Attention: Cross-Layer Value Mixing for Language Models

Boyi Zeng, Yiqin Hao, Zitong Wang, Shixiang Song+6

LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University · Sun Yat-sen University · Shanghai Innovation Institute+2

Layer-Wise Cross-Layer Interactions Key-Value Cache Value

Open paper AskPDF arXiv

Self-attention selects information freely across the sequence, but across depth, Transformers merely add each layer's output to the residual stream, so later layers cannot selectively reuse earlier-layer representations. Recent cross-layer methods improve this flow but operate on hidden states outside attention, adding state beyond the key-value cache at inference--a cost that becomes increasingly salient as modern LLMs compress the cache with grouped-query and multi-head latent attention. We introduce Depth-Attention, which performs this selection inside the attention module itself: before a layer attends over the sequence, its query attends over the keys of earlier layers at the same token position and mixes their values into the value that self-attention then reads. Because Depth-Attention reuses the standard attention queries, keys, and value-cache slots, storing depth-mixed values in place of the original values, it adds no parameters and introduces no persistent inference state beyond the standard key-value cache--the same cache size as a vanilla decoder and less than hidden-state-based cross-layer methods. On Qwen3-style decoders at 1.5B and 3B parameters, Depth-Attention attains the lowest perplexity and the highest average downstream accuracy, improving over the vanilla Transformer by up to 2.3 accuracy points and surpassing strong cross-layer baselines in perplexity and average accuracy, while adding under 0.01% extra arithmetic FLOPs and no additional persistent inference state. The gains hold from 360M to 3B parameters and extend to looped Transformers.

Abstract

Explore similar work

Depth-Attention: Cross-Layer Value Mixing for Language Models

You Do Not Fully Utilize Transformer's Representation Capacity

The Recurrent Transformer: Greater Effective Depth and Efficient Decoding