💡 The paper addresses the issue of load balancing in vision-language mixture-of-experts models, where the number of image and text tokens varies across batches. The standard token-level auxiliary loss, Std-Aux, balances the mixed load but can lead to large image and text load errors that cancel each other out. The authors find that the same trained routers show more than a five-fold change in load imbalance across image resolutions.

To address this issue, the authors propose a new method called Relax Within, Balance Across, or ReBA. ReBA takes into account the distinct regions occupied by image and text tokens and the strong grouping of visual tokens by source image. This motivates separate image and text terms and equal-weight routing per image.

The authors evaluate ReBA across four split backbones and find that it lowers the load on every reported benchmark input while keeping the mean task accuracy comparable to Std-Aux. ReBA also lowers the average load over the tested range and worst physical load under resolution and tiling shifts. The results demonstrate the effectiveness of ReBA in improving load balancing in vision-language mixture-of-experts models. The code for ReBA is available at https://github.com/ZiangWu-77/ReBA.