Technical Summary: Expert Merging in Sparse Mixture of Experts with Nash Bargaining

Problem Statement

Sparse Mixture of Experts (SMoE) architectures have become a cornerstone for scaling neural networks efficiently by activating only a subset of expert modules per input. While significant research has focused on improving routing mechanisms (e.g., load balancing, stability), the complementary problem of expert merging—combining all expert parameters into a unified model for deployment or inference—remains underexplored.

Existing merging strategies typically rely on heuristic weighting schemes, such as simple input-independent averaging or top-$k$ aggregation. These methods often ignore the intricate cooperative and competitive dynamics between experts. Furthermore, recent approaches like Curvature-aware Merging of Experts (CAMEx) and its dynamic variant (EP-CAMEx), which utilize natural gradients to align experts, have shown limitations. Specifically, EP-CAMEx often underperforms its static counterpart, likely due to insufficient coordination among expert contributions and slow convergence in propagation. The core challenge is to interpret expert merging not merely as a linear average, but as a structured interaction that balances the unique contributions of diverse experts.

Methodology: NAMEx

The authors propose NAMEx (Nash Merging of Experts), a framework that reinterprets expert merging through the lens of game theory, specifically utilizing the Nash Bargaining Solution (NBS).

1. Expert Merging as a Bargaining Game

The authors frame the merging process as a cooperative-competitive game among experts.

• Players: Each expert $i$ in the SMoE layer.
• Utility Function: Defined as $u_i(\Delta E) = \tau_i^\top \Delta E$, where $\tau_i = E_i - E_m$ is the "domain-vector" representing the deviation of the $i$-th expert from the base expert $E_m$. This vector captures the expert's unique contribution.
• Disagreement Point: Set to 0, corresponding to no update to the base expert.
• Objective: Find an update direction $\Delta E$ within a ball $B_\epsilon$ that maximizes the product of individual gains over the disagreement point:
  $$\Delta E^* = \arg\max_{\Delta E \in B_\epsilon} \prod_{i=1}^N (\Delta E^\top \tau_i)$$
  Under the assumption of linearly independent domain vectors, the solution is derived as a weighted sum of domain vectors:
  $$\Delta E = \sum_{i=1}^N \alpha_i \tau_i$$
  where the weights $\alpha$ satisfy the Nash Bargaining equation: $G^\top G \alpha = 1/\alpha$, with $G$ being the matrix of domain vectors and $1/\alpha$ denoting element-wise reciprocals.

2. Integration with Expert Propagation

NAMEx is integrated into the EP-CAMEx framework. Instead of using a fixed gradient descent step for the base expert propagation, NAMEx computes the optimal update direction $\Delta E^{(l)}$ at each layer $l$ using the Nash solution. The base expert $E_m$ is updated as:
$$E_m^{(l+1)} = E_m^{(l)} + \gamma \sum_{i=1}^N \alpha_i^{(l)} \tau_i^{(l)}$$
where $\gamma$ is the step size.

3. Complex Momentum

To address the slow convergence observed in EP-CAMEx, the authors integrate complex momentum into the NAMEx framework (NAMEx-Momentum).

• A complex momentum buffer $\mu^{(j)} \in \mathbb{C}^d$ accumulates directional updates.
• The update rule involves a complex momentum coefficient $\beta \in \mathbb{C}$:
  $$\mu^{(j+1)} = \beta \mu^{(j)} + \Delta E^{(j)}$$
  $$E_m^{(j+1)} = E_m^{(j)} + \Re(\gamma \mu^{(j+1)})$$
• Theoretical analysis (Proposition 3.5 and Theorem C.3) provides convergence guarantees, showing that under mild conditions, the spectral radius of the system matrix can be kept less than 1, ensuring linear convergence.

Key Contributions

1. NAMEx Framework: A novel expert merging method that integrates the Nash Bargaining optimization framework into expert propagation, treating expert domain vectors as utility functions to derive fair and efficient merging coefficients.
2. Complex Momentum Integration: The incorporation of complex momentum into NAMEx to accelerate convergence and enhance stability, particularly in scenarios with adversarial or conflicting expert dynamics. The authors provide theoretical proofs for the convergence rate of NAMEx-Momentum.
3. Scalability and Empirical Validation: Demonstration of NAMEx's effectiveness across diverse tasks (language modeling, text classification, image classification) and its scalability to large-scale systems (Qwen1.5-MoE 14B, DeepSeek-MoE 16B).

Experimental Results

The authors evaluated NAMEx against baselines (SMoE, CAMEx, EP-CAMEx) across multiple benchmarks:

• Language Modeling (WikiText-103): NAMEx-Full-Mom achieved the lowest validation and test perplexities for both small and medium-scale pre-training, outperforming all baselines.
• Text Classification (GLUE): On the GLUE benchmark using T5-Base, NAMEx-Full-Mom achieved the best results across all tasks (SST-2, MRPC, CoLA, etc.), showing consistent improvements over SMoE and curvature-aware methods.
• Image Classification (ImageNet-1k & Robustness): On Swin-MoE, NAMEx-Mom outperformed all baselines in Top-1 accuracy. Crucially, it demonstrated superior robustness on corrupted datasets (ImageNet-A, ImageNet-O, ImageNet-R), with momentum variants showing the strongest gains.
• Large-Scale Models: When applied to DeepSeek-MoE (16B) and Qwen1.5-MoE (14B), NAMEx-Full consistently outperformed baselines and EP-CAMEx in both zero-shot and fine-tuned (SmolTalk) settings on MMLU, GSM8K, and ARC benchmarks.
• Ablation Studies:
  • Momentum: Complex momentum significantly improved convergence speed and performance compared to real-valued momentum.
  • Iterations: Solving the Nash system with 2 iterations per layer was found to be efficient and effective, with diminishing returns at higher iteration counts.
  • Disagreement Point: Using 0 as the disagreement point was found to be robust and stable compared to using the mean.

Significance and Claims

The paper claims that NAMEx offers a principled alternative to heuristic expert merging. By modeling expert interactions as a bargaining game, NAMEx ensures a Pareto-efficient merging process that accounts for both cooperation and competition among experts. The integration of complex momentum addresses convergence bottlenecks in dynamic merging scenarios.

The authors emphasize that NAMEx:

• Improves Performance: Consistently outperforms existing merging strategies across vision and language tasks.
• Enhances Robustness: Demonstrates superior generalization under data corruption and distribution shifts.
• Scales Effectively: Proves viable for large-scale MoE architectures (up to 16B parameters) without compromising computational efficiency significantly, as inference costs remain unchanged and training overhead is manageable.
• Theoretical Grounding: Provides convergence guarantees for the momentum-enhanced variant, moving beyond purely empirical heuristics.

The work suggests that viewing expert merging through game theory opens new avenues for optimizing modular deep learning systems, with potential applications in federated learning and transfer learning where fair expert combination is critical.