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

Abstract: 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.

• The paper introduces WhiteMatter, a Transformer architecture that dynamically mixes key–value states across all source layers for each consumer layer, enabling deep-to-shallow feedback and content-dependent connectivity.
• WhiteMatter achieves 19.968 perplexity with a full cache—an 8.2% improvement over an equal-depth vanilla model—and its half-cache version reaches 20.377 perplexity while reducing cache size by 50%.
• The cyclic Gauss–Seidel schedule makes fixed-point prefill practical, but training and prefill require roughly 2.3–2.5× and 3.1–3.3× vanilla FLOPs respectively, while decoding compute remains near parity.

Motivation and architectural gap

In a standard autoregressive Transformer, each attention layer can only read key–value (KV) entries produced at its own depth, even though the model has already computed a full stack of hidden states for every past token. The authors argue that this restriction wastes information the model has already produced and limits effective computational depth. Prior work relaxes it only partially: feedback architectures such as the Feedback Transformer and LCKV expose deeper past-token states to shallow layers but use a single connection pattern shared across all consumer layers (Wu et al., 2024), while feedforward cross-layer methods (DenseFormer, MUDDFormer, Hyper-Connections, FusedKV) provide consumer-specific connectivity within or across earlier layers of past tokens but never deep-to-shallow access to past representations (Lin et al., 3 Dec 2025). WhiteMatter combines the two axes: every consumer layer receives connections to all $L$ source depths of each past token, with weights that vary per consumer layer and adapt to the source token's content.

The design is motivated by cortical connectivity, where white-matter fibers link distant areas with area-specific, dynamically modulated pathways; the paper maps this onto four properties—direct long-range connections, deep-to-shallow feedback, consumer-specific connectivity, and dynamic modulation.

Method

WhiteMatter retains standard decoder blocks but replaces the $L$ per-layer KV projections with a cross-layer KV pool of $k \le L$ shared channels. At each token position, a linear router reads RMS-normalized hidden states entering each layer (optionally subsampled every $p$th layer) and produces signed mixing weights $$\alpha^K[i], \alpha^V[i] \in \mathbb{R}^{k\times L}$$ that combine all source states into $k$ channels. Shared projection pairs convert channels into keys and values, with per-channel QK-norm and RoPE applied before caching. Each consumer layer reads exactly one channel via a fixed cyclic assignment ($\ell \bmod k$), preserving one KV read per layer rather than streaming all channels from HBM. Because the cache stores $k$ instead of $L$ channel pairs, the KV-cache footprint scales as $k/L$ relative to vanilla.

Two implementation details matter for correctness and efficiency. First, since a token's channels are computed only after its full forward pass, its own queries must not attend to them; masking the diagonal would break FlashAttention-2 compatibility, so a learned dummy token is prepended to offset the cache by one position. Second, parallel training and prefill face a circular dependency between a token's hidden states and earlier tokens' KV. The paper formulates execution as a fixed-point problem $L$0, $L$1 and resolves it with a cyclic Gauss–Seidel schedule: tokens are partitioned into $L$2 strided groups evaluated in order, so later groups read current-pass updates from earlier groups while retaining intra-group parallelism. This interpolates between Jacobi iteration ($L$3) and exact autoregressive evaluation ($L$4); gradients are carried only through the last passes (truncated backpropagation).

Main results

All models were pretrained from scratch on 8B tokens of FineWeb-Edu using the Qwen3 decoder at width 512, trained with Muon plus AdamW on eight RTX A6000 GPUs.

Model	Held-out PPL	Non-emb. params	Cache vs. vanilla
Vanilla 16L	21.747	51.9M	1.00×
Vanilla 24L	20.181	—	1.00×
LCKV $L$5	21.461	49.7M	0.50×
WhiteMatter $L$6	20.377	50.6M	0.50×
WhiteMatter $L$7	19.968	54.1M	1.00×

Full-cache WhiteMatter reduces perplexity by 8.2% over the same-depth vanilla baseline and outperforms a vanilla model with 50% more layers. Half-cache WhiteMatter ($L$8) retains most of the gain (6.3% reduction) and is 5.0% below equal-cache LCKV $L$9. Downstream, both WhiteMatter variants beat the 32-layer vanilla model on LAMBADA perplexity (60.73 and 71.58 vs. 79.39), and full-cache WhiteMatter leads all 16-layer models on LAMBADA, WikiText, PIQA, and HellaSwag. Decoding FLOPs are essentially unchanged from vanilla (~1.0×), so these gains come without decode-time compute overhead.

Prefill convergence and cost

Using a controlled 4-layer model trained with exact autoregressive execution, cyclic Gauss–Seidel with $k \le L$0 reaches within 1% of autoregressive perplexity in 4 passes, achieving 0.01245 s/sequence versus 0.1393 s for Jacobi (75 passes) and 0.1729 s for exact autoregressive evaluation—13.9× faster than exact prefill and 11.2× faster than Jacobi. A larger 8-layer cyclic-trained model shows similar behavior (5 passes, 15.5× faster than autoregressive rollout). One anomalous finding stands out: Jacobi requires more than twice as many passes as cyclic $k \le L$1 despite performing an equivalent number of sequential updates, and its perplexity oscillates across iterations; the authors state they have not identified the cause.

The costs are asymmetric. Training requires roughly 2.3–2.5× vanilla FLOPs and three-pass prefill about 3.1–3.3×, whereas decoding remains near parity. The method therefore targets inference-time quality and memory efficiency at the expense of training and prefill compute.

Analysis

Three ablations isolate the design's components. First, a schedule sweep over gradient-carrying passes, no-gradient passes, and group counts shows up to 32% perplexity spread between the strongest and weakest schedules, with diminishing returns toward the fixed point; models trained farther from the fixed point degrade when iterated beyond their training schedule, including under autoregressive decoding, while converged-schedule models are stable but need more inference iterations. Second, scaling $k \le L$2 from 1 to 16 yields monotone improvement with diminishing returns; notably, even $k \le L$3 beats the vanilla baseline with a 16× KV-cache compression and 7.3% perplexity reduction. Third, removing deep-to-shallow feedback (mixing only over layers $k \le L$4) leaves the model 7.5% worse than full-cache WhiteMatter and worse than the $k \le L$5 model despite a 16× larger cache, establishing feedback—not merely dynamic mixing—as the dominant contributor. Static learnable mixing weights (no router) cost about 2% perplexity, confirming the value of content-dependent routing.

Limitations and open questions

The authors are explicit about scope. All main results use small models (~50M non-embedding parameters) at an 8B-token budget, so neither the quality gains nor the systems trade-offs are established at scale. No optimized end-to-end decoding benchmark is provided—only FLOP counts and cache sizes—and the unexplained Jacobi oscillation phenomenon remains open. The dependence on truncated backpropagation and short iteration schedules raises the question of whether models trained far from the fixed point would behave differently under longer training. More efficient fixed-point solvers or a dedicated prefill encoder are suggested but not explored.

Conclusion

WhiteMatter demonstrates that exposing all past-token layer states to every consumer layer through content-dependent KV mixing yields substantial language-modeling gains at fixed depth—matching or exceeding models with 50–100% more layers—while the channel count $k \le L$6 provides a direct knob for KV-cache compression. The cyclic Gauss–Seidel schedule makes iterative training and prefill practical, though at 2.3–3.3× vanilla FLOPs. The central open question is whether these trade-offs hold at realistic model and data scales.