循环语言模型不再交一次循环的税:同一份权重一次循环追平非循环基线,四次循环多出7.25分
核心概要
循环语言模型过去在一次循环时明显落后同规模非循环模型,现在同一份权重在一次循环下与同数据训练的非循环基线相当,四次循环在十任务平均上再涨7.25分。
深度剖析
循环语言模型可以只保留一份权重,用循环次数换取推理算力与精度,而不必在一次循环时落后于同规模的非循环模型。 此前的弹性深度模型在一次循环时明显低于同深度非循环模型,Huginn 在 ARC-Easy 上以 34.9 对 46.4 落后,Parcae 在一次循环低于 Transformer 基线、八次循环才反超,这一差距被称为一次循环税。 在 4B 稠密与 30B-A3B 混合专家两个规模上,ELR 与同数据、同参数、同训练步数的非循环基线对比,一次循环的十任务平均分别为 65.13 对 65.65 和 66.08 对 65.04,四次循环分别再涨 7.25 与 4.75 分。
让一次循环端点存活下来的关键,是状态更新方式与循环次数在优化器更新内部的混合方式,而不是把循环范围缩小。 按步采样单一循环次数、跨步变化,即使三分之一的时间抽到一次循环,也保不住一次循环端点;在一到五次循环上均匀采样同样让一次循环端点落后于基线。 0.5B 消融中,同步加权采样的一次循环损失为 1.6871,加入 15% 一次循环秩覆盖后降到 1.6830,独立秩采样为 1.6774,而基线为 1.6852。
循环次数可以在预训练后段才引入,条件化头在切换点零初始化,使这一晚引入的日程反而优于从头混合深度训练。 从头做混合深度训练的训练算力约为基线的 1.62 倍,而在最后四分之一才切换只需约 1.16 倍;在 0.5B 消融中,75% 切换点把估计的额外训练算力从约 79% 降到约 20%。 在 0.5B 上,25%、50%、75% 三个切换点分别保留五次循环相对基线增益的 97%、90%、84%,75% 切换时五次循环损失只增加 0.009;条件化后五次循环验证损失为 1.6147,低于从头混合深度训练的 1.6273。
循环混合专家模型可以复用第一次循环的路由结果,推理时只保留最后一次循环的注意力 KV 缓存,从而压低训练与显存开销。 每个循环都重新执行路由器会重复计算,标准推理则为每个循环各存一份注意力 KV 缓存,显存随深度增长。 跨循环路由复用在混合深度阶段把更新耗时降低 11.7%,四次循环验证损失只上升 0.0038;Koi-4B 上最后一次循环 KV 缓存把每 128K token 请求的注意力 KV 显存从 7.5 GiB 压到 3.0 GiB,四次循环得分相对逐循环缓存下降 0% 到 3%。
启示与展望
这套配方在 4B 稠密混合模型与 30B-A3B 混合专家模型上验证,Koi-4B 走完预训练、长上下文中间训练、监督微调与可验证奖励强化学习全流程,评测在自建评测框架内完成,与公开榜单数字可能不同。Koi-4B 没有同数据训练的非循环基线,一次循环精度是与同规模开源模型比较。作者提出的下一步是把 ELR 与按请求或按 token 的自适应深度结合,并对混合专家版本做后训练。
增益在第二次循环之后递减,作者认为这与后续循环出现在更少训练样本中一致,但没有测试重新加权循环次数分布能否找回这部分增益。条件化在中间深度上缩小了与固定深度专才模型的差距,但在两端没有,作者尚未分离原因。服务设置固定每次请求的循环次数,按请求或按 token 的自适应深度会与预算条件化相互牵制。最后一次循环 KV 缓存的统计只覆盖单次运行内的问题间差异,不覆盖运行间差异,且 GDPval 因超时任务在两个缓存策略下覆盖不同任务子集而未纳入比较。
