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

WhiteMatter. All-to-All Cross-Layer Connections via KV Source Mixing. By Wenbo Zhang and one other, from the University of Southern California. Published on the twenty-seventh of September. Its code repository is worked out: about one hundred and eighty source files, with a readme and a licence.

A Transformer generating text builds a representation of every past token at every layer. Yet each layer can normally read only the representations from its own depth. A middle layer attends to what that same layer produced, and nothing else. The authors see a loss in this: the model has already computed information at other depths, and the restriction stops it from fully reusing what it has.

WhiteMatter lifts the restriction. Every layer is allowed to draw on past-token representations from any depth. A learned mixer decides, for the current context, which depths are most useful, and combines their representations into shared key-value cache channels. The mixing happens at the source of the keys and values. In effect the cache is no longer a private record per layer but a common pool that any layer can consult. And because those channels are shared across layers, the cache itself can shrink, so the same idea that widens what a layer can see also narrows what has to be stored.

Connecting layers across depth has a cost, and the paper names it. The new connections introduce dependencies that slow down training and prompt processing. The answer is cyclic iteration. Tokens are split into interleaved groups, the groups are updated in turn, and the tokens inside each group are processed in parallel. On a reference model trained with exact autoregressive execution, cyclic iteration converges twelve point five times faster than standard Jacobi iteration.

Two comparisons carry the case for the architecture itself. At the same number of training tokens, WhiteMatter with a full-size cache performs comparably to a standard Transformer with fifty percent more layers. And with half the key-value cache, it outperforms matched standard Transformers at two model scales, reaching up to one point three billion parameters. So a smaller cache does not have to mean a weaker model here.

One way to hear WhiteMatter is as a trade between resources. Memory is saved because depths share a cache, and quality holds because any layer can reach any depth. The price is sequential dependency between layers, and cyclic iteration is the authors' way of paying it down.