基于分块注意力的循环记忆增强 Transformer 用于长上下文语言建模

我们提出了一种用于长上下文语言建模的 Transformer 架构，将全局注意力与两种生物学启发的组件相结合：分块局部注意力和门控 FIFO 记忆机制。该统一注意力块允许模型在不增加注意力成本呈二次增长的情况下高效处理短程和长程依赖。记忆模块采用类循环网络启发的门控更新机制持久存储过去的 token 表示。每个注意力头均应用旋转位置编码，以实现方向性解耦、尺度不变的位置信号。该架构完全从 PyTorch 中实现，无需依赖高级库，实现透明且模块化的实验。我们的模型为诸如对话建模、代码补全和文档理解等任务提供了轻量且可扩展的设计。