Riepilogo Tecnico: Quando il Merging dei Modelli Rompe il Routing: Calibrazione Training-Free per MoE

1. Definizione del Problema

Il merging dei modelli è diventato una strategia conveniente per consolidare le capacità di molteplici Large Language Models (LLM) senza i costi proibitivi del riaddestramento congiunto. Sebbene le tecniche esistenti (ad esempio, Task Arithmetic, TIES-Merging) riescano a fondere architetture dense tramite aritmetica lineare dei parametri, esse falliscono quando applicate ad architetture Mixture-of-Experts (MoE).

Gli autori identificano un modo di guasto critico denominato routing breakdown (rottura del routing). Nei modelli MoE, un router smista dinamicamente i token a un sottoinsieme di esperti utilizzando un'attivazione softmax non lineare seguita da una selezione discreta Top-k. Il paper sostiene che i metodi standard di merging lineare siano inadatti a questo scenario perché:

1. Discrepanza di Non-linearità: L'interpolazione lineare dei pesi del router ($W_m = \sum \lambda_i W_i$) non si traduce in un'interpolazione lineare delle probabilità di routing a causa della convessità della funzione softmax.
2. Sensibilità: Piccole perturbazioni nei pesi del router vengono amplificate esponenzialmente dai meccanismi softmax e Top-k.
3. Vincoli di Load-Balancing: Gli obiettivi di pretraining che impongono il bilanciamento del carico rendono l'assegnazione degli esperti altamente sensibile ai cambiamenti di gating.

Evidenza Empirica: Esperimenti preliminari fondendo due modelli OLMoE (specializzati in matematica e codice) hanno rivelato che il 50,11% delle selezioni degli esperti deviava dai modelli sorgente immediatamente dopo il merging. Questa instabilità si propaga attraverso la profondità della rete, con il 99,98% dei token che sperimenta una divergenza di routing su 16 layer. Fondamentalmente, la valutazione incrociata ha mostrato che il ripristino del router originale nel modello fuso recuperava le prestazioni, mentre l'applicazione del router fuso al modello originale le degradava, provando che preservare i parametri degli esperti è inutile senza un comportamento di routing coerente.

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

Per affrontare il routing breakdown, gli autori propongono HARC, un framework training-free che riallinea il router fuso utilizzando informazioni sulla curvatura del secondo ordine.

2.1. Formulazione Teorica

L'obiettivo è trovare una matrice dei pesi del router unificata $W_m$ che minimizzi la divergenza di Kullback-Leibler (KL) tra la distribuzione di output del router fuso e le distribuzioni dei router sorgente attraverso tutti i task.
$$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'ottimizzazione diretta è intrattabile a causa della non-linearità della softmax. Gli autori derivano un'approssimazione quadratica della divergenza KL utilizzando un'espansione di Taylor del secondo ordine (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)$$
dove $H_i = \text{diag}(r_i) - r_i r_i^\top$ è l'Hessiana della funzione di partizione logaritmica. Questa Essiana cattura la dinamica competitiva tra gli esperti (termini fuori diagonale) e la sensibilità degli esperti attivi (termini diagonali).

2.2. Soluzione in Forma Chiusa

Sostituendo la relazione lineare $z = Wx$ nell'approssimazione quadratica, il problema si trasforma in un sistema di minimi quadrati pesati. I pesi ottimali $w^*_m$ (forma vettorizzata di $W_m$) soddisfano il sistema lineare:
$$\left( \sum_{i=1}^D \Sigma_i \right) w^*_m = \sum_{i=1}^D \Sigma_i w_i$$
dove $\Sigma_i = \mathbb{E}_{x \sim \mathcal{D}_i} [H_i \otimes xx^\top]$ è la matrice di precisione specifica per il task.

2.3. Computazione Efficiente (Matrix-Free)

La costruzione esplicita e l'inversione della matrice di precisione aggregata $\Sigma$ sono computazionalmente proibitive (complessità di spazio $O(K^2d^2)$). HARC impiega due strategie chiave per garantire la scalabilità:

1. Regolarizzazione Diagonale: Per gestire la scarsità di rango nell'Hessiana (dovuta al gating sparso Top-k), gli autori introducono un termine di regolarizzazione che interpola tra la matrice di precisione completa e la sua approssimazione diagonale. Ciò stabilizza la soluzione senza assumere l'indipendenza tra gli esperti.
2. Conjugate Gradient (CG) Matrix-Free: Invece di materializzare $\Sigma$, l'algoritmo utilizza un operatore matrix-free per calcolare i prodotti matrice-vettore $\Sigma w$ al volo utilizzando identità di prodotto di Kronecker. Ciò riduce la complessità di spazio a $O(Kd)$ e consente l'uso del metodo del Gradiente Coniugato per risolvere il sistema in modo efficiente.

2.4. Algoritmo

HARC opera come uno step di post-processing su un modello fuso non calibrato. Per ogni layer MoE:

1. Passaggio forward dei dati di calibrazione attraverso il modello fuso fino al layer corrente per generare le feature di input.
2. Calcolo degli output di routing e delle Essiane per i modelli sorgente.
3. Accumulo delle statistiche per formare il sistema lineare.
4. Risoluzione per i pesi del router calibrati utilizzando il solver CG matrix-free.
5. Aggiornamento dei pesi del router nel modello fuso.

3. Contributi Chiave

1. Identificazione del Routing Breakdown: Il paper stabilisce che l'instabilità del routing, piuttosto che l'interferenza dei parametri degli esperti, è la causa primaria del degrado delle prestazioni nel merging MoE.
2. Framework HARC: Un nuovo metodo di calibrazione training-free che sfrutta l'informazione dell'Essiana del secondo ordine per riallineare le distribuzioni del router, offrendo una soluzione in forma chiusa al problema del merging non lineare.
3. Ottimizzazione Scalabile: La derivazione di un solver Conjugate Gradient matrix-free che consente a HARC di scalare su grandi architetture MoE senza la materializzazione esplicita dell'Essiana.
4. Validazione Empirica: Dimostrazione che HARC mitiga efficacementamente il routing breakdown attraverso diversi baseline di merging (Weight Averaging, TIES-Merging, WUDI, ecc.) e task (Math, Code).

4. Risultati Sperimentali

Gli esperimenti sono stati condotti sul modello OLMoE-1B-7B-0125, fondendo varianti specializzate nel ragionamento matematico (GSM8K, MATH500) e nella generazione di codice (HumanEval+, MBPP+).

• Miglioramento Consistente: HARC ha migliorato costantemente le prestazioni in tutti i baseline di merging testati. Ad esempio, applicato al forte baseline WUDI-Merging, il punteggio complessivo è aumentato da 38,08 a 38,96.
• Mitigazione della Sensibilità al Dominio: Il merging ingenuo ha degradato severamente le prestazioni di generazione del codice (ad esempio, HumanEval+ è sceso da 33,50 a 20,27 con Weight Averaging). HARC ha parzialmente recuperato questa perdita, dimostrando la sua capacità di preservare le capacità specifiche del dominio correggendo la selezione degli esperti.
• Studi di Ablazione:
  • Struttura dell'Essiana: L'uso dell'intera Essiana ha superato le approssimazioni diagonali o identità, confermando che modellare le dinamiche competitive fuori diagonale è cruciale.
  • Regolarizzazione: Il fattore di smorzamento $\alpha$ (che sopprime i termini fuori diagonale) si è dimostrato più stabile dello scaling diagonale $\gamma$.
  • Composizione dei Dati: La calibrazione utilizzando le risposte corrette filtrate della sorgente (SrcCrt) ha prodotto i risultati migliori. Sorprendentemente, l'uso dei soli prompt ha ottenuto risultati comparabili alle risposte complete, suggerendo che le decisioni di routing iniziali sono fondamentali.
• Efficienza dei Dati: HARC ha raggiunto prestazioni quasi ottimali utilizzando solo il 40% dei campioni di calibrazione, indicando una rapida convergenza delle statistiche del secondo ordine.
• Analisi per Layer: La visualizzazione della divergenza KL ha mostrato che gli errori di routing nei metodi baseline si accumulano esponenzialmente nei layer più profondi. HARC ha soppresso con successo questa tendenza, mantenendo l'allineamento nei layer profondi critici.

5. Significato e Rivendicazioni

Il paper afferma che preservare la logica di routing è essenziale quanto allineare i pesi dei parametri nelle architetture sparse. Affrontando il modo di guasto specifico del routing breakdown, HARC offre una via robusta ed efficiente per consolidare i modelli MoE senza riaddestramento.

Gli autori sottolineano che il loro metodo:

• Non richiede riaddestramento o backpropagation del gradiente attraverso gli esperti.
• Si generalizza attraverso diversi strategie di merging (aritmetica lineare, ottimizzazione) come un leggero step di post-processing.
• Riduce i costi computazionali e le barriere alla distribuzione di modelli capaci, consentendo il riutilizzo affidabile di molteplici modelli di esperti.

Il lavoro conclude che, sebbene le ipotesi di merging lineare esistenti (Linear Mode Connectivity) valgano per i layer densi, esse sono fragili per i router MoE. HARC colma questa lacuna affrontando esplicitamente le dinamiche di gating non lineari inerenti alle architetture Mixture-of-Experts.