用分块记忆机制替代全局注意力,提升长文本建模效率
BCMT: Blockwise Causal Memory Transformer
- 将序列分块处理,每块独立计算局部注意力
- 通过指数因果记忆聚合信息,实现高效长程依赖传播
- 训练速度更快、内存更低,适合超长文本场景
Transformer架构依赖密集自注意力建模长程依赖,但其复杂度随序列长度呈二次增长。本文提出BCMT(分块因果记忆Transformer),将局部交互与全局上下文传播解耦:在局部块内使用密集因果自注意力,每个块通过指数因果记忆生成自适应摘要,并回注到词元表示中,实现无需显式全局注意力的长程信息传递。与标准Transformer和循环记忆架构不同,BCMT既不依赖远距离词元的密集交互,也不使用可学习的记忆状态,其记忆机制完全并行化且兼容标准自注意力实现。在最长1024个词元的文本建模任务上,BCMT性能接近密集Transformer,同时显著提升训练吞吐量并降低内存消耗。消融实验验证了该改进源于所提出的记忆机制。结果表明,由块摘要构建的指数因果记忆,是长序列语言建模中稠密全局注意力的有效替代方案。
原文摘要 · Abstract (English)
Transformer architectures rely on dense self-attention to model long-range dependencies, but this mechanism exhibits quadratic complexity with respect to sequence length. We introduce BCMT (Blockwise Causal Memory Transformer), an architecture for long-context language modeling that decouples local token interactions from global context propagation. Dense causal self-attention is applied independently within local blocks, while each block produces an adaptive summary aggregated through an exponential causal memory. This memory is subsequently injected back into the token representations, enabling efficient propagation of long-range contextual information without relying on explicit global attention. Unlike standard Transformers and recurrent memory architectures, BCMT maintains neither dense interactions between distant tokens nor learned memory states. Its memory mechanism is fully parallelizable and remains compatible with standard implementations of dense self-attention. Experiments on language modeling with context lengths of up to 1024 tokens show that BCMT achieves validation performance comparable to that of Dense Transformers while significantly improving training throughput and reducing memory consumption. An ablation study further confirms that these improvements arise from the proposed memory mechanism. These results demonstrate that an exponential causal memory constructed from block summaries provides an effective alternative to dense global attention mechanisms for long-context language modeling.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。