Technical Summary: When Model Merging Breaks Routing: Training-Free Calibration for MoE

1. Problem Statement

Model merging has become a cost-effective strategy for consolidating the capabilities of multiple Large Language Models (LLMs) without the prohibitive costs of joint retraining. While existing techniques (e.g., Task Arithmetic, TIES-Merging) successfully merge dense architectures via linear parameter arithmetic, they fail when applied to Mixture-of-Experts (MoE) architectures.

The authors identify a critical failure mode termed routing breakdown. In MoE models, a router dynamically dispatches tokens to a subset of experts using a non-linear softmax activation followed by a discrete Top-k selection. The paper argues that standard linear merging methods are ill-suited for this setting because:

1. Non-linearity Mismatch: Linear interpolation of router weights ($W_m = \sum \lambda_i W_i$) does not translate to linear interpolation of routing probabilities due to the convexity of the softmax function.
2. Sensitivity: Small perturbations in router weights are exponentially amplified by the softmax and Top-k mechanisms.
3. Load-Balancing Constraints: Pretraining objectives that enforce load balancing make expert assignments highly sensitive to gating shifts.

Empirical Evidence: Preliminary experiments merging two OLMoE models (fine-tuned for math and code) revealed that 50.11% of expert selections deviated from source models immediately after merging. This instability compounds through network depth, with 99.98% of tokens experiencing routing divergence across 16 layers. Crucially, cross-evaluation showed that restoring the original router to the merged model recovered performance, while applying the merged router to the original model degraded it, proving that preserving expert parameters is futile without consistent routing behavior.

2. Methodology: Hessian-Aware Router Calibration (HARC)

To address routing breakdown, the authors propose HARC, a training-free framework that realigns the merged router using second-order curvature information.

2.1. Theoretical Formulation

The objective is to find a unified router weight matrix $W_m$ that minimizes the Kullback-Leibler (KL) divergence between the merged router's output distribution and the source routers' distributions across all tasks.
$$W^*_m = \arg \min_{W_m} \sum_{i=1}^D \mathbb{E}_{x \sim \mathcal{D}_i} [\text{KL}(r_i(x) \parallel r_m(x; W_m))]$$

Direct optimization is intractable due to the non-linearity of softmax. The authors derive a quadratic approximation of the KL divergence using a second-order Taylor expansion (Lemma 4.1):
$$\text{KL}(r_i \parallel r_m) \approx \frac{1}{2}(z_m - z_i)^\top H_i (z_m - z_i)$$
where $H_i = \text{diag}(r_i) - r_i r_i^\top$ is the Hessian of the log-partition function. This Hessian captures the competitive dynamics between experts (off-diagonal terms) and the sensitivity of active experts (diagonal terms).

2.2. Closed-Form Solution

Substituting the linear relationship $z = Wx$ into the quadratic approximation transforms the problem into a weighted least-squares system. The optimal weights $w^*_m$ (vectorized form of $W_m$) satisfy the linear system:
$$\left( \sum_{i=1}^D \Sigma_i \right) w^*_m = \sum_{i=1}^D \Sigma_i w_i$$
where $\Sigma_i = \mathbb{E}_{x \sim \mathcal{D}_i} [H_i \otimes xx^\top]$ is the task-specific precision matrix.

2.3. Efficient Computation (Matrix-Free)

Explicitly constructing and inverting the aggregate precision matrix $\Sigma$ is computationally prohibitive ($O(K^2d^2)$ space complexity). HARC employs two key strategies to ensure scalability:

1. Diagonal Regularization: To handle rank-deficiency in the Hessian (due to sparse Top-k gating), the authors introduce a regularization term that interpolates between the full precision matrix and its diagonal approximation. This stabilizes the solution without assuming independence between experts.
2. Matrix-Free Conjugate Gradient (CG): Instead of materializing $\Sigma$, the algorithm uses a matrix-free operator to compute matrix-vector products $\Sigma w$ on the fly using Kronecker product identities. This reduces space complexity to $O(Kd)$ and allows the use of the Conjugate Gradient method to solve the system efficiently.

2.4. Algorithm

HARC operates as a post-processing step on an uncalibrated merged model. For each MoE layer:

1. Forward pass calibration data through the merged model up to the current layer to generate input features.
2. Compute routing outputs and Hessians for source models.
3. Accumulate statistics to form the linear system.
4. Solve for the calibrated router weights using the matrix-free CG solver.
5. Update the router weights in the merged model.

3. Key Contributions

1. Identification of Routing Breakdown: The paper establishes that routing instability, rather than expert parameter interference, is the primary cause of performance degradation in MoE merging.
2. HARC Framework: A novel, training-free calibration method that leverages second-order Hessian information to realign router distributions, offering a closed-form solution to the non-linear merging problem.
3. Scalable Optimization: The derivation of a matrix-free Conjugate Gradient solver that enables HARC to scale to large MoE architectures without explicit Hessian materialization.
4. Empirical Validation: Demonstration that HARC effectively mitigates routing breakdown across diverse merging baselines (Weight Averaging, TIES-Merging, WUDI, etc.) and tasks (Math, Code).

4. Experimental Results

Experiments were conducted on the OLMoE-1B-7B-0125 model, merging variants fine-tuned on mathematical reasoning (GSM8K, MATH500) and code generation (HumanEval+, MBPP+).

• Consistent Improvement: HARC consistently improved performance across all tested merging baselines. For instance, when applied to the strong WUDI-Merging baseline, the overall score increased from 38.08 to 38.96.
• Mitigation of Domain Sensitivity: Naive merging severely degraded code generation performance (e.g., HumanEval+ dropped from 33.50 to 20.27 with Weight Averaging). HARC partially recovered this loss, demonstrating its ability to preserve domain-specific capabilities by correcting expert selection.
• Ablation Studies:
  • Hessian Structure: Using the full Hessian outperformed diagonal or identity approximations, confirming that modeling off-diagonal competitive dynamics is crucial.
  • Regularization: The damping factor $\alpha$ (suppressing off-diagonals) proved more stable than diagonal scaling $\gamma$.
  • Data Composition: Calibration using filtered correct source responses (SrcCrt) yielded the best results. Surprisingly, using prompts only performed comparably to full responses, suggesting that initial routing decisions are paramount.
• Data Efficiency: HARC reached near-optimal performance using only 40% of the calibration samples, indicating rapid convergence of second-order statistics.
• Layer-wise Analysis: Visualization of KL divergence showed that routing errors in baseline methods accumulate exponentially in deeper layers. HARC successfully suppressed this trend, maintaining alignment in critical deeper layers.

5. Significance and Claims

The paper claims that preserving routing logic is as essential as aligning parameter weights in sparse architectures. By addressing the specific failure mode of routing breakdown, HARC offers a robust and efficient path for consolidating MoE models without retraining.

The authors emphasize that their method:

• Does not require retraining or gradient backpropagation through the experts.
• Generalizes across different merging strategies (linear arithmetic, optimization-based) as a lightweight post-processing step.
• Reduces computational costs and barriers to deploying capable models by enabling the reliable reuse of multiple expert models.

The work concludes that while existing linear merging hypotheses (Linear Mode Connectivity) hold for dense layers, they are brittle for MoE routers. HARC bridges this gap by explicitly accounting for the non-linear gating dynamics inherent to Mixture-of-Experts architectures.