Abstract
A different allocation
for a shared model.
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 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 1021 FLOPs, LatentLoop consistently lowers validation NLL. At a fitted NLL of 1.988, predicted compute savings are 19.0% relative to MoE Base and 7.9% relative to MoE Loop. 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.
Shared weights.
Asymmetric updates.
We give each visit a learnable residual coefficient, initialized at one. At both model sizes, the first visit learns a much smaller update scale than the second.
The visits reuse parameters, but prefer different update strengths.
A smaller first update.
A better loss.
Fix one visit’s coefficient at one and vary the other. Reducing the first coefficient preserves or improves NLL; applying the same reduction to the second hurts performance.
The best tested setting at both sizes: approximately 0.008 lower NLL at 1.3B and 0.003 at 3B relative to (1, 1).
Could the two visits also benefit from
different compute allocations?
Compress once.
Refine at full
resolution.
The first visit operates on a shorter latent sequence. The second processes every token. We reallocate the saved computation to a wider model.
- 01
Encode the sequence
A simple mean encoder aggregates pairs of states into T/2 latent positions, preserving the channel width.
- 02
Compute in latent space
The first visit processes fewer positions, using the full model width and the same number of active experts per position.
- 03
Calibrate the writeback
Learned coefficients control how strongly the latent increment updates each channel and receiver position, with causal alignment.
- 04
Revisit the full sequence
The second visit processes all T positions. MoE blocks are shared; attention blocks have independent parameters.
Stored parameters · KV cache
FLOPs per step · Total training FLOPs
Better allocation.
More efficient scaling.
LatentLoop achieves lower validation NLL at all 15 matched budget endpoints. Fitted compute-optimal frontiers quantify the computation needed to reach the same loss.
At the same fitted NLL of 1.988, within the observed model and data allocation range. LatentLoop uses approximately 1.85 × 10²⁰ FLOPs.
The improvement also carries over to nine downstream benchmarks, with higher average scores at all three sizes in both 0-shot and 5-shot evaluation.