SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers

🤗 Upvotes: 53 | cs.LG

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

Title:
SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers

Arxiv:
http://arxiv.org/abs/2609.01343v1

Abstract:
Looped Transformers increase effective depth by iterating a shared block of layers, but most evaluations compare at fixed model size, conflating architectural advantage with extra FLOPs. We study looping on Mixture-of-Experts Transformers while closely matching per-token FLOPs, total non-embedding parameters, and KV cache. Through a series of ablations, we arrive at a recipe we call SMELT (Sparse MoE Transformer, middle layers Loop Twice), which loops the middle half of layers twice while matching the unlooped Baseline on all three budgets. We scale SMELT across four sizes up to 54B non-embedding parameters and fit a separate Chinchilla-style scaling law for each architecture. SMELT's loss drops faster with compute, saving 6.8--18.0\% of training FLOPs on the compute-optimal frontier. The advantage transfers to downstream benchmarks beyond what validation loss predicts, is largest on Code, and grows with sample length and the number of in-context examples. Mechanistic analysis shows that the second visit reduces the attention sink and redirects mass toward content-relevant tokens, an inductive bias that may underlie the observed performance gains. These results show that looping can improve Transformers even under budget matching, offering a practical recipe that turns depth reuse into measurable gains.

What is Daily Paper Cast?

We update every weekday to discuss highest-voted papers from Huggingface Daily Paper (https://huggingface.co/papers). Both the podcast scripts and audio are generated by AI. Feedback and suggestions are welcome! Email us: [e-mail]

Creator:
Jingwen Liang, 3D ML, https://www.linkedin.com/in/jingwen-liang/
Gengyu Wang, LLM ML, http://wanggengyu.com

Listen on:
Spotify: https://open.spotify.com/show/21nrhmdaA8qoBiH8q03NXL
Apple Podcast: https://podcasts.apple.com/us/podcast/daily-paper-cast/id1777620236

Cover Image by Kawen Kuang https://kawen.art

Evan: Welcome to Daily Paper Cast.

Ashley: Today we're discussing a paper featured on the Hugging Face daily paper list of September 2, 2026, with 53 upvotes.

Evan: The title of the paper is 'SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers'.

Ashley: The first two authors are Shaowen Wang and Ge Zhang.

The corresponding author is Shaowen Wang from Tsinghua University.

Evan: Alright Ashley, let's dive right into the Introduction.

What problem does this paper address?

Ashley: Looped Transformers have been shown to increase a model’s effective depth by repeating a shared block of layers instead of stacking new ones.

This tactic has generated significant interest because it allows models to match or exceed the performance of much larger unlooped models on tasks such as arithmetic, multi-hop induction, and mathematical reasoning.

However, past evaluations often compare models with a fixed parameter count, which can conflate the architectural benefits of looping with the extra floating-point operations per second, or FLOPs, involved in repeating layers.

Evan: So, it sounds like there might be a significant computational cost associated with looping that hasn't been fully isolated in previous studies.

What exactly are the authors investigating?

Ashley: Exactly.

The authors aim to disentangle the architectural advantages of looping from the computational overhead.

They do this by studying Mixture-of-Experts (MoE) Transformers and carefully matching the per-token FLOPs, total non-embedding parameters, and key-value (KV) cache between the looped and unlooped models.

Evan: That sounds comprehensive.

How do they go about this comparison?

Ashley: Through a series of ablations, they develop a method called SMELT: Sparse MoE Transformer, middle layers Loop Twice.

In SMELT, the middle half of the layers loop twice.

This design is compared to an unlooped baseline while keeping the computational budgets closely matched.

They scale SMELT across four different sizes, up to 54 billion non-embedding parameters, and establish separate Chinchilla-style scaling laws for each architecture.

Evan: So, they keep the computational budgets constant and analyze the performance.

What are the main findings?

Ashley: They find that SMELT's loss decreases more rapidly with increased compute, allowing it to save between 6.8% and 18.0% of training FLOPs on the compute-optimal frontier.

Moreover, these gains extend beyond validation loss predictions, particularly in downstream tasks with longer sample lengths and more in-context examples.

A key mechanistic insight is that the second visit in the loop reduces the concentration of attention on initial tokens and redirects it towards content-relevant tokens.

This inductive bias might underlie the observed performance improvements.

Evan: That’s fascinating.

So, this paper essentially establishes a practical recipe for improving Transformers by reusing layers without increasing the overall computational budget.

Are there any specific areas where the performance gains are especially noticeable?

Ashley: Yes, the advantage is most prominent in tasks involving code and tends to grow with longer sample lengths and a higher number of in-context examples.

This indicates that looping can become more beneficial as the input data becomes richer in structure.

Evan: Great, that gives us a solid understanding of the background and objectives of their study.

It sounds like SMELT could be a valuable technique for improving deep learning models under strict computational constraints.

Ashley: Indeed.

And with those insights, we've reached the end of the Introduction section.

Evan: Ashley, now that we have a clear understanding of the background and objectives, let's dive into the methodology.

How do the authors approach the problem and what methods do they propose?

Ashley: The methodology is quite detailed.

The authors define the Baseline as a standard unlooped Mixture-of-Experts, or MoE, Transformer.

For their experiments, they use a looped Transformer variant that repeats a contiguous block of layers multiple times.

They perform several ablations to explore different aspects such as the loop span, loop count, and depth-to-width ratio.

Evan: Interesting, so they are essentially experimenting to find the optimal configuration for looping.

Could you break down their experimental setup and the key variables they adjust in their models?

Ashley: Of course.

They start by defining key terms and hyperparameters.

In the Looped Transformer, each sublayer’s residual update within the looped span is scaled by the reciprocal of the loop count to prevent the residual stream from inflating.

Each model is a decoder-only Transformer with sparse MoE feed-forward layers where each token is routed to its top-8 experts using a grouped-query attention mechanism.

Evan: Alright, so they maintain the residual stream's stability by scaling it appropriately.

What are the four scales they use for the experiments?

Ashley: They use four scales: 100 million, 200 million, 600 million, and 1.6 billion non-embedding parameters.

Each of these is matched against an internal baseline in terms of active non-embedding parameters to ensure a fair comparison.

Evan: And how do they ensure that the comparison remains fair in terms of computational budget?

Ashley: Excellent point.

They maintain three budgets: per-token FLOPs, total parameters, and the key-value cache (KV cache).

The matching process involves adjusting the hidden dimension or number of layers to keep the per-token FLOPs constant, and modifying the expert count to maintain approximate parameter parity.

Evan: That sounds complex but necessary for a rigorous comparison.

How do they measure and validate the performance of these configurations?

Ashley: They use the token-weighted cross-entropy loss on a held-out validation set that comprises 39 individual sources grouped into five categories: Code, Math and STEM, Knowledge, Finance, and Web.

They also measure DCLM Core suite and MMLU benchmarks with 10 random few-shot seeds per task and BOS-preserving truncation for consistency.

In addition, they evaluate the token-weighted cross-entropy on gold answer tokens for the DCLM Completion metric, which captures performance in free-form answer tasks.

Evan: Right.

Having these structured and unstructured benchmarks provides a comprehensive view of the model performance.

What did the ablations reveal about looping?

Ashley: The ablations provided three main insights: first, looping the middle 50% of the layers, rather than the entire stack, resulted in better performance.

Second, the optimal depth-to-width ratio for the Looped Transformer was larger than for the Baseline, indicating a preference for more layers relative to width.

Third, looping twice was found to be preferable to looping three or four times, as additional loops led to diminishing returns and a thinner model.

Evan: That makes sense.

Did they find any specific configurations that stood out in terms of their matching criteria and performance?

Ashley: Yes, they arrived at what they call the SMELT recipe: loop the middle half of layers twice, narrow the hidden dimension, raise the expert count to recover total parameters, and scale the looped residuals by one-half.

This configuration was tested across the four model sizes and successfully outperformed the Baseline in nearly all conditions.

Evan: Now, I'm curious about the scaling laws they fitted for the different architectures.

Can you explain how they used the data and what the results showed?

Ashley: Certainly.

They fitted separate Chinchilla-style scaling laws for the Baseline and SMELT across the different scales and sparsity levels.

These fits were used to predict how the loss evolves with increasing compute, allowing them to project the compute-efficiency gains at higher budgets.

Specifically, they showed that SMELT's loss decreased faster with both compute and data, giving it a significant edge in compute efficiency— a 6.8% to 18.0% saving in training compute at different budgets and sparsity levels.

Evan: So, SMELT not only performed better under matched budgets but also did so more efficiently.

How did they verify these results?

Ashley: They used a rigorous fitting procedure that employed Huber-loss minimization on log-loss with L-BFGS-B optimization.

Separate scaling surfaces were fitted for the Baseline and SMELT, ensuring that the comparison isolated the architectural effect from the compute allocation.

They ran 32 training runs, each with six decay branches, yielding 192 evaluation endpoints to inform their scaling analysis.

Evan: That’s quite thorough.

Finally, aside from summary metrics like loss, did they find any task-specific or domain-specific insights?

Ashley: Yes, they did.

The advantage of SMELT was found to be domain-specific as well.

It performed exceptionally well in structured tasks like code and STEM.

Moreover, the gain was further amplified on longer samples and with a higher number of in-context examples, suggesting that the architectural benefits of looping are more pronounced in more complex and structured inputs.

Evan: Alright, so to summarize, the authors present SMELT as a method to enhance Transformer efficiency by reusing layers without increasing the overall computational cost.

They rigorously matched computational budgets for fair comparisons, developed sophisticated scaling laws, and found the method particularly effective in structured tasks and scenarios involving long or complex input data.

Ashley: Exactly.

And that wraps up the Method section.

Evan: Now that we've covered the methodology, let's delve into the experiments and results.

How did the authors validate their method and what did they find, Ashley?

Ashley: The experiments were designed to validate the performance of the SMELT method across different model scales and sparsity levels.

They implemented the SMELT recipe and compared it to the Baseline across 16 different configurations – four scales each with four different sparsity levels.

Evan: Can you remind us what these scales and sparsity levels are?

Ashley: Sure.

The scales used were 100 million, 200 million, 600 million, and 1.6 billion non-embedding parameters.

For each scale, they tested the model with a dense-reference control at zero percent sparsity and three levels of compute-equivalent sparsity approximately 85%, 95%, and 97%.

Evan: Alright, and they used these configurations to monitor training progress and evaluate performance.

How did they track the metrics during the training?

Ashley: They tracked the training loss, especially highlighting the stable (constant learning rate) phase and the cosine-decay branches, which were forked at six different checkpoints.

They plotted these metrics against cumulative training tokens to monitor progress and determine efficiency.

Evan: Got it.

And what were the main findings regarding the training loss and compute efficiency?

Ashley: SMELT consistently showed lower training loss compared to the Baseline across all configurations.

For instance, at the 1.6 billion parameter scale with 97% sparsity, SMELT's training loss curve stayed below that of the Baseline throughout the training.

The compute-efficient sparsity fits showed that SMELT's loss dropped faster with both compute and data, allowing for a 6.8% to 18.0% saving in training compute on the compute-optimal frontier depending on the scale and sparsity level.

Evan: So, SMELT was more efficient across the board.

How does this translate to downstream performance on specific benchmarks?

Ashley: In terms of downstream performance, SMELT outperformed the Baseline in 96 of 96 matched pairs on the DCLM Completion metric, 83 of 96 on DCLM Core, and 29 of 30 pairs on MMLU where the Baseline scored at least 10 percentage points above chance.

This indicates that SMELT's improvements in training loss translated well into superior downstream task performance.

Evan: Those are impressive results.

Did the authors find that SMELT's performance gains were uniform across different domains and tasks, or were there specific areas where it excelled?

Ashley: The performance gains were most pronounced in structured data, particularly in the Code domain, followed by Finance and Math/STEM.

Additionally, SMELT showed greater benefits for longer samples and tasks with in-context learning, suggesting that its architecture is particularly advantageous for complex, structured inputs with long-range dependencies.

Evan: It's clear that the method shines with structured and complex tasks.

Do the authors provide any mechanistic insights as to why SMELT performs better, especially on these intricate tasks?

Ashley: Yes, they conducted a mechanistic analysis which revealed that the second visit in SMELT helps reduce the attention sink—where attention mass gets concentrated towards initial tokens—by redistributing attention towards content-relevant tokens.

This redirection likely enhances the model's capacity for structured reasoning and in-context learning, where it's crucial to focus attention on specific parts of the input data.

Evan: It's fascinating how a seemingly simple architectural tweak can lead to such substantial gains.

How did SMELT compare to the Baseline in terms of compute efficiency specifically?

Ashley: SMELT exhibited superior compute efficiency across all tested scales and sparsity levels.

The paper reports that SMELT's loss curves were generally steeper, indicating faster improvement per unit of compute used.

The authors quantified these savings and reported that SMELT saves up to 18.0% in training compute for the same performance level.

Evan: Impressive.

Before we wrap up this segment, were there any long-term trends or future projections discussed based on their scaling laws?

Ashley: Indeed.

SMELT's computed scaling laws suggest that its advantages become more pronounced as compute budgets grow.

This indicates potential for even greater efficiency gains at larger scales and longer training durations, paving the way for more resource-efficient large-scale models in the future.

Evan: To sum up, the experiments robustly validated SMELT's approach, demonstrating consistent training efficiency gains, superior downstream performance, and particular effectiveness on structured and complex tasks.

Ashley: Exactly.

And that brings us to the end of the Experiment section.

Evan: Transitioning now to the related work that sets the stage for SMELT.

Ashley, can you explain how this research builds on or contrasts with previous studies?

Ashley: Evan.

The paper discusses a range of prior work in multiple related areas, with a particular focus on two main threads: Looped Transformers and Mixture-of-Experts, or MoE, architectures.

Evan: Let's start with Looped Transformers.

What have previous studies found in this area?

Ashley: Looped Transformers were first introduced with the Universal Transformer, which applied weight sharing across depth by iteratively using a shared block of layers.

This concept was scaled up by models like Huginn, which employs a prelude-recur-coda layout, and Ouro, which loops the entire stack of layers.

These studies demonstrated that looped models could match or exceed the performance of significantly larger unlooped models, particularly on reasoning benchmarks.

Evan: So, what are the limitations of these earlier approaches?

Ashley: One major limitation is that many of these studies keep the parameter count fixed while increasing recurrences, which inflates the FLOPs without isolating the architectural advantage of looping.

Other studies compared looped models to parameter-matched untied models, but these too often conflate architectural efficiency with added computational cost.

Evan: I see.

SMELT addresses this by holding computational budgets constant.

What about the second main thread of related work—MoE architectures?

Ashley: MoE architectures decouple the total parameter count from per-token FLOPs by activating only a subset of the model's parameters for each token.

Previous work in this area includes studies like MoEUT, which matched standard Transformers in language modeling, and LoopMoE, which applied MoE principles to looped architectures.

These studies showed how sparse routing could be leveraged to scale models efficiently.

Evan: Does SMELT bring any new insights to MoE designs?

Ashley: SMELT builds on this by combining MoE and looping while ensuring budget matching across three axes: per-token FLOPs, total parameters, and KV cache.

By doing so, the authors isolate and elevate the architectural advantages of looping, showing that this approach provides substantial gains even under strict budget constraints.

Evan: Very interesting.

What insights do prior studies on scaling laws and attention mechanisms contribute to this work?

Ashley: The authors cite Kaplan et al.'s work on power-law relations between loss, model size, and training tokens, and Hoffmann et al.'s compute-optimal allocation theory, which provides a theoretical framework for the scaling laws they fit for SMELT.

They also reference works on attention sinks, which have shown that autoregressive Transformers tend to concentrate attention mass on initial tokens.

SMELT’s design mitigates this by redistributing attention to more content-relevant tokens in its second visit.

Evan: How do these mechanisms enhance the effectiveness of SMELT?

Ashley: By addressing attention sinks and optimizing token routing, SMELT effectively utilizes its architecture to ensure that the second pass through the looped layers refines and reinforces the model’s grasp on content-specific tokens.

This dual-pass mechanism appears to underpin its improved performance, particularly in structured and complex tasks.

Evan: It sounds like SMELT's approach is a thoughtfully orchestrated synthesis of several advanced techniques.

Before we wrap up this section, are there any other notable works that the authors build on or contrast with?

Ashley: Definitely.

They discuss works like RELAX, which introduces low-rank adapters to ease strict weight tying, and Mixture-of-Recursions, which dynamically adjusts recursion depth on a token level.

These variations point to the versatility in optimizing deep learning architectures and reinforce SMELT’s position as a robust, compute-efficient design.

Evan: Thank you, Ashley.

It’s clear that SMELT stands on the shoulders of several significant advances in the field, pushing the envelope even further.

And that brings us to the end of the Related Work section.

Evan: Alright Ashley, let's wrap this up.

Can you summarize the key contributions and takeaways of the SMELT paper?

Ashley: Evan.

The paper 'SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers' makes several significant contributions.

Firstly, it introduces the SMELT method, which loops the middle layers twice in Mixture-of-Experts Transformers while keeping the computational budget consistent.

This approach isolates the architectural advantages of looping from the computational costs.

Evan: Right, and they ensure fair comparisons by matching per-token FLOPs, total parameters, and key-value cache across different configurations.

Ashley: Exactly.

The experiments showed that SMELT consistently reduces training loss faster than a Baseline across multiple scales and sparsity levels.

It also demonstrated significant improvements in downstream benchmarks, particularly in structured tasks and those requiring long context windows or in-context learning.

Evan: We also saw that SMELT saves between 6.8% and 18.0% in training compute at various compute-optimal points, making it a more efficient architecture.

The authors provided mechanistic insights into how the second pass helps reduce attention sink and improve focus on relevant tokens.

Ashley: Indeed.

These findings suggest that SMELT offers a practical and efficient way to enhance the performance of Transformers without increasing computational requirements.

It builds on existing work in looped and Mixture-of-Experts architectures while providing a new direction for efficient deep learning.

Evan: That's a great summary, Ashley.

Thanks to our listeners for joining us on Daily Paper Cast.

We hope you found today's discussion insightful.

Ashley: Join us again next time as we continue to explore the latest in AI research.

Be sure to tune in.

Bye for now!

Evan: Bye everyone!