LatentLoop: Scaling Looped MoE Models with Asymmetric Compute Allocation

LatentLoop. Scaling Looped MoE Models with Asymmetric Compute Allocation. By Zefeng He and seven others, from Shanghai Jiao Tong University, the Shanghai AI Lab, and three other institutions. Published on the fourteenth of October. Its code repository is public and real: thirteen source files.

Looped Mixture-of-Experts models pass the input through the same block more than once, combining parameter reuse with sparse computation to do better under matched budgets. Earlier work has shown that different visits to the loop can play different roles, even though they share parameters. If the visits do different jobs, then spending the same compute on each of them may not be the best use of the budget. The question is which visit deserves what.

To probe that, the authors take a model that applies a shared block twice and learn a separate residual coefficient for each visit. The first visit ends up with a smaller scale. Fixing the coefficients by hand sharpens the picture: reducing the first visit's scale preserves or even improves performance, while the same reductions applied to the second visit hurt it. That asymmetry motivates asymmetric compute allocation, and LatentLoop is the resulting design.

LatentLoop compresses the sequence during the first visit, keeps full resolution during the second, and reallocates the computation saved to increase the hidden size. On top of that, the model learns how strongly the first visit updates the full sequence. The MoE weights are shared across visits, while each visit gets its own attention parameters, to accommodate the different sequence resolutions the two visits work at.

Put together, the first pass is cheaper and coarser, working over a compressed sequence, and how much it changes the full sequence is itself learned rather than fixed. The second pass sees everything at full resolution. The compute that the compression frees is not simply banked as savings. It is spent again, on a wider model.

The evaluation spans fifteen configurations, from one point three billion to seven billion parameters, with the largest budget exceeding ten to the power twenty-one FLOPs. Across all of them, at small scale and large, LatentLoop consistently lowers validation NLL.

At a fitted NLL of one point nine eight eight, the predicted compute savings are nineteen point zero percent relative to vanilla MoE, and seven point nine percent relative to the looped baseline. LatentLoop also improves average scores on nine benchmarks at every size, in both zero-shot and five-shot settings.

The authors frame asymmetric compute allocation as a promising direction for looped MoE models. The path to it is the part to hold onto: a learned coefficient hinted that the first visit mattered less at full strength, a controlled reduction confirmed it, and the design follows directly. Spend less on the visit that needs less, and give the savings to width.