PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Wenbo Zhang, Xiang Ren

arXiv:2608.18486v1cs.CLcs.LG

TL;DR

Transformers cannot directly use deeper past-token representations at shallower layers during autoregressive decoding. WhiteMatter mixes all layer states into shared, token-adaptive KV channels, and full-cache WhiteMatter lowers perplexity by 8.2% versus a same-depth vanilla baseline while outperforming a deeper vanilla model.

  • Problem

    Transformers restrict each layer to same-depth KV from past tokens, limiting access to information already produced at deeper layers during decoding.

  • Method

    WhiteMatter uses a router to mix all layer states into k shared KV channels, enabling consumer-specific, source-token-adaptive cross-layer connections.

  • Results

    8.2% lower perplexity than a same-depth vanilla baseline, with full-cache WhiteMatter also outperforming a 24-layer vanilla model.

  • Takeaways & Limitations

    A half-sized KV cache retains most of WhiteMatter’s improvement, achieving a 6.3% perplexity reduction relative to the 16-layer vanilla model.

  • Takeaways & Limitations

    The experiments use small models and an 8B-token budget, leaving scaling with model size, data, and optimized end-to-end decoding unresolved.

Abstract

from arXiv · show

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.

1 INTRODUCTION

WhiteMatter addresses the restriction that each Transformer layer can attend only to same-depth KV by dynamically connecting consumer layers to representations from all source depths. It uses shared KV channels to enable consumer-specific, source-token-dependent mixing while reducing cache size, and improves perplexity over vanilla and LCKV baselines.

  • Motivation: Each Transformer layer can attend only to KV from its own depth, preventing use of information produced at other layers for past tokens.The model produces a stack of hidden states for each token, but subsequent layers access only same-depth KV.
  • WhiteMatter: WhiteMatter routes all L layer states into k shared KV channels, allowing each consumer layer to select a different channel.The router mixes source-layer hidden states at each token position, and consumer layers receive different connections to source depths.
  • WhiteMatter: The router makes source-to-KV connection weights depend on token content, providing dynamic connectivity rather than a static shared mixture.Because the router reads hidden states, its connection weights adapt to the source token.
  • Results: 19.968 held-out perplexity at 16 layers and k=16 is 8.2% lower than a same-depth vanilla model’s 21.747 and slightly below a 24-layer vanilla model’s 20.181.All models were pretrained from scratch on 8B FineWeb-Edu tokens with matched data, token budget, and optimizer settings.
  • Results: 20.377 perplexity with k=8 is 5.0% below an LCKV baseline with the same cache size.Halving the cache from k=16 to k=8 retains most of WhiteMatter’s full-cache gain.
  • WhiteMatter: k < L reduces KV-cache size by sharing channels among consumer layers.The number of channels controls cache size, so fewer channels than layers compress the cache.

2 RELATED WORK

Prior work connects Transformer layers through shared or fixed deep-to-shallow KV pathways, feedforward residual mixing, and other KV-routing schemes. Related research also explores repeated computation by recycling hidden states or latent inputs.

  • Deep-to-shallow feedback connections: Deep-to-shallow feedback methods either share one connection pattern across consumers or use a fixed connection from a single deep source.Fan et al. (2021) pools all layer states into shared KV, while Tu (2024) uses top-layer KV and Cai et al. (2026) uses a cached middle-layer state.
  • Feedforward cross-layer connections: Feedforward cross-layer methods mix earlier layers’ outputs within the residual stream, with some making mixing content-dependent or learning parallel-stream connections.MUDDFormer separately aggregates Q, K, V, and residual streams using content-dependent weights.
  • KV and value pathways: Other methods form cross-layer connections through value-residual additions or grouped KV sharing across layers.Examples include per-layer coefficients or per-token gates for value residuals, and grouped schemes such as CLA, MLKV, and YOCO.
  • Latent reasoning via repeated computation: Latent-reasoning methods recycle hidden states, input embeddings, or inserted latent positions to enable repeated computation.Some PonderLM variants use adaptive halting, and inserted-position variants append latent inputs after selected observed tokens.

3 METHOD

WhiteMatter replaces per-layer KV projections with a cross-layer pool that dynamically mixes all L layer states into k shared KV channels, reducing cache size to k/L of a standard cache. Fixed per-layer channel selection preserves one KV read per layer, while autoregressive decoding and parallel training use procedures that resolve the resulting feedback dependency.

  • Cross-layer KV pool: WhiteMatter’s cross-layer KV pool mixes all L layer states at each past token into k ≤ L shared channels before projection and caching.The router produces data-dependent mixtures independently for keys and values.
  • Cross-layer KV pool: k/L of a standard L-layer cache is required because the cache stores k key and value channels per token.Stored keys are rotated and K-normalized, while values remain raw; each cache contains k · T · Hkv · d elements.
  • Cross-layer KV pool: Signed, position-dependent mixing weights let the key and value branches select different combinations of source-layer representations.The router may read every pth source layer to reduce its input size while producing weights for all L source layers.
  • Per-layer channel selection: Fixed cyclic selection assigns layer ℓ to channel ℓ mod k, so each layer performs standard causal attention using one channel.When k=1 all layers share one channel; when k=L each layer reads its corresponding channel, preserving one KV read per layer in between.
  • Execution schedules: Autoregressive decoding appends a token’s pooled channels only after its layer stack, while parallel training formulates the feedback dependency as a fixed-point problem.Jacobi uses token-parallel passes with cost linear in the number of passes; cyclic Gauss–Seidel instead updates strided groups sequentially within each pass.

4 EXPERIMENTS

WhiteMatter improves held-out language modeling over same-depth vanilla and deeper baselines, while half-cache compression retains most of the gain. It also transfers to downstream tasks and enables substantially faster converged prefill with cyclic evaluation.

  • Language modeling: 8.2%: Full-cache WhiteMatter reduces perplexity from 21.747 to 19.968 versus the 16-layer vanilla model and beats 24-layer vanilla at 20.181.Full-cache WhiteMatter uses 54.1M non-embedding parameters versus 51.9M for the 16-layer vanilla baseline.
  • Language modeling: 6.3%: Half-cache WhiteMatter reaches 20.377 perplexity, retains most of the full-cache improvement, and comes within 1.0% of the 24-layer model.It uses 50.6M non-embedding parameters, slightly fewer than the 16-layer vanilla baseline.
  • Baseline comparison: 5.0%: At equal cache size, WhiteMatter k=8 lowers perplexity to 20.377 versus 21.461 for LCKV w=7.Half-cache WhiteMatter also outperforms equal-cache LCKV on both language-modeling benchmarks and every reported multiple-choice task except ARC-Easy and OpenBookQA.
  • Downstream evaluation: Full-cache WhiteMatter has the best language-modeling perplexity among 16-layer models and the highest PIQA and HellaSwag accuracy.Both WhiteMatter variants outperform the 32-layer vanilla model on LAMBADA, with 60.73 and 71.58 versus 79.39 perplexity.
  • Prefill convergence: 13.9×: Cyclic g=16 reaches the quality threshold in 4 passes at 0.01245 s/sequence, versus 0.1729 s for autoregressive evaluation and 0.1393 s for Jacobi.Increasing to g=32 does not reduce the pass count and is slower because each pass costs more.
  • Compute cost: 2.5×: WhiteMatter costs around 2.5× the vanilla training FLOPs and 3.3× the prefill FLOPs, while decoding computation remains nearly identical across methods.The training multiplier is lower than the pass count because of truncated backpropagation.

5 ANALYSIS

The analysis shows that training schedules closer to convergence improve quality, while WhiteMatter’s performance depends on channel capacity, deep-to-shallow feedback, and dynamic routing. Even compressed configurations outperform vanilla baselines under the reported conditions.

  • Training schedule: 32% lower perplexity separates the strongest evaluated training schedule from the weakest, with added gradient or no-gradient passes improving performance at diminishing returns.Schedules farther from the fixed point degrade when iterated beyond training, including during autoregressive decoding.
  • Pool rank: 7.3% perplexity reduction is achieved by k=1 versus the vanilla baseline, while providing 16× KV-cache compression.Increasing the number of channels improves performance with diminishing returns.
  • Deep-to-shallow feedback: 7.5% higher perplexity remains for WhiteMatter without deep-to-shallow feedback than for full-cache WhiteMatter, despite a 16× larger KV cache than k=1.The feedback-free model still outperforms vanilla through dynamic KV mixing but has costs similar to vanilla and does not require iteration.
  • Dynamic routing: About 2% higher perplexity occurs in both static-mixing models compared with their dynamically routed counterparts.The comparison includes static models with k=16 and k=1.

LIMITATIONS

WhiteMatter improves decode-time performance and KV-cache efficiency but incurs substantially higher training and prefill costs. Evidence remains limited to small models and an 8B-token budget without optimized end-to-end decoding benchmarks.

  • Training and prefill costs: 2.3–2.5× vanilla FLOPs are required for WhiteMatter training, while three-pass prefill requires 3.1–3.3× vanilla FLOPs.Training and prefill require autoregressive processing or multiple parallel iterations, although decoding uses similar FLOPs to vanilla decoding and can reduce memory consumption.
  • Empirical scope: Main results come from small models trained with an 8B-token budget, leaving scaling with model size and data unestablished.The experiments do not establish how quality or systems trade-offs scale.
  • Empirical scope: The paper reports cache size and schedule convergence but lacks an optimized end-to-end decoding benchmark, leaving larger-scale evaluation for future work.Evaluating larger models and optimized end-to-end decoding remains future work.

A WHITE-MATTER CONNECTIVITY

Figure 7 presents a population-averaged human structural connectome reconstructed from diffusion MRI, showing whole-brain white-matter tracts spanning distant regions. Sagittal and coronal views depict fiber orientation through color coding.

  • A WHITE-MATTER CONNECTIVITY: Figure 7 shows whole-brain white-matter tractography reconstructed from diffusion MRI.The image depicts a population-averaged human structural connectome.
  • A WHITE-MATTER CONNECTIVITY: The fiber tracts span the brain and arc between distant regions in every direction.The visualization emphasizes long-range structural connectivity across the brain.
  • A WHITE-MATTER CONNECTIVITY: Sagittal and coronal views use color to encode local fiber orientation: red left–right, green anterior–posterior, and blue superior–inferior.The rendering was produced with DSI Studio using a population-averaged human template.

B CONVERGENCE OF A LARGER CYCLIC-TRAINED MODEL

For the larger cyclic-trained model, cyclic g=8 converges in five passes and substantially reduces per-sequence inference time versus Jacobi and autoregressive rollout. The evaluated model has 8 layers, D=1024, k=8, and T=4096.

  • Model and evaluation: The evaluated model has 8 layers, D=1024, k=8, and T=4096.Figure 8 reports convergence timing for this larger cyclic-trained model.
  • Training setup: The model was trained for 122,000 steps on approximately 4.0B tokens using a cyclic g=8 schedule.Training used global batch size 8 and sequence length 4096.
  • Convergence and timing: 5 passes and 0.159 s/sequence make cyclic g=8 7.7× faster than Jacobi and 15.5× faster than autoregressive rollout.Jacobi required 52 passes and took 1.220 s/sequence, while autoregressive rollout took 2.470 s/sequence.

C FULL DOWNSTREAM RESULTS

Table 3 presents complete zero-shot downstream results from the lm-evaluation-harness suite, using accuracy-based metrics for most tasks and perplexity for LAMBADA and WikiText. Task inclusion depends on exceeding random-choice baselines, with BoolQ evaluated against a 62.17% majority-class baseline.

  • Evaluation setup: Table 3 covers every task in the zero-shot lm-evaluation-harness suite, reporting normalized accuracy for PIQA, HellaSwag, ARC-Easy, ARC-Challenge, and OpenBookQA.WinoGrande and BoolQ use accuracy.
  • Baselines and inclusion: 62.17% is the majority-class baseline used for BoolQ, while WinoGrande and ARC-Challenge remain near 50% and 25% random-choice baselines, respectively.Multiple-choice tasks appear in the main table when at least one model exceeds random choice by two estimated standard errors.
  • Evaluation setup: LAMBADA and WikiText report perplexity, while all other Table 3 columns report accuracy in percent.This metric distinction applies across the complete zero-shot downstream-results table.
Loading 2608.18486v1…