Search papers, labs, and topics across Lattice.
This paper introduces WhiteMatter, a novel architecture that enables all attention layers in a Transformer to access key-value (KV) representations from all previous layers, allowing for more flexible and adaptive connections. By implementing a router that mixes layer states into multiple KV channels, WhiteMatter reduces memory usage while maintaining performance. The results show that WhiteMatter outperforms a standard Transformer with significantly more layers, achieving comparable performance with a 50% reduction in KV-cache size.
WhiteMatter achieves superior performance with fewer resources by allowing each attention layer to dynamically access all previous layer representations, challenging traditional fixed connection patterns.
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.
