Papers
arxiv:2606.18524

Loop-μP: Residual Scaling and Hyperparameter Transfer in Looped Transformers

Published on Oct 18
Authors:
,
,
,
,
,
,

Abstract

Loop-μP scales residual branches by loop and depth factors to stabilize training and preserve hyperparameters across looped Transformers.

Looped Transformers have drawn growing attention because looping improves multi-step reasoning and increases the executed depth of a model without adding parameters. They achieve this by reusing a stack of L layers for N loops. However, increasing the loop count N not only destabilizes training but also shifts the best hyperparameters, so they must be searched again for every loop count. Inspired by Depth-μP (Yang et al., 2024), which multiplies each residual branch by L^{-1/2} for depth L so that a learning rate tuned on a shallow untied network transfers to a deep one, we study how the residual branches and the learning rate should scale with the loop count N. We find that loops cannot simply be treated as extra depth, because the loops share weights. At initialization, distinct layers add nearly uncorrelated increments to the residual stream, so their sum grows as L, whereas a shared layer adds positively correlated increments at every loop, so their sum grows linearly in N. For a single layer looped N times, we show in the infinite-width limit that multiplying the residual branch by 1/N and keeping the learning rate independent of N is the only parameterization with both a stable forward pass and maximal residual-stream updates. We then generalize the analysis to a stack of L distinct layers looped N times to obtain a parameterization we call Loop-μP, in which the loop factor 1/N combines with the depth factor L^{-1/2}. Empirically, we find in language modeling that the base learning rate tuned on the smallest model stays best under Loop-μP for loop counts up to 8 and executed depths up to 384 layers. These results make Loop-μP a practical recipe for scaling looped Transformers.

Community

Sign up or log in to comment

Models citing this paper 0

No model linking this paper

Cite arxiv.org/abs/2606.18524 in a model README.md to link it from this page.

Datasets citing this paper 0

No dataset linking this paper

Cite arxiv.org/abs/2606.18524 in a dataset README.md to link it from this page.

Spaces citing this paper 0

No Space linking this paper

Cite arxiv.org/abs/2606.18524 in a Space README.md to link it from this page.

Collections including this paper 1