Does Balance-Aware Tree Routing Improve Expert Utilization in Mixture-of-Experts Models?

Tree routing replaces the flat expert list

A conventional mixture-of-experts router treats experts as an unstructured collection: a gating matrix produces a score per expert, softmax converts scores to probabilities, and the top-K experts are dispatched. The expert indices carry no topological meaning, and nothing in the router itself prevents some experts from being selected far more often than others. BRANCH-MoE changes the assignment rule rather than the expert networks. It places the E experts at the leaves of a balanced binary tree of depth log2 E, so each token descends through a sequence of binary decisions and the leaf scores are products of branch probabilities along each root-to-leaf path.

The paper's motivation is that a flat router can yield imbalanced expert utilization, and that balance is usually restored by an auxiliary load-balancing loss or a batch-coupled assignment. BRANCH-MoE instead builds balance into the geometry of the decision. At each internal node the branching probability is centered on the arrival-weighted mean score of the traffic reaching that node, so the decision boundary sits in the middle of the distribution actually arriving there. The authors describe this as promoting utilization of both child subtrees without an auxiliary load-balancing loss.

Arrival-weighted anchors set each split

The mechanism is an anchor, not a load-error feedback variable. At internal node j the bias is bj = -νj, the negated exponential moving average of the arrival-weighted mean score at that node. During training the anchor is the arrival-weighted mean score of the current batch at that node; at inference it is a non-trainable exponential moving average with decay 0.99, updated at every training step. The paper contrasts this with DeepSeek-V3's dynamic per-expert bias, which increases for underloaded experts and decreases for overloaded ones. BRANCH-MoE's bias anchors the node's decision boundary at the center of its own incoming score distribution rather than reacting to a load error.

The exponential moving average introduces a trade-off the authors make explicit. Theorem 4.2 shows that under stationary anchor estimates the initialization error decays geometrically while the asymptotic mean-square fluctuation is at most (1-alpha)s^2/(1+alpha), so a larger alpha reduces the sampling-noise floor but makes the EMA slower to respond. Proposition 4.3 extends this to a drifting anchor, where large alpha suppresses sampling noise and small alpha reduces lag. The paper also notes that under soft routing every node receives positive mass in exact arithmetic, but deep-node arrival weights can be highly concentrated, increasing variance and the risk of numerical underflow.

Anti-collapse holds for linear routing under log-concave arrivals

The paper proves that arrival-weighted anchoring prevents routing-mass collapse under specific conditions. Lemma 4.4 states that if the input distribution has a log-concave density on its affine support and every ancestor of a node uses an affine score, then the arrival-weighted law at that node is log-concave. Corollary 4.5 then shows that for a linear node map with a nonzero weight vector, at the exact population anchor the smaller of the two child probabilities is at least 1/(2e). With an arbitrary active anchor, the bound becomes (1/2)(1/e - A|delta|)+, where A bounds the centered projected score density and delta is the active-anchor error.

Repeated conditioning along the path to a leaf gives a lower bound on expected leaf mass. In the linear log-concave case with exact population anchors, every leaf has expected routing mass at least (2e)^-D. The authors are careful about what this does and does not imply: these bounds do not by themselves imply that every leaf is retained by deterministic top-K truncation. The paper states that the soft-mass bound removes an important mechanism of expert starvation and should make realized execution frequencies larger and more stable than under unconstrained learned routing, but it does not formally imply a positive top-K execution probability.

Execution frequency sets the optimization bottleneck

A separate result connects utilization to training dynamics. Under a frozen router, the paper defines pe as the actual execution probability of expert e after top-K selection, which is distinct from the expected leaf mass qe. Theorem 4.7 shows that under frozen-router expert SGD with a strongly convex and smooth conditional risk, the expected squared distance to the expert's optimum contracts at a rate governed by pe times the learning rate and strong-convexity constant. If every expert has execution probability at least pmin, the transient falls below epsilon once the number of global training steps reaches roughly (1/(pmin gamma mu)) log(initial distance squared over epsilon).

The implication the authors draw is that realized execution frequency is the optimization bottleneck for sparsely activated experts: a larger minimum frequency yields faster worst-expert convergence in global training steps. BRANCH-MoE is designed to improve this quantity by anchoring each internal decision to the population reaching that node, discouraging persistent high-level imbalances that would suppress every expert in an entire subtree. The paper is explicit that uniform random routing remains the natural benchmark for marginal balance, and that BRANCH-MoE seeks comparable utilization while retaining task-dependent specialization and tree-aligned communication locality.

Confident root decisions bound cross-device communication

The tree structure also has a communication consequence when experts are assigned to devices by tree prefix. The paper analyzes a hardware boundary whose domains align with depth-r routing subtrees, where a token is transmitted once to every contacted domain and may then be replicated locally. Theorem 4.8 shows that the leaf mass outside the preferred depth-r subtree is at most the sum over the first r levels of exp(-gamma_i), where gamma_i is the learned score margin at level i. The number of depth-r communication domains containing at least one selected expert is bounded by the minimum of K and 1 plus the ratio of that outside mass to the smallest retained probability among the top-K experts.

The practical threshold is stated directly: if the sum of exp(-gamma_i) over the first r levels is below the smallest retained top-K probability, then all top-K experts lie in the preferred subtree, the token contacts exactly one domain, and communication cost is just the token representation size times the per-byte transmission cost. If the first r score margins are each at least gamma, the outside mass is at most r exp(-gamma), so gamma greater than log(r/tau_K) is sufficient for one-domain dispatch at the depth-r boundary. The authors note this per-token result is compatible with the population-level anti-collapse result: a token may assign most of its probability to one branch while other tokens favor the opposite branch, so broad expert utilization across tokens can coexist with concentrated, structured assignments for each token.

Task quality stays competitive across four benchmarks

The evaluation compares BRANCH-MoE against four routing paradigms: Switch softmax with an auxiliary balancing loss, DeepSeek-V3 dynamic-bias, Skywork logit-normalized gating, and deterministic hash routing. All conditions use E = 16, top-4 routing, and capacity factor 1.25, and differ only in the gate, so comparisons are at matched capacity. The benchmarks are Criteo click-through-rate prediction as the primary large-scale task, plus Forest Covertype, HIGGS, and YearPredictionMSD as tabular benchmarks that admit many-seed replication. Every condition is replicated over five seeds.

On Criteo, logloss and AUC are the best held-out values over 15 evaluations, and the paper reports that at 5 versus 5 seeds differences below about 3.7e-4 in logloss or 2.7e-4 in AUC are not resolvable by a two-sided Welch t-test at p < 0.05. Criteo comparisons are unpaired because TPU execution is not bitwise reproducible; re-running an identical configuration with the same seed changes held-out logloss about as much as changing the seed does. On the UCI benchmarks the seed governs initialization but not data order, and all arms within a dataset share the identical seed set, so comparisons are paired by seed. The abstract's summary is that hierarchical routing can preserve task quality and balanced utilization while inducing a topology that supports localized expert co-activation and reduced communication.

One preprocessing detail and a single-layer scope

A concrete result on Forest Covertype shows how sensitive tree routers can be to preprocessing. Forty of the 54 features are binary soil-type indicators, and the rarest fires on only 2 of the 464,809 training rows. Standard z-scoring with an epsilon of 1e-7 denominator guard maps a column that is constant in the sample used to fit the normalizer to spikes of magnitude around 1e7 on the rows where it does fire, which corrupts running statistics and induces bimodal expert collapse in tree routers. Pinning any column with adapted variance below 1e-4 to unit variance eliminates the collapse and reduces the seed-to-seed accuracy spread of BRANCH-MoE from 6.2 to 0.46 percentage points.

The study's scope is a single BRANCH-MoE layer, which the authors state does not capture interactions among routing decisions across multiple layers. The theory is also conditional: the anti-collapse guarantee is proved for linear node maps and log-concave arrival distributions, and the convergence result assumes a frozen router with i.i.d. examples and a strongly convex, smooth conditional risk. The communication bound assumes prefix-aligned placement and depends on learned score margins rather than being automatic. The paper's own priority is the trade-off between population-level balance and per-token concentration: arrival-weighted anchoring prevents the token population as a whole from placing arbitrarily little routing mass on either child, but does not require each token to split its mass evenly. That distinction, rather than a claim of uniform per-token routing, is what the evidence supports.

Sources used in this answer