Lesson 7b · Quantile Balancing

Lesson 7 fixed exploding activations with a bounded activation function. This is the other failure mode extreme sparsity creates: keeping nearly a thousand experts per layer evenly loaded. It is the last piece of Stable LatentMoE, and the width axis is complete once it lands.

Why imbalance is expensive

The router learns which experts suit which tokens, and nothing stops it developing favourites. Left alone it will route heavily to a popular few while others receive almost nothing. Both ends of that distribution are bad.

    WHY IMBALANCE IS EXPENSIVE

    overloaded expert            starved expert
    ─────────────────            ──────────────
    its device becomes the       receives too few tokens
    slowest in the batch;        to learn anything useful;
    every other device waits     its parameters are dead weight

    "Imbalanced routing slows expert-parallel training and may
     leave some experts poorly trained."          — §2.3.3

Two symptoms, one cause. With 896 experts per layer, both get worse: more experts means more chances for one to overheat, and more chances for one to die unnoticed.

The standard fix is to add a per-expert bias to the router's scores, used only for choosing experts and deliberately excluded from the mixture weights. That exclusion is the clever part of the existing method: the bias steers dispatch without touching the router's gradients, so balancing does not fight learning.

The question is how to set that bias. The method K3 inherited nudges it by a fixed step: if an expert is above target, push its bias down a little; if below, push it up. It works, and the report names its weakness precisely: the step size trades off slow adaptation against load oscillation. Too small and imbalance persists for thousands of steps; too large and loads swing back and forth. At 896 experts, the report says this exceeds the regime in which existing auxiliary-loss-free bias updates remain well behaved.¹

Try all three settings below. Start with no balancing to see the problem, then step the fixed-step method a few times, then switch to Quantile Balancing and press the button once.

The difference is not incremental. Fixed-step nudging is still badly imbalanced after six steps; Quantile Balancing drops imbalance by an order of magnitude in one, and revives every dead expert.

Why one step is enough

The insight is a reframing. Instead of asking "which direction should this bias move," QB asks "what bias would have produced exactly the target load?" That question has a direct answer, because it is a question about a quantile.

    THE REFRAMING

    fixed-step               quantile
    ──────────               ────────
    "expert 7 is too busy    "expert 7 should serve 64 tokens.
     — nudge it down a bit    Among all tokens' margins for expert 7,
     and see what happens"    the 64th largest is the threshold that
                              admits exactly 64. Set the bias there."

    iterative, tuned          direct, no step size to tune


    HOW THE MARGIN IS OBTAINED

    for each token, route with Top-(k+1) instead of Top-k:

        rank:  1    2    3   …   k   │  k+1
               ●    ●    ●       ●   │   ○
               └── actually used ──┘ │   └── the CUTOFF: the score an
                                     │       expert must beat to get in

    margin for expert j on token i  =  its score  −  that token's cutoff

    bias for expert j  =  −(the (1 − k/n) quantile of its margins)

Quantile Balancing (§2.3.3, eqs. 13–14). Taking the cutoff from Top-(k+1) routing is what makes this cheap: the information needed already falls out of the routing pass, so no separate computation over tokens is required. The 64-token figure here is illustrative, matching the router simulator above, not a value from the report.

What makes this more than a trick is where it comes from. Appendix C derives QB from the optimal balanced assignment problem: maximise total router score subject to every token getting exactly k experts and every expert getting exactly mk/n tokens. That is an integer program, but its linear relaxation is exact for this structure, and in the dual, the expert-side multipliers are the biases.¹ So QB approximates the solution to a problem whose optimum is known, which is a stronger footing than a heuristic that merely works in practice.

show me the actual equation

Routing, from §2.3.3 eq. 13:

s_i = Sigmoid( W_r x_i ) router scores T_i = argtopk( s_i + b ) selection USES the bias p_{i,j} = s_{i,j} / Σ_{r∈T_i} s_{i,r} weights EXCLUDE the bias

The update, eq. 14, where α_i is token i's Top-(k+1) cutoff:

b̂_j^(t+1) ← − quantile_{1−k/n} ( s_{:,j} − α^(t) ) b^(t+1) ← b̂^(t+1) − mean( b̂^(t+1) ) · 1

Three details that each matter:

The quantile is 1 − k/n because the target load is q = mk/n out of m tokens, and q/m = k/n. Setting the threshold at the (q+1)-th largest margin admits exactly q tokens.
The second line subtracts the mean bias. A common offset added to every expert changes no Top-k decision, so removing it keeps the biases from drifting.
The update takes effect only in the next step, so a batch is never routed using a bias derived from itself. The report calls this causality; it prevents the balancing from peeking at its own outcome.

At inference the biases are frozen: balancing is purely a training-time concern.

The engineering problem underneath

An exact quantile needs all the margins in one place, and at K3's scale they number in the millions, spread across many devices and gradient-accumulation steps. Gathering them is not viable. So QB builds a histogram of each expert's margins instead: one all-reduce sums the per-device bin counts, and the quantile is read from the pooled histogram. Because counts add, the result reflects the true global batch regardless of how tokens were sharded, at a cost of only a few hundred bins per expert. This is the kind of detail that separates a method that works on paper from one that ships at 2.8T.¹

Check yourself

What does Quantile Balancing do differently from the fixed-step bias method?

It computes the bias yielding the target load directly
It applies the bias to the mixture weights as well
It adds an auxiliary loss penalising unbalanced routing
It gives popular experts a higher target load to serve

Fixed-step asks which direction to move and by how much, so it needs a tuned step size that trades slow adaptation against oscillation. QB asks what bias would have produced the target load, which is a quantile of the router-score margins and has a direct answer. Appendix C derives it from optimal balanced assignment via linear-programming duality.

Why does the QB bias affect expert selection but never the mixture weights?

So balancing steers dispatch without disturbing the router's gradients
So the bias can be safely discarded once the model is deployed
So selection and weighting can use different score functions
So the histogram estimate stays accurate across sharded devices

If the bias entered the mixture weights, balancing would fight learning: nudging an expert's load would also change how much its output counts. Keeping it out of the weights lets dispatch be balanced while the router's gradient-based optimisation proceeds untouched.

The width axis, complete

    STABLE LATENTMOE — every piece and its job

    LatentMoE          896 experts affordable at half width      L6
    RMSNorm before W↑  routed branch arrives at a stable scale    L6
    SiTU-GLU           activations bounded at 100, for 4-bit      L7
    Quantile Balancing 896 experts balanced in one step           L7b
                       ────────────────────────────────────
                       the "Stable" in Stable LatentMoE


    ✅ SEQUENCE   KDA · chunkwise · g_min · Gated MLA · NoPE      L2–L3
    ✅ DEPTH      Block AttnRes · N = 8 · pseudo-queries          L4
    ✅ WIDTH      LatentMoE · SiTU-GLU · Quantile Balancing       L6–L7b
    ▢  INPUT      MoonViT-V2 native vision                       L8
    ▢  OPTIMISE   Per-Head Muon · scaling law                    L9
    ▢  SYSTEM     serving · cache · quantization                  L11
    ▢  BEHAVIOUR  post-training · effort levels · limitations     L12

All three axes done. You can now account for every architectural component in Figure 2 of the report, and derive its parameter count.

Lesson 8 leaves the backbone for the input: MoonViT-V2, and the decision to train a vision encoder from scratch rather than initialise it from a contrastively pretrained model. The report's justification is a training-stability argument backed by a gradient-norm plot, and it lands on a conclusion worth noticing: contrastive pretraining may simply be unnecessary at this scale.