IAAILearned Stochastic Stopping
Stabilizing Extrapolation in Looped Transformers via Learned Stochastic Stopping
Traces the brittle length generalization of looped transformers to a spurious correlation between sequence length and loop count, and shows that randomizing loop counts in training, or learning a stochastic schedule with RL-Halting, stabilizes predictions across inference-time loop counts.
Inside the method
[ Shared block ] × r, r sampled (heuristic or RL-Halting) during training
Simplified conceptual schematic. Consult the paper for the complete architecture.
- Recurrence family
- Adaptive depth
- Depth control
- Stochastic stopping at training time
- KV / state strategy
- See paper
Reading note
On binary addition, Dyck-1, Unique Set and Copy, learned stochastic stopping often improves the accuracy-stability trade-off but can also stabilize a suboptimal computation.
Sources checked 2026-10-28. This catalog does not imply independent reproduction.
Related works
Cite this work
@misc{stochasticstopping2026,
title = {Stabilizing Extrapolation in Looped Transformers via Learned Stochastic Stopping},
author = {Hsun-Yu Kuo and El Mahdi Chayti and Patrik Reizinger and Wieland Brendel and Martin Jaggi},
year = {2026},
eprint = {2606.29983},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2606.29983}
}