PaperlayerView exact revision on arXiv

Source-linked AI summary

LatentLoop: Scaling Looped MoE Models with Asymmetric Compute Allocation

Zefeng He, Yunkai Chai, Xiaoyang Chen, Tong Zhu, Zongliang Wu, Yun Luo, Yafu Li, Yu Cheng

arXiv:2610.19975v1cs.LG

TL;DR

Looped MoE models reuse parameters, but equal-budget gains are not established by fixed-parameter comparisons alone. LatentLoop compresses the first loop visit and reallocates saved computation to model width while calibrating writeback and retaining shared MoE weights. It lowers validation NLL across matched configurations and improves average scores across downstream benchmarks and settings.

  • Problem

    Looping can improve language modeling without storing separate weights for each visit, but each visit still costs computation, so fixed-parameter gains do not establish an advantage under equal resource budgets.

    Source: §1 · p. 1

  • Method

    LatentLoop compresses the first visit’s sequence, reallocates saved computation to model width, calibrates latent writeback with learned coefficients, and shares MoE routers and experts while keeping attention blocks independent.

    Source: §1 · p. 2, §3.2 · p. 6, §3.3 · p. 6

  • Results

    LatentLoop achieves lower validation NLL than MoE Base and MoE Loop across 15 matched configurations and higher average scores across nine downstream benchmarks in all size-and-shot settings.

    Source: §1 · p. 2, §6 · p. 13

  • Takeaways & Limitations

    The results suggest asymmetric compute allocation is a promising direction for improving looped MoE models.

    Source: §1 · p. 2

    The mechanisms behind learned residual-coefficient patterns, particularly with three and four loop visits, remain incompletely understood.

    Source: Limitations · p. 13

Abstract

from arXiv · show

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.

1. Introduction

The introduction identifies an asymmetry in the two visits’ preferred residual scales and uses it to motivate LatentLoop, which compresses the first visit while retaining full resolution in the second. Across matched configurations, LatentLoop lowers validation NLL and improves average downstream scores over the baselines.

1. Introduction

  • Because each loop visit still requires computation, gains at a fixed parameter count do not establish an advantage under equal resource budgets.The study matches parameter count, KV cache, FLOPs per step, and total training FLOPs within 1%.Source: §1 · p. 1
  • The first visit learns a substantially smaller residual coefficient than the second, and reducing its scale preserves or improves validation NLL while the same reduction to the second visit increases it.This asymmetry motivates exploring different compute allocations across visits.Source: §1 · pp. 1–2
  • LatentLoop compresses the first visit’s sequence and reallocates saved computation to model width, while learnable writeback coefficients calibrate latent updates and attention blocks remain independent.MoE routers and experts are shared across visits; stored expert count is slightly reduced to offset the added attention parameters while the active count remains unchanged.Source: §1 · p. 2
  • LatentLoop achieves lower validation NLL than MoE Base and MoE Loop across all 15 configurations, and the highest average score on nine benchmarks in all six size-and-shot settings.At a fitted NLL of approximately 1.988, scaling curves predict 19.0% less training computation than MoE Base and 7.9% less than MoE Loop.Source: §1 · p. 2

2. Preliminaries

The paper follows SMELT’s looped MoE architecture and budget-matching framework. Its middle block is visited twice with shared parameters, and comparisons match cache, per-step computation, stored parameters, and total training computation.

2. Preliminaries

  • Comparisons match KV cache, computation per step, stored parameter budget, and total training FLOPs.The matching framework follows SMELT.Source: §2 · p. 2

2.1. Looped MoE Architecture

  • The entry and exit blocks run once, while the middle block runs twice with shared parameters, yielding 3L/2 layer executions.The middle visits share attention blocks, routers, and experts; both residual coefficients are set to one as a reference.Source: §2.1 · p. 3

2.2. Budget Matching

  • Budget matching adjusts attention projection dimensions, model width, stored expert count, and optimizer steps, in that order.Query heads, KV groups, and active experts remain fixed within each comparison.Source: §2.2 · p. 3, Figure 1 · p. 3

3. LatentLoop

The two visits favor different residual update scales, motivating LatentLoop’s compressed first visit and full-resolution second visit. The section also develops calibrated writeback and selective parameter sharing for these different sequence resolutions.

3.1. Asymmetry in Loop Updates

  • Learned scales favor a smaller first visit (0.27, 0.87 at 1.3B; 0.36, 1.09 at 3B), and reducing only its coefficient lowers NLL while matching reductions to the second raise it.The (0.25, 1) setting improves over (1, 1) by about 0.008 and 0.003 NLL at 1.3B and 3B, respectively.Source: §3.1 · p. 4, Figure 2 · p. 4

3.2. Asymmetric Compute Allocation

  • Among four first-visit reductions, sequence compression gives the lowest validation NLL; LatentLoop keeps the second visit full-resolution and reallocates saved computation to model width.The first visit processes fewer latent positions, while the second processes the original sequence resolution.Source: §3.2 · p. 5, Figure 3 · p. 5

3.3. Calibrated Writeback

  • LatentLoop retains a parameter-free mean encoder and learns decoder writeback scales per channel and receiver position; learning either alone improves NLL, while learning both adds no improvement over decoder-only learning.With compression factor 2, the scaled decoder adds 2d parameters.Source: §3.3 · p. 6

3.4. Selective Parameter Sharing

  • Independent attention blocks with shared MoE routers and experts improve validation NLL over full parameter sharing; in the matched 3B configuration, reducing stored experts from 76 to 75 accommodates the extra attention parameters.Attention FLOPs and KV storage remain unchanged.Source: §3.4 · p. 7

4. Experiments

Across matched-budget experiments, LatentLoop lowers validation NLL and improves downstream averages; ablations examine sequence compression, writeback calibration, and selective parameter sharing.

4.1. Experimental Setup

  • The 15 matched-budget configurations span 1.3B–7B parameters, and the largest training budget exceeds 10^21 FLOPs.Baselines are standard MoE executed once and SMELT-style MoE Loop with two shared middle-block visits.Source: §4.2 · p. 7, §4.1 · p. 7

4.2. Scaling Analysis

  • At NLL 1.988, LatentLoop reaches that loss with 1.85 × 10^20 FLOPs, saving 19.0% versus MoE Base and 7.9% versus MoE Loop.Validation NLL is lower than both baselines at all 15 matched budget endpoints.Source: §4.2 · pp. 8–9, Figure 4 · p. 8

4.3. Downstream Evaluation

  • LatentLoop has the highest average score in all six size-and-shot settings, with gains over both baselines across 0-shot and 5-shot evaluations.Against MoE Loop, gains range from 0.73 to 1.18 percentage points in 5-shot and 0.17 to 0.67 in 0-shot.Source: Table 2 · p. 8, §4.3 · p. 9

4.4. Ablation Studies

  • Among four compute-reduction strategies, sequence aggregation with direct writeback gives the lowest NLL; scaled-decoder writeback lowers NLL to 2.0199 from 2.0268 for direct writeback.Learning both encoder convolution and scaled-decoder coefficients gives 2.0209, so it does not improve on the scaled decoder alone.Source: §4.4 · p. 9, Table 3 · p. 9, Table 4 · p. 10
  • A single learned writeback scalar gives NLL 2.0210, while normalized coefficients give 2.0228 and the scaled decoder gives 2.0199.The authors report that the scalar recovers most of the improvement, with channel and position adjustments adding a smaller reduction.Source: §4.4 · pp. 9–10
  • With expert counts reduced to match parameter budgets, independent attention improves both compressed-visit variants but raises ordinary MoE Loop NLL from 2.0293 to 2.0319.This comparison accounts for the extra attention parameters, unlike the cited Wang et al. (2026c) comparison.Source: §4.4 · p. 10, Table 5 · p. 10

4.5. Analysis

  • With three or four visits, learned residual coefficients increase with visit index; at 90% training, consecutive ratios are about 2.4× and 2.0×, respectively.For writeback, both receiver positions learn scales below their unit initialization, with final channel means of approximately 0.199 and 0.160; the nearer receiver learns the higher mean.Source: §4.5 · p. 11, Figure 5 · p. 11

5. Related Work

Related work links recurrent depth and sparse expert computation to budget-matched looped MoE; LatentLoop uses sequence compression to allocate computation asymmetrically across visits.

Mixture-of-experts Transformers

  • Sparse MoE activates only a subset of experts per token, allowing parameter capacity to grow without proportional per-token computation.The cited background traces this gating approach to Shazeer et al. (2017) and sparse MoE Transformers to Fedus et al. (2022) and Liu et al. (2024).Source: §5 · pp. 11–12
  • LatentLoop revisits an ordered stack of middle layers, sharing routers and experts across visits while keeping attention blocks independent.MoRE (Qiu et al., 2026) instead lets adjacent layers draw from a shared expert pool while retaining separate attention and routers.Source: §5 · p. 12

Looped Transformers

  • Looped-model comparisons can favor extra computation or reduce stored parameters, motivating looped MoE comparisons that jointly match parameter and compute budgets.Prior recurrent approaches include Universal Transformers (Dehghani et al., 2018), deep equilibrium models, Huginn, and Ouro (Zhu et al., 2025).Source: §5 · p. 12

Looped MoE Transformers

  • SMELT (Wang et al., 2026b) matches per-token FLOPs, stored backbone parameters, and KV cache; LatentLoop builds on its setting to change compute allocation within a two-visit loop.Prior looped-MoE work also studies routing divergence, iteration-conditioned normalization, and recurrence-sparsity tradeoffs (Lee et al., 2026; Chen et al., 2026a,b).Source: §5 · p. 12

Compressed computation

  • Unlike SpiralFormer’s multiresolution schedule for hierarchical dependencies (Yu et al., 2026), LatentLoop uses compression to shift computation across visits and increase model width.The paper compares aggregation and writeback operators with SpiralFormer-style variants and reports lower validation NLL for its design.Source: §5 · p. 12

6. Conclusion

The paper studies compute allocation across loop visits under matched budgets and proposes LatentLoop, which compresses the first visit and reallocates saved computation to model width. Across the tested configurations, LatentLoop improves validation NLL and downstream averages over both baselines.

6. Conclusion

  • Across 15 configurations at the 1.3B, 3B, and 7B tiers, LatentLoop lowers validation NLL than MoE Base and MoE Loop and scores higher on nine benchmarks in both shot settings.The largest training budget exceeds 10^21 FLOPs.Source: §6 · p. 13
  • LatentLoop compresses the first visit’s sequence, processes the full sequence on the second, and reallocates saved computation to model width.Learned writeback coefficients calibrate the latent increment; MoE routers and experts are shared while attention blocks remain independent.Source: §6 · p. 13
  • Learned and fixed residual-coefficient experiments reveal different preferred update scales across visits, motivating asymmetric compute allocation.Reducing the first visit’s scale can preserve or improve performance, whereas the same reductions to the second visit hurt performance.Source: §6 · p. 13

Limitations

The mechanisms underlying LatentLoop’s learned residual-coefficient asymmetry remain incompletely understood, particularly for models with three or four loop visits.

Limitations

  • The mechanisms behind the learned asymmetry remain incompletely understood, especially the residual-coefficient patterns with three and four visits.The paper leaves a more detailed investigation of these patterns to future work.Source: Limitations · p. 13
  • Further study of the mechanisms and roles of different loop visits may guide more effective compute allocation for looped MoE models.Source: Limitations · p. 13
  • LatentLoop’s asymmetric allocation is motivated by observed residual-coefficient asymmetry, but the underlying mechanisms are not fully explained.Source: Limitations · p. 13

A. Model Configurations and Matched Budgets

The appendix describes the model configurations and matching procedure used to compare resource budgets, and reports small deviations from the MoE Base targets.

A. Model Configurations and Matched Budgets

  • Each tier’s architecture remains unchanged across token budgets, and configurations match KV cache, per-step computation, parameter budget, and total training FLOPs.The matching procedure adjusts attention dimensions, model width, stored experts, and finally optimizer steps.Source: A. Model Configurations and Matched… · p. 17
  • Backbone parameters include attention, experts, routers, normalization, and any learned decoder, while total parameters also include embeddings and the vocabulary output projection.Shared parameters are counted once, and LatentLoop’s independent attention parameters are included.Source: A. Model Configurations and Matched… · p. 17
  • Across 45 endpoints, maximum absolute deviations from matched targets are 0.566% for backbone parameters, 0.704% for total parameters, 0.781% for KV storage, and 0.901% for FLOPs per step.Integer step adjustments reduce the maximum total-training-FLOPs deviation to 0.00943%.Source: Table 7 · p. 17, A. Model Configurations and Matched… · p. 18
  • At each tier’s largest endpoint, total computation includes the shared training prefix and that endpoint’s annealing branch, counting the prefix once.Source: A. Model Configurations and Matched… · p. 18, Table 8 · p. 18
  • Training FLOPs count three times the forward estimate, including attention, activated expert matrices, routers, vocabulary projection, and learned decoder operations.Elementwise operations, optimizer updates, communication, padding overhead, and recomputation are excluded.Source: A. Model Configurations and Matched… · p. 18

B. Additional Implementation Details

The appendix specifies implementation and evaluation procedures, including encoder and decoder details, incremental decoding, and benchmark settings. It also reports lower GPU time than MoE Loop at larger tiers, with similar GPU time at 1.3B.

Architecture and data

  • The decoder-only Transformer uses sparse MoE layers, with 12, 16, and 20 backbone layers at the 1.3B, 3B, and 7B tiers, respectively.All models use the same FineWeb-Edu training stream and a separate validation set.Source: B. Additional Implementation Details · pp. 18–19

Training and validation

  • Training uses the same ordered token stream across methods, and final validation evaluates 24,415 fixed windows of length 4096.Validation reports token-weighted cross-entropy without router auxiliary loss.Source: B. Additional Implementation Details · p. 19

Scaling budgets and branches

  • Each architecture shares a mainline across token budgets; annealing branches resume from saved checkpoints and include their shared prefix plus their own annealing computation.The reference budgets provide 15 validation NLL endpoints per architecture.Source: B. Additional Implementation Details · p. 19

Convolution encoder

  • The convolution encoder replaces fixed-mean aggregation with a causal convolution applied independently to each channel, initialized to reproduce mean aggregation.Learnable coefficients adjust local aggregation weights and latent-state scale.Source: B. Additional Implementation Details · p. 19

Decoder parameterization

Causal incremental decoding

  • During cached decoding, the first visit processes completed token pairs and writes increments to their full-resolution positions; the second visit uses the full-sequence cache.Caches are separate by layer execution, and incomplete groups receive no latent update until completed.Source: B. Additional Implementation Details · p. 20

Downstream evaluation

  • Evaluation covers nine tasks in both 0-shot and 5-shot settings, with identical task data, prompts, sampling, and scoring configurations across methods.The reported average weights the nine benchmark scores equally.Source: B. Additional Implementation Details · pp. 20–21

Training time

C. Additional Ablations

Additional ablations favor compressing only the first visit while retaining full resolution in the second, and support shared MoE weights with independent attention blocks. They also indicate that stronger compression and delayed writeback worsen validation NLL.

Attention and MoE Sharing

  • 2.0199 is LatentLoop's NLL with shared MoE and independent attention, versus 2.0237 for full sharing; untying MoE yields 2.0348 or 2.0340.Budget matching reduced expert counts per pool from 76 to 51 with shared attention and from 75 to 50 with independent attention; the ablations use the 3B/40B reference setting.Source: Table 12 · p. 22, C. Additional Ablations · p. 22

Latent Position

  • 2.0199 is the lowest NLL when only the first visit is compressed; compressing the second gives 2.0386, and compressing both gives 2.0467.The full/full control with independent attention blocks has NLL 2.0319.Source: Table 13 · p. 22, C. Additional Ablations · p. 22

Compression Factor and Writeback Delay

  • 2.0257 and 2.0297 are the NLLs for compression factors k=3 and k=4, versus 2.0199 for k=2; delaying receivers one token raises NLL to 2.0269.The authors attribute worse results partly to increased writeback delay; the delayed model's receiver coefficient means are 0.161 and 0.138, below the original values.Source: Table 14 · p. 23, C. Additional Ablations · p. 23, Figure 8 · p. 23

D. Comparison with SpiralFormer

The paper compares its largest LatentLoop model with the largest dense SpiralFormer-L model on nine benchmarks in 0-shot and 5-shot settings. Differences in training data, tokenizers, budgets, and implementation details complicate interpretation of score differences.

D. Comparison with SpiralFormer

  • The largest LatentLoop model is compared with the largest dense SpiralFormer-L model on nine downstream benchmarks in 0-shot and 5-shot settings.The table reports benchmark scores in percent and an unweighted mean of the nine scores.Source: D. Comparison with SpiralFormer · p. 23, Table 15 · p. 24
  • The systems use different training data, tokenizers, training budgets, and other implementation details, which can contribute to observed differences.SpiralFormer scores are quoted from its Appendix H, Table 3 because its implementation was not publicly available.Source: D. Comparison with SpiralFormer · pp. 23–24

E. Full Scaling Results

Across the 15 matched budgets, LatentLoop achieves the lowest validation NLL among the compared methods.

E. Full Scaling Results

  • LatentLoop achieves the lowest validation NLL at all 15 matched budgets across the 45 model and budget configurations.Values use the same fixed 0.1B-token validation set; token counts are adjusted to match total training FLOPs.Source: Table 16 · p. 24, E. Full Scaling Results · p. 24

Related papers

Loading 2610.19975v1…