Scaling Linear Mode Connectivity and Merging to Billion Parameter Pretrained Transformers

Sintesi Tecnica: Scalare la Connettività Modale Lineare e la Fusione a Transformer Pre-addestrati da Miliardi di Parametri

Problematica
La Connettività Modale Lineare (LMC) postula che reti neurali addestrate indipendentemente possano essere connesse tramite un percorso lineare a bassa perdita nello spazio dei pesi, consentendo una fusione efficiente dei modelli senza ricorrere a ensemble o ri-addestramenti. Tuttavia, i metodi esistenti faticano a scalare questo concetto ai Transformer pre-addestrati con miliardi di parametri. L'ostacolo principale è che i modelli addestrati indipendentemente occupano spesso regioni distanti dello spazio dei pesi a causa di simmetrie che preservano la funzione (ad esempio, permutazioni dei neuroni, rotazioni dei pesi) e differenze nell'inizializzazione o nella dinamica di addestramento. L'interpolazione lineare ingenua tra tali modelli incontra barriere di perdita elevate. Sebbene le tecniche di allineamento precedenti (Weight Matching, Learned Matching) abbiano mostrato potenziale per reti più piccole, esse tipicamente si affidano a un'ottimizzazione unidirezionale (allineando un modello a un riferimento fisso) e spesso non riescono a sfruttare appieno i complessi gruppi di simmetria inerenti alle architetture Transformer, limitandone l'efficacia su larga scala.

Metodologia
Gli autori propongono il Dual Learned Matching (LMC-DM), un framework scalabile progettato per rivelare percorsi lineari a bassa perdita nei grandi Transformer pre-addestrati, ottimizzando congiuntamente le trasformazioni di simmetria per entrambi i modelli terminali.

1. Simmetrie che Preservano la Funzionalità: Il documento formula sistematicamente una vasta famiglia di simmetrie specifiche per i Transformer, estendendosi oltre le semplici permutazioni dei neuroni. Queste includono:

   • Assorbimento della Normalizzazione: Conversione di LayerNorm/RMSNorm con parametri affini in operatori privi di parametri, assorbendo i guadagni (gain) e i bias in strati lineari adiacenti.
   • Simmetria Globale dello Spazio Residuo: Applicazione di rotazioni ortogonali globali e scalature allo stream residuo, che sono preservate da RMSNorm.
   • Permutazioni delle Teste di Attenzione: Riordinamento delle teste di attenzione.
   • Simmetrie dei Circuiti: Trasformazioni invertibili interne nei circuiti Query-Key (QK) e Value-Output (OV) delle teste di attenzione.
   • Simmetrie FFN: Permutazioni dei canali nascosti e simmetrie di riscalamento di tipo GLU.

2. Parametrizzazione: Per rendere queste simmetrie apprendibili tramite discesa del gradiente, gli autori parametrizzano le trasformazioni continue utilizzando:

   • La Trasformazione di Cayley per le matrici ortogonali (efficiente e stabile).
   • La Decomposizione Polare (combinando la trasformazione di Cayley per le parti ortogonali e la decomposizione di Cholesky per le parti simmetriche definite positive) per le matrici invertibili.
   • Le simmetrie discrete (permutazioni di teste/FFN) sono inizializzate tramite Weight Matching e mantenute fisse durante la fase di apprendimento per garantire l'efficienza.

3. Procedura di Apprendimento Duale: A differenza degli approcci precedenti "unidirezionali" che ottimizzano un modello ($\Theta_B$) per corrispondere a un target fisso ($\Theta_A$), LMC-DM introduce un'ottimizzazione bidirezionale. Entrambi i modelli apprendono i propri parametri di trasformazione ($\pi_A$ e $\pi_B$) verso un percorso di interpolazione lineare condiviso:
   $$\Theta_{interp} = \lambda \pi_A(\Theta_A) + (1 - \lambda) \pi_B(\Theta_B)$$
   I parametri di simmetria sono ottimizzati direttamente per minimizzare la perdita lungo il percorso di interpolazione, permettendo a entrambi i modelli di muoversi verso un bacino comune a bassa perdita piuttosto che forzare uno a conformarsi al sistema di coordinate dell'altro.

Contributi Chiave

• Formulazione Esaustiva delle Simmetrie: Il documento definisce una vasta famiglia di trasformazioni che preservano la funzionalità per i Transformer, incluse l'assorbimento della normalizzazione, le rotazioni dello spazio residuo e le specifiche simmetrie dei circuiti di attenzione/FFN.
• Dual Learned Matching: L'introduzione di un framework di ottimizzazione bidirezionale in cui entrambi gli estremi sono trasformati, risultando in un panorama di ottimizzazione meglio condizionato e barriere di interpolazione inferiori rispetto ai metodi unidirezionali.
• Dimostrazione di Scalabilità: L'applicazione con successo della fusione basata su LMC a modelli con fino a miliardi di parametri, una scala precedentemente inesplorata per una connettività quasi priva di barriere.

Risultati Sperimentali
Gli autori valutano LMC-DM su Vision Transformers (ViT-S, ViT-B, ViT-L) e Modelli di Linguaggio (Pythia e OLMo) che vanno da 14M a 6.9B di parametri.

• Barriere di Perdita: Su WikiText, LMC-DM raggiunge barriere di perdita quasi nulle per modelli da 160M di parametri (0.18), superando significativamente Weight Matching (0.27) e il singolo-sided Learned Matching (0.27). Per modelli da 1B di parametri, la barriera rimane bassa (0.23).
• Dominio della Visione: ViT-L mantiene oltre il 69% di accuratezza top-1 su ImageNet lungo l'intero percorso di interpolazione, un miglioramento sostanziale rispetto all'interpolazione ingenua (che subisce un calo significativo) e ai metodi unidirezionali.
• Comportamento di Scalabilità: Mentre i modelli Pythia fino a 6.9B di parametri mostrano piccole barriere di perdita, OLMo-7B mostra barriere più elevate, suggerendo che la LMC nei grandi modelli è influenzata da fattori che vanno oltre la semplice scala, come i dati di addestramento e le dinamiche di ottimizzazione.
• Studi di Ablazione: Rimuovere le simmetrie QK o OV aumenta significativamente le barriere di perdita, in particolare per i ViT. L'approccio di apprendimento duale supera costantemente le varianti unidirezionali, e la parametrizzazione basata su Cayley per le simmetrie ortogonali si dimostra più efficace degli approcci non vincolati o basati sull'esponenziale di matrice.

Significatività e Rivendicazioni
Il documento sostiene di fornire la prima dimostrazione di connettività lineare quasi priva di barriere per modelli di linguaggio con parametri di medie dimensioni (160M) e di mostrare che gli LLM da miliardi di parametri possono essere connessi con solo piccole barriere di perdita se opportunamente allineati.

Gli autori sostengono che i loro risultati suggeriscono che i grandi Transformer pre-addestrati non sono punti isolati nello spazio dei pesi, ma sono "rappresentazioni correlate per simmetria" che possono essere efficacemente allineate e fuse attraverso semplici percorsi lineari. Risolvendo le simmetrie dei parametri tramite l'ottimizzazione duale, il metodo consente la fusione e la composizione affidabile di modelli a scale precedentemente ritenute intrattabili per la connettività modale lineare. Il lavoro colma il divario tra la geometria teorica del panorama di perdita e il riutilizzo pratico dei modelli, offrendo una base per la composizione scalabile di grandi modelli pre-addestrati.