LatentLoop makes looped Mixture of Experts models more compute-efficient by having the first loop pass over a sequence compressed 2× and spending the saved compute on a wider model, reaching the same validation loss with 19% less training compute than a standard MoE baseline.
If you want a bigger-feeling model without storing more weights, one trick is to reuse the same layers twice. That’s a looped Transformer: run the middle block, then run it again with the same parameters. You get more depth of computation for free in parameter count, but not in FLOPs: each extra pass still costs compute.
Recent work combined this with sparse Mixture of Experts layers so you can grow the stored expert pool to compensate for the parameter savings from reuse. The specific recipe LatentLoop builds on is SMELT, which loops the middle half of layers twice and carefully matches parameters, KV cache, per-step FLOPs, and total training FLOPs against a non-looped baseline. SMELT showed looping helps under those matched budgets. The question this paper asks: inside that two-visit loop, are both passes doing equally useful work, and if not, should they get equal compute?
The authors first run a diagnostic. They attach a learnable scalar in front of each loop visit’s residual update, so the model can decide how strongly each pass writes into the residual stream. Across model sizes, the first visit learns a much smaller coefficient (around 0.27 to 0.36) while the second stays near 1. Forcing the first visit’s scale down by hand doesn’t hurt and sometimes helps; doing the same to the second visit hurts. So the first pass apparently wants to contribute less.
That motivates the main idea: if the first pass contributes less, give it less compute, and reinvest the savings. They try four ways to shrink the first visit (narrower residual stream, narrower expert FFN, fewer active experts, shorter sequence) and sequence compression wins. In LatentLoop the first visit runs on a sequence that’s half as long, formed by averaging adjacent token states. The second visit runs on the full-length sequence. Saved FLOPs are poured back into model width.
Two design details follow. First, the increment produced on the short sequence has to be written back onto the full sequence. A per-channel, per-receiver-position learnable scale (the scaled decoder) controls how strongly. It starts at 1 and typically learns values around 0.16 to 0.20, echoing the earlier residual-asymmetry finding. Second, because the two visits now see different sequence resolutions, they share Mixture of Experts routers and experts (where most parameters live) but get independent attention blocks. One stored expert per layer is dropped to pay for the extra attention parameters.
H = entry_block(tokens) # full-length states
Z = mean_pool_pairs(H) # length T/2
delta = first_visit(Z) - Z # shared MoE, own attention
H_tilde = H + scaled_writeback(delta, w) # causally aligned, learned w
Y = second_visit(H_tilde) # same shared MoE, own attention
out = exit_block(Y)
Across 15 matched-budget configurations from 1.3B to 7B parameters, with the biggest training run above 10^21 FLOPs, LatentLoop beats both the standard MoE baseline (MoE Base) and the looped baseline (MoE Loop, their SMELT reimplementation) on validation Negative log-likelihood (NLL) at every single endpoint. Fitting a scaling law per architecture, at a loss of 1.988 LatentLoop needs 19.0% less compute than MoE Base and 7.9% less than MoE Loop.
On nine downstream benchmarks (including HellaSwag, PIQA, WinoGrande, ARC, and MMLU in continuation form) evaluated via lm-evaluation-harness, LatentLoop has the highest average score in all six combinations of model size and 0-shot/5-shot. Gains over MoE Loop average roughly 0.17 to 1.18 points depending on setting. These are averages across benchmarks, not per-task wins.
Ablations isolate the pieces at the 3B scale. The scaled writeback alone drops NLL from 2.0268 (direct writeback) to 2.0199. Untying attention helps the compressed-first-visit setup but hurts vanilla MoE Loop under matched budgets, because the comparison honestly pays for the untied attention by cutting one expert. Compressing the second visit instead of the first is clearly worse, and compressing both is worse still. The paper also notes LatentLoop is about 9 to 13% faster in wall-clock GPU time at 3B and 7B than MoE Loop at matched theoretical FLOPs, likely because wider matrices use hardware better.
If you’re pretraining a MoE from scratch and already considering layer reuse, the practical takeaway is that the two visits of a looped block don’t deserve equal compute. Testing a smaller residual coefficient on the first visit costs nothing and gave a small improvement at both sizes they tried. That’s a cheap thing to add to an existing looped MoE recipe.
If you’re going further and want the full LatentLoop design, the authors released code and a project page. Note the asymmetry direction: compress the first pass, not the second. Their ablations show that reversing this is clearly worse under the same budget.
For people not training MoEs: the broader lesson worth sitting with is that under budget-matched comparisons, where to spend FLOPs matters as much as how many you spend. The paper’s budget-matching protocol (KV cache, per-step FLOPs, stored params, total training FLOPs all within ~1%) is a useful template for comparing architectural variants without accidentally handing one of them extra resources.
All results are validation NLL and standard zero/few-shot academic benchmarks on models trained from scratch on FineWeb-Edu. There’s no instruction tuning, no coding or long-context evaluation, and no comparison against open-weight production models. Transfer of the compute savings to those regimes is not demonstrated.
The authors are explicit that they do not understand why the first visit prefers a smaller residual scale, including the regular ratio they see with three and four visits. The design is motivated empirically from that observation. Finally, the headline 19.0% / 7.9% numbers come from fitted scaling-law curves at one specific loss value within the observed range, not from a single head-to-head run at a fixed budget. The per-endpoint wins in the full table are real but smaller in absolute terms.