发表机构
University of Southern California（南加州大学）
机构由 AI 辅助整理，请以论文原文为准。AI 中文总结
WhiteMatter通过将每个注意力层连接到所有过去token的所有层表示，以可变连接权重实现全对全跨层KV混合，在预训练中性能优于多50%层数的普通Transformer，缓存压缩50%仍保留大部分增益。
AI 中文摘要
在Transformer中，每个层仅通过自身深度产生的KV来关注过去的token，尽管在自回归解码过程中存在更深层的表示。反馈架构允许浅层的消费层关注更深层过去token表示产生的KV，但为所有消费层提供与源层相同的固定连接模式。我们提出WhiteMatter，它将每个注意力层连接到每个过去token的所有层的表示，连接权重可在消费层之间变化并适应源token。对于每个token，一个路由机制通过将其L层状态混合为k个KV通道来实现这些连接，这些通道会被缓存以供后续token使用；每个消费层关注其中一个通道。通道数k控制KV缓存的大小，设置k<L可减少缓存的内存占用。在我们的预训练实验中，WhiteMatter的性能优于层数多50%的普通Transformer，且在KV缓存压缩50%时仍保留了大部分增益。
英文摘要
When generating text, a Transformer produces representations of past tokens at every layer, but each layer can normally use only representations from the same depth. This restriction prevents the model from fully reusing information it has already computed. We introduce WhiteMatter, which allows every layer to draw on past-token representations from any depth. A learned mixer selects the most useful depths for the current context and combines their representations into shared key-value (KV) cache channels. Sharing these channels across layers can reduce the cache size. Given the same number of training tokens, WhiteMatter with a full-size cache performs comparably to a standard Transformer with 50% more layers. With half the KV cache, WhiteMatter outperforms matched standard Transformers at two model scales, up to 1.3B parameters. Cross-layer connections, however, introduce dependencies that slow training and prompt processing. We address this problem with cyclic iteration, which updates interleaved groups of tokens in turn while processing the tokens within each group in parallel. On a reference model trained with exact autoregressive execution, cyclic iteration converges 12.5x faster than standard Jacobi iteration.
CommentsCode available at https://github.com/Cy-47/White-Matter
