Source-linked AI summary
Sparse Layers are Critical to Scaling Looped Language Models
Ryan Lee, Jacob Biloki, Edward J. Hu, Jonathan May
TL;DR
Language models are costly to store and slow to run, and looped models have underperformed equal-compute baselines. This paper compares dense and MoE transformers with and without looping, finding that sparse experts address the scaling deficit and looped models offer better early-exit compute-quality trade-offs.
Problem
Language models are expensive to store and slow at inference, while parameter reuse through looping has underperformed baselines at equal training compute.
Source: §1 · p. 1Method
The paper conducts a controlled comparison of dense, MoE, looped, and Looped-MoE transformers, using top-k token-choice routing in sparse layers.
Source: §2.2 · p. 2, §8 · p. 9Results
Sparse experts resolve looped models’ scaling deficit, and looped models have better compute-quality trade-offs than standard transformers with training-free early exits.
Source: §1 · p. 1, §5.3 · p. 5, §8 · p. 9Takeaways & Limitations
Replacing dense FFNs with top-k MoE layers in looped models is presented as a practical route to fewer stored weights at the same accuracy and more compute savings at loop-boundary exits.
Source: §1 · p. 2The study does not measure end-to-end throughput gains with an optimized inference engine, and its four architectures were not scaled beyond 305M/711M active/stored parameters.
Source: §8 · p. 9
Abstract
from arXiv · showhide
Looped language models repeat a set of transformer layers through depth, reducing memory costs and providing natural early-exit points at loop boundaries. However, looped models do not scale as favorably as standard transformers with unique layers. We compare standard and Mixture-of-Experts (MoE) transformers, with and without looping, and find two main results. First, we find Looped-MoE models scale better than the standard baseline while dense looped models do not. We trace this to routing divergence between loops: in Looped-MoE models, different experts are activated on each pass through the same shared layers, recovering expressivity without additional parameters. Our second finding is that looped models have better compute-quality trade-offs with early exits than standard models. Because each loop ends with the same layers that produce the final output, loop boundaries are superior exit points, as confirmed by earlier output convergence at these points. In sum, we provide a clear direction for scaling looped models: a Looped-MoE model with early exits can not only beat standard transformers at scale, but also enable significant memory and inference savings with minimal degradation in quality.
1 Introduction
The introduction frames parameter reuse as a response to language-model storage and inference costs, while noting that prior looped approaches underperform equal-compute baselines. It proposes sparse experts and loop-boundary exits as routes to better scaling and compute-quality trade-offs.
1 Introduction
- Controlled comparisons find that sparse experts resolve looped models’ scaling deficit and achieve the best downstream accuracy.The introduction presents sparse experts as a way to recover expressiveness lost through weight tying.Source: §1 · p. 1
- Training-free early exits give looped models better compute-quality trade-offs than standard transformers, with loop boundaries as especially accurate exit points.More tokens reach near-final predictions after the first loop layers.Source: §1 · pp. 1–2
- Different experts are selected on successive passes through the same shared layers, enabling distinct computations across iterations.The introduction identifies loop-specific expert routing as a property underlying sparse looped models’ scaling and early-exit trade-offs.Source: §1 · p. 2
- The practical recommendation is to replace dense FFNs with top-k MoE layers in looped language models.The resulting architecture is described as storing fewer weights for the same accuracy and enabling more compute savings through loop-boundary exits.Source: §1 · p. 2
2 Model Descriptions
The paper compares transformer variants with dense or sparse FFNs and unique or repeated layers, using a shared decoder-only backbone. It describes the looping and MoE designs and reports a µP transfer check for the scaling study.
2 Model Descriptions
- All variants share a decoder-only transformer backbone, while the study varies whether FFNs are dense or sparse and whether layers are looped.The scaling configurations are summarized in Table 1.Source: §2 · p. 2, Table 1 · p. 2
2.1 Layer Looping
- Looped models repeat L unique transformer layers R times, giving effective depth L × R; the scaling study uses L=8 and R=2.This matches the 16-layer effective depth of non-looped counterparts.Source: §2.1 · p. 2
2.2 Sparse Mixture-of-Experts
- Sparse MoE layers replace the SwiGLU FFN with smaller expert FFNs, and a learned router assigns each token embedding to its top-k experts.The router maps token embeddings x to expert outputs y.Source: §2.2 · p. 2
- In the µP Transfer test, the smallest model’s best learning rate remains optimal across larger sizes; when a width’s optimal rate differs from the base learning rate, the loss difference is <1%.Source: Figure 2 · p. 3
- The µP scaling rules extend to router, expert, and unembedding weights, while excluding the embedding weights from the table’s rules.The extension removes the original input/output multipliers.Source: Table 2 · p. 3
- All MoE configurations use k=2 active experts out of E=8 and apply load-balancing and router z-losses to prevent expert collapse.The auxiliary losses encourage uniform token assignment and penalize large router logits.Source: §2.2 · p. 3
3 Experimental Set-up
The setup applies width-scaled parameterization across dense, sparse, and looped architectures, then validates its transfer across model widths. Compute comparisons use active parameters per forward pass, distinguishing them from stored parameters.
3.1 Maximal Update Parameterization for Looped and Sparse Layers
- µP scales initialization variance and learning rates inversely with width ratio, transferring a small-width tuned learning rate to larger models without retuning each width.For looped layers, no additional modification is needed; expert weights and the router receive width-scaled initialization and learning rates.Source: §3.1 · p. 3
- At widths 128, 256, 512, and 1024, the smallest-width best learning rate remained near-optimal, with at most 0.8% loss difference across architectures.The transfer test used models with 4-layer effective depth.Source: §3.1 · p. 3
- Table 3 distinguishes active from stored non-embedding parameters, accounting for inactive expert weights in MoE.It defines effective layers as counting repeated layers and notes that embedding parameters are included in total Nactive but not shown.Source: Table 3 · p. 4
3.2 FLOPs Accounting for Looped and Sparse Layers
- For FLOPs accounting, Nactive counts parameters used in one forward pass, including repeated looped layers and only active MoE experts; Nunique counts stored parameters.All four architectures have the same Nactive, neglecting small router weights, at matched token counts; embeddings and unembeddings are included in Nactive.Source: §3.2 · p. 4
4 Experiments
The experiments select compute-optimal models across fixed training-compute budgets and evaluate both scaling and downstream performance. They also measure training-free early-exit compute-quality trade-offs using entropy thresholds and architecture-specific exit points.
4.1 IsoFLOP Scaling Study
- For each of four compute budgets, the study trains several widths, selects the quadratic-fit loss minimum, and fits a power law through the resulting compute-optimal points.The budgets are 5 × 10^16, 2 × 10^17, 5 × 10^17, and 10^18 FLOPs; training uses a 10B-token FineWeb-Edu sample with GPT-2 tokenization.Source: §4.1 · p. 4
- The scaling study evaluates compute-optimal Base, Looped, and Looped-MoE models at the largest budget on the AI2 OLMES benchmark suite.This tests whether sparse routing restores expressiveness beyond test loss.Source: §4.1 · p. 4
4.2 Early Exit Evaluation
- Early exit uses output-distribution entropy thresholds; looped models may exit only at loop boundaries, while non-looped models may exit at any intermediate layer.Sweeping the threshold traces compute saved versus perplexity; experiments use compute-optimal models at 1018 training FLOPs, and end-to-end throughput is not measured.Source: §4.2 · p. 4
4.1 IsoFLOP Scaling Study
- Figure 3 marks compute-optimal sizes with stars and compares fitted loss-scaling exponents of 0.076 for Base and 0.077 for Looped-MoE against Kaplan et al. [4] at 0.076.The left and right panels show Base and Looped-MoE IsoFLOP curves, respectively.Source: Figure 3 · p. 5
- At 1018 FLOPs, Table 4 reports that Looped-MoE has the highest average AI2 OLMES score with fewer stored parameters than Base.The benchmark evaluates compute-optimal models.Source: Table 4 · p. 5
5 Results
Sparse experts improve looped-model scaling and downstream accuracy, while looped architectures show more favorable early-exit compute-quality trade-offs. The results connect loop-specific routing with scaling and identify loop boundaries as useful exit points.
5.1 IsoFLOP Scaling Laws
- Looped-MoE scales better than Base, while dense Looped scales strictly worse; fitted loss exponents are 0.077 for Looped-MoE and 0.076 for Base.The Base exponent matches Kaplan et al. [4]; sparse routing is presented as restoring expressiveness lost to weight tying.Source: §5.1 · p. 5
5.2 Downstream Evaluation
- At 1018 FLOPs, Looped-MoE achieves a 39.6 Core 9 average, versus 38.7 for Base and 37.4 for Looped, while storing 216M versus Base’s 246M unique parameters.Table 4 reports AI2 OLMES results for compute-optimal models.Source: §5.2 · p. 5
5.3 Early Exit
- At 10% FLOPs saved, perplexity is 50.2 for Looped and 51.0 for Looped-MoE, versus 55.4 for Base and 75.7 for MoE.Looped models degrade less than their non-looped counterparts under the training-free entropy criterion; the figure identifies loop boundaries as exit points.Source: §5.3 · p. 5, Figure 4 · p. 6
6 Analysis
The analysis examines why MoE substitution changes looped scaling and early-exit behavior, using models at 1018 training FLOPs on 800K test tokens. It attributes these effects to loop-varying expert assignments and earlier convergence at loop boundaries, and tests whether more loops improve exits.
6.1 Why do MoE Layers Improve Looped Scaling Laws?
- Across layers 1–6 and 8, 25–53% of tokens receive non-overlapping expert assignments across passes, while only 4–14% receive identical assignments.Layer 7 is an outlier, with 37% exact matches and only 3% zero overlap; its proposed stabilizing role is speculative.Source: §6.1 · p. 6, Figure 5 · p. 7
6.3 Does More Looping Improve Early-Exit Efficiency?
- All three Looped-MoE configurations achieve a better early-exit compute-quality Pareto frontier than non-looped MoE, and more loops further improve the tradeoff.For Looped-MoE, the 4 × 4 curve lies below 8 × 2, and 2 × 8 lies below both; loop boundaries offer more exit opportunities and are high-quality exit points.Source: Figure 7 · p. 8, §6.3 · p. 8
6.2 Why are Early Exits for Looped Transformers Favorable?
- By the end of the first loop, most tokens have near-final output distributions, with looped models showing a sharp convergence increase at loop boundaries.The proposed mechanism is architectural: each loop ends with layers also trained to produce the final output, unlike equivalent-depth exits in non-looped models.Source: Figure 6 · p. 7, §6.2 · p. 8
7 Related Works
Prior work explores looping, sparse layers, and early exits through varied architectural changes. This paper focuses on standard top-k MoE in looped models and on their early-exit properties by construction.
Looped Transformer Architectures
- MoEUT combines a looping Universal Transformer with σ-MoE [25], SwitchHead [26], and modified LayerNorm [27], whereas this study isolates standard token-choice top-k MoE.The paper contrasts its fixed-compute isoFLOP comparisons with Ouro [2], which compares by unique parameters while using 4× the effective depth.Source: §7 · p. 9
Early Exits and Layer Skipping
- LayerSkip [28] trains standard transformers for early exits, while this paper studies looped models whose shared layer stack already handles final activations during training.Mixture of Depths [29] and Mixture of Recursions [30] use more complex token-skipping or recursive-depth mechanisms.Source: §7 · p. 9
8 Conclusion and Limitations
The study finds that top-k MoE layers resolve looped transformers’ scaling deficit, with routing divergence allowing shared layers to compute differently across iterations. Looping also supports training-free early exits with better compute-quality trade-offs, but the study’s early-exit savings are theoretical and its scaling experiments are limited in size.
8 Conclusion and Limitations
- Replacing dense FFNs with top-k MoE layers resolves the scaling deficit of looped transformers.MoE routing diverges across loop iterations, allowing shared layers to implement distinct computations per loop.Source: §8 · p. 9
- Earlier output convergence at loop boundaries enables training-free early exits with better compute-quality trade-offs.The study concludes that Looped-MoE models are a practical path to models cheaper to store, faster to run, and competitive in quality.Source: §8 · p. 9
- Experiments scale only to 305M active and 711M stored parameters, with larger-scale validation planned at 1B and 7B.The study relies on the principle that compute-optimal scaling laws fitted at smaller scales predict larger-model performance.Source: §8 · p. 9
- Early-exit results show theoretical compute savings, but end-to-end throughput gains with an optimized inference engine were not measured.Source: §8 · p. 9
- The study holds depth constant and scales width because µP transfer does not hold when model depth changes.The authors identify depth-scaling extensions such as CompleteP as a possible direction for future validation.Source: §8 · p. 9
A Appendix
The appendix supplies compute and architecture context alongside scaling, early-exit, and benchmark comparisons. Its figures describe scaling-law fits and architecture offsets, while the tables report early-exit perplexity and full benchmark results.
Appendix
- All experiments used NVIDIA H100 GPUs with 80GB; the reported total compute was approximately 2,000 GPU-hours.The isoFLOP study used approximately 1,200 GPU-hours, µP validation 200, early-exit and loop-depth variants 300, and analysis 100.Source: A Appendix · p. 13
- The compared architectures use a pre-norm transformer layer with SwiGLU FFNs and RMS Norm without learnable gain.Source: Figure 8 · p. 13
- Figure 9 plots Looped and MoE isoFLOP curves, with power-law fits over compute budgets from 5 × 10^16 to 10^18 FLOPs.Dashed lines show fitted scaling relations, and a dotted line shows the Kaplan et al. scaling exponent.Source: Figure 9 · p. 14
- Figure 10 shows roughly equal scaling slopes across architectures, with MoE best, followed by Looped-MoE, Base, and Looped.The caption notes that looping the same architecture appears to reduce expressivity.Source: Figure 10 · p. 14
- At 10% FLOPs saved, Looped-MoE 2 × 8 reaches perplexity 42.0 versus 55.4 for Base and 75.7 for MoE.Across savings levels, Looped has lower perplexity than Base and MoE, and the advantage grows with the number of loops; MoE degrades faster than Base.Source: Table 7 · p. 15
- At 10^18 FLOPs, MoE scores 36.4 on average in the AI2 OLMES benchmarks despite achieving the best test loss.The passage hypothesizes narrower per-token expert access as a reason for the average-score result.Source: Table 6 · p. 15