On the Residual Scaling of Looped Transformers: Stability and Transferability

2026-06-16

AI Summary

This work addresses the training instability and poor hyperparameter transferability in Looped Transformers, which arise from highly correlated residual updates across iterations due to weight sharing. The authors demonstrate that conventional depth-scaling strategies (e.g., ε = 1/√L) are insufficient in such recurrent architectures and reveal, for the first time, that residual correlations necessitate a stronger 1/N scaling with respect to the number of loop iterations N. They propose a composite scaling strategy, ε = λ/(N√L), which decouples the learning rate from N and makes it dependent solely on the effective depth L. Grounded in theoretical analysis and multi-layer loop block modeling, experiments confirm that this approach significantly enhances training stability, consistently outperforms 1/√N scaling across varying N, and enables direct transfer of large-N models without re-tuning hyperparameters.

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.

Problem

• Looped Transformers
• Residual Scaling
• Weight Tying
• Trainability
• Hyperparameter Transfer

Innovation

• Looped Transformers
• Residual Scaling
• Weight Tying
• Hyperparameter Transferability
• Stability Analysis