removing the balancers - training nanogpt-2 in moe

24 Oct, 2026

Abstract

Every mixture-of-experts (MoE) model ships with a load balancer. The idea of load balancer is to prevent dead experts - a condition where a few experts take most of the tokens and the rest go idle. In this blogpost we ask whether the balancer is really needed if the router simply starts in the right place?

Here we tried to upcycle nanoGPT's small GPT-2 style model on character-level Tiny Shakespeare into an 8-expert MoE and start its router from the structure of the model's own hidden states - using k-means, and a split that finds groups without being told how many.

With no balancer at all, the data-started router reaches a validation loss of 1.888 and it beats the standard moe training recipe - random initialisation + DeepSeek's loss-free balancer which resulted in 1.912 on all three seeds, beats our dense model benchmark 1.905 and lands close to nanoGPT's published dense result (1.88). The only drawback was a few experts still got idle.

1. Introduction

An MoE layer replaces one big MLP with many smaller ones called experts and a small router decides which expert each token goes to [1]. But the router has a well-known failure. In early training a few experts get slightly more tokens, so they get more gradient, so they get better, so the router sends them even more tokens. The rest ends up starving. This is commonly known expert collapse and leave a large part of the model doing nothing.

The standard industry trick everyone uses here is a load balancer. Switch Transformer adds an auxiliary loss that penalises uneven routing [2]. DeepSeek went a step further with a loss-free balancer - every expert gets a bias that is nudged up when the expert is underused and down when it is overused, and the bias only changes which expert gets picked [3, 4]. Both work in practise but however, they end up pushing towards the same goal - an equal share of tokens for every expert.

But the problem is that data are not split equal. Some characters and patterns are simply more common than others. A balancer that forces equal shares is fighting the data a little bit at every single step. So we try something else - no balancer at all, just start the router in the right place based on the what data we are using, and see what happens.

2. Experiment Setup

We keep the experiment small (I am GPU poor independent researcher).I used my macbook for training the model - about 3 minutes per MoE - and used public number for benchmark.

2.1 Data and model

We use the character-level Tiny Shakespeare setup from nanoGPT [5] - 1.1M characters, 65 unique symbols, and nanoGPT's 90/10 train/validation split. The model is nanoGPT's "MacBook" configuration, a small GPT-2 style transformer.

Table 1: Model and training configuration (nanoGPT's MacBook config).

2.2 Going from dense to MoE

Training an MoE from scratch at 0.8M parameters is noisy, so we use upcycling [6]. We train the dense model for 500 steps, then copy every layer's MLP into 8 identical experts with top-1 routing, and keep training until step 2,000.

2.3 Turning the data into a router

A random router hands each expert an arbitrary slice of the token space, and the race for tokens starts from there. But at the moment of upcycling we already have a trained dense model. So we can look at the exact vectors the router is going to see - before it routes a single token - and let each expert start by owning a natural cluster.

1. Collect what the router sees - from the step-500 dense model, we take the hidden states of 4,096 tokens as they enter each layer's MLP, right after LayerNorm.
2. Group them - we tried two ways.
   • k-means with 8 groups, one per expert. Simple, and it works.
   • A spectral split, where we don't tell it how many groups to find. If the hidden states were pure noise, the eigenvalues of their covariance would follow the Marchenko–Pastur law [8], and the largest one would almost never cross the Tracy–Widom edge [9, 10]. So any direction whose eigenvalue clears that edge is real structure, not noise. We project the tokens onto those directions, plus the least-Gaussian directions found by ICA [11], and cut at the emptiest valley of the density. Then we repeat inside each piece until nothing is left above the edge - until every group looks like noise. It found 31–36 groups per layer, and we keep the 8 biggest (Ignore if you don't understand - I am just blabbering).
3. Point each expert at a group - expert e's router row becomes its group's centre μ_e, rescaled to the length s of a random router row: w_e = s · μ_e / ‖μ_e‖. A token's router score is now simply how much it looks like each group, so from the very first step every expert already owns a region of the token space.

2.4 Baselines and metrics

• Dense - the same model with no experts, trained to step 2,000. This is the GPT-2 baseline.
• Random router, no balancer - plain upcycling with a random router (normal init, std 0.02).
• Random router + DeepSeek balancer - the standard recipe today. After every step, each expert's selection bias moves by γ = 0.001 towards an even share, b_e ← b_e + γ · sign(mean load − load_e) [3].
• Data-started router + DeepSeek balancer - to see whether the start still helps when a balancer is on.

Every setup is run with 3 seeds. We report validation loss on 200 batches (and perplexity, which is just e^loss). An expert counts as dead when it gets under 3% of the tokens over the last 500 steps - a fair share is 12.5%. With 8 experts in each of 4 layers, there are 32 experts per model.

3. Results

3.1 Final validation loss

Table 2: Final validation loss (mean ± std over 3 seeds), perplexity, and dead experts out of 32.

With no balancer, the k-means start lands at 1.888 and the spectral start at 1.890. The standard recipe - random router + DeepSeek's balancer - lands itself at 1.912. That gap of about 0.024 is where we are trying to win :D. Interestingly he standard recipe ends up worse than plain dense (1.905) with dead experts itself.

The start also helps when a balancer is on. k-means + DeepSeek beats random + DeepSeek on every seed (by 0.015–0.025), and keeps almost every expert alive - 0–2 dead instead of 3–5.

3.2 What happens after the switch

Every MoE pays a small price right after the switch for a few hundred steps all of them are worse than the dense model at the same step. However, the data-started routers recover and overtake the dense model around step 1,400–1,500, and the gap keeps growing until the end.

3.3 Dead experts

So did we solve collapse? Not fully. I do have few ideas around it but need to build a solid loss function and weird maths! Without a balancer 4–8 of the 32 experts still end up dead with data start. That is better than a random router without a balancer (8–11 dead) but worse than having DeepSeek's balancer on (3–5 dead with a random start, 0–2 with a data start).

The data start doesn't seems to stop collapse for now but it does make collapse cheaper. The experts that survive the training owns real group of the data and performance.

3.4 Dead experts show up in the weights

Since this whole thing started from random matrix, we posed another question - Can you tell an expert is dead just from its weights without looking at the routing at all?

We tracked two numbers per expert - the heavy-tail exponent α of its weight eigenvalue spectrum, the same quantity used in heavy-tailed self-regularisation [12] (lower α means more learned structure), and how far its spectrum has moved since the switch, measured with the Stieltjes transform.

Table 3: Weight-spectrum signals for dead and live experts (runs without a balancer, last 100 steps, 63 dead out of 288 experts).

Dead experts' spectra move less after the switch and keep lighter tails, while live experts grow heavier ones. From the weights alone, α picks out dead experts with an AUROC of 0.90. So in principle you could audit an MoE checkpoint for dead experts without any routing logs at all :P

4. Discussion

Why does the data start work? The best guess is that collapse is kind of a race and the random router is what starts it but when each expert begins by owning a coherent group of tokens, the gradients it gets are about one kind of token, so it specialises fast, and the router has a real reason to keep sending those tokens there.

A balancer instead referees the race at every step and in this setup the refereeing costs us training loss - it ends up pushing tokens to experts where it is not required.

The results behind Spectral vs k-means are basically tied more or less (1.890 vs 1.888, within seed noise). I would need to further interate on Spectral theory to improve the clustering. The spectral split's real advantage is something we haven't used yet, it decides the number of groups from the data. Here it found 31–36 natural groups per layer but we threw most of it other than top 8.

It is also worth mentioning about the size of the experiment. 0.8M parameters, 1,500 steps after the switch. DeepSeek's balancer is designed for training runs that are orders of magnitude longer, and its gentle step of 0.001 had not revived every expert in our short runs. Idle experts also matter much more at scale - 4–8 idle experts out of 32 would be a lot of wasted parameters in a big model.

The idea of starting a router from cluster centres is not new on its own - Cluster-Aware Upcycling [7] sets the router to cluster centroids too and initialises the experts from each cluster.

What is new here is the question - can the right start replace the balancer - and a split that finds its groups from a noise test rather than a preset k.

5. Future Research

• Does the gap hold at scale - say a GPT-2 small sized MoE on OpenWebText, trained long enough that DeepSeek's balancer has time to do its job?
• If the data has 31–36 natural groups per layer, why use 8 experts? Could the number of experts be read straight off the spectrum?
• Can we keep the data start's loss and still have zero idle experts - for example with a balancer whose target shares come from the size of the data's groups, instead of equal shares for everyone?
• Does the weight-spectrum signal hold up on released MoE checkpoints, where nobody has the routing logs?

6. TL;DR

• No balancer, data-started router: 1.888 validation loss - better than the standard recipe (random router + DeepSeek balancer, 1.912) on all 3 seeds.
• The standard recipe was worse than dense here (1.905) - the data-started MoE was better.
• k-means and the random matrix theory split are tied - the spectral split picks its own number of groups.
• Collapse is reduced, not removed - 4–8 of 32 experts still go idle.
• Dead experts are visible in the weights - ESD α flags them with an AUROC of 0.90.

7. Acknowledgement

Thanks to Andrej Karpathy's nanoGPT, which made the baseline easy to reproduce and gave us a public number to compare against. Claude Opus 5.5 for writing down my experiments and my shower thoughts to the ideas that Opus formulated maths for.

8. References

[1] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, Jeff Dean. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. ICLR, 2017. Paper Link

[2] William Fedus, Barret Zoph, Noam Shazeer. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. JMLR 23(120):1–39, 2022. Paper Link

[3] Lean Wang, Huazuo Gao, Chenggang Zhao, Xu Sun, Damai Dai. Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts. arXiv:2408.15664, 2024. Paper Link

[4] DeepSeek-AI. DeepSeek-V3 Technical Report. arXiv:2412.19437, 2024. Paper Link

[5] Andrej Karpathy. nanoGPT. GitHub, 2022. Link

[6] Aran Komatsuzaki, Joan Puigcerver, James Lee-Thorp, Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay, Mostafa Dehghani, Neil Houlsby. Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints. ICLR, 2023. Paper Link

[7] Sanghyeok Chu, Pyunghwan Ahn, Gwangmo Song, Seung Hwan Kim, Honglak Lee, Bohyung Han. Enhancing Mixture-of-Experts Specialization via Cluster-Aware Upcycling. CVPR, 2026, pp. 11283–11292. Paper Link

[8] V. A. Marchenko, L. A. Pastur. Distribution of eigenvalues for some sets of random matrices. Mathematics of the USSR-Sbornik 1(4):457–483, 1967.

[9] Craig A. Tracy, Harold Widom. Level-spacing distributions and the Airy kernel. Communications in Mathematical Physics 159:151–174, 1994. Paper Link

[10] Iain M. Johnstone. On the distribution of the largest eigenvalue in principal components analysis. Annals of Statistics 29(2):295–327, 2001.

[11] Aapo Hyvärinen, Erkki Oja. Independent component analysis: algorithms and applications. Neural Networks 13(4–5):411–430, 2000.

[12] Charles H. Martin, Michael W. Mahoney. Implicit Self-Regularization in Deep Neural Networks: Evidence from Random Matrix Theory and Implications for Learning. JMLR 22(165):1–73, 2021. Paper Link