Simply Stabilizing the Loop via Fully Looped Transformer
- URL: http://arxiv.org/abs/2605.18797v1
- Date: Mon, 11 May 2026 07:21:53 GMT
- Title: Simply Stabilizing the Loop via Fully Looped Transformer
- Abstract summary: 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.
- Score: 41.240805541680395
- License: http://creativecommons.org/licenses/by/4.0/
- Abstract: Scaling model performance typically requires increasing model size. Looped Transformer offers a compelling alternative by iteratively reusing the same Transformer blocks, trading additional computation for improved performance without increasing parameter count or context length. Because the number of loop iterations can be adjusted at inference, it also provides a natural mechanism for balancing performance and test-time compute. However, Looped Transformer still suffers from training instability when the number of loop iterations increases. Our analysis reveals that this instability stems from two sources: gradient oscillation and residual explosion. To address these two problems, we propose the Fully Looped Transformer, which introduces two parameter-free modifications: (1) Fully Looped Architecture, which distributes inter-loop signals across all layers to mitigate residual explosion; (2) Attention Injection, which reuses the existing attention block to suppress gradient oscillation. These modifications stabilize training dynamics, enabling the Fully Looped Transformer to be trained stably up to 12 loop iterations, whereas other baseline looped models collapse in this regime. In milder settings where Looped Transformer does not collapse, Fully Looped Transformer still improves average downstream-task performance by up to 13.2\%. Overall, our experiments demonstrate that Fully Looped Transformer improves training stability, enhances downstream performance, and provides preliminary adaptability under different test-time compute budgets by varying loop iterations at inference.
Related papers
- SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers [50.00050394179417]
Looped Transformers increase effective depth by iterating a shared block of layers, but most evaluations compare at fixed model size.<n>We study looping on Mixture-of-Experts Transformers while closely matching per-token FLOPs, total non-embedding parameters, and KV cache.<n>Through a series of ablations, we arrive at a recipe we call SMELT, which loops the middle half of layers twice while matching the unlooped Baseline on all three budgets.
arXiv Detail & Related papers (2026-09-01T14:52:57Z) - Training-Free Hidden-State Refinement for Flow-Matching Image Generators [13.43479682921888]
We introduce a training-free looping framework that repeatedly applies selected transformer layers inside each denoising call.<n>Across two Scale-RAE model scales, loop variants improve primary and auxiliary quality metrics with competitive quality--efficiency trade-offs.
arXiv Detail & Related papers (2026-08-29T09:17:02Z) - Stabilizing Extrapolation in Looped Transformers via Learned Stochastic Stopping [56.14767235650558]
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.
arXiv Detail & Related papers (2026-06-29T08:58:09Z) - 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) - Parcae: Scaling Laws For Stable Looped Language Models [35.9547796403241]
Traditional fixed-depth architectures scale quality by increasing training FLOPs, typically through increased parameterization, at the expense of a higher memory footprint, or data.<n>A potential alternative is looped architectures, which instead increase FLOPs by sending activations through a block of layers in a loop.<n>While promising, existing recipes for training looped architectures can be unstable, suffering from residual explosion and loss spikes.<n>We propose Parcae, a novel stable, looped architecture that constrains the spectral norm of the injection parameters via discretization of a negative diagonal parameterization.
arXiv Detail & Related papers (2026-04-14T16:43:37Z) - 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) - Investigating Recurrent Transformers with Dynamic Halt [64.862738244735]
We study the inductive biases of two major approaches to augmenting Transformers with a recurrent mechanism.<n>We propose and investigate novel ways to extend and combine the methods.
arXiv Detail & Related papers (2024-02-01T19:47:31Z) - 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)
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.