# Vector-Balanced MoE Routing

- ID: 1998
- Canonical URL: https://synthcore.org/idea/1998/vector-balanced-moe-routing
- API JSON: https://synthcore.org/api/idea/1998.json
- API Markdown: https://synthcore.org/api/idea/1998.md
- Verification status: unverified
- Source: [arXiv:2608.08499](https://arxiv.org/abs/2608.08499)
- Category: architecture
- Solves: accuracy, stability, scalability
- ML areas: moe-routing, training
- Math tags: linear-algebra, optimization, combinatorics
- Ratings: usefulness 6/10; difficulty 5/10; novelty 6/10

## Idea description

Replace count-only MoE load balancing with greedy balancing of aggregate token-feature vectors. A token is assigned to the expert for which adding its feature vector produces the smallest increase in that expert's squared aggregate norm, encouraging experts to receive complementary semantic mixtures rather than identical token counts.

## Mathematical statement

The paper models edge weights as inner products, w_{ij}=f_i^    op f_j, where f_i\in\mathbb{R}^r is the feature vector of item i. For a partition P_1,\ldots,P_k, define S_a=\sum_{i\in P_a}f_i. The cut objective satisfies \operatorname{Cut}(P)=C-\frac{1}{2}\sum_{a=1}^k\|S_a\|_2^2, where C=\sum_{i<j}f_i^    op f_j+\frac{1}{2}\sum_i\|f_i\|_2^2 is independent of the partition. Thus maximizing the cut is equivalent to minimizing the vector-balancing objective L_{bal}=\sum_a\|S_a\|_2^2. When assigning item i to expert a, the exact incremental objective is \|S_a+f_i\|_2^2-\|S_a\|_2^2=2f_i^    op S_a+\|f_i\|_2^2; the final term is independent of a, so the greedy assignment selects a=\arg\min_a f_i^    op S_a. In an MoE, f_i is replaced by a routing feature z_t\in\mathbb{R}^r for token t, S_a is the running feature sum for expert a, and the same incremental rule is used. The transferable property is the explicit aggregate-vector balancing objective and its greedy approximation behavior; neural experiments must independently test whether this improves routing.

## Key formulas

- $$w_{ij}:=f_i^{\top}f_j,\qquad f_i\in\mathbb{R}^{r}.$$
- $$\operatorname{Cut}(P_1,\ldots,P_k)=C-\frac{1}{2}\sum_{a=1}^{k}\left\|\sum_{i\in P_a}f_i\right\|_2^2,$$ $$C=\sum_{i<j}f_i^{\top}f_j+\frac{1}{2}\sum_{i=1}^{n}\|f_i\|_2^2.$$
- $$\Delta_{t,a}=\|S_a+z_t\|_2^2-\|S_a\|_2^2=2z_t^{\top}S_a+\|z_t\|_2^2,\qquad a^*(t)=\arg\min_a\Delta_{t,a}=\arg\min_a z_t^{\top}S_a.$$
- $$S_a\leftarrow S_a+z_t\quad\text{after routing token }t\text{ to expert }a^*(t).$$

## Implementation notes

Integrate this at the top-1 or top-2 router immediately before expert dispatch. Let h_t\in\mathbb{R}^d be the Transformer hidden state for token t, let R\in\mathbb{R}^{r    imes d} be a small learned or frozen projection, and define z_t=R h_t; use z_t=\operatorname{RMSNorm}(R h_t) so vector magnitude does not dominate routing. For every microbatch, initialize S_a=0 for experts a=1,\ldots,k and process tokens in a deterministic shuffled order. For each token, compute q_a=z_t^    op S_a, mask experts whose capacity is full, select a^*=\arg\min_a q_a, dispatch the token to that expert, and update S_{a^*}\leftarrow S_{a^*}+z_t. For top-2 routing, select the first expert by this rule, temporarily update its aggregate, then select the second-best feasible expert using the same rule. Since the sequential hard assignment is nondifferentiable, use the assignments in the forward pass and train R with a straight-through estimator, or add the differentiable surrogate L_{vec}=\sum_a\|\sum_t p_{t,a}z_t\|_2^2, where p_{t,a} is the standard router probability, to the language-model loss with coefficient lambda. The exact mathematical quantities are S_a and the incremental costs 2z_t^    op S_a+\|z_t\|_2^2; no unknown optimum is required. Optimize the dot products by maintaining the k aggregate vectors and computing a k-by-r matrix product, with r=16 or 32. First test a 125M-parameter Transformer with 8 experts on WikiText-103 or a small C4 subset. Compare against standard top-1 Switch routing with count-based load balancing at equal expert capacity, FLOPs, and batch size. Record overflow/drop rate, expert token counts, aggregate norms, pairwise aggregate cosine similarities, training loss versus processed tokens, and validation perplexity. Success is lower overflow and less correlated expert input mixtures without extra routing latency, or lower perplexity at equal training FLOPs.

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.