Maglev: Sliding Recurrent Memory

技术摘要：Maglev —— 滑动循环记忆

问题陈述
当前的 Transformer 架构面临着记忆容量与计算效率之间的根本权衡。全因果自注意力（Full causal self-attention）能够实现强大的非线性检索和上下文计算，但其成本随上下文长度线性增长。相反，滑动窗口注意力（Sliding-window attention）限制了注意力和 KV 缓存的成本，却丢弃了远距离信息，无法提供持久的长期记忆。经典的循环模型（如 LSTM）虽然压缩了历史信息，但其状态传递是顺序性的，限制了容量和序列并行性。虽然记忆 Transformer（Memory Transformers）和线性注意力模型试图弥补这一差距，但它们通常会引入人为的分段边界、依赖结构化的线性更新，或在训练期间需要顺序展开（sequential unrolling）。核心挑战在于开发一种实用的、针对 Token 的非线性循环记忆机制，使其既能用于 Transformer，又能实现并行的预训练，同时保持有界的推理成本。

方法论：Maglev 架构
Maglev 通过将训练期间的记忆构建与推理期间的记忆传播解耦来解决这一问题。该架构由两个耦合的模型组成：

1. 预填充器 (Prefiller, $Q$)： 一个因果 Transformer，通过较大的感受野（例如，交错的全注意力与滑动窗口注意力）并行处理观测到的序列。它为每个 Token 位置生成一个目标记忆向量轨迹 $m'_t$。
2. 解码器 (Decoder, $P$)： 一个滑动窗口 Transformer，用于预测下一个 Token 并生成自身的记忆向量 $m_t$。至关重要的是，$P$ 无法访问完整的历史记录；它只能看到最近的有限窗口 Token 以及一个偏移后的记忆向量窗口。

训练机制：提升并行训练 (Lifted Parallel Training)
为了避免顺序展开（这会阻碍并行化），Maglev 采用了两阶段训练策略：

• 第一阶段： 预填充器 $Q$ 并行生成完整的目标记忆轨迹 $m'_{1:T}$。
• 第二阶段： 解码器 $P$ 结合输入 Token 消费偏移后的目标轨迹 $m'_{0:T-1}$，以预测下一个 Token 并产生自身的记忆 $m_{1:T}$。
• 目标函数： 模型使用复合损失函数进行优化：
  • 下一 Token 预测损失 ($L_{CE}$)： 标准的语言模型交叉熵损失。
  • 记忆一致性损失 ($L_{cons}$)： 一个用于将解码器的输出记忆 $m_t$ 与预填充器的目标 $m'_t$ 对齐的项。这教会了解码器如何在不进行顺序展开的情况下，重现下一步所需的记忆状态。

推理机制
在推理阶段，预填充器 $Q$ 被丢弃。解码器 $P$ 以循环方式运行，将其先前生成的记忆 $m_{t-1}$ 反馈回系统中。

• 循环 K/V 注入： Maglev 并非使用专门的记忆 Token，而是通过学习到的投影和门控机制，将偏移后的记忆状态注入到注意力机制的键/值（K/V）特征中。
• 有界复杂度： 解码器维护一个固定大小的 Token 和记忆滑动窗口。循环状态通过 K/V 缓存向前传递，确保推理成本（计算量和缓存大小）保持有界且独立于总序列长度，符合标准滑动窗口注意力的特性。

核心贡献

1. 逐 Token 非线性循环： Maglev 引入了一种机制，使每个 Token 的解码器状态都能作为记忆，遍历 Transformer 的完整非线性深度，而无需额外的记忆 Token 或人为的分段边界。
2. 循环模型的并行训练： 通过使用预填充器并行生成目标，Maglev 实现了非线性循环模型的并行预训练，避免了传统 RNN 训练中的顺序瓶颈。
3. 参数共享： 作者证明了预填充器 $Q$ 和解码器 $P$ 可以共享大部分参数（特别是 Transformer 块）而不显著降低性能，从而减少了训练设置中的参数内存开销。
4. 有界推理： 部署的模型保留了滑动窗口注意力的效率，同时通过循环 K/V 注入获得了持久的记忆能力。

实验结果
在 nanochat d20 架构（435M 参数，在 43.52B Token 上训练）上进行评估，并将 Maglev 与交错全注意力/滑动窗口 Transformer、纯滑动窗口模型以及潜在循环 Transformer (LRT) 基线进行了对比。

• 性能： 最优的 Maglev 变体（独立参数，$\lambda=1$）将 FineWeb-Edu 验证集每字节比特数 (BPB) 从 0.7413（滑动窗口基线）提升至 0.7251。
• 下游任务： 它实现了 56.4% 的平均下游准确率，超越了匹配的滑动窗口 Transformer (54.1%) 和匹配的 LRT 基线。
• 一致性权重： 研究指出，最优的一致性权重 ($\lambda$) 取决于参数共享策略；较高的 $\lambda$ 有利于独立参数模型，但可能会限制共享参数模型。

意义与主张
论文声称，Maglev 提供了一条实用的路径，为非线性 Transformer 配备持久的、逐 Token 的记忆，且该记忆可以在每个 Token 之后进行重写。它成功地弥合了全注意力的表达能力与有界上下文模型的高效性之间的差距。通过在训练期间分离记忆生成与记忆使用，Maglev 实现了非线性循环架构的并行预训练，这在以往是难以实现的。作者将 Maglev 定位为对该设计空间的一次初步探索，并指出虽然它在测试设置中优于滑动窗口和潜在循环基线，但仍需进一步工作来扩展该方法、优化部署时的算子效率，并探索其他的记忆注入机制（例如残差流注入或交叉注意力）。