用分块注意力和记忆机制,让Transformer高效处理长文本。
Recurrent Memory-Augmented Transformers with Chunked Attention for Long-Context Language Modeling
- 分块局部注意力+门控循环记忆,兼顾长短距离依赖
- 内存模块持久存储历史表示,降低计算开销
- 适合对话、代码补全等长序列任务
我们提出一种用于长上下文语言建模的Transformer架构,结合全局注意力与两种受生物启发的组件:分块局部注意力和门控FIFO记忆机制。该统一注意力模块可在不使注意力成本二次增长的情况下,有效处理短程和长程依赖。记忆模块通过受循环网络启发的门控更新机制,持久存储过去的词元表示。每注意力头应用旋转位置编码,以实现方向解耦、尺度不变的位置信号。整个架构在PyTorch中从零实现,无需高层库依赖,支持透明且模块化的实验。模型设计轻量且可扩展,适用于对话建模、代码补全和文档理解等任务。
原文摘要 · Abstract (English)
We present a Transformer architecture for long-context language modeling that combines global attention with two biologically inspired components: chunked local attention and a gated FIFO memory mechanism. This unified attention block allows the model to efficiently handle both short-range and long-range dependencies without increasing attention cost quadratically. The memory module persistently stores past token representations using a gated update mechanism inspired by recurrent networks. Rotary positional encoding is applied per attention head to enable directionally disentangled, scale-invariant positional signals. The architecture is implemented entirely from scratch in PyTorch, with no reliance on high-level libraries, enabling transparent and modular experimentation. Our model offers a lightweight and extensible design for tasks such as dialogue modeling, code completion, and document understanding.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。