想让模型变强,最直接的办法是加参数——加层、加宽度、加注意力头。但这招越来越贵了。
Looped Transformer 想走另一条路:同一组 Transformer 层反复用。推理时多循环几次,效果就更好。参数不增加,上下文不拉长,只多花计算时间。
听起来很美好。但之前的人发现一个头疼的问题——循环次数一多,训练就崩了。
为什么循环 Transformer 会崩
这篇论文先做了诊断,找到了两个罪魁祸首:
梯度振荡。 每次循环后梯度方向剧烈变化,像开车时方向盘左右乱打,越走越偏。
残差爆炸。 循环过程中残差信号逐层累积,到后面直接溢出。
这两个问题导致之前的 Looped Transformer 在循环次数超过某个阈值后就训练不下去——模型输出变成 NaN,一切归零。
Fully Looped Transformer 的解法
两个改动,都不需要额外参数:
第一,完全循环架构。 把循环间的信号分散到所有层,而不是只加在某个固定位置。这缓解了残差爆炸——信号不再从一个点涌入,而是均匀分布在整个网络里。
第二,注意力注入。 复用已有的注意力模块来抑制梯度振荡。不是加新模块,而是改变现有注意力的使用方式,让它同时承担稳定训练的功能。
结果?之前的循环模型在 12 次循环时崩溃,Fully Looped Transformer 稳定跑完。在较温和的设置下,即使 baseline 不会崩溃,Fully Looped Transformer 仍然把下游任务的平均性能提升了 13.2%。
推理时算力 = 可调旋钮
这篇论文最让我感兴趣的一点不是性能数字,而是它暗示的可能性。
因为循环次数在推理时可以调整——这意味着你可以用一个模型,根据查询的复杂度动态决定花多少计算量。
简单问题,循环 1-2 次就够。复杂推理题,循环 8-12 次。不需要部署多个不同规模的模型,只需要调整一个超参数。
这不就是 test-time compute 的朴素版本吗?不需要 o1 那样的复杂推理训练,只需要一个训练稳定的循环架构。
当然,论文自己也承认,目前只在中等规模任务上做了验证,离生产环境还有距离。但至少方向是对的——用架构改动来解锁推理时灵活性,比堆参数更可持续。
论文:Simply Stabilizing the Loop via Fully Looped Transformer 代码:尚未公开