How to Loop MoE: Flatten the Experts, Untie the Attention
Foil flattens experts and unties attention so looped mixture-of-experts models train with lower loss.
Foil studies how to loop a sparse mixture-of-experts model with expert parameters and per-token expert compute held fixed. It halves expert layers, doubles experts per layer and the number of passes, and gives each pass separate attention parameters while experts and routers stay shared. At 20 billion tokens every Foil model has lower pretraining loss than the unflattened looped baseline. At 100 billion tokens loss improves with more flattening, and the most flattened Foil ends 0.012 nat below the baseline at equal parameters and compute, with downstream accuracy on par or better.
- Foil flattens experts and unties attention while sharing routers.
- Every Foil model beats the baseline loss at 20B tokens.
- Most flattened Foil is 0.012 nat lower at 100B tokens.
- Downstream accuracy is on par or better at equal compute.
- Routing confidence tracks expert use better than load balance.
From the original article
1 min read at huggingface.co
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; while sparse mixture-of-experts (MoE) models activate only a few of many experts for each token. Looped MoE bridges these two design philosophies and gives MoE models new potential for better expert usage, but it raises a question: how to loop a MoE?
By Shouren Wang, Chuang Ma, Mohsen Hariri · huggingface.co
A short excerpt; the article and its images belong to the publisher. Summary, key facts and assessment above are AI-written from the article and can be wrong.