On the Residual Scaling of Looped Transformers: Stability and Transferability

Abstract: Looped (weight-tied) Transformers apply a shared residual block $N$ times ($h \leftarrow h + \varepsilon\,f(h)$, same $f$ at each step), increasing effective depth without adding parameters. Prior depth-scaling analyses prescribe $\varepsilon = 1/!\sqrt{L}$ for depth-$L$ residual networks. We show that this is insufficient for looped architectures: weight sharing makes residual updates correlated across iterations, requiring the stronger scaling $\varepsilon = 1/N$. For multi-layer blocks ($L$ unique layers looped $N$ times), we derive a factored parameterization $\varepsilon = λ/(N!\sqrt{L})$ that separates the two sources of growth: $1/N$ controls the within-layer loop correlation, and $1/!\sqrt{L}$ controls the across-layer variance. A key consequence is that the optimal learning rate depends only on the number of unique layers $L$, not on the loop count $N$, enabling direct hyperparameter transfer from small to large $N$ without retuning. Experiments on looped Transformers confirm that $1/N$ scaling improves trainability and yields better loss than $1/!\sqrt{N}$ scaling across loop counts.

• The paper establishes that correct residual scaling in looped Transformers requires adapting the standard 1/L rule to a 1/N prescription due to weight sharing.
• It derives a factored scaling law that decouples the learning rate from the loop count by using 1/N and 1/√L adjustments, and confirms this through rigorous analysis and empirical validation.
• Experimental results demonstrate that proper scaling leads to stable residual norms and improved performance, simplifying hyperparameter transfer across varying loop counts.

On the Residual Scaling of Looped Transformers: Stability and Transferability

Problem Context and Motivation

Looped or weight-tied Transformers, wherein a shared residual block is applied multiple times ($N$-fold reuse), offer a parameter-efficient strategy to augment effective depth, appearing in variants such as Universal Transformers and ALBERT. However, this architectural motif induces unique challenges in optimization stability not addressed by classical residual scaling analyses based on non-shared (deep) stacks. In particular, norm growth in residual-streams and amplified sensitivity to learning rates have impaired the reliable scaling of such models, hindering practical exploration of looped architectures as an efficient scaling axis.

Main Theoretical Contributions

The paper systematically dissects the accumulation of activations in looped residual networks, with the central finding that residual scaling must be adapted from the standard $1/L$ prescription to $1/N$ in single-layer looped architectures, due to correlation between successive residual updates induced by weight sharing. This core insight is analytically established by contrasting the random-walk behavior of independent-residual deep stacks with the linear, reinforced accumulation in looped stacks.

Generalizing to multi-layer blocks consisting of $L$ unique layers, each looped $N$ times, the authors derive a factored scaling law:

$\epsilon = \frac{\lambda}{N \sqrt{L}}$

where $1/N$ corrects for the within-layer looped constructive interference, and $1/\sqrt{L}$ addresses the across-layer variance. This parameterization, directly motivated by rigorous variance and alignment analyses, is shown to control both the output norm and learning rate stability.

A significant byproduct is the demonstration that, under proper residual scaling, the optimal learning rate depends exclusively on the number of unique layers $L$, not on the loop count $N$, enabling effective hyperparameter transfer from shallow to deep (large $1/L$0) looped regimes without costly retuning. This property, confirmed in empirical sweeps, is non-trivial and critical for model development efficiency.

Empirical Validation and Numerical Results

Comprehensive experiments substantiate the theoretical predictions. Diagnostics on decoder-only Transformers with looped stacks trained on FineWeb-Edu confirm:

• Standard $1/L$1 scaling is insufficient: Forward-pass residual norms grow uncontrollably for large $1/L$2, breaking stability.
• Linear $1/L$3 scaling stabilizes the residual stream: Both theoretical predictions and measured norms confirm bounded residuals for all practical $1/L$4.
• Learning rate transfer is realized: With factored scaling, a learning rate optimal at $1/L$5 remains close to optimal for $1/L$6.
• Performance improves at large loop counts: For example, with $1/L$7 and $1/L$8, the minimum validation loss at large $1/L$9 is lower under $1/N$0 scaling (by up to 0.025 nats) than with weaker scaling rules.

The pairwise cosine similarity (loop-step increments) evidences persistently dense positive alignment (not present in standard deep stacks), confirming the theoretical mechanism of constructive accumulation. Importantly, post-training analyses show the $1/N$1 growth in residual norm is not an initialization artifact but persists after extended optimization.

Implications and Theoretical Ramifications

This work extends and corrects existing theory on depth scaling in residual architectures (Dey et al., 2 May 2025) to the case of shared-parameter (looped) stacks. The breakdown of the independence assumption is shown to fundamentally alter norm and learning dynamics, compelling new scaling laws. This result also aligns with and completes the tensor-program maximal-update and depth-transfer theory (Yang et al., 2022, Dey et al., 2 May 2025), integrating loop count as a clean, decoupled scaling axis.

By decoupling loop count from optimizer stability, weight reuse becomes a tunable dimension: looped architectures can be scaled at training and inference without repeated grid-searches for normalization and learning rates, addressing a critical bottleneck for practical deployment. For foundational models or research seeking to explore parameter-efficient scaling, this work provides exact instructions for maintaining trainability and convergent dynamics.

Limitations and Future Directions

The analysis relies on simplifying abstractions such as single-branch MLPs and ReLU (rather than real-world SwiGLU) activations; however, empirical results suggest the core mechanisms persist in practical architectures. The present treatment omits optimizer state dynamics, normalization variant effects, and feature-learning or data-dependent complications, which may all influence stability in large-scale, deployed models.

Future extensions should target alternate normalization schemes (especially Post-Norm), interactions with regularization and optimizer momentum/bias correction, and potential regime changes in extreme-depth or heterogeneous-loop settings. There is also scope to validate the recipe in the context of new Transformer variants with dynamic or content-based parameter sharing, and in tasks beyond standard language modeling.

Conclusion

This paper provides a mathematically principled and empirically validated framework for residual scaling in looped (weight-tied) Transformers, showing that constructive accumulation due to parameter reuse necessitates $1/N$2 (and, in multi-layer stacks, factored) residual scaling and learning-rate laws. These results resolve a major roadblock in stable looped network training and yield immediately applicable recipes for scalable, hyperparameter-transferable Transformer architectures. Addressing the inherent instability of looped designs, these scaling laws broaden the practical and theoretical toolkit for efficient deep learning model construction and open robust avenues for research on recursive and parameter-efficient neural architectures.

Reference: "On the Residual Scaling of Looped Transformers: Stability and Transferability" (2606.18524)