论文

ELR:消除循环语言模型的单轮精度损失

One Model to Loop Them All: Elastic Latent Recurrence for Language Models

模型架构模型推理递归架构测试时计算扩展KV Cache

摘要

潜在循环语言模型重复共享层块,可通过循环次数使同一检查点在推理计算和准确率间权衡。但单轮通常明显落后于同规模非循环模型,低延迟时另用非循环模型更好。弹性潜在循环ELR消除这一代价:同数据训练时,单轮服务成本等于同规模非循环模型,准确率相近,增加循环换取更多计算与准确率。三项选择是带可学习步长的门控状态更新、每次优化器更新内混合循环数,以及按循环索引和预算条件化每次运行。前两项保单轮准确率,第三项降低额外循环损失。可在预训练后期引入混合深度降低成本,在切换时加入零初始化条件,可比从头混合训练更低损失。MoE后续循环复用首轮专家路由使混合深度更新时间缩短11.7%,准确率无可测变化。推理仅保留末轮注意力KV缓存,内存维持单轮水平且保留大部分额外循环收益。4B稠密与30B-A3B MoE混合模型十任务均分中,ELR单轮相对匹配基线分别低0.52点、高1.04点,四轮增加7.25和4.75点。后训练得到Koi-4B,同一检查点在Artificial Analysis Intelligence Index单轮估计18分,与Qwen3.5-4B相当,四轮29分。

ELR:消除循环语言模型的单轮精度损失:论文配图
图1:弹性潜在循环。共享块ℛθ在输入ℐ和输出𝒪间运行k次。每轮i重新注入h_in,利用从s̃_i计算的门按自适应delta规则(式2)更新状态;k=1时每独立层恰运行一次。