Input Domain Aware MoE: Decoupling Routing Decisions from Task Optimization in Mixture of Experts

TL;DR

The paper tackles the specialization–balance dilemma in sparse MoE for vision-language models by decoupling routing from task optimization. It introduces IDA-MoE, which grounds routing in the input distribution using a multi-component Gaussian Mixture Model and a component-based routing mechanism, complemented by a component reactivation strategy. The approach yields more decisive routing and better expert utilization without auxiliary load-balancing losses, achieving superior performance on VLM benchmarks with 2B-activated parameters and stable training. A key insight is that decoupled routing fosters natural load balance, improving robustness to distribution shifts and routing noise. Overall, IDA-MoE advances scalable, efficient MoE routing by aligning routing decisions with input structure rather than task-driven loss signals, improving both performance and resource utilization.

Abstract

Sparse Mixture of Experts (sMoE) has become a pivotal approach for scaling large vision-language models, offering substantial capacity while maintaining computational efficiency through dynamic, sparse activation of experts. However, existing routing mechanisms, typically based on similarity scoring, struggle to effectively capture the underlying input structure. This limitation leads to a trade-off between expert specialization and balanced computation, hindering both scalability and performance. We propose Input Domain Aware MoE, a novel routing framework that leverages a probabilistic mixture model to better partition the input space. By modeling routing probabilities as a mixture of distributions, our method enables experts to develop clear specialization boundaries while achieving balanced utilization. Unlike conventional approaches, our routing mechanism is trained independently of task-specific objectives, allowing for stable optimization and decisive expert assignments. Empirical results on vision-language tasks demonstrate that our method consistently outperforms existing sMoE approaches, achieving higher task performance and improved expert utilization balance.

Figures (6)

• Figure 1: Benchmark performance of IDA-MoE (ours) across visual question answering and multimodal reasoning tasks. The left panel (\ref{['fig:ablation_radar']}) illustrates the effectiveness of different MoE routing implementations, while the right panel (\ref{['fig:sota_radar']}) demonstrates comparative performance against models of similar computational scale.
• Figure 2: Illustration of load balancing loss effects. Experts colored circles and their corresponding token assignments colored dots in matching colors. The transition from \ref{['fig:before_lb']} to \ref{['fig:after_lb']} simulates the effect of load balancing loss. Token embeddings are pushed toward and accumulated around decision boundaries.
• Figure 3: Routing probability heatmaps of the last MoE layer for the same input sequence. The x-axis represents token indices, the y-axis represents expert indices, and the color intensity indicates the probability of assigning the $i$-th token to the $j$-th expert. Note how IDA-MoE (b) exhibits significantly more decisive routing patterns with clearer token-expert affinities compared to the conventional approach (a).
• Figure 4: Comparison between IDA-MoE and traditional MoE routing. (Left) IDA-MoE: Input token representations ($\textbf{u}_1,\textbf{u}_2$ ) are projected into a lower-dimensional routing space where their distribution is modeled into distinct clusters. Tokens are routed to specific experts based on their cluster index, effectively decoupling routing decisions from task-specific optimization. (Right) Traditional MoE: A learned router module directly calculates token-expert affinities (classification-style) to determine expert assignments, typically trained end-to-end with the task objective.
• Figure 5: $\mathbf{\mathrm{CV}}_\mathrm{mean}$ change over training.