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

Self-Routing: Parameter-Free Expert Routing from Hidden States
Apr 1, 2026Jama Hussein Mohamud, Drew Wagner, Mirco RavanelliMixture-Of-Experts Large Language Model Routing
Mila - Quebec AI Institute · Universite de Montreal · Concordia University
ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts
Jun 1, 2026Heng Zhao, Zilei Shao, Guy Van den Broeck +1Large Language Model Routing Probabilistic Inference
Department of Computer Science, University of Virginia, Charlottesville, USA · Department of Computer Science, University of California, Los Angeles, Los Angeles, USA
Elbow-Based MoE Routing: A Training-Free Inference Time Plugin for Expert Selection
Aug 5, 2026Robin Pan, Raymond Liu, Daniel Fang +2Mixture-Of-Expert Inference Mixture-Of-Experts
John A. Paulson School of Engineering and Applied Sciences, Harvard University · Department of Mathematics, Harvard University

Apr 1, 2026, cs.AIJ/K move · Enter open · S save

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

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

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

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