线性注意力架构:机制、权衡和跨层路由
Linear Attention Architectures: Mechanisms, Trade-offs, and Cross-Layer Routing
摘要
自注意力使每个标记从完整上下文中检索信息,但其序列长度的二次成本限制了长上下文中的训练和推理。本文介绍了 softmax 注意力和最近四种循环线性注意力架构的比较研究:DeltaNet、Gated DeltaNet、Kimi Delta Attention 和 Gated DeltaNet-2。我们用通用的循环记忆符号来表达这些机制,明确它们在表达性、记忆衰减、擦除和写入控制、训练吞吐量和实现复杂性方面的差异。我们的实验以针对 15B 词元训练的 3.5 亿参数模型为中心,包括优化器和学习率比较、混合与纯堆栈比较、序列长度运行时间测量、在 1.3B 和 3B 参数下运行的较大 DeltaNet 以及一小组下游评估。报告的速度结果衡量训练吞吐量和迭代时间;我们不提供经验推理速度基准。在报告的 350M 参数、15B 词元扫描中,使用 Muon 的 Kimi Delta Attention 达到了最低的最终验证损失,使用 AdamW 训练的纯门控 DeltaNet 堆栈具有最高的归一化训练吞吐量,混合堆栈通常以吞吐量成本改善损失,并且在我们评估的匹配架构设置中,Muon 相对于 AdamW 始终降低最终验证损失。我们介绍并评估了 DeltaNet 式存储器的轻量级跨层路由机制。最自然的受 DeltaNet 启发的公式,将较低层的增量规则写入错误转发到下一层的值目标中,并不会比匹配的基线有所改进。相反,路由到对齐的隐藏流并转发写入值会在我们报告的匹配运行中产生适度的改进:跨层值路由(CLVR)降低了 DeltaNet 和门控 DeltaNet 的最终验证损失。