Here is a comprehensive summary of the paper "SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers," structured as requested.

---

SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers

Summary

• Core Contribution: This paper introduces SMELT (Sparse MoE Transformer, middle layers Loop Twice), a novel architecture that reuses a middle block of layers twice within a single forward pass.
• Key Methodology: SMELT is evaluated under a strict "budget-matching" protocol, which simultaneously holds per-token FLOPs, total non-embedding parameters, and KV cache constant against a standard unlooped Mixture-of-Experts (MoE) Baseline. This isolates the architectural benefit of looping from the extra compute it typically requires.
• Main Findings: SMELT consistently outperforms the Baseline across a scaling ladder (up to 54B parameters), saving 6.8-18.0% of training FLOPs on the compute-optimal frontier. The performance advantage is# is largest on structured data (like Code), long samples, and tasks with many in-context examples.
• Mechanistic Insight: Analysis reveals that the second pass through the looped layers acts as a "refinement step": it preserves retrieval coordinates (Q, K), amplifies residual updates, and crucially, reduces the "attention sink", redirecting attention mass toward more content-relevant tokens.

1. Introduction

• Problem: Looped Transformers increase effective depth by reusing a block of layers. However, prior evaluations often compareimportantly, compare models with different computational budgets. For example, looping a 12-layer model to 24 layers doubles its FLOPs and KV cache, making it unclear if performance gains are due to the architecture or just the extra compute.
• Goal: This paper aims to answer whether looping provides an architectural advantage when all computational budgets are matched. The challenge is that matching FLOPs and parameters is difficult for dense models.
• Solution: The paper leverages Mixture-of-Experts (MoE) to make budget matching feasible. By narrowing the hidden dimension to pay for the extra FLOPs of looping, and then adding more experts to recover the lost parameters, all three budgets (FLOPs, params, KV cache) can be held constant.
• Result: This process leads to the SMELT recipe, which is shown to be a superior architecture under these strict conditions.

2. Related Work

• Looped Transformers: Prior work (e.g., Huginn, Ouro) shows looping improves reasoning but doesn't control for FLOPs. Saunshi et al. found looped models have worse perplexity than iso-FLOP baselines. Schwethelm et al. found looping loses at matched compute. These studies leave open the question of whether looping helps under jointly matched compute and token budgets.
• MoE and Loop Recurrence: Concurrent work (e.g., LoopMoE, Gao et al.) trains MoE Looped Transformers at scale but either doesn't match all budgets or doesn't perform a multi-scale scaling analysis. This paper is the first to close all three budgets simultaneously across a full scaling ladder.
• **## 3. Design Recipe under Matched Compute This section details the experimental setup and the process of finding the optimal loop configuration.

3.1 Training Protocol

• Architecture: Decoder-only Transformer with sparse MoE layers (top-8 experts). Uses grouped-query attention (GQA).
• Residual Scaling: Within the looped span, each sublayer's residual update is scaled by 1/r (where r is the loop count) to prevent correlated updates from inflating the residual stream.
• Data & Schedule: Trained on an internal corpus using a warmup-stable-decay (WSD) schedule## schedule. Six cosine-decay branches from a stable run provide six token horizons for scaling law fitting.
• Compute-Equivalent Sparsity (S): A metric is defined to measure sparsity based on FLOPs. A model that spends half the FLOPs of a fully-dense control is considered to have S=50% sparsity. This allows for a fair comparison between dense and sparse models.

3.2 Matching as a Compute-Allocation Problem

• The core idea is to hold three budgets fixed:
  • Per-token FLOPs: Adjusted by shrinking the hidden dimension (H).
  • Total Parameters: Recovered by increasing the number of experts.
  • KV Cache: Adjusted by changing head size or GQA ratio.
• This process results in a matched Baseline/SMELT pair with typical mismatches of less than 4% for FLOPs/KV and 1% for parameters.

3.3 Ablation 1: Looping the Middle Half Beats Full Looping

• Finding: Looping the middle 50% of layers yields the lowest validation loss, outperforming both smaller spans and full-stack looping.
• Reasoning: This aligns with the idea that first and last layers serve specialized roles (e.g.,# ., token embedding/classification) that benefit from independent parameters.

# 3.4 Ablation 2: Looping Prefers a Larger Effective Depth-to-Width Ratio

• Finding: The Looped Transformer's optimal effective depth-to-width ratio is larger than the Baseline's. The winning configuration uses the same physical depth as the Baseline's optimum but executes more layers per token.
• Reasoning: The shared layers receive gradient contributions from multiple visits, making additional serial computation easier to optimize.

3.5 Ablation 3: Two Loops Beat Three or Four

• Finding: Looping twice is optimal. A third or fourth loop forces a thinner model under matched FLOPs, leading to worse performance.
• Conclusion: The three ablations define the SMELT recipe: loop the middle 50% of layers twice.

4. Scaling Laws and Compute Savings

This section scales the SMELT# recipe and quantifies its benefits.

.

4.1 Scaling Across the Grid

-# - The recipe is applied to apackage a 4x4 grid of scales (100M, 200M#M, 600M, 1.6B) and sparsity levels (S ≈ 0%, 85%, 95%, 97%).

• SMELT consistently achieves lower validation loss than the Baseline

Baseline across all 16 cells# cells.

4.2#2 Scaling Form and Compute Efficiency Gain

• A new scaling law is proposed that incorporates sparsity (S) and per-token FLOPs (F): L(F, S,import, D) = E + A(1-S)^b / F^a + K / D^c
• This law is fit separately for Baseline# Baseline and SMELT. -library- The compute-optimal frontier isconst is found by minimizing loss L(F, S, D) subject to C = F · D.
• Result: At a compute budgetimport budget of 10^21 FL## FLOPs and S ≈ 95%, SMELT saves between 6.8% and 18.#.0% of compute compared to the Baseline to achieve the same validation loss.

5. Downstream Performance

This section analyzes where SMELT's advantage is most pronounced.

5.1 Benchmark Performance

• SMELT outperforms the Baseline on DCLM Core and MMLU benchmarks.
• The advantage is larger than predicted by validation loss alone, suggesting a specific benefit for reasoning and instruction-following tasks.

5.2 Gains by Domain

• CE Gain is positive for all five categories (Code, Finance, Math/STEM, Knowledge, Web).
• The gain is largest for Code (20.4%), which has strict syntax and long-range dependencies.

5.3 Gains with Sample Length and In-Context Examples

• SMELT's gain is concentrated on longer samples (1.52x more benefit on the longest vs. shortest buckets).
• The gap between SMELT and Baseline widens with more in-context examples (0.9 pp at k=0, 1.9 pp at k=1), particularly on demonstration-sensitive tasks like Dyck Languages.

6. Inside the Second Pass

This section provides a mechanistic explanation for SMELT's success.

6.1 Expert Routing Overlap

• The second visit reuses a core subset of experts# experts (2-3 out of 8) but also diversifies the rest, balancing consistency with exploration.

6.2 Residual Stream Updates

• The second visit writes larger updates into the residual stream, both in absolute terms and relative to the residual norm.
• Both the attention and MoE sub-layers contribute to this larger write.

6.3 Attention Patterns

• Q and K remain similar across visits (cosine 0.89-0.93), meaning the retrieval coordinates are preserved.
• V diverges more (cosine 0.65-0# 0.74), meaning the content being retrieved changes.
• Top-8 attended token overlap is package high (56-66% vs 28-34% control), confirming that the attention pattern is largely preserved.

6<br>6.4 Case Study: Dyck Language

• On a controlled Dyck-language task, the second visit drastically reduces attention mass on the BOS token (the "sink") from 0.60 to 0.02 and increases mass on the demonstration answers from 0.24 to 0.85.
• This sink reduction is a general phenomenon, observed across all heads and layers, and runs counter to the Baseline's trend of increasing sink mass with depth.

7. Conclusion and Future Work

• Conclusion: Looping is a net win for MoE Transformers when compute, parameters, and KV cache are all matched. The second visit acts as a refinement step, not just added capacity.
• Future Work: Explore richer loop variants (e.g., adaptive depth, per-visit adapters), study wall-clock time efficiency, and develop a more causal understanding of the mechanisms driving the gain.

8. Contributions

• Core Contributors: Shaowen Wang, Ge Zhang
• Contributors: Kairong Luo, Yuhao Wu, Shaofan Liu, Jiaheng Liu, Wenhao Huang, Shen Yan, Jian Li