Technische Samenvatting: Wanneer Model Merging de Routing Verstoort: Training-vrije Kalibratie voor MoE

1. Probleemstelling

Model merging is een kosteneffectieve strategie geworden voor het consolideren van de capaciteiten van meerdere Large Language Models (LLMs) zonder de prohibitieve kosten van gezamenlijke hertraining. Hoewel bestaande technieken (bijv. Task Arithmetic, TIES-Merging) succesvol dense architecturen kunnen samenvoegen via lineaire parametrische aritmetiek, falen ze wanneer ze worden toegepast op Mixture-of-Experts (MoE) architecturen.

De auteurs identificeren een kritieke foutmodus die zij routing breakdown noemen. In MoE-modellen stuurt een router tokens dynamisch naar een subset van experts met behulp van een niet-lineaire softmax-activatie gevolgd door een discrete Top-k selectie. Het artikel betoogt dat standaard lineaire merging-methoden ongeschikt zijn voor deze setting omdat:

1. Niet-lineaire mismatch: Lineaire interpolatie van routergewichten ($W_m = \sum \lambda_i W_i$) vertaalt zich niet naar lineaire interpolatie van routing-waarschijnlijkheden vanwege de convexiteit van de softmax-functie.
2. Gevoeligheid: Kleine perturbaties in routergewichten worden exponentieel versterkt door de softmax en Top-k mechanismen.
3. Load-Balancing Constraints: Pretraining-objectieven die load balancing afdwingen, maken de toewijzing van experts zeer gevoelig voor verschuivingen in de gating.

Empirisch Bewijs: Voorlopige experimenten waarbij twee OLMoE-modellen (fijngetuned voor wiskunde en code) werden samengevoegd, toonden aan dat 50,11% van de expert-selecties direct na het mergen afweek van de bronmodellen. Deze instabiliteit stapelt zich op door de netwerkdiepte, waarbij 99,98% van de tokens routing-divergentie ervaart over 16 lagen. Cruciaal was dat het herstellen van de originele router in het samengevoegde model de prestaties herstelde, terwijl het toepassen van de samengevoegde router op het originele model de prestaties verslechterde; dit bewijst dat het behouden van expert-parameters zinloos is zonder consistente routing-consistentie.

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

Om routing breakdown aan te pakken, stellen de auteurs HARC voor, een training-vrij framework dat de samengevoegde router heraligneert met behulp van tweede-orde kromminginformatie.

2.1. Theoretische Formulering

Het doel is om een verenigde router-gewichtmatrix $W_m$ te vinden die de Kullback-Leibler (KL) divergentie tussen de outputdistributie van de samengevoegde router en de distributies van de bronrouters over alle taken minimaliseert.
$$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))]$$

Directe optimalisatie is onhandelbaar vanwege de niet-lineariteit van de softmax. De auteurs leiden een kwadratische benadering af van de KL-divergentie met behulp van een tweede-orde Taylor-expansie (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)$$
waarbij $H_i = \text{diag}(r_i) - r_i r_i^\top$ de Hessian is van de log-partition functie. Deze Hessian vangt de competitieve dynamiek tussen experts (off-diagonal termen) en de gevoeligheid van actieve experts (diagonale termen) op.

2.2. Closed-Form Oplossing

Door de lineaire relatie $z = Wx$ in de kwadratische benadering te substitueren, wordt het probleem getransformeerd naar een gewogen kleinste-kwadraten systeem. De optimale gewichten $w^*_m$ (gevectoriseerde vorm van $W_m$) voldoen aan het lineaire systeem:
$$\left( \sum_{i=1}^D \Sigma_i \right) w^*_m = \sum_{i=1}^D \Sigma_i w_i$$
waarbij $\Sigma_i = \mathbb{E}_{x \sim \mathcal{D}_i} [H_i \otimes xx^\top]$ de taakspecifieke precisie-matrix is.

2.3. Efficiënte Berekening (Matrix-Free)

Het expliciet construeren en inverteren van de geaggregeerde precisie-matrix $\Sigma$ is computationeel onhaalbaar ($O(K^2d^2)$ ruimtecomplexiteit). HARC gebruikt twee belangrijke strategieën om schaalbaarheid te garanderen:

1. Diagonale Regularisatie: Om rangtekort in de Hessian te behandelen (door sparse Top-k gating), introduceren de auteurs een regularisatieterm die interpoleert tussen de volledige precisie-matrix en de diagonale benadering. Dit stabiliseert de oplossing zonder uit te gaan van onafhankelijkheid tussen experts.
2. Matrix-Free Conjugate Gradient (CG): In plaats van $\Sigma$ te materialiseren, gebruikt het algoritme een matrix-vrije operator om matrix-vector producten $\Sigma w$ on-the-fly te berekenen met behulp van Kronecker-product identiteiten. Dit reduceert de ruimtecomplexiteit naar $O(Kd)$ en maakt het gebruik van de Conjugate Gradient methode mogelijk om het systeem efficiënt op te lossen.

2.4. Algoritme

HARC werkt als een post-processing stap op een ongekalibreerd samengevoegd model. Voor elke MoE-laag:

1. Voer een forward pass uit met kalibratiedata door het samengevoegde model tot de huidige laag om input features te genereren.
2. Bereken de routing-outputs en Hessians voor de bronmodellen.
3. Acumuleer statistieken om het lineaire systeem te vormen.
4. Los de gekalibreerde router-gewichten op met de matrix-vrije CG-solver.
5. Update de router-gewichten in het samengevoegde model.

3. Belangrijkste Bijdragen

1. Identificatie van Routing Breakdown: Het artikel stelt vast dat routing-instabiliteit, in plaats van expert-parameter interferentie, de primaire oorzaak is van prestatievermindering bij MoE merging.
2. HARC Framework: Een nieuw, training-vrij kalibratiemethode dat tweede-orde Hessian-informatie gebruikt om router-distributies te heraligneren, wat een closed-form oplossing biedt voor het niet-lineaire merging-probleem.
3. Schaalbare Optimalisatie: De afleiding van een matrix-vrije Conjugate Gradient solver die HARC in staat stelt om te schalen naar grote MoE-architecturen zonder expliciete materialisatie van de Hessian.
4. Empirische Validatie: Demonstratie dat HARC effectief routing breakdown mitigeert over diverse merging baselines (Weight Averaging, TIES-Merging, WUDI, etc.) en taken (Math, Code).

4. Experimentele Resultaten

Experimenten werden uitgevoerd op het OLMoE-1B-7B-0125 model, waarbij varianten werden samengevoegd die zijn fijngetuned op wiskundige redenering (GSM8K, MATH500) en codegeneratie (HumanEval+, MBPP+).

• Consistente Verbetering: HARC verbeterde consequent de prestaties over alle geteste merging baselines. Zo verhoogde het, wanneer toegepast op de sterke WUDI-Merging baseline, de algemene score van 38,08 naar 38,96.
• Mitigatie van Domeingevoeligheid: Naïef mergen degradeerde de prestaties voor codegeneratie ernstig (bijv. HumanEval+ daalde van 33,50 naar 20,27 met Weight Averaging). HARC herstelde een deel van dit verlies, wat de capaciteit bewijst om domeinspecifieke vaardigheden te behouden door de expert-selectie te corrigeren.
• Ablatie Studies:
  • Hessian Structuur: Het gebruik van de volledige Hessian presteerde beter dan diagonale of identiteits-benaderingen, wat bevestigt dat het modelleren van off-diagonale competitieve dynamiek cruciaal is.
  • Regularisatie: De dempingsfactor $\alpha$ (die off-diagonals onderdrukt) bleek stabieler dan diagonale schaling $\gamma$.
  • Data Compositie: Kalibratie met gefilterde correcte bron-responses (SrcCrt) leverde de beste resultaten op. Verrassend genoeg presteerde het gebruik van alleen prompts vergelijkbaar met volledige responses, wat suggereert dat de initiële routing-beslissingen doorslaggevend zijn.
• Data Efficiëntie: HARC bereikte bijna optimale prestaties met slechts 40% van de kalibratie-samples, wat wijst op een snelle convergentie van tweede-orde statistieken.
• Laag-specifieke Analyse: Visualisatie van de KL-divergentie toonde aan dat routing-fouten in baseline-methoden exponentieel toenemen in diepere lagen. HARC slaagde erin deze trend te onderdrukken en behield de alignment in kritieke diepere lagen.

5. Betekenis en Claims

Het artikel claimt dat het behouden van de routing-logica even essentieel is als het uitlijnen van parametergewichten in sparse architecturen. Door de specifieke foutmodus van routing breakdown aan te pakken, biedt HARC een robuust en efficiënt pad voor het consolideren van MoE-modellen zonder hertraining.

De auteurs benadrukken dat hun methode:

• Geen hertraining vereist of gradient backpropagation door de experts.
• Generaliseert over verschillende merging-strategieën (lineaire aritmetiek, optimalisatie-gebaseerd) als een lichtgewicht post-processing stap.
• Computationele kosten en barrières verlaagt voor het inzetten van capabele modellen door het betrouwbaar hergebruik van meerdere expert-modellen mogelijk te maken.

Het werk concludeert dat hoewel bestaande lineaire merging-hypothesen (Linear Mode Connectivity) gelden voor dense lagen, ze broos zijn voor MoE routers. HARC overbrugt deze kloof door expliciet rekening te houden met de niet-lineaire gating-dynamiek die inherent is aan Mixture-of-Experts architecturen.