Riepilogo Tecnico: Fusione di Modelli tramite Retrofitting

Enunciato del Problema

La fusione di modelli mira a combinare reti neurali addestrate in modo indipendente in un unico modello senza riaddestramento, offrendo una soluzione per scenari in cui il riaddestramento è proibitivamente costoso o non fattibile a causa di vincoli sulla privacy dei dati (ad esempio, Federated Learning). Tuttavia, questo compito è complicato dalla divergenza rappresentazionale causata dall'invarianza permutazionale, dall'inizializzazione casuale e da distribuzioni eterogenee dei dati di addestramento.

I metodi esistenti, come OTFusion (Singh & Jaggi, 2020) e Git Re-Basin (Ainsworth et al., 2023), si basano sulla congettura della Connessione Modale Lineare (LMC), che postula che le reti addestrate in modo indipendente risiedano nello stesso bacino a bassa perdita. Sebbene la LMC valga per modelli addestrati su dati identici, studi precedenti indicano che anche modelli addestrati sugli stessi dati possono apprendere caratteristiche interne divergenti. Di conseguenza, i metodi di fusione esistenti faticano significativamente in ambienti zero-shot (fusione senza ulteriore addestramento) e sotto regimi di dati non-IID (non Indipendenti e Identicamente Distribuiti), collassando spesso verso prestazioni quasi casuali. Inoltre, gli approcci attuali spesso mancano di generalità tra le architetture (ad esempio, faticano con i Transformer) e si affidano a modelli di base che potrebbero non essere completamente ottimizzati, offuscando il vero potenziale degli algoritmi di fusione.

Metodologia

Gli autori propongono una famiglia di algoritmi di fusione centrata sui neuroni che riformula la fusione di modelli come un problema di corrispondenza rappresentazionale principiato. Invece di semplicemente mediare i pesi o trovare matrici di permutazione, il metodo tratta la fusione come un processo in due fasi: raggruppare i neuroni dai modelli genitori in rappresentazioni target e quindi addestrare le sotto-reti del modello fuso per approssimare questi target.

Framework Principale

1. Costo Rappresentazionale: L'obiettivo è definito dalla distanza $L_2$ al quadrato tra le attivazioni dei neuroni del modello di base e i neuroni più vicini nel modello fuso, ponderata per punteggi di importanza dei neuroni (punteggi di attribuzione). Questo penalizza la mancata rappresentazione delle caratteristiche salienti.
2. Decomposizione (Teorema 1): Gli autori introducono un vettore target ausiliario $T$ per decomporre il costo rappresentazionale in due componenti interpretabili:
   • Errore di Raggruppamento: Misura quanto bene i neuroni del modello di base si raggruppano attorno ai centroidi ponderati per importanza ($T_k$).
   • Errore di Approssimazione: Misura quanto bene i neuroni di output del modello fuso approssimano questi centroidi.
3. Algoritmo in Due Fasi:
   • Fase 1 (Raggruppamento): I neuroni dai modelli di base vengono raggruppati per minimizzare l'errore di raggruppamento. Ciò produce centroidi ponderati per importanza che servono come target.
   • Fase 2 (Retrofitting): I pesi del modello fuso al livello corrente vengono ottimizzati (tramite soluzione in forma chiusa per i livelli lineari o SGD per i livelli non lineari) per minimizzare l'errore di approssimazione rispetto a questi target. Questo processo viene eseguito dal basso verso l'alto attraverso i livelli della rete.

Varianti dell'Algoritmo

Il paper introduce due algoritmi principali basati sulla strategia di raggruppamento:

• Fusione Ungherese (HF): Progettata per modelli di dimensioni uguali con corrispondenza uno-a-uno a livello di livello. Utilizza l'algoritmo ungherese per trovare l'assegnazione ottimale.
• Fusione K-means (KF): Progettata per il caso generale, inclusi modelli di larghezze e profondità arbitrarie. Utilizza il clustering euristico K-means per raggruppare i neuroni.

Entrambi gli algoritmi hanno varianti Lineari (soluzione in forma chiusa per i livelli affini) e Basate sul Gradiente (SGD/Adam per i livelli non lineari).

Innovazioni Chiave

• Attribuzione dei Neuroni: A differenza dei lavori precedenti, questo approccio incorpora punteggi di importanza dei neuroni (ad esempio, Conductance, DeepLIFT) per orientare l'allineamento verso caratteristiche salienti, migliorando le prestazioni in ambienti eterogenei.
• Flessibilità Architettonica: Il metodo è applicabile a qualsiasi architettura modularizzabile come un Grafo Aciclico Diretto (DAG) di livelli, validato empiricamente su VGG, ResNet e Vision Transformers (ViT).
• Capacità Zero-Shot: Ottimizzando esplicitamente l'approssimazione delle rappresentazioni raggruppate, il metodo ottiene prestazioni solide senza l'esteso fine-tuning richiesto dai metodi precedenti.

Risultati Sperimentali

Gli autori hanno valutato i loro metodi in tre regimi di addestramento: Sharded (insiemi di classi disgiunti), Non-IID (distribuzioni di classi sbilanciate) e Dataset Completo.

• Regimi Sharded/Non-IID: In ambienti non-IID estremi (ad esempio, divisioni in 6 classi su CIFAR-100), i metodi esistenti come OTFusion e Git Re-Basin sono collassati verso una precisione quasi casuale (spesso <5%). Al contrario, KF-Gradient e HF-Gradient hanno mantenuto una robusta precisione zero-shot (ad esempio, 34,7% su CIFAR-100 per divisioni in 6 classi contro <1,5% per le baseline).
• Regimi Dataset Completo: Anche quando i modelli sono addestrati sull'intero dataset, i metodi proposti superano le baseline nella fusione zero-shot. Ad esempio, su CIFAR-100, KF-Gradient ha raggiunto una precisione zero-shot del 63,0% rispetto al 4,3% di Transformer OTFusion. Dopo il fine-tuning, i metodi proposti si sono avvicinati costantemente alle prestazioni dell'ensemble mantenendo i costi di inferenza di un singolo modello.
• Compressione ed Eterogeneità: Il metodo ha fuso con successo modelli di diverse larghezze e ha compresso modelli più grandi (ResNet-50) in modelli più piccoli (ResNet-18) con una precisione significativamente superiore alla Distillazione della Conoscenza (KD) standard quando i dati erano limitati.
• Robustezza: L'approccio ha dimostrato robustezza rispetto alle dimensioni variabili del dataset di fusione, con l'aumento dei dati che ha efficacemente mitigato la necessità di grandi set di fusione.

Significato e Affermazioni

Il paper afferma di colmare tre lacune chiave nella ricerca precedente:

1. Riproducibilità: Fornire una reimplementazione open-source flessibile che facilita il benchmarking sistematico tra le architetture.
2. Qualità del Modello di Base: Dimostrare che i metodi di fusione possono produrre miglioramenti significativi anche quando i modelli di base sono ben addestrati, sfidando l'idea che i risultati di fusione scadenti siano dovuti esclusivamente a baseline deboli.
3. Dati Eterogenei: Mostrare che i metodi esistenti falliscono nella fusione zero-shot in condizioni non-IID, mentre l'approccio centrato sui neuroni proposto integra con successo modelli addestrati su distribuzioni di dati disgiunte o sbilanciate.

Gli autori sottolineano che il loro approccio è il primo a incorporare con successo i punteggi di importanza dei neuroni nella fusione di modelli, portando a miglioramenti coerenti rispetto alle baseline esistenti, in particolare negli scenari zero-shot e non-IID dove altri metodi falliscono. Posizionano il metodo come una soluzione pratica per l'apprendimento federato e l'aggregazione di modelli che preserva la privacy, dove il riaddestramento non è un'opzione. Pur riconoscendo che le varianti basate sul gradiente sono computazionalmente più costose delle baseline leggere, gli autori sostengono che questo sovraccarico è giustificato dalla capacità di produrre modelli fusi vitali in regimi in cui i concorrenti producono risultati inutilizzabili.