Résumé Technique : Quand la fusion de modèles brise le routage : Calibration sans entraînement pour les MoE

1. Énoncé du problème

La fusion de modèles est devenue une stratégie rentable pour consolider les capacités de multiples grands modèles de langage (LLM) sans les coûts prohibitifs d'un réentraînement conjoint. Bien que les techniques existantes (par exemple, Task Arithmetic, TIES-Merging) réussissent à fusionner des architectures denses via l'arithmétique linéaire des paramètres, elles échouent lorsqu'elles sont appliquées aux architectures de type Mixture-of-Experts (MoE).

Les auteurs identent un mode de défaillance critique appelé rupture de routage (routing breakdown). Dans les modèles MoE, un routeur distribue dynamiquement les jetons à un sous-ensemble d'experts en utilisant une activation softmax non linéaire suivie d'une sélection Top-k discrète. L'article soutient que les méthodes de fusion linéaire standard sont mal adaptées à ce contexte car :

1. Incompatibilité de non-linéarité : L'interpolation linéaire des poids du routeur ($W_m = \sum \lambda_i W_i$) ne se traduit pas par une interpolation linéaire des probabilités de routage en raison de la convexité de la fonction softmax.
2. Sensibilité : De petites perturbations dans les poids du routeur sont amplifiées exponentiellement par les mécanismes softmax et Top-k.
3. Contraintes d'équilibrage de charge : Les objectifs de pré-entraînement qui imposent un équilibrage de charge rendent l'assignation des experts extrêmement sensible aux changements de lissage (gating).

Preuve empirique : Des expériences préliminaires fusionnant deux modèles OLMoE (fine-tunés pour les mathématiques et le code) ont révélé que 50,11 % des sélections d'experts divergeaient des modèles sources immédiatement après la fusion. Cette instabilité s'accentue à travers la profondeur du réseau, avec 99,98 % des jetons subissant une divergence de routage sur 16 couches. Crucialement, l'évaluation croisée a montré que la restauration du routeur d'origine dans le modèle fusionné permettait de récupérer les performances, tandis que l'application du routeur fusionné au modèle d'origine dégradait celles-ci, prouvant que préserver les paramètres des experts est futile sans un comportement de routage cohérent.

2. Méthodologie : Calibration du routeur sensible à l'Hessienne (HARC)

Pour remédier à la rupture de routage, les auteurs proposent HARC, un cadre sans entraînement qui réaligne le routeur fusionné en utilisant des informations de courbure de second ordre.

2.1. Formulation théorique

L'objectif est de trouver une matrice de poids de routeur unifiée $W_m$ qui minimise la divergence de Kullback-Leibler (KL) entre la distribution de sortie du routeur fusionné et les distributions des routeurs sources à travers toutes les tâches.
$$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))]$$

L'optimisation directe est intraitable en raison de la non-linéarité de la softmax. Les auteurs dérivent une approximation quadratique de la divergence KL en utilisant une expansion de Taylor de second ordre (Lemme 4.1) :
$$\text{KL}(r_i \parallel r_m) \approx \frac{1}{2}(z_m - z_i)^\top H_i (z_m - z_i)$$
où $H_i = \text{diag}(r_i) - r_i r_i^\top$ est la Hessienne de la fonction de partition logarithmique. Cette Hessienne capture la dynamique compétitive entre les experts (termes hors-diagonaux) et la sensibilité des experts actifs (termes diagonaux).

2.2. Solution en forme fermée

En substituant la relation linéaire $z = Wx$ dans l'approximation quadratique, le problème se transforme en un système de moindres carrés pondérés. Les poids optimaux $w^*_m$ (forme vectorisée de $W_m$) satisfont le système linéaire :
$$\left( \sum_{i=1}^D \Sigma_i \right) w^*_m = \sum_{i=1}^D \Sigma_i w_i$$
où $\Sigma_i = \mathbb{E}_{x \sim \mathcal{D}_i} [H_i \otimes xx^\top]$ est la matrice de précision spécifique à la tâche.

2.3. Calcul efficace (sans matrice explicite)

La construction explicite et l'inversion de la matrice de précision agrégée $\Sigma$ sont prohibitivement coûteuses ($O(K^2d^2)$ en complexité d'espace). HARC emploie deux stratégies clés pour assurer la scalabilité :

1. Régularisation diagonale : Pour gérer la déficience de rang de l'Hessienne (due au gating Top-k creux), les auteurs introduisent un terme de régularisation qui interpole entre la matrice de précision complète et son approximation diagonale. Cela stabilise la solution sans supposer l'indépendance entre les experts.
2. Gradient Conjugué (CG) sans matrice : Au lieu de matérialiser $\Sigma$, l'algorithme utilise un opérateur sans matrice pour calculer les produits matrice-vecteur $\Sigma w$ à la volée en utilisant des identités de produit de Kronecker. Cela réduit la complexité spatiale à $O(Kd)$ et permet d'utiliser la méthode du Gradient Conjugué pour résoudre efficacement le système.

2.4. Algorithme

HARC fonctionne comme une étape de post-traitement sur un modèle fusionné non calibré. Pour chaque couche MoE :

1. Passer les données de calibration par le modèle fusionné jusqu'à la couche actuelle pour générer les caractéristiques d'entrée.
2. Calculer les sorties de routage et les Hessiennes pour les modèles sources.
3. Accumuler les statistiques pour former le système linéaire.
4. Résoudre pour les poids du routeur calibrés en utilisant le solveur CG sans matrice.
5. Mettre à jour les poids du routeur dans le modèle fusionné.

3. Contributions clés

1. Identification de la rupture de routage : L'article établit que l'instabilité du routage, plutôt que l'interférence des paramètres des experts, est la cause principale de la dégradation des performances lors de la fusion de MoE.
2. Cadre HARC : Une nouvelle méthode de calibration sans entraînement qui exploite l'information de la Hessienne de second ordre pour réaligner les distributions de routage, offrant une solution en forme fermée au problème de fusion non linéaire.
3. Optimisation scalable : La dérivation d'un solveur de Gradient Conjugué sans matrice qui permet à HARC de passer à l'échelle sur de grandes architectures MoE sans matérialisation explicite de l'Hessienne.
4. Validation empirique : Démonstration que HARC atténue efficacement la rupture de routage à travers diverses bases de fusion (Weight Averaging, TIES-Merging, WUDI, etc.) et diverses tâches (Math, Code).

4. Résultats expérimentaux

Les expériences ont été menées sur le modèle OLMoE-1B-7B-0125, fusionnant des variantes spécialisées dans le raisonnement mathématique (GSM8K, MATH500) et la génération de code (HumanEval+, MBPP+).

• Amélioration constante : HARC a systématiquement amélioré les performances de toutes les bases de fusion testées. Par exemple, appliqué à la base solide WUDI-Merging, le score global est passé de 38,08 à 38,96.
• Atténuation de la sensibilité au domaine : La fusion naïve dégrade sévèrement les performances de génération de code (par exemple, HumanEval+ chute de 33,50 à 20,27 avec Weight Averaging). HARC a partiellement récupéré cette perte, démontrant sa capacité à préserver les capacités spécifiques au domaine en corrigeant la sélection des experts.
• Études d'ablation :
  • Structure de l'Hessienne : L'utilisation de l'Hessienne complète a surpassé les approximations diagonales ou identité, confirmant que la modélisation des dynamiques compétitives hors-diagonales est cruciale.
  • Régularisation : Le facteur d'amortissement $\alpha$ (supprimant les termes hors-diagonaux) s'est avéré plus stable que l'échelle diagonale $\gamma$.
  • Composition des données : La calibration utilisant les réponses sources correctes filtrées (SrcCrt) a donné les meilleurs résultats. Étonnamment, l'utilisation de prompts uniquement a donné des résultats comparables aux réponses complètes, suggérant que les décisions de routage initiales sont primordiales.
• Efficacité des données : HARC atteint des performances proches de l'optimal en utilisant seulement 40 % des échantillons de calibration, indiquant une convergence rapide des statistiques de second ordre.
• Analyse par couche : La visualisation de la divergence KL a montré que les erreurs de routage dans les méthodes de base s'accumulent exponentiellement dans les couches plus profondes. HARC a réussi à supprimer cette tendance, maintenant l'alignement dans les couches profondes critiques.

5. Signification et affirmations

L'article affirme que préserver la logique de routage est aussi essentiel que d'aligner les poids des paramètres dans les architectures creuses. En traitant le mode de défaillance spécifique de la rupture de routage, HARC offre une voie robuste et efficace pour la consolidation des modèles MoE sans réentraînement.

Les auteurs soulignent que leur méthode :

• Ne nécessite pas de réentraînement ni de rétropropagation du gradient à travers les experts.
• Se généralise à travers différentes stratégies de fusion (arithmétique linéaire, optimisation) en tant qu'étape de post-traitement légère.
• Réduit les coûts de calcul et les barrières au déploiement de modèles capables en permettant la réutilisation fiable de multiples modèles d'experts.

L'étude conclut que si les hypothèses de fusion linéaire existantes (Connectivité de Mode Linéaire) tiennent pour les couches denses, elles sont fragiles pour les routeurs MoE. HARC comble cette lacune en tenant explicitement compte de la dynamique de lissage non linéaire inhérente aux architectures Mixture-of-Experts.