arXiv cs.AI· Hsun-Yu Kuo, El Mahdi Chayti, Patrik Reizinger, Wieland Brendel, Martin Jaggi·· 2 天前AI 评分22
通过可学习随机停止稳定 Looped Transformer 的外推能力
Stabilizing Extrapolation in Looped Transformers via Learned Stochastic Stopping
AI 导读
Looped Transformer 在训练序列长度之外进行外推时方差较高，研究将其归因于序列长度与循环次数之间的伪相关。训练时引入随机循环次数可显著降低 OOD 方差，而 RL-Halting 作为可学习随机调度通常能改善准确率与稳定性的权衡。在二元加法、Dyck-1、Unique Set 和 Copy 任务上，学习式随机停止有时也会稳定次优计算，表明"何时停止"应作为训练时设计选择。
正文 · 原文
Abstract: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.
| Comments: | Published at NeurIPS 2026 |
| Subjects: | Machine Learning (cs.LG); Artificial Intelligence (cs.AI) |
| Cite as: | arXiv:2606.29983 [cs.LG] |
| (or arXiv:2606.29983v2 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2606.29983 arXiv-issued DOI via DataCite |
Submission history
From: Hsun-Yu Kuo [view email]
[v1]
Mon, 29 Jun 2026 08:58:09 UTC (677 KB)
[v2]
Wed, 14 Oct 2026 21:15:48 UTC (683 KB)
来源：arXiv cs.AI · arxiv.org