循环模型正确之道（二）：在不动点处重新思考 Towards Looped Models Done Right, Part II: Rethinking at Fixed Points
不动点让循环语言模型在训练、解码、预填充与强化学习中跳过完整递归轨迹
前置知识
循环语言模型 (Looped LM)
将同一块共享参数的 Transformer 层重复应用 $R$ 次的网络结构，如 Huginn 的 1+2×R+1 布局：一个 prelude 块、两块共享核心迭代 $R$ 次、一个 coda 块，逻辑深度 $2+2R$ 个块但物理参数只有 4 个块。循环次数可动态调整，用测试时计算换取更强的潜在推理能力。
本文的全部讨论对象是循环语言模型，所有成本分析（训练激活、KV 缓存、预填充、RL 重放）都源于每次循环都要付费这一结构特性。
不动点 (Fixed Point)
满足 $z^\star = F_\theta(z^\star; x)$ 的递归状态。随着循环深度增加，状态逐渐趋近不动点（包括有限深度的近似平稳态），此时继续循环状态几乎不变。不动点的本质是端点可以代替整条轨迹。
这是论文所有效率捷径（截断反传、KV 共享、蒸馏预填充、RL 复用）共同的理论根基，也解释了为什么这些捷径对 Huginn 有效、对 Ouro 失效。
截断反向传播 (TBPTT)
训练时只对最后 $b$ 次递归计算梯度而非全部 $R$ 次，显存和计算量随 $b$ 而非 $R$ 缩放。其梯度在访问状态收敛到 $z^\star$ 时自动逼近平衡梯度，且从初始化起就与完整 BPTT 一样稳定，无需热身阶段。
论文证明截断本身就能塑造不动点（$b=2$ 时困惑度最多涨 2%），这解释了循环模型无需隐式梯度也能安全使用终端 KV 共享。
KV 缓存与终端共享
自回归解码时缓存每个前缀 token 的 key/value 以避免重算。终端 KV 共享只保留每个前缀 token 最后一次递归写入的终端 KV 库，后续所有递归都读取它，使缓存不再随循环深度线性增长，本文中缓存缩小 3 倍。
终端共享是解码阶段的关键省内存捷径，但它只对不动点塑造良好的模型无损，这直接引出本文对深度先验的改进。
隐式微分与 Neumann 级数
对不动点方程微分得到平衡敏感度 $D^\star = (I-J^\star)^{-1}B^\star$，其 Neumann 截断 $D_b = \sum_{k=0}^{b-1}(J^\star)^k B^\star$ 可只从端点估计梯度，无需展开整条轨迹，但要求固定点存在且 $J^\star$ 谱半径小于 1。
论文的 RL rollout 状态复用正是把 Neumann-(b-1) 循环 $v \leftarrow g + J_R^\top v$ 应用在保存的端点上，实现约 2 倍加速。
深度先验 (Depth Prior)
训练时采样递归次数 $R$ 的概率分布。Huginn 用固定的平移 Poisson–对数正态（PLN）分布，均值等于目标深度 $\bar R$。先验决定监督落在哪些深度上，直接影响目标深度的预测质量与不动点塑形（即 KV 共享兼容性）之间的权衡。
深度先验是本文两大改进之一：固定深度训练破坏 KV 共享，固定宽先验稀释目标深度的监督，学习的先验用预测反馈加熵正则平衡两者。
输入注入 (Input Injection)
每次递归把 prelude 输出的嵌入 $q_t$ 加到状态上的机制，使不动点保留输入信息，常见形式为 $\zeta_r^t = M h_r^t + q_t$。Parcae 用对角衰减矩阵 $\Lambda$（元素在 0 到 1 之间）控制状态 carryover 以稳定训练。
论文发现 carryover 中沿输入方向的分量会按符号放大或抵消每次注入的信号，正交注入（OrthoInj）把该分量投影掉，是本文第二个核心改进。
研究动机
循环语言模型把共享层块重复 $R$ 次，但典型实现中每次循环都要付费：训练要存储并反传 $R$ 层激活，解码的 KV 缓存随深度线性增长，预填充要对每个 prompt 逐层迭代，强化学习还要重放整条轨迹来给 rollout 打分。是否应该训练不动点在社区中一直有争议：一派（如 Ouro）固定小循环数、放弃不动点，另一派（如 Huginn）塑造不动点、循环更深。实证上这条分界线很残酷：对 Huginn 做终端 KV 共享几乎不掉 GSM8K 精度，而对 4 次循环的 Ouro 做同样操作会使其崩溃，差别就在于状态是否收敛到不动点。更糟的是当前训练方式未必产生正确的几何：固定深度 $R=5$ 训练的模型外推到 $R=64$ 时困惑度上升超过 25%，强制终端 KV 共享使其上升 21%，在 1.6B 规模下 GSM8K 从 50.6 崩到 21.2；Huginn 的固定宽先验虽支持共享，却稀释了目标深度的监督。此外，Huginn-Linear、Parcae-Decay 等注入方案中，carryover 状态沿输入方向的分量会按其符号放大或抵消每次注入的信号，使不动点保留的输入量依赖状态本身。
本文的目标是本文的目标是系统论证并兑现不动点让循环模型在每一个阶段都更便宜这一命题。具体分三层：概念层，证明不动点方程 $z^\star = F_\theta(z^\star; x)$ 只涉及端点不涉及路径，因此在不动点附近端点可以代替整条轨迹，由此统一解释两个已有捷径（截断反传、终端 KV 共享）并推导出两个新捷径（蒸馏预填充、RL rollout 状态复用），每个捷径都消除一项随深度增长的成本；训练层，改进塑造不动点的两个组件：把固定的 Poisson–对数正态深度先验换成从预测反馈学习的分类先验，用熵正则保持其宽度以支撑 KV 共享，并用正交投影移除注入方案中状态沿输入的分量，保证每次递归收到等强输入；实证层，在 100M、400M、1.6B 三个规模上验证：学习的先验在同等甚至 3 倍更小的 KV 缓存下匹配或超越固定深度训练与 Untied 基线，蒸馏学生在 8K 提示上预填充提速至 1.79 倍，RL 复用把每次更新的打分与反传时间砍半。
与已有工作不同的是，已有工作对不动点的用法要么是解方程：DEQ 直接求解不动点并用隐式梯度，但需要展开热身阶段才能稳定训练；要么是决定深度：稳定不动点、把监督状态设为吸引子、收敛即停、学习退出分布。本文的切入角度不同：把不动点当作可复用的端点，推导出一组贯穿训练、推理、后训练全流程的计算捷径，并反过来用这些捷径是否成立来诊断模型是否真正收敛（例如终端 KV 共享对 Huginn 几乎无损、对 Ouro 致命）。在训练侧，与 Ouro 和 RL-Halting 从循环状态学习深度分布不同，本文从预测损失反馈学习先验且以不动点塑形为目标，它们每步都要展开到最大深度而本文不需要；与固定深度训练和固定 PLN 先验的二元选择不同，学习的先验加熵正则给出了可调的中间地带。在注入侧，先前方案（Huginn-Linear、Parcae-Decay、DEQ-QKV）都没有约束状态沿输入的分量，本文首次指出该分量会按符号放大或抵消注入，并用正交投影移除之，同时保留 Parcae 的收缩保证。
核心方法
直觉上，不动点附近继续循环什么也不变，所以凡是沿着轨迹做的计算都可以换成只碰端点。论文先在第三节建立四个捷径：训练用截断 BPTT，只反传最后 $b$ 步，内存与计算随 $b$ 而非 $R$ 缩放，其梯度随状态收敛自动逼近平衡敏感度 $D^\star=(I-J^\star)^{-1}B^\star$；解码用终端 KV 共享，每个前缀 token 只留最后一次递归的 KV 库，Lemma 1 证明冻结前缀迭代与联合迭代收敛到同一 $h^\star$；预填充用蒸馏学生 $\hat z = S_\psi(e) \approx z_R^\theta$ 直接预测端点，之后一次教师递归即可建好所有 KV 库；RL 用保存的 rollout 端点 $z^R$ 配合 Neumann-(b-1) 循环 $v \leftarrow g + J_R^\top v$ 计算梯度，省去重放。第四节再改进训练本身：把深度先验参数化为可学习 logits $p_\phi = \mathrm{softmax}(\phi)$，用分离的交叉熵算奖励 $u = e^{-\ell}$、PopArt 式优势 $A_n$、熵项 $-\lambda_H H(p_\phi)$ 和预算惩罚 $\lambda_m(E[R]-\bar R)^2$ 更新；输入注入改为正交投影 $\zeta_r^t = Q_{q_t}\Lambda h_r^t + q_t$，其中 $Q_{q_t} = I - q_tq_t^\top/(\|q_t\|_2^2+\epsilon)$。
核心创新一是截断本身塑造不动点这一反直觉发现：固定深度 $R=5$ 加完整 BPTT 训练的模型外推到 $R=64$ 时 PPL 涨 25%、终端 KV 共享涨 21%；而只反传最后两步（$b=2$）的 TBPTT 分别只涨 2% 和 3%，即不需要显式的不动点目标或隐式梯度的两阶段热身，单一估计器全程稳定且自动产生支撑 KV 共享的不动点几何。在此基础上的学习先验把在哪些深度监督变成可学习的分类分布：逆困惑度奖励 $u=e^{-\ell}$ 拉高预测好的深度的概率，熵项 $H(p_\phi)$ 抵抗监督向单一深度坍缩、保持向 $R_{max}=64$ 的宽度（这正是终端 KV 共享所需要的），预算项把期望深度钉在 $\bar R$，因此 $\lambda_H$ 成为性能与内存效率的直接旋钮。核心创新二是 OrthoInj：Parcae 的 $\Lambda h_r^t$ 可能带有沿 $q_t$ 的分量，其符号决定这次递归的输入被放大还是抵消，导致不动点保留的输入量依赖状态本身；OrthoInj 用 $Q_{q_t}=I-q_tq_t^\top/(\|q_t\|_2^2+\epsilon)$ 精确移除该分量，使状态沿输入方向的分量每次恰好等于 $q_t$，同时 $\|Q_{q_t}\Lambda\|_2\le\|\Lambda\|_2<1$ 保持收缩性，正交部分自由演化。
方法步骤详情
模型遵循 Huginn 的 1+2×R+1 布局：prelude 块把输入嵌入为 $e=P_\theta(x)$，两块共享核心迭代 $R$ 次并每次递归注入 $e$，coda 块输出 logits；规模为 S/M/L（100M 配 21.5B、400M 配 85.9B、1.6B 配 343.6B tokens）。四个捷径的流程：第一步，预训练用 TBPTT，仅对最后 $b$ 次递归反传，梯度形如 $\tilde D_b=\sum_{k=0}^{b-1}\prod_j J_{R-j}B_{R-1-k}$，在访问状态 $z_r$ 上求雅可比，内存与计算随 $b$ 而非 $R$ 缩放；第二步，解码只缓存终端 KV 库（4 个物理层各一，而非 12 个逻辑层）；第三步，蒸馏学生为两块网络、从教师共享核心初始化，以归一化隐状态回归加下一 token KL（式 4）对着冻结教师的深度 R 端点在线训练，token 量为教师的四分之一，推理时学生预测 $\hat z$、再跑一次教师递归和 coda；第四步，RL 按 Algorithm 1：保存端点 $z^R$，对其做一次前向记录计算图 $G$，读出反传得 $g$，Neumann 循环 $b-1$ 次更新 $v\leftarrow g+J_R^\top v$（复用 $G$），最后把 $v$ 反传到 $\theta$，更新成本与 $R$ 无关。深度先验学习按 Algorithm 2：每个 microbatch 采 $r\sim p_{\phi_n}$，CE 梯度只通过最后 $\min(r,b_{max})$ 次递归累积；EMA 统计中心化并重标奖励得优势 $A_i=s_{n-1}(u_i-\hat u_{n-1})$；按 $J_n=-E[\mathrm{stopgrad}(A_i)\log p_\phi(r)]-\lambda_HH(p_\phi)+\lambda_m(E_p[R]-\bar R)^2$ 用 Adam 更新 $\phi$。评估统一在 R=5 加终端 KV 共享。
技术新颖性
新颖性体现在四处。第一，TBPTT 被重新解释为不动点塑形器而非单纯省显存技巧：论文用固定深度训练作对照，证明截断是终端 KV 共享近无损的必要条件，且无需 DEQ 的展开热身加隐式梯度两阶段训练，单一估计器端到端稳定。第二，深度先验从固定分布变成从预测反馈学习：与 Ouro 和 RL-Halting 的区别在于它们从循环状态学退出分布且每步展开到最大深度、不以不动点塑形为目标；本文的先验学习不增加任何前向开销（奖励来自分离的 CE），熵项专门为终端 KV 共享保留宽度，先验可从 PLN-5 初始化继续学习而非搜索分布族。第三，OrthoInj 首次识别并消除注入方案中的状态-输入对齐分量，比 Parcae 的对角衰减更进一步且保留其收缩保证，还顺手去掉了 prelude 的 RMSNorm，两处增益可叠加。第四，两个新捷径都属首次：非循环学生只预测端点的蒸馏预填充（已有循环模型蒸馏的学生仍是循环的，prompt-KV 预测器则要预测每次递归的 KV，而本文学生只预测端点、一次教师递归即建全部 KV 库）；从保存的 rollout 端点直接打分的 RL 更新（重算控制实验显示其与复用的梯度余弦相似度达 0.997-0.999，干净地分离了状态复用与 Neumann 近似各自的贡献）。
实验结果
五个实验块逐一分析。深度先验（表 2）：终端 KV 共享下，学习先验（λH=0.01）在 S/M/L 的验证 PPL 为 5.67/3.93/3.07，全面低于 Fixed PLN-5 的 5.73/4.00/3.10，训练 FLOPs 至多多 1.6%；下游平均 41.65/49.07/60.00 全超 PLN-5；对同参数同缓存的 Untied 4，下游平均领先从 S 的 2.2 点扩到 L 的 8.9 点（L 的 GSM8K 47.92 vs 7.20，领先 40.7 点）；以 3 倍更小 KV 缓存在 L 上反超固定深度训练的全缓存下游平均（60.00 vs 56.90），距 3 倍参数的 Untied 12（61.47）仅 1.5 点。熵项在 L 关键：去掉后 GSM8K 从 47.92 掉到 36.47。注入（表 3）：OrthoInj 在所有规模取得最低训练末损失、最低验证 PPL（5.67/3.94/3.08）和最高下游平均（41.65/49.16/59.26），比最强基线 Parcae-Decay 的 PPL 低 0.3-1.4%。RL 复用（表 4）：Neumann-4 把打分加反传时间从 2.77s 降到 1.39s（GSM8K，1.99 倍）、1.03s 降到 0.51s（MBPP+，2.02 倍）；GSM8K pass@1 61.65% vs 完整 BPTT 63.20%，差距在 1.6 点运行波动内，均从 RL 前的 40.63% 大幅提升；增益迁移到未训练的 MATH500（+3.9-5.4 点）。蒸馏预填充（表 5）：8K 提示上 L 延迟 830.3ms vs 教师 1482.1ms（1.79 倍提速），下游平均落后 1.0/2.3/4.5 点（S/M/L），同延迟下比教师 R=2 高 0.5-0.9 点；单次递归再快 12-23% 但下游损失是蒸馏的 2.7-3.5 倍。截断反传：b=2 时外推与 KV 共享的 PPL 涨幅上限为 2% 和 3%，b≥4 与完整 BPTT 差距在 1% 内。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| GSM8K 强化学习（L 模型，Dr. GRPO，R=6） | pass@1 (%) 与每次更新打分+反传时间 | 61.65%，1.39s（Reuse + Neumann-4） | 完整 BPTT：63.20%，2.77s；RL 前 40.63% | 时间 1.99 倍提速，pass@1 差距 1.6 点在运行间波动内 |
| WikiText-103 语言建模（L 规模，R=5，终端 KV 共享） | PPL（越低越好） | 8.68（学习先验 λH=0.01） | 8.76（Fixed PLN-5）；8.92（Fixed R=5） | 低约 0.9%，训练 FLOPs 仅多 1.6% |
| 下游 9 项任务平均（L 规模） | 平均准确率 (%) | 60.00（3 倍更小 KV 缓存） | 56.90（Fixed R=5 全缓存）；61.47（Untied 12，3 倍参数与非嵌入参数） | 反超固定深度训练 3.1 点；距 Untied 12 仅 1.5 点 |
| 预填充延迟（L 规模，8K 提示，batch 8） | 平均延迟 ms | 830.3（蒸馏学生，教师初始化） | 1482.1（教师 R=5）；848.6（教师 R=2） | 1.79 倍提速，下游平均仅落后教师 4.5 点，同延迟下高于 R=2 约 0.5-0.9 点 |
| MBPP+ 强化学习后 pass@1 | accuracy (%) | 39.38（Reuse + Neumann-4），0.51s | 完整 BPTT：40.88%，1.03s；RL 前 30.88% | 时间 2.02 倍提速，精度差 1.5 点 |
| 验证集困惑度（OrthoInj 对比注入方案，S/M/L） | PPL（越低越好） | 5.67 / 3.94 / 3.08 | Parcae-Decay：5.75 / 3.96 / 3.09 | 所有规模最低，低 0.3-1.4%，且训练末损失、下游平均均最优 |
局限与改进
作者承认的局限：实验只做到 1.6B 参数且每个配置只用一个随机种子；消融没有覆盖完整设计空间；更大规模和其他模型配置未测试；这些技术在真实应用中的可行性需要进一步研究。我的补充观察：第一，蒸馏预填充的质量差距随规模扩大（S 落后 1.0 点、M 2.3 点、L 4.5 点），两块固定学生可能成为容量瓶颈，且 L 的 GSM8K 从 47.92 掉到 32.90，生成任务的退化比选择题更严重。第二，学习先验对超参敏感：熵系数 $\lambda_H$ 从 0.01 降到 0 使 L 的 GSM8K 损失 11.45 点（47.92 到 36.47），说明预测质量与 KV 共享兼容性之间的平衡相当脆弱，需要额外调参。第三，Lemma 1 的理论保证依赖映射 $\kappa$-收缩的假设，实际 Transformer 的收缩性只有经验证据支撑，未给出谱半径的度量。第四，RL 实验限于 GSM8K 400 题和 MBPP+ 272 题的小规模设定，rollout 生成本身仍占训练时间大头，复用只加速打分与反传部分。第五，终端 KV 共享下预填充收益在中等深度就饱和（附录 C.6 指出预 fill 与解码可在不同深度预算运行），深层循环的收益集中在解码侧。
独立分析的弱点
独立分析的弱点及改进方向：第一，学生容量问题——蒸馏差距随规模从 1.0 点扩大到 4.5 点，固定两块学生显然不足；改进方向是作者提到的混合专家层学生、更大学生网络，或把蒸馏 token 预算按层与任务难度自适应分配。第二，学习先验的奖励设计单一（逆困惑度），与下游生成任务（GSM8K 等）可能错位，M 规模下游平均仅比 PLN-5 高 0.01 点也提示信号偏弱；改进方向是把下游任务反馈纳入先验学习，或在不同训练阶段动态调整熵预算。第三，RL 复用的收益上界受限于打分加反传占总时间的比例，rollout 生成为主时收益有限；改进方向是把端点复用也用于生成阶段的 warm-start（附录 E.4 已有相关讨论）。第四，每配置单种子使 1 点级别的结论不可靠，例如 M 规模学习先验与 PLN-5 的下游平均差距在噪声范围内；改进方向是至少 3 个种子报告均值与方差。第五，终端 KV 共享可能在注意力模式上引入系统性偏差（长程检索、多跳推理任务），论文未按任务类型分析共享损失的分布；改进方向是建立共享安全性随任务特征变化的预测指标，例如用收敛深度分布直接估计。第六，OrthoInj 的投影矩阵随 $q_t$ 变化，虽计算量小但破坏了注入的平移不变性，对极长序列的影响未评估。
未来方向
作者明确提出的方向：更好的蒸馏目标函数；混合专家（MoE）层等学生架构；蒸馏 token 预算的分配策略；学生作为独立语言模型或投机解码 drafter 的可能性；循环状态精炼与基于流的语言建模之间的联系；更大规模与更多配置的验证。基于本文成果可以延伸的方向：第一，把学习深度先验与测试时自适应循环数（halting、提前退出）结合，让先验随输入难度动态调整而非固定分布。第二，把 RL rollout 端点复用推广到多轮对话 RLHF 和长链推理（如过程奖励模型逐步打分），端点复用的收益会随轮数和深度放大。第三，正交注入的去对齐思想可迁移到其他带状态传递的架构，如线性注意力与状态空间模型的输入条件化。第四，蒸馏学生天然适合作为投机解码的 drafter，因为它正是预测教师深循环结果的廉价近似，论文已把这一点列为开放问题。第五，用固定点诊断工具（收敛深度分布、雅可比谱半径）作为循环模型的健康指标，指导推理引擎自动在完整缓存与终端共享之间切换。第六，端点复用加 warm-start 求解器（附录 E.4）可能把相关输入之间的预填充成本进一步压低。
复现评估
复现条件评估：代码已在 GitHub 开源（github.com/ifm-ai/xllm-loop），附录 B 给出完整训练配方、学习率扫描与实现细节，附录 D 给出全部评估协议，透明度较高。训练数据为公开的 TxT360 多语言混合，tokenizer 为公开的 Jais64k。算力门槛分级明显：S 规模 100M 参数、21.5B tokens（约 7.3×10^19 FLOPs）单机多卡可复现；M 规模约 9.2×10^20 FLOPs 需要中小集群；L 规模 1.6B 参数、343.6B tokens（约 1.27×10^21 FLOPs）需要相当规模的多节点训练，普通实验室难以完整复现。RL 实验只在 GSM8K 400 题或 MBPP+ 272 题上微调已有模型，成本较低、最易复现。主要复现风险：每配置单种子使细粒度数字有方差；学习先验引入多个新超参（$\nu, \lambda_H, \lambda_m, \bar R, b_{max}$）和 EMA 统计、PopArt 重标定等工程细节；终端 KV 共享与蒸馏学生需要修改推理引擎的缓存逻辑。总体判断：S 规模复现难度中等，核心结论（截断塑造不动点、KV 共享近无损、注入方案对比）在 S 规模即可独立验证，适合作为入门复现起点。
论文图表
展示 Huginn 中 token 收敛的异质性：(a) 相对隐状态变化 $\rho_{r,t}$ 的热图（收敛阈值 $\tau=2\%$），不同 token 在不同深度收敛；(b) 局部顺序反转的例子，位置 13 和 10 的 token 先于位置 59 收敛；(c) 位置与收敛深度的关系，中位数线显示两者只有弱相关（Spearman 相关 0.14）。
该图论证了收敛深度因 token 而异且与位置弱相关，说明序列并行训练时采样深度决定监督是否看到已收敛的前缀上下文，是采样深度先验与终端 KV 共享之间联系的经验基础。