LatentLoop: Scaling Looped MoE Models with Asymmetric Compute Allocation

AI summary

Overview

• Research area: Efficient language model architecture — specifically looped Mixture-of-Experts (MoE) Transformers, compute allocation across loop visits, and compute-optimal scaling laws.
• Technical level: Advanced. The paper assumes familiarity with MoE routing, KV cache accounting, FLOP-matched training budgets, and power-law loss-surface fitting.
• Scope: The paper diagnoses an asymmetry between the two visits of a looped MoE block, then designs and evaluates LatentLoop, an architecture that compresses the first visit's sequence and spends the saved computation on model width, across 15 matched-budget configurations from 1.3B to 7B parameters.

What This Paper Is About

Looped models raise effective depth by running the same shared layers more than once, and MoE layers let the number of stored experts be adjusted separately from the number activated per token. However, every extra visit still costs computation, so gains at a fixed parameter count do not by themselves prove an advantage under equal resource budgets. This paper asks whether the two visits of a looped MoE block deserve equal amounts of computation, and answers by showing that the first visit prefers a much smaller update scale — motivating an architecture that gives the first visit a compressed sequence and reallocates the savings to model width.

Key Contributions

1. Identifying residual-scale asymmetry across visits. Using learnable residual coefficients and fixed-coefficient experiments, the authors show that reducing the first visit's scale can preserve or improve performance, while the same reductions applied to the second visit hurt it.
2. Proposing LatentLoop. An architecture that compresses the first visit's sequence, keeps full resolution in the second, and reallocates the saved computation to model width under matched resource budgets (parameter count, KV cache, FLOPs per step, and total training FLOPs).
3. Introducing calibrated writeback and selective parameter sharing. Learnable scaled-decoder coefficients calibrate how strongly the latent increment is written back to full-resolution states, while MoE routers and experts are shared across visits and attention blocks are kept independent.
4. Demonstrating consistent gains under matched budgets. LatentLoop achieves lower validation NLL than both baselines across all 15 configurations and higher average scores on nine benchmarks at all three model sizes in both 0-shot and 5-shot settings.

Main Findings

• Learned coefficients diverge. When each visit receives its own learnable residual coefficient starting at one, the first visit's coefficient falls well below the second's. Final checkpoint coefficients are approximately (0.27, 0.87) for 1.3B and (0.36, 1.09) for 3B.
• Fixed coefficients confirm the direction. Reducing the first coefficient does not hurt performance and can lower validation NLL, whereas corresponding reductions to the second increase it. The setting (0.25, 1) achieves the lowest NLL at both sizes, improving over (1, 1) by about 0.008 for 1.3B and 0.003 for 3B. This asymmetric setting also beats the symmetric (0.5, 0.5) scaling of Wang et al. (2026a) at both sizes.
• Sequence compression wins among four reduction strategies. At the 3B/40B reference setting, sequence compression reaches an NLL of 2.0303, versus 2.0361 for reducing residual width, 2.0364 for halving expert intermediate width, and 2.0349 for cutting active experts from eight to four. For comparison, MoE Base scores 2.0426 and MoE Loop 2.0293 at that setting.
• The final design is strongest in ablation. LatentLoop with mean aggregation and the scaled decoder reaches 2.0199, versus 2.0268 for mean plus direct writeback, 2.0264 for SpiralFormer-style softmax aggregation and allocation, 2.0219 for a convolution encoder with direct writeback, and 2.0209 for convolution plus the scaled decoder.
• Most of the writeback gain comes from overall scale. Replacing all channel and position coefficients with a single learned scalar gives 2.0210, close to the full scaled decoder's 2.0199. Normalizing the coefficients by their RMS gives 2.0228, above direct writeback's 2.0268 but worse than the full scaled decoder, suggesting that letting the overall scale adapt matters.
• Independent attention helps compressed visits but not the uncompressed loop. Untying attention while reducing stored experts from 92 to 91 worsens ordinary MoE Loop from 2.0293 to 2.0319. The same change from 76 to 75 experts improves the compressed variants, from 2.0303 to 2.0268 with mean aggregation plus direct writeback, and from 2.0237 to 2.0199 with mean aggregation plus the scaled decoder.
• Scaling results favor LatentLoop at every matched endpoint. Across 15 configurations spanning 1.3B–7B parameters, with the largest budget exceeding 10^21 FLOPs, LatentLoop achieves lower validation NLL than MoE Base and MoE Loop. At a fitted NLL of 1.988, LatentLoop reaches that loss with 1.85 × 10^20 FLOPs, corresponding to predicted compute savings of 19.0% relative to MoE Base and 7.9% relative to MoE Loop.
• Downstream scores improve across the board. LatentLoop has the highest average score in all six size and shot settings. Relative to MoE Loop, average gains range from 0.73 to 1.18 percentage points in 5-shot and 0.17 to 0.67 in 0-shot; relative to MoE Base, gains range from 0.67 to 1.82 percentage points in 5-shot and 0.80 to 1.39 in 0-shot. At 7B/140B, the 5-shot average is 63.40 for LatentLoop, 62.55 for MoE Loop, and 62.19 for MoE Base.
• Asymmetry persists with more visits. With three and four visits at both 1.3B and 3B, coefficients increase with visit index: the first visit has the smallest scale and the last the largest. At 90% training progress, ratios between consecutive coefficients are approximately 2.4× for three visits and 2.0× for four visits.
• Writeback coefficients stay well below one. At 3B/40B, the learned scaled-decoder channel means settle at approximately 0.199 and 0.160. After warmup, the nearer receiver (2j) consistently learns a higher mean coefficient than the farther receiver (2j+1).

Methodology in Plain English

The authors start from SMELT's looped MoE recipe, which splits a Transformer into an entry block, a middle block, and an exit block. The entry and exit blocks run once, while the middle block (half of the layers) runs twice using shared parameters — giving 3L/2 layer executions for L stored layers.

To find out whether the two visits matter equally, they attach a learnable scalar coefficient to each visit's residual branches and train it alongside the model. Seeing that the first visit learns a much smaller coefficient, they then run a controlled sweep: fix one visit's coefficient at one and vary the other over {0.125, 0.25, 0.5, 2}, training each setting from scratch.

That asymmetry motivates testing four ways to cut first-visit computation: projecting the residual stream to half width, halving expert intermediate width, activating four experts instead of eight, and compressing the sequence into fewer latent positions. Sequence compression produces the lowest validation NLL, so it becomes the basis of LatentLoop.

In LatentLoop, the first visit operates on a compressed sequence of latent positions (compression factor k = 2, using mean aggregation of adjacent states), while the second visit runs at the original sequence resolution. The saved computation is reinvested in a wider model. A learnable scaled decoder adds just 2d parameters to control how strongly the latent increment is written back to full-resolution states, with separate coefficients per channel and per receiver position and causal alignment so that updates only reach positions where all source tokens are available. MoE routers and experts are shared across the two visits, while attention blocks are made independent to accommodate the two different sequence resolutions; at the matched 3B configuration, dropping one stored expert per MoE layer (from 76 to 75) pays for the extra attention parameters.

Comparisons hold stored backbone parameters, KV cache, FLOPs per step, and total training FLOPs within tight margins. Budgets are matched in a fixed order — attention projection dimensions, then model width, then stored expert count, then an integer number of optimizer steps. Maximum absolute deviations are 0.6% (backbone parameters), 0.8% (KV cache), and 0.9% (FLOPs per token), compared with 1.0%, 3.6%, and 3.9% for SMELT. All methods train on the same FineWeb-Edu stream with the TinyLlama tokenizer and a 0.1B-token validation set, using sparse MoE Transformers in the style of Qwen3 with eight active experts per token. Scaling is analyzed by fitting a separate additive power-law loss surface per architecture, with five parameters fit by L-BFGS-B from 48 initializations using Huber loss with δ = 10^-3.

Why This Matters

Impact on research. The paper reframes looped-model design from "how many visits" to "how much computation per visit." It provides direct evidence that the visits of a shared block do not contribute equally, and shows that a simple residual-coefficient change can improve validation NLL without adding trainable parameters or changing the computational budget. The budget-matching discipline — controlling parameters, KV cache, FLOPs per step, and total training FLOPs simultaneously — sets a higher bar than fixed-parameter comparisons that can silently give looped models extra compute.

Real-world applications:

• Training more capable language models under a fixed compute or memory ceiling, where every percentage point of compute saving translates into lower cost per model.
• Deployment in memory-constrained settings, since the approach targets a smaller KV cache alongside matched FLOPs per step (LatentLoop's KV figures at each tier are 47.81, 87.50, and 145.31 MiB in the reported configurations).
• Long-context or sequence-heavy workloads, where processing part of the network at a compressed resolution is directly relevant.
• Serving architectures that share expert weights across depth, which reduces the parameter footprint that must be stored or sharded.

Industry relevance. The comparison targets are matched budgets, not ideal hardware layouts; the authors acknowledge that the resulting widths, head dimensions, and expert counts "may be less favorable for hardware execution" because they were chosen for fairness of comparison. That is a candid framing of what remains between this research result and a production system, and it points to architectural co-design as the natural next step.

Future Directions

• Explaining the asymmetry. The authors state that the mechanisms underlying the learned asymmetry remain incompletely understood and leave a more detailed investigation to future work.
• Understanding multi-visit coefficient patterns. The consistent ratios between consecutive coefficients — approximately 2.4× for three visits and 2.0× for four visits at 90% training progress — are reported but explicitly left unexplained.
• Extending the scaling analysis. Results cover 15 configurations from 1.3B to 7B parameters with the largest budget exceeding 10^21 FLOPs; whether the fitted savings hold at larger scales is not established in the provided content.
• Hardware co-design. Because budget matching produces widths, head dimensions, and expert counts that may be less favorable for hardware execution, translating the compute savings into realized wall-clock gains is an open engineering question.

Target Audience

Researchers and engineers working on efficient Transformer architectures, Mixture-of-Experts training, looped or recurrent-depth models, and compute-optimal scaling laws. The paper is most useful to readers already comfortable with MoE gating, KV cache accounting, FLOP-matched experimental design, and loss-surface fitting; readers seeking an introduction to efficient model architectures would need background reading first. Practitioners focused on training-cost reduction and parameter-efficient deployment are also part of the intended audience.

Authors’ abstract

Looped Mixture-of-Experts (MoE) models combine parameter reuse with sparse computation to improve performance under matched budgets. Prior work demonstrates that different loop visits can play different roles even when they share parameters. To probe the contributions of different loop visits, we learn separate residual coefficients in a model that applies a shared block twice, observing a smaller scale in the first visit. Fixed coefficient experiments show that reducing the first visit's scale preserves or improves performance, whereas the same reductions to the second hurt it. This asymmetry motivates us to explore asymmetric compute allocation. We propose \textbf{LatentLoop}, which compresses the first visit's sequence, retains full resolution in the second, and reallocates saved computation to increase the hidden size. We further learn how strongly the first visit updates the full sequence, and share MoE weights across visits while using separate attention parameters to accommodate their different sequence resolutions. Across 15 configurations spanning 1.3B--7B parameters, with the largest budget exceeding $10^{21}$ FLOPs, LatentLoop consistently lowers validation NLL. At a fitted NLL of 1.988, predicted compute savings are 19.0\% and 7.9\% relative to vanilla MoE and the looped baseline. It also improves average scores on nine benchmarks at all sizes in 0-shot and 5-shot settings. These results suggest asymmetric compute allocation is a promising direction for looped MoE models.