Paper deep dive
BRANCH-MoE: Balance-Aware Tree Routing for Large Embedding Models
Gang Fu, Adel Javanmard, MohammadHossein Bateni, Vahab Mirrokni
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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.
Tags
Links
- Source: https://arxiv.org/abs/2610.06725v1
- Canonical: https://arxiv.org/abs/2610.06725v1
Trouble viewing inline? Open PDF directly →
Full Text
89,788 characters extracted from source content.
Expand or collapse full text
BRANCH-MoE: Balance-Aware Tree Routing for Large Embedding Models Gang Fu Affiliation: Google Research, thomasfu,bateni,mirrokni@google.com Adel Javanmard Affiliation: Google Research, thomasfu,bateni,mirrokni@google.com Affiliation: University of Southern California, ajavanma@usc.edu MohammadHossein Bateni Affiliation: Google Research, thomasfu,bateni,mirrokni@google.com Vahab Mirrokni Affiliation: Google Research, thomasfu,bateni,mirrokni@google.com 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 BRANCH-MoE, a routing architecture that places E experts at the leaves of a binary decision tree of depth log2E _2E. 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=16E=16, top-44 routing, and five random seeds. BRANCH-MoE matches the predictive performance of the strongest baselines, attaining the lowest mean held-out loss on Forest Covertype and showing no statistically significant difference from the best-performing methods on Criteo, HIGGS, or YearPredictionMSD across random seeds. It also maintains balanced expert load without an auxiliary penalty. Unlike flat routers, BRANCH-MoE induces input-conditional co-routing locality: for each input token, the experts selected together tend to lie closer to one another in the routing tree. On all three UCI benchmarks, linear node maps yield a normalized tree distance of 0.7100.710–0.7180.718 between co-selected experts, compared with 0.8170.817 for random expert pairs—a reduction of 1212–13%13\%. Every flat-routing baseline remains within 3.5%3.5\% of the random-pair distance. Together, these 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 Mixture-of-experts (MoE; Jacobs et al., 1991) layers increase parameter capacity while activating only a small subset of expert networks for each token (Shazeer et al., 2017; Lepikhin et al., 2021; Fedus et al., 2022). A standard token-choice layer uses a flat E-way router followed by top-K selection. This design has enabled very large sparse models, but it leaves several properties to losses, constraints, or systems heuristics rather than building them into the routing geometry. Limitations of flat routing. First, a flat softmax or top-K gate provides no distribution-free positive lower bound on the mass assigned to every expert: small initial routing advantages can produce highly unequal update frequencies, and an underused expert, receiving little data, can improve slowly (Eigen et al., 2013; Shazeer et al., 2017). Auxiliary load-balancing losses, capacity constraints, balanced assignment, and expert-choice routing mitigate this problem (Shazeer et al., 2017; Fedus et al., 2022; Lewis et al., 2021; Zhou et al., 2022), but the protection does not follow from the geometry of a flat gate. Second, a flat expert label is an unstructured categorical index that does not encode which experts share a device, node, rack, or other hardware cut, so the router offers no prefix locality for a distributed implementation to exploit. Third, preventing early routing collapse requires an auxiliary load-balancing penalty that competes with the task objective or an external bias controller (Wang et al., 2024; DeepSeek-AI, 2024). We do not claim that every conventional MoE collapses or incurs the same communication cost, only that these guarantees do not arise from flat routing alone. Our proposal, BRANCH-MoE, changes the routing architecture itself. Its novelty is the combination of three coupled mechanisms within one routing tree: (i) a depth-D binary routing tree for E=2DE=2^D experts; (i) EMA score anchors and induced offsets that center every routing decision on the tokens arriving at its node; and (i) a routing-mask cascade whose prefixes can serve as hierarchical expert or device addresses. The cascade determines the incoming mask used to update each node’s anchor, the anchor determines the routing offset, and the final mask selects the experts each token is dispatched to, with the tree path giving each leaf a structured binary address rather than an arbitrary flat label. Our contributions. We summarize our contributions as follows: 1. Architecture. BRANCH-MoE couples a soft mask cascade with one scalar EMA score anchor per node, updated with mask weights and a single shared decay α (Section 3). The anchor induces the inference-time bias bj=−νjb_j=- _j, so no auxiliary balancing loss, separately learned bias, or d-dimensional centroid is needed; the full soft router evaluates E−1E-1 internal nodes in O(log2E)O( _2E) sequential depth, and inference uses the stored anchors as fixed biases without EMA updates. 2. Theory (Section 4). (1) The mask-weighted constant-decay EMA tracks each node’s incoming mean score with an explicit noise–lag trade-off. (2) For linear routing under log-concave arrival distributions, anchoring prevents vanishing expected branch mass and yields a positive soft leaf-mass bound that degrades quantitatively with anchor error. (3) Under a frozen router and strongly convex conditional expert risks, the realized expert-execution frequency controls the convergence rate of expert SGD (the soft-mass bound does not itself guarantee top-K execution). (4) We formalize the communication locality induced by tree routing: a token’s experts are determined by its path, so co-routed tokens share path prefixes and confident upper-level decisions concentrate a token’s routing mass within a shared-prefix subtree; under tree-aligned expert placement, the depth at which two paths diverge bounds whether their dispatches cross a hardware boundary, giving explicit bounds on cross-domain communication (Section 4.4). 3. Experiments. We compare BRANCH-MoE against flat softmax (Switch), auxiliary-loss-free (DeepSeek-V3), logit-normalized (Skywork; Wei et al., 2024), and deterministic hash routing on four tabular benchmarks, namely Criteo click-through-rate prediction (Criteo AI Lab, 2015), Forest Covertype (Blackard and Dean, 1999; Blackard, 1998), HIGGS (Baldi et al., 2014; Whiteson, 2014), and YearPredictionMSD (Bertin-Mahieux et al., 2011; Bertin-Mahieux, 2011) (see Appendix D for description of datasets). As experiments show BRANCH-MoE matches the strongest baselines, with the lowest mean held-out loss on Forest Covertype and ties within seed noise on Criteo, HIGGS, and MSD. On all three UCI benchmarks, it is the only router whose co-routing statistics fall materially below the analytical random-pair baseline, reducing normalized tree distance to 0.7100.710–0.7180.718 against 0.8170.817 (1212–13%13\% lower) and the leaf index gap by 1818–23%23\%, while every flat router stays within 3.5%3.5\% of random (Table 2). 2 Background and Related Work We first recall the flat token-choice router that serves as the reference template. Let Fe(⋅,θe)e=1E\F_e(·; _e)\_e=1^E be expert networks, where θe _e collects the trainable parameters of expert e, and let Wg∈ℝdin×EW_g ^d_in× E be a flat gating matrix. A standard token-choice router forms the score vector s=x⊤Wgs=x W_g, whose component sf=x⊤(Wg):,fs_f=x (W_g)_:,f is the score of expert f, and computes G(x)=softmax(x⊤Wg).G(x)=softmax(x W_g). (1) After retaining the K largest components, the routed output is y(x)=∑e=1ETopK(G(x))eFe(x).y(x)= _e=1^ETopK(G(x))_eF_e(x). (2) Prior work modifies different aspects of this template to improve sparsity, load balance, or routing stability. Switch Transformer specializes it to top-11 routing and combines an auxiliary load-balancing loss with fixed expert capacities (Fedus et al., 2022). DeepSeek-V3 instead adjusts expert selection through a dynamic per-expert bias: the bias increases for underloaded experts and decreases for overloaded experts, while the gate multiplying the expert output is computed from the original, unbiased affinity score (DeepSeek-AI, 2024). Its primary balancing mechanism is therefore auxiliary-loss-free. However, its implementation includes a complementary sequence-wise balancing loss to encourage an even distribution of expert load within each sequence. BRANCH-MoE also uses routing biases, but for a different purpose. Its bias bj=−νjb_j=- _j is the negated EMA of the arrival-weighted mean score at node j, and anchors that node’s decision boundary at the center of its own incoming score distribution; it is not a load-error feedback variable. Other methods alter the assignment rule more fundamentally. BASE constructs a balanced batch assignment through a matching problem (Lewis et al., 2021), while Sinkhorn-BASE approximates such an assignment using entropy-regularized optimal transport and Sinkhorn iterations (Clark et al., 2022). Expert Choice reverses the usual assignment direction: each expert selects a fixed-size bucket of high-scoring tokens, controlling expert workloads while allowing different tokens to be processed by different numbers of experts (Zhou et al., 2022). Hash Layers remove the learned router and assign inputs to experts using a fixed input-dependent hash, thereby avoiding router parameters, balancing losses, and batchwise matching (Roller et al., 2021). These methods respectively impose balance through batch-coupled assignment, expert-side capacity control, or a fixed routing rule. BRANCH-MoE instead retains learned, representation-dependent token routing without solving a globally coupled assignment problem. BRANCH-MoE is closely related to hierarchical conditional-computation architectures, including hierarchical mixtures of experts (Jordan and Jacobs, 1994; Shazeer et al., 2017), soft decision trees (İrsoy et al., 2012; Frosst and Hinton, 2017), tree-based sparse gating mechanisms (Ibrahim et al., 2023), and fast feedforward networks (Belcak and Wattenhofer, 2023). Hierarchical mixtures of experts arrange routing decisions in a tree structure (Jordan and Jacobs, 1994). Similarly, the soft decision tree of Frosst and Hinton uses learned sigmoid splits and path probabilities to distill a neural network into a tree whose leaves represent class distributions (Frosst and Hinton, 2017). BRANCH-MoE employs the hierarchy differently: its soft, fully differentiable branch decisions route tokens to trainable neural experts, and the same cascaded masks used for routing also determine the tokens used to update each node’s score anchor. Notably, BRANCH-MoE does not enforce exact per-batch load equality, as do assignment- constrained methods. Instead, under the conditions stated in our analysis, its anchored hierarchy provides a population-level nonstarvation guarantee while preserving adaptive sparse routing. The resulting tree paths also give experts hierarchical addresses, complementing systems work on token permutation, distributed communication, and topology-aware routing (Rajbhandari et al., 2022; Hwang et al., 2023; He et al., 2022; Chen et al., 2022). Notably, prior topology-aware models generally rely on post-hoc routing penalties, dynamic replication, or systems-level scheduling to minimize cross-node traffic, whereas BRANCH-MoE embeds topological locality intrinsically into the geometric structure of the router. 3 Architecture: BRANCH-MoE 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)O(E) work and treats expert indices as unstructured labels. (b) A depth-D binary routing tree with the experts as leaves gives O(log2E)O( _2E) 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−1E-1 internal nodes. (c) BRANCH-MoE centers each internal split j on its arrival-weighted EMA score anchor (pj=σ(fj(x)−νj)p_j=σ(f_j(x)- _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. Figure 1 illustrates the architecture, a self-balancing routing tree whose leaves are the experts; Table 1 summarizes the notation. Let T be a complete binary tree of depth D with internal-node set J and leaf set ℰE; an internal node (d,j)(d,j) has depth d∈0,…,D−1d∈\0,…,D-1\ and index j∈0,…,2d−1j∈\0,…,2^d-1\ within that depth, and for any nonroot node or leaf v, path(v)path(v) denotes the internal nodes strictly preceding v on its root-to-v path. To avoid overloading the depth index, dind_ in denotes the hidden dimension. Table 1: Principal notation. Symbol Definition B,dinB,d_in Number of resident tokens and hidden dimension D,E=2D,\ E=2^D Routing-tree depth and number of leaf experts ,,ℰT,J,E Binary tree, internal nodes, and leaves W∈ℝ(2D−1)×dinW ^(2^D-1)× d_in Matrix of learned node directions νj,bj=−νj _j,\ b_j=- _j Node score anchor and induced routing bias x¯j x_j Arrival-weighted centroid of the tokens reaching node j P(d)P^(d) Routing probabilities at tree depth d M(d)∈ℝB×2dM^(d) ^B× 2^d Incoming routing masks at depth d Mexperts=M(D)M_ experts=M^(D) Final expert routing mask X∈ℝB×dinX ^B× d_in Input token representations α∈(0,1)α∈(0,1) Single EMA decay used for anchor tracking KS_K Top-K expert set selected for a token Fe(⋅,θe),θeF_e(·; _e), _e Expert e and its trainable parameter vector 3.1 Phase 1: adaptive geometric anchoring Each internal node (d,j)(d,j) carries a trainable scalar-valued node map fj,t(d):ℝdin→ℝf_j,t^(d):R^d_ in , evaluated at training step t, and a scalar anchor νj,t(d)∈ℝ _j,t^(d) . The parameters of fj,t(d)f_j,t^(d) are optimized by ordinary backpropagation through the routing probabilities, whereas νj,t(d) _j,t^(d) is a non-gradient state variable maintained by the EMA Equation 7; there is no independently learned routing bias or additional scalar calibration variable, since the anchor plays that role and is refreshed whenever the node statistics change. Node maps: linear and nonlinear projections. We instantiate fj,t(d)f_j,t^(d) either as a linear projection with a routing direction Wj,t(d)∈ℝdinW_j,t^(d) ^d_ in and no bias term, or as a nonlinear per-node two-layer MLP of hidden width h: fj,t(d)(x) f_j,t^(d)(x) =(Wj,t(d))⊤x, = (W_j,t^(d) ) x, (3) fj,t(d)(x) f_j,t^(d)(x) =(vj,t(d))⊤φ(Uj,t(d)x+cj,t(d)),Uj,t(d)∈ℝh×din,vj,t(d),cj,t(d)∈ℝh, = (v_j,t^(d) ) \! (U_j,t^(d)x+c_j,t^(d) ), U_j,t^(d) ^h× d_ in, v_j,t^(d),c_j,t^(d) ^h, (4) with φ applied elementwise (ReLU in our experiments); the matrix Wt∈ℝ(E−1)×dinW_t ^(E-1)× d_ in stacks the (Wj,t(d))⊤ (W_j,t^(d) ) as its rows. Every node owns a private MLP whose output layer carries no bias, since anchoring removes any additive constant; taking φ=id =id and h=1h=1 recovers Equation 3 up to reparameterization, so linear routing is the degenerate case of nonlinear routing. The added cost is O(hdin)O(hd_ in) parameters per node with h small (h∈2,4,8h∈\2,4,8\ here), and the router can still be implemented as a fused batched computation. Arrival-weighted anchoring. The anchor statistics and routing masks are constructed jointly in a root-to-leaf pass. At training step t, let mj,t(d)(x)m_j,t^(d)(x) denote the probability that token x reaches node (d,j)(d,j), and collect the depth-d masks into mt(d)(x):=(m0,t(d)(x),…,m2d−1,t(d)(x))⊤∈ℝ2dm_t^(d)(x):= (m_0,t^(d)(x),…,m_2^d-1,t^(d)(x) ) ^2^d (abbreviated m(d)(x)m^(d)(x) when the training index is clear); every token reaches the root with unit mass, so m0,t(0)(x)=1m_0,t^(0)(x)=1. Given the depth-d masks and a batch Xt=(x1,…,xB)X_t=(x_1,…,x_B), the normalized arrival weight of token xbx_b at node (d,j)(d,j), the batch anchor statistic (the arrival-weighted mean of the node scores), and the persistent single-decay EMA anchor are ωj,t(d)(xb) _j,t^(d)(x_b) :=mj,t(d)(xb)∑b′=1Bmj,t(d)(xb′), := m_j,t^(d)(x_b) _b =1^Bm_j,t^(d)(x_b ), (5) ν^j,t(d) ν_j,t^(d) :=∑b=1Bωj,t(d)(xb)fj,t(d)(xb), := _b=1^B _j,t^(d)(x_b)\,f_j,t^(d)(x_b), (6) νj,t+1(d) _j,t+1^(d) =αtνj,t(d)+(1−αt)ν^j,t(d),αt=min(α,t−1t),t≥1, = _t _j,t^(d)+(1- _t) ν_j,t^(d), _t= (α, t-1t), t≥ 1, (7) where α∈(0,1)α∈(0,1) is the architectural decay rate and the cap (t−1)/t(t-1)/t is a warmup correction that removes initialization bias from the first updates. Every token thus contributes to every node’s anchor in proportion to the routing mass reaching that node, so the statistic at (d,j)(d,j) is dominated by the token population selected by its ancestors; at the root, all arrival weights equal 1/B1/B and ν^0,t(0) ν_0,t^(0) is the ordinary batch mean of the root scores. Following the batch-normalization convention, the active anchor differs between training and inference; the centered node logit and its soft right-branch probability are aj,t(d) a_j,t^(d) :=ν^j,t(d),during training,νj,t(d),during inference, := cases ν_j,t^(d),&during training,\\[2.0pt] _j,t^(d),&during inference, cases (8) zj,t(d)(x) z_j,t^(d)(x) =fj,t(d)(x)−aj,t(d),pj,t(d)(x)=σ(zj,t(d)(x)), =f_j,t^(d)(x)-a_j,t^(d), p_j,t^(d)(x)=σ\! (z_j,t^(d)(x) ), (9) with σ(⋅)σ(·) the sigmoid function. The live batch statistic makes training-time centering responsive to changes in the upstream representations, while the stored EMA makes inference deterministic and independent of batch composition. Hierarchical mask propagation. The incoming masks for depth d+1d+1 follow from the branch probabilities at depth d: m2j,t(d+1)(x) m_2j,t^(d+1)(x) =mj,t(d)(x)(1−pj,t(d)(x)), =m_j,t^(d)(x) (1-p_j,t^(d)(x) ), m2j+1,t(d+1)(x) m_2j+1,t^(d+1)(x) =mj,t(d)(x)pj,t(d)(x),j=0,…,2d−1. =m_j,t^(d)(x)p_j,t^(d)(x), j=0,…,2^d-1. (10) Consequently, mj,t(d)(x)m_j,t^(d)(x) is the product of the branch probabilities along the unique root-to-node path ending at (d,j)(d,j). For finite logits, every sigmoid probability lies strictly between zero and one, so mj,t(d)(x)∈(0,1]m_j,t^(d)(x)∈(0,1], with equality only at the root: in exact arithmetic, soft routing sends every token to every node with strictly positive mass rather than assigning it to a single subtree. The cascade also preserves total mass, ∑j=02d−1mj,t(d)(x)=1 _j=0^2^d-1m_j,t^(d)(x)=1 for every depth d, by induction from m0,t(0)(x)=1m_0,t^(0)(x)=1, as the two children of (d,j)(d,j) receive masses summing to mj,t(d)(x)m_j,t^(d)(x). For the batch Xt=(x1,…,xB)X_t=(x_1,…,x_B), define [Mt(d)]b,j:=mj,t(d)(xb),[Pt(d)]b,j:=pj,t(d)(xb),[M_t^(d)]_b,j:=m_j,t^(d)(x_b), [P_t^(d)]_b,j:=p_j,t^(d)(x_b), (11) so that Mt(d),Pt(d)∈ℝB×2d,Mt(0)=B.M_t^(d),P_t^(d) ^B× 2^d, M_t^(0)=1_B. The dimensionally explicit batch form of Equation 10 is [Mt(d+1)]:,2j [M_t^(d+1)]_:,2j =[Mt(d)]:,j⊙(1−[Pt(d)]:,j), =[M_t^(d)]_:,j (1-[P_t^(d)]_:,j ), [Mt(d+1)]:,2j+1 [M_t^(d+1)]_:,2j+1 =[Mt(d)]:,j⊙[Pt(d)]:,j,j=0,…,2d−1. =[M_t^(d)]_:,j [P_t^(d)]_:,j, j=0,…,2^d-1. (12) Equivalently, after applying the fixed column permutation that groups the left children together and the right children together, Mt(d+1)=[Mt(d)⊙(1−Pt(d))∥Mt(d)⊙Pt(d)].M_t^(d+1)= [M_t^(d) (1-P_t^(d))\ \|\ M_t^(d) P_t^(d) ]. (13) The denominator in Equation 5 is strictly positive since the probabilities are strictly positive. In reduced precision, however, a depth-d mask is a product of d sigmoid factors and may underflow to zero. The implementation therefore clamps the denominator away from zero and skips the corresponding EMA update if the pre-clamp arrival mass vanishes, leaving νj,t(d) _j,t^(d) unchanged. This is a numerical safeguard rather than a modification of the routing rule. What is anchored, and in which space. Substituting Equation 6 into the training branch of Equation 9 gives ∑b=1Bωj,t(d)(xb)zj,t(d)(xb)=0. _b=1^B _j,t^(d)(x_b)\,z_j,t^(d)(x_b)=0. (14) This identity holds for any node map fj,t(d)f_j,t^(d), whether linear or nonlinear: each node centers its scores relative to the population that actually reaches it, so any routing imbalance must arise from variation of the scores around this population-specific mean rather than from a shared offset in the incoming representations. For the linear node map fj,t(d)(x)=(Wj,t(d))⊤x,f_j,t^(d)(x)= (W_j,t^(d) ) x, define the arrival-weighted centroid x¯j,t(d):=∑b=1Bωj,t(d)(xb)xb. x_j,t^(d):= _b=1^B _j,t^(d)(x_b)x_b. (15) Linearity gives ν^j,t(d)=(Wj,t(d))⊤x¯j,t(d), ν_j,t^(d)= (W_j,t^(d) ) x_j,t^(d), (16) and hence the training-time logit becomes zj,t(d)(x)=(Wj,t(d))⊤(x−x¯j,t(d)).z_j,t^(d)(x)= (W_j,t^(d) ) (x- x_j,t^(d) ). (17) Thus, during training, the separator is the hyperplane with normal Wj,t(d)W_j,t^(d) passing through the centroid of its incoming population. During inference, zj,t(d)(x)=(Wj,t(d))⊤x−νj,t(d),z_j,t^(d)(x)= (W_j,t^(d) ) x- _j,t^(d), so the stored anchor induces the scalar bias bj,t(d)=−νj,t(d).b_j,t^(d)=- _j,t^(d). (18) Scalar anchoring is therefore exactly equivalent to centroid anchoring for the live statistic in the linear case, while requiring only O(1)O(1) persistent state per node rather than O(dim(x))O( (x)).11 1 For the inference statistic, the equivalence is affected by the drift of Wj,t(d)W_j,t^(d) over the EMA window: νj,t(d) _j,t^(d) averages past values of (Wj,τ(d))⊤x¯j,τ(d) (W_j,τ^(d) ) x_j,τ^(d) computed with potentially different router parameters. For a nonlinear node map, the anchored quantity is the scalar node score rather than a point in token space. At training time, the decision boundary is the level set x:fj,t(d)(x)=ν^j,t(d). \x:f_j,t^(d)(x)= ν_j,t^(d) \. This boundary may be curved, piecewise linear, or disconnected, depending on the node map, and need not pass through a centroid in the original token space. What remains exact is the score-centering identity Equation 14. Accordingly, geometric arguments that require halfspace sublevel sets apply directly to linear routing, or to routing that is linear after a shared feature map, but not to an arbitrary nonlinear node map. Gradient geometry. During training, the arrival weights in Equation 5 are detached when they are used to form the live anchor. Their forward values determine which tokens contribute to the node statistic, but gradients from the anchor are not propagated through the routing masks. The live anchor remains differentiable through the node scores themselves. Therefore, for the parameters θj,t(d) _j,t^(d) of node (d,j)(d,j), ∇θj,t(d)zj,t(d)(xb)=∇θj,t(d)fj,t(d)(xb)−∑b′=1Bωj,t(d)(xb′)∇θj,t(d)fj,t(d)(xb′). _ _j,t^(d)z_j,t^(d)(x_b)= _ _j,t^(d)f_j,t^(d)(x_b)- _b =1^B _j,t^(d)(x_b ) _ _j,t^(d)f_j,t^(d)(x_b ). (19) Thus, each node receives parameter gradients centered by the arrival-weighted mean gradient of the population it serves. For the linear map, ∇Wj,t(d)zj,t(d)(xb)=xb−x¯j,t(d). _W_j,t^(d)z_j,t^(d)(x_b)=x_b- x_j,t^(d). (20) For a nonlinear map, the same centering occurs in the corresponding Jacobian features. Anchoring therefore removes the shared mean component that could otherwise dominate the router gradient when the incoming representations contain a large common offset. 3.2 Phase 2: expert execution Phase 1 produces the leaf-mask matrix Mexperts=M(D)∈ℝB×EM_experts=M^(D) ^B× E, with [M(D)]b,e=me(D)(xb)[M^(D)]_b,e=m_e^(D)(x_b). Each leaf probability is the product of the D branch probabilities along its root-to-leaf path, and every internal node splits its incoming mass between its two children with probabilities p and 1−p1-p, so the router outputs a normalized probability distribution over the E leaf experts rather than unnormalized expert scores. The router returns both this distribution and the indices of its K largest entries, leaving top-K sparsification to the expert layer: the same routing distribution thus supports both dense and sparse implementations, and routing statistics and intended expert loads can be measured before capacity constraints alter the realized assignments. Gating. Let K(x)⊆ℰS_K(x) index the K largest entries of m(D)(x)m^(D)(x), with retained mass RK(x):=∑e′∈K(x)me′(D)(x)R_K(x):= _e _K(x)m_e ^(D)(x). The gate applied to expert e is given by m~e(x)=me(D)(x) 1e∈K(x)RK(x)η,η∈0,1, m_e(x)= m_e^(D)(x)\,1\e _K(x)\R_K(x)^η, η∈\0,1\, (21) With η=1η=1, each token’s retained gates are renormalized, after top-K masking, to sum to one; with η=0η=0, they retain their original path probabilities. Importantly, normalization is applied only after top-K masking; unselected experts therefore have no effect on the relative weights of the selected experts. Regarding the choice of η in Equation 21, note that renormalization preserves the relative weights within K(x)S_K(x). For K>1K>1, this produces a convex combination of the selected expert outputs and stabilizes the scale of the mixture. For K=1K=1, however, renormalization would force the unique nonzero gate to equal one, discarding the router’s only scalar measure of confidence in the selected leaf. We therefore use η=K>1η=1\K>1\ as the default choice. Execution. The BRANCH-MoE output is the gated combination of the selected experts, y^(x)=∑e∈K(x)m~e(x)Fe(x). y(x)= _e _K(x) m_e(x)\,F_e(x). (22) The sparse implementation evaluates only the K selected experts, grouping tokens by expert and scattering the gated outputs back to their original positions. Ignoring dispatch overhead, this requires a Θ(K/E) (K/E) fraction of the expert computation of dense evaluation. 4 Theoretical Analysis We analyze four components of BRANCH-MoE: tracking of the arrival-weighted score anchors, prevention of routing-mass collapse, optimization of sparsely selected experts, and communication locality. Proofs are deferred to Appendix A. 4.1 Arrival-weighted population and anchor tracking Fix a node (d,j)(d,j) and temporarily freeze the router and the node score. To avoid excessive subscripts, write ν^t=ν^j,t(d) ν_t= ν_j,t^(d) and νt=νj,t(d) _t= _j,t^(d). Assumption 4.1 (Stationary anchor estimates). There is a deterministic scalar ν¯ ν such that ν^t=ν¯+ξt,[ξt∣ℱt]=0,[ξt2∣ℱt]≤s2. ν_t= ν+ _t, [ _t _t]=0, [ _t^2 _t]≤ s^2. (23) Here ℱtF_t contains the model state and all batches preceding the t-th update, so νt _t is ℱtF_t-measurable and ν^t ν_t is computed from a fresh batch. The initialization ν0 _0 is deterministic. Theorem 4.2 (Stationary behavior of the anchor EMA). Under Assumption 4.1, consider the constant-decay phase of Equation 7, after the warmup has reached αt=α _t=α. Reindexing the first update in this phase as t=0t=0, the anchor satisfies [νt−ν¯] [ _t- ν] =αt(ν0−ν¯), =α^t( _0- ν), (24) [(νt−ν¯)2] [( _t- ν)^2] ≤α2t(ν0−ν¯)2+1−α1+α(1−α2t)s2. ≤α^2t( _0- ν)^2+ 1-α1+α(1-α^2t)s^2. (25) Thus, the initialization error decays geometrically, while the asymptotic mean-square fluctuation is at most (1−α)s2/(1+α)(1-α)s^2/(1+α). The theorem makes the role of α transparent. A larger α reduces the sampling-noise floor but makes the EMA slower to respond. Under soft routing every node receives positive mass in exact arithmetic, but deep-node arrival weights can be highly concentrated, increasing the variance parameter s2s^2 and the risk of numerical underflow. Proposition 4.3 (Tracking a drifting anchor). Under the same constant-decay update, suppose [ν^t∣ℱt]=ν¯t,[(ν^t−ν¯t)2∣ℱt]≤s2,|ν¯t+1−ν¯t|≤v,E[ ν_t _t]= ν_t, [( ν_t- ν_t)^2 _t]≤ s^2, | ν_t+1- ν_t|≤ v, (26) where (ν¯t)( ν_t) and ν0 _0 are deterministic. Then [(νt−ν¯t)2]≤ [( _t- ν_t)^2]≤ (αt|ν0−ν¯0|+v(1−αt)1−α)2 (α^t| _0- ν_0|+ v(1-α^t)1-α )^2 +1−α1+α(1−α2t)s2. + 1-α1+α(1-α^2t)s^2. (27) The two nontransient terms exhibit the tracking trade-off of the single EMA: large α suppresses sampling noise, whereas small α reduces lag. 4.2 Score anchoring and routing-mass anti-collapse Fix the soft router, including its active anchors, and a node (d,j)(d,j), and suppress the fixed step index. Its expected incoming mass is qj(d):=[mj(d)(X)]>0q_j^(d):=E[m_j^(d)(X)]>0. For every measurable set A in the input space, define the arrival-weighted law by j(d)():=[mj(d)(X)X∈]/qj(d)D_j^(d)(A):=E[m_j^(d)(X)1\X \]/q_j^(d). Let νj(d),⋆:=j(d)[fj(d)(X)] _j^(d), :=E_D_j^(d)[f_j^(d)(X)] be the population score anchor and let δj(d):=aj(d)−νj(d),⋆ _j^(d):=a_j^(d)- _j^(d), denote the active-anchor error. The relative masses sent to the right and left children are, respectively, πj,+(d):=j(d)[σ(fj(d)(X)−aj(d))] _j,+^(d):=E_D_j^(d)[σ(f_j^(d)(X)-a_j^(d))] and πj,−(d):=j(d)[σ(aj(d)−fj(d)(X))] _j,-^(d):=E_D_j^(d)[σ(a_j^(d)-f_j^(d)(X))]. Lemma 4.4 (Log-concavity under soft linear routing). Suppose the input distribution has a log-concave density on its affine support and every ancestor of node (d,j)(d,j) uses an affine score. Then the arrival-weighted law j(d)D_j^(d) is log-concave. Corollary 4.5 (Linear routing under a log-concave arrival law). Suppose j(d)D_j^(d) is log-concave on its affine support and fj(d)(x)=(Wj(d))⊤xf_j^(d)(x)= (W_j^(d) ) x (28) with Wj(d)W_j^(d) nonzero on the linear span of that support. At the exact population anchor, minπj,+(d),πj,−(d)≥12e. \ _j,+^(d), _j,-^(d)\≥ 12e. (29) If the centered projected score has density bounded by Aj(d)A_j^(d), then for an arbitrary active anchor, minπj,+(d),πj,−(d)≥12(1e−Aj(d)|δj(d)|)+. \ _j,+^(d), _j,-^(d)\≥ 12 ( 1e-A_j^(d)| _j^(d)| )_+. (30) For a fixed soft router, let qe:=[me(D)(X)]q_e:=E[m_e^(D)(X)] (31) be the expected routing mass of leaf e. If the density condition in Corollary 4.5 holds at every node on the path to e, define ρj(d):=12(1e−Aj(d)|δj(d)|)+, _j^(d):= 12 ( 1e-A_j^(d)| _j^(d)| )_+, (32) then repeated conditioning along the path to e gives qe≥q¯e:=∏(d,j)∈path(e)ρj(d).q_e≥ q_e:= _(d,j) (e) _j^(d). (33) In the linear log-concave case with exact population anchors, q¯e≥(2e)−Dfor every leaf e. q_e≥(2e)^-D every leaf e. (34) These bounds do not, by themselves, imply that every leaf is retained by deterministic top-K truncation; this distinction is addressed next. 4.3 Expert optimization under sparse top-K routing Freeze the router, the top-K rule, and any capacity policy. Let Ie,t:=expert e executes example t,pe:=ℙ(Ie,t=1).I_e,t:=1\expert $e$ executes example $t$\, p_e:=P(I_e,t=1). (35) Thus pep_e is the actual execution probability after top-K selection, not the expected leaf mass qeq_e in Equation 31. Let θe _e collect the trainable parameters of expert FeF_e, and let ℒeL_e denote its conditional risk on the examples it executes. Assumption 4.6 (Frozen-router expert SGD). The examples are i.i.d. and ℙ(Ie,t=1∣ℋt)=peP(I_e,t=1 _t)=p_e, where ℋtH_t is the history before example t, and θe,0 _e,0 is deterministic. The conditional risk ℒeL_e is μe _e-strongly convex and LeL_e-smooth with minimizer θe⋆ _e . Conditional on Ie,t=1I_e,t=1, the stochastic gradient satisfies [ge,t∣ℋt,Ie,t=1] [g_e,t _t,I_e,t=1] =∇ℒe(θe,t), = _e( _e,t), [∥ge,t−∇ℒe(θe,t)∥2∣ℋt,Ie,t=1] [\|g_e,t- _e( _e,t)\|^2 _t,I_e,t=1 ] ≤σe2. ≤ _e^2. (36) The global-example update is θe,t+1=θe,t−γeIe,tge,t. _e,t+1= _e,t- _eI_e,tg_e,t. (37) Theorem 4.7 (Expert progress under intermittent execution). Under Assumption 4.6, if pe>0p_e>0 and 0<γe≤1/Le0< _e≤ 1/L_e, then [‖θe,T−θe⋆‖2]≤ [\| _e,T- _e \|^2]≤ (1−peγeμe)T‖θe,0−θe⋆‖2+γeσe2μe (1-p_e _e _e)^T\| _e,0- _e \|^2+ _e _e^2 _e ≤ ≤ exp(−peγeμeT)‖θe,0−θe⋆‖2+γeσe2μe. (-p_e _e _eT)\| _e,0- _e \|^2+ _e _e^2 _e. (38) Consequently, if pe≥pmin>0p_e≥ p_ >0 for all experts, then, for 0<ε<‖θe,0−θe⋆‖20< <\| _e,0- _e \|^2, the transient is at most ε once T≥1pminγeμelog(‖θe,0−θe⋆‖2ε).T≥ 1p_ _e _e \! ( \| _e,0- _e \|^2 ). (39) Implication for BRANCH-MoE. Theorem 4.7 identifies the realized execution frequency as the optimization bottleneck for sparsely activated experts: a larger minimum frequency pminp_ 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, thereby discouraging persistent high-level routing imbalances that would simultaneously suppress every expert in an entire subtree. Although the soft-mass bound in Equation 33 does not formally imply a positive top-K execution probability, it removes an important mechanism of expert starvation and should make realized execution frequencies larger and more stable than under unconstrained learned routing. Thus, the theorem explains how the improved utilization encouraged by BRANCH-MoE can translate into faster optimization of the least frequently selected experts. Uniform random routing remains the natural benchmark for marginal balance; BRANCH-MoE seeks comparable utilization while retaining task-dependent specialization and tree-aligned communication locality. 4.4 Communication locality from learned tree routing We finally analyze communication at a hardware boundary whose domains are aligned with depth-r routing subtrees; a token is transmitted once to every contacted domain and may then be replicated locally. The argument depends only on the centered node scores, so it applies to both linear and nonlinear node maps. Fix a token x, follow the more probable child for r levels, and let u0,u1,…,uru_0,u_1,…,u_r be the resulting path, where ui=(i,ji)u_i=(i,j_i) and u0u_0 is the root. Let ℒ(ur)L(u_r) be the leaves below uru_r, assumed to contain at least K experts, and let K(x)S_K(x) be the top-K leaf set. Under prefix-aligned placement, let Nr(x)N_r(x) denote the number of depth-r communication domains containing at least one selected expert, and suppose a token representation contains stoks_ tok bytes and transmission across this boundary costs crc_r per byte. The learned score margin at uiu_i, the leaf mass outside the prefix subtree, the smallest retained probability of K(x)S_K(x), and the communication cost are γi(x) _i(x) :=|zji,t(i)(x)|=|fji,t(i)(x)−aji,t(i)|, := |z_j_i,t^(i)(x) |= |f_j_i,t^(i)(x)-a_j_i,t^(i) |, (40) εr(x) _r(x) :=∑e∉ℒ(ur)me,t(D)(x), := _e (u_r)m_e,t^(D)(x), (41) τK(x) _K(x) :=mine∈K(x)me,t(D)(x), := _e _K(x)m_e,t^(D)(x), (42) Commr(x) _r(x) :=stokcrNr(x). :=s_ tokc_rN_r(x). (43) At uiu_i, the preferred child receives probability σ(γi(x))σ( _i(x)), while the other child receives σ(−γi(x))≤e−γi(x)σ(- _i(x))≤ e^- _i(x). Theorem 4.8 (Learned-margin communication locality). For every token x, εr(x)=1−∏i=0r−1σ(γi(x))≤∑i=0r−1e−γi(x). _r(x)=1- _i=0^r-1σ( _i(x))≤ _i=0^r-1e^- _i(x). (44) Moreover, |K(x)∖ℒ(ur)|≤⌊εr(x)τK(x)⌋, |S_K(x) (u_r) |≤ _r(x) _K(x) , (45) and Nr(x)≤minK,1+⌊εr(x)τK(x)⌋.N_r(x)≤ \K,1+ _r(x) _K(x) \. (46) In particular, if ∑i=0r−1e−γi(x)<τK(x), _i=0^r-1e^- _i(x)< _K(x), (47) then all top-K experts lie in ℒ(ur)L(u_r) and Nr(x)=1,Commr(x)=stokcr.N_r(x)=1, _r(x)=s_ tokc_r. (48) If the first r score margins satisfy γi(x)≥γ _i(x)≥γ, then εr(x)≤re−γ. _r(x)≤ re^-γ. (49) Thus γ>log(r/τK(x))γ> (r/ _K(x)) is sufficient for one-domain dispatch at the depth-r boundary. Interpretation. An upper-level decision acts on an entire subtree: a large learned margin suppresses the total probability of every leaf in the nonpreferred branch, so after r confident decisions the routing mass is concentrated in a subtree of E/2rE/2^r experts; if that subtree is placed within one hardware domain, the lower levels refine the expert choice locally. This per-token result is compatible with the population-level anti-collapse result of Section 4.2: arrival-weighted anchoring prevents the token population as a whole from placing arbitrarily little routing mass on either child of an internal node, but does not require each token to split its mass evenly. A token may assign most of its probability to one branch while other tokens favor the opposite branch, so BRANCH-MoE can maintain broad expert utilization across tokens while producing concentrated, structured assignments for each token. 5 Experiments We evaluate BRANCH-MoE against representative MoE routers on four benchmarks, testing (i) whether it is competitive on task quality without an auxiliary balancing loss and (i) whether it induces the expert locality analyzed in Section 4.4. In short, quality is at parity with the strongest baselines on all four benchmarks, locality appears on the three UCI benchmarks but not on Criteo, and load stays balanced without an auxiliary loss. 5.1 Setup We use Criteo (≈41.5≈41.5M examples, 3939 fields), Forest Covertype (581581K, 5454 features), HIGGS (600600K, 2828), and YearPredictionMSD (515515K, 9090). All MoE conditions share one block of E=16E=16 SwiGLU experts (top-44, capacity factor 1.251.25) between a minimal input stage and a linear head, and differ only in the gate. The baselines are softmax top-K with an auxiliary loss (Switch), dynamic-bias routing (DeepSeek-V3), logit-normalized softmax with an auxiliary loss (Skywork), and a hash-based control (Hash). Node maps are linear or two-layer MLPs of width h∈2,4,8h∈\2,4,8\. Every condition has five seeds; comparisons are paired by seed on UCI and unpaired (Welch) on Criteo, where TPU execution is not bitwise reproducible. We report the normalized tree distance L/DL/D between co-selected experts (L=bit_length(i⊕j)L=bit\_length(i j), D=log2ED= _2E; random pairs score 0.8170.817, lower is more local) and the load CV (std/mean of per-expert load). Details: Appendix B. Table 2: Main results (E=16E=16, top-44, 55 seeds; mean ± 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.8170.817; bold: most local on UCI) / load CV of the same arms. †Not read as locality (Section C.1). Per-map results: Appendix C. Method Criteo (logloss) Covertype (CE) HIGGS (logloss) MSD (MSE) (a) Best held-out loss ↓ BRANCH-MoE (best map) 0.445960.44596± 0.00015 h4 0.23720.2372± 0.0013 h8 0.52560.5256± 0.0014 h8 0.67670.6767± 0.0027 h8 BRANCH-MoE (linear) 0.446150.44615± 0.00025 0.25450.2545± 0.0041 0.52590.5259± 0.0008 0.68430.6843± 0.0035 Skywork 0.446030.44603± 0.00030 lin 0.24060.2406± 0.0036 h8 0.52470.5247± 0.0005 lin 0.67570.6757± 0.0015 h2 DeepSeek-V3 0.446050.44605± 0.00005 h8 0.24210.2421± 0.0033 h8 0.52640.5264± 0.0013 lin 0.67460.6746± 0.0043 h8 Switch 0.446230.44623± 0.00019 lin 0.24060.2406± 0.0028 h8 0.52760.5276± 0.0003 h4 0.67660.6766± 0.0038 h8 Hash 0.447070.44707± 0.00028 0.44650.4465± 0.0011 0.55740.5574± 0.0007 0.68990.6899± 0.0024 (b) Tree distance ↓ / load CV ↓ of the same arms BRANCH-MoE (best map) 0.7570.757 / 0.2680.268 0.7420.742 / 0.1120.112 0.7360.736 / 0.2310.231 0.7490.749 / 0.2190.219 BRANCH-MoE (linear) 0.8130.813 / 0.1150.115 0.7180.718 / 0.0930.093 0.7100.710 / 0.2270.227 0.7180.718 / 0.1770.177 Skywork 0.8180.818 / 0.0040.004 0.8170.817 / 0.1370.137 0.8120.812 / 0.1830.183 0.8130.813 / 0.5360.536 DeepSeek-V3 0.8190.819 / 0.0040.004 0.8130.813 / 0.0010.001 0.8160.816 / 0.0020.002 0.8130.813 / 0.0030.003 Switch 0.8210.821 / 0.0030.003 0.8190.819 / 0.1410.141 0.8140.814 / 0.1040.104 0.8000.800 / 0.2780.278 Hash 0.8170.817 / 0.0020.002 0.8170.817 / 0.0030.003 0.8170.817 / 0.0020.002 0.8170.817 / 0.0040.004 5.2 Results Quality. BRANCH-MoE is at parity with the strongest baselines (Table 2a). On Criteo the learned routers span 1.9×10−41.9× 10^-4 logloss with the linear map, about half the smallest resolvable difference; at its best map (h=4h=4) BRANCH-MoE has the lowest mean, but its leads over Skywork (p=0.63p=0.63) and DeepSeek-V3 (p=0.26p=0.26) are within seed noise, and DeepSeek-V3 has the higher AUC. On Covertype BRANCH-MoE (h=8h=8) has the lowest mean cross-entropy (0.23720.2372); its 0.00340.0034 margin over Switch is resolved only by the paired test (t4=−2.86t_4=-2.86, critical value 2.7762.776), the equal margin over Skywork by neither test, and the 0.00490.0049 margin over DeepSeek-V3 by the paired test (t4=−3.25t_4=-3.25). BRANCH-MoE ranks second on HIGGS (0.00090.0009 behind Skywork, within seed spread) and fourth on MSD, where the four learned routers span 0.00210.0021 against seed sds of 0.00150.0015–0.00430.0043. Hash is the worst MoE arm everywhere. Locality. On the three UCI benchmarks BRANCH-MoE is the only router materially below random (Table 2b). With the linear map its tree distance is 0.7100.710–0.7180.718, 1212–13%13\% below random (1818–23%23\% on index distance), and 0.7360.736–0.7490.749 at h=8h=8; every flat router stays within 3.5%3.5\% of random (closest: Switch on MSD, 0.8000.800). On Criteo BRANCH-MoE shows no locality: with the linear map its tree distance (0.8130.813) is indistinguishable from random (p=0.48p=0.48), although an untrained tree scores 0.7250.725. The lower values with MLP maps (†) also appear for the flat Skywork router (0.750.75–0.770.77), so we do not read them as locality. Criteo differs from UCI in input representation, batch size, and training length; we have not isolated which matters (Section C.1). Balance. BRANCH-MoE uses no auxiliary loss or bias controller. On UCI its linear-map load CV (0.090.09–0.230.23) is comparable to that of the auxiliary-loss routers at their best maps (Switch 0.100.10–0.280.28, Skywork 0.140.14–0.540.54); on Covertype the only dead experts are Skywork’s at h=2h=2 (66 of 8080 slots). On Criteo its CV of 0.1150.115 exceeds the flat routers’ (≤0.006≤ 0.006) but costs nothing measurable: it drops 0.08%0.08\% of dispatches. linh2h4h80.240.240.260.260.280.280.30.3Covertype CE(a) loss vs. node mapBRANCH-MoESkyworkSwitchDeepSeek-V3linh8h2h40.650.650.70.70.750.750.80.8BRANCH-MoE tree dist.(b) locality vs. node mapCovertypeHIGGSMSD Figure 2: Node-map effect. (a) Covertype loss of the learned routers vs. node map, ± (Hash, at 0.4470.447, is off-scale). (b) BRANCH-MoE’s tree distance vs. node map on UCI; dashed: random pairs (0.8170.817); shaded: flat routers at their best maps. Node maps. The node map can matter more than the routing rule (Figure 2a). On Covertype, replacing the linear map by an h=8h=8 MLP lowers cross-entropy by 0.0170.017–0.0280.028 for every learned router (paired t from −7.99-7.99 to −21.44-21.44), 3.53.5–5.8×5.8× the largest gap between routers at h=8h=8 (0.00490.0049; 0.0160.016 at linear), while Hash, which ignores the node map, moves by only 0.000130.00013. The sign depends on the dataset: MSD agrees; on HIGGS the linear map is best for Skywork and DeepSeek-V3 and ties for BRANCH-MoE; on Criteo no router’s loss changes significantly (≤4.1×10−4≤ 4.1× 10^-4, p>0.1p>0.1) except Switch’s, which fails with every MLP map (≈68%≈68\% of dispatches dropped). The MLP also loosens BRANCH-MoE’s structure: on UCI its tree distance rises by 0.0240.024–0.0310.031 (Figure 2b; same sign on all three datasets, not resolved on Covertype alone), and on Criteo its load CV rises from 0.1150.115 to 0.210.21–0.380.38 (p≤0.01p≤ 0.01 at h=4h=4 and 88). Takeaway. Balance without an auxiliary loss and locality on UCI are supported, and quality is at parity; locality on Criteo and whether MLP maps erode locality remain open (Appendix C). Conclusion and limitations. BRANCH-MoE combines arrival-weighted score anchoring with hierarchical routing to encourage balanced population-level utilization and structured per-token expert selection without an auxiliary balancing loss. Our analysis characterizes anchor tracking, establishes soft-mass anti-collapse for linear routing under log-concave arrivals, relates realized expert-selection frequency to optimization progress, and bounds communication locality under prefix-aligned placement. The experiments demonstrate competitive task quality and structured locality across several benchmarks. The present study is limited to a single BRANCH-MoE layer and therefore does not capture interactions among routing decisions across multiple layers. Future work will extend both the theory and empirical evaluation to multilayer architectures and develop fused, topology-aware implementations for large-scale distributed settings. AI use statement In this work, we used generative AI tools to improve the manuscript’s clarity, presentation and writing. We have not used generative AI tools for conducting experiments. All AI-assisted material was carefully reviewed and verified by the authors. We accept full responsibility for the final content of this work, including all text, claims, and artifacts produced with the assistance of generative AI. References Baldi et al. (2014) P. Baldi, P. Sadowski, and D. Whiteson Searching for exotic particles in high-energy physics with deep learning. Nature Communications 5, p. 4308. External Links: Document Cited by: Appendix D, item 3. Belcak and Wattenhofer (2023) P. Belcak and R. Wattenhofer Fast feedforward networks. arXiv preprint arXiv:2308.14711. Cited by: §2. Bertin-Mahieux et al. (2011) T. Bertin-Mahieux, D. P. W. Ellis, B. Whitman, and P. Lamere The million song dataset. In Proceedings of the 12th International Society for Music Information Retrieval Conference, p. 591–596. Cited by: Appendix D, item 3. Bertin-Mahieux (2011) T. Bertin-Mahieux Year prediction msd. UCI Machine Learning Repository. Note: Dataset External Links: Document, Link Cited by: Appendix D, item 3. Blackard and Dean (1999) J. A. Blackard and D. J. Dean Comparative accuracies of artificial neural networks and discriminant analysis in predicting forest cover types from cartographic variables. Computers and Electronics in Agriculture 24 (3), p. 131–151. External Links: Document Cited by: Appendix D, item 3. Blackard (1998) J. Blackard Covertype. UCI Machine Learning Repository. Note: Dataset External Links: Document, Link Cited by: Appendix D, item 3. Chen et al. (2022) C. Chen, M. Li, Z. Wu, D. Yu, and C. Yang TA-MoE: topology-aware large scale mixture-of-expert training. In Advances in Neural Information Processing Systems, Vol. 35, p. 22173–22186. Cited by: §2. Clark et al. (2022) A. Clark, D. de las Casas, A. Guy, A. Mensch, M. Paganini, J. Hoffmann, B. Damoc, B. Hechtman, T. Cai, S. Borgeaud, G. van den Driessche, E. Rutherford, T. Hennigan, M. J. Johnson, A. Cassirer, C. Jones, E. Buchatskaya, D. Budden, L. Sifre, S. Osindero, O. Vinyals, M. Ranzato, J. Rae, E. Elsen, K. Kavukcuoglu, and K. Simonyan Unified scaling laws for routed language models. In Proceedings of the 39th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 162, p. 4057–4086. Cited by: §2. Criteo AI Lab (2015) Criteo AI Lab Criteo 1tb click logs dataset. Note: https://ailab.criteo.com/download-criteo-1tb-click-logs-dataset/Accessed: 2026-10-04 Cited by: Appendix D, item 3. DeepSeek-AI (2024) DeepSeek-AI DeepSeek-V3 technical report. arXiv preprint arXiv:2412.19437. Cited by: item 2, §1, §2. Eigen et al. (2013) D. Eigen, M. Ranzato, and I. Sutskever Learning factored representations in a deep mixture of experts. arXiv preprint arXiv:1312.4314. Cited by: §1. Fedus et al. (2022) W. Fedus, B. Zoph, and N. Shazeer Switch Transformers: scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research 23 (120), p. 1–39. Cited by: item 1, §1, §1, §2. Frosst and Hinton (2017) N. Frosst and G. Hinton Distilling a neural network into a soft decision tree. arXiv preprint arXiv:1711.09784. Cited by: §2. He et al. (2022) J. He, J. Zhai, T. Antunes, H. Wang, F. Luo, S. Shi, and Q. Li FasterMoE: modeling and optimizing training of large-scale dynamic pre-trained models. In Proceedings of the 27th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, p. 120–134. Cited by: §2. Hwang et al. (2023) C. Hwang, W. Cui, Y. Xiong, Z. Yang, Z. Liu, H. Hu, Z. Wang, R. Salas, J. Jose, P. Ram, H. Chau, P. Cheng, F. Yang, M. Yang, and Y. Xiong Tutel: adaptive mixture-of-experts at scale. In Proceedings of Machine Learning and Systems, Vol. 5, p. 269–287. Cited by: §2. Ibrahim et al. (2023) S. Ibrahim, W. Chen, H. Hazimeh, N. Ponomareva, Z. Zhao, and R. Mazumder COMET: learning cardinality constrained mixture of experts with trees and local search. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, p. 832–844. Cited by: §2. İrsoy et al. (2012) O. İrsoy, O. T. Yıldız, and E. Alpaydın Soft decision trees. In Proceedings of the 21st International Conference on Pattern Recognition, p. 1819–1822. Cited by: §2. Jacobs et al. (1991) R. A. Jacobs, M. I. Jordan, S. J. Nowlan, and G. E. Hinton Adaptive mixtures of local experts. Neural Computation 3 (1), p. 79–87. Cited by: §1. Jordan and Jacobs (1994) M. I. Jordan and R. A. Jacobs Hierarchical mixtures of experts and the EM algorithm. Neural Computation 6 (2), p. 181–214. Cited by: §2. Lepikhin et al. (2021) D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen GShard: scaling giant models with conditional computation and automatic sharding. In International Conference on Learning Representations, Cited by: §1. Lewis et al. (2021) M. Lewis, S. Bhosale, T. Dettmers, N. Goyal, and L. Zettlemoyer BASE layers: simplifying training of large, sparse models. In Proceedings of the 38th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 139, p. 6265–6274. Cited by: §1, §2. Rajbhandari et al. (2022) S. Rajbhandari, C. Li, Z. Yao, M. Zhang, R. Y. Aminabadi, A. A. Awan, J. Rasley, and Y. He DeepSpeed-MoE: advancing mixture-of-experts inference and training to power next-generation AI scale. In Proceedings of the 39th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 162, p. 18332–18346. Cited by: §2. Roller et al. (2021) S. Roller, S. Sukhbaatar, A. Szlam, and J. Weston Hash layers for large sparse models. In Advances in Neural Information Processing Systems, Vol. 34, p. 17555–17566. Cited by: item 4, §2. Shazeer et al. (2017) N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean Outrageously large neural networks: the sparsely-gated mixture-of-experts layer. In International Conference on Learning Representations, Cited by: §1, §1, §2. Wang et al. (2024) L. Wang, H. Gao, C. Zhao, X. Sun, and D. Dai Auxiliary-loss-free load balancing strategy for mixture-of-experts. arXiv preprint arXiv:2408.15664. Cited by: §1. Wei et al. (2024) T. Wei, B. Zhu, L. Zhao, C. Cheng, B. Li, W. Lü, P. Cheng, J. Zhang, X. Zhang, L. Zeng, X. Wang, Y. Ma, R. Hu, S. Yan, H. Fang, and Y. Zhou Skywork-MoE: a deep dive into training techniques for mixture-of-experts language models. arXiv preprint arXiv:2406.06563. Cited by: item 3, item 3. Whiteson (2014) D. Whiteson HIGGS. UCI Machine Learning Repository. Note: Dataset External Links: Document, Link Cited by: Appendix D, item 3. Zhou et al. (2022) Y. Zhou, T. Lei, H. Liu, N. Du, Y. Huang, V. Zhao, A. M. Dai, Z. Chen, Q. V. Le, and J. Laudon Mixture-of-experts with expert choice routing. In Advances in Neural Information Processing Systems, Vol. 35, p. 7103–7114. Cited by: §1, §2. Appendix A Proofs Proof of Theorem 4.2. Let et:=νt−ν¯e_t:= _t- ν. From the constant-decay update and Equation 23, et+1=αet+(1−α)ξt.e_t+1=α e_t+(1-α) _t. Taking conditional expectation given ℱtF_t gives [et+1∣ℱt]=αetE[e_t+1 _t]=α e_t. Iterating proves Equation 24. For the second moment, [et+12∣ℱt] [e_t+1^2 _t] =α2et2+2α(1−α)et[ξt∣ℱt]+(1−α)2[ξt2∣ℱt] =α^2e_t^2+2α(1-α)e_tE[ _t _t]+(1-α)^2E[ _t^2 _t] ≤α2et2+(1−α)2s2. ≤α^2e_t^2+(1-α)^2s^2. Taking expectations and unrolling the recursion yields [et2]≤α2te02+(1−α)2s2∑k=0t−1α2k.E[e_t^2]≤α^2te_0^2+(1-α)^2s^2 _k=0^t-1α^2k. Using ∑k=0t−1α2k=(1−α2t)/(1−α2) _k=0^t-1α^2k=(1-α^2t)/(1-α^2) gives Equation 25. ∎ Proof of Proposition 4.3. Write ξt:=ν^t−ν¯t _t:= ν_t- ν_t and et:=νt−ν¯te_t:= _t- ν_t. Then et+1=αet+(ν¯t−ν¯t+1)+(1−α)ξt.e_t+1=α e_t+( ν_t- ν_t+1)+(1-α) _t. Unrolling gives et= e_t= αte0+∑k=0t−1αt−1−k(ν¯k−ν¯k+1) α^te_0+ _k=0^t-1α^t-1-k( ν_k- ν_k+1) +(1−α)∑k=0t−1αt−1−kξk. +(1-α) _k=0^t-1α^t-1-k _k. The absolute value of the deterministic first line is at most αt|e0|+v∑k=0t−1αt−1−k=αt|e0|+v(1−αt)1−α.α^t|e_0|+v _k=0^t-1α^t-1-k=α^t|e_0|+ v(1-α^t)1-α. The martingale-difference property makes distinct noise terms orthogonal and eliminates their cross term with the deterministic part. Hence the noise variance is at most (1−α)2s2∑k=0t−1α2k=1−α1+α(1−α2t)s2.(1-α)^2s^2 _k=0^t-1α^2k= 1-α1+α(1-α^2t)s^2. Combining the two bounds proves Equation 27. ∎ Proof of Lemma 4.4. Along the path to node (d,j)(d,j), its mask can be written as mj(d)(x)=∏u∈path((,,,))σ(suzu(x)),m_j^(d)(x)= _u ((d,j))σ(s_uz_u(x)), where su∈−1,+1s_u∈\-1,+1\ selects the appropriate branch and zu(x)z_u(x) is an affine centered score. The function y↦logσ(y)=−log(1+exp(−y))y σ(y)=- (1+ (-y)) is concave. Therefore each factor x↦σ(suzu(x))x σ(s_uz_u(x)) is log-concave, and so is their product. The unnormalized arrival density is the product of this mask and the log-concave input density, hence is log-concave. Normalization preserves log-concavity. ∎ Proof of Corollary 4.5. Let μj(d):=j(d)[X] _j^(d):=E_D_j^(d)[X]. Linearity gives νj(d),⋆=(Wj(d))⊤μj(d), _j^(d), = (W_j^(d) ) _j^(d), so the population-centered score is Z=(Wj(d))⊤(X−μj(d)).Z= (W_j^(d) ) (X- _j^(d)). Grünbaum’s inequality for log-concave measures implies ℙ(Z≥0)≥1e,ℙ(Z≤0)≥1e.P(Z≥ 0)≥ 1e, (Z≤ 0)≥ 1e. At the exact population anchor, πj,+(d)=[σ(Z)] _j,+^(d)=E[σ(Z)] and πj,−(d)=[σ(−Z)] _j,-^(d)=E[σ(-Z)], where both expectations are under j(d)D_j^(d). Since σ(y)≥1/2σ(y)≥ 1/2 for y≥0y≥ 0, each expectation is at least 1/(2e)1/(2e), proving Equation 29. For an arbitrary active anchor, the two expectations become [σ(Z−δj(d))]E[σ(Z- _j^(d))] and [σ(δj(d)−Z)]E[σ( _j^(d)-Z)]. They are bounded below by one half of ℙ(Z≥δj(d))P(Z≥ _j^(d)) and ℙ(Z≤δj(d))P(Z≤ _j^(d)), respectively. If the density of Z is bounded by Aj(d)A_j^(d), moving the threshold from zero to δj(d) _j^(d) changes either tail probability by at most Aj(d)|δj(d)|A_j^(d)| _j^(d)|. Combining this fact with Grünbaum’s bounds and truncating the resulting lower bound at zero proves Equation 30. ∎ Proof of Theorem 4.7. Write ut=θe,t−θe⋆u_t= _e,t- _e . Conditional on ℋtH_t, unbiasedness and Equation 37 imply [‖ut+1‖2∣ℋt]= [\|u_t+1\|^2 _t]= ‖ut‖2−2peγe⟨ut,∇ℒe(θe,t)⟩ \;\|u_t\|^2-2p_e _e u_t, _e( _e,t) +peγe2[∥ge,t∥2∣ℋt,Ie,t=1]. +p_e _e^2E[\|g_e,t\|^2 _t,I_e,t=1]. Strong convexity gives ⟨ut,∇ℒe(θe,t)⟩≥μe‖ut‖2 u_t, _e( _e,t) ≥ _e\|u_t\|^2. Smoothness, optimality of θe⋆ _e , and Equation 36 give [∥ge,t∥2∣ℋt,Ie,t=1]≤∥∇ℒe(θe,t)∥2+σe2.E[\|g_e,t\|^2 _t,I_e,t=1]≤\| _e( _e,t)\|^2+ _e^2. Because a differentiable convex LeL_e-smooth function has an Le−1L_e^-1-cocoercive gradient, ‖∇ℒe(θe,t)‖2≤Le⟨ut,∇ℒe(θe,t)⟩\| _e( _e,t)\|^2≤ L_e u_t, _e( _e,t) . Therefore, using γe≤1/Le _e≤ 1/L_e, [‖ut+1‖2∣ℋt]≤(1−peγeμe)‖ut‖2+peγe2σe2.E[\|u_t+1\|^2 _t]≤(1-p_e _e _e)\|u_t\|^2+p_e _e^2 _e^2. Iteration of this recursion and summation of the geometric series yield ‖uT‖2≤(1−peγeμe)T‖u0‖2+γeσe2μe.E\|u_T\|^2≤(1-p_e _e _e)^T\|u_0\|^2+ _e _e^2 _e. Finally, 1−x≤exp(−x)1-x≤ (-x) gives the second line of Equation 38. If pe≥pminp_e≥ p_ , replacing pep_e by pminp_ in the exponential transient and solving for T gives Equation 39. ∎ Proof of Theorem 4.8. The leaf masks below ur=(r,jr)u_r=(r,j_r) sum to its incoming mask: ∑e∈ℒ(ur)me,t(D)(x)=mjr,t(r)(x). _e (u_r)m_e,t^(D)(x)=m_j_r,t^(r)(x). (50) This follows by induction from the cascade, because the two children of a node with incoming mask m receive masses m(1−p)m(1-p) and mpmp, whose sum is m. Along the preferred path, the probability assigned at level i is σ(γi(x))σ( _i(x)). Therefore, ∑e∈ℒ(ur)me,t(D)(x)=mjr,t(r)(x)=∏i=0r−1σ(γi(x)). _e (u_r)m_e,t^(D)(x)=m_j_r,t^(r)(x)= _i=0^r-1σ( _i(x)). Since the full leaf mask sums to one, εr(x)=1−∏i=0r−1σ(γi(x)). _r(x)=1- _i=0^r-1σ( _i(x)). For numbers ai∈[0,1]a_i∈[0,1], 1−∏i(1−ai)≤∑iai1- _i(1-a_i)≤ _ia_i. Taking ai=σ(−γi(x))a_i=σ(- _i(x)) gives εr(x)≤∑i=0r−1σ(−γi(x))≤∑i=0r−1e−γi(x), _r(x)≤ _i=0^r-1σ(- _i(x))≤ _i=0^r-1e^- _i(x), which proves Equation 44. Let hr(x):=|K(x)∖ℒ(ur)|.h_r(x):=|S_K(x) (u_r)|. Every selected leaf has probability at least τK(x) _K(x), and hence εr(x)≥∑e∈K(x)∖ℒ(ur)me,t(D)(x)≥hr(x)τK(x). _r(x)≥ _e _K(x) (u_r)m_e,t^(D)(x)≥ h_r(x) _K(x). Because hr(x)h_r(x) is integer, this proves Equation 45. All selected experts inside ℒ(ur)L(u_r) occupy one domain, and every selected expert outside it can add at most one further domain. Thus Nr(x)≤1+hr(x)N_r(x)≤ 1+h_r(x), while trivially Nr(x)≤KN_r(x)≤ K. This proves Equation 46. Finally, Equation 47 implies εr(x)<τK(x) _r(x)< _K(x). If a selected expert lay outside ℒ(ur)L(u_r), its probability would be at least τK(x) _K(x), which would contradict the definition of εr(x) _r(x). Hence every selected expert lies in the same depth-r domain. ∎ Appendix B Experimental Details B.1 Datasets Table 3 summarizes the four benchmarks. Criteo is the primary large-scale task; the three UCI datasets are tabular benchmarks that admit many-seed replication at modest cost. Table 3: Benchmark suite. All MoE conditions use E=16E=16 experts with top-44 routing. “Examples” is the full corpus size (train ++ held-out); “Seeds” is the number of independent replications per condition. Dataset Task Examples Features Seeds Metric Criteo (small) CTR (binary) ≈41.5≈41.5M22 2 The held-out split is exactly 4.424.42M examples (one full evaluation pass of 2,1582,158 steps at batch 20482048). The training split is not counted directly: we infer ≈37≈37M from the ratio of on-disk sizes (5.935.93 GB train against 0.710.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. 3939 fields 55 Logloss / AUC Forest Covertype 77-class 581581K 5454 55 Cross-entropy / Acc. HIGGS Binary 600600K 2828 55 Logloss / AUC YearPredictionMSD Regression 515515K 9090 55 MSE / RMSE (yr) Criteo. Each example comprises 1313 integer (count) features and 2626 categorical features, for 3939 fields in total, with a binary click label. Categorical fields are mapped to integer IDs by a string lookup built from a vocabulary with a minimum-frequency cutoff of 100100; post-cutoff vocabulary sizes V range from 44 to 20,13620,136. Integer fields are discretized into 4141 bins by the boundaries bi=1.5i−1−0.49b_i=1.5^\,i-1-0.49 for i=0,…,39i=0,…,39, converting magnitude into an ordered categorical ID. Embeddings are not shared across fields: integer fields use width 1010, categorical fields use df=min(⌊6Vf1/4⌋,Vf)∈[4,71]d_f= ( 6V_f^1/4 ,V_f)∈[4,71], and the concatenation yields a 11371137-dimensional input. UCI benchmarks. Forest Covertype (5454 cartographic features, 77 cover classes, 80/2080/20 split) and HIGGS (2828 particle-physics features, the first 600600K of the 1111M-event corpus, 500500K/100100K split) are standardized. YearPredictionMSD (9090 timbre statistics) uses the canonical 463,715/51,630463,715/51,630 split, which prevents artist leakage; its features are likewise standardized, the target is the standardized release year, and we report unstandardized RMSE in years alongside the MSE objective. HIGGS and MSD inputs are additionally clipped at ±10σ± 10σ; Covertype inputs are not. On Forest Covertype one preprocessing detail proves critical: 4040 of the 5454 features are binary soil-type indicators, the rarest firing on only 22 of the 464,809464,809 training rows. Standard z-scoring with an ϵ=10−7ε=10^-7 denominator guard maps a column that is constant in the sample used to fit the normalizer to spikes of magnitude ∼107 10^7 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 10−410^-4 to unit variance eliminates the collapse and reduces the seed-to-seed accuracy spread of BRANCH-MoE from 6.26.2 to 0.460.46 percentage points. B.2 Baselines We benchmark against four representative MoE routing paradigms, covering auxiliary-loss, auxiliary-loss-free, normalization-based, and deterministic mechanisms. 1. Top-K Softmax Routing (Switch). The canonical learned paradigm (Fedus et al., 2022): a learned gating matrix, a softmax, and top-K dispatch, with multiplicative uniform jitter [1−ϵ,1+ϵ][1-ε,1+ε], ϵ=10−2ε=10^-2, during training. Load balance is enforced by the auxiliary loss ℒbalance=αE∑efePeL_balance=α E _ef_eP_e with α=0.01α=0.01. 2. Auxiliary-Loss-Free Routing (DeepSeek-V3). Affinities are computed with an independent sigmoid per expert, so experts do not compete for probability mass (DeepSeek-AI, 2024). A non-trainable additive bias beb_e enters top-K selection only and is updated after each step by a fixed increment γ=10−3γ=10^-3 in the direction that relieves congestion; the combining gates remain the unbiased affinities. 3. Gating Logit Normalization (Skywork). Logits are standardized across the expert dimension before the softmax, which counteracts representation collapse and sharpens the gate distribution (Wei et al., 2024). A fixed auxiliary balancing coefficient is retained. We include this baseline because it proves to be the strongest competitor on two of the four benchmarks and ties Switch as the runner-up on a third. 4. Deterministic Balanced Routing (Hash). A fixed assignment mapping each example to experts by a 3232-bit avalanche hash of its feature vector, following Hash Layers (Roller et al., 2021). The assignment is uniform by construction but carries no dependence on content beyond the hash, isolating the value of learned gating against a perfectly balanced but semantically uninformed one, at zero router parameter cost. 5. BRANCH-MoE (ours). The E experts are the leaves of a balanced binary tree of depth log2E _2E. At internal node n the probability of branching right is σ(fn()−μn)σ(f_n(x)- _n). During training μn _n is the arrival-weighted mean score of the current batch at n; at inference it is νn _n, a non-trainable exponential moving average (decay 0.990.99) of that batch mean, updated at every training step. Centering each split on the mean score of its own traffic places the decision boundary through the middle of the distribution actually arriving there, keeping both children populated without a balancing penalty. Leaf scores are products of branch probabilities along each root-to-leaf path; all E−1E-1 node scores are evaluated, and the products are formed in log2E _2E vectorized steps. B.3 Node Maps, Architecture, and Protocol Node maps. The score fnf_n at each node, and correspondingly the logit map of each flat router, is one of two forms, swept as a factor in its own right: • Linear: fn()=n⊤f_n(x)=w_n x. • MLP: a two-layer network of hidden width h∈2,4,8h∈\2,4,8\ per node, fn()=n⊤ReLU(n+n)f_n(x)=v_n ReLU(U_nx+c_n). This factor is swept across all four learned routers and the Hash control on the three UCI benchmarks, and across all four learned routers on Criteo. Architecture and protocol. To isolate routing, we adopt a minimal backbone: an input stage, a single MoE block, and a linear head. On Criteo, the input stage projects the 3939 embeddings to 256256 dimensions, followed by SiLU and layer normalization; on the UCI datasets, it is a feature-standardization layer alone with no projection. The MoE block has no residual connection. Within this block, each expert is a SwiGLU network defined as down(SiLU(gate())⊙up())down(SiLU(gate(x)) (x)), where the upup layer expands the input into a higher-dimensional space and the downdown layer compresses it back to the baseline hidden size. Finally, a post-MoE layer normalization precedes the head on Criteo, HIGGS, and MSD (the Covertype model has none). All conditions use E=16E=16 and top-44 routing with capacity factor 1.251.25, and differ only in the gate g(⋅)g(·), so comparisons are at matched capacity. Criteo trains for 37,50037,500 steps at batch 20482048 with Adam (lr=5×10−4lr=5× 10^-4), organized as 1515 evaluation intervals of 2,5002,500 steps; since an interval covers 5.125.12M examples against a ≈37≈37M training split, this is ≈2.1≈2.1 passes over the data rather than 1515. The UCI datasets train for 2020 true epochs with Adam (lr=10−3lr=10^-3) at batch 512512 (Covertype) or 256256 (HIGGS, MSD). Every condition is replicated over 55 seeds. On the UCI benchmarks the seed governs initialization but not the data order, which is shuffled without a seed; all arms within a dataset share the identical seed set, so comparisons are paired by seed. On Criteo the seed likewise governs initialization only, and TPU execution is not bitwise reproducible, so Criteo comparisons are unpaired (two-sided Welch t-tests). Re-running an identical Criteo configuration with the same seed changes held-out logloss about as much as changing the seed does (standard deviation 1.51.5–2.8×10−42.8× 10^-4 within a seed against 1.91.9–3.0×10−43.0× 10^-4 across seeds). B.4 Locality Metrics For each token routed to experts (i,j)(i,j) we measure two index-based distances, averaged over all (k2) k2 co-selected pairs. 1. Tree distance L. Reading an expert index as a D-bit root-to-leaf path (D=log2ED= _2E), the number of levels to the lowest common ancestor is L(i,j)=bit_length(i⊕j)∈1,…,DL(i,j)=bit\_length(i j)∈\1,…,D\. Under block placement across P=2mP=2^m devices, two experts share a device if and only if L(i,j)≤D−mL(i,j)≤ D-m, making L an exact topological proxy for cross-device fanout. 2. Index gap |i−j||i-j|. The linear gap between expert indices. Both require a reference: for uniformly random distinct pairs, [L]=(D−1+2−D)/(1−2−D)E[L]=(D-1+2^-D)/(1-2^-D) and [|i−j|]=(E+1)/3E[|i-j|]=(E+1)/3, giving 3.26673.2667 and 5.66675.6667 at E=16E=16 (D=4D=4). Normalized by D and E−1E-1 these are 0.81670.8167 and 0.37780.3778 — random routing already scores 0.820.82 on the tree metric, so locality is only evidenced below these values, and an unreferenced number is uninterpretable. We also report the load coefficient of variation CV=std(f)/mean(f)CV=std(f)/mean(f) of the per-expert dispatch distribution. Tree distance is also the more sensitive of the two index-based measures: it separates power-of-two boundary pairs such as (7,8)(7,8), which have |i−j|=1|i-j|=1 but lie at maximal tree separation, and which the index gap therefore scores as maximally local when they are maximally remote. Appendix C Additional Results C.1 Criteo: Click-Through-Rate Prediction Table 4: Criteo, E=16E=16, top-44, capacity factor 1.251.25, 55 seeds (mean ± std). Logloss is binary cross-entropy; logloss and AUC are the best held-out values over the 1515 evaluations. Load CV and the normalized tree and index distances are training-time routing statistics; random baselines are 0.81670.8167 and 0.37780.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 55 vs. 55 seeds, differences below ≈3.7×10−4≈3.7× 10^-4 in logloss or ≈2.7×10−4≈2.7× 10^-4 in AUC are not resolvable (Welch t-test, p<0.05p<0.05). Method Best val. logloss ↓ Best val. AUC ↑ Load CV ↓ Tree ↓ Index ↓ Linear node map Skywork 0.4460340.446034 0.8058270.805827 0.00400.0040 0.81840.8184 0.37780.3778 ±0.000298± 0.000298 ±0.000184± 0.000184 DeepSeek-V3 0.4460820.446082 0.8057990.805799 0.00570.0057 0.81880.8188 0.38490.3849 ±0.000270± 0.000270 ±0.000174± 0.000174 BRANCH-MoE (ours) 0.4461490.446149 0.8057070.805707 0.11520.1152 0.81260.8126 0.38430.3843 ±0.000250± 0.000250 ±0.000175± 0.000175 Switch softmax 0.4462260.446226 0.8055950.805595 0.00250.0025 0.82120.8212 0.38120.3812 ±0.000185± 0.000185 ±0.000204± 0.000204 Hash 0.4470650.447065 0.8047160.804716 0.00230.0023 0.81670.8167 0.37780.3778 ±0.000277± 0.000277 ±0.000116± 0.000116 Best MLP node map BRANCH-MoE (h=4h=4) 0.4459590.445959 0.8058650.805865 0.26760.2676 0.75720.7572 0.32550.3255 ±0.000153± 0.000153 ±0.000022± 0.000022 DeepSeek-V3 (h=8h=8) 0.4460490.446049 0.8059700.805970 0.00380.0038 0.81900.8190 0.38240.3824 ±0.000047± 0.000047 ±0.000100± 0.000100 Table 4 reports the benchmark. With the common linear node map, the four learned routers are statistically tied: their mean held-out logloss spans 1.9×10−41.9× 10^-4, about half of the smallest resolvable difference, and BRANCH-MoE ranks third, 1.1×10−41.1× 10^-4 behind Skywork (Welch p=0.53p=0.53). When each router uses its best node map, BRANCH-MoE with h=4h=4 has the lowest mean logloss, but its margins over Skywork (7.5×10−57.5× 10^-5, p=0.63p=0.63) and DeepSeek-V3 at h=8h=8 (9.0×10−59.0× 10^-5, p=0.26p=0.26) are well inside seed noise; choosing the best of four node maps after the fact also favours it, and DeepSeek-V3 at h=8h=8 has the higher AUC (+1.05×10−4+1.05× 10^-4, p=0.08p=0.08). We therefore claim parity, not superiority. The only unambiguous gap is to Hash, which trails every other row of the table by 0.00080.0008–0.00110.0011 in logloss and 0.00090.0009–0.00130.0013 in AUC (Welch p<0.002p<0.002 in every comparison). All Criteo tests are unpaired, because TPU execution is not bitwise reproducible (Section B.3). Balance. BRANCH-MoE keeps every expert in use without an auxiliary loss or bias controller, but balances load less tightly than the flat routers: with the linear node map its load CV is 0.1150.115 (largest and smallest expert shares 1.17×1.17× and 0.72×0.72× the mean), against at most 0.0060.006 for every flat router. At capacity factor 1.251.25 this has no measurable cost: every linear-map condition drops at most 0.11%0.11\% of dispatches at capacity, and BRANCH-MoE drops 0.08%0.08\%. Locality does not appear on Criteo. With the linear node map, BRANCH-MoE’s normalized tree distance is 0.813±0.0120.813± 0.012, indistinguishable from the 0.81670.8167 random baseline (one-sample t-test, p=0.48p=0.48), and its index distance (0.3840.384) is marginally above random (p=0.04p=0.04). The tree does not start at random: an untrained BRANCH-MoE scores 0.7250.725, and training moves it to the random value. The flat routers stay within 2%2\% of both random baselines with linear node maps, as expected for arbitrary expert labels. Lower tree distances do appear with MLP node maps (0.7570.757 for BRANCH-MoE at h=4h=4), but they vary strongly across seeds, are lowest in the most imbalanced runs, and appear also for the flat Skywork router (0.750.75–0.770.77, p≤0.05p≤ 0.05 at every width), whose expert indices carry no meaning. We have not diagnosed the Skywork effect; it points to an index-order artifact in expert selection, such as tie-breaking among exactly equal logits, and we read neither result as locality. Criteo is therefore the one benchmark on which the locality of Section C.4 does not replicate. It differs from the UCI benchmarks in input representation (learned embeddings of categorical fields), batch size (20482048 against 256256–512512), and training length (≈2.1≈2.1 passes over the data rather than 2020 epochs); we have not isolated which of these is responsible. C.2 Forest Covertype: Multi-Class Classification Table 5: Forest Covertype, E=16E=16, top-44, 2020 epochs, batch 512512, 55 paired seeds (mean ± 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. Best val. cross-entropy ↓ Method linear h=2h=2 h=4h=4 h=8h=8 BRANCH-MoE (ours) 0.254510.25451 0.252060.25206 0.243940.24394 0.237200.23720 ±0.00408± 0.00408 ±0.00487± 0.00487 ±0.00356± 0.00356 ±0.00125± 0.00125 Skywork 0.258620.25862 0.251690.25169 0.242770.24277 0.240620.24062 ±0.00365± 0.00365 ±0.00324± 0.00324 ±0.00494± 0.00494 ±0.00358± 0.00358 Switch softmax 0.261500.26150 0.256460.25646 0.242860.24286 0.240620.24062 ±0.00136± 0.00136 ±0.00713± 0.00713 ±0.00180± 0.00180 ±0.00281± 0.00281 DeepSeek-V3 0.270320.27032 0.266200.26620 0.254170.25417 0.242070.24207 ±0.00225± 0.00225 ±0.00301± 0.00301 ±0.00196± 0.00196 ±0.00326± 0.00326 Hash (control) 0.446650.44665 0.447710.44771 0.447070.44707 0.446520.44652 ±0.00087± 0.00087 ±0.00084± 0.00084 ±0.00162± 0.00162 ±0.00108± 0.00108 Table 5 reports the node-map sweep. BRANCH-MoE with h=8h=8 is the best arm overall (0.237200.23720 cross-entropy, 90.21%90.21\% accuracy, best on both), but Skywork and Switch tie for second at the same h (0.240620.24062 each), and the margin of −0.00342-0.00342 over them is not clean. Against Switch a paired t-test resolves it only marginally (t4=−2.86t_4=-2.86 against a critical value of 2.7762.776) and the unpaired Welch test does not (t=−2.49t=-2.49); against Skywork neither test resolves it (paired t4=−1.64t_4=-1.64, Welch t=−2.01t=-2.01). The margin over DeepSeek-V3 at h=8h=8 is cleaner (−0.00487-0.00487, paired t4=−3.25t_4=-3.25). With linear node maps BRANCH-MoE again has the lowest mean: its gap to Switch is larger and resolved (−0.00699-0.00699, paired t4=−3.25t_4=-3.25), but its gap to Skywork, the runner-up, is not (−0.00411-0.00411, paired t4=−1.26t_4=-1.26). At h=2h=2 and h=4h=4 Skywork has the lowest mean, by margins of 0.00040.0004 and 0.00010.0001 that are far inside seed noise. The four Hash arms are four runs of one identical configuration, so the spread of their seed means isolates pure run-to-run nondeterminism: sd=0.00053sd=0.00053. The BRANCH-MoE-vs-Switch gap at h=8h=8 is well above this floor yet not resolved by the unpaired test, which is precisely why pairing is required to see it at all; the equal gap to Skywork is not resolved even with pairing. On balance BRANCH-MoE attains CV=0.093CV=0.093 with a 1.174×1.174× peak load at linear (0.1120.112 / 1.240×1.240× at h=8h=8), against 0.1210.121 / 1.210×1.210× for Switch and 0.2480.248 / 1.366×1.366× for Skywork, the least balanced learned router at linear (0.1370.137 / 1.247×1.247× at h=8h=8); DeepSeek-V3’s controller is tighter still (0.0010.001 / 1.002×1.002×), as expected from a mechanism that optimizes balance directly. The only dead experts (below 1%1\% of uniform load) are Skywork’s at h=2h=2: 66 of its 8080 expert slots across the five seeds. C.3 HIGGS and YearPredictionMSD Table 6: HIGGS (logloss) and YearPredictionMSD (MSE), E=16E=16, top-44, 55 paired seeds (mean ± std). Each router is shown at its best node map, given in parentheses; for BRANCH-MoE we show both h=8h=8 (best loss) and linear (best locality), since the two differ by 0.00030.0003 on HIGGS. Tree and index distances are normalized against the 0.81670.8167 / 0.37780.3778 random baselines. Hash ignores the node map, so its four arms are one configuration run four times; the best is shown. HIGGS YearPredictionMSD Method Logloss ↓ Tree ↓ Index ↓ MSE ↓ Tree ↓ Index ↓ Skywork 0.52470.5247 (lin) 0.8120.812 0.3700.370 0.67570.6757 (h=2h=2) 0.8130.813 0.3650.365 ±0.0005± 0.0005 ±0.0015± 0.0015 BRANCH-MoE (h=8h=8) 0.52560.5256 0.7360.736 0.3220.322 0.67670.6767 0.7490.749 0.3240.324 ±0.0014± 0.0014 ±0.0027± 0.0027 BRANCH-MoE (linear) 0.52590.5259 0.7100.710 0.2920.292 0.68430.6843 0.7180.718 0.3010.301 ±0.0008± 0.0008 ±0.0035± 0.0035 DeepSeek-V3 0.52640.5264 (lin) 0.8160.816 0.3800.380 0.67460.6746 (h=8h=8) 0.8130.813 0.3710.371 ±0.0013± 0.0013 ±0.0043± 0.0043 Switch softmax 0.52760.5276 (h=4h=4) 0.8140.814 0.3740.374 0.67660.6766 (h=8h=8) 0.8000.800 0.3650.365 ±0.0003± 0.0003 ±0.0038± 0.0038 Hash 0.55740.5574 (h=4h=4) 0.8170.817 0.3780.378 0.68990.6899 (h=4h=4) 0.8170.817 0.3780.378 ±0.0007± 0.0007 ±0.0024± 0.0024 Table 6 reports the two remaining benchmarks, and the outcome is different. On HIGGS BRANCH-MoE ranks second behind Skywork (+0.0009+0.0009, within the combined seed spread); on MSD it ranks fourth, behind DeepSeek-V3, Skywork, and Switch, though the four are separated by only 0.00210.0021 MSE — 0.0140.014 years of RMSE — against seed standard deviations of 0.00150.0015–0.00430.0043. The honest summary is that on these two datasets all learned routers are statistically tied, and BRANCH-MoE is not the best point estimate. C.4 Routing Locality Across the Suite The locality prediction of Section 4.4 is the claim that survives most cleanly. On all three UCI benchmarks BRANCH-MoE is the only router whose co-routing statistics depart materially from the random baseline (Table 7). Table 7: Normalized co-routing distances (E=16E=16, top-44; mean over 55 seeds). Flat routers: range over Switch, DeepSeek-V3, Skywork, and Hash, each at the arm shown in Table 2. Bold: most local. BRANCH-MoE (linear) BRANCH-MoE (h=8h=8) Flat routers Dataset Tree Index Tree Index Tree Index Forest Covertype 0.7180.718 0.3120.312 0.7420.742 0.3340.334 0.8130.813–0.8190.819 0.3710.371–0.3850.385 HIGGS 0.7100.710 0.2920.292 0.7360.736 0.3220.322 0.8120.812–0.8170.817 0.3700.370–0.3800.380 YearPredictionMSD 0.7180.718 0.3010.301 0.7490.749 0.3240.324 0.8000.800–0.8170.817 0.3650.365–0.3780.378 Criteo 0.8130.813 0.3840.384 0.8040.804 0.3700.370 0.8170.817–0.8210.821 0.3780.378–0.3820.382 Random baseline 0.8170.817 0.3780.378 0.8170.817 0.3780.378 0.8170.817 0.3780.378 On the three UCI datasets the reduction is 1212–13%13\% below baseline on tree distance for linear BRANCH-MoE (88–10%10\% at h=8h=8) and 1818–23%23\% on index distance (1212–15%15\% at h=8h=8), while every flat router sits within 3.5%3.5\% of baseline on both — the negative control behaving exactly as expected, since for a flat router the expert index is an arbitrary label. Criteo is the exception: there BRANCH-MoE stays at the random baseline, and with MLP node maps even the flat Skywork router falls below it, an artifact discussed in Section C.1. Appendix D Description of Datasets We evaluate the considered methods on four widely used tabular benchmarks spanning binary classification, multiclass classification, and regression. These datasets vary substantially in size, feature type, and prediction task. Criteo. The Criteo 1TB Click Logs dataset is a large-scale benchmark for click-through-rate (CTR) prediction. Each instance corresponds to a display advertisement served by Criteo, and the binary label indicates whether the ad was clicked. The training data consist of 24 chronologically ordered daily files and contain 13 numerical, predominantly count-based features and 26 anonymized categorical features. The categorical values are hashed into 32-bit integers, and some entries may be missing. We use this dataset as a large-scale binary classification task (Criteo AI Lab, 2015). Forest Covertype. The Forest Covertype dataset consists of 581,012581,012 observations of 30×3030× 30 meter forest patches from the Roosevelt National Forest in Colorado. The task is to predict the dominant forest cover type among seven classes from 54 cartographic variables, including elevation, slope, aspect, hillshade measurements, wilderness-area indicators, and soil-type indicators. It therefore provides a multiclass classification benchmark with a mix of continuous and binary features (Blackard and Dean, 1999; Blackard, 1998). HIGGS. The HIGGS dataset is a binary classification dataset derived from simulated high-energy particle-collision events. Its goal is to distinguish events corresponding to a signal process that produces Higgs bosons from background events. The dataset contains 1111 million instances, each represented by 28 real-valued features: 21 low-level kinematic variables measured by particle detectors and seven high-level features derived from the low-level variables. This dataset is a standard large-scale benchmark for binary classification (Baldi et al., 2014; Whiteson, 2014). YearPredictionMSD. The YearPredictionMSD dataset is a regression benchmark constructed from the Million Song Dataset. The task is to predict a song’s release year from 90 audio features that summarize timbre characteristics. It contains 515,345515,345 songs, primarily Western commercial recordings released between 1922 and 2011. We use it to assess performance on a large-scale continuous prediction problem (Bertin-Mahieux et al., 2011; Bertin-Mahieux, 2011).