Stabilizing Extrapolation in Looped Transformers via Learned Stochastic Stopping

Hsun-Yu Kuo, El Mahdi Chayti, Patrik Reizinger, Wieland Brendel, Martin Jaggi

Keywords: Length Generalization, Looped Transformers, Recurrent Models, Learning when to stop

Abstract:

Looped Transformers – which repeatedly apply a shared 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 stabilises predictions across inference-time loop counts.
To improve upon heuristic randomisation schemes, we further analyse RL-Halting as a learned stochastic schedule and find it generally improves the accuracy--stability trade-off.
We provide comparisons across binary addition, Dyck-1, Unique Set, and Copy, showing that learned stochastic stopping often improves this trade-off but can also stabilise a suboptimal computation.
Overall, our results suggest that ``when to stop'' should be treated as a training-time design choice in looped architectures, not merely as an inference-time compute-allocation rule.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 55