Summary (Overview)

• Problem: Vision-language Mixture-of-Experts (MoE) models suffer from load imbalance that varies with the image-text token mix. Standard token-level Switch auxiliary loss (Std-Aux) balances only the mixed load, allowing image and text load errors to cancel at one composition but diverge as the token ratio changes.
• Key Finding: The same trained router shows more than a fivefold change in load imbalance across image resolutions. The paper derives an exact "composition-shift law" showing that the image-text load gap controls sensitivity to token mix changes.
• Proposed Method: ReBA (Relax Within, Balance Across) — a geometry-guided load balancing approach that uses separate image and text auxiliary loss terms, and treats each image as one equal-weight routing instance.
• Results: Across four split backbones, ReBA lowers load on every reported benchmark input while maintaining comparable mean task accuracy to Std-Aux. It also reduces average and worst-case physical load under resolution and tiling shifts.
• Contributions: (1) Workload-wide evaluation across token mixes, (2) exact composition-shift law for fixed conditional profiles, (3) geometry-guided ReBA method.

---

Introduction and Theoretical Foundation

Background

Sparse MoE layers replace dense FFNs with $N$ parallel experts $\{E_i\}_{i=1}^N$. For token $t$, $x_t$ is the router input and $W_r x_t$ gives router logits. Efficient expert-parallel execution requires balanced expert load. The standard Switch auxiliary loss is:

$$\mathcal{L}_{\mathrm{aux}} = N \cdot \sum_{i=1}^{N} f_i \cdot p_i, \tag{1}$$

where $f_i = \frac{1}{T} \sum_t \frac{1}{k} \mathbf{1}[i \in \mathrm{top}\text{-}k(W_r x_t)]$ is the normalized hard assignment frequency and $p_i = \frac{1}{T} \sum_t \mathrm{softmax}(W_r x_t)_i$ is the average soft gate mass.

Key Theoretical Insight

The paper identifies that Std-Aux constrains only the mixed load of the current batch. Image and text can each remain imbalanced, with their expert-wise errors pointing in opposite directions and canceling near one token mix. When the token mix changes, the cancellation breaks.

Routing Geometry

Probing router-input states reveals two structural boundaries:

1. Modality boundary: Image and text occupy distinct regions of router-input space. In native routers, 52% and 61% of experts receive more than twice as much dispatch from one modality as from the other.
2. Image boundary: Visual tokens group strongly by source image, with a measurable inter-image gap exceeding the within-image gap. Text has roughly twice the cross-sample gap, justifying pooled text treatment.

These scales give two design rules:

$$\text{modality distinction} \gg \text{inter-image distinction},$$ $$\text{inter-image distinction} > \text{within-image distinction}.$$

---

Methodology

Composition-Shift Law

For fixed conditional profiles, let $q^I$ and $q^T$ be normalized image and text expert-load profiles, and $a$ be the image-token fraction. The mixed profile is $q(a) = a q^I + (1-a) q^T$. If the profile is uniform at reference fraction $a_0$:

$$\boxed{q(a) - u = (a - a_0)(q^I - q^T)} \tag{2}$$

For the squared aggregate RMS-CV $R(a) = L^{-1} \sum_l \mathrm{CV}^2(q_l(a))$:

$$\boxed{R(a) = R(a^{\star}) + \kappa(a - a^{\star})^2, \quad \kappa = \frac{N}{L} \sum_l \|q_l^I - q_l^T\|_2^2} \tag{3}$$

The load curve is a parabola: the image-text load gap controls curvature $\kappa$, determining sensitivity to token mix changes.

ReBA Loss Design

Let $g(x_t)$ be router logits, $p_{t,i} = \mathrm{softmax}(g(x_t))_i$ the soft probability, and $f_{t,i} = \frac{1}{k}\mathbf{1}[i \in \mathrm{top}\text{-}k\, g(x_t)]$ the normalized hard dispatch. For a routing instance $m$ with $S_m$ tokens:

$$P_{m,i} = \frac{1}{S_m} \sum_{t \in m} p_{t,i}, \quad F_{m,i} = \frac{1}{S_m} \sum_{t \in m} f_{t,i} \tag{4}$$

Equal-instance averages over the routing instance set $\mathcal{M}$:

$$\bar{P}_i = \frac{1}{|\mathcal{M}|} \sum_{m \in \mathcal{M}} P_{m,i}, \quad \bar{F}_i = \frac{1}{|\mathcal{M}|} \sum_{m \in \mathcal{M}} F_{m,i} \tag{5}$$

The instance balance loss:

$$\mathcal{L}_{\mathrm{inst}}(\mathcal{M}) = N \sum_{i=1}^{N} \bar{F}_i \bar{P}_i \tag{6}$$

The complete ReBA objective uses separate image and text terms weighted by observed token fractions:

$$\mathcal{L}_{\mathrm{ReBA}} = \lambda_{\mathrm{img}} \mathcal{L}_{\mathrm{inst}}(\mathcal{M}_{\mathrm{img}}) + \lambda_{\mathrm{txt}} \mathcal{L}_{\mathrm{inst}}(\mathcal{M}_{\mathrm{txt}}) \tag{7}$$

Experimental Setup

• Primary testbed: Split-Qwen3VL-4B (N=4, top-2 experts, 36 decoder layers)
• Additional backbones: Split-Qwen2.5VL-3B, Split-Qwen2VL-7B, Split-InternVL3-8B
• Training: One epoch on Cambrian-737K, varying only the auxiliary loss
• Load metrics: Mean layer CV (layers weighted equally) and RMS-CV (weights imbalanced layers more)

---

Empirical Validation / Results

Benchmark-Input Load and Accuracy

Table 1: Task accuracy and benchmark-input load across four backbones (mean layer CV, lower is better)

ReBA lowers load on every benchmark and backbone. Mean accuracy differs from Std-Aux by +0.3, +0.2, −0.2, and +0.6 points across the four backbones.

Conditional Load Improvement

How Image Balance Is Realized

Table 2: Within-image CV and between-profile share

ReBA uses both routes: tokens spread more evenly within each image, and different images develop varied expert profiles.

Ablation Study

Table 3: Ablation mean layer CV on fixed probe

Key findings: Image instances alone are insufficient (Coupled-ImgInst raises text CV to 0.640). Separate modality terms alone are insufficient (Decoupled-Matched reaches only 0.154). Both choices are necessary for the full benefit.

Physical Load Under Shifts

• ReBA reduces curvature $\kappa$ by 96.3% vs. Std-Aux
• Under resolution shifts (Split-Qwen3VL-4B): RMS-CV falls from 0.351 (Std-Aux) to 0.115 (ReBA) at highest resolution
• Under tiling shifts (InternVL): ReBA remains below Std-Aux at every tile count

Expert-Parallel Compute Proxy

The idealized speedup ratio $S_{\mathrm{ideal}} = T_{\mathrm{Std-Aux}}^{\mathrm{proxy}} / T_{\mathrm{ReBA}}^{\mathrm{proxy}}$ where $T_{\mathrm{expert}}^{\mathrm{proxy}} = \sum_l \max_e n_{l,e}$, gives 1.23–1.25× improvement at medium and high visual loads.

---

Theoretical and Practical Implications

Theoretical Contributions

1. Exact composition-shift law: The paper provides a rigorous mathematical characterization of how load imbalance varies with token mix, separating the effect of token-ratio change ($a - a_0$) from the conditional profile gap ($q^I - q^T$). This is the first exact law of its kind for vision-language MoEs.

2. Routing geometry diagnosis: The paper identifies that linear routers can easily exploit modality-separated input regions to learn opposing conditional loads, explaining why Std-Aux converges to a "checkpoint-specific balance point" rather than uniform balance.

3. Hidden cancellation problem: The work formally demonstrates that mixed-load objectives can appear successful while hiding large conditional errors that cancel at one composition.

Practical Implications

1. Deployment robustness: ReBA provides load balance that generalizes across changing token mixes—critical for real-world inference where image resolutions, counts, and prompt lengths vary dynamically.

2. Training-time solution: Unlike post-hoc inference optimizations, ReBA changes conditional routing profiles during training, addressing the root cause of imbalance.

3. Broad applicability: The routing geometry findings extend to native sparse models (Qwen3-VL-MoE-30B, Qwen3.5-MoE-35B), suggesting ReBA may apply beyond split-MoE conversion scenarios.

---

Conclusion

The paper addresses a critical gap in vision-language MoE load balancing: standard auxiliary losses balance only the mixed load at a single composition, leading to severe imbalance when the image-text token ratio changes. The authors derive an exact composition-shift law showing that the image-text conditional load gap determines sensitivity to token mix changes.

ReBA (Relax Within, Balance Across) implements two geometry-guided design choices:

1. Separate modality terms: Prevents image and text errors from canceling in a mixed loss
2. Image-level routing instances: Treats each image as one equal-weight routing instance, limiting repeated influence from correlated patches

Across four split backbones, ReBA consistently lowers load on every benchmark input while maintaining comparable task accuracy, and demonstrates robustness under resolution and tiling shifts. The idealized expert-compute proxy suggests 1.23–1.25× speedup potential at medium and high visual loads.

Future directions include applying ReBA to native sparse MoE training, extending the framework to handle conditional profile changes (not just token-ratio effects), and validating the expert-compute speedup in real latency measurements.