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

Hugging Face Daily Papers Papers
mixture-of-experts load-balancing vision-language routing multimodal

Summary

This paper presents ReBA (Relax Within, Balance Across), a geometry-guided load-balancing method for vision-language mixture-of-experts models. It addresses modality-specific routing imbalance by balancing image and text separately and grouping correlated visual tokens, improving robustness across resolution and composition shifts.

Vision-language MoE batches contain different numbers of image and text tokens. Image resolution, image count, tiling, and prompt length all change this token mix. We call the standard token-level Switch auxiliary loss Std-Aux. Std-Aux balances only the mixed load, so large image and text load errors can cancel at one mix. On our main model, the same trained router shows more than a fivefold change in load imbalance across image resolutions. We hold the image and text load profiles fixed and derive the exact load curve as the token mix varies. The image-text load gap controls sensitivity to the token mix. Physical preprocessing can also change the conditional profiles. The fixed-profile law excludes such changes. To design a remedy, we examine the router input structure. Image and text occupy distinct regions, while visual tokens group strongly by source image. The modality boundary motivates separate image and text terms. The image boundary motivates one equal-weight routing instance per image. ReBA, or Relax Within, Balance Across, implements both choices. Across four split backbones, ReBA lowers load on every reported benchmark input while keeping mean task accuracy comparable to Std-Aux. ReBA also lowers average load over the tested range and worst physical load under resolution and tiling shifts. Code is available at https://github.com/ZiangWu-77/ReBA.
Original Article
View Cached Full Text

Cached at: 08/04/26, 09:38 AM

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

Source: https://huggingface.co/papers/2608.00574 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 toReBA — 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.

Similar Articles

$\phi$-Balancing for Mixture-of-Experts Training

arXiv cs.LG

This paper proposes φ-balancing, a principled framework for load balancing in Mixture-of-Experts models that directly targets population-level expert balance using convex duality and mirror descent, achieving more stable expert utilization and outperforming prior methods on reasoning and code generation benchmarks.

Rank-Aware Hyperbolic Alignment for Vision-Language Dataset Distillation

Hugging Face Daily Papers

This paper proposes Rank-Aware Hyperbolic Alignment (RAHA), a method for vision-language dataset distillation that leverages hyperbolic geometry and alignment capacity control to efficiently compress large image-text datasets into high-quality synthetic pairs.

Balancing Multimodal Learning through Label Space Reshaping

arXiv cs.LG

Introduces Balanced Multimodal Label Reshaping (BMLR), a method that addresses modality imbalance in multimodal learning by reshaping the label space to equalize mapping difficulty across modalities, improving performance across various architectures.