PaperlayerView exact revision on arXiv

Source-linked AI summary

Parcae: Scaling Laws For Stable Looped Language Models

Hayden Prairie, Zachary Novack, Taylor Berg-Kirkpatrick, Daniel Y. Fu

arXiv:2604.12946v1cs.LG

TL;DR

Looped architectures can scale compute without increasing parameter count, but prior training methods suffer instability. Parcae models looping as a dynamical system, constrains recurrent injection parameters for stability, and achieves stronger quality with predictable training and test-time scaling laws.

  • Problem

    Prior looped architectures offer parameter-efficient compute scaling but suffer residual explosion, loss spikes, and unstable training.

  • Method

    Parcae recasts looping as a nonlinear time-variant dynamical system, linearizes it for stability analysis, and constrains A through discretization of a negative diagonal parameterization.

  • Results

    Parcae reduces validation perplexity by up to 6.3% versus prior RDMs and improves Core and Core-Extended by 2.99 and 1.18 points versus parameter-matched Transformers.

  • Takeaways & Limitations

    FLOP-optimal training increases looping and data together following predictable power laws, while test-time looping follows a saturating exponential decay law.

  • Takeaways & Limitations

    The reported scaling-law observations are limited to small architectures, leaving their behavior at large FLOP budgets and parameterizations unresolved.

Abstract

from arXiv · show

Traditional fixed-depth architectures scale quality by increasing training FLOPs, typically through increased parameterization, at the expense of a higher memory footprint, or data. A potential alternative is looped architectures, which instead increase FLOPs by sending activations through a block of layers in a loop. While promising, existing recipes for training looped architectures can be unstable, suffering from residual explosion and loss spikes. We address these challenges by recasting looping as a nonlinear time-variant dynamical system over the residual stream. Via a linear approximation to this system, we find that instability occurs in existing looped architectures as a result of large spectral norms in their injection parameters. To address these instability issues, we propose Parcae, a novel stable, looped architecture that constrains the spectral norm of the injection parameters via discretization of a negative diagonal parameterization. As a result, Parcae achieves up to 6.3% lower validation perplexity over prior large-scale looped models. Using our stable looped architecture, we investigate the scaling properties of looping as a medium to improve quality by increasing FLOPs in training and test-time. For training, we derive predictable power laws to scale FLOPs while keeping parameter count fixed. Our initial scaling laws suggest that looping and data should be increased in tandem, given a fixed FLOP budget. At test-time, we find that Parcae can use looping to scale compute, following a predictable, saturating exponential decay. When scaled up to 1.3B parameters, we find that Parcae improves CORE and Core-Extended quality by 2.99 and 1.18 points when compared to strong Transformer baselines under a fixed parameter and data budget, achieving a relative quality of up to 87.5% a Transformer twice the size.

1 Introduction

Looped models offer a way to scale compute without proportionally increasing parameters, but prior training recipes can be unstable. Parcae stabilizes looping and shows predictable training and test-time scaling behavior.

  • Motivation: Looped architectures increase FLOPs by repeatedly sending activations through a layer block, potentially matching larger fixed-depth models without increasing parameters.They also offer potential for latent reasoning and per-token adaptive compute.
  • Motivation: Prior looped models suffer residual-state explosion, loss spikes, and sensitivity to hyperparameter selection and residual normalization.Variable-depth training can produce additional late-training loss spikes.
  • Approach: Parcae recasts looping as a nonlinear time-variant dynamical system and uses a linear approximation to connect instability with injection-parameter spectral norms.The framework identifies unstable parameterizations of A as inducing residual-stream explosion.
  • Approach: Parcae constrains A through a discretized negative diagonal parameterization, normalizes input injection, and uses per-sequence depth sampling to reduce training instability.These changes target residual explosion and loss spikes while preserving looped computation.
  • Results: 6.3% lower validation perplexity is achieved against parameter- and data-matched RDMs, while 1.3B-parameter Parcae improves Core and Core-Extended by 2.99 and 1.18 points versus matched Transformers.The latter result matches Transformers up to twice the size.
  • Results: FLOP-optimal training increases looping and data together through empirical power laws, while test-time looping follows predictable exponential decay with an irreducible loss.The study evaluates end-to-end quality, training FLOP scaling, and test-time scaling.

2 Background

The paper situates looped models within recurrent and linear dynamical-system frameworks and reviews scaling-law methods for allocating parameters, data, and compute.

  • Looped Architectures: Prior looped architectures vary in loop placement, halting mechanism, topology, and differentiation; this work focuses on implicit-halting middle-looped models with explicit differentiation.The review distinguishes pre-, mid-, and post-looping, explicit routers and stochastic depth, and single-block and hierarchical topologies.
  • Looped Architectures: A middle-looped architecture consists of a prelude P, recurrent unit R, and coda C, with the recurrent unit repeatedly transforming the embedded input.The formulation maps input s to e = P(s), iterates h, and produces output p = C(h_T).
  • Looped Architectures: Training samples recurrence depth T from a distribution Λ, while truncated backpropagation limits the backward pass to a constant depth.Variable-depth training is intended to support stronger test-time scaling.
  • Dynamical Systems: The paper uses an LTI dynamical system as a tractable linear surrogate for analyzing complex nonlinear looped models.The surrogate describes hidden-state evolution under input-driven dynamics.
  • Dynamical Systems: Discrete LTI stability depends on the spectral radius of A: systems are stable when ρ(A) < 1 and unstable when ρ(A) > 1.The reviewed regimes are stable, marginally stable, and unstable.
  • Scaling Laws: Scaling-law analysis fits loss as a function of model size and training tokens under fixed FLOP budgets using parabolic and parametric methods.The parametric form includes irreducible loss and power-law terms for parameters and data.

3 Understanding Instability in Looped Architectures

The instability analysis models the residual stream as a dynamical system and identifies unconstrained recurrent-state dynamics, especially the spectral radius of A, as the critical stability issue.

  • Training Instability: Residual normalization can mitigate state explosion but requires sensitive hyperparameter tuning, and large-scale runs can still exhibit late-training loss spikes.The paper reports such spikes after 170k steps and addresses them by normalizing the input to B.
  • State Explosion from Unconstrained A and B: Table 1 compares prior loop update rules by their LTI stability properties, emphasizing whether their A parameterizations satisfy the spectral-radius condition.The comparison covers addition and concatenation injection within the same framework.
  • Dynamical System over Residual Stream: An exact nonlinear residual-stream update is linearized into a discrete LTI system, making stability conditions analytically tractable.The full recurrent update includes transformer operations through R, while the linearized form drops R.
  • State Explosion from Unconstrained A and B: Prior looped architectures use A parameterizations with ρ(A) = 1 or unconstrained ρ(A), corresponding to marginally stable or unstable systems.These choices do not enforce the conditions required for bounded recurrent dynamics.
  • State Explosion from Unconstrained A and B: Divergent runs learn ρ(A) ≥ 1, whereas convergent runs maintain ρ(A) < 1, empirically linking spectral-radius constraints to residual stability.The analysis identifies ρ(A) as the critical factor governing instability.

4 Parcae: A Stable Looped Architecture

Parcae enforces stable recurrent dynamics through constrained continuous-to-discrete parameterizations and further reduces variance from variable-depth training with per-sequence sampling.

  • Architecture and Training: Parcae explicitly satisfies the dynamical-system stability constraints and adds per-sequence depth sampling to stabilize variable-depth training.The method targets both recurrent instability and variance introduced by sampled depth.
  • Stable Recurrent Parameterization: Continuous A and B are discretized with a learned step size using ZOH or Euler schemes, yielding discrete recurrent parameters.The resulting update uses A and B in the loop recurrence.
  • Stable Recurrent Parameterization: A is parameterized as a negative diagonal matrix before discretization, ensuring negative eigenvalues while keeping the stability mechanism simple and inexpensive.The learnable vector is transformed through Diag(−exp(·)); B remains unconstrained and the input e is normalized.
  • Architecture and Training: Parcae normalizes the embedded input e before recurrence and computes outputs by applying the coda projection after T loops.Its recurrent update is ht+1 = Aht + Be + R(ht, e).
  • Architecture and Training: Parcae matches prior RDM and Transformer setups while differing through residual normalization choices and dynamical-system parameters such as A, B, C, and ∆.Against Transformers, the middle third of layers is looped and a prelude norm is included.
  • Stable Training Algorithm: Per-sequence depth sampling provides more depth samples within a micro-batch, and the training design is intended to reduce loss spikes and improve recurrence extrapolation.The method parameterizes Λ using µ_rec and sets µ_bwd = ⌈µ_rec^2⌉.

5 Results

Parcae is evaluated against prior looped models and fixed-depth Transformers on end-to-end quality and stability. It improves perplexity and downstream scores under matched parameter and data settings.

  • Parcae is compared with parameter- and data-matched RDMs and Transformers using held-out perplexity and downstream benchmarks.
  • Comparison against RDMs: Parcae reduces perplexity by up to 6.2% and 9.1% on validation and WikiText, respectively, versus prior RDMs.
  • Ablations: Each Parcae modification contributes to stability or quality: constrained A enables convergence at high T, sampling stabilizes lower depths, and prelude normalization improves quality.
  • Comparison Against Transformers: Parcae reduces validation perplexity by 4.3–9.2% and improves Core and Core-Extended scores by up to 2.99 and 1.18 points versus Transformers.
  • Comparison Against Transformers: A 770M Parcae model achieves Core quality comparable to a 1.3B Transformer with roughly half the parameters.

Parametric IsoLoss Contours for Looping IsoFLOP Curves for Looping

Looping provides an additional compute-scaling axis for both training and inference. Training benefits from increasing recurrence and data together, while test-time gains follow a predictable saturating exponential pattern.

  • Training FLOP Scaling: Compute-optimal training increases mean recurrence µrec and data in tandem, establishing looping as an orthogonal scaling axis.
  • Training FLOP Scaling: At 140M and 370M scales, increasing µrec while proportionally reducing tokens lowers validation loss under fixed FLOPs.
  • Training FLOP Scaling: Optimal recurrence and token budgets follow predictable power laws with exponents γµ ≈0.40 and γD ≈0.78.
  • Training FLOP Scaling: The scaling-law fit predicts held-out validation loss with 1.3% error at 140M and 0.8% error at 370M.
  • IsoFLOP Comparison: The optimal looped frontier achieves strictly lower loss than fixed-depth models and yields 1.2–2.0 points higher Core scores.
  • Test-Time Scaling: Test-time looping follows a saturating exponential decay, with gains plateauing near the training recurrence µrec.
  • Unified Scaling Law: A unified fit predicts held-out test-time loss within 0.85–1.31% average error, or 0.1–0.17% when using empirical loss at T = µrec.

6 Discussion and Future Work

The paper’s scaling conclusions remain bounded by the scale and design choices studied. Larger architectures, deeper recurrence, alternative parameterizations, and efficient inference require further investigation.

  • Looped Architectures: Larger-scale investigation is needed for loop-unit placement, recurrent-unit composition, and extreme looping at deeper recurrence depths.
  • Looped Architectures: Different discretizations, full-rank parameterizations, and recurrent update rules may enable recurrence at larger depths.
  • Scaling: The study’s scaling observations are limited to small architectures, so performance at large FLOP budgets and parameterizations remains unresolved.
  • Scaling: As µrec increases, more test-time steps are required to achieve equivalent quality, motivating methods that reduce inference steps.

7 Conclusion

The paper frames looped-model stability through dynamical systems and introduces Parcae to constrain residual dynamics. It also develops FLOP accounting and implementation procedures for recurrent computation.

  • Conclusion: Parcae prevents residual explosion by constraining the spectral norm of looped-model injection parameters.The paper presents this as the central stability mechanism for its looped architecture.
  • Dynamical-systems formulation: Looped-model recurrence can be represented as a nonlinear time-variant dynamical system over the residual stream.The framework identifies A, B, and C as dynamical-system parameters and R as the nonlinear transformer update.
  • Dynamical-systems formulation: Linearizing the recurrence connects stability to the spectral properties of the residual-state transition matrix.For Pre-Norm additive injection, the linearized transition has A = I, making the system marginally stable.
  • FLOP estimation: Parcae’s effective training FLOPs account separately for parameters that are and are not backpropagated through recurrent depth.The formulation is C = (2 N̂1 + 6 N̂2)D, with embedding parameters excluded from N̂ and attention FLOPs additionally estimated.
  • Implementation: The implementation provides forward-pass and per-sequence stochastic-depth procedures for efficiently training variable recurrent depths.Batched sequences use the maximum sampled depth while avoiding initial state updates to preserve efficient gradient flow.

F Additional Stability Ablations

Additional ablations examine learning-rate stability, sampling granularity, and truncated-recurrence sampling. They find that per-sequence sampling reduces loss spikes and that correcting distributional truncation improves recurrence extrapolation.

  • Stability across learning rates: Parcae remains stable across the tested learning-rate sweep for baseline and residual-normalized RDM comparisons.The sweep covers learning rates from 2e−4 through 1e−3.
  • Per-sequence sampling: Per-sequence sampling eliminates loss spikes that occur with per-micro-batch sampling in 350M-parameter Parcae training.The spikes correspond to overly large recurrent residual jumps at the final recurrence.
  • Per-sequence sampling: Per-sequence depth sampling supports stronger fixed-point behavior by reducing recurrent residual jumps and improving the training objective’s estimate.Recurrent residual and state-norm metrics are used to assess this behavior.
  • Truncated recurrence sampling: The proposed recurrence-sampling method decouples the sampled recurrence distribution from the backpropagation-depth parameter.Prior sampling truncated recurrence based on µbwd, producing a shifted or compressed distribution relative to the desired pretraining distribution.
  • Truncated recurrence sampling: At T = µrec = 8, the ablated sampling methods have nearly identical training and validation trajectories, but the proposed method performs better at T ∈ [4, 16, 64].The reported difference is attributed to the prior method’s constricted effective recurrence range during pretraining.

I Selecting µrec and µbwd

The ablations show that recurrent depth and backpropagation depth are coupled during pretraining. Increasing µrec alone can hurt validation performance, while increasing µbwd generally improves deeper-recurrence performance until gains become minimal.

  • Varying µrec: µrec = 8 performs best when µbwd = 4, including at higher test-time recurrences.Models with larger µrec performed worse at both lower and higher test-time recurrence values in the reported experiment.
  • Varying µbwd: Increasing µbwd improves validation performance at T = 16 and T = 64 when µrec = 20.The improvement is observed across the tested lower and higher recurrence settings, although lower µbwd can perform better at lower validation recurrences.
  • Open question: The authors leave FLOP-optimal choices of µrec and µbwd for future work.The reported experiments establish coupling between the two quantities but do not determine their jointly optimal settings.

J Ablation of Prelude Normalization

The 1.3B Parcae run developed late-stage instability, but diagnostics localized the problem primarily to prelude-induced residual explosion. Normalizing the prelude output stabilized recurrent dynamics and enabled stable training, with clearer quality gains at smaller model sizes.

  • The 1.3B run was stable for 150k optimizer steps before exhibiting state explosion and loss spikes.
  • Spectral norms of A and B remained low, while C’s norm grew but produced a relatively small, slowly decreasing empirical expansion ratio.
  • The recurrent state norm jumped on the first recurrence and then remained relatively stable, indicating that the main explosion occurred before recurrent iteration.
  • Prelude-output normalization stabilized recurrent norms and enabled the stable 1.3B training run.
  • A single layer in the prelude created a residual-norm explosion, while the nonlinear recurrent blocks had little effect on the explosion.
  • Prelude normalization improved convergence for 140M and 370M models, but yielded only negligible improvement for 770M and 1.3B models.

L Fitting Parametric Functions to Test-Time Looping

The paper fits parametric laws for test-time looping and finds that validation loss follows saturating exponential decay with a finite floor. A unified fit combines training scaling with test-time decay and generalizes to held-out models with low error.

  • The authors compare four functional forms for the saturating relationship between test-time recurrence T and validation loss, fitting each curve independently.
  • The exponential decay form achieves the lowest Huber loss, performing 2.3× better in-distribution and 7.1× better under extrapolation than the shifted power law.
  • Removing the irreducible loss floor increases error by over 40×, showing that test-time scaling saturates to a finite loss determined by training.
  • The fitted irreducible loss closely matches empirical loss at T = µrec, motivating the training scaling law as the unified law’s floor.
  • Conditioning the decay rate on training recurrence is necessary: removing µrec conditioning increases training error by 3.5× and held-out error by approximately 33%.
  • The unified fit predicts held-out validation loss within 0.85–1.31% average error, while replacing its floor with an oracle reduces error to 0.10–0.17%.

O Expanded Setup For Training and Test-Time Scaling Laws

The scaling-law experiments vary recurrence under controlled FLOP, parameter, and data budgets, then evaluate test-time recurrence and unified-law generalization. The setup spans multiple model sizes, baselines, benchmarks, and training configurations.

  • The isoFLOP setup fixes training FLOPs and varies µrec while adjusting token budgets, using 36 models at 140M and 18 at 370M.
  • The evaluation includes Core and Core-Extended downstream tasks, with three seeds used for baseline-transformer comparisons.
  • Fixed-depth Parcae models with µrec = 1 provide baselines for comparing looping frontiers against pure data scaling.
  • The test-time setup trains 140M and 370M models on 11.2B tokens with µrec from 2 to 12, then evaluates recurrences T from 1 through 24.
  • The unified-law setup reuses isoFLOP models, evaluates ten test-time recurrence values, and fits all eight parameters jointly using Huber loss and L-BFGS.
  • Parcae and baseline Transformers use separate architectural setups, while the comparison reports parameter definitions and parameter counts explicitly.
  • Parcae training uses specified warmup, cooldown, optimizer, learning-rate, clipping, precision, and FineWeb-edu data configurations.
  • Baseline Transformers receive learning-rate sweeps, whereas Parcae does not; the authors expect a more optimal Parcae learning rate could further improve performance.
Loading 2604.12946v1…