PaperlayerView exact revision on arXiv

Source-linked AI summary

Scaling Laws for Looped Mixture of Experts

Yanbei Chen, Anirudh Goyal, Raghuraman Krishnamoorthi

arXiv:2609.40316v1cs.LGcs.AIcs.CL

TL;DR

Existing scaling laws model recurrence and MoE sparsity separately, although these are complementary routes to efficient scaling. This paper introduces a bounded, sparsity-conditional law that models both axes, and reports parameter-efficient downstream performance gains, including at matched training compute on reasoning benchmarks.

  • Problem

    Existing scaling laws model recurrence or MoE sparsity in isolation, leaving their joint scaling behavior uncharacterized.

    Source: §1 · p. 1

  • Method

    Loop Scaling Laws jointly model model size, data, recurrence, and sparsity using a bounded recurrence mapping whose effective-parameter gain depends on sparsity.

    Source: §1 · p. 2, §3.3 · p. 5

  • Results

    Across downstream benchmarks, sparsity lets MoE surpass larger dense models with 1/3 the active parameters, while recurrence lets looped models match non-looped models with 1/2 the total parameters.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The fitted law selects recurrence and expert count under compute and memory budgets, and the paper argues both axes should be co-designed to advance the performance–efficiency frontier.

    Source: §5 · p. 11

    The expert-count parameterization defines E as the expert count under top-1 routing and uses Eroute/k as effective expert expansion for top-k routing.

    Source: §3.3 · p. 6

Abstract

from arXiv · show

Looped transformers and Mixture-of-Experts (MoE) offer complementary routes to efficient scaling: recurrence increases computational depth at fixed parameters, while MoE sparsity expands total capacity at fixed active compute. Yet existing scaling laws model recurrence or sparsity in isolation. In this work, we introduce Loop Scaling Laws, the first scaling law to jointly model recurrence and sparsity alongside model size and data. At its core is a bounded, sparsity-conditional recurrence mapping that characterizes the effective-parameter gain from looping and how sparsity raises this gain. The laws predict the held-out loss of looped models more accurately than prior alternatives, and recover the standard dense and MoE scaling laws as special cases. Beyond prediction, the fitted laws provide a principled foundation for designing looped MoE models under compute and memory constraints. Downstream evaluations further demonstrate the complementary benefits of the two axes: sparsity delivers ~3x active-parameter efficiency, recurrence yields ~2x total-parameter efficiency on reasoning, and joint scaling further advances the performance frontier. As a practical extension, we show these gains hold at trillion-token scale: at matched training compute, a looped MoE with law-derived recurrence matches a ~2x larger non-looped MoE on the reasoning benchmarks, while enabling test-time scaling through recurrence.

1 Introduction

Existing scaling laws model recurrence or MoE sparsity separately, leaving their joint behavior uncharacterized. Loop Scaling Laws combine both axes, and evaluations report parameter-efficient reasoning gains and test-time scaling through recurrence.

1 Introduction

  • Prior scaling laws model recurrence or sparsity separately, while recent looped MoE analyses fix recurrence in their primary scaling analyses, leaving the interaction unmodeled.The paper frames the missing joint law around how recurrence and sparsity shape effective capacity.Source: §1 · p. 1
  • Loop Scaling Laws jointly model model size, data, recurrence, and sparsity with a bounded recurrence gain whose asymptote rises with expert count.The law recovers standard dense, dense-looped, and non-looped MoE scaling laws as special cases and predicts held-out loss more accurately than prior alternatives.Source: §1 · p. 2
  • Across 14 downstream benchmarks, sparsity lets an MoE surpass larger dense models with 1/3 the active parameters, while recurrence lets a looped model match non-looped models with 1/2 the total parameters.Jointly scaling both axes further advances the performance frontier beyond either axis alone.Source: §1 · p. 2
  • At matched training compute, a 0.3B-active/1.3B-total looped MoE matches the reasoning performance of a 0.6B-active/2.9B-total non-looped MoE, with approximately 2× parameter efficiency.The looped model uses additional inference compute, and recurrence can be varied at inference for on-demand test-time scaling.Source: §1 · p. 2

2 Related Work

Prior work explores recurrent depth and MoE sparsity as distinct ways to scale language models efficiently, while existing scaling laws generally model one axis at a time. This paper places recurrence and sparsity in a unified scaling law.

2 Related Work

  • Looped transformers reuse shared layers to increase computational depth without proportional growth in stored parameters, and prior work reports parameter-efficient training- and test-time scaling, especially on reasoning tasks.Architectures include single-layer recurrence, full-model recurrence, and middle-block reuse.Source: §2 · pp. 2–3
  • MoE increases total capacity without proportionally increasing per-token computation by routing each token to a sparse subset of experts; prior studies report complementary advantages when combining MoE and looping.Examples span cloud and edge deployment, and prior looped MoE studies examine benefits in both directions.Source: §2 · p. 3
  • Existing scaling laws incorporate model size, data, recurrence, or MoE sparsity, but prior laws model recurrence or sparsity separately rather than their joint effect.The unified law models model size, data, recurrence, and sparsity, while recovering prior dense and MoE laws as reduced forms.Source: §2 · p. 3

3 Loop Scaling Laws

The paper builds from standard model-and-data scaling laws to a bounded recurrence mapping, then conditions that mapping on MoE sparsity. The resulting law captures saturation and sparsity-dependent gains, predicts held-out loss more accurately than alternatives, and recovers standard scaling laws as special cases.

3.1 Preliminaries

  • The preliminaries distinguish active and total parameters, relate recurrence to unrolled parameter count, and specify that training and per-token inference compute scale with that count.At R=1, looped models recover the corresponding non-looped baseline; recurrence trades compute for quality at fixed parameters, while MoE sparsity trades parameters for quality at fixed compute.Source: §3.1 · pp. 3–4

3.2 Loop Scaling Law

  • Dense loop scaling replaces model size N in the standard scaling law with recurrence-dependent effective parameters Neff(R), distinct from the unrolled parameter count that determines compute.The bounded mapping gives diminishing returns and approaches a finite effective-parameter bound; its coefficients determine the asymptotic gain and how quickly it is approached.Source: §3.2 · pp. 4–5
  • The bounded dense recurrence mapping captures the loss plateau and achieves the lowest held-out RMSE across model size, data, and recurrence among the compared recurrence mappings.Figure 2 compares mappings fitted with R≤8 and evaluated on held-out recurrence up to R=16; linear mappings overestimate gains and power-law mappings remain overly optimistic.Source: Figure 2 · p. 4, §3.2 · p. 5

3.3 MoE Loop Scaling Law

  • For MoE, the active-parameter ratio m conditions recurrence gain: decreasing m raises the asymptotic effective-parameter gain and sustains it over more passes.At fixed m>0 the gain is bounded; in the extreme-sparsity limit m→0, it grows linearly with R, an asymptotic property rather than a guarantee beyond the observed range.Source: §3.3 · pp. 5–6
  • The MoE loop scaling law recovers dense loop scaling at E=1, standard MoE scaling at R=1, and standard dense scaling at R=1 and E=1.The law replaces the active-parameter term with a sparsity-conditional effective-parameter count.Source: §3.3 · p. 6
  • The sparsity-conditional bounded mapping achieves the lowest held-out RMSE across recurrence, expert count, model size, and data axes among the compared MoE recurrence mappings.Fit uses runs with R≤8 and evaluation includes held-out R=16; Figure 3 compares predictive loss and RMSE.Source: Figure 3 · p. 5, §3.3 · p. 6

4 Experiments and Findings

The experiments use a fitted loop-scaling law to characterize effective-capacity gains and predict resource-optimal recurrence and sparsity. Downstream evaluations show complementary gains from recurrence and sparsity, including a trillion-token comparison.

4.1 Scaling Experiments and Analysis

  • The scaling sweep varies active model size, training tokens, expert count, and recurrence, using middle-cycle looping over the middle block.The first two prelude layers and last two coda layers remain unshared.Source: §4.1 · pp. 6–7
  • Higher sparsity improves the IsoFLOP frontier at matched compute, while higher recurrence can lower loss at the same active size when given more compute.For both dense and E=8 models, R=2 at 10^21 FLOPs achieves lower loss than R=1 at 5×10^20 FLOPs.Source: Figure 4 · p. 7, §4.1 · p. 7
  • Recurrence raises the effective-parameter multiplier toward a finite asymptote, while greater sparsity raises that asymptote and delays saturation.The asymptote increases with expert count.Source: §4.1 · p. 7

4.2 Compute- and Memory-Optimal Recurrence and Sparsity

  • Larger training-compute budgets and greater sparsity favor higher compute-optimal recurrence.The recurrence criterion counts an added pass only when its predicted loss reduction meets the tolerance set to the fitted-law RMSE.Source: §4.2 · p. 8, Figure 5 · p. 8
  • Larger weight-memory budgets and higher recurrence favor greater memory-optimal sparsity.The optimization considers weight memory; KV cache and other runtime state are outside its scope.Source: §4.2 · pp. 8–9
  • At fixed compute, tight memory favors recurrence and smaller expert counts, whereas more memory favors sparsity and less recurrence.The fitted law selects the lowest-predicted-loss configuration among candidates meeting compute and memory constraints.Source: §4.2 · p. 9

4.3 Downstream Scaling with Recurrence and Sparsity

  • ≈3× active-parameter efficiency from sparsity on overall performance and ≈2× total-parameter efficiency from recurrence on reasoning summarize their complementary downstream gains.Jointly increasing sparsity and recurrence consistently outperforms the dense baseline at matched training compute.Source: Figure 7 · p. 10, §4.3 · p. 10
  • At matched training compute, the 0.3B-active, 1.3B-total LoopMoE at R ∈ {4, 5} matches the larger MoE on reasoning but trails on Overall performance.The comparison uses an approximately 2× larger non-looped MoE and trades additional inference compute for parameter efficiency.Source: Table 2 · p. 10, §4.3 · p. 10

5 Conclusion

Loop Scaling Laws jointly model recurrence, sparsity, model size, and data with a bounded, sparsity-conditional account of effective capacity. The fitted law guides resource-constrained design, while appendix results show that standard dense and MoE scaling laws emerge as special cases.

5 Conclusion

  • The fitted law selects recurrence and expert count under compute and memory constraints, while the paper identifies improved looping strategies, broader model scales, and inference optimizations as future directions.The proposed inference directions include optimizing recurrent states or adding early-exit gating to scale recurrence only when needed.Source: §5 · p. 11

A.1 Dense Loop Scaling Law

A.2 MoE Loop Scaling Law

  • The MoE loop scaling law recovers the dense loop law at E=1, the non-looped MoE law at R=1, and the standard dense law when E=1 and R=1.The reduced non-looped MoE form is equivalent to Ludziejewski et al. (2025).Source: §A.2 · pp. 14–15

B Expert-Path Diversity Across Recurrence

The paper defines expert-path diversity to measure how many distinct experts tokens reach across recurrent passes. Its routing measurements show that diversity increases with expert count at fixed recurrence.

Definition and measurement

  • Expert-path diversity deduplicates experts selected across recurrent passes, averages over tokens and looped layers, and normalizes by the number of active experts.It equals 1 for a single pass or for a single routed expert.Source: B Expert-Path Diversity Across… · p. 15

Empirical measures of Ψ

  • At fixed recurrence, expert-path diversity Ψ increases consistently with expert count, indicating that greater sparsity lets each pass reach a broader set of expert parameters.The statistic was computed on 800 validation prompts spanning tasks such as math, code, and knowledge.Source: B Expert-Path Diversity Across… · p. 15, Table B.1 · p. 15
  • The authors use Ψ as an empirical routing statistic and therefore model the modulated effective-parameter gain by conditioning on sparsity.The table reports Ψ across varying recurrence R and expert count E as mean ± standard deviation.Source: B Expert-Path Diversity Across… · p. 15

C.1 Training Details

The scaling experiments use models trained from scratch on a shared data mixture, with fixed training and architecture conventions. A model ladder and sweep vary active size, data, expert count, and recurrence, and the design-space evaluation extends some settings beyond the fitted range.

C.1 Training Details

  • Scaling-sweep models use the same open-licensed, web-heavy training mixture supplemented with math, code, knowledge, and science data, with a 2,048-token context and 3,072-sequence global batch.The sweep uses 8 nodes with 64 NVIDIA H100 96 GB GPUs and approximately 6.29 million tokens per batch.Source: §C.1 · p. 16
  • For looped MoE, the scaling sweep applies next-token prediction loss to the final recurrent pass, while the Table 2 model adds per-loop supervision for evaluation at different test-time recurrence.The setup also specifies sigmoid gating, per-token top-k normalization, auxiliary-loss-free load balancing, router z-loss, and drop-and-pad dispatch.Source: §C.1 · p. 16

C.2 Model Ladder and Scaling Sweep

  • The model ladder has five active scales from 0.3B to 2.4B, with each rung fixing its configuration and middle-block looping that leaves the first and last two layers unshared.Nact and Ntotal include embeddings; Nloop counts only recurrent-block parameters.Source: §C.2 · p. 16, Table C.1 · p. 16
  • The scaling sweep varies active size, training tokens, expert count, and recurrence; the optimization evaluation covers 300 configurations, with some larger scales, expert counts, and recurrences outside the fitted range.The sweep uses Nact of 0.3B, 0.6B, and 1.0B; D from 100B to 500B; E from 1 to 16; and R from 1 to 8.Source: §C.2 · p. 16

C.3 Parametric Fitting

The fitting procedure compares alternative recurrence mappings within dense and MoE scaling laws, using staged fits and held-out slices. The selected sparsity-conditional MoE formulation shares a sparsity exponent because a less constrained version did not improve fit error.

Fitting the dense loop scaling laws

  • Dense loop fitting compares linear, power-law, and bounded recurrence mappings, with the latter two modeling effective-parameter gain from looping.The linear mapping coincides with the unrolled parameter count; the power-law and bounded forms add fitted recurrence coefficients.Source: §C.3 · p. 17
  • Dense fitting first estimates base coefficients on the R=1 subset, then fits recurrence jointly; held-out evaluation uses R=16, N=1.0B, and D in (400, 500]B.The formulations are compared on the same held-out runs, with RMSE reported for held-out slices.Source: §C.3 · p. 17

Fitting the MoE loop scaling laws

  • MoE loop fitting compares linear, power-law, bounded, and sparsity-conditional mappings, with only the latter making the recurrence mapping depend on sparsity.For the sparsity-conditional form, κj(m) = κj m^-θ for j in {1, 2}; other mappings are independent of m.Source: §C.3 · p. 17
  • MoE fitting first establishes the non-looped law on R=1 runs, then jointly estimates coefficients across recurrence, model size, data, and expert count; held-out slices are R=16, E=16, Nact=1.0B, and D in (400, 500]B.The final law is fit using more than 1,000 loss observations.Source: §C.3 · p. 17

Sparsity exponent ablation

  • RMSE remains 0.0037 when the sparsity-conditional mapping uses separate exponents, so the authors adopt a shared exponent.The extra exponent yields no meaningful gain.Source: §C.3 · p. 18

C.4 Bootstrap Results of the MoE Loop Scaling Law

Bootstrap refits indicate stable coefficients for the sparsity-conditional recurrence mapping and stable architecture choices under compute and memory budgets.

C.4 Bootstrap Results of the MoE Loop Scaling Law

  • The protocol refit the law on 100 subsets, each sampling 90% of observations without replacement, and reported full-data estimates with 90% percentile confidence intervals.The intervals use the 5th and 95th percentiles across refits.Source: §C.4 · p. 18
  • Across 100 refits, coefficient intervals were narrow and budget-selected configurations remained unchanged, supporting robust architecture optimization.The joint optimum was evaluated under compute and memory budgets with INT4 weights.Source: Table C.3 · p. 18, §C.4 · p. 18
  • The fitted coefficients were κ1 = 0.3682 [0.3658, 0.3723], κ2 = 1.4037 [1.3896, 1.4243], and θ = 0.3286 [0.3183, 0.3350].Source: §C.4 · p. 18

D.1 Downstream Evaluation Protocols

The downstream evaluation covers 14 benchmarks across five competencies and specifies the evaluation harness, inference settings, and benchmark-level setups.

D.1 Downstream Evaluation Protocols

  • Table D.1 lists the 14 benchmark setups, with task names and metrics following lm-eval.Source: Table D.1 · p. 18, §D.1 · p. 19
  • The evaluation spans 14 benchmarks across reasoning, science, commonsense, reading, and knowledge.The benchmarks include BBH and GSM8K for reasoning, and MMLU, Natural Questions, and TriviaQA for knowledge.Source: §D.1 · p. 19
  • Models are evaluated with lm-eval and the vLLM backend, using bfloat16 precision and greedy decoding for generation tasks.Source: §D.1 · p. 19

D.2 Compute-Matched Trillion-Token Experiment

The experiment selects recurrence by comparing candidate loop counts at matched training compute, then evaluates the selected model across recurrence levels at test time.

D.2 Compute-Matched Trillion-Token Experiment

  • The recurrence-selection procedure tests candidate loop counts on an IsoFLOP basis, predicts loss, and chooses the compute-optimal recurrence subject to compute and loss-change constraints.The target token budget is derived from the compute budget and unrolled parameter count.Source: §D.2 · p. 19, R L(Nact, DR · p. 19

R L(Nact, DR

  • The loss-improvement threshold is set to the fitted-law RMSE, so gains below predictive discrepancy do not justify another recurrent pass.Source: R L(Nact, DR · p. 19
  • The compute-matched comparison uses an A0.3B-1.3B LoopMoE with R=5 and an A0.6B-2.9B non-looped MoE, trained for approximately 3T and 6T tokens, respectively.The training runs correspond to approximately 1.5×10^22 FLOPs.Source: R L(Nact, DR · p. 19
  • At test time, the LoopMoE varies recurrence from 1 to 5, allowing inference effort to be adjusted on demand because it produces valid output at each pass.Source: R L(Nact, DR · p. 19
  • Inference compute is compared using a relative ratio based on the unrolled parameter count at each recurrence.The stated inference-compute relation is Finf(R) = 2Nunroll(R).Source: R L(Nact, DR · p. 19

Related papers

Loading 2609.40316v1…