不动点让终点足以代表整条轨迹
循环语言模型会反复应用同一组共享层。每一次额外循环都会增加成本:训练时要存储激活值,解码时要缓存键和值,预填充每个提示需要计算,强化学习中为每个 rollout 打分还需要重放。论文的出发点是:当循环状态接近不动点时,这些成本会缩小,因为不动点方程只涉及终点,而不涉及到达终点的路径。
论文用终端 KV 共享来检验这一说法,即只保留每个前缀 token 的终端键值对,而不是每次访问都保留一个缓存库。复用终端 KV 缓存几乎不改变 Huginn 的零样本 GSM8K 准确率,但在四次循环时会让 Ouro 的准确率崩塌。论文把这一差异归因于收敛性:Huginn 的状态接近不动点,Ouro 的状态则没有。这是论文对观察到的对比给出的解释,而不是一项独立的对照实验。
同一个不动点也为其余成本提供了捷径。只对最后几次循环做截断反向传播,就能近似平衡态梯度,而不需要隐式梯度所需的预热。一个被蒸馏来直接预测不动点的学生模型,可以在预填充时替代这些循环。在强化学习中,从保存的不动点状态为 rollout 打分,可以避免重放循环。每种捷径都消除了一个与深度相关的成本,因此每增加一次循环,节省就更多。
深度先验决定监督看到什么
深度先验是训练深度的分布。它之所以重要,是因为不同 token 在不同深度收敛,因此采样到的深度决定了训练时能否看到终端 KV 共享在推理时提供的已收敛前缀上下文。在单一深度上训练的模型在该深度预测最好,但会破坏 KV 共享。Huginn 固定的宽泛先验支持共享,但在目标深度上对监督的稀释程度超过了共享所需。
论文从预测反馈中学习先验,把监督拉向预测良好的深度,并加入熵项使其宽泛到足以支持共享。在预算约束下,预测反馈会偏向那些已经预测良好的深度,因此若不加干预,它会把监督集中在很窄的深度范围内,并可能削弱不动点的塑造。熵项抵消这一倾向;由于深度从 1 开始,更宽的分配主要向平均深度以上、朝最大深度方向延伸。
因此熵系数是平衡预测性能与内存效率的直接旋钮。在匹配的 M 规模扫描中,把系数从 0 提高到 0.1,平均深度仍接近 5,而深度 5 的概率质量从 67.4% 降到 14.0%。训练深度 5 上的 WikiText 困惑度从 12.89 升到 13.20,而在深度 32 上,系数 0.01 和 0.05 以 12.93 并列最低困惑度。
学习到的先验优于固定深度与 PLN 采样
论文比较了三种选择训练深度的方式。固定深度训练在其训练深度 R = 5 上最好,但在其他深度上急剧退化。PLN 采样在所有深度上表现都不错,但在 R = 5 上落后。学习到的先验介于两者之间,在 R = 5 上缩小差距,同时在各深度上保持稳定。
最清楚的后果出现在 KV 共享上。用每次访问 KV 和共享 KV 分别给每个检查点打分显示,在 1,024 token 预填充后,共享会把固定深度训练的 WikiText 续写困惑度从 12.19 提高到 16.71。PLN-5 和所有学习到的先验则保持或略微改善该指标。更高的熵并不会在 WikiText 和 LAMBADA 上单调改善共享,因此熵项并不是简单的越大越好。
在不复用 KV 时,系数 0 在 M 和 L 规模上的验证困惑度略低于 0.01(4.55 对 4.57,3.53 对 3.54),在 S 规模上则与 0.01 持平。但在终端 KV 共享下,它在 L 规模的 GSM8K 上落后,为 36.47 对 47.92。因此同一批检查点会因是否共享缓存而排名不同,这也是论文同时报告两种布局的原因。
正交注入移除状态中沿输入方向的分量
输入注入把同一个嵌入送入每一次循环,使不动点以输入为条件。在 Huginn 和 Parcae 使用的加性形式中,被携带的状态会加到注入的输入上。Parcae 使用取值在 (0,1) 的对角携带矩阵,因此携带状态可能含有沿注入输入方向的分量,从而放大或抵消注入。不动点最终保留的输入便取决于状态,论文报告这会损害性能。
正交注入通过把携带状态投影到注入输入的正交方向上来移除该分量。在输入非零且稳定项为零时,注入量沿输入方向的分量在每次循环中都恰好等于输入,包括在不动点处,而正交部分自由演化。正的稳定项使这一性质变为近似成立。由于投影的范数至多为 1,投影后的携带矩阵保持收缩性。
实际效果是每次循环都以相同强度接收输入。论文报告正交注入的困惑度低于 Huginn 和 Parcae 的注入方案,并且这一结果在 1 亿到 16 亿参数范围内成立。正交注入还去掉了 Parcae 在注入前对前奏输出施加的 RMSNorm;论文把该消融放在附录中。
各规模结果与缩小 3 倍的缓存
从 1 亿到 16 亿参数,学习到的先验和正交注入在每个规模上都分别相对 Huginn 的先验和现有注入方案降低了困惑度。最核心的对比在 16 亿参数:学习到的先验配合缩小 3 倍的 KV 缓存,可匹配使用完整缓存的固定深度训练的下游平均表现。
缓存缩减来自五次循环时的终端 KV 共享,它把 KV 缓存缩小 3 倍。论文报告共享不会降低任何学习到的先验的平均任务准确率,平均变化为 +0.25 到 +0.70 个百分点,而 R = 5 的固定训练损失 0.61 个百分点。这些配对数字来自 M 规模熵扫描,预填充和解码深度均为 5,因此描述的是该设定,而非所有规模。
论文还报告了另外两项捷径。蒸馏学生模型的预填充速度最高快 1.79 倍,强化学习更新从保存的 rollout 状态计算梯度,比通过重放轨迹反向传播快 2 倍。论文把蒸馏预填充描述为用下游准确率换取预填充延迟,且差距随规模增大,因此预填充加速并非无代价的改进。
这些结果的适用范围止于何处
证据覆盖最多 16 亿参数的模型,每个配置只用一个随机种子,且消融并未穷尽设计空间。更大规模和其他模型配置仍未经测试,论文也指出这些技术在真实应用中的可行性仍需进一步研究。因此规模对比支持的是关于已测试范围的结论,而不是关于循环语言模型整体的结论。
若干报告数字与特定协议绑定。共享惩罚是同一检查点内终端 KV 与完整 KV 困惑度的比值,而不是相对原生 MixVal 的差值。配对任务准确率使用每个任务 2,000 次配对自助抽样,任务权重相等,其区间量化的是评估不确定性,而不是训练种子波动或系数之间的差异。预训练未做文档级去污染。
熵扫描同样是单一规模的结果。M 规模的各次运行只在熵系数上不同,模型、数据、种子、学习率调度和控制器都相同,这使该对比干净但范围狭窄。关于更高熵会分散先验并降低共享惩罚的说法,应被理解为 M 规模的发现,而不是缩放规律。
最该先调的旋钮是熵系数
论文最具决策相关性的结果是一种权衡,而不是单一最优设置。提高熵系数会把学习到的先验向更晚的深度分散,从而降低 KV 共享惩罚,从系数 0 时的 +0.96% 降到系数 0.1 时的 -0.22%,但会把训练深度 5 上的困惑度从 12.89 提高到 13.20。该系数就是这一平衡的直接控制量。
共享 MixVal 的绝对困惑度在系数 0.01 时最低,为 4.48,与系数 0 和 0.05 相差在 0.01 以内。因此当目标是共享缓存质量而非最大深度分散时,0.01 是自然的起点。在 0.01 和 0.1 之间如何选择,取决于优先考虑训练深度困惑度,还是更晚深度上的共享惩罚。
第二个优先级是注入方案。正交注入是在已测试规模上相对现有方案降低困惑度的改动,且独立于深度先验的选择。对于决定先采用什么的读者,论文证据指向正交注入为规模支持最广的改动,而熵系数则是必须与目标缓存机制相匹配的设置。
本文引用的文献
Towards Looped Models Done Right, Part II: Rethinking at Fixed Points
从 1 亿到 16 亿参数,学习到的深度先验和正交输入注入在每个规模上都相对 Huginn 的先验和现有注入方案降低了困惑度;在 16 亿参数下,学习到的先验配合缩小 3 倍的 KV 缓存,可匹配使用完整缓存的固定深度训练的下游平均表现。
