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

TL;DR

WhiteMatter dynamically mixes all layer states into KV channels; 16-layer full-cache PPL reaches 19.968, 8.2% below vanilla.

cs.CL 🔴 Advanced 2026-08-19 21 views
Wenbo Zhang Xiang Ren
Transformer KV cache cross-layer connectivity feedback architectures language modeling

Key Findings

Methodology

WhiteMatter feeds every past token’s L layer states into content-dependent routers that independently produce k shared Key and Value channels. Each consumer layer reads one channel through a fixed cyclic assignment, allowing arbitrary source-depth mixtures while preserving one KV read per layer. Training and prefill are formulated as a fixed point and approximated with cyclic Gauss–Seidel iteration; autoregressive decoding directly appends the mixed channels after processing each token.

Key Results

  • After 8B-token FineWeb-Edu pretraining, 16-layer WhiteMatter with k=16 achieved 19.968 held-out perplexity, an 8.2% relative reduction from 16-layer vanilla at 21.747, and slightly outperformed 24-layer vanilla at 20.181.
  • With k=8, the KV cache was compressed by 50% and perplexity was 20.377. This was 5.0% better than equal-cache LCKV w=7 at 21.461; LAMBADA and WikiText perplexities were 71.58 and 44.40.
  • In a separate exact-autoregressive study, cyclic Gauss–Seidel with g=16 reached within 1% of reference perplexity in four passes, while converged prefill was 13.9× faster than exact autoregressive evaluation and 11.2× faster than Jacobi.

Significance

The paper addresses a structural bottleneck in Transformer decoding: shallow layers cannot access deeper representations of past tokens, even though those representations have already been computed. WhiteMatter combines cross-layer feedback, consumer-specific connectivity, and token-adaptive routing without requiring additional recurrent depth during decoding. Its results suggest that better inter-layer communication can substitute for some brute-force depth, with potential benefits for long-context inference, memory-constrained deployment, and high-throughput prefill.

Technical Contribution

The main contributions are a separate Key/Value routing mechanism, cyclic fixed-point execution, and cache-size control. RMS-normalized states are mixed with signed, token-dependent weights into k channels, then projected by channel-specific KV matrices. Fixed ℓ mod k reading preserves one KV stream per consumer layer instead of streaming all channels from HBM. Cyclic Gauss–Seidel updates interleaved token groups, while truncated backpropagation limits the cost of multiple passes. Cache size becomes k/L of a standard L-layer cache.

Novelty

Unlike Feedback Transformer’s shared static mixture, LCKV’s top-layer source, and FusedKV’s static feedforward fusion, WhiteMatter combines all source depths, consumer-specific channel access, and token-dependent source-to-channel weights. Its fundamental innovation is to implement dynamic all-depth feedback through cacheable shared KV channels rather than storing every layer state or requiring dense consumer-side reads.

Limitations

  • The evidence is limited to Qwen3-style models with D=512 and 16 layers trained on 8B tokens. It does not establish scaling behavior for billion-parameter models, substantially longer contexts, or other modalities.
  • The fixed ℓ mod k consumer policy restricts learned read-side flexibility; cyclic training is about 1.5× more expensive than vanilla, and approximate fixed-point execution can create train–serve discrepancies.
  • Only k=8 and k=16 were systematically evaluated, with routers reading every second source layer. More aggressive compression, sparse routing, and alternative initialization remain underexplored.

Future Work

Future work should study learned consumer-side routing, adaptive channel counts, sparse access, and adaptive iteration stopping. Validation on larger models, long-context and multimodal tasks is needed, together with end-to-end measurements of throughput, memory bandwidth, energy, and latency across hardware. Joint optimization with FlashAttention, KV quantization, and hardware-aware caching could determine whether the architectural gains translate into production systems.

AI Executive Summary

Autoregressive Transformers normally let each layer attend to past tokens only through Key and Value states produced at the same depth. Deeper representations are computed but become inaccessible to shallow layers. Feedback Transformer and LCKV introduce cross-token feedback, yet their source patterns are shared or fixed; FusedKV mainly offers static, feedforward fusion. The unresolved question is how to expose all past-token depths without making the cache prohibitively large.

WhiteMatter addresses this with a cross-layer KV pool. For each token, content-dependent routers independently mix the L layer states into k shared Key and Value channels. Each consumer layer reads one channel through a fixed cyclic assignment. Thus, shallow layers can use deeper historical information, while k<L reduces cache memory. Training and prefill are expressed as the fixed point P=Pool(H), H=States(X;P), and solved approximately with cyclic Gauss–Seidel iteration.

On 8B FineWeb-Edu tokens, 16-layer full-cache WhiteMatter reached 19.968 perplexity, versus 21.747 for 16-layer vanilla and 20.181 for 24-layer vanilla. With half the cache, it reached 20.377 and beat equal-cache LCKV by 5.0%. A separate controlled study found 13.9× faster converged prefill than exact autoregressive evaluation. The costs are roughly 1.5× vanilla training and limited scaling evidence, but the work demonstrates a credible alternative to simply adding layers: improve how existing layers communicate.

Deep Analysis

Background

A standard Transformer decoder stores per-layer KV, so layer ℓ at a new position reads only layer-ℓ KV from earlier positions. Feedback Transformer adds static shared cross-depth mixtures; LCKV reuses top-layer KV and uses Jacobi iteration; DenseFormer and MUDDFormer primarily add feedforward cross-layer paths within the current token. WhiteMatter targets the remaining gap: allowing shallow consumers to use deep representations of past tokens in a dynamic, cache-efficient way.

Core Problem

A token’s complete stack of hidden states is computed, but later tokens cannot freely access that stack across depth. Giving every consumer all source layers would increase memory traffic and cache cost; using one shared fixed mixture prevents different consumers or tokens from choosing different sources. The desired architecture must therefore provide all-depth feedback, consumer-specific connectivity, content dependence, and controllable memory overhead simultaneously.

Innovation

WhiteMatter moves cross-layer connectivity into KV production. A router reads all layer states for one token and dynamically forms shared channels; each consumer reads one channel. k=L preserves full channel capacity, whereas k<L compresses the cache. Relative to LCKV’s single-source feedback and FusedKV’s static fusion, the method permits deep-to-shallow paths and token-dependent mixtures. Cyclic Gauss–Seidel supplies a practical approximation to the resulting feedback fixed point.

Methodology

  • �� Input: hℓ[i], the state entering layer ℓ at position i, is separately RMS-normalized for Key and Value branches.
  • �� Routing: linear routers output αK[i], αV[i]∈R^{k×L}; each mixed channel is ˜h_j=Σℓα[j,ℓ]ĥℓ. Weights are signed and token-dependent.
  • �� Projection: channel-specific W^K_j and W^V_j plus RMSNorm produce KV; Keys receive QK normalization and RoPE.
  • �� Reading: consumer layer ℓ reads channel ℓ mod k and performs standard causal SDPA.
  • �� Training: solve P=Pool(H), H=States(X;P). Partition positions into g strided groups, update them sequentially with cyclic Gauss–Seidel, and backpropagate only through the final passes.
  • �� Decoding: generate a token’s layer stack first, append its channels afterward, and use a learned dummy cache entry to prevent same-position self-reading.

Experiments

All models used a Qwen3 decoder with D=512, intermediate size 1536, six query heads, three KV heads, and head dimension 96. Data came from FineWeb-Edu, tokenized with the Qwen3-0.6B tokenizer, packed to length 2048, and trained for 30,518 steps or 8.0B tokens. Baselines were vanilla 16/24/32-layer models and LCKV w=4/7; WhiteMatter used k=8/16. Metrics included held-out perplexity, LAMBADA, WikiText, PIQA, HellaSwag, ARC-E, and OBQA. WhiteMatter used g=8 with one no-gradient and two gradient-carrying passes.

Results

Full-cache 16-layer WhiteMatter achieved 19.968 PPL, versus 21.747 for vanilla 16L and 20.181 for vanilla 24L. Half-cache k=8 achieved 20.377, a 6.3% reduction from vanilla 16L and 5.0% below equal-cache LCKV w=7 at 21.461. LAMBADA/WikiText were 60.73/43.28 for k=16 and 71.58/44.40 for k=8. In the exact-reference experiment, four cyclic passes reached the 1% criterion and produced a 13.9× prefill speedup.

Applications

The method is directly relevant to autoregressive LLM serving, where KV cache often dominates memory for long contexts and large batches. Choosing k<L can reduce memory while retaining quality; cross-layer feedback may improve the effective depth of compact models. Deployment requires kernels that support shared multi-channel caches, and g and k should be tuned to hardware bandwidth, sequence length, and latency targets.

Limitations & Outlook

The study does not cover very large models, extremely long contexts, multilingual or multimodal workloads, or broad hardware configurations. Dynamic routing adds computation, cyclic training costs about 1.5× vanilla, and fixed consumer assignment may limit expressivity. Fixed-point approximation, dummy-token masking, and iteration counts must be aligned between training and serving. Future work should measure end-to-end systems and investigate soft reads, sparse routing, quantized caches, and adaptive stopping.

Plain Language Accessible to non-experts

Imagine a 16-floor archive. In the ordinary system, when the next worker arrives, the worker on floor 3 may read only the file prepared on floor 3 for the previous document. A detailed report from the basement or a useful summary from the top floor is inaccessible. Feedback Transformer lets everyone read one fixed blended file, but the same blend is used for every worker and document.

WhiteMatter adds an intelligent sorting desk. For each document, the desk examines its versions on all floors and mixes them into several new packets. Different workers may receive different packets, and the mixing changes with the document. If there are 16 floors but only 8 packets, workers share packets and the storage requirement is halved.

There is a circular dependency: preparing one document uses packets from earlier documents, while those packets depend on completed work. The paper solves this by repeatedly updating interleaved groups, like teams revising a shared spreadsheet in sequence. The full-cache model reduced perplexity from 21.747 to 19.968, and the half-cache version reached 20.377. The approach is promising, but sorting costs more during training and still needs testing in much larger archives.

ELI14 Explained like you're 14

Picture a video game with 16 levels. Every turn creates clues at every level. In a normal Transformer, the next turn’s level 3 can read only the clue from level 3 of the previous turn. Level 16 may have discovered the best strategy, but level 3 is not allowed to see it. Weird rule, right?

WhiteMatter is like a clever teammate who collects all 16 clues and packs them into a few strategy bundles. Different levels can grab different bundles, and the teammate changes the packing depending on the enemy. If the game normally stores 16 bundles, it can store only 8 and let several levels share them, saving memory.

Training is tricky because new clues depend on earlier bundles. The researchers update groups of turns one after another, so later groups can immediately use fresher information. The results are strong: perplexity fell from 21.747 to 19.968 for the 16-layer model, beating a normal 24-layer model at 20.181. Even with half the cache, it reached 20.377.

Still, it is not free power! Training costs about 1.5 times as much, and the fixed bundle-sharing rule may not always be perfect. Researchers also need to test giant models, very long conversations, and real phone or cloud hardware. The big idea is simple: instead of building more floors, help the existing floors talk better!

Glossary

KV cache

The stored Keys and Values used to avoid recomputing attention for earlier tokens. Its size grows with layers, sequence length, and KV heads.

WhiteMatter stores k shared KV channels per token instead of one set for every layer.

Cross-layer KV pool

A mechanism that mixes several layer states from one token before projecting them into Keys and Values. It places cross-layer communication at cache production.

This is WhiteMatter’s central module, with separate Key and Value branches.

Content-dependent router

A learned linear module that generates mixing weights from the current token’s normalized multi-layer states. The weights vary by token and may be signed.

The routers produce αK and αV for source-to-channel mixing.

Cyclic Gauss–Seidel

An iterative schedule that updates interleaved token groups in order, allowing later groups to use current-pass updates. It trades parallelism for faster convergence.

WhiteMatter uses it to approximate feedback dependencies during training and prefill.

Fixed-point iteration

Repeatedly applying state-to-KV and KV-to-state transformations until they become mutually consistent. It resolves the circular dependency created by feedback.

Jacobi and cyclic Gauss–Seidel are alternative solvers for P=Pool(H), H=States(X;P).

RoPE

Rotary positional encoding injects position through rotations of attention vectors. It preserves positional information while remaining compatible with causal attention.

WhiteMatter applies RoPE to normalized Keys before storing them.

Open Questions Unanswered questions from this research

  • 1 It remains unknown whether dynamic mixing remains stable in much larger models and contexts of hundreds of thousands of tokens, or whether new quality–compression trade-offs appear.
  • 2 The fixed ℓ mod k read policy may be suboptimal. Learned consumer routing, sparse access, and hardware-aware assignment could improve flexibility but may increase bandwidth and kernel complexity.
  • 3 The paper emphasizes perplexity and benchmark accuracy, leaving limited system evidence on throughput, energy, latency, and memory behavior across GPUs and batch sizes.

Applications

Immediate Applications

Long-context LLM serving

Serving providers can test k=8 to reduce a 16-layer model’s KV cache by 50%. This is attractive for memory-limited GPUs, large batches, and long prompts, but kernels and iteration settings must be tuned for the target hardware.

Quality improvement for compact models

Edge assistants and specialized models can gain effective depth through cross-layer feedback instead of simply adding layers. The FineWeb-Edu results show a 16-layer WhiteMatter model outperforming a 24-layer vanilla model, making it a candidate architecture for compact deployment.

Long-term Vision

Hardware-aware dynamic memory

Future systems could select channel counts and access patterns based on token content, layer importance, and device bandwidth. Such adaptable memory could jointly control quality, capacity, and latency, although routing overhead and stability remain obstacles.

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.

cs.CL cs.LG