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
LatentLoop improves looped MoE efficiency by compressing the first visit's sequence and reallocating compute to width.
How can we improve the compute efficiency of looped Mixture-of-Experts (MoE) models by using asymmetric compute allocation across loop visits?
Looped Mixture-of-Experts (MoE) models reuse parameters to increase depth, but they often waste compute by treating every loop visit as equally important. LatentLoop identifies that the first visit prefers a smaller residual update scale than the second, so it compresses the first visit's sequence and reallocates the saved computation to increase model width. Across 15 configurations, this asymmetric allocation consistently lowers validation loss and improves downstream benchmark scores while matching total resource budgets.
Paper Primer
The authors observe that in a two-visit looped MoE, the first visit consistently learns a smaller residual update coefficient than the second. LatentLoop exploits this asymmetry: it processes the first visit at a lower sequence resolution, then uses the saved FLOPs to widen the model's hidden layers. The core mechanism is a "scaled decoder" that learns how to write back these compressed latent updates into the full-resolution sequence, ensuring the model effectively integrates the two different processing resolutions.
LatentLoop achieves superior compute efficiency compared to standard looped MoE baselines.
Fitted scaling frontiers across 15 matched-budget configurations. At a fixed validation NLL of 1.988, LatentLoop requires 7.9% less training computation than the looped MoE baseline and 19.0% less than a standard MoE.
Validation loss gains translate directly to downstream reasoning and language tasks.
Average scores across nine benchmarks (e.g., MMLU, HellaSwag, ARC) in 0-shot and 5-shot settings. LatentLoop achieved the highest average score in all six tested size-and-shot settings.
Why does the model use independent attention blocks if it shares MoE experts?
The two visits process different sequence resolutions, which may favor different attention patterns. By untying attention blocks and slightly reducing the stored expert count, the authors maintain the total parameter budget while allowing each visit to optimize its own attention mechanism.
Is this approach limited to two-visit loops?
No; the authors observed a consistent asymmetry in residual scales even when extending the model to three or four visits, suggesting that asymmetric compute allocation remains a viable strategy for deeper loops.
Introduction
We expose and exploit asymmetry in looped MoE visits to improve efficiency.
Looped MoE models reuse parameters by applying the same layer multiple times, but each visit consumes the same amount of compute. This symmetric allocation wastes resources because the two visits do not contribute equally to the final representation.
Looped MoE repeatedly applies a shared expert layer to a sequence, letting the model deepen its computation without storing separate weights for each pass.
Empirically, the first visit learns a much smaller residual coefficient than the second, and fixing the first visit’s scale improves or preserves validation NLL while reducing the second’s scale harms performance. These findings reveal a clear asymmetry that suggests allocating compute unevenly between visits.
Symmetric loop visits waste compute; allocating compute asymmetrically can yield efficiency gains.
Budget Matching Framework
Defines Budget Matching and the MoE baseline before introducing the asymmetric compute trick.
We adopt SMELT’s looped MoE architecture as a baseline and use its budget‑matching framework to compare configurations. The following sections introduce the core Budget Matching mechanism that aligns compute and parameter budgets across model stages.
A Mixture‑of‑Experts layer activates only a small subset of expert sub‑networks per token, keeping overall compute sparse.
We sequentially adjust attention projection size, model width, expert count, and training steps so that the Looped MoE model respects a target budget at each stage.
Step 1: Reduce the projection dimension from 4 to 2 to fit the KV cache, halving the per‑token attention cost.
Step 2: Decrease the model width from 8 to 4, cutting the matrix‑multiply FLOPs per layer by 50 %.
Step 3: Lower the expert count from 4 to 2, reducing the parameter budget by roughly 40 %.
Step 4: Choose 5 optimizer steps (instead of 10) so total training FLOPs drop by half.
Result: The combined adjustments keep the model’s functional capacity while meeting a target budget that is 30 % lower than the original configuration.
Budget Matching changes each component independently, so the model retains expressive power where it matters while trimming compute elsewhere.
How does Budget Matching differ from uniformly scaling all layers?
Uniform scaling shrinks every component by the same factor, which can under‑allocate capacity to critical parts such as attention projections. Budget Matching adjusts each component separately, preserving the relative capacity of attention while reducing compute where it is less sensitive, leading to a more efficient trade‑off.
**Figure 1.** The four steps of budget matching. We adjust attention projection dimensions, model width, the number of stored experts, and training steps in this order.
Table 1 reports the maximum absolute budget deviations (in %) for backbone parameters, KV cache, and FLOPs per token. Across all three categories, our method stays within 0.9 % of the target, improving on SMELT’s deviations of up to 3.9 %.
LatentLoop Architecture
LatentLoop compresses the first loop’s sequence and shares MoE weights, improving efficiency without extra parameters.
Looped models exhibit a hidden asymmetry: the first visit’s residual scale collapses while the second stays near unity, indicating that the two passes do not contribute equally to the final prediction.
**Figure 2.** Residual scale asymmetry at 1.3B and 3B. (a) Learned coefficients for the two visits. (b) Validation NLL with fixed coefficients; rows and columns correspond to the first and second visits. Colors show differences from the same size (1, 1) baseline using scores rounded to three decimals.
Allocate fewer FLOPs to the first loop by compressing its representation, then redirect the saved budget to the second loop where full resolution is needed.
Start from baseline NLL = 2.05.
Apply $\alpha$₁ = 0.25 → NLL drops to 2.042.
Keep $\alpha$₂ = 1 → no further change.
Reducing the first‑visit coefficient improves the loss without any extra parameters, confirming that the first loop is over‑parameterized.
How does this differ from the symmetric “1/N” scaling used in prior work?
Symmetric scaling forces both visits to use the same coefficient, which the data shows is sub‑optimal; asymmetric allocation lets the model learn a small coefficient for the first pass while preserving full strength for the second.
Four reduction ideas were evaluated: halving residual width, halving expert intermediate width, halving active experts, and compressing the sequence. The sequence‑compression variant kept full model width and expert count, reallocating the saved compute to a wider backbone, and achieved the best validation NLL.
**Figure 3.** LatentLoop Architecture. LatentLoop compresses the first visit's sequence and retains full resolution in the second, sharing MoE weights with independent attention blocks.
The first loop operates on a shortened latent sequence, writes a correction back to the original states, and the second loop refines the full‑resolution sequence using the same MoE weights.
Compute Z₁ = [(h₁₁ + h₂₁)/2, (h₁₂ + h₂₂)/2]; Z₂ = [(h₃₁ + h₄₁)/2, (h₃₂ + h₄₂)/2].
Run the first visit F₁ on Z, obtain $\Delta$ = F₁(Z) − Z.
Decode: add w₀ ⊙ $\Delta$₁ to positions 2 and 3, add w₁ ⊙ $\Delta$₂ to positions 4 (if it exists).
Resulting H is updated at positions 2‑4; position 1 receives no update.
The mean encoder preserves the average of each token pair, and the learnable decoder lets the model control how strongly each latent correction influences its two source tokens.
Why not simply drop the first visit altogether instead of compressing it?
Removing the first visit eliminates a useful early‑stage refinement; compressing it retains a cheap correction that improves the representation before the full‑resolution pass, yielding lower NLL without extra compute.
Calibrated Writeback uses a parameter‑free mean encoder and a learnable decoder with two channel‑wise vectors w₀, w₁. With k = 2, only 2d extra parameters are added, and the decoder rescales the latent increment before adding it back to H.
Downstream Performance
LatentLoop consistently beats MoE baselines on validation NLL and downstream scores at equal compute budgets.
LatentLoop achieves lower validation NLL than both MoE Base and MoE Loop across all 15 matched‑budget endpoints.
At the compute‑optimal frontier, LatentLoop reaches NLL = 1.988 with 1.85×10²⁰ FLOPs, saving 19.0 % of computation versus MoE Base and 7.9 % versus MoE Loop.
Implementation follows the sparse MoE Transformer style of Qwen3 with eight active experts per token. LatentLoop uses $k=2$, mean aggregation, and the scaled decoder while sharing MoE weights across layers. Baselines include MoE Base (standard MoE with 64 stored experts per layer) and MoE Loop (the looped architecture of SMELT, visiting the middle half of layers twice). All methods are trained on the same FineWeb‑Edu stream and evaluated on a 0.1 B‑token validation set.
**Table 2.** Downstream benchmark scores (%). Params denote approximate total/active sizes; token budgets are FLOPs matching reference labels. LD-O/LD-S denote LAMBADA OpenAI/Standard, HS is HellaSwag, WG is WinoGrande, and MMLU uses continuation scoring.
The scaling analysis fits an additive power‑law loss surface $L(F,D)=E+A\,F^{-α}+K\,D^{-β}$ to each architecture’s FLOP‑token frontier. Optimizing over $F$ and $D$ yields the predicted compute‑optimal NLL curves shown in the appendix.
LatentLoop outperforms MoE Loop at matched budgets.
Scaling Frontiers
LatentLoop beats baselines on compute efficiency and downstream scores.
LatentLoop reallocates compute asymmetrically: the first MoE visit is compressed while the second retains full resolution, yielding a more efficient looped MoE.
LatentLoop reduces validation NLL by 19.0 % relative to MoE Base at comparable FLOPs.
Figure 4 shows the fitted frontier for LatentLoop lying 19.0 % below the MoE Base curve.
LatentLoop attains the highest average downstream score across all model sizes and shot settings, improving up to 1.18 percentage points over MoE Loop in 5‑shot.
Table 2 reports the average scores; the largest 5‑shot gain is 1.18 pp.
**Figure 4.** Fitted compute-optimal NLL. Solid segments lie within the convex hull of observed (log F, log D) points; dashed segments involve extrapolation.
Ablation studies (Table 3) compared four initial compute‑reduction strategies; sequence aggregation with direct writeback achieved the lowest NLL and was selected for further refinement.
Parameter‑sharing experiments (Table 5) show that reducing the expert count while untying attention blocks raises NLL for ordinary MoE Loop (2.0319) but lowers it for LatentLoop variants (down to 2.0199), confirming the benefit of shared routers with separate attention.
Figure 5 reveals that residual‑scale coefficients grow with visit index; the ratio between consecutive coefficients is about 2.4× for three visits and 2.0× for four visits, indicating a consistent asymmetry across deeper loops.
Figure 6 shows learned receiver‑position scales well below unity (≈0.199 for the nearer receiver and ≈0.160 for the farther one), reflecting the model’s preference to weight updates from the nearer position more heavily.
Related Work
Survey of prior MoE and looped Transformer approaches relevant to our asymmetric allocation.
Mixture-of-experts (MoE) Transformers distribute computation across multiple expert networks via a learned gating mechanism. Recent sparse MoE Transformers, such as the Qwen3 style, activate only a small subset of experts per token, enabling parameter capacity to grow without proportional compute increase. Works like MoRE share expert pools across adjacent layers while keeping attention and routers independent, a design we also adopt.
Looped Transformers revisit the same block multiple times, decoupling computational depth from stored parameters. Prior models include Universal Transformers with adaptive halting, deep equilibrium models representing repeated computation as a fixed point, and recent recurrent approaches like Huginn and Ouro that demonstrate reasoning gains from additional loop visits. Looped MoE models extend this idea to sparse architectures, matching both parameter and compute budgets.
Compressed computation has been explored in works such as Barrault et al., Hwang et al., and Yu et al., with SpiralFormer applying sequence compression within looped Transformers. In contrast, our approach leverages the MoE architecture to allocate compute asymmetrically across visits, using compression on the first visit and reallocating saved resources to model width, alongside learned writeback calibration and selective parameter sharing. Our experiments compare SpiralFormer‑style aggregation and writeback operators, showing lower validation NLL and superior performance against the largest SpiralFormer‑L.
Ablation Studies
Key ablation results show why asymmetric compute allocation matters.
Recall that Looped MoE models treat every loop visit the same, wasting compute. LatentLoop instead compresses the first visit while keeping the second at full resolution, allocating compute asymmetrically.
**Figure 5.** Learned residual coefficients with three and four visits, for 1.3B and 3B.
**Figure 6.** Learned scaled decoder coefficients at 3B/40B. (a) Channel means during training, with 25th–75th percentile bands. (b) Final channel distributions.
Asymmetric compute allocation reduces validation NLL to 2.0199, a modest but consistent improvement over fully shared configurations.
Full sharing yields NLL 2.0237; the asymmetric variant achieves 2.0199, a $\Delta$ 0.0038.
Table 12 compares four attention‑MoE sharing choices. Sharing both attention and MoE gives NLL 2.0237, while independent attention with shared MoE improves to 2.0199. Untying MoE degrades performance (≈2.0340), confirming that allocating parameters to larger expert pools is more beneficial than separating routers.
Table 13 evaluates latent resolution across the two loop visits. Compressing only the first visit attains the lowest NLL 2.0199. Moving compression to the second visit raises NLL to 2.0386, and compressing both visits further degrades to 2.0467, despite identical budgets, showing that the second‑visit compression is the dominant source of error.
Table 14 explores different compression factors $k$ and a delayed writeback. The baseline $k=2$ (our default) reaches NLL 2.0199. Larger groups ($k=3$ and $k=4$) increase NLL to 2.0257 and 2.0297 respectively, likely because writeback delay grows with group size. Delaying the writeback by one token for $k=2$ also raises NLL to 2.0269, reinforcing the importance of timely updates.
Comparison with SpiralFormer
Scaling results compare LatentLoop to dense baselines across benchmarks and validation NLL.
We now examine how LatentLoop scales compared to dense baselines. Two result sets are shown: downstream benchmark performance and validation NLL across model sizes.
**Table 15.** System comparison on downstream benchmarks (%). Params denote approximate total/active parameter counts. Benchmark abbreviations follow Table 2; Avg is the unweighted mean of the nine scores.
Across all scales, LatentLoop delivers lower NLL and higher downstream scores than dense baselines.
Model Configurations
Provides detailed model, resource, and training configurations for all experiments.
This appendix enumerates the exact configurations used for each method and model size, ensuring that all comparisons are made on matched compute and memory budgets.
**Table 6.** Actual architecture configurations. The Layers column gives stored/executed depth; $d$, $d_e$, and $d_h$ are model width, expert intermediate width, and attention head dimension. $E$ is the stored expert count per layer. Parameter counts are in billions.
**Table 10.** Reference token budgets.
**Table 11.** Downstream task identifiers and scoring metrics. Both shot settings use the same metric. `acc_norm` is the harness length-normalized accuracy.
Across all 45 scaling endpoints, the maximum absolute deviations are 0.566 % for backbone parameters, 0.704 % for total parameters, 0.781 % for KV storage, and 0.901 % for per‑step FLOPs. Integer‑step adjustments further reduce total training FLOPs deviation to 0.00943 %.
FLOPs are counted by treating each multiply‑add as two operations. We include attention projections, causal query–key/value pairs, activated expert matrices, routers, the vocabulary projection, and learned decoder ops. Sequence compression affects both projection costs and the number of attention pairs. Training FLOPs are estimated as three times the forward count; optimizer updates, communication, padding, and recomputation are excluded.
Implementation Details
Implementation specifics for the decoder‑only Transformer with sparse MoE layers.
We employ a decoder‑only Transformer augmented with sparse MoE layers following the Qwen3 design. Each token activates eight experts with normalized routing weights, and the backbone uses RMSNorm, rotary position embeddings, grouped query attention (with query and key normalization), and SwiGLU experts.
Attention operates with eight query heads and two KV groups; the 1.3B, 3B, and 7B tiers contain 12, 16, and 20 backbone layers respectively. All models train on the FineWeb‑Edu sample350BT stream, tokenized with the 32K TinyLlama tokenizer, and follow a warmup‑stable‑decay (WSD) learning‑rate schedule.
Training uses a global batch size of 256 sequences, mixed‑precision BF16, and AdamW ($\beta$₁=0.9, $\beta$₂=0.95) with weight decay 0.1 and gradient clipping norm 1.0. The schedule consists of a 500‑step warmup, peak learning rate $η_{\text{max}}=10^{-3}$, minimum $η_{\text{min}}=10^{-4}$, and a final annealing phase covering the last 20 % of each endpoint.
**Figure 7.** GPU time normalized to MoE Base at each model size. GPU time is elapsed training time multiplied by GPU count.
Scaling budgets are defined in Table 10; each tier has four, five, or six reference token budgets, yielding 15 validation NLL endpoints per architecture. Budgets are labeled by the MoE Base token count, and other architectures adjust token counts to match total FLOPs.
For each architecture we share a mainline that runs to 80 % of the target token budget, then saves a checkpoint. Branches resume from this checkpoint and anneal over the remaining budget, while the mainline continues at constant learning rate for larger endpoints.
The convolution encoder applies a causal kernel per channel: $\text{Enc}(H)_j = \sum_{r \in R} a_r \odot H_{k\,j-r}$ for $1 \le j \le l$, where $R$ are non‑negative offsets, $a_r \in \mathbb{R}^d$ are learnable coefficients, and $\odot$ denotes element‑wise multiplication.
Decoder parameterization uses effective coefficients $w_s = 1 + \delta_s$, with each $\delta_s \in \mathbb{R}^d$ initialized to zero. AdamW optimizes $\delta_s$ with weight decay; the coefficients are shared across token groups and remain unconstrained.
Causal incremental decoding indexes the latent cache by completed token groups; RoPE uses the zero‑based position ID $2j-1$ for each group’s final token, with the one‑based index $2j$. The first visit writes the increment to positions $2j$ and $2j+1$, while the second visit reads the stored increment when the next token arrives.
Downstream evaluation uses lm‑evaluation‑harness 0.4.13 with a fixed adapter; we report token‑weighted cross‑entropy without the router auxiliary loss. Tasks are evaluated in both 0‑shot and 5‑shot settings, using seed 1234 and a single evaluation per setting.
Table 11 lists the nine benchmark tasks (LAMBADA, HellaSwag, PIQA, WinoGrande, ARC‑Easy, ARC‑Challenge, SciQ, MMLU) and the metric (accuracy or normalized accuracy) reported for each.
Despite matching theoretical FLOPs, LatentLoop reduces GPU time by roughly 9 % at 3B and 13 % at 7B compared to MoE Loop, while the 1.3B tier shows comparable GPU time. The savings stem from reallocating the first‑visit computation to model width, enabling larger matrix operations and better hardware utilization.
Training‑time measurement starts the timer immediately before the optimizer loop and includes validation and checkpoint writes; it excludes initialization, queue waits, and post‑run auditing. Resumed mainlines combine the original prefix with the continuation, interpolating elapsed time at resume boundaries.