Stabilizing Extrapolation in Looped Transformers via Learned Stochastic Stopping¶

会议: NeurIPS 2026
arXiv: 2606.29983
代码: https://github.com/Hsun-Yu/looped_stopping
领域: 其他
关键词: 循环Transformer、长度泛化、随机停止策略、强化学习停止准则、外推稳定性

一句话总结¶

本文发现循环 Transformer 在简单算法任务中长度外推的极高方差源于训练序列长度与循环次数的伪相关，提出在训练阶段引入随机循环深度与基于强化学习的自适应停止（RL-Halting），将停止时机视为训练期归纳偏置，显著抑制了分布外预测震荡并提升了准确率与稳定性的权衡边界。

研究背景与动机¶

循环 Transformer（Looped Transformers）通过在深度方向重复复用相同的 Transformer 模块，在保持参数量固定的同时动态扩展计算图深度，天然契合多位数加法、Dyck 括号匹配、路径规划等可复用迭代算法任务。理论上，只要增加推理时的循环迭代次数，模型即可覆盖超出训练长度的更长输入；然而在实践中，这种长度泛化与外推能力极其脆弱。研究人员频繁观察到：即便在训练集内（In-Distribution, ID）具有近乎一致的高准确率，不同随机种子的模型在分布外（Out-of-Distribution, OOD）表现上却呈现巨大方差——部分模型能平滑外推至远超训练长度的样本，而另一些模型刚脱离训练长度便迅速崩溃。

这种脆弱性的根源在于传统训练方式所制造的伪相关与计算轨迹选择问题。现有的循环模型训练多采用确定性停止调度（Deterministic Schedule），例如固定循环步数 $K=20$ 或与输入长度严格绑定的长度匹配调度 $K=n$。在确定性监督下，模型在每次更新中仅在预先指定的单一切面接受监督，导致参数优化严重依附于该特定深度；更严重的是，算法任务中输入长度与循环次数之间形成了虚假锁定，网络无法形成在不同迭代步数下一致稳定的收敛轨迹。因此，循环结构并不天然赋予可靠的外推归纳偏置，反而引入了轨迹选择的严峻挑战：梯度下降容易收敛到仅在局部切面成立、但脱离该步数就剧烈震荡的退化解。

针对该痛点，本文提出不应仅将“何时停止”视为推理阶段的算力分配机制，而应将其视为训练阶段决定模型收敛轨迹的核心正则手段。与其在确定性单点施加监督，不如通过暴露局部邻域循环步数来消除对单一深度的病态敏感度，进而进一步通过任务损失引导的强化学习来学习自适应随机分布。核心 idea：将停止调度重构为训练期的随机正则机制，通过单样本单次采样的随机循环深度（及由任务负损失驱动的 RL-Halting 策略梯度）迫使隐状态在不同迭代步间保持表征一致并平滑梯度残差，从根本上收敛分布外方差并拓宽推理稳定区间。

方法详解¶

整体框架¶

模型以长度为 $m$ 的符号序列 $x$ 为输入，通过嵌入层获得初始状态 $H^{(0)} = E(x) \in \mathbb{R}^{m \times d}$。模型维护单一共享的 3 层 Transformer 模块 $\mathrm{Block}_{\theta}$，在每个循环步 $t = 1, \dots, K$ 中执行带输入注入（Input Injection）的循环迭代更新： $$H^{(t)} = \mathrm{Block}_{\theta}(H^{(t-1)} + H^{(0)})$$ 在训练阶段，算法不再将循环总步数 $K$ 限制为常量或确定性长度函数，而是根据停止策略 $\pi_{\mathrm{stop}}(K \mid x)$ 为每个样本独立采样一个随机停止步数 $\tau$。前向传播展开至该批次所需步数后，仅在每个样本采样的深度 $\tau$ 处施加单点任务交叉熵监督 $\ell(\theta; x, y, \tau) = \mathrm{CE}(\hat{y}_{\tau}(x), y)$。模型支持启发式均匀局部窗口随机化，亦可通过附加的轻量停止头与策略梯度端到端学习自适应停止分布 $\pi_{\mathrm{stop}, \phi}$。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入序列 x 与嵌入状态 H(0)"] --> Loop["循环Transformer前向迭代<br/>输入注入: H(t) = Block(H(t-1) + H(0))"]
    Loop --> StopCond{"训练停止调度策略"}
    StopCond -->|手制随机调度| RandWin["局部窗口采样<br/>K = clip(n + Δ, 1, Tmax)"]
    StopCond -->|自适应停止| RLHead["RL-Halting 危险率网络<br/>rt = σ(q_ϕ(ht)), ht = Pool(H(t))"]
    RLHead --> SampleTau["采样单一切面深度 τ ~ π_stop,ϕ"]
    RandWin --> SingleSup["单切面单点监督<br/>ℓ(θ) = CE(y_hat_τ, y)"]
    SampleTau --> SingleSup
    SingleSup --> DualOpt["解耦反向传播<br/>主网络梯度 ∇_θ ℓ + 停止头 REINFORCE ∇_ϕ J"]

关键设计¶

1. 采样深度随机调度：以期望梯度平滑打破单一深度依赖

传统的确定性调度 $K=n$ 或 $K=c$ 迫使模型在固定的循环深度拟合目标，导致网络学习到依赖特定深度的特异性梯度残差。针对这一问题，本文设计了局部窗口随机化调度：对于长度为 $n$ 的输入，采样偏移量 $\Delta \sim \mathrm{Unif}\{-w, \dots, w\}$，实际训练深度设定为 $K = \mathrm{clip}(n + \Delta, 1, T_{\max})$。尽管每个 SGD 步仅在单个采样深度 $\tau$ 处进行反向传播，但在期望意义下，优化目标等价于邻域深度的平均损失 $J_w(\theta; z) = \frac{1}{|S_w(x)|} \sum_{k \in S_w(x)} L_k(\theta; z)$。

若将深度 $k$ 的梯度分解为跨深度共享方向 $\bar{G}$ 与深度特异残差 $\xi_k$（满足 $\|\xi_k\|^2 \le B^2$ 及跨深度弱对齐 $\langle \xi_i, \xi_j \rangle \le \rho B^2$），则采样深度调度能有效抑制残差分量。数学推导表明其期望梯度与共享分量的误差上界满足： $$\|G_\pi(\theta; z) - \bar{G}(\theta; z)\|^2 \le B^2 \left( \rho + \frac{1-\rho}{m} \right)$$ 其中 $m = |S_w(x)|$ 为采样窗口大小。通过增大局部窗口，跨深度的特异扰动被统计平均抵消，有效阻断了模型对特定循环步数的脆弱过拟合。

2. RL-Halting 自适应停止：负任务损失驱动的策略梯度学习

启发式均匀窗口随机化虽能降低方差，但其均匀分布与任务复杂度无关，无法根据具体样本动态分配算力。为此，本文提出了自适应随机停止机制 RL-Halting。在每个循环步 $t$，模型对隐状态执行均值池化得到汇总向量 $h_t = \mathrm{Pool}(H^{(t)})$，并通过由单层线性层构成的轻量停止头 $q_\phi$ 预测该步的停止危险率（hazard）$r_t = \sigma(q_\phi(h_t))$。由此在时间范围 $1 \le t < T_{\max}$ 内诱导出累积停止概率分布： $$\pi_{\mathrm{stop}, \phi}(\tau = t \mid x) = r_t \prod_{i=1}^{t-1} (1 - r_i)$$ 剩余概率质量全部分配给最大截止步数 $T_{\max}$。训练时为每个样本按该分布采样单一切面深度 $\tau \sim \pi_{\mathrm{stop}, \phi}(\cdot \mid x)$，以任务负损失作为强化学习奖励 $R(x, y, \tau) = -\ell(\theta; x, y, \tau)$，采用带滑动平均基线 $b$ 的 REINFORCE 算法与探索熵正则化更新停止头参数 $\phi$： $$\mathcal{L}_{\mathrm{halt}}(\phi) = -\mathrm{sg}(R - b) \log \pi_{\mathrm{stop}, \phi}(\tau \mid x) - \lambda_{\mathrm{ent}} \mathcal{H}(\pi_{\mathrm{stop}, \phi}(\cdot \mid x))$$ 为防止特征污染，送入停止头的池化隐状态 $h_t$ 经过截断梯度（detach）处理，确保停止头策略优化不干扰 Transformer 主干参数 $\theta$ 的正常梯度更新。

3. 循环表征一致性：约束推理期翻转率与平坦吸引子

模型在推理阶段面临循环次数选择困境，若相邻步数的输出频繁震荡，策略准确率将不可避免大幅落后于预言机（Oracle）准确率。随机多深度监督在理论上构筑了强烈的循环表征一致性先验：假设任务存在最优目标表征 $Z^\star(x)$ 且局部损失具有 Polyak-Łojasiewicz 型下界 $L_k(\theta; x, y) \ge \frac{\mu}{2} \|H_k^\theta(x) - Z^\star(x)\|_F^2$，则当随机调度平均损失满足 $J_S(\theta) \le \varepsilon$ 时，任意两步隐状态间差异的期望严格受控： $$\mathbb{E}_{(x, y) \sim \mathcal{D}} \left[ \frac{1}{|S|^2} \sum_{i, j \in S} \|H_i^\theta(x) - H_j^\theta(x)\|_F^2 \right] \le \frac{8\varepsilon}{\mu}$$ 该约束促使网络在多步循环中形成渐进稳定的吸引子，不仅极大地拉宽了推理时能够正确解答的循环步数区间，更显著降低了相邻步数间的预测翻转率（Prediction Flip Rate）$\Pr_x[\hat{y}_{K+1}(x) \neq \hat{y}_K(x)]$，将预言机与固定策略之间的泛化落差（Oracle-Policy Gap）大幅压缩至极低水平。

损失函数 / 训练策略¶

系统总联合训练损失函数为： $$\mathcal{L}(\theta, \phi) = \mathrm{CE}(\hat{y}_\tau(x), y) - \mathrm{sg}(R - b) \log \pi_{\mathrm{stop}, \phi}(\tau \mid x) - \lambda_{\mathrm{ent}} \mathcal{H}(\pi_{\mathrm{stop}, \phi}(\cdot \mid x))$$ 其中 $\mathrm{sg}(\cdot)$ 为 stop-gradient 算子。优化超参数设置： - 优化器采用 AdamW，批大小设为 64，共训练 100k 步； - 初始学习率设为 $10^{-4}$，并伴随整个训练过程采用余弦衰减（Cosine Schedule）退火至 0； - 奖励基线采用滑动平均更新：$b \leftarrow 0.9 b + 0.1 \bar{R}$（$\bar{R}$ 为批次平均奖励）； - 探索熵正则化系数设定为 $\lambda_{\mathrm{ent}} = 0.001$（消融实验证实过小会导致策略退化坍缩，过大则阻碍自适应收敛）； - 训练截断步数 $T_{\mathrm{train}} = 20$，推理截断步数 $T_{\mathrm{test}} = 30$；测试阶段策略评估直接采用最大概率步数 $K_\pi(x) = \arg\max_{1 \le t \le T_{\mathrm{test}}} \pi_{\mathrm{stop}, \phi}^{(T_{\mathrm{test}})}(t \mid x)$。

实验关键数据¶

主实验¶

实验统一在受控长度外推协议下进行：所有模型在长度 $n < 20$ 的样本上进行在分布（ID）训练，在长度 $n \in \{20, 25, 30, \dots, 60\}$ 上进行分布外（OOD）评估。加法任务上评估 32 个独立随机种子，其余任务各评估 18 个独立随机种子。评估指标包含：OOD 预言机平均准确率（Oracle OOD, %）、达到 $\ge 90\%$ 准确率的最远外推前沿长度（Front.@90）、跨种子标准差（Std.，百分点，越低越稳定）以及预言机与策略评估准确率之间的落差（Gap，百分点）。

任务	训练调度方法	Oracle OOD (%) ↑	Front.@90 ↑	Std. (pp) ↓	Gap (pp) ↓
Binary Addition	固定 $K=20$	34.2	25	7.4	16.4
	长度匹配 $K=n$	22.1	20	9.9	7.8
	随机化 $K=n \pm 5$	30.9	25	5.9	1.3
	RL-Halting	45.0	30	2.7	5.8
Dyck-1	固定 $K=20$	84.6	30	23.8	32.8
	长度匹配 $K=n$	66.6	25	33.3	32.4
	随机化 $K=n \pm 5$	88.2	35	18.7	2.8
	RL-Halting	97.5	60	3.6	4.3
Unique Set	固定 $K=20$	66.1	20	34.0	41.7
	长度匹配 $K=n$	78.5	35	17.5	14.0
	随机化 $K=n \pm 5$	73.2	40	9.3	4.4
	RL-Halting	82.8	40	3.3	5.3
Copy	固定 $K=20$	62.3	25	27.7	29.3
	长度匹配 $K=n$	68.7	35	13.2	24.9
	随机化 $K=n \pm 5$	55.9	30	8.5	1.5
	RL-Halting	43.2	25	2.7	1.9

消融实验¶

论文针对控制变量（固定初始化/数据顺序）、停止策略变体、监督方式（单深度随机采样 vs 确定性多深度平均）以及关键超参数在二进制加法上展开了深度消融分析。

消融配置 / 机制验证	Oracle OOD (%) ↑	Policy OOD (%) ↑	Std. (pp) ↓	结论与机制说明
基线: 固定深度 $K=20$	34.2	17.8	7.4	确定性监督方差大，Gap 高达 16.4 pp
固定算力随机: $K=20 \pm 5$	34.2	-	4.7	均值不变但方差从 7.4 降至 4.7，证明随机化普适性
PonderNet 概率加权损失	29.4	-	15.1	易发生确定性坍缩，训练方差极高
多深度确定性平均监督	27.6	26.2	3.9	显示平均梯度可直接压降方差，证明非采样偶发效应
单深度随机采样 ($K=n \pm 5$)	30.9	29.6	5.9	保持梯度平均的同时带来隐式正则化，准确率高 3.3%
RL-Halt ($\lambda_{\mathrm{ent}}=0$)	12.51	-	-	缺乏熵正则导致策略迅速坍缩至极浅步数
RL-Halt ($\lambda_{\mathrm{ent}}=0.001$, 默认)	45.00	39.2	2.7	保持自适应性的同时维持适度探索，达最优综合权衡
RL-Halt ($\lambda_{\mathrm{ent}}=0.1$)	31.35	-	-	强熵约束迫使分布过度均匀，破坏长步数针对性收敛
RL-Halt ($T_{\max}=5$)	26.26	-	-	训练深度受限，无法覆盖复杂进位计算
RL-Halt ($T_{\max}=10$)	48.70	-	-	短步数预算下外推表现优异，算力性价比极高

关键发现¶

随机停止有效压降方差但并不等价于单调增益：在二进制加法中，从确定性 $K=n$ 切换至 $K=n \pm 5$ 使跨种子方差从 9.9 骤降至 5.9；但在 Copy 任务中，随机化虽将方差压降至 2.7，但 Oracle 准确率反而由 68.7% 掉至 43.2%。这表明稳定性与外推最优解并非同一维度，稳定化可能锁死在一个次优的计算轨迹上。
PonderNet 式加权损失存在严重的自强化坍缩风险：PonderNet 将所有深度的损失按停止概率加权求和，早期训练批次赋予某深度微小的概率优势，便会导致该深度获得更大的反向传播权重，进而自我强化坍缩为单一极度确定性浅层策略（方差高达 15.1 pp）。相反，RL-Halting 坚持“单点采样监督 + REINFORCE 更新”，成功绕开了权重自强化陷阱。
训练开销几乎与确定性模型持平：批处理实现中，前向展开步数取决于当前批次内样本采样的最大深度 $\max_i K_i$。在 batch size 64 下，RL-Halting 平均每步消耗 20.00 次模块前向评估，仅比 $K=n$（18.97 次）多 5.4%，与 $K=n \pm 5$（19.99 次）相当；在 50k 步壁钟时间对比中，所有方案耗时均在 2.5 小时左右（差距在 10 分钟内）。

亮点与洞察¶

将停止机制从推理侧算力分配前移为训练期归纳偏置：过去文献几乎全部将早停或自适应停止（如 ACT、PonderNet）作为推理阶段节省 FLOPs 的工具，本文深刻揭示了循环深度在训练时对梯度几何的重塑效应，提出了停止调度决定模型“学到哪条解题轨迹”的核心命题。
单样本单切面采样监督配合 REINFORCE 的极简设计：不依赖昂贵的跨深度全展开损失累加，仅在单个采样深度进行常规反向传播，将停止头与主干网络梯度更新彻底解耦，既杜绝了梯度污染，又以微乎其微的算力代价实现了策略自适应。
梯度残差平均与表征一致性的双重理论视角：不仅通过实验验证现象，更从优化几何（弱对齐残差相消）与隐空间动力学（吸引子收敛）两个正交视角建立了严谨的 toy 分析，为循环网络的动力学分析提供了清晰的理论抓手。

局限与展望¶

任务范畴局限于合成符号算法：实验集中于二进制加法、括号匹配和数组拷贝等小规模算法任务，虽然能干净地隔离控制变量，但尚未在大语言模型的大规模预训练和复杂自然语言多步推理（如 GSM8k、MATH、代码生成）中验证有效性。
稳定化可能固化次优解（Stability-Optimality Mismatch）：在 Copy 等任务中，随机化机制虽然将跨种子方差压缩至最低，却牺牲了极限外推能力。如何设计兼顾稳定性与最大外推容量的正则目标仍属开放课题。
未来方向：探索与深层架构稳定性（如 Deep Equilibrium Models 的谱半径正则化或路径独立性）的结合，并将自适应随机停止机制向 Looped LM 潜在思维空间推理方向扩展。

评分¶

新颖性: ⭐⭐⭐⭐ [重新诠释了循环停止策略的训练期正则属性，打破了仅在推理期讨论自适应算力的思维惯性]
实验充分度: ⭐⭐⭐⭐⭐ [跨多任务多达 32 个随机种子，涵盖确定性多深度干预、梯度残差分析及详尽的方差消融]
写作质量: ⭐⭐⭐⭐⭐ [逻辑层层递进，动机与理论机制阐述透彻，图表指标设计极具诊断说服力]
价值: ⭐⭐⭐⭐ [为循环大模型、潜在状态推理及循环架构的长度外推训练提供了至关重要的正则化指导]

任务	训练调度方法	Oracle OOD (%) ↑	Front.@90 ↑	Std. (pp) ↓	Gap (pp) ↓
Binary Addition	固定 \(K=20\)	34.2	25	7.4	16.4
	长度匹配 \(K=n\)	22.1	20	9.9	7.8
	随机化 \(K=n \pm 5\)	30.9	25	5.9	1.3
	RL-Halting	45.0	30	2.7	5.8
Dyck-1	固定 \(K=20\)	84.6	30	23.8	32.8
	长度匹配 \(K=n\)	66.6	25	33.3	32.4
	随机化 \(K=n \pm 5\)	88.2	35	18.7	2.8
	RL-Halting	97.5	60	3.6	4.3
Unique Set	固定 \(K=20\)	66.1	20	34.0	41.7
	长度匹配 \(K=n\)	78.5	35	17.5	14.0
	随机化 \(K=n \pm 5\)	73.2	40	9.3	4.4
	RL-Halting	82.8	40	3.3	5.3
Copy	固定 \(K=20\)	62.3	25	27.7	29.3
	长度匹配 \(K=n\)	68.7	35	13.2	24.9
	随机化 \(K=n \pm 5\)	55.9	30	8.5	1.5
	RL-Halting	43.2	25	2.7	1.9

消融配置 / 机制验证	Oracle OOD (%) ↑	Policy OOD (%) ↑	Std. (pp) ↓	结论与机制说明
基线: 固定深度 \(K=20\)	34.2	17.8	7.4	确定性监督方差大，Gap 高达 16.4 pp
固定算力随机: \(K=20 \pm 5\)	34.2	-	4.7	均值不变但方差从 7.4 降至 4.7，证明随机化普适性
PonderNet 概率加权损失	29.4	-	15.1	易发生确定性坍缩，训练方差极高
多深度确定性平均监督	27.6	26.2	3.9	显示平均梯度可直接压降方差，证明非采样偶发效应
单深度随机采样 (\(K=n \pm 5\))	30.9	29.6	5.9	保持梯度平均的同时带来隐式正则化，准确率高 3.3%
RL-Halt (\(\lambda_{\mathrm{ent}}=0\))	12.51	-	-	缺乏熵正则导致策略迅速坍缩至极浅步数
RL-Halt (\(\lambda_{\mathrm{ent}}=0.001\), 默认)	45.00	39.2	2.7	保持自适应性的同时维持适度探索，达最优综合权衡
RL-Halt (\(\lambda_{\mathrm{ent}}=0.1\))	31.35	-	-	强熵约束迫使分布过度均匀，破坏长步数针对性收敛
RL-Halt (\(T_{\max}=5\))	26.26	-	-	训练深度受限，无法覆盖复杂进位计算
RL-Halt (\(T_{\max}=10\))	48.70	-	-	短步数预算下外推表现优异，算力性价比极高