BRANCH-MoE: Balance-Aware Tree Routing for Large Embedding Models

What it claims

A Google Research team (Fu, Javanmard, Bateni, Mirrokni) reports that tree routing matches flat routers on task loss. Across 5 seeds with E=16 and top-4, BRANCH-MoE has the lowest mean held-out loss on Forest Covertype. On Criteo, HIGGS and YearPredictionMSD it shows no statistically significant difference from the best baseline. The result the summary omits is co-routing locality: experts picked together for one input sit closer in the tree.

• With linear node maps, normalized tree distance between co-selected experts is 0.710–0.718 on all three UCI sets, against an analytical random-pair value of 0.817. That is 12–13% lower.
• Leaf-index gap drops 18–23%.
• Every flat baseline (Switch, DeepSeek-V3, Skywork, hash) stays within 3.5% of random.

This locality is what would let prefix-based device placement cut cross-device traffic.

How it works

• Tree: depth D, E = 2^D leaf experts, E−1 internal nodes.
• Node maps: each node has a scalar score f_j(x). It is either a bias-free linear projection or a private 2-layer ReLU MLP with hidden width h ∈ {2, 4, 8}. The output layer has no bias because anchoring removes additive constants.
• Branch probability: p_j = σ(f_j(x) − a_j).
• Mask cascade: a node's left and right children get mass m·(1−p) and m·p. Each token's mass therefore sums to 1 at every depth, and every node gets strictly positive mass.
• Anchor: an arrival-weighted mean of node scores, where each token is weighted by the mass reaching that node. It is not a load-error feedback signal, unlike DeepSeek-V3's bias.
• Train vs. inference: this works like batch norm.
• Training centers on the live batch statistic.
• Inference uses the stored EMA ν_{t+1} = α_t ν_t + (1−α_t) ν̂_t, with warmup α_t = min(α, (t−1)/t) and one shared α.
• Inference is therefore deterministic and independent of batch composition.
• Cost: sequential depth is O(log₂E), but the full soft router still evaluates all E−1 nodes.
• Numerical safeguard: in reduced precision, deep masks can underflow. The implementation clamps the denominator and skips the EMA update when arrival mass vanishes.

flowchart TD
  X[Token x] --> F[Node score f_j x]
  M[Incoming mask m_j] --> S[Arrival-weighted batch mean]
  F --> S
  S -->|training| A[Anchor = batch mean]
  S --> E[EMA update ν_j]
  E -->|inference| A2[Anchor = stored ν_j]
  A --> P["p_j = σ(f_j − anchor)"]
  A2 --> P
  P --> C[Child masks m·(1−p), m·p]
  C -->|next depth| M

Evidence

• Benchmarks: four tabular datasets (Criteo CTR, Forest Covertype, HIGGS, YearPredictionMSD), E=16, top-4, 5 seeds.
• Loss: wins on Covertype only. The other three are ties within seed noise.
• Balance: expert load stays balanced without an auxiliary penalty.
• Locality: the co-routing figures above are reported for linear node maps on the three UCI sets. The fetched text gives no Criteo locality figure.

Limits

• Benchmark scope: despite the title's "Large Embedding Models", every reported benchmark is tabular at E=16. Nothing in the fetched text covers language models, larger E, or wall-clock or communication measurements. The communication benefit is a theoretical bound, not a measured speedup.
• Theory conditions: the guarantees hold only under stated conditions.
• Non-starvation needs linear node maps and log-concave arrival distributions, and it weakens with anchor error.
• The SGD-rate link needs a frozen router and strongly convex expert risks.
• The authors note that the soft-mass bound does not guarantee top-K execution of an expert.
• Balance is population-level only: there is no exact per-batch load equality, unlike BASE, Sinkhorn or Expert Choice.
• Missing details: the fetched text cuts off in Section 3. Ablations, nonlinear-map locality numbers, hyperparameters such as α, and any code release are not visible here.