PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Gang Fu, Adel Javanmard, MohammadHossein Bateni, Vahab Mirrokni

arXiv:2610.06725v1cs.LGcs.AI

TL;DR

Flat routers can leave expert utilization and topology-dependent locality to auxiliary mechanisms or systems heuristics. BRANCH-MoE uses arrival-weighted EMA score anchors in a binary routing tree; across four tabular benchmarks it matches the strongest baselines while producing more localized co-routing on the UCI benchmarks.

  • Problem

    Flat routing provides no distribution-free positive lower bound on every expert’s mass and gives expert labels no topology that encodes hardware locality.

    Source: §1 · p. 1

  • Method

    BRANCH-MoE places experts at binary-tree leaves and centers each node’s routing decision on an arrival-weighted EMA score anchor, with tree prefixes providing hierarchical addresses.

    Source: §1 · p. 2

  • Results

    BRANCH-MoE matches the strongest baselines in held-out loss, leads on Forest Covertype, ties within seed noise on Criteo, HIGGS, and MSD, and reaches normalized tree distance 0.710–0.718 versus the 0.817 random-pair baseline on UCI benchmarks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Under the stated conditions, anchored routing provides a population-level nonstarvation guarantee, while tree paths support localized co-routing and bounds on cross-domain communication under tree-aligned placement.

    Source: §1 · p. 2

    Geometric arguments requiring halfspace sublevel sets apply directly to linear routing or routing linear after a shared feature map, not arbitrary nonlinear node maps.

    Source: §3.1 · p. 7

Abstract

from arXiv · show

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.

1 Introduction

Flat routing leaves utilization guarantees and hardware locality outside the routing geometry. BRANCH-MoE builds anchoring and hierarchical addresses into a binary tree, with theoretical guarantees and benchmark results supporting its design.

Motivation

  • Flat routing offers no distribution-free positive lower bound on expert mass, leaves expert labels without hardware topology, and does not itself guarantee early-collapse prevention.Underused experts can receive too little data and improve slowly; the authors do not claim that every conventional MoE collapses or incurs the same communication cost.Source: §1 · pp. 1–2

Architecture

  • BRANCH-MoE combines a binary tree, arrival-weighted EMA score anchors, and a routing-mask cascade whose prefixes act as hierarchical expert or device addresses.The cascade determines incoming masks for anchor updates; the anchors set routing offsets, and the final masks select dispatched experts.Source: §1 · p. 2
  • The architecture uses one scalar EMA anchor per node and needs no auxiliary balancing loss, separately learned bias, or d-dimensional centroid.Its full soft router evaluates E −1 internal nodes with O(log2 E) sequential depth; inference uses stored anchors as fixed biases.Source: §1 · p. 2

Theory

  • For linear routing with log-concave arrivals, anchoring gives a positive soft leaf-mass bound; under a frozen router, realized expert execution frequency controls expert-SGD convergence.The soft-mass bound does not itself guarantee top-K execution; tree-aligned placement also yields explicit bounds on cross-domain communication.Source: §1 · p. 2

Experiments

  • The evaluation compares BRANCH-MoE with Switch, DeepSeek-V3, Skywork, and deterministic hash routing.The benchmarks are Criteo click-through-rate prediction, Forest Covertype, HIGGS, and YearPredictionMSD.Source: §1 · p. 2
  • BRANCH-MoE matches the strongest baselines in held-out loss, leads on Forest Covertype, and ties within seed noise on Criteo, HIGGS, and MSD.On the three UCI benchmarks, its normalized tree distance is 0.710–0.718 versus the analytical random-pair baseline of 0.817; flat routers remain within 3.5% of random.Source: §1 · p. 2

2 Background and Related Work

Prior routing approaches address sparsity, balance, or stability through varied router or assignment changes, while topology-aware systems work has often handled locality outside the router. BRANCH-MoE instead combines learned token routing with a tree structure that supplies hierarchical addresses and prefix locality.

2 Background and Related Work

  • Switch combines top-1 routing with auxiliary balancing loss and fixed capacities, while DeepSeek-V3 uses dynamic expert biases and a complementary sequencewise balancing loss.BRANCH-MoE’s bias instead negates the arrival-weighted mean score at a node; it is not a load-error feedback variable.Source: §2 · p. 3
  • BASE and Sinkhorn-BASE use batch-coupled balanced assignment, Expert Choice controls workloads with expert-side selection, and Hash Layers use fixed input-dependent routing.BRANCH-MoE retains learned, representation-dependent token routing without globally coupled assignment.Source: §2 · p. 3
  • BRANCH-MoE uses a hierarchy of differentiable branch decisions to route tokens to trainable experts, with the same masks also determining tokens used to update node anchors.Unlike prior topology-aware approaches relying on penalties, replication, or systems scheduling, its tree structure embeds topological locality in the router.Source: §2 · pp. 3–4
  • Tree paths give experts hierarchical addresses, while Figure 1 contrasts unstructured flat labels, unanchored tree splits, and anchored BRANCH-MoE routing.The tree has O(log2 E) sequential routing depth, although a full soft leaf distribution still evaluates all E −1 internal nodes.Source: §2 · p. 3, Figure 1 · p. 4

3 Architecture: BRANCH-MoE

BRANCH-MoE routes through a complete binary tree, anchoring each node’s decision to the score distribution of tokens reaching it. The resulting leaf masks drive top-K expert execution, while linear and nonlinear node maps differ in the geometric interpretation of anchoring.

3 Architecture: BRANCH-MoE

  • A complete binary tree places experts at leaves, with each internal node identified by its depth and within-depth index and each leaf described by its preceding path.The root-to-leaf path supplies a structured address for each expert.Source: §3 · p. 4

3.1 Phase 1: adaptive geometric anchoring

  • Each node’s anchor is an EMA of its arrival-weighted mean score, so tokens contribute in proportion to the routing mass reaching that node.The cap (t −1)/t is a warmup correction that removes initialization bias from early updates.Source: §3.1 · p. 5
  • The active anchor centers the node score: training uses a live batch statistic, while inference uses the stored EMA for deterministic routing independent of batch composition.The anchor is a non-gradient state variable rather than an independently learned routing bias.Source: §3.1 · pp. 4, 6
  • For linear node maps, scalar score anchoring is equivalent to centroid anchoring and requires O(1) persistent state per node rather than O(dim(x)).Anchoring also removes the shared mean component that could dominate router gradients when incoming representations have a large common offset.Source: §3.1 · p. 7
  • For arbitrary nonlinear maps, the centered decision boundary may be curved, piecewise linear, or disconnected, so halfspace geometric arguments do not directly apply.Those arguments apply directly to linear routing or routing that is linear after a shared feature map.Source: §3.1 · p. 7

3.2 Phase 2: expert execution

  • The tree produces a normalized probability distribution over leaf experts; the expert layer selects the K largest entries and can renormalize retained gates when K > 1.Sparse execution evaluates only selected experts, requiring Θ(K/E) of dense expert computation when dispatch overhead is ignored.Source: §3.2 · p. 8

4 Theoretical Analysis

Section 4 analyzes anchor tracking, routing-mass collapse, sparsely selected-expert optimization, and communication locality.

4.1 Arrival-weighted population and anchor tracking

  • The constant-decay EMA trades sampling noise against response lag, while initialization error decays geometrically and asymptotic mean-square fluctuation is bounded by (1−α)s^2/(1+α).Deep-node arrival weights can be concentrated, increasing the variance parameter and risk of numerical underflow.Source: §4.1 · p. 9

4.2 Score anchoring and routing-mass anti-collapse

  • Under log-concave arrival laws and exact population anchors, the linear-routing analysis bounds leaf mass away from collapse.The analysis distinguishes soft routing-mass bounds from whether deterministic top-K truncation retains every leaf.Source: §4.2 · pp. 9–10

4.3 Expert optimization under sparse top-K routing

  • With a frozen router, realized top-K execution probability—not soft leaf mass—sets an expert's convergence rate; higher minimum frequency speeds worst-expert convergence.The stated optimization implication does not claim that a soft-mass bound guarantees positive top-K execution probability.Source: §4.3 · pp. 10–11

4.4 Communication locality from learned tree routing

  • At a depth-r boundary, sufficiently large learned margins can confine top-K experts to one prefix-aligned domain; a sufficient condition is γ > log(r/τK(x)).This per-token locality result is compatible with population-level anti-collapse: different tokens may favor different branches.Source: §4.4 · pp. 11–12

5 Experiments

Across four benchmarks, BRANCH-MoE matches the strongest baselines in task quality and maintains balanced loads without an auxiliary loss; locality appears on UCI but not Criteo.

5.1 Setup

  • The evaluation uses Criteo, Forest Covertype, HIGGS, and YearPredictionMSD, with E=16 experts, top-4 routing, and five seeds.MoE conditions use SwiGLU experts and vary only the gate; node maps are linear or MLPs of width 2, 4, or 8.Source: §5.1 · p. 12

5.2 Results

  • Across four benchmarks, BRANCH-MoE matches the strongest baselines in task quality, maintains balanced loads without an auxiliary loss, and shows locality on UCI but not Criteo.Table 2 reports the main quality, tree-distance, and load-CV comparisons.Source: §5 · p. 12, Table 2 · p. 13
  • 0.2372 cross-entropy is BRANCH-MoE’s lowest Covertype mean; its 0.0034 margin over Switch is resolved by the paired test.The 0.0049 margin over DeepSeek-V3 is also resolved by the paired test, while the equal margin over Skywork is not.Source: §5.2 · p. 12, Table 2 · p. 13
  • 0.710–0.718 tree distance with linear maps is 12–13% below random on the three UCI benchmarks; flat routers remain within 3.5% of random.On Criteo, BRANCH-MoE’s linear-map distance is 0.813, indistinguishable from random; MLP-map values are not read as locality.Source: §5.2 · p. 13
  • On UCI, BRANCH-MoE’s linear-map load CV is 0.09–0.23, comparable to best-map Switch (0.10–0.28) and Skywork (0.14–0.54).On Criteo, its CV is 0.115 versus at most 0.006 for flat routers, with 0.08% of dispatches dropped.Source: §5.2 · p. 13
  • On Covertype, h=8 MLP maps lower cross-entropy by 0.017–0.028 for every learned router, 3.5–5.8× the largest h=8 router gap.MLPs raise BRANCH-MoE’s UCI tree distance by 0.024–0.031; on HIGGS, linear maps are best for Skywork and DeepSeek-V3 and tie for BRANCH-MoE.Source: §5.2 · pp. 13–14, Figure 2 · p. 13

A Proofs

The proofs bound the tracking error of the node anchors and show how routing geometry supports nonvanishing branch mass, expert learning, and localized communication. These guarantees rely on stated distributional and routing conditions.

A Proofs

  • The mask-weighted EMA error combines deterministic lag with martingale noise, whose orthogonality bounds the noise contribution.The recursion also gives E[e_{t+1} | F_t] = αe_t, which yields the stated geometric tracking relation.Source: A Proofs · p. 18
  • Log-concave arrival densities remain log-concave after path masking, supporting positive branch-mass bounds that weaken with anchor error.The proof uses log-concave measure bounds to keep expected branch mass positive.Source: A Proofs · pp. 18–19
  • Under the stated expert-risk conditions, the convergence bound improves with realized execution frequency, and a lower frequency bound yields an exponential transient.The result assumes a frozen router and strongly convex conditional expert risks.Source: A Proofs · p. 19
  • The communication bound limits selected experts to at most 1 + h_r(x) domains, and sufficiently confident preferred-path decisions place them in one depth-r domain.The domain bound is also capped by the top-K expert count.Source: A Proofs · p. 20

B.1 Datasets

The experiments use four tabular benchmarks and compare routing paradigms with different balancing and selection mechanisms. The BRANCH-MoE baseline centers tree decisions on the traffic arriving at each node.

B.1 Datasets

  • Forest Covertype, HIGGS, and YearPredictionMSD are UCI benchmarks; Criteo is the primary large-scale task.The datasets differ in features, splits, and preprocessing, with HIGGS and MSD inputs clipped at ±10σ.Source: §B.1 · pp. 20–21
  • On Covertype, pinning near-constant features to unit variance reduced BRANCH-MoE’s seed-to-seed accuracy spread from 6.2 to 0.46 percentage points.The issue arose from rare binary soil indicators producing spikes under guarded z-scoring.Source: §B.1 · p. 21

B.2 Baselines

  • The baselines cover auxiliary-loss softmax and logit-normalized routing, dynamic-bias selection, and uniform hash assignment independent of semantic content.Switch and Skywork retain auxiliary balancing losses; DeepSeek-V3 adjusts selection bias, while Hash has no router parameters.Source: §B.2 · p. 21
  • BRANCH-MoE places experts at balanced binary-tree leaves and centers each node’s decision on its arriving traffic’s mean score.Training uses the current batch mean; inference uses a non-trainable EMA with decay 0.99, avoiding an auxiliary balancing penalty.Source: §B.2 · pp. 21–22

B.3 Node Maps, Architecture, and Protocol

The protocol isolates routing by varying gate designs within a minimal MoE backbone. It uses matched expert capacity and five seeds, while Criteo’s TPU execution limits reproducibility.

B.3 Node Maps, Architecture, and Protocol

  • Each node map is linear or a two-layer MLP with hidden width h ∈ {2, 4, 8}, swept across routers and datasets.The factor is swept across learned routers and the Hash control on UCI, and across learned routers on Criteo.Source: §B.3 · p. 22
  • The minimal backbone uses an input stage, one MoE block of SwiGLU experts, and a linear head, with dataset-specific input processing.The MoE block has no residual connection.Source: §B.3 · p. 22
  • All conditions use 16 experts, top-4 routing, capacity factor 1.25, and five seeds; only the gate differs.Criteo and UCI use different training schedules and batch sizes.Source: §B.3 · p. 22
  • Criteo comparisons are unpaired because TPU runs are not bitwise reproducible, and same-seed reruns vary about as much as changing seeds.UCI comparisons are paired by seed.Source: §B.3 · p. 22

B.4 Locality Metrics

The section defines tree and index distances for expert co-routing and calibrates them against random routing. Tree distance captures hierarchy more faithfully than linear index gap.

B.4 Locality Metrics

  • Tree distance L measures levels to the lowest common ancestor; under block placement, L ≤ D − m means experts share a device.Thus L is an exact topological proxy for cross-device fanout under that placement.Source: §B.4 · p. 22
  • At E = 16, random pairs have normalized reference distances 0.8167 on the tree metric and 0.3778 on index gap, so lower values indicate locality.The section also reports load coefficient of variation across experts.Source: §B.4 · p. 23
  • Tree distance distinguishes adjacent indices across a power-of-two boundary as maximally separated, unlike linear index gap.For example, experts 7 and 8 have index gap 1 but maximal tree separation.Source: §B.4 · p. 23

C.1 Criteo: Click-Through-Rate Prediction

On Criteo, BRANCH-MoE matches the strongest learned routers rather than establishing superiority; it also maintains expert use, though its load is less tightly balanced than flat routers, and locality does not appear.

C.1 Criteo: Click-Through-Rate Prediction

  • BRANCH-MoE is statistically at parity with the strongest learned routers on Criteo; its best-map lead over Skywork and DeepSeek-V3 is within seed noise, while Hash trails every other router.BRANCH-MoE’s best-map AUC is below DeepSeek-V3’s; the only unambiguous gap is to Hash.Source: §C.1 · p. 23
  • With the linear node map, BRANCH-MoE’s load CV is 0.115 versus at most 0.006 for every flat router, although it keeps every expert in use without an auxiliary loss or bias controller.At capacity factor 1.25, BRANCH-MoE drops 0.08% of dispatches at capacity.Source: §C.1 · p. 24
  • With the linear map, BRANCH-MoE’s normalized tree distance is 0.813±0.012, indistinguishable from the 0.8167 random baseline; Criteo locality therefore does not replicate.The lower tree distance with MLP maps varies across seeds and also appears for flat Skywork, so it is not read as locality.Source: §C.1 · p. 24

C.2 Forest Covertype: Multi-Class Classification

Across the supplied results, BRANCH-MoE’s benchmark standing varies: Forest Covertype results depend on node-map width, HIGGS and MSD learned routers are statistically tied, and UCI co-routing locality distinguishes BRANCH-MoE from flat routers.

C.2 Forest Covertype: Multi-Class Classification

  • At h = 8, BRANCH-MoE has Forest Covertype’s best reported cross-entropy (0.23720), but its margin over Skywork is not resolved by either test and its margin over Switch is resolved only marginally.With linear maps, it again has the lowest mean, but its gap to runner-up Skywork is not resolved; at h = 2 and h = 4, Skywork has the lowest mean within seed noise.Source: §C.2 · p. 24
  • On Forest Covertype, BRANCH-MoE’s load CV is 0.093 with a 1.174× peak load at linear maps, while DeepSeek-V3’s controller is tighter at 0.001 / 1.002×.Skywork is the least balanced learned router at linear maps, with CV 0.248 and peak load 1.366×.Source: §C.2 · p. 25

C.3 HIGGS and YearPredictionMSD

  • On HIGGS, BRANCH-MoE ranks second behind Skywork; on MSD it ranks fourth, and all learned routers are statistically tied on both datasets.The four MSD estimates span only 0.0021 MSE, against seed standard deviations of 0.0015–0.0043.Source: §C.3 · p. 25, Table 6 · p. 25

C.4 Routing Locality Across the Suite

  • On all three UCI benchmarks, BRANCH-MoE is the only router whose co-routing statistics depart materially from the random baseline.Table 7 reports normalized co-routing distances for E = 16 and top-4 routing, averaged over five seeds.Source: §C.4 · p. 25, Table 7 · p. 26
  • On the three UCI datasets, linear BRANCH-MoE reduces tree distance 12–13% and index distance 18–23% below baseline, while every flat router stays within 3.5% of baseline on both.Criteo is the exception: BRANCH-MoE stays at the random baseline there.Source: §C.4 · p. 25

D Description of Datasets

The evaluation covers four tabular datasets spanning binary classification, multiclass classification, and regression, with varied sizes, features, and prediction tasks.

D Description of Datasets

  • Criteo provides a large-scale click-through-rate classification task, with 24 daily training files containing numerical and anonymized categorical features.The dataset has 13 numerical and 26 categorical features; categorical values are hashed to 32-bit integers, and some entries may be missing.Source: D Description of Datasets · p. 26
  • Forest Covertype is a seven-class prediction task over 581,012 forest patches, using 54 cartographic variables that include continuous and binary features.The observations are 30 × 30 meter patches from Colorado’s Roosevelt National Forest.Source: D Description of Datasets · p. 26
  • HIGGS is a binary classification task distinguishing simulated Higgs-signal events from background, using 28 real-valued features.The dataset contains 11 million instances, with 21 low-level and seven derived high-level features.Source: D Description of Datasets · p. 26
  • YearPredictionMSD is a regression task that predicts song release year from 90 audio features and contains 515,345 songs.The recordings are primarily Western commercial releases from 1922 to 2011.Source: D Description of Datasets · p. 26

Related papers

Loading 2610.06725v1…