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

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 ${\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. 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.710\)–\(0.718\) between co-selected experts, compared with \(0.817\) for random expert pairs—a reduction of \(12\)–\(13%\). Every flat-routing baseline remains within \(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.