Relax Within, Balance Across: Geometry-Guided Load Balancing for Vision-Language Mixture-of-Experts

This work started from several routing phenomena that we believe may be useful beyond ReBA itself.

First, mixed balance can hide large modality-specific imbalance.
The standard token-level auxiliary loss only observes the combined image–text load. Image and text can therefore develop large, opposing expert preferences that cancel near one training composition. The router may appear balanced at that particular image–text ratio, but the cancellation breaks when image resolution, image count, dynamic tiling, or prompt length changes.

Second, the sensitivity to composition shift is controlled by the image–text load gap.
When the conditional image and text routing profiles are fixed, the expert-load curve over the image-token fraction is exactly quadratic. A larger gap between the two modality profiles produces a steeper curve and a narrower low-load region. This suggests that balancing one training composition is not sufficient; the modality-specific profiles themselves must be brought closer to balance.

Third, router-input geometry exposes two natural routing boundaries.
Image and text tokens occupy distinct router-input regions, making modality-specific routing shortcuts easy to learn. Within the visual modality, patches from the same image form a strongly correlated routing bloc, while different images retain a measurable separation. This indicates that the natural balancing unit for vision is an image rather than an individual patch token.

These observations lead to ReBA — Relax Within, Balance Across:

• balance image and text with separate modality objectives;
• average correlated routing decisions within each image;
• balance equal-weight image routing profiles across the batch.

Across four split vision-language MoE backbones, ReBA consistently reduces expert-load imbalance while maintaining comparable mean task accuracy. It also improves robustness under image-resolution, dynamic-tiling, and image–text composition shifts.

💻 Code: https://github.com/ZiangWu-77/ReBA

We would be especially interested in discussions about whether similar modality and instance boundaries appear in other native multimodal MoEs, and how these routing phenomena could guide future router objectives.