论文

三月:使用内容路由状态锚扩展循环内存

MARCH: Scaling Recurrent Memory with Content-Routed State Anchors

摘要

Transformer 强大的长上下文检索能力很大程度上归功于随上下文长度增长的 词元级 内存。然而,这种灵活性会在训练期间带来二次计算复杂性,并且在自回归推理期间导致键值缓存线性增长。循环替代方案通过将整个历史压缩为固定大小的状态来提供有效的解码,但在召回密集型任务上通常表现不佳,因为早期的关联通常会被后续更新覆盖,并且仅保留最新的上下文信息。在本文中,我们介绍了跨上下文历史的记忆锚路由(MARCH),这是一种网络架构,可以有效地将状态空间模型扩展到固定大小的维度之外,同时保持长序列的计算效率。 MARCH 定期缓存累积的循环状态检查点作为状态锚,并将每个锚与一个紧凑的、内容调节的锚密钥相关联。这使得 MARCH 能够维护一个内存库,该内存库可以随着上下文长度的增加而增长,从而在历史分辨率和内存成本之间提供可控的权衡。在每个标记处,MARCH 都会生成一个锚点查询来参与所有因果上可用的状态锚点,并且输出被计算为沿当前状态的所有历史锚点的注意力式聚合。我们表明,经过标准预训练后,MARCH 在常识推理、LongBench 和上下文检索方面始终优于多个线性注意力变体。这些结果表明,内容路由状态缓存大大增强了循环远程内存,同时保留了其本机计算路径。