DMT用选择性动态递归记忆与注意力读取实现常量状态开销

Transformers 保留了序列的明确历史记录，允许注意力直接检索过去的表示。然而，KV-Cache 的大小和新 token 的成本随着上下文长度的增加而增加。状态空间模型将过去压缩为固定大小的循环状态，但代价是降低了可寻址性。我们提出了动力学记忆 Transformer (DMT)，这是一种将序列编码与信息读取分开的混合架构。编码器使用多个存储单元，每个存储单元均由具有复数对角循环矩阵的动态系统实现。依赖于输入的泄漏率控制着它们的更新。然后，解码器使用注意力来读取并组合它们的输出，而不需要遍历 token 历史记录。在维度固定时，该递推形式支持总工作量为 O(T)、并行深度为 O(log T) 的结合扫描。在自回归推理期间，循环状态和卷积缓存的大小保持固定。我们在认知和语言建模任务上将 DMT 与 Mamba 和 Transformer 进行比较，并展示其优点和局限性：DMT 在多项认知任务上取得了最佳结果，语言建模性能接近 Transformer，但在语言建模方面仍落后于 Mamba。在选择性复制方面，DMT 还表现出比 Mamba 和 Transformer 更强的长度外推能力。这些结果表明，并行动态记忆单元是保留和使用序列信息的有前途的机制。