BRANCH-MoE: Balance-Aware Tree Routing for Large Embedding Models
[Uczenie maszynowe] Autorzy: Gang Fu, Adel Javanmard, MohammadHossein Bateni i 1 więcej
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-ma…
- Kategoria arxiv
- cs.LG
- Kategoria
- Uczenie maszynowe