InfiGFusion: Graph-on-Logits Distillation via Efficient Gromov-Wasserstein for Model Fusion

Riepilogo Tecnico: InfiGFusion

Enunciato del Problema
I recenti sforzi per fondere Modelli Linguistici su larga scala (LLM) eterogenei in sistemi unificati hanno fatto affidamento prevalentemente su metodi di distillazione basati sui logit. Sebbene tali approcci mantengano l'efficienza di inferenza, trattano tipicamente le dimensioni del vocabolario in modo indipendente, allineando le distribuzioni dei token tramite obiettivi come la divergenza KL o la distanza di Wasserstein. Questo allineamento token per token trascura le dipendenze semantiche—le interazioni tra dimensioni e i pattern di co-attivazione che codificano i processi di ragionamento interno di un modello. Di conseguenza, i metodi esistenti faticano ad allineare modelli con comportamenti di generazione divergenti, in particolare in compiti che richiedono ragionamento multi-step, inferenza relazionale e allineamento semantico fine.

Metodologia: InfiGFusion
Gli autori propongono InfiGFusion, un framework di fusione consapevole della struttura che modella esplicitamente queste dipendenze semantiche. L'innovazione centrale è la funzione di perdita Graph-on-Logit Distillation (GLD), che opera come segue:

1. Costruzione del Grafo dei Logit: Invece di trattare i logit come vettori indipendenti, il metodo costruisce un grafo globale di co-attivazione. Per ogni output, vengono mantenuti i top-$k$ logit per garantire la sparsità. I prodotti esterni di questi logit attraverso le posizioni della sequenza vengono aggregati per formare un grafo in cui i nodi rappresentano i canali del vocabolario e gli archi quantificano le loro attivazioni congiunte (somiglianza).
2. Allineamento Consapevole della Struttura: Per allineare i grafi dei logit dei modelli sorgente con un modello pivot, il framework impiega la distanza Gromov-Wasserstein (GW). GW è scelta per la sua capacità di far corrispondere distribuzioni basandosi sulla somiglianza strutturale interna piuttosto che sulla corrispondenza diretta dei nodi.
3. Approssimazione Efficiente: Il calcolo standard della distanza GW ha una complessità proibitiva di $O(n^4)$, rendendolo non fattibile per vocabolari di grandi dimensioni ($n > 10^5$). Gli autori introducono una nuova approssimazione in forma chiusa basata sull'ordinamento che riduce la complessità a $O(n \log n)$.
   • Compressione delle Caratteristiche del Grafo: La matrice di somiglianza viene compressa in caratteristiche scalari a livello di nodo (specificamente, medie delle righe o caratteristiche di tipo grado).
   • Corrispondenza Basata sull'Ordinamento: L'allineamento viene eseguito ordinando queste caratteristiche e facendo corrispondere la voce $i$-esima più grande della sorgente con la voce $i$-esima più grande del target.
   • Garanzie Teoriche: Il documento fornisce un limite di errore di approssimazione dimostrabile di $\frac{n-1}{n^2} + \frac{m-1}{m^2}$, che decade come $O(1/n)$. Inoltre, gli autori dimostrano che la perdita guidata da GW ha una costante di Lipschitz significativamente più bassa rispetto alle perdite KL e di Wasserstein standard, suggerendo limiti di generalizzazione più stretti e un addestramento più stabile.
4. Obiettivo Unificato: L'obiettivo finale di addestramento combina la perdita GLD consapevole della struttura con una perdita di distillazione universale dei logit a grana grossa (ULD, per l'allineamento della distribuzione a livello di token) e segnali standard di Fine-Tuning Supervisionato (SFT).

Contributi Chiave

• Primo Framework di Fusione Consapevole della Struttura: InfiGFusion è il primo metodo a modellare le dipendenze semantiche tramite grafi di logit a livello di caratteristica e ad allinearli utilizzando una nuova perdita GLD, andando oltre l'allineamento indipendente dei token.
• Approssimazione GW Scalabile: Gli autori propongono un'approssimazione in forma chiusa per la distanza Gromov-Wasserstein che riduce la complessità computazionale da $O(n^4)$ a $O(n \log n)$ con limiti di errore dimostrabili, rendendo fattibile l'allineamento basato su grafi per vocabolari su scala LLM.
• Superiorità Empirica nel Ragionamento: Esperimenti estesi su 11 benchmark dimostrano che InfiGFusion supera costantemente le baseline di fusione all'avanguardia (ad es. MiniLogit, FuseLLM, FuseChat) e le baseline SFT.

Risultati
Il metodo è stato valutato su 11 benchmark che coprono ragionamento, codifica, matematica e seguito delle istruzioni, fondendo molteplici modelli sorgente (Qwen2.5, Mistral, Qwen2.5-Coder) in un modello pivot Phi-4.

• Miglioramenti nel Ragionamento Complesso: InfiGFusion ha mostrato miglioramenti significativi in compiti che richiedono inferenza multi-step e relazionale. In particolare, ha ottenuto un miglioramento di +35.6 punti sull'aritmetica multi-step e un miglioramento di +37.06 punti sul giudizio causale rispetto alle baseline SFT.
• Prestazioni Complessive: Il modello ha superato il modello pivot e i modelli sorgente in diversi domini, ottenendo un punteggio medio di 83.85 rispetto ai 79.54 del pivot.
• Efficienza: Nonostante l'allineamento strutturale, il metodo ha mantenuto un'efficienza competitiva, richiedendo solo 195 ore GPU per l'addestramento, significativamente meno rispetto alle pipeline di fusione multi-step come FuseChat (650 ore GPU).
• Studi di Ablazione: Gli esperimenti hanno confermato che sia il componente ULD (allineamento a grana grossa) che il componente GLD (allineamento strutturale) sono necessari, con GLD che guida specificamente i guadagni nei compiti di ragionamento.

Significato
Il documento afferma che InfiGFusion affronta una lacuna critica nella fusione dei modelli: l'incapacità degli obiettivi a livello di token di catturare le strutture relazionali alla base del ragionamento del modello. Allineando esplicitamente la "forma" dello spazio dei logit e non solo le distribuzioni marginali, il framework consente l'integrazione di LLM eterogenei senza aumentare i costi di inferenza. Gli autori ipotizzano che questo allineamento che preserva la struttura sia essenziale per avanzare le capacità collettive degli LLM, in particolare in scenari di ragionamento complesso dove la coerenza logica e l'inferenza multi-hop sono fondamentali. Il lavoro suggerisce che i progressi futuri nella fusione dei modelli dipendono dal superare il semplice abbinamento superficiale dei token per preservare le dipendenze semantiche interne che guidano un ragionamento coerente.