FlashPapers Explorable
Architecture & Scaling Laws • Oct 2026

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
SJTU • Shanghai AI Lab • USTC • CUHK • NTU

TL;DR

Looped MoE models improve parameter reuse by passing representations through shared middle layers twice. However, the first visit naturally produces much smaller residual updates than the second. Rather than treating both visits symmetrically, LatentLoop compresses the sequence 2× during the first visit, keeps full resolution in the second, and reallocates the freed compute to widen the entire model backbone—slashing training compute by up to 19% under strictly matched budgets.

19.0%
Predicted (fitted) compute saved vs vanilla MoE at target 1.988 NLL
3× Scale Delta
First loop learns residual weights ~0.3×; loop two learns ~1.0×
k = 2 Latents
Mean pair aggregation with causal scaled writeback decoder
< 1% Budgets
Strictly matched parameters, KV cache, FLOPs/step & total FLOPs

Why This Matters

Scaling large language models is hitting thermal, memory, and hardware walls. Two paradigms have emerged to squeeze more intelligence out of each parameter:

  • Mixture-of-Experts (MoE): Activates only a tiny subset of experts per token (e.g., 8 out of 64+), enabling vast parameter capacity without exploding inference latency.
  • Looped Transformers: Re-execute shared layer blocks (e.g., visiting the middle layers twice) to gain the expressive computational depth of a deep model without loading more weights into memory.

Prior state-of-the-art methods like SMELT proved that looped MoE works well when resource-matched. However, all existing looped architectures operate under a fundamental assumption: both passes through the repeated block should be symmetric in computation and sequence resolution.

This paper overturns that assumption. It shows that the initial pass through a looped network performs coarse, low-magnitude refinement, while the second pass performs high-fidelity synthesis. Forcing both passes to consume the exact same FLOPs is wasteful.

The Big Idea

“Probe the contributions of loop visits and allocate compute where it actually matters: compress the sequence by half during the first visit ($T/2$ tokens), run full resolution during the second visit ($T$ tokens), calibrate how latent updates are written back with causal alignment, and reinvest the saved FLOPs into making the entire model wider.”

How It Works: Step 1. The Asymmetry Discovery

To probe what each loop pass does, the authors inserted learnable scalar coefficients $\alpha_1$ and $\alpha_2$ on the residual branch of the shared middle block:

$$h' = h + \alpha_v f_\ell(h), \quad v \in \{1, 2\}$$

Across both 1.3B and 3B models, the optimizer drove $\alpha_1$ down to 0.27–0.36, while $\alpha_2$ remained strong at 0.87–1.09 (a ~3× ratio!). In fixed experiments, reducing the first visit's scale to 0.25 consistently lowered validation NLL, whereas shrinking the second visit devastated performance.

Interactive Explorer: Empirical Residual Weight Grid

Explore validation NLL changes ($\Delta$ NLL vs baseline $(1,1)$) from Figure 2(b).

Model:
Second-Loop Coefficient ($\alpha_2$)
0.125
0.250
0.500
1.000
2.000
First-Loop Coefficient ($\alpha_1$)
Selected Setting
$(\alpha_1=0.250, \, \alpha_2=1.000)$
Validation NLL: 2.208
$\Delta$ vs Baseline $(1.0, 1.0)$: -0.008 (Improvement)
Interpretation: Optimal Asymmetry

Suppressing visit 1 residual scale allows stable representation building while visit 2 drives the final prediction.

How It Works: Step 2. Asymmetric Sequence Compression

If the first loop performs smaller updates, why give it equal compute? LatentLoop tests four reduction strategies:

1. Residual Width Projecting Projects $d \to d/2$. NLL: 2.0361 (Loses model capacity).
2. Expert Intermediate Halving Halves MLP intermediate $d_e$. NLL: 2.0364 (Incomplete expert capacity).
3. Fewer Active Experts Reduces top-$k$ from 8 to 4. NLL: 2.0349 (Restricts routing diversity).
4. Sequence Compression ($k=2$) ⭐ Compresses tokens $T \to T/2$. NLL: 2.0303 (Preserves full width $d$ & expert pool).

Under matched resource budgets, Sequence Compression saves massive compute without restricting the active parameter count per position. LatentLoop reallocates those saved FLOPs to expand the global model width $d$ across all layers!

Interactive Architecture Flow

Hover or click any block to reveal its exact formulation and budget impact.

k = 2 Compression
Input: $T$ Tokens Entry Block ($L/4$ layers) H (d wide) Encoder: Mean Agg ($T/2$) Visit 1: Latent Pass ($F_1$) Untied Attn 1 MoE (Shared) Processes $T/2$ latents Scaled Dec (w₀, w₁) + Visit 2: Full Pass ($F_2$) Untied Attn 2 MoE (Shared) Processes full $T$ tokens Exit Block ($L/4$ layers) Output $T$ Tokens
Interactive Component Inspector

Click or hover over any module in the SVG diagram above to inspect how LatentLoop processes states, handles parameter sharing, and respects causal boundaries.

How It Works: Step 3. Calibrated Causal Writeback

Compressing tokens is simple, but writing the latent updates back into the full-resolution sequence requires careful mathematical hygiene. Two challenges emerge:

  1. Causal Alignment: In causal language modeling, position $t$ cannot attend to future tokens. For $k=2$, latent position $j$ aggregates tokens $H_{2j-1}$ and $H_{2j}$. Hence, the update $\Delta_j$ can only be written to positions $2j$ (the second token of the pair) and $2j+1$ (the immediate next token). Token 1 receives zero update.
  2. Magnitude Calibration: Direct writeback forces latent additions to have scale $1.0$. Guided by the residual scaling discovery, LatentLoop introduces a Scaled Decoder (SD) with learned channel vectors $w_0, w_1 \in \mathbb{R}^d$:
    $$\text{Dec}(\Delta)_{kj+s} = w_s \odot \Delta_j, \quad 0 \le s < k$$

Interactive Causal Alignment & Weight Attenuation Demo

Select a token in the sequence to see its source group, writeback receiver targets, and empirical learned weights ($w_0 \approx 0.199$, $w_1 \approx 0.160$).

Tokens:
Token Selected: Pos 2 (Group 1, Token 2)
Latent Representation $Z_j$: $Z_1 = \text{mean}(H_1, H_2)$
Update Received: $\tilde{H}_2 = H_2 + w_0 \odot \Delta_1$
Learned Scaling Factor ($w$): Mean $w_0 \approx 0.199$

Position 2 is the near receiver ($2j$). It contains the full source pair context, so the network assigns it a robust update weight of 0.199.

Results & Empirical Scaling Laws

The authors trained and evaluated 45 model configurations spanning 1.3B, 3B, and 7B tiers across 15 matched FLOP budgets up to $10^{21}$ FLOPs, strictly adhering to compute-matched constraints.

Validation NLL vs Training Compute

Actual 15 scaling endpoints from Table 17 (App. E).

2.35 2.25 2.15 2.00
MoE Base MoE Loop (SMELT) LatentLoop (Ours)
LatentLoop achieves the lowest validation NLL at every token budget.

Downstream Benchmark Performance

Validation improvements translate directly into superior performance across nine standard benchmarks (LAMBADA, HellaSwag, PIQA, WinoGrande, ARC-E/C, SciQ, and MMLU):

Setting Tokens MoE Base MoE Loop LatentLoop Gain vs Loop
1.3B (5-shot) 40B 51.24% 51.18% 51.91% +0.73%
3B (5-shot) 70B 56.69% 57.33% 58.51% +1.18%
7B (5-shot) 140B 62.19% 62.55% 63.40% +0.85%
3B (0-shot) 70B 55.61% 56.02% 56.69% +0.67%

Component Ablation Analysis

Why does LatentLoop succeed where other modifications stumble? The paper isolates each individual architectural component at the 3B / 40B reference budget:

Untying Attention vs MoE

Untying attention blocks across passes costs very few parameters (offset by just 1 fewer stored expert per layer, 76 → 75), but lowers NLL from 2.0237 → 2.0199. Conversely, untying MoE weights requires sacrificing 25 experts (75 → 50) and degrades performance to 2.0340.

Writeback Latency & Lag

Delaying writeback by just one token shifts receivers from $(2j, 2j+1)$ to $(2j+1, 2j+2)$ and causes NLL to jump from 2.0199 → 2.0269. Empirical learned weights show that more distant context receives substantially less weight.

Which Loop to Compress?

Compressing Visit 1 alone gives 2.0199. Compressing Visit 2 instead degrades to 2.0386. Compressing both visits drops sharply to 2.0467 even though model width grows to 1440! The final loop must preserve full resolution.

Limitations & Open Questions

Underlying Mechanism of Asymmetry Still Unproven

While learned residual coefficients systematically favor small updates in early visits across 2, 3, and 4 loops (exhibiting consistent ~2.0× to 2.5× ratios between consecutive passes), the exact mathematical and mechanistic driver for this emergent hierarchy remains an open question.

Non-Standard Hardware Dimensions: To match parameters and FLOPs to within <1% across all baseline comparisons, intermediate widths (e.g., $d=1296$, $d_h=140$) and expert counts ($E=75$) were customized. LatentLoop's GPU-hours (elapsed time × GPU count) were roughly 9% (3B) and 13% (7B) lower, in a confounded comparison.

Fixed Aggregation Stride: The paper focuses primarily on factor $k=2$ compression. Larger strides ($k=3, 4$) suffered from writeback lag and higher NLL, suggesting that multi-scale hierarchical compression schedules across deeper recurrence iterations require further investigation.

Glossary

Looped Transformer

A Transformer architecture where an identical block of layers is executed multiple times in sequence during the forward pass. This increases computational depth and reasoning ability without storing duplicate weight tensors in GPU memory.

Mixture-of-Experts (MoE)

A sparse neural network architecture that contains a large pool of feed-forward expert networks. For each token, a lightweight gating router dynamically selects a subset of experts (e.g., 8 active out of 64 or 75 total), decoupling parameter count from computational cost per step.

Residual Scaling ($\alpha$)

A scalar multiplier applied to the sub-layer update before adding it back to the highway residual connection: $h' = h + \alpha f(h)$. In LatentLoop, learning or tuning $\alpha_v$ reveals that the first visit prefers much smaller update magnitudes than subsequent visits.

Scaled Decoder (SD)

A parameter-efficient writeback module using learned channel vectors $w_0, w_1 \in \mathbb{R}^d$ ($2d$ parameters total) that scales the latent delta update before writing it into the receiver tokens, calibrating update strength based on causal lag.

Citation

@article{he2026latentloop,
  title   = {LatentLoop: Scaling Looped MoE Models with Asymmetric Compute Allocation},
  author  = {Zefeng He and Yunkai Chai and Xiaoyang Chen and Tong Zhu and Zongliang Wu and Yun Luo and Yafu Li and Yu Cheng},
  journal = {arXiv preprint arXiv:2610.19975},
  year    = {2026}
}
Made with Flash Papers — make your own ↗