记忆高效的循环 Transformer：在循环语言模型中解耦计算与显存 Memory-Efficient Looped Transformer: Decoupling Compute from Memory in Looped Language Models
用门控潜变量把循环Transformer的KV缓存压成常量，解耦推理深度与显存
前置知识
Transformer 与 KV 缓存
Transformer 自回归生成时，每一层的注意力都要为已生成的 token 缓存其 Key 和 Value 矩阵，避免重复计算。KV 缓存大小为 $O(N \times L \times d)$，其中 $N$ 是层数、$L$ 是序列长度、$d$ 是每头维度。
本文的核心就是改造 KV 缓存策略，必须先理解它在线生成时的作用和成本结构。
循环 Transformer (Looped Transformer)
把同一组 Transformer 层在时间维度上重复 $T$ 次，相当于增加每 token 的计算深度而不增加参数量。LoopLM/Ouro 是近期代表，在数学推理上能与近 2 倍大的普通 Transformer 持平。
MELT 正是要改进 LoopLM 的显存瓶颈，必须理解其权重共享但缓存累加的设计。
潜在推理 (Latent Reasoning)
模型不显式生成 "思考" token，而是在隐空间中多次迭代更新表示以完成推理。与 Chain-of-Thought 相比省去了输出长度，但需要在架构上提供额外的迭代计算能力。
MELT 的循环机制是潜在推理的工程实现，其显存问题直接限制了推理深度 $T$ 的可扩展性。
门控更新机制 (Gated Update)
类似 GRU 的更新方式：先用 sigmoid 门 $z_t$ 决定保留多少旧信息，再以 $(1-z_t)$ 写入新信息。形式上 $h_t = z_t \odot h_{t-1} + (1-z_t) \odot x_t$，可在固定大小下融合多步信息。
这是 MELT 替代 append-only KV 缓存的核心算子，决定了哪些历史信息被保留或覆盖。
知识蒸馏与表征对齐
用预训练大模型 (teacher) 监督小模型 (student) 的输出或中间表征。本文采用层级的 post-attention 表征对齐损失：$\|o_{\text{MELT}}^{(l,t)} - sg(o_{\text{LoopLM}}^{(l,t)})\|^2$。
MELT 的两阶段训练深度依赖冻结的 LoopLM 作为 teacher，必须理解 stop-gradient 和层级蒸馏的写法。
研究动机
循环 Transformer（如 LoopLM/Ouro）通过把同一组层循环 $T$ 次来增强推理，参数几乎不增加但每 token 的计算量线性提升。然而这种架构在自回归生成时，每经过一次循环就要给每一层追加一份新的 KV 条目，单层缓存尺寸随推理深度 $T$ 线性增长，整体显存占用为 $M_{\text{LoopLM}} \propto O(N \times L \times T)$。论文报告的具体数字是：生成 32k token 时 Ouro-1.4B-Thinking 的 KV 缓存高达 25.17 GB，模型权重仅 2.87 GB，KV 缓存占用了 90% 显存，约束了循环次数的可扩展性。已有方案如 MQA/GQA 跨层共享 (CLA) 或循环间压缩在短任务上有效，但在长思维链上会显著掉点。
本文的目标是提出 MELT (Memory-Efficient Looped Transformer)，把循环深度 $T$ 与 KV 缓存大小彻底解耦，使每层只维护一份固定大小的 KV 缓存（与 $T$ 无关），整体显存回到 $M_{\text{MELT}} \propto O(N \times L)$。在仅 130 GPU-hours 后训练的前提下，让 1.6B 参数的 MELT 在数学与通用推理基准上全面超过 Qwen3-1.7B、Gemma4-E2B 等同尺寸普通 Transformer，并逼近显存高出数倍的 Ouro-1.4B-Thinking。
与已有工作不同的是，已有循环 Transformer 的 KV 共享方案要么沿用 append-only 策略仅做局部压缩，要么用静态平均/EMA 聚合多步信息，缺乏可学习的、按 token 自适应的门控机制，且没有配套的从 LoopLM 平稳迁移的训练流程。本文独到之处是：(1) 用一个 latent state $h_t$ 加 element-wise 门控 $z_t$ 跨循环更新，并在该状态上重新投影出 K、V，从而保持 query-key 对齐；(2) 提出 chunk-wise 训练 + 插值过渡 + 注意力对齐蒸馏的两阶段后训练，使架构大幅改动后仍能继承 Ouro 的预训练成果。
核心方法
MELT 的核心直觉是 "不要把循环当成 KV 的 append，而要把循环当成一个固定大小的内存写入过程"。每一层维护一个与序列长度同长、与循环深度无关的 latent state $h_t \in \mathbb{R}^d$，循环 $T$ 次时该 state 通过门控被反复更新，K、V 始终由当前 $h_t$ 投影得到并写入 KV 缓存的同一行。这样每层每 token 的 KV 缓存就只占一份，而不是 $T$ 份。训练时由于 token 间的 KV 计算变成顺序依赖，作者设计了 chunk-wise 训练（段内并行、段间串行）来近似自回归推理，并用两阶段策略（插值过渡 + 注意力对齐蒸馏）从预训练的 LoopLM 平滑迁移过来。
与已有方法的本质区别有三层。第一，KV 缓存策略不同：LoopLM 是 append-only 的 $O(L \times T)$；CLA/MLA 是层间共享但不跨循环压缩；MELT 是单层单 token 缓存 + 门控跨循环更新。第二，状态表征不同：其他方法直接在 K、V 空间做聚合；MELT 显式引入 latent state $h$，把 "写什么" 和 "怎么检索" 解耦，从而保护 query-key 语义对齐。第三，训练机制不同：已有循环模型要么从头训练要么直接做架构替换；MELT 配套设计了 chunk-wise 训练 + 插值过渡（$KV = \alpha KV_{\text{MELT}} + (1-\alpha) KV_{\text{base}}$，$\alpha: 0 \to 1$）+ 冻结 teacher 的层级注意力对齐蒸馏，三者缺一不可（去掉 chunk-wise 训练会直接得到 0% 准确率）。
方法步骤详情
前向流程：算 $z_t = \sigma(x_t W_z + h_{t-1} U_z + b_z)$，按 $h_t = z_t \odot h_{t-1} + (1-z_t) \odot x_t$ 更新 latent state；$h_t$ 投影 $k_t = h_t W_K$、$v_t = h_t W_V$ 原地覆写该层 KV 缓存（不追加），单 token 缓存恒为 $d$；$q_t$ 与 $K_t, V_t$ 做 attention 过 FFN，循环 $T$ 次后 lmhead 输出。训练两阶段：第一阶段 chunk-wise（段内并行、段间串行）+ $KV$ 插值过渡 $KV = \alpha KV_{\text{MELT}} + (1-\alpha) KV_{\text{base}}$，$\alpha: 0 \to 1$，配 LoopLM 逐循环 KD 损失；第二阶段冻结 LoopLM teacher，加层级对齐损失 $\beta \sum_{l,t}\|o_{\text{MELT}}^{(l,t)} - sg(o_{\text{LoopLM}}^{(l,t)})\|^2$。
技术新颖性
技术新颖性体现在四个层面：(1) 显式 latent state + element-wise 门控的循环 KV 缓存结构，把 "状态演化" 与 "K/V 投影" 解耦，区别于直接修改 K、V 的早期工作；(2) 严格的常数显存保证 $M_{\text{MELT}} \propto O(N \times L)$，理论上证明了推理深度 $T$ 不再进入显存复杂度；(3) chunk-wise 训练方法显式建模了 "段内可并行、段间需串行" 的自回归结构，是对 LoopLM 训练范式的工程革新；(4) 两阶段后训练把 "架构替换" 转化为 "表征对齐" 问题，配套的插值过渡和层级注意力蒸馏是首次系统性地用于循环 Transformer 的架构改造。
实验结果
MELT-1.6B 在 10 个推理基准上全面跑赢同尺寸非循环模型。数学方向 AIME24 pass@1 达 $46.7$（超 Qwen3-1.7B 的 $43.1$、Gemma4-E2B 的 $40.6$、DeepSeek-R1-1.5B 的 $32.1$），MATH-500 准确率 $93.4\%$、OlympiadBench $64.7\%$ 均居首。通用推理上 HumanEval $81.7\%$ 反超 Ouro 的 $76.8\%$。显存方面，生成 32k token 时 MELT 总占 9.49 GB，比 Ouro 的 27.97 GB 节省 $2.95\times$，KV 缓存缩小 $4\times$。门控消融（Table 3）显示完整 element-wise 门控 AIME24 pass@1 取得 $44.8$，远优于 Mean ($29.0$)、EMA ($30.2$)、Last ($33.7$)、Single-gated ($34.4$)；组件消融（Table 4）表明去掉 chunk-wise 训练直接得到 $0.0\%$ 准确率。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| AIME24 | pass@1 | 46.7 | Ouro-1.4B-Thinking 50.2 / Qwen3-1.7B 43.1 | 超 Qwen3-1.7B 3.6 个点，逼近 Ouro |
| AIME25 | pass@1 | 33.3 | Ouro 36.7 / Qwen3-1.7B 33.1 | 与 Qwen3-1.7B 持平 |
| AIME26 | pass@1 | 41.0 | Ouro 44.0 / Qwen3-1.7B 31.7 | 超 Qwen3-1.7B 9.3 个点 |
| AMC23 | pass@1 | 80.2 | Ouro 81.2 / Qwen3-1.7B 79.2 / Gemma4-E2B 82.7 | 居中，略低于 Gemma4 |
| MATH-500 | accuracy | 93.4 | Ouro 94.4 / Qwen3-1.7B 90.6 / Gemma4-E2B 87.6 | 超非循环基线 2.8-6.0 个点 |
| OlympiadBench | accuracy | 64.7 | Ouro 67.5 / Qwen3-1.7B 63.5 | 超 Qwen3-1.7B 1.2 个点 |
| GPQA | accuracy | 42.6 | Ouro 40.8 / Qwen3-1.7B 37.3 / Qwen3.5-2B 45.1 | 反超 Ouro 1.8 个点 |
| MMLU-Red | accuracy | 74.2 | Ouro 74.2 / Gemma4-E2B 75.3 | 与父模型持平 |
| HumanEval | accuracy | 81.7 | Ouro 76.8 / Qwen3-1.7B 71.3 | 反超 Ouro 4.9 个点 |
| 32k token 生成总显存 | GB | 9.49 | Ouro 27.97 / Qwen3-1.7B 7.07 | 较 Ouro 节省 2.95×，接近 Qwen3 |
局限与改进
作者承认三点限制：(1) 循环次数 $T$ 在推理时固定，无法根据输入难度自适应分配思考深度，而 MELT 的常数显存实际上让动态深度扩展变得更容易，是未来工作；(2) 当前实现未引入 Grouped-Query Attention (GQA)，因此在 KV 缓存绝对量上仍比 Qwen3-1.7B 多 2.5 GB；(3) 训练时 KV 顺序依赖迫使 chunk-wise 训练，并行度受限。我们补充一点：MELT 仍依赖 LoopLM 的预训练权重作为初始化和 teacher，对没有 LoopLM 起点的新模型族落地成本较高；Table 1 显示在 HLE (Humanity's Last Exam) 这种极难推理上绝对值仅 $2.0\%$，说明深层语义推理能力的提升空间仍大。
独立分析的弱点
独立分析下有三个可改进的弱点。第一，循环次数 $T$ 静态化是性能天花板，若不解决则 MELT 仍受 Ouro 推理深度上限的拖累，改进方向是引入早退机制 (early exit) 或基于置信度的循环门。第二，HLE 上仅 $2.0\%$ 准确率说明 latent reasoning 的 "思考量" 仍不够，可考虑在 latent state 上引入更显式的 "思考步骤" 监督，或借鉴 test-time scaling 思路增加循环内反思机制。第三，论文报告 MELT 的门控值训练稳定但未分析不同循环步的注意力分布，未来可可视化 latent state 的可解释性，例如用 probe 检查 $h_t$ 是否编码了 Chain-of-Thought 的关键节点，从而把门控变成可解释的思考控制器。
未来方向
作者点明的方向包括 (1) 自适应循环深度，让 $T$ 根据输入动态调整；(2) 引入 GQA，把 KV 缓存再压一半，潜在关闭与 Qwen3 的最后 2.5 GB 差距；(3) 设计更可并行的训练策略，使 MELT 能扩展到更大模型和更长推理链。基于论文成果还可延伸：(4) 把 latent state $h_t$ 接到外部工具调用 (tool use)，形成 latent tool-use 框架；(5) 用 chunk-wise 训练作为通用顺序依赖模块的工程范式，推广到状态空间模型 (SSM) 等其他有 token 依赖的架构；(6) 将层级注意力对齐蒸馏迁移到其他架构替换场景（如 Mamba-to-Transformer 升级），把 "teacher 引导的架构迁移" 变成可复用的方法论。
复现评估
作者承诺 "所有代码和模型将很快发布"，但论文发表时仓库尚未公开（仅给出 Ouro-1.4B-Thinking 和 AceReason-1.1-SFT、OpenThoughts3 等公共数据）。复现所需算力较高：完整两阶段训练 130 小时 8 卡 H100 80GB（合计 1040 GPU-hours），消融实验还需额外算力（附录 D 提到）。实验细节较完整：评测统一用 LightEval v0.8.1、温度 1.0、top-p 0.7、最大生成长度 32k；训练超参与数据集在 Table 6 列出。整体复现难度中等偏上——架构清晰、数据公开、评测标准化是优势，但 GPU 预算和尚未上线的代码仓库是主要门槛。
论文图表
