# Stabilizing Extrapolation in Looped Transformers via Learned Stochastic Stopping

arXiv: https://arxiv.org/abs/2606.29983 · HTML: https://mlsanity.com/paper/2606.29983/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 33 references that are in the map.
- First posted on arXiv 29 Jun 2026.

## Builds on (most cited first, 25 of 25)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Networks with Stochastic Depth](https://mlsanity.com/paper/1603.09382/) (2016)
- [Adaptive Computation Time for Recurrent Neural Networks](https://mlsanity.com/paper/1603.08983/) (2016)
- [Reducing Transformer Depth on Demand with Structured Dropout](https://mlsanity.com/paper/1909.11556/) (2019)
- [Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach](https://mlsanity.com/paper/2502.05171/) (2025)
- [The Impact of Positional Encoding on Length Generalization in Transformers](https://mlsanity.com/paper/2305.19466/) (2023)
- [On The Power of Curriculum Learning in Training Deep Networks](https://mlsanity.com/paper/1904.03626/) (2019)
- [Reasoning with Latent Thoughts: On the Power of Looped Transformers](https://mlsanity.com/paper/2502.17416/) (2025)
- [Looped Transformers as Programmable Computers](https://mlsanity.com/paper/2301.13196/) (2023)
- [What Algorithms can Transformers Learn? A Study in Length Generalization](https://mlsanity.com/paper/2310.16028/) (2023)
- [Hierarchical Reasoning Model](https://mlsanity.com/paper/2506.21734/) (2025)
- [Scaling Latent Reasoning via Looped Language Models](https://mlsanity.com/paper/2510.25741/) (2025)
- [Looped Transformers are Better at Learning Learning Algorithms](https://mlsanity.com/paper/2311.12424/) (2023)
- [Can You Learn an Algorithm? Generalizing from Easy to Hard Problems with Recurrent Networks](https://mlsanity.com/paper/2106.04537/) (2021)
- [Randomized Positional Encodings Boost Length Generalization of Transformers](https://mlsanity.com/paper/2305.16843/) (2023)
- [Looped Transformers for Length Generalization](https://mlsanity.com/paper/2409.15647/) (2024)
- [Transformers Can Achieve Length Generalization But Not Robustly](https://mlsanity.com/paper/2402.09371/) (2024)
- [Generalization on the Unseen, Logic Reasoning and Degree Curriculum](https://mlsanity.com/paper/2301.13105/) (2023)
- [Path Independent Equilibrium Models Can Better Exploit Test-Time Computation](https://mlsanity.com/paper/2211.09961/) (2022)
- [End-to-end Algorithm Synthesis with Recurrent Networks: Logical Extrapolation Without Overthinking](https://mlsanity.com/paper/2202.05826/) (2022)
- [Self-Improving Transformers Overcome Easy-to-Hard and Length Generalization Challenges](https://mlsanity.com/paper/2502.01612/) (2025)
- [Position Coupling: Improving Length Generalization of Arithmetic Transformers Using Task Structure](https://mlsanity.com/paper/2405.20671/) (2024)
- [Position: Understanding LLMs Requires More Than Statistical Generalization](https://mlsanity.com/paper/2405.01964/) (2024)
- [Extrapolation by Association: Length Generalization Transfer in Transformers](https://mlsanity.com/paper/2506.09251/) (2025)
- [On Vanishing Variance in Transformer Length Generalization](https://mlsanity.com/paper/2504.02827/) (2025)

## Built on by (most cited first, 4 of 4)

- [Loop the Loopies!](https://mlsanity.com/paper/2607.16051/) (2026)
- [Improving Test-Time Scaling with Adaptive Looped Transformers](https://mlsanity.com/paper/2609.35748/) (2026)
- [Towards Looped Models Done Right, Part II: Rethinking at Fixed Points](https://mlsanity.com/paper/2610.06833/) (2026)
- [Learning to Compute Is Not Learning to Retain or Reuse in Looped Transformers](https://mlsanity.com/paper/2610.19768/) (2026)

## Abstract (arXiv)

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.