Thinking with Looped Flows
Ayhan Suleymanzade, Chanhyuk Lee, Floor Eijkelboom, Nicholas M. Boffi, İsmail İlkan Ceylan, Jinwoo Kim
cs.LG, cs.AI
2026-09-11
循环流用递减噪声局部训练循环状态,同架构ARC-AGI-1达58.8%、ARC-AGI-2达12.2%,均超TRM。
循环模型靠共享参数反复更新隐状态,推理时把有效深度拉长。HRM、TRM、FPRM 这一路已经能用几百万参数在 Sudoku、迷宫和 ARC-AGI 上做出像样结果。训练才是卡点:完整 BPTT 沿时间反传,内存随步数线性涨,梯度还容易爆或消失。实践里每一步都会 stop-gradient,后面步的损失到不了前面的状态。早期更新学不会为后续计算铺路,循环就不稳:不收敛,或者收敛到解不对的假吸引子。
流模型把「从噪声走到答案」拆成一串局部去噪任务,每一步独立监督,不用沿整条轨迹反传。但这条轨迹本身不擅长某些串行计算。Looped Flows 把两件事接起来:用局部去噪当课程,诱导出能跨步复用的循环状态。
去噪器做成有状态的。每一步同时干两件事:根据当前带噪样本预测干净答案,并更新循环状态 z。推理时按预测还原概率流的速度,用 ODE 或 SDE 推进一步。
训练时从 [0,1] 抽 k+1 个时间再排序,得到递减的噪声水平 (1-t)。同一道题、同一份噪声、同一个答案在这 k 步里共用。每一步对干净答案算交叉熵,步与步之间切断梯度。共享噪声让相邻去噪任务对得上号,循环状态才有理由把前面算过的特征带下去。k 取 16,和 TRM 训练时的最大步数对齐。
噪声水平往下走,损失会单调变小,后面几步可能已经没什么可学。于是加了一个 ACT 头:二分类判断这一步是不是已经答对,分数过半就提前停,避免在已饱和的步上过拟合。
推理可以从比训练更细的时间网格积分,步数就是额外算力。也可以从不同初始噪声各走一条轨迹,用 ACT 头给出的成功概率做 best-Q 集成。γ>0 的 SDE 积分比纯 Euler 更稳一点。
架构直接改 TRM。Sudoku 用 5M 的 MLP-Mixer,其余任务 7M 的非因果 Transformer。去噪器内部仍是 TRM 那套双状态 (h, ℓ):预测从 h 解码,ℓ 负责更新。一次去噪调用内部重复三次 TRM 式循环,前两次不算梯度,只在最后一次反传;Sudoku 内层 6 次,其他任务 4 次。只多了噪声插值的投影和时间嵌入,大约多 0.27M 参数,Sudoku 从 5.03M 涨到 5.30M。
六个基准,四个单解、两个多解。单轨迹对照如下。
| 方法 | Sudoku-Extreme | Maze-Hard | ARC-AGI-1 | ARC-AGI-2 |
| TRM(同架构) | 87.4 | 85.3 | 44.6 | 7.8 |
| FPRM | 94.2 | 87.0 | 47.5 | 6.2 |
| GRAM | - | - | 52.0 | 11.1 |
| Looped Flows | 97.9 | 86.7 | 58.8 | 12.2 |
迷宫上没超过 FPRM 的 87.0,差 0.3 个点。ARC 上领先最明显,相对 TRM 分别多 14.2 和 4.4 个点。5 条轨迹的 best-Q 集成把 Sudoku 推到 99.3,ARC-AGI-1 到 59.5;EqR 用 128 条才到 Sudoku 99.8。ARC-AGI-2 的均值停在 12.2,只是方差从 ±1.9 收到 ±0.9。
Sudoku 上把积分步数从 8 加到 128,准确率从 74.5% 涨到 97.9%,32 步就超过 GRAM。约 6.5 万道测试题上,TRM 失败率 12.6%,其中 88.3% 不收敛、11.7% 掉进假吸引子。Looped Flows 救回不收敛案例的 89.9%、假吸引子的 98.0%,合计消化掉 TRM 失败的 90.9%。
多解任务每题抽 20 次。8×8 N-Queens 准确率 99.9%、覆盖 91.4%,GRAM 是 99.7% / 90.3%。10×10 上差距拉开:准确率 94.4% 对 89.7%,覆盖 61.5% 对 57.5%。图着色冲突更少,8 顶点 0.7 对 GRAM 的 2.7,10 顶点 1.0 对 3.3,覆盖也略高。
消融最狠的一刀是同时去掉时间和插值:ARC-AGI-1 从 58.8 掉到 43.6,几乎退回 TRM。只去掉递减噪声到 51.6,去掉噪声共享到 56.4。无循环的纯流模型 FLM 会在训练集上背答案;自条件化能改善泛化,跨噪声水平滚循环状态仍然最好。SDE 相对 ODE 在各任务上有 0.2 到 1.3 个点的小幅增益。
对做循环推理的人,这篇给了一个可操作的训练改法:别再让每一步都从同一份干净输入重复预测,改成沿着噪声课程走。同架构、同量级参数,ARC-AGI-1 多了十四个点。循环稳定性的改善比分数更有信息量,TRM 那类不收敛和假吸引子,大部分能被这条课程消化。
多解任务上,换初始噪声就能采样不同合法解,不必像 GRAM 那样专门做变分推断。Sudoku 训练大约 1–5 小时,ARC 要 1–2 天,门槛不高,适合跟进复现。
这仍然是小模型、结构化谜题上的渐进改进,不是通用语言模型的推理方案。
迷宫上没有赢 FPRM。ARC-AGI-2 的 12.2% 离「能做抽象推理」还很远,5 条轨迹集成只稳住方差、均值没动。单解且数据少的任务必须靠伪目标正则:Sudoku 只用 1000 条训练,后期插值里的真解要换成模型上一步的预测,否则会过拟合。
训练并没有摆脱多步展开。k=16,每步仍要前向,只是梯度截断。作者把「免仿真训练」列为后续工作。理论上共享噪声存在捷径,可以从相邻插值线性解出答案;他们用「加步数还能涨分」来排除模型只在复读第一步,但没有直接探测循环状态里有没有缓存噪声。也没有和带思维链的大语言模型对照。