Generalizing the Geometry of Model Merging Through Frechet Averages

1. Enunciazione del Problema

Il model merging (fusione di modelli) mira a combinare più reti neurali addestrate in un unico modello senza ulteriore addestramento end-to-end. Sebbene la semplice media nello spazio dei parametri (ad esempio, Model Soups) funzioni bene per modelli fine-tuned a rango pieno, fallisce significativamente quando applicata a metodi di Fine-Tuning Efficiente dei Parametri (PEFT) come LoRA (Low-Rank Adaptation).

Il problema centrale è la simmetria architetturale:

• I parametri delle reti neurali non sono coordinate canoniche; impostazioni distinte dei parametri possono rappresentare la stessa identica funzione input-output a causa di simmetrie (ad esempio, permutazioni dei neuroni, ridimensionamento o simmetrie di gauge nelle fattorizzazioni a basso rango).
• La media ingenua (euclidea) ignora queste simmetrie. Può "uscire" dall'orbita geometrica corretta, portando a un'interferenza distruttiva in cui il modello fuso collassa o performa male.
• I metodi di allineamento esistenti (come KnOTS) tentano di risolvere questo problema allineando gli aggiornamenti dei task tramite SVD, ma spesso fanno affidamento su scelte specifiche di gauge o euristiche dipendenti dalle coordinate che non rispettano pienamente la geometria sottostante.

2. Metodologia: GeoMerge

Gli autori propongono GeoMerge, un framework che riformula il model merging come il calcolo di una media di Fréchet su una varietà Riemanniana.

Svolta Concettuale di Base

Invece di mediare i parametri nello spazio euclideo, GeoMerge tratta lo spazio dei parametri come una varietà quoziente ($\mathcal{M} = \bar{\mathcal{M}} / G$), dove $G$ è il gruppo di simmetria (ad esempio, $GL(r)$ per LoRA).

• Media Intrinseca: L'obiettivo è trovare un punto sulla varietà quoziente che minimizzi la somma delle distanze geodetiche rispetto ai modelli di input.
• Invarianza alla Simmetria: Operando sul quoziente, la fusione diventa invariante rispetto alla specifica parametrizzazione (gauge) dei modelli di input.

Formulazione Matematica

1. Media di Fréchet: Dati i modelli $\theta_1, \dots, \theta_T$, il modello fuso $\mu^*$ è il minimizzatore del funzionale di Fréchet:
   $$\mu^* \in \arg\min_{\mu \in \mathcal{M}} \sum_{i=1}^T w_i d_{\mathcal{M}}(\mu, \theta_i)^2$$
   dove $d_{\mathcal{M}}$ è la distanza geodetica sulla varietà.
2. Geometria del Quoziente per LoRA:
   • Gli aggiornamenti LoRA sono rappresentati come $\Delta W = GH^\top$. Questa fattorizzazione possiede una simmetria $GL(r)$:$(G, H) \sim (GA^{-1}, HA^\top)$.
   • Gli autori propongono l'uso di una Fattorizzazione Polare ($\Delta W = UBV^\top$) per mappare il problema su un prodotto di varietà di Stiefel ($St$) e varietà Simmetriche Positive Definite (SPD).
   • Definiscono metriche Riemanniane specifiche:
     • Metrica affine-invariante per il fattore SPD ($B$).
     • Metrica canonica di Stiefel per i fattori ortogonali ($U, V$).
3. Algoritmo (GeoMerge del Quoziente):
   L'algoritmo alterna due passaggi per calcolare la media:
   • Allineamento (Minimizzazione dell'Orbita): Per una stima corrente $\mu_t$, si trova la trasformazione di simmetria $G_i$ per ogni input $\theta_i$ che minimizza la distanza rispetto a $\mu_t$. Questo allinea gli input allo stesso "gauge" (spazio orizzontale).
   • Media Intrinseca: Si calcola la mappa logaritmica (vettori tangenti) da $\mu_t$ agli input allineati. Questi vettori sono garantiti essere orizzontali (ortogonali alle orbite di simmetria). La media viene aggiornata muovendosi lungo la geodetica definita dalla somma pesata di questi vettori orizzontali.

3. Contributi Chiave

• Unificazione Geometrica: Il documento stabilisce che il model merging è fondamentalmente un problema di calcolo delle medie di Fréchet su varietà Riemanniane. Unifica varie strategie di fusione sotto questa lente geometrica.
• Fusione Consapevole delle Simmetrie: Fornisce una soluzione principiale al problema del "disallineamento" nel PEFT modellando esplicitamente lo spazio dei parametri come una varietà quoziente, garantendo che la fusione sia invariante rispetto alle scelte di gauge.
• Connessione con Fisher Merging: Gli autori dimostrano che la fusione pesata con Fisher è un'approssimazione trattabile della media di Fréchet sotto la geometria Fisher-Rao (geometria dell'informazione), specificamente quando si utilizza un surrogato quadratico e si ignorano i cambiamenti di covarianza. Tuttavia, notano che la fusione Fisher standard soffre ancora di problemi di simmetria perché il meccanismo di media stesso non è invariante alla simmetria.
• Algoritmo Pratico per LoRA: Derivano un algoritmo concreto per gli adattatori LoRA utilizzando la fattorizzazione polare e metriche Riemanniane specifiche, permettendo un calcolo efficiente delle geodetiche e dell'allineamento senza dati.
• Analisi su Modello Giocattolo: Forniscono una soluzione analitica in forma chiusa per un semplice modello con simmetria di ridimensionamento, dimostrando che la media euclidea ingenua fallisce (collassando a zero) mentre GeoMerge recupera correttamente il predittore intrinseco.

4. Risultati Sperimentali

Gli autori hanno valutato GeoMerge su Llama-3 8B con adattatori LoRA fine-tuned su sei dataset di Inferenza sul Linguaggio Naturale (NLI) (SNLI, MNLI, SICK, QNLI, RTE, SciTail).

• Baseline: Hanno confrontato con KnOTS, un metodo all'avanguardia che allinea gli aggiornamenti LoRA tramite SVD congiunto prima della fusione.
• Prestazioni: GeoMerge ha costantemente superato KnOTS su tutti i dataset in termini di accuratezza normalizzata (il rapporto tra l'accuratezza del modello fuso e quella del miglior specialista).
  • Accuratezza Normalizzata Media: GeoMerge ha raggiunto il 91,88% contro il 90,38% di KnOTS.
• Efficienza: Una scoperta critica è che GeoMerge ottiene prestazioni superiori con un rango effettivo molto più basso.
  • KnOTS tipicamente aumenta il rango del modello fuso a $T \times r$ (dove $T$ è il numero di task).
  • GeoMerge mantiene il rango originale $r$, suggerendo che utilizza la capacità dei pesi in modo più efficiente rispettando la geometria della varietà piuttosto che espandendo lo spazio dei parametri.

5. Significato e Impatto

• Avanzamento Teorico: Il documento sposta il paradigma del model merging da una "media euristica dei parametri" a un'"ottimizzazione geometrica su varietà". Evidenzia che la scelta della geometria (metrica e varietà) è la scelta progettuale più critica nella fusione.
• Robustezza: Affrontando le simmetrie in modo intrinseco, GeoMerge offre una soluzione più robusta per la fusione di modelli specializzati, in particolare nell'era del PEFT dove gli aggiornamenti a basso rango sono lo standard.
• Direzioni Future: Il framework apre la porta all'integrazione di altri concetti geometrici (ad esempio, mediane Riemanniane, medie di Huber) e all'estensione ad altri metodi PEFT o modelli a rango pieno dove sono presenti simmetrie. Suggerisce inoltre che l'"allineamento" non è solo un passaggio di pre-elaborazione, ma una parte integrante del processo di media sullo spazio quoziente.

In sintesi, GeoMerge risolve la fragilità del model merging sotto le simmetrie architetturali trattando il problema come un compito di media geometrica su una varietà quoziente, ottenendo prestazioni all'avanguardia per la fusione LoRA con una maggiore efficienza dei parametri.