通过学习随机停止策略，提升循环Transformer在长序列上的预测稳定性。
Stabilizing Extrapolation in Looped Transformers via Learned Stochastic Stopping
- 训练时引入随机循环次数，缓解长度泛化中的预测波动。
- 在二进制加法等任务上，显著降低分布外误差，提升稳定性。
- 适合关注算法推理模型鲁棒性的研究者与工程师。
循环Transformer通过重复应用共享的Transformer模块，天然适用于可变长度的算法任务。尽管其在训练序列长度之外表现出较强的长度泛化能力，但该行为极为脆弱，导致分布外（OOD）预测方差高，即使在表现良好的分布内解中也是如此。我们发现，这种方差源于简单算法任务中序列长度与循环次数之间的虚假相关性。在训练中引入循环次数的随机性，可显著降低分布外方差，并稳定不同推理循环次数下的预测结果。为进一步改进启发式随机化方案，我们分析了基于强化学习的停机机制（RL-Halting），发现其通常能改善准确率与稳定性的权衡。在二进制加法、Dyck-1、Unique Set和Copy任务上，学习到的随机停止策略往往优化了这一权衡，但也可能稳定一个次优计算过程。本工作表明，‘何时停止’应作为训练阶段的设计选择，而不仅是推理时的计算分配规则。
原文摘要 · Abstract (English)
Looped Transformers, which repeatedly apply a shared transformer block, are an architecturally natural fit for variable-length algorithmic tasks. Although they can exhibit strong length generalization beyond the length of training sequences, this behavior is brittle, yielding high out-of-distribution (OOD) variance, even across well-performing in-distribution solutions. We trace this variance to the spurious correlation in simple algorithmic tasks between sequence length and number of loops. Introducing stochasticity into the number of loops during training sharply reduces OOD variance and stabilizes predictions across inference-time loop counts. To improve upon heuristic randomization schemes, we further analyze RL-Halting as a learned stochastic schedule and find that it generally improves the accuracy-stability trade-off. Across binary addition, Dyck-1, Unique Set, and Copy, learned stochastic stopping often improves this trade-off but can also stabilize a suboptimal computation. Our work suggests that "when to stop" should be treated as a training-time design choice, not merely an inference-time computation-allocation rule.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。