PaperlayerView exact revision on arXiv

Source-linked AI summary

DeepLoop: Depth Scaling for Looped Transformers

Shuzhen Li, Yifan Zhang, Jiacheng Guo, Quanquan Gu, Mengdi Wang

arXiv:2607.13491v2cs.LGcs.AI

TL;DR

Looped Transformers reuse physical blocks, but existing residual-scaling analyses do not account for repeated parameter visits. DeepLoop derives a visit-aware scaling rule and improves validation loss and downstream accuracy when recurrent depth is activated.

  • Problem

    Standard residual-scaling analyses assume untied depth, although looped Transformers repeatedly visit shared parameters and aggregate gradient contributions across rounds.

  • Method

    DeepLoop derives a loop-aware first-order perturbation bound using visit-alignment coefficient κ_R and sets α=(2N)^1/2 and β=(8N)^−1/2.

  • Results

    DeepLoop is neutral without block revisits and improves validation loss and downstream accuracy as recurrent depth increases.

  • Takeaways & Limitations

    Residual scaling for recurrent depth should account for physical-block revisits rather than nominal depth alone.

  • Takeaways & Limitations

    The rule’s scope remains uncertain at larger scales, under alternative parameterizations, and without direct measurements of κ_R or cross-round gradient alignment.

Abstract

from arXiv · show

Looped Transformers scale sequential computation by applying a compact stack of physical blocks for multiple rounds, increasing unrolled depth without increasing stored parameters. This reuse changes the residual-scaling problem: in an untied Transformer, each residual branch receives and applies its own parameter update, whereas in a looped Transformer one shared update aggregates gradients from repeated visits and is read back by those same visits in the next linearized forward pass. We formalize this tied-depth effect through a first-order perturbation bound controlled by a visit-alignment coefficient $κ_R$. The bound recovers the DeepNorm exponent when visits decorrelate, but in the conservative aligned regime it requires the exponent to increase from $1/4$ to $1/2$ as loop count grows at fixed physical depth. The resulting method, \textbf{DeepLoop}, keeps the Post-LN DeepNorm architecture and sets $α=(2N)^{1/2}$ and $β=(8N)^{-1/2}$ for unrolled depth $N$. On GPT-style looped language models at GPT-2 small and GPT-2 medium scale, DeepLoop is neutral when no physical block is revisited and improves validation loss and downstream accuracy once recurrent depth is activated. These results show that stable recurrent depth requires residual scaling rules that account for parameter visits, not only nominal layer count.

1 Introduction

DeepLoop analyzes how revisiting shared physical blocks changes residual-scaling stability in looped Transformers. It derives a visit-aware scaling rule and reports improvements in validation loss and downstream performance when recurrent depth is activated.

  • Looped Transformers apply K physical blocks for R rounds, producing unrolled depth N = KR while storing only K blocks.Weight sharing causes updates to aggregate gradients across repeated visits and be read back through those same visits.
  • p = 1/4 is recovered when visits decorrelate, while aligned visits with fixed K require p = 1/2 as loop count grows.The loop-aware bound is controlled by the visit-alignment coefficient κ_R and differs from the untied DeepNorm condition.
  • DeepLoop adds no gates, learned residual coefficients, auxiliary losses, or architecture-specific tuning constants.It is a deterministic residual parameterization for increasing effective depth by revisiting shared blocks.
  • R = 1 is essentially neutral, while larger loop counts improve final validation loss and generally favor downstream averages.Controlled GPT-style ablations cover GPT-2 small and medium scales; downstream results have one exception at an intermediate loop count.
  • DeepLoop uses α = (2N)^1/2 and β = (8N)^-1/2 in a Post-LN looped Transformer with normalized branch inputs.The method is presented as a one-line conservative parameterization derived from the tied-depth analysis.

2 Background

Looped Transformers reuse K physical blocks for R rounds, yielding effective depth N = KR without increasing stored parameters. This tied-depth setup changes residual-scaling behavior because shared parameters aggregate gradient contributions across visits and are read by those visits during unrolled computation.

  • Looped Transformers: A looped Transformer applies K physical blocks for R rounds, giving effective depth N = KR while storing each block’s parameters only once.The studied regime keeps K fixed or small and increases R, raising test-time compute and unrolled depth without adding physical blocks.
  • Looped Transformers: Universal Transformer and ALBERT represent related parameter-reuse extremes, while this work focuses on increasing R with K fixed or small.Universal Transformer reuses the same transition across depth; ALBERT ties Transformer parameters across layers.
  • Normalization and DeepNorm: Pre-LN improves optimization stability, whereas Post-LN can preserve a more expressive residual stream; recurrent-depth models may combine them through sandwich placement.Section 3.1 adopts sandwich placement and states that it leaves the perturbation analysis unchanged.
  • Normalization and DeepNorm: DeepNorm’s useful first-order condition is 2N(β/α)2 = 1/2 for an untied depth-N Transformer.Here N is the number of blocks and M = 2N residual sublayer applications; β is an initialization gain rather than a runtime residual multiplier.
  • Tied-depth effect: Shared looped parameters receive gradient sums over R visits, and aligned visits can make the tied update R times larger than the decorrelated case.The same shared update is read by all R visits during unrolled computation, motivating a smaller β/α ratio to keep its first-order effect bounded.

3 DeepLoop Transformer

DeepLoop adapts Post-LN DeepNorm to looped Transformers by accounting for repeated residual-sublayer parameter visits through the alignment coefficient κ_R. It conservatively raises the scaling exponent from 1/4 to 1/2, using α=(2N)^1/2 and β=(8N)^−1/2 at fixed physical depth.

  • Architecture: A looped Transformer reuses K physical blocks for R rounds, giving unrolled depth N=KR and M=2N residual-sublayer visits.Each physical block contains attention and MLP residual sublayers, and the same physical parameters are reused across visits.
  • Architecture: RMSNorm makes the relevant residual sensitivity ratio β/α because the residual branch enters the normalized direction through a factor of 1/α.DeepLoop retains the post-normalized sandwich block with inner and outer RMSNorms.
  • Stability analysis: The visit-alignment coefficient satisfies 0≤κ_R≤R, with κ_R=O(1) for decorrelated visits and κ_R=Θ(R) for fully aligned visits.Parameter tying causes the shared update to be read through all visits, producing the tied first-order perturbation bound.
  • Scaling rule: The scaling threshold is p=1/4 when visits decorrelate and p=1/2 in the conservative aligned case; at fixed K, DeepNorm’s p=1/4 leaves residual growth of Θ(R).The aligned tied-depth condition is MR(β/α)^2=O(1).

4 Application to Hierarchical Recurrent Reasoners

DeepLoop extends to hierarchical recurrent reasoners by replacing forward visits with gradient-visible visits under one-step truncation and decomposing the stability bound across modules. The binding aligned-inner-cycle regime requires p = 1/2, while decorrelated or expanding-physical-depth regimes permit p ≥ 1/4.

  • Scaling regimes: p ≥ 1/4 suffices when inner-cycle alignment remains bounded, including decorrelated inner visits or fixed CL with growing physical depth.In these cases, ¯κg = O(1) or Θ(1), respectively.
  • Scaling regimes: p ≥ 1/2 is required when physical depths remain fixed and CL →∞, because Mg = Θ(CL) and ¯κg = Θ(CL).This is the fixed-physical-depth tied-loop threshold along the inner-cycle axis.
  • Binding regime: The aligned inner module is the binding case for hierarchical reasoners, so a shared exponent p = 1/2 satisfies both the recurrent L module and the truncated H module.The framework presents this as a worst-case prediction rather than a claim that κ(L)g = Θ(CL) is empirically attained.
  • Gradient-visible analysis: One-step gradient truncation makes the perturbation bound depend on gradient-visible visits rather than all forward visits.Earlier-cycle states are detached, while forward values remain unchanged.
  • Per-module scaling: The hierarchical bound decomposes into per-module summands, allowing independent residual-scaling assignments (αH, βH) and (αL, βL).Each module’s admissible exponent is determined by its own summand rather than the global visit count.

5 Experiments

Experiments compare DeepLoop with a standard pre-LN looped Transformer baseline across GPT-2 scales and a hierarchical recurrent reasoner. DeepLoop matches the baseline without revisiting blocks, then improves validation loss, downstream accuracy, and HRM accuracy when recurrent depth is activated.

  • GPT-2 language modeling: DeepLoop compares the complete post-normalized sandwich-block parameterization against a standard pre-LN looped Transformer with no residual scaling.Both methods share looped residual-block sharing, tied input-output embeddings, and input-embedding RMSNorm.
  • GPT-2 language modeling: At R=1, validation-loss curves coincide within noise, with differences of +0.0004 nats at small scale and +0.0011 nats at medium scale.The runs use FineWeb-Edu 50BT and report final-checkpoint validation cross-entropy.
  • GPT-2 language modeling: DeepLoop beats the baseline at every R ≥3: by −0.016 nats at small scale for R=3, peaking at −0.023 nats for R=5, and reaching −0.028 nats at medium scale by R=7.At small scale, the gap remains −0.019 nats at R=7; these are single-seed runs.
  • Downstream accuracy: On the GPT-2 medium eight-task suite, DeepLoop is tied at R=1, opens the gap at R=3, and achieves the best Avg at R=7 in both shot settings.The 1-shot R=7 Avg is 55.20%; the 0-shot Avg gaps at R=3, R=5, and R=7 are +0.43, +1.06, and +0.93.
  • Hierarchical recurrent reasoner: DeepLoop raises HRM paper-protocol two-vote accuracy from 36.50% to 39.75%, a +3.25 pp gain, and improves every voting-budget column.A four-seed control gives approximately 0.5 pp per-seed standard deviation at K=2, making the gain roughly a 6σ effect.

6 Related Work

DeepLoop builds on shared-depth Transformers, recurrent computation, and established Transformer parameterizations. It specifically corrects residual scaling for correlations caused by reusing residual-branch parameters across loop visits.

  • Looped and universal Transformers: Universal Transformers and ALBERT introduced depth-wise parameter sharing to increase effective computation without increasing parameter count.
  • Looped and universal Transformers: DeepLoop complements in-context, algorithmic, and recurrent-depth architectures by prescribing residual and initialization scales for revisited blocks within a post-normalized sandwich block.
  • Adaptive computation and iterative reasoning: Adaptive Computation Time, PonderNet, recurrent extrapolation, and recent language models use repeated shared computation as differentiable or latent test-time computation.
  • Residual scaling and parameterization: DeepNorm/DeepNet derives α = (2N)1/4 and β = (8N)−1/4 for untied Post-LN Transformers, while DeepLoop corrects the exponent for loop-specific parameter reuse.
  • Residual scaling and parameterization: Existing width-and-depth parameterizations such as µP and Depth-µP address feature learning and hyperparameter transfer but do not explicitly model loop-visit correlation.

7 Conclusion · Appendix

DeepLoop addresses the residual-scaling effects of repeated parameter reuse in Looped Transformers through visit alignment and a conservative exponent p = 1/2. Its scaling rule is neutral without revisits and improves looped GPT-style models as recurrent depth increases, while future work targets direct alignment measurement and broader validation.

  • 7 Conclusion: Repeated visits accumulate a shared residual update and feed it back through the same visits on the next linearized forward pass.This tied-depth effect changes the residual-scaling problem relative to parameter-nonreusing depth.
  • 7 Conclusion: The visit-alignment coefficient κ_R formalizes the tied-depth effect in the residual-scaling analysis.The coefficient captures how repeated parameter visits influence the bound.
  • 7 Conclusion: Decorrelated visits recover the DeepNorm exponent, whereas aligned visits at fixed physical depth require the conservative exponent p = 1/2.The conservative boundary increases as looped depth effects become aligned.
  • 7 Conclusion: DeepLoop uses α = (2N)^1/2 and β = (8N)−1/2 in the post-normalized sandwich block.The correction is implemented as a one-line scaling rule in Eq. (3.1).
  • 7 Conclusion: DeepLoop is neutral when no physical block is revisited and improves validation loss and downstream accuracy as recurrent depth increases.These empirical effects are reported for looped GPT-style language models.
  • Appendix: At R=3, a p-sweep places the empirical stability transition in the vicinity of p = 1/2.The sweep provides empirical support for the conservative exponent boundary.
  • 7 Conclusion: Future work should measure κ_R or cross-round gradient alignment directly and test whether the same boundary holds at larger scale or under alternative parameterizations.The paper also proposes encouraging decorrelated visits to safely use less conservative exponents.

A Proofs 20 … A Proofs

The proofs establish first-order stability conditions for depth-untied and tied looped Transformers, showing that visit alignment introduces an additional loop-count factor and requires p ≥ (1 + γ)/4. For sandwich blocks with inner normalization, the same bounds and worst-case p = 1/2 exponent apply unchanged.

  • A Proofs: Lemma A.1 shows that RMS normalization exposes the residual branch through 1/α under RMS(z)/α ≤ c < 1.The proof uses a first-order Taylor expansion with a remainder uniform under the same condition.
  • A Proofs: Lemma A.2 shows that inner normalization is inert on a post-normalized stream and has Jacobian operator norm 1 when RMS(x) = 1.Its Jacobian is an orthogonal projection scaled by 1/RMS(x).
  • A Proofs: For the sandwich block, every visit has RMS(x_i) = 1, so inner normalization evaluates the branch and preserves the established stability propositions.The same visit-wise bounds give 0 ≤ κ_R ≤ R, with worst-case κ_R = Θ(R) and p = 1/2 unchanged.
  • A Proofs: Depth-untied Transformers satisfy M(β/α)^2 = O(1) as a sufficient first-order stability condition.Here M = 2N counts residual-sublayer visits, and the result is independent of whether token embeddings are tied.
  • A Proofs: Tied looped Transformers satisfy Mκ_R(β/α)^2 = O(1), where κ_R captures visit alignment across repeated physical sublayer visits.The bound follows by summing each physical sublayer’s contribution over J = 2K physical residual sublayers and R visits.
  • A Proofs: In the worst-case aligned regime κ_R = Θ(R), the tied-depth condition reduces to MR(β/α)^2 = O(1).This alignment factor distinguishes tied loop perturbations from the depth-untied bound.
  • A Proofs: At fixed physical depth K, uniform stability as R → ∞ requires p ≥ (1 + γ)/4.Substitution yields N-dependence Θ(N^(1 + γ − 4p)), which remains bounded exactly when 1 + γ − 4p ≤ 0.

B Small-scale downstream evaluation

At GPT-2 small scale, DeepLoop matches the baseline at R=1 but achieves higher eight-task average accuracy at R=3 and R=5. Its headline result is 49.66% for DeepLoop R=5 1-shot, outperforming the same-R baseline on several tasks.

  • Evaluation setup: GPT-2 small uses a 768-wide, 12-layer backbone, the same FineWeb-Edu 50BT schedule, and an eight-task lm-evaluation-harness protocol.Training runs on 4×H200 141 GB GPUs.
  • Results: At R=1, the two methods remain within ±0.25 points on both 0-shot and 1-shot settings.This establishes near-neutral performance without recurrent depth.
  • Results: At R=3, DeepLoop’s Avg exceeds the baseline by +1.31 points for 0-shot and +0.79 points for 1-shot.These are eight-task unweighted average accuracy differences.
  • Results: At R=5, DeepLoop’s Avg exceeds the baseline by +0.15 points for 0-shot and +0.50 points for 1-shot.The strongest reported small-scale cell is DeepLoop R=5 1-shot at 49.66%.
  • Task-level comparison: Against base R=5 1-shot at 49.16%, DeepLoop wins on ARC-C, ARC-E, HellaSwag, and WinoGrande, stays within 0.11 points on three tasks, and loses clearly only on OpenBookQA.The three close tasks are PIQA, SciQ, and SIQA.

C Empirical p-sweep at fixed loop count · D Compute resources

At R=3 on GPT-2 small, the p-sweep places the empirical stability transition near p=1/2, while compute scales roughly linearly with loop count and totals approximately 10,700 H200 GPU-hours project-wide.

  • C Empirical p-sweep at fixed loop count: At fixed R=3, the GPT-2 small sweep varies only the residual-scaling exponent p while holding architecture, data, optimization, and seed protocol constant.The cells use the sandwich block, tied embeddings, input-embedding RMSNorm, FineWeb-Edu 50BT, fixed batch size and learning rate, and the stated seed protocol.
  • C Empirical p-sweep at fixed loop count: The sweep covers p ∈ {0.30, ..., 0.60} and classifies attempts by whether validation loss escapes the unigram-frequency floor of ≈7.67 nats.Figure 6 reports validation-loss snapshots at step 2000, with red crosses for non-escaping attempts and blue circles for training attempts.
  • C Empirical p-sweep at fixed loop count: The empirical transition lies around p=1/2: failures concentrate below it, while the loss penalty for exceeding it increases with p.The paper adopts p=1/2 as the DeepLoop default based on the worst-case analysis.
  • C Empirical p-sweep at fixed loop count: Choices p < 1/2 achieve lower trained-attempt mean loss when they train, but at least 3/5 seeds fail to escape the floor.This shows the tradeoff between stronger learning signals and training stability below the selected exponent.
  • C Empirical p-sweep at fixed loop count: The p-sweep is limited to GPT-2 small, R=3, and one optimizer-step budget, leaving larger scales, other normalization placements, and longer training for future work.Therefore, the observed p=1/2 boundary is not established across those untested settings.
  • D Compute resources: Experiments ran on NVIDIA H200 141 GB GPUs in a SLURM cluster, using 4×H200 for GPT-2 small cells and 8×H200 for GPT-2 medium cells.Per-cell wall-clock time comes from WandB run records, while aggregate GPU-hours are rounded to the nearest hundred.
  • D Compute resources: Per-cell wall-clock time scales roughly linearly with loop count R.
  • D Compute resources: The reported 16-cell FineWeb-Edu matrix totals approximately 3,700 H200 GPU-hours, while unreported exploratory experiments add approximately 7,000 GPU-hours.The matrix spans base and DeepLoop models at small and medium scales with R ∈ {1, 3, 5, 7}.
Loading 2607.13491v2…