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

Gang Fu; Adel Javanmard; MohammadHossein Bateni; Vahab Mirrokni

cs.LGOct 5, 2026

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

Authors: Gang Fu, Adel Javanmard, MohammadHossein Bateni, Vahab Mirrokni

Organizations: Google Research · University of Southern California

Abstract

Mixture-of-experts (MoE) layers increase model capacity without a proportional increase in per-example computation. However, conventional flat routers can yield imbalanced expert utilization and treat experts as an unstructured collection, whose indices carry no topological meaning. We introduce {\bf BRANCH-MoE}, a routing architecture that places $E$ experts at the leaves of a binary decision tree of depth $\log_2 E$ . At each internal node the branching probability is centered on the arrival-weighted mean score of the traffic reaching that node. This mean is estimated using an exponential moving average, which promotes utilization of both child subtrees without an auxiliary load-balancing loss. We show that this moving-average estimate admits an explicit noise-lag trade-off. We prove that for linear node maps and log-concave arrival distributions, this mechanism prevents routing-mass collapse. We further establish that, under a frozen router, an expert's execution frequency controls its stochastic-gradient convergence rate, and that confident decisions near the root bound cross-device communication when experts are assigned to devices by tree prefix. We evaluate BRANCH-MoE against Switch softmax, DeepSeek-V3 dynamic-bias, Skywork logit-normalized, and deterministic hash routing on Criteo click-through-rate prediction, Forest Covertype, HIGGS, and YearPredictionMSD, using $E=16$ , top- $4$ routing, and five random seeds. Our results show that hierarchical routing can preserve task quality and balanced utilization while inducing a topology that supports localized expert co-activation and reduced communication.

Figures & tables

Figure 1: Architectural progression from conventional flat routing to BRANCH-MoE. (a) A flat Top-

K

router scores all

E

experts independently in

O(E)

work and treats expert indices as unstructured labels. (b) A depth-

D

binary routing tree with the experts as leaves gives

O(\log_{2}E)

sequential routing depth and a structured root-to-leaf binary address per leaf ( 00 , 01 , …), but unanchored internal splits can starve entire subtrees; the full soft distribution over the leaves still requires evaluating all

E-1

internal nodes. (c) BRANCH-MoE centers each internal split

j

on its arrival-weighted EMA score anchor (

p_{j}=\sigma(f_{j}(x)-\nu_{j})

), supporting non-starvation under the conditions of Section 4.2 without an auxiliary load-balancing loss and inducing prefix-aligned co-routing locality within hardware devices.

Symbol	Definition
$B,d_{in}$	Number of resident tokens and hidden dimension
$D,\ E=2^{D}$	Routing-tree depth and number of leaf experts
$\mathcal{T},\mathcal{J},\mathcal{E}$	Binary tree, internal nodes, and leaves
$W\in\mathbb{R}^{(2^{D}-1)\times d_{in}}$	Matrix of learned node directions
$\nu_{j},\ b_{j}=-\nu_{j}$	Node score anchor and induced routing bias
$\bar{x}_{j}$	Arrival-weighted centroid of the tokens reaching node $j$

Table 1: Principal notation.

(a) Best held-out loss $\downarrow$
Method	Criteo (logloss)	Covertype (CE)	HIGGS (logloss)	MSD (MSE)
BRANCH-MoE (best map)	$\mathbf{0.44596}$ $\pm$ 0.00015 h4	$\mathbf{0.2372}$ $\pm$ 0.0013 h8	$0.5256$ $\pm$ 0.0014 h8	$0.6767$ $\pm$ 0.0027 h8
BRANCH-MoE (linear)	$0.44615$ $\pm$ 0.00025	$0.2545$ $\pm$ 0.0041	$0.5259$ $\pm$ 0.0008	$0.6843$ $\pm$ 0.0035
Skywork	$0.44603$ $\pm$ 0.00030 lin	$0.2406$ $\pm$ 0.0036 h8	$\mathbf{0.5247}$ $\pm$ 0.0005 lin	$0.6757$ $\pm$ 0.0015 h2
DeepSeek-V3	$0.44605$ $\pm$ 0.00005 h8	$0.2421$ $\pm$ 0.0033 h8	$0.5264$ $\pm$ 0.0013 lin	$\mathbf{0.6746}$ $\pm$ 0.0043 h8
Switch	$0.44623$ $\pm$ 0.00019 lin	$0.2406$ $\pm$ 0.0028 h8	$0.5276$ $\pm$ 0.0003 h4	$0.6766$ $\pm$ 0.0038 h8

Table 2: Main results (

E{=}16

, top-

4

5

seeds; mean

\pm

sd). (a) Best held-out loss, each router at its best of four node maps (tag), selected on the reporting split; BRANCH-MoE also with the linear map covered by the theory. Bold: best arm. Hash ignores the node map (UCI: best of its four identical runs). (b) Normalized tree distance (random pairs

0.817

; bold: most local on UCI) / load CV of the same arms. † Not read as locality ( Section C.1 ). Per-map results: Appendix C .

Panel of Figure 2 — Figure 2: Node-map effect. (a) Covertype loss of the learned routers vs. node map, $\pm$ sd (Hash, at $0.447$ , is off-scale). (b) BRANCH-MoE’s tree distance vs. node map on UCI; dashed: random pairs ( $0.817$ ); shaded: flat routers at their best maps.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Task	Examples	Features	Seeds	Metric
Criteo (small)	CTR (binary)	${\approx}41.5$ M 2 2 2 The held-out split is exactly $4.42$ M examples (one full evaluation pass of $2{,}158$ steps at batch $2048$ ). The training split is not counted directly: we infer ${\approx}37$ M from the ratio of on-disk sizes ( $5.93$ GB train against $0.71$ GB held-out), which assumes a common mean record size across the two splits. All reported results depend on the training step count, not on this figure, so the corpus size should be read as approximate.	$39$ fields	$5$	Logloss / AUC
Forest Covertype	$7$ -class	$581$ K	$54$	$5$	Cross-entropy / Acc.
HIGGS	Binary	$600$ K	$28$	$5$	Logloss / AUC
YearPredictionMSD	Regression	$515$ K	$90$	$5$	MSE / RMSE (yr)

Appendix

Table 3: Benchmark suite. All MoE conditions use

E=16

experts with top-

4

routing. “Examples” is the full corpus size (train

+

held-out); “Seeds” is the number of independent replications per condition.

Linear node map
Method	Best val. logloss $\downarrow$	Best val. AUC $\uparrow$	Load CV $\downarrow$	Tree $\downarrow$	Index $\downarrow$
Skywork	$\mathbf{0.446034}$	$\mathbf{0.805827}$	$0.0040$	$0.8184$	$\mathbf{0.3778}$
	$\pm 0.000298$	$\pm 0.000184$
DeepSeek-V3	$0.446082$	$0.805799$	$0.0057$	$0.8188$	$0.3849$
	$\pm 0.000270$	$\pm 0.000174$
BRANCH-MoE (ours)	$0.446149$	$0.805707$	$0.1152$	$\mathbf{0.8126}$	$0.3843$

Appendix

Table 4: Criteo,

E=16

, top-

4

, capacity factor

1.25

5

seeds (mean

\pm

std). Logloss is binary cross-entropy; logloss and AUC are the best held-out values over the

15

evaluations. Load CV and the normalized tree and index distances are training-time routing statistics; random baselines are

0.8167

and

0.3778

. Upper block: every router with the linear node map. Lower block: the two routers whose lowest-loss node map is an MLP, at that map (for Skywork and Switch it is the linear map). Bold marks the best logloss and AUC. At

5

vs.

5

seeds, differences below

{\approx}3.7\times 10^{-4}

in logloss or

{\approx}2.7\times 10^{-4}

in AUC are not resolvable (Welch

t

-test,

p<0.05

	Best val. cross-entropy $\downarrow$
Method	linear	$h{=}2$	$h{=}4$	$h{=}8$
BRANCH-MoE (ours)	$\mathbf{0.25451}$	$0.25206$	$0.24394$	$\mathbf{0.23720}$
	$\pm 0.00408$	$\pm 0.00487$	$\pm 0.00356$	$\pm 0.00125$
Skywork	$0.25862$	$\mathbf{0.25169}$	$\mathbf{0.24277}$	$0.24062$
	$\pm 0.00365$	$\pm 0.00324$	$\pm 0.00494$	$\pm 0.00358$
Switch softmax	$0.26150$	$0.25646$	$0.24286$	$0.24062$

Appendix

Table 5: Forest Covertype,

E=16

, top-

4

20

epochs, batch

512

5

paired seeds (mean

\pm

std).

h

is the MLP node-map hidden width ( Section B.3 ). Hash ignores the node map by construction; its four arms are the same configuration run four times and serve as an exact negative control.

	HIGGS			YearPredictionMSD
Method	Logloss $\downarrow$	Tree $\downarrow$	Index $\downarrow$	MSE $\downarrow$	Tree $\downarrow$	Index $\downarrow$
Skywork	$\mathbf{0.5247}$ (lin)	$0.812$	$0.370$	$0.6757$ ( $h{=}2$ )	$0.813$	$0.365$
	$\pm 0.0005$			$\pm 0.0015$
BRANCH-MoE ( $h{=}8$ )	$0.5256$	$0.736$	$0.322$	$0.6767$	$0.749$	$0.324$
	$\pm 0.0014$			$\pm 0.0027$
BRANCH-MoE (linear)	$0.5259$	$\mathbf{0.710}$	$\mathbf{0.292}$	$0.6843$	$\mathbf{0.718}$	$\mathbf{0.301}$

Appendix

Table 6: HIGGS (logloss) and YearPredictionMSD (MSE),

E=16

, top-

4

5

paired seeds (mean

\pm

std). Each router is shown at its best node map, given in parentheses; for BRANCH-MoE we show both

h=8

(best loss) and linear (best locality), since the two differ by

0.0003

on HIGGS. Tree and index distances are normalized against the

0.8167

0.3778

random baselines. Hash ignores the node map, so its four arms are one configuration run four times; the best is shown.

	BRANCH-MoE (linear)		BRANCH-MoE ( $h{=}8$ )		Flat routers
Dataset	Tree	Index	Tree	Index	Tree	Index
Forest Covertype	$\mathbf{0.718}$	$\mathbf{0.312}$	$0.742$	$0.334$	$0.813$ – $0.819$	$0.371$ – $0.385$
HIGGS	$\mathbf{0.710}$	$\mathbf{0.292}$	$0.736$	$0.322$	$0.812$ – $0.817$	$0.370$ – $0.380$
YearPredictionMSD	$\mathbf{0.718}$	$\mathbf{0.301}$	$0.749$	$0.324$	$0.800$ – $0.817$	$0.365$ – $0.378$
Criteo	$0.813$	$0.384$	$0.804$	$0.370$	$0.817$ – $0.821$	$0.378$ – $0.382$
Random baseline	$0.817$	$0.378$	$0.817$	$0.378$	$0.817$	$0.378$

Appendix

Table 7: Normalized co-routing distances (

E=16

, top-

4

; mean over

5

seeds). Flat routers: range over Switch, DeepSeek-V3, Skywork, and Hash, each at the arm shown in Table 2 . Bold: most local.

Explore similar work

Apr 1, 2026cs.AI

Self-Routing: Parameter-Free Expert Routing from Hidden States

Mixture-of-Experts (MoE) layers increase model capacity by activating only a small subset of experts per token, and typically rely on a learned router to map hidden states to expert assignments. In this work, we ask whether a dedicated learned router is strictly necessary for MoE routing. We propose Self-Routing, a parameter-free routing mechanism that uses a designated subspace of the token hidden state directly as expert logits, eliminating the router projection entirely while leaving the rest of the MoE layer unchanged. We evaluate Self-Routing on language modeling across different expert counts and model scales, and on ImageNet-1K classification by comparing it against a standard learned router, random-routing baselines, and dense non-MoE baselines. Our results show that Self-Routing remains competitive with the learned-router baseline while removing all dedicated routing parameters, and yields more balanced expert utilization, with about 17 % higher average normalized routing entropy and no explicit load-balancing loss. On ImageNet-1K with DeiT-S/16, Self-Routing also slightly improves over the corresponding learned-router MoE. These findings suggest that effective MoE routing can emerge from the hidden representation itself without requiring a separate learned router module.

Jama Hussein Mohamud, Drew Wagner, Mirco Ravanelli

Mila - Quebec AI Institute · Universite de Montreal · Concordia University

Mixture-Of-Experts Large Language Model Routing Hidden States

Ask PDF

Jun 1, 2026cs.LG

ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

Mixture-of-Experts (MoE) models scale by activating only a small subset of experts per token. However, training such models remains challenging because top-

k

routing is discrete and non-differentiable, requiring gradient estimators for expert selection whose design remains a central open problem. We introduce ProbMoE, a probabilistic routing framework that models expert selection as a distribution over cardinality-constrained expert subsets and formulates routing as probabilistic inference in this discrete subset space. We first propose ProbMoE Exact-

k

routing, which samples

k

-expert subsets in the forward pass, and the backward pass uses gradients through each expert's exact marginal probability as a tractable surrogate for the true gradient. ProbMoE naturally generalizes to a dynamic-

k

routing setting, where both training and inference constrain the routing cardinality to the same predefined range, allowing adaptive expert allocation per token. Across benchmarks and model backbones, ProbMoE Exact-

k

achieves strong performance compared to competitive baselines, with improved expert utilization and routing diversity; ProbMoE Dynamic-

k

achieves comparable performance with fewer activated experts.

Heng Zhao, Zilei Shao, Guy Van den Broeck +1

Department of Computer Science, University of Virginia, Charlottesville, USA · Department of Computer Science, University of California, Los Angeles, Los Angeles, USA

Large Language Model Routing Probabilistic Inference Cardinality

Ask PDF

Aug 5, 2026cs.LG

Elbow-Based MoE Routing: A Training-Free Inference Time Plugin for Expert Selection

Mixture-of-Experts (MoE) models enable model scaling while maintaining low inference-time compute by activating only a subset of experts per token. However, conventional routing relies on a fixed top-k selection, forcing the model to spend the same compute regardless of how many experts are relevant. We introduce elbow-based routing, a training-free inference-time modification that dynamically adjusts the number of experts on a per-token basis. Our method examines the sorted router probability distribution and identifies an elbow point that separates high- and low-probability experts. We find that most router distributions exhibit clear inflection points suitable for this strategy, and we show both theoretically and empirically that elbow-based routing preserves expert load balance. Experiments on a state-of-the-art MoE model demonstrate an average latency reduction of 5.3% while maintaining accuracy across six benchmarks.

Robin Pan, Raymond Liu, Daniel Fang +2

John A. Paulson School of Engineering and Applied Sciences, Harvard University · Department of Mathematics, Harvard University

Mixture-Of-Expert Inference Mixture-Of-Experts Load Balancing Experts

Ask PDF