Transport and Merge: Cross-Architecture Merging for Large Language Models

1. Il Problema

I Large Language Models (LLM) ad alte prestazioni richiedono enormi risorse computazionali e dati di addestramento massicci, rendendo difficile il loro utilizzo in scenari reali con vincoli di risorse (es. lingue a bassa risorsa come il malese o il cantonese, o domini specializzati come medicina e finanza).
Esistono modelli più piccoli e adattati a queste risorse limitate, ma spesso mancano delle conoscenze profonde possedute dai modelli "sorgente" ad alta risorsa.
Le tecniche di trasferimento di conoscenza esistenti, come la distillazione, richiedono spesso addestramenti costosi o dati aggiuntivi. Inoltre, la maggior parte dei metodi di fusione di modelli (Model Merging) attuali presuppone che i modelli sorgente e target condividano la stessa architettura, permettendo operazioni dirette nello spazio dei pesi. Questo limita drasticamente la possibilità di trasferire conoscenza da un LLM grande (es. LLaMA-3 8B) a un modello target eterogeneo e più piccolo (es. un modello da 1B con architettura diversa), poiché non è possibile allineare direttamente i neuroni.

2. Metodologia Proposta

Gli autori propongono un framework di fusione Cross-Architecture basato sul Trasporto Ottimale (Optimal Transport - OT). L'obiettivo è trasferire i parametri da un modello sorgente ad alta risorsa a un modello target a bassa risorsa con architettura diversa, utilizzando solo un piccolo insieme di dati di calibrazione.

Il processo si articola in tre fasi principali:

A. Allineamento nello Spazio delle Attivazioni

Poiché non è possibile allineare direttamente i pesi di architetture diverse, il metodo allinea prima le attivazioni interne (feature) dei due modelli.

1. Estrazione delle Attivazioni: Si passa un piccolo dataset $D$ attraverso entrambi i modelli (sorgente e target) per registrare le attivazioni intermedie nei vari strati.
2. Corrispondenza Neurale tramite OT: Per ogni strato target $\ell$ e strato sorgente $m$, si calcola una matrice di costo basata sulla correlazione tra i canali di attivazione. Si risolve un problema di Trasporto Ottimale regolarizzato entropicamente per trovare una matrice di trasporto $Q^{\ell m}$. Questa matrice definisce come le attivazioni del modello sorgente possono essere ricombinate per corrispondere a quelle del modello target, inferendo relazioni "neurone-a-neurone" anche tra architetture diverse.
3. Allineamento a Livello di Strato: Si aggregano le corrispondenze a livello di feature per calcolare un costo tra gli strati e si risolve un secondo problema OT per ottenere una matrice globale $P$ che mappa gli strati del modello sorgente a quelli del target.

B. Fusione dei Pesi (Weight Fusion)

Le relazioni di trasporto trovate nello spazio delle attivazioni vengono "sollevate" (lifted) allo spazio dei pesi per fondere i parametri.

• Operatore di Trasporto: Le matrici di trasporto $Q$ agiscono come operatori di trasformazione che mappano i pesi del modello sorgente nel sistema di coordinate del modello target.
• Sostituzione Selettiva (Top-k): Per garantire robustezza ed evitare interferenze, la fusione non avviene su tutti i parametri. Si selezionano solo i top-k neuroni più attivi nel modello target sul dataset di calibrazione.
• Formula di Fusione: I pesi fusi sono una combinazione del peso originale del target e del residuo trasportato dal sorgente, pesato da un coefficiente $\alpha$ e mascherato sui neuroni selezionati:
  $$W^{fused} = W^{base} + \alpha \cdot M \odot (\text{Trasporto}(W^{sorgente}) - W^{base})$$

C. Adattamento Residuo Congelato (Residual-Frozen Adaptation)

Dopo la fusione, viene eseguita una fase di adattamento leggera:

• I parametri residui trasportati (il contributo del modello sorgente) vengono congelati.
• Vengono aggiornati solo i parametri base originali del modello target.
• Questo approccio permette al modello di adattarsi alla distribuzione del compito specifico senza sovrascrivere la conoscenza trasferita, agendo come un'inizializzazione stabile.

3. Contributi Chiave

1. Fusione Cross-Architettura Diretta: Il primo framework che permette la fusione diretta dei pesi tra modelli LLM con architetture eterogenee (diversi numeri di parametri, diverse dimensioni di embedding o strutture) senza distillazione.
2. Trasporto Ottimale per l'Allineamento Neurale: L'uso innovativo dell'OT per inferire corrispondenze strutturate tra neuroni di modelli diversi basandosi sulle statistiche delle attivazioni, trasformando queste corrispondenze in operatori di fusione dei pesi.
3. Efficienza dei Dati: Il metodo richiede solo un piccolo insieme di dati di calibrazione (non necessita di ri-addestramento su grandi corpus) per trasferire conoscenza.
4. Interpretazione Teorica: Dimostrano che il trasporto dei pesi è equivalente a una mappatura delle feature nello spazio di rappresentazione del modello sorgente, applicazione della trasformazione lineare e ritorno allo spazio target.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su diverse lingue a bassa risorsa e domini specializzati:

• Lingue a Bassa Risorsa: Il metodo è stato testato su Malese, Indonesiano, Cantonese e Thai, fondendo modelli sorgente (es. LLaMA-3 8B o Qwen2.5-7B) con modelli target da 1B.
  • Risultati: La fusione ha portato a miglioramenti consistenti rispetto ai modelli target base. Ad esempio, sul benchmark MalayMMLU, l'approccio con adattamento ha migliorato le prestazioni in tutte le categorie (Umanità, STEM, Scienze Sociali) con guadagni fino al +6.51%.
  • Confronto: Supera sia l'addestramento SFT (Supervised Fine-Tuning) che la distillazione tradizionale in termini di performance normalizzata, specialmente in scenari cross-dominio.
• Domini Specializzati (Finanza e Medicina): La fusione ha trasferito efficacemente conoscenze di dominio da modelli generalisti a modelli specializzati, migliorando i risultati su benchmark come MMLU (Business Ethics, Accounting, Anatomy).
• Robustezza: Il metodo si dimostra robusto rispetto alla scelta del modello sorgente (funziona bene sia con LLaMA che con Qwen) e al coefficiente di fusione $\alpha$, purché sia in un range moderato (0.05-0.15).
• Capacità Generali: La fusione non degrada le capacità generali del modello target, mantenendo le prestazioni su task generici (ragionamento, senso comune) simili al modello originale.

5. Significato e Impatto

Questo lavoro risolve una delle principali barriere all'adozione di LLM in contesti reali con risorse limitate: la difficoltà di trasferire conoscenza da modelli grandi a modelli piccoli quando le architetture non sono compatibili.

• Praticità: Offre un'alternativa efficiente alla distillazione, riducendo drasticamente i requisiti computazionali e di dati.
• Accessibilità: Permette di potenziare modelli per lingue minoritarie o domini di nicchia utilizzando modelli foundation open-source esistenti, senza bisogno di ri-addestrarli da zero.
• Fondamentale: Stabilisce un nuovo paradigma per la fusione di modelli, dimostrando che l'allineamento nello spazio delle attivazioni tramite Trasporto Ottimale può guidare operazioni dirette nello spazio dei pesi, aprendo la strada a tecniche di "Model Soups" tra architetture eterogenee.

In sintesi, "Transport and Merge" fornisce un meccanismo principiato e pratico per colmare il divario tra modelli LLM ad alta risorsa e modelli a bassa risorsa, rendendo l'intelligenza artificiale più accessibile ed efficiente per applicazioni globali e specializzate.