How to Loop MoE: Flatten the Experts, Untie the Attention

How to Loop MoE. Flatten the Experts, Untie the Attention. By Shouren Wang and eight others, from Case Western Reserve University, Kyoto University, and one other institution. Published on the twenty-eighth of September. Its code repository is worked out: about ninety-five source files, with a readme and a licence.

Two design philosophies meet here. Looped Transformers reuse one block of layers several times. By spending extra computation, they push a model of fixed size further, and so use its parameters more fully. Sparse mixture-of-experts models, or MoE models, take a different route: they activate only a few of many experts for each token. A looped MoE bridges the two, and it gives MoE models new potential for making better use of their experts. But it raises a practical question, which the title states directly. How should a MoE be looped? Which parts should repeat, which should be shared, and how should the experts be laid out across the loop?

The answer offered is Foil, built from two changes. Both are made with the expert parameters, and the expert compute per token, held fixed, so the comparison isolates structure rather than size. The first change flattens the experts. The second unties the attention.

Flattening halves the number of expert layers, doubles the number of experts in each layer, and doubles the number of passes. The expert budget stays the same, but every routing decision now chooses from a larger pool. Untying gives each pass its own attention parameters, while the experts and the routers remain shared across passes. Put together: shared experts, picked from a wider pool on each pass, with attention that is free to differ from one pass to the next.

Foil clearly outperforms the unflattened looped baseline. At twenty billion tokens, every Foil model has lower pretraining loss than that baseline. At one hundred billion tokens, the loss improves monotonically with the degree of flattening, and the most flattened Foil ends zero point zero one two nat below the baseline, at equal parameters and equal compute. Downstream accuracy is on par or better. Untying the attention, separately, yields routing that is more balanced and more confident at an equal shape.

Ablations analyse why Foil works, and turn the findings into design guidance for looped MoE. Three points come out of them. The returns from looping and the returns from widening the expert layers amplify each other, rather than competing. Routing confidence tracks healthy expert use better than load balance does. And so a sparse looped MoE should use more experts per layer, and more passes.

The title compresses the recipe into two instructions, and the results stand behind each of them. Flatten, so that a fixed budget of experts is spread wider and visited more often. Untie, so that each pass attends in its own way while the experts stay common to all of them. Beyond the recipe, the ablations leave one more practical pointer: when judging whether experts are being used well, routing confidence is a better signal than load balance.