Research map / Architectures and efficiency / Attention and transformers

Stabilizing Extrapolation in Looped Transformers via Learned Stochastic Stopping

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Transformer expressivity and learnability · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Shows that Looped Transformers' length generalization is unstable across runs, and that randomizing (including learned RL-Halting) training loop depth reduces OOD variance and stabilizes predictions, though stability does not guarantee extrapolation.

Problem

Length generalization of Looped Transformers on algorithmic tasks is brittle: runs with similar in-distribution accuracy differ widely in OOD behavior, due to the training-time choice of loop depth.

Approach

Treat the loop-depth stopping schedule as a training-time design choice: sample one supervised depth per example (hand-designed windows around input length, or a learned hazard-based stopping distribution trained with REINFORCE using negative task loss as reward) to reduce OOD variance and flip rates across inference-time loop counts.

Key results

• On binary addition, RL-Halting achieves OOD accuracy 45.0 with std 2.7, versus 22.1 (std 9.9) for length-matched K=n and 34.2 (std 7.4) for fixed K=20.
• Length-matched K=n looped runs have OOD std 0.122 versus 0.009 and 0.005 for 3-layer and 60-layer Transformers.
• On Copy, RL-Halting reduces std from 13.2 to 2.7 but lowers OOD accuracy to 43.2 from 68.7 for length-matched K=n.

Surprising findings

• On Copy, RL-Halting stabilizes variance but stabilizes a suboptimal, non-extrapolating computation.
• PonderNet-style halting collapses toward near-deterministic stopping distributions and is less stable than RL-Halting.

Limitations

• Stability does not guarantee extrapolation; RL-Halting can stabilize suboptimal computations (e.g., Copy).
• Experiments limited to simple algorithmic tasks with small models trained on short lengths.
• Theoretical analysis is a local toy argument under an assumed shared-gradient model.

Experimental setup

Datasets: binary addition Dyck-1 Unique Set Copy

Benchmarks: binary addition Dyck-1 Unique Set Copy

Metrics: oracle-over-iterations accuracy Front.@90 standard deviation across runs prediction flip rate oracle-policy gap

Baselines: Fixed K=20 looped transformer length-matched K=n looped transformer K=n±5 random looped transformer 3-layer Transformer 60-layer Transformer PonderNet

Builds on: Adaptive Computation Time PonderNet Fan et al. looped transformer for length generalization REINFORCE input injection

Closest papers by meaning

• Stability and Generalization in Looped Transformers (NeurIPS 2026 · similarity 0.91)
• On Expressive Power of Looped Transformers: Theoretical Analysis and Enhancement via Timestep Encoding (ICML 2025 · similarity 0.87)
• From Shortcut to Induction Head: How Data Diversity Shapes Algorithm Selection in Transformers (NeurIPS 2025 · similarity 0.86)
• Linear Transformers Implicitly Discover Unified Numerical Algorithms (NeurIPS 2025 · similarity 0.86)
• Algorithmic Capabilities of Random Transformers (NeurIPS 2024 · similarity 0.86)
• Looped Transformers for Length Generalization (ICLR 2025 · similarity 0.86)