Stabilizing Extrapolation in Looped Transformers via Learned Stochastic Stopping

• URL: http://arxiv.org/abs/2606.29983v1
• Date: Mon, 29 Jun 2026 08:58:09 GMT
• Title: Stabilizing Extrapolation in Looped Transformers via Learned Stochastic Stopping
• Abstract summary: Looped Transformers, which repeatedly apply a shared transformer block, are an architecturally natural fit for variable-length computation tasks.<n>We trace this variance to the spurious correlation in simple algorithmic tasks between sequence length and number of loops.<n>Our work suggests that "when to stop" should be treated as a training design choice, not merely an inference-time-allocation rule.
• Score: 56.14767235650558
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Direct Acceleration of Stochastic Root-Finding Without Variance Reduction and Regularization [15.2286904549704]
  We show that the dual-anchor mechanism, namely the dual-anchor mechanism, extends to the iteration setting without such error accumulation.<n>For strongly monotone operators, the same algorithm attains a sharper complexity, nearly matching the lower bound in terms of $$-dependence.
  arXiv  Detail & Related papers  (2026-08-12T13:25:55Z)
• LoopCoder-v2: Only Loop Once for Efficient Test-Time Computation Scaling [72.71005779366162]
  This study studies loop-count selection through a gain--cost view.<n>We instantiate this study by training LoopCoder-v2 with different loop counts, from scratch on 18T tokens, followed by matched instruction tuning and evaluation.<n> Empirically, the two-loop variant delivers broad gains over the non-looped baseline across code generation, code reasoning, agentic software engineering, and tooluse benchmarks.
  arXiv  Detail & Related papers  (2026-06-16T15:03:05Z)
• Beyond Similarity: Temporal Operator Attention for Time Series Analysis [5.084091872064936]
  A persistent paradox in time-series forecasting is that structurally simple and linear models often outperform high-capacity Transformers.<n>We propose $textbf Operator Attention (TOA)$, a framework that augments attention with explicit learnable sequence-space operators.<n>To make dense $N times operators practical, we introduce a high-variance dropout mechanism that stabilizes training and prevents trivial memorization.
  arXiv  Detail & Related papers  (2026-05-11T22:13:53Z)
• Simply Stabilizing the Loop via Fully Looped Transformer [41.240805541680395]
  Looped Transformer suffers from training instability when the number of loop iterations increases.<n>Our experiments demonstrate that Fully Looped Transformer improves training stability, enhances downstream performance, and provides preliminary adaptability under different test-time compute budgets.
  arXiv  Detail & Related papers  (2026-05-11T07:21:53Z)
• Learning to Shuffle: Block Reshuffling and Reversal Schemes for Stochastic Optimization [20.63685754619351]
  We use a large language model (LLM)-guided program evolution framework to discover an effective shuffling rule for without replacing SGD.<n>We show that block reshuffling strictly reduces prefix-gradient variance constants within the unified shuffling, yielding provable improvements over random reshuffling under mild conditions.<n>We also show that paired reversal symmetrizes the epoch map and cancels the leading order-dependent second-order term, reducing order sensitivity from quadratic to cubic in the step size.
  arXiv  Detail & Related papers  (2026-03-31T21:40:43Z)
• Weights to Code: Extracting Interpretable Algorithms from the Discrete Transformer [65.38883376379812]
  We propose the Discrete Transformer, an architecture engineered to bridge the gap between continuous representations and discrete symbolic logic.<n> Empirically, the Discrete Transformer not only achieves performance comparable to RNN-based baselines but crucially extends interpretability to continuous variable domains.
  arXiv  Detail & Related papers  (2026-01-09T12:49:41Z)
• Enhancing Auto-regressive Chain-of-Thought through Loop-Aligned Reasoning [47.06427150903487]
  Chain-of-Thought (CoT) prompting has emerged as a powerful technique for enhancing language model's reasoning capabilities.<n>Looped Transformers possess remarkable length generalization capabilities, but their limited generality and adaptability prevent them from serving as an alternative to auto-regressive solutions.<n>We propose RELAY to better leverage the strengths of Looped Transformers.
  arXiv  Detail & Related papers  (2025-02-12T15:17:04Z)
• Timer-XL: Long-Context Transformers for Unified Time Series Forecasting [67.83502953961505]
  We present Timer-XL, a causal Transformer for unified time series forecasting.<n>Based on large-scale pre-training, Timer-XL achieves state-of-the-art zero-shot performance.
  arXiv  Detail & Related papers  (2024-10-07T07:27:39Z)
• Relaxed Quantile Regression: Prediction Intervals for Asymmetric Noise [51.87307904567702]
  Quantile regression is a leading approach for obtaining such intervals via the empirical estimation of quantiles in the distribution of outputs.<n>We propose Relaxed Quantile Regression (RQR), a direct alternative to quantile regression based interval construction that removes this arbitrary constraint.<n>We demonstrate that this added flexibility results in intervals with an improvement in desirable qualities.
  arXiv  Detail & Related papers  (2024-06-05T13:36:38Z)
• Finetuning Pretrained Transformers into RNNs [81.72974646901136]
  Transformers have outperformed recurrent neural networks (RNNs) in natural language generation. A linear-complexity recurrent variant has proven well suited for autoregressive generation. This work aims to convert a pretrained transformer into its efficient recurrent counterpart.
  arXiv  Detail & Related papers  (2021-03-24T10:50:43Z)
• Variance Regularization for Accelerating Stochastic Optimization [14.545770519120898]
  We propose a universal principle which reduces the random error accumulation by exploiting statistic information hidden in mini-batch gradients. This is achieved by regularizing the learning-rate according to mini-batch variances.
  arXiv  Detail & Related papers  (2020-08-13T15:34:01Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.