PaperlayerView exact revision on arXiv

Source-linked AI summary

How Much Is One Recurrence Worth? Iso-Depth Scaling Laws for Looped Language Models

Kristian Schwethelm, Daniel Rueckert, Georgios Kaissis

arXiv:2604.21106v3cs.LGcs.CL

TL;DR

Looped transformers may not substitute for unique blocks at matched compute, but evidence has been limited on the value of parameter sharing. This paper fits an iso-depth joint scaling law and finds that one recurrence is worth roughly half a unique block, with φ = 0.46.

  • Problem

    Evidence is limited on whether a transformer block looped r times can replace r unique blocks at matched compute.

  • Method

    The paper runs an iso-depth pretraining sweep over recurrence counts r ∈ {1, 2, 4, 8} and fits a joint scaling law with recurrence-equivalence exponent φ.

  • Results

    φ = 0.46, so each shared recurrence is worth roughly half a unique block; hyperconnections raise φ to 0.65, while truncated BPTT lowers it to 0.38.

  • Takeaways & Limitations

    ∆φ distinguishes architecture-side capacity gains from training-side efficiency gains that raw validation loss cannot separate, while reflecting deployment inference cost.

  • Takeaways & Limitations

    The fitted φ is architecture-, optimizer-, and recipe-dependent, and values from different non-looped baselines are not comparable.

Abstract

from arXiv · show

We measure how much one recurrence is worth to a looped (depth-recurrent) transformer, in equivalent unique parameters. From an iso-depth pretraining sweep across recurrence counts $r \in \{1, 2, 4, 8\}$ spanning ${\sim}50\times$ in training compute, we fit a joint scaling law $L = E + A\,(N_\text{once} + r^{\varphi} N_\text{rec})^{-α} + B\,D^{-β}$ and measure a recurrence-equivalence exponent $\varphi = 0.46$. Intuitively, $\varphi$ tells us whether looping a block $r$ times is equivalent in validation loss to $r$ unique blocks of a non-looped model (full equivalence, $\varphi{=}1$) or to a single block run repeatedly with no capacity gain ($\varphi{=}0$). Our $\varphi = 0.46$ sits in between, so replacing unique blocks with shared recurrences increases validation loss at matched training compute. For example, at $r{=}4$ a 410M looped model performs on par with a 580M non-looped model, but incurs the training cost of a 1B non-looped one. We demonstrate the utility of $\varphi$ as a diagnostic tool on two case studies: commonly used truncated backpropagation lowers $\varphi$ to $0.38$, indicating that the loop mechanism is poorly trained under truncation, even though validation loss decreases. Conversely, hyperconnections raise $\varphi$ to $0.65$, a genuine capacity gain. Our method separates true loop improvements from training-side gains, a distinction raw validation loss cannot make.

1 Introduction

The paper introduces a recurrence-equivalence exponent φ to measure how many unique parameter blocks one recurrence is worth under matched compute. A baseline φ = 0.46 shows partial capacity equivalence, while diagnostic shifts distinguish training-side from architecture-side gains.

  • Iso-depth scaling sweep: r ∈ {1, 2, 4, 8} defines an iso-depth sweep whose variants execute the same forward layers per token and incur the same per-token FLOPs at matched width.Here, r=1 is the non-looped baseline, while unique non-embedding parameters drop by 3.2× as recurrence increases.
  • Reference points: φ = 1 represents full parameter equivalence, whereas φ = 0 represents no capacity gain from repeatedly running a shared block.The joint law uses effective parameters Nonce + r^φNrec to quantify the contribution of recurrent parameters relative to unique ones.
  • Diagnostic applications: Δφ separates training-side from architecture-side gains: truncated backpropagation lowers φ to 0.38, while hyperconnections raise φ to 0.65 despite both decreasing validation loss.Truncated backpropagation saves ∼30% training FLOPs, whereas hyperconnections add parallel residual streams between loops.
  • Recurrence-equivalence exponent: φ = 0.46 means each recurrence is worth roughly half a unique parameter block at the same FLOPs.The exponent is defined through a joint scaling law over recurrence count, shared parameters, single-use parameters, and training tokens.

2 Related Work

Prior work established looped language models through depth-wise weight sharing and explored their use for latent reasoning and test-time compute scaling. Related scaling-law work fixes unique parameter count, whereas this paper studies recurrence under increasing depth, inference FLOPs, and KV-cache costs.

  • Weight sharing across depth was introduced by the Universal Transformer as the foundation for looped language models.
  • ∼3B parameters and trillion-token training budgets were reached by Huginn and Ouro, which reported strong downstream results for looped language models.
  • Prairie et al. [10] fit scaling laws at fixed unique parameter count, tracing compute-optimal recurrence under a limited memory budget.
  • In contrast, this setup allows depth, per-token inference FLOPs, and KV cache memory to grow with recurrence count, measuring per-recurrence sharing value.

3 Methodology

The methodology compares non-looped and looped transformers at fixed effective depth and approximately matched per-token compute, while quantifying loop capacity through an effective parameter count. This framework separates parameter savings from the loss contribution of repeated recurrent blocks.

  • Iso-depth comparison: Four variants use 20 effective layers: a non-looped baseline at r=1 and looped models at r ∈ {2, 4, 8}.At matched width, per-token training and inference FLOPs are approximately equal across recurrence counts, apart from a small input-injection correction.
  • Architecture: The prelude-recur-coda design fixes (ℓprelude, ℓcoda) = (2, 2), leaving ℓrecur = 16/r recurrent layers executed r times.This gives (8, 4, 2) recurrent layers for r ∈ {2, 4, 8}.
  • Parameter accounting: At the same width, looped models retain ∼61%, ∼41%, and ∼31% of the parameters of a non-looped model for r ∈ {2, 4, 8}.At s=10, the corresponding parameter counts are N ∈ {98.3, 59.8, 40.2, 30.3} M.
  • Compute matching: Ffwd(r) = 2(Nonce + r Nrec) ≈ Ffwd(1) = 2 ℓeff nb = 2N, so compute is approximately iso-FLOPs across variants.The injection-layer overhead is r/120 ∈ {1.7%, 3.3%, 6.7%} at r ∈ {2, 4, 8}; training FLOPs satisfy Ftrain(r) = 3 Ffwd(r).
  • Recurrence-equivalence model: The extended scaling law defines Neff ≡ Nonce + r^φNrec as the effective parameter count, with r^φ amplifying recurrent parameters’ contribution to loss reduction.Neff represents the parameter count of a non-looped model matching the looped model’s loss at the same D.
  • Recurrence-equivalence model: φ = 0 means repeated execution adds FLOPs without capacity gain, whereas φ = 1 corresponds to full equivalence with unique parameter blocks under iso-FLOPs.Values 0 < φ < 1 quantify partial recovery of unique-parameter capacity.

4 Iso-Depth Scaling Laws

Across fixed-compute sweeps, looped architectures require wider models yet achieve a progressively worse compute-optimal loss frontier as recurrence count increases. The joint scaling law estimates partial recurrence-equivalence, with φ = 0.46: recurrence adds capacity, but less than unique parameters.

  • Fixed-compute scaling: At fixed compute, larger recurrence counts produce upward-shifted, flatter iso-FLOPs curves, with looped minima at wider widths than the non-looped baseline.The standard Chinchilla law is fit separately for each architecture.
  • Compute-optimal allocation: At r=2, r=4, and r=8, the compute-optimal loss frontier trails the baseline by [0.03, 0.06], [0.05, 0.08], and [0.09, 0.12] nats, respectively.The gap grows monotonically with r, widens at lower budgets, and changes by ≤0.006 nats between 10^19 and 2.15 × 10^19 FLOPs.
  • Joint scaling law: φ = 0.46 places recurrence between no capacity gain (φ = 0) and full equivalence with unique blocks (φ = 1).The joint law uses six shared parameters across all 116 runs and scales recurrent parameters by r^0.46.
  • Joint scaling law: At r=4 and r=8, recurrent blocks contribute ∼1.9× and ∼2.6× their unique parameter counts, respectively.Because φ > 0 but φ < 1, recurrence adds partial capacity and allows fewer tokens at equal performance and parameter count.
  • Robustness: The 95% block-bootstrap CI for φ is [0.41, 0.53], with no resample reaching φ = 0 or φ = 1.Per-r residual RMSE is comparable across architectures at 0.009–0.011 nats.

5 Case Studies

The case studies show that recurrence-equivalence φ distinguishes training-side improvements from genuine architectural capacity gains. Truncated BPTT lowers validation loss but reduces φ, whereas hyperconnections improve both validation loss and recurrence equivalence.

  • Truncated BPTT: 30%: Truncated BPTT approximately reduces per-token training cost by skipping backward passes for early recurrences.The setup detaches the recurrent state for all but the last rbwd loops, with rbwd = ⌈r/2⌉.
  • Truncated BPTT: 0.38: Truncated BPTT lowers φ from 0.45 despite substantially reducing validation loss across runs.With rbwd = ⌈r/2⌉, early recurrences skip backward passes, likely receiving less accurate learning signals.
  • Hyperconnections: 0.65: Hyperconnections raise φ from 0.45 while substantially lowering validation loss across looped runs.The intervention uses K=2 parallel residual lanes with full BPTT.
  • Hyperconnections: r=2: Hyperconnections can match or beat the r=1 baseline at some budgets, but this does not establish φ = 1.Full equivalence would require all four architectures to lie on the same compute-optimal frontier.
  • Hyperconnections: Hyperconnections move the compute-optimal allocation to narrower widths, lowering per-token inference FLOPs, unlike the widening seen under truncated BPTT.They replace linear input injection with scaled and mixed K=2 residual lanes across loops, and are applied only at the loop boundary.

6 Discussion

The discussion frames φ as a recipe-dependent measure that distinguishes architectural capacity from training-efficiency gains, while linking it to inference cost and outlining empirical and modeling limitations. It also finds no above-noise reasoning signal at the tested budgets and identifies several extensions for future study.

  • Per-recurrence value: φ = 0.46 means that at r=4, the shared block recovers 40.46 ≈1.89 unique blocks of capacity, about 47% of full equivalence.The result depends jointly on architecture, optimiser, and training recipe, rather than looping alone.
  • Downstream validation: Reasoning-heavy downstream tasks show no above-noise architectural signal at the tested budgets, including ∼20× extrapolation runs, leaving φ’s link to scaled reasoning quality untested.Parametric-knowledge tasks track validation loss directly.
  • ∆φ as a development metric: ∆φ separates token-side training gains from architecture-side capacity gains, which raw validation loss cannot distinguish.Truncated BPTT and hyperconnections illustrate the two pure cases at opposite ends.
  • Inference cost: Higher φ enables fewer unique parameters at the same loss, narrower compute-optima, and lower per-token inference FLOPs; lower φ has the opposite effect.The freed compute under higher φ can instead be spent on more training tokens.
  • Limitations: The study fixes 20 effective layers with (ℓprelude, ℓcoda) = (2, 2), caps recurrences at rmax = 16, and treats r^φ as a pre-saturation local approximation.Different depth allocations or prelude/coda sizes may shift φ, and a single φ cannot capture mixed outcomes across recurrence counts.

7 Conclusion … C.1 Implementation Details

The paper introduces the recurrence-equivalence exponent φ to quantify the parameter-sharing cost of looped language models, finding φ = 0.46 on the main baseline. It also shows that design choices change φ, while documenting related work, compute resources, and implementation details.

  • 7 Conclusion: φ = 0.46 on the prelude-recur-coda baseline, making each shared recurrence worth roughly half a unique block.Hyperconnections raise φ to 0.65, whereas truncated BPTT lowers it to 0.38, even though both lower validation loss.
  • A Extended Related Work: The paper extends prior scaling-law analyses to looped architectures.Earlier work established power-law relations among loss, model size, and training tokens, while later studies examined compute-optimal allocation, learning-rate transfer, and inference-aware scaling.
  • A Extended Related Work: The method separates unique parameters N_once from recurrent parameters N_rec, unlike iso-parameter accounting that multiplies the full parameter count by recurrence.The contrasted framework uses N_eff = μ_recN, including prelude and coda parameters.
  • A Extended Related Work: The main grid uses fixed architectural recurrence counts and full BPTT, while truncated BPTT is evaluated separately because it reduces training FLOPs.This keeps training and inference FLOPs aligned with the matched non-looped baseline.
  • A Extended Related Work: 8 loops are worth ≈2.6 unique blocks at matched depth, quantifying why additional loops can add less capacity than unique parameters.Prior compute-matched studies commonly used r ≤4 because each additional loop carries a large training-FLOPs cost.
  • B Compute Resources: Approximately 5,000 GPU-hours covered the full project, including the 116-run iso-depth grid, two case studies, exploratory configurations, failed runs, and side experiments.Experiments used a mix of A100-80GB and H100-80GB GPUs.
  • C Model Architecture: The architecture maintains shared effective depth 20 through ℓprelude + r · ℓrecur + ℓcoda = 20, with (ℓprelude, ℓcoda) = (2, 2) for r > 1.For r ∈ {2, 4, 8}, ℓrecur = 16/r evaluates to {8, 4, 2}; the recurrent block writes its output back into the latent state via the injection layer.
  • C.1 Implementation Details: Each looped architecture applies input injection at every recurrence iteration, with W_inject initialized as [I ∥0] so u(0) ≈ e at training start.Three model-level RMSNorms control residual-stream scale after embedding, after each recurrence iteration, and before the lm_head.

C.2 Input-Injection Ablation … F Scaling Law Fit Diagnostics

The paper selects linear input injection for the main scaling-law runs, establishes η*=0.014 as a transferable base learning rate, and evaluates iso-compute configurations across recurrence counts. Scaling-law diagnostics use robust multi-start fitting and residual, bootstrap, and budget-half stability checks.

  • C.2 Input-Injection Ablation: At r=4 and target compute 10^18 FLOPs, parameter-free injection alternatives are compared with the default linear map.Passthrough uses no injection, while Additive provides parameter-free residual injection; both share the reference configuration and target FLOPs budget.
  • C.2 Input-Injection Ablation: 973M vs. 955M training tokens gives parameter-free alternatives a ~2% data advantage, yet Passthrough fails to train and Additive trails linear injection by 0.004 nats.The linear injection is therefore adopted for all reported scaling-law runs.
  • D.1 Learning Rate Sweep: η*=0.014 is the shared optimum for both architectures in the s=10 learning-rate sweep and becomes the base learning rate for subsequent experiments.The sweep uses ratio 10, B = 256K, and eight values in η ∈[0.008, 0.024].
  • D.2 Transfer Validation: 0.004 nats is the maximum width-transfer regret when using ηbase = 0.014, with all minima clustering near that value across s ∈{8, 10, 14}.At s=18, ηbase = 0.014 is marginally better than 0.012, yielding 2.473 vs. 2.476 nats.
  • D.2 Transfer Validation: Below 0.005 nats, regret at ηbase = 0.014 remains low across ratios {10, 20, 40} for both architectures.The data-scaling validation spans a 4× range in training tokens and tests the D^-0.32 correction.
  • E Iso-Depth Grid: For each compute budget and r ∈{1, 2, 4, 8}, model width is swept to identify the compute-optimal point, recording unique non-embedding parameters and training tokens.Looped models train on slightly fewer tokens because of input-injection compute overhead; empty cells are untested widths.
  • F Scaling Law Fit Diagnostics: Residuals, joint-fit statistics, block-bootstrap confidence intervals for φ, and stability across budget halves are used to assess scaling-law robustness.Per-architecture fit quality is also examined through predicted-versus-actual validation loss, while Table 7 reports joint-law residuals in nats across 29 runs per architecture.
  • F Scaling Law Fit Diagnostics: 500 random L-BFGS-B restarts are used for both per-architecture and joint fits because the Huber-on-log objective is non-convex.Parameters start uniformly within bounded boxes, with a per-restart cap of 10,000 iterations.

F.1 Per-Architecture and Joint Fit Residuals … G.2 Fit Quality

The scaling-law fits are accurate across architectures and compute scales, while the r=4 example quantifies recurrence’s parameter and training-cost tradeoff. Case-study fits remain strong overall, but truncated backpropagation introduces architecture-specific residuals that hyperconnections avoid.

  • F.1 Per-Architecture and Joint Fit Residuals: 0.007 nats is the maximum residual for the four per-architecture Chinchilla fits, with no systematic bias across N or D.Predicted and actual validation losses cluster tightly around the diagonal.
  • F.1 Per-Architecture and Joint Fit Residuals: 0.036 nats is the joint fit’s maximum absolute residual, with pooled RMSE = 0.010 nats and R2 = 0.997 across 116 runs.The six-parameter joint law matches all four architectures with comparable RMSE across r (0.009–0.011 nats).
  • F.2 Bootstrap Procedure: φ = 0.46 has a 95% CI of [0.41, 0.53], estimated by block-bootstrap resampling of 200 budget–architecture cell samples.The procedure preserves experimental block structure by resampling non-empty cells rather than individual runs.
  • F.3 Stability Across Budget Halves: φ = 0.44 on the low-budget half and φ = 0.49 on the high-budget half show no drift with scale within the compute window.The bootstrap CI comfortably contains both half-window estimates.
  • F.4 Example: Equivalent Model Sizes at r=4: 0.40× is the r=4 variant’s unique-parameter ratio, corresponding to approximately 410M parameters including the injection term.This comparison uses matched width and effective depth against a 1B r=1 configuration.
  • F.4 Example: Equivalent Model Sizes at r=4: ≈580M is the r=4 variant’s effective-parameter equivalent, while its per-step training cost matches the 1B non-looped baseline up to ∼3% injection overhead.The executed-layer count is identical at fixed dmodel.
  • G.1 Methods: 30% fewer training FLOPs per token from truncated backpropagation are reinvested as more tokens, yielding Dtrunc/Dfull = 1.315.The method detaches earlier recurrent states and backpropagates only through the final rbwd = ⌈r/2⌉ iterations.
  • G Case Study Details: R2 = 0.983 for truncated BPTT is driven by systematic r=2 residuals, whereas hyperconnections show uniformly small residuals across r.Refitting truncated BPTT on r ∈{4, 8} raises R2 from 0.983 to 0.996 without changing φ; hyperconnections’ main limitations are fewer runs and a narrower compute span.

G.3 Compute-Optimal Allocation Under the Case Studies · H Downstream Evaluation Suite · H.1 Setup

Compute-optimal width shifts track changes in recurrence-equivalence φ, with truncated BPTT widening models and hyperconnections narrowing them. The downstream suite evaluates five capability axes using standardized few-shot settings across four architectures.

  • G.3 Compute-Optimal Allocation Under the Case Studies: Higher φ lowers effective parameter amplitude and compute-optimal width at recurrence r > 1.The joint law links gr, Aeff, and N*(C), making φ relevant to compute-optimal allocation.
  • G.3 Compute-Optimal Allocation Under the Case Studies: 0.38 φ under truncated BPTT widens compute-optimal models and trades training FLOPs for inference FLOPs.This follows φ falling from 0.45, with compute-optimal stars shifting right relative to full BPTT.
  • G.3 Compute-Optimal Allocation Under the Case Studies: 0.65 φ under hyperconnections contracts compute-optimal width and lowers per-token inference FLOPs at matched compute.The hyperconnections stars sit at smaller N* than corresponding full-BPTT linear-injection stars.
  • G.3 Compute-Optimal Allocation Under the Case Studies: Δφ alone determines the direction of compute-optimal width shifts across both case studies, supporting the joint law as a budget-allocation tool.The paper observes both shift directions within the same architecture family.
  • H.1 Setup: The downstream suite partitions tasks into five mechanistically motivated axes to isolate capability dimensions and expose architectural biases.Tasks come from CORE, the Saunshi suite, and in-house probes.
  • H.1 Setup: All four architectures use the same shot count and prompts on every task, with CoQA reduced to 1-shot because full-passage prompts exceed the 2,048-token context.Few-shot counts match or approximate canonical source-benchmark settings.
  • H.1 Setup: The suite covers parametric knowledge, reading comprehension, math word problems, reasoning primitives, and compositional symbolic manipulation.Examples include TriviaQA, Lambada-OpenAI, SVAMP, induction-head and variable-assignment probes, and BigBench symbolic tasks.

H.2 Compute-Optimal Per-Axis Results … H.5 Accuracy versus Continuation Loss

Compute-optimal downstream behavior separates into capacity-bound, gap-closing, and unresolved regimes across recurrence counts, while continuation loss provides a more reliable signal than accuracy. Per-task and validation-loss analyses reveal architecture-specific orderings that aggregate axes can obscure.

  • H.2 Compute-Optimal Per-Axis Results: Compute-optimal evaluation uses the lowest-validation-loss checkpoint at each FLOPs budget and reports per-token continuation loss across five downstream axes.Recurrence counts are r ∈ {1, 2, 4, 8}.
  • H.2 Compute-Optimal Per-Axis Results: Parametric knowledge is capacity-bound: r=1 leads at every compute budget, with the gap reaching 0.28 nats at r=8.Its ordering follows validation loss, consistent with φ = 0.46 leaving less unique-parameter capacity for knowledge storage as recurrences increase.
  • H.2 Compute-Optimal Per-Axis Results: Reading comprehension and compositional symbolic largely close the gap: r ∈ {2, 4} match r=1 on reading comprehension, while symbolic aggregates are roughly tied.On reading comprehension, only r=8 trails by 0.05–0.18 nats; symbolic per-task outcomes are mixed.
  • H.2 Compute-Optimal Per-Axis Results: Reasoning primitives and math word problems do not resolve a per-r signal at the tested budgets.The r=1 baseline leads at nearly every reasoning-primitives budget, while math-word-problem separation remains within noise.
  • H.3 Per-Task Continuation Loss: Per-task results reproduce the axis-level patterns: r=1 is lowest on all parametric-knowledge tasks, while reading-comprehension orderings vary by task.TydiQA-GoldP, SQuADv2, DROP, and CoQA favour r=4, whereas Lambada-OpenAI is monotone in r=1.
  • H.4 Per-Axis Continuation Loss versus Validation Loss: Per-axis continuation-loss curves can reorder recurrence counts relative to the main compute-optimal figure because architectures reach a given validation loss with different (N, D) allocations.Figure 12 plots continuation loss against validation loss for every iso-Depth checkpoint.
  • H.5 Accuracy versus Continuation Loss: Continuation loss tracks validation loss nearly linearly, whereas aggregate accuracy is noisier and flat at small scales.Small-scale tasks often sit near the random-chance accuracy floor, making accuracy a coarse, bimodal signal.

I Extrapolation Beyond the Grid

An extrapolation run at s=34 and 47B tokens tests the iso-depth findings at approximately 20× the grid’s maximum training compute. The r=4 model still trails r=1 in validation loss, while downstream results preserve the earlier three-regime pattern.

  • Extrapolation setup: At s=34 and 47B tokens, matched-token r=1 and r=4 runs extend the test to approximately 20× the grid’s maximum training compute.Both runs use dmodel = 2,176; the r=4 model receives an approximately 3% training-FLOPs advantage from its injection layer.
  • Extrapolation result: 0.061 nats is the r=4 validation-loss gap relative to r=1, remaining within the 0.05–0.08 nats r=4 band measured on the iso-FLOPs grid.The reported gap is r=4 minus r=1, so positive values indicate that the looped model trails.
  • Downstream behavior: Downstream results preserve the three-regime pattern: parametric knowledge retains a capacity cost, open-book axes track validation loss, and reasoning primitives show no signal favoring r=4.These downstream patterns match those reported in Section H.2.
Loading 2604.21106v3…