Stabilizing Extrapolation in Looped Transformers via Learned Stochastic Stopping

2026-06-29

AI Summary

This work addresses the instability of Looped Transformers in variable-length algorithmic tasks, where spurious correlations between sequence length and loop count lead to poor out-of-distribution extrapolation and high prediction variance. To mitigate this, the authors propose injecting randomness into the number of computation steps during training and introduce RL-Halting—a learnable stochastic halting mechanism grounded in reinforcement learning—that treats “when to stop” as a core trainable component rather than merely an inference-time computational heuristic. This study presents the first systematic investigation into how learnable stochastic stopping enhances extrapolation stability in Looped Transformers. Experiments on tasks such as binary addition and Dyck-1 demonstrate that the approach substantially reduces out-of-distribution variance and improves the trade-off between accuracy and stability, albeit occasionally converging to suboptimal yet stable computation paths.

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.

Problem

• Looped Transformers
• length generalization
• out-of-distribution variance
• stochastic stopping
• algorithmic tasks

Innovation

• Looped Transformers
• stochastic stopping
• length generalization
• OOD stability
• RL-Halting