RATTENTION：向局部-全局注意力模型获取最小滑动窗口大小迈进

局部-全局注意力模型最近涌现为标准Transformer的有力替代方案，承诺在训练和推理效率方面取得改进。然而，窗口大小的关键选择存在帕累托权衡：较大的窗口在保持类似完整注意力性能方面效果良好，但在短上下文场景中几乎没有效率收益，而较小的窗口则可能导致性能下降。当前模型如 Gemma2 和 Mistral 采用保守的窗口大小(例如 4096 相对于 8192 的预训练长度)以保留性能。本工作旨在探索将这种帕累托前沿向前移动的策略，使局部-全局模型即使在短上下文情境下也能实现效率收益。我们的核心动机是解决局部注意力的固有限制——其完全忽略定义窗口范围之外的标记。我们探索了 RATTENTION，这是一种局部注意力的变体，集成了旨在捕获这些超出窗口标记信息的专用线性注意力机制。在 3B 和 12B 规模的预训练实验中，RATTENTION 实现了性能与效率之间的优越帕累托权衡。作为理想点，窗口大小仅为 512 的 RATTENTION 在各种设置下始终匹配完整注意力模型的性能。此外，RATTENTION 中线性注意力组件所固有的循环特性有助于增强长程上下文性能，经 RULER 框架验证。至关重要的是，这些改进并未妥协训练效率；得益于专用内核实现和缩小窗口大小，RATTENTION 在训练速度方面保持与现有SOTA方法的持平水平。我们开源了 Pallas kernels 以及模型代码，以促进进一步的研究。