Maglev：滑动窗口循环记忆 Maglev: Sliding Recurrent Memory
用并行预填充器生成记忆目标，一致性损失训练滑窗解码器，实现可并行预训练的固定内存循环语言模型
前置知识
滑动窗口注意力（Sliding-Window Attention, SWA）
一种稀疏注意力形式：每个位置只对最近 $W$ 个 token 计算注意力，而不是整个前缀。注意力计算量与 KV 缓存大小因此与序列长度无关，始终保持常数 $W$。Longformer、Mistral 等模型广泛使用。代价是窗口之外的信息被直接丢弃，模型没有任何持久的长期记忆。
Maglev 的部署模型本质上就是一个滑动窗口 Transformer，只有理解 SWA 的成本结构和信息丢弃问题，才能理解论文要补上的是什么。
KV 缓存
自回归推理时为避免重复计算，把每个 token 在每层的 key 和 value 向量缓存下来，新 token 只需与缓存中的 K/V 做注意力。缓存大小随序列长度线性增长，是长上下文推理的主要显存与带宽瓶颈。
Maglev 的核心卖点之一是推理时缓存与普通滑动窗口注意力完全相同——记忆通过混入既有 K/V 条目携带，不占任何额外缓存位置。
线性注意力与状态空间模型（Mamba 等）
一类用固定大小循环状态替代完整 KV 缓存的架构：线性注意力把注意力改写为可递推的核函数形式，状态空间模型（S4、Mamba）用线性或仿射的状态转移方程压缩历史。它们可并行训练、推理常数内存，但记忆更新是结构化的线性运算，表达力受限。
这是 Maglev 最重要的对比类别：Maglev 同样固定状态大小，但每步记忆更新走完整 Transformer 深度的非线性计算。
一致性损失与特征蒸馏
让学生模型的中间表示去对齐教师模型对应表示的损失项，典型形式是均方误差 $\|m_t - m'_t\|^2$。不同于只匹配最终输出分布的知识蒸馏，特征级监督能把教师“中间怎么想”的信息传递给学生，常用于压缩和表示对齐。
Maglev 的训练目标就是下一词交叉熵加一致性项，正是靠它让学生解码器学会自己产生下一步推理所需的记忆。
循环模型的并行训练难题与 teacher forcing
训练循环模型时，若每步输入依赖上一步的隐藏状态，就必须顺序展开，无法并行。NLP 中经典做法 teacher forcing 用真实数据代替模型自身输出，使各时间步可独立计算。但隐状态轨迹没有真值可供 forcing，这正是非线性循环记忆难以并行预训练的根源。
理解了这个矛盾，才能明白 Maglev 的“提升平行训练”——用一个因果预填充器并行生成记忆真值——到底巧妙在哪里。
RoPE 旋转位置编码
通过把 query 和 key 向量按其绝对位置对应的角度做复数旋转，使注意力得分只依赖相对位置的编码方案，是现代大模型的标配位置编码。
论文实现细节中提到：启用 RoPE 时需对局部 key 和注入的循环 key 都先做旋转再混合，读懂方法部分会用到。
研究动机
Transformer 的记忆策略是保留所有 token：全因果注意力让每次预测都能回看整个前缀，支撑强大的非线性检索与上下文计算，代价是注意力和 KV 缓存随上下文线性增长。滑动窗口注意力把两者都限定在最近 $W$ 个 token 内，但窗口外信息被直接丢弃，模型没有任何持久的长期记忆。经典循环网络（LSTM）用非线性逐 token 更新压缩历史，但状态太小且必须顺序传播，容量与并行性双双受限。Memory Transformer 一类方法（Transformer-XL、Compressive Transformer）虽有更丰富的状态，却只在段或块边界更新记忆，人为切分了记忆粒度。线性注意力和状态空间模型（Mamba 等）移除了边界并可并行，但它们的记忆更新被限制为结构化的线性或仿射运算。缺的正是这样一个实用方案：让非线性 Transformer 在每个 token 之后都能重写一份持久记忆。
本文的目标是本文要构建一个带逐 token 隐状态记忆 $m_t$ 的语言模型，同时满足三个性质。第一，逐 token 的非线性记忆：每次记忆更新穿越完整 Transformer 深度，相当于自带 looped Transformer 式的循环深度，且不需要专门的记忆 token 或缓存，不引入人工的段/块边界。第二，有界循环推理：处理 $x_t$ 后写下 $m_t$，仅保留固定的记忆窗口 $M_t = \{m_{t-W+1}, \dots, m_t\}$，推理时 token 窗口与记忆窗口都不随序列增长，缓存成本与滑动窗口注意力持平，形式化接口为 $p_\Theta(x_{t+1}, m_t \mid x_{t-W+1:t}, M_{t-1})$。第三，可并行的大规模预训练：虽然推理是循环的，预训练应能对所有位置并行处理，保持标准 Transformer 的吞吐量。前两条由接口构造直接保证，第三条并不自动成立，是全文要攻克的核心问题。
与已有工作不同的是，核心困难在于：文本只提供 $x_{1:T}$，不提供隐状态轨迹 $m_{1:T}$；由于每步更新要消费前一步的记忆，直接训练必须顺序展开。Maglev 的独特切入是把“预训练时的记忆构建”与“推理时的记忆传播”解耦：训练时引入一个感受野更强但仍因果的预填充器 $Q$，从观测序列并行生成全部目标记忆 $m'_{1:T}$；可部署的滑动窗口解码器 $P$ 消费平移后的 $m'_{0:T-1}$ 做下一词预测，同时产出自己的记忆 $m_{1:T}$，用一致性损失把 $m_t$ 对齐到 $m'_t$。推理时直接丢弃 $Q$，让 $P$ 用自己的记忆闭环。这与并发的 LRT（用并行精炼轮次训练隐循环）以及 SMT、Rec2PM（用模型生成的标签监督循环记忆）思路相近，但 Maglev 把该思想用于语言建模 Transformer：通过 K/V 通道注入平移记忆，缓存与计算曲线与普通滑动窗口注意力完全一致，并支持 $Q$、$P$ 参数共享。
核心方法
直觉上，Maglev 是一个“教师—学生”方案：让一个能看到全部历史的强模型当老师，指出每个位置“合理的记忆应该长什么样”；再让一个只能看 512 token 窗口的学生模仿老师写记忆，学会之后学生就能自己给自己喂记忆，独立部署。技术路线上，Maglev 由两个耦合的 Transformer 组成。预填充器 $Q$ 采用 SLSL 层模式（滑动窗口与全注意力交错，滑窗 512、全注意力覆盖 2048 训练上下文），一次因果并行前向产出全部目标记忆 $m'_{1:T}$。解码器 $P$ 采用 SSSS 模式（每层都是 $W=512$ 的滑动窗口），把平移一格的记忆目标通过“循环 K/V 注入”混入自己的注意力缓存，输出下一词预测并写下记忆 $m_{1:T}$。训练损失是下一词交叉熵加上一致性项 $\lambda\|m_t - m'_t\|^2/\sqrt{d}$。推理时 $Q$ 被丢弃，$P$ 递归地消费自己上一步的记忆，缓存大小与普通滑动窗口注意力完全一致。
核心创新是“提升平行训练”（lifted parallel training）：把长度为 $T$ 的顺序展开替换为两次序列并行的前向。关键洞察是——虽然 $P$ 在推理时必须自食其果，但训练时可以让一个更强的因果模型 $Q$ 一次并行前向就把整条记忆轨迹构造出来，再用一致性损失教会 $P$ 复现它；这样 $P$ 在“消费平移记忆”通道上见过的输入分布与推理时自喂的分布完全一致，闭环因此稳定。与已有方法的本质区别有三点。其一，记忆注入方式：把平移记忆经共享投影变成 K/V 特征，混入普通滑窗缓存的既有条目，不占额外序列位置，因此部署模型的缓存与注意力成本和 SWA 完全相同，而 RMT、Block-Recurrent 等都要靠专门的记忆 token。其二，表达力：每步记忆更新走完解码器全部 20 层的非线性计算，而 Mamba 等的更新是逐层局部的线性/仿射运算。其三，$Q$ 和 $P$ 可共享绝大多数参数（仅解码路径用独立的残差缩放），一致性信号直接绑在同一个 Transformer 上，而非完全独立的师生系统。
方法步骤详情
训练分两步。第一步预填充：$Q_\phi(x_{1:T}) = m'_{1:T}$，$Q$ 用 SLSL 模式在 2048 上下文内做一次因果并行前向，所有 Transformer 块加最终 RMSNorm 后的状态即为目标记忆。第二步解码：$P_\theta(x_{1:T}, m'_{0:T-1}) = m_{1:T}$，其中 $m'_0 = 0$。在解码器第 $\ell$ 层，共享投影把平移记忆映射为循环 K/V 特征：$k^{\mathrm{rec}}_t = \mathrm{RMSNorm}(W^{\mathrm{rec}}_k m'_{t-1})$、$v^{\mathrm{rec}}_t = W^{\mathrm{rec}}_v m'_{t-1}$，再由逐层门控与本地特征混合：$\bar{k}_t = g^{\mathrm{loc}} \odot k_t + g^{\mathrm{rec}} \odot k^{\mathrm{rec}}_t$（$\bar{v}$ 同理），其中 $g = 2\sigma(\cdot)$ 保证零预激活时门控呈中性值 1。滑窗注意力在混合后最近 $W$ 个 K/V 上进行，启用 RoPE 时对两类 key 都先旋转。输出状态 $m_t$ 同时送入 LM 头 $p(x_{t+1}|x_{\le t}) = \mathrm{softmax}(15\tanh(Um_t/15))$。总损失 $\mathcal{L}(\Theta) = \frac{1}{|I|}\sum_{t\in I}[\mathrm{CE} + \lambda\|m_t - m'_t\|^2/\sqrt{d}]$。推理时丢弃 $Q$，按 $m_t = P_\theta(x_{t-W+1:t}, M_{t-1})$ 递归运行并把 $m_t$ 追加进循环窗口。
技术新颖性
技术新颖性可从四个维度评估。架构层面：现有循环记忆方案要么用专门的记忆 token（RMT、Block-Recurrent），要么压缩旧段激活（Transformer-XL、Infini-attention），要么把状态藏在注意力机制之外（Mamba），Maglev 则让每个 token 的解码器最终状态直接充当记忆——通过它自己的普通 K/V 条目——循环因此不需要任何专用缓存结构。训练层面：非线性循环的并行化是公认难题，Maglev 用“辅助轨迹 + 一致性损失”绕开而非求解它，监督信号显式作用在平移记忆通道上，且师生可共享参数。表达力层面：记忆更新穿越全部 20 层 Transformer，比线性注意力/SSM 的局部线性更新强得多，还带来 looped-Transformer 式的额外有效深度而不增加迭代次数。工程层面：预填充器只是训练机制而非固定架构——任何能并行产出 $m'_t$ 的因果模型都可以当 $Q$，甚至可以换成预训练好的大模型来提供记忆目标，为蒸馏式压缩打开空间。
实验结果
实验在 nanochat 预训练栈进行：$L=20$ 层、宽度 $d=1280$、10 个注意力头（头维度 128）、上下文 2048、窗口 $W=512$，非嵌入参数约 435M，训练 43.52B token（不含嵌入时约为 100× Chinchilla 预算），MuonAdamW 配方。核心结果（Table 1）：SSSS 滑窗基线 FW BPB 0.7413、下游平均 54.1；共享参数 Maglev（$\lambda=0.1$）达 0.7295 BPB、56.2 平均分；分离参数 Maglev（$\lambda=1$）最优，BPB 0.7251、平均 56.4——在推理只用滑窗的情况下同时超过 SLSL 全/滑混合基线（0.7373、54.5）和对应的 LRT（SSSS 0.7331、55.0；SLSL 0.7292、55.9）。分项上最亮眼的是 BoolQ：64.0 对基线 56.2（+7.8）；LAMBADA 困惑度从 8.54 降至 8.06。一个反直觉发现是 $\lambda$ 并非单调：$\lambda=1$ 时共享模型明显受损（平均 53.9，BoolQ 跌至 51.1），分离模型却受益——说明共享参数时过强的一致性约束会挤压解码器自身表示，而 $Q$ 有独立容量时更强的目标才成为有效训练信号。
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| FineWeb-Edu 验证集语言建模 | BPB（越低越好） | 0.7251（Maglev 分离参数 λ=1） | 0.7413（SSSS 滑动窗口注意力） | 降低 0.0162（约 2.2%） |
| FineWeb-Edu 验证集语言建模 | BPB（越低越好） | 0.7295（Maglev 共享参数 λ=0.1） | 0.7331（LRT SSSS） | 降低 0.0036 |
| 8 项下游基准平均（PIQA/HellaSwag/WinoGrande/ARC-E/ARC-C/SIQA/BoolQ/LAMBADA） | 平均准确率（%） | 56.4（分离参数 λ=1） | 54.1（SWA SSSS）；54.5（SLSL 全注意力混合） | +2.3 / +1.9 |
| BoolQ | 准确率（%） | 64.0 | 56.2（SWA SSSS） | +7.8 |
| LAMBADA | 困惑度 PPL（越低越好） | 8.06 | 8.54（SWA SSSS） | 降低 0.48 |
| LAMBADA | 准确率（%） | 47.4（分离 λ=1）；48.7（分离 λ=0.1） | 46.2（SWA SSSS） | +1.2 ~ +2.5 |
局限与改进
作者明确承认这是受算力限制的初步研究而非该设计空间的定论：规模单一（仅 435M / 43.52B token），scaling 行为未知；递归注入在部署时缺乏高效 kernel；$Q$ 目前从零与 $P$ 联合训练，最优共享模式未探明。我自己的观察：其一，论文没有做任何长上下文检索/复制类评测（如 needle-in-a-haystack），“记忆能携带窗口外信息”这一核心主张只有 BPB 与标准基准的间接支撑；其二，训练需要 $Q$、$P$ 两次前向，论文未报告训练吞吐或 FLOPs 对比，实际成本可能接近普通 SWA 的两倍；其三，$\lambda$ 超参敏感且与参数共享方式强耦合（共享 λ=1 时 BoolQ 从 63.6 崩到 51.1），工程上调参成本高；其四，未见多随机种子与方差报告，1–2 个点的下游差距可能不显著；其五，评测上下文只有 2048，$W=512$ 的记忆窗口在如此短的序列里其长期记忆能力未必被充分激发。
独立分析的弱点
独立分析四个弱点并给出改进方向。第一，长程能力未被直接验证：全文没有任何合成长程任务（复制、检索、多跳推理）实验，建议补充 8K–64K 上下文的 needle 评测，并扫描 $W$ 从 512 到 4096 的记忆保持曲线。第二，训练成本翻倍：两次全深度前向使训练 FLOPs 约为单模型的两倍，可尝试训练后期冻结或降频 $Q$、用浅层/低秩 $Q$ 提供目标，研究目标质量与成本的 Pareto 前沿。第三，一致性与下一词目标的冲突：$\lambda$ 实为两个目标间的权衡旋钮，且效果依赖 Q/P 是否共享，可探索 $\lambda$ 退火、不确定性加权或按层差异化的一致性权重来缓解。第四，分布偏移隐患：训练时 $P$ 总是消费“教师质量”的记忆，推理时消费自己可能有偏差的记忆，长序列上误差可能累积——这是 teacher forcing 的经典暴露偏差问题，可用计划采样式混合训练（部分位置用 $P$ 自己的记忆）来鲁棒化。第五，只在因果语言建模上验证，该记忆机制在指令跟随、长文档问答、agent 场景中的价值完全未知。
未来方向
作者提出四个方向：一是研究 scaling 时预填充器强度与解码器容量的权衡，并为循环注入实现高效 kernel；二是让 $Q$ 不必从零训练——用一个预训练或轻量微调的语言模型提供记忆目标，只训练 $P$，把大模型表示蒸馏进紧凑的循环解码器；三是系统研究参数共享谱系（共享嵌入、MLP、注意力投影或部分层，同时保留任务专属组件）；四是比较其他记忆暴露方式：残差流注入、循环 token、交叉注意力、逐层记忆投影，刻画它们在表达力、稳定性、参数成本、缓存大小和推理吞吐上的取舍。基于本文成果还可延伸：把框架推广到多模态序列的固定内存压缩；与混合层栈结合——保留少量全注意力层做精确检索、其余用 Maglev 层做压缩记忆；对 $m_t$ 做探针分析，理解一致性目标下记忆究竟编码了什么信息。
复现评估
复现基础较好。实现完全建立在 Karpathy 的 nanochat 开源栈上（d20 配置、MuonAdamW、tokenizer 与评测脚本均现成），论文对架构的数学描述完整：门控注入、损失函数、LM 头的 $\mathrm{softmax}(15\tanh(\cdot/15))$ 温度设计都给出了公式，关键超参齐全——$W=512$、上下文 2048、$\lambda\in\{0.1, 1.0\}$、SLSL/SSSS 层模式、435M 非嵌入参数、每步 524,288 token。nanochat d20 的定位就是单机 8×GPU 数天量级可训完，有 GPU 资源的学术组能跑通全量复现，估算成本数千美元以内；算力受限者可按 nanochat 惯例缩小到 d16 级别先验证相对排序。不确定因素：论文是 preprint，正文未明确给出 Maglev 本身的官方代码链接，注入模块与一致性损失需自行实现，约几百行改动；下游评测沿用 nanochat 脚本，无额外障碍。综合难度中等：工程量小，但全量训练需要多卡环境。
论文图表
主结果表，汇总 8 个模型在 43.52B token 训练后的表现：FineWeb-Edu 验证 BPB、LAMBADA 困惑度与准确率，以及 PIQA、HellaSwag、WinoGrande、ARC-Easy、ARC-Challenge、SocialIQA、BoolQ 共 7 项下游基准和平均分。基线包括 SLSL 全/滑混合 Transformer（BPB 0.7373、平均 54.5）、SSSS 滑窗（0.7413、54.1）及两种层模式的 LRT（55.9/55.0）；Maglev 四个变体中，分离参数 λ=1 最优（BPB 0.7251、平均 56.4），共享参数 λ=0.1 次之（0.7295、56.2），而共享 λ=1 明显退化（53.9，BoolQ 仅 51.1）。
论文所有定量结论的出处：既证明 Maglev 超过滑窗基线、LRT 甚至全注意力混合基线，也完整暴露了 λ 与参数共享的交互效应，是评估与批评本文贡献的核心证据。