Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
BRANCH-MoE introduces a balance-aware tree routing mechanism for Mixture-of-Experts (MoE) models, where experts are organized in a binary decision tree. By estimating branching probabilities using an arrival-weighted mean score with an exponential moving average, it promotes balanced expert utilization without auxiliary losses and prevents routing-mass collapse for log-concave distributions. Empirical evaluation on large-scale datasets demonstrates that this hierarchical routing preserves task quality and balanced utilization while enabling localized expert co-activation and reduced communication overhead.
Hierarchical MoE routing with a novel arrival-weighted mean estimation prevents expert collapse and enables localized co-activation, outperforming flat routing on 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 {\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.