BALANCED BILATERAL ROUTING FOR MIXTURE-OF-EXPERTS: AN EMPIRICAL STUDY

Abstract

Mixture-of-Experts (MoE) layers enable massive language models to scale efficiently by activating only a subset of parameters per token. However, load imbalance – where some experts receive far more tokens than others – remains a core challenge, often addressed via auxiliary losses or explicit capacity constraints. We introduce Balanced Bilateral Routing (BBR), a novel routing mechanism that combines: (i) a differentiable Sinkhorn‑Knopp assignment enforcing per‑expert capacity constraints, (ii) learnable token‑expert embeddings with exponential moving average (EMA) prototypes, and (iii) an anti‑collapse loss that pushes prototypes apart. BBR guarantees load balance by design and is compatible with autoregressive generation. We compare BBR against a strong Switch‑Transformer baseline on a 16‑expert language model trained on WikiText‑103. In a fair setting where both routers use top‑4 soft combination, the standard router achieves a final perplexity of 345.68 (lower is better) and an expert utilization entropy of 0.9949, while BBR achieves 374.18 perplexity and 0.9850 entropy. Despite the negative result, our analysis reveals that the Sinkhorn capacity constraint is not perfectly enforced at low temperature, and the standard auxiliary loss remains highly effective up to 16 experts. We provide a detailed ablation, discuss the limitations of Sinkhorn‑based balancing, and release our code and logs for reproducibility.