Scaling Linear Mode Connectivity and Merging to Billion Parameter Pretrained Transformers

Technische Samenvatting: Schalen van Linear Mode Connectivity en Merging naar miljarden-parameters pre-trained Transformers

Probleemstelling
Linear Mode Connectivity (LMC) stelt dat onafhankelijk getrainde neurale netwerken verbonden kunnen worden via een pad met een lage loss in de gewichtsruimte, wat efficiënte model-merging mogelijk maakt zonder ensembles of hertraining. Echter, bestaande methoden hebben moeite om dit concept te schalen naar miljarden-parameters pre-trained Transformers. De primaire hindernis is dat onafhankelijk getrainde modellen vaak in verre regio's van de gewichtsruimte verblijven vanwege functie-behoudende symmetrieën (bijv. neuron-permutaties, gewichtsrotaties) en verschillen in initialisatie of trainingsdynamiek. Naïeve lineaire interpolatie tussen dergelijke modellen stuit op hoge loss-barrières. Hoewel eerdere alignment-technieken (Weight Matching, Learned Matching) veelbelovend zijn gebleken voor kleinere netwerken, vertrouwen ze doorgaans op eenzijdig optimalisatie (het alignen van één model met een vaste referentie) en slagen ze er vaak niet in om de complexe symmetriegroepen die inherent zijn aan Transformer-architecturen volledig uit te buiten, wat hun effectiviteit op schaal beperkt.

Methodologie
De auteurs stellen Dual Learned Matching (LMC-DM) voor, een schaalbaar framework dat is ontworpen om paden met een lage loss in grote pre-trained Transformers te ontsluiten door gezamenlijk de symmetrie-transformaties van beide eindmodellen te optimaliseren.

1. Functie-behoudende Symmetrieën: Het artikel formuleert systematisch een brede familie van symmetrieën specifiek voor Transformers, die verder gaan dan eenvoudige neuron-permutaties. Deze omvatten:

   • Normalization Absorption: Het omzetten van LayerNorm/RMSNorm met affine parameters naar parameter-vrije operatoren door gains en biases te absorberen in aangrenzende lineaire lagen.
   • Global Residual-Space Symmetry: Het toepassen van globale orthogonale rotaties en schalingen op de residual stream, die behouden blijven door RMSNorm.
   • Attention Head Permutations: Het herordenen van attention heads.
   • Circuit Symmetries: Interne inverteerbare transformaties binnen de Query-Key (QK) en Value-Output (OV) circuits van attention heads.
   • FFN Symmetries: Verbergen van kanaal-permutaties en GLU-stijl rescaling symmetrieën.

2. Parameterisatie: Om deze symmetrieën leerbaar te maken via gradiëntafdaling, parameteriseren de auteurs continue transformaties met behulp van:

   • Cayley Transform voor orthogonale matrices (efficiënt en stabiel).
   • Polar Decomposition (combinatie van Cayley voor orthogonale delen en Cholesky voor symmetrische positief-definiete delen) voor inverteerbare matrices.
   • Discrete symmetrieën (head/FFN permutaties) worden geïnitialiseerd via Weight Matching en gehouden als vast tijdens de leerfase om efficiëntie te waarborgen.

3. Dual Learning Procedure: In tegen tegenstelling tot eerdere "eenzijdige" benaderingen die één model ($\Theta_B$) optimaliseren om een vaste doelstelling ($\Theta_A$) te matchen, introduceert LMC-DM een bidirectionele optimalisatie. Beide modellen leren hun eigen transformatieparameters ($\pi_A$ en $\pi_B$) richting een gedeeld lineair interpolatiepad:
   $$\Theta_{interp} = \lambda \pi_A(\Theta_A) + (1 - \lambda) \pi_B(\Theta_B)$$
   De symmetrie-parameters worden direct geoptimaliseerd om de loss langs het interpolatiepad te minimaliseren, waardoor beide modellen naar een gedeelde low-loss basin bewegen in plaats van één model te dwingen zich aan het coördinatensysteem van de ander te conformeren.

Belangrijkste Bijdragen

• Omvattende Symmetrie-formulering: Het artikel definieert een brede familie van functie-behoudende transformaties voor Transformers, inclus\nslief normalization absorption, residual-space rotaties en specifieke attention/FFN circuit-symmetrieën.
• Dual Learned Matching: De introductie van een bidirectioneel optimalisatieframework waarbij beide eindpunten worden getransformeerd, wat resulteert in een beter geconditioneerd optimalisatie-landschap en lagere interpolatie-barrières vergeleken met eenzijdige methoden.
• Demonstratie van Schaalbaarheid: De succesvolle toepassing van LMC-gebaseerde merging op modellen met tot miljarden parameters, een schaal die voorheen onverkend was voor bijna barrière-vrije connectiviteit.

Experimentele Resultaten
De auteurs evalueren LMC-DM op Vision Transformers (ViT-S, ViT-B, ViT-L) en Language Models (Pythia en OLMo) variërend van 14M tot 6.9B parameters.

• Loss Barrières: Op WikiText bereikt LMC-DM bijna nul loss-barrières voor 160M parameter modellen (0.18), wat significant beter presteert dan Weight Matching (0.27) en single-sided Learned Matching (0.27). Voor 1B parameter modellen blijft de barrière laag (0.23).
• Visueel Domein: ViT-L behoudt meer dan 69% ImageNet top-1 accuraatheid over het gehele interpolatiepad, een substantiële verbetering ten opzichte van naïeve interpolatie (die significant daalt) en single-sided methoden.
• Schaalgedrag: Terwijl Pythia-modellen tot 6.9B parameters kleine loss-barrières vertonen, laat OLMo-7B hogere barrières zien, wat suggereert dat LMC in grote modellen wordt beïnvloed door factoren buiten alleen schaal, zoals trainingsdata en optimalisatiedynamiek.
• Ablatie Studies: Het verwijderen van QK of OV symmetrieën verhoogt de loss-barrières aanzienlijk, met name voor ViTs. De dual-learning benadering presteert consistent beter dan single-sided varianten, en Cayley-gebaseerde parameterisatie voor orthogonale symmetrieën blijkt effectiever dan ongecontroleerde of matrix-exponentiële benaderingen.

Betekenis en Claims
Het artikel claimt de eerste demonstratie te bieden van bijna barrière-vrije lineaire connectiviteit voor taalmodellen met gemiddelde grootte (160M) en te laten zien dat miljarden-parameters LLM's verbonden kunnen worden met slechts kleine loss-barrières wanneer ze correct zijn gealigneerd.

De auteurs stellen dat hun resultaten suggereren dat grote pre-trained Transformers geen geïsoleerde punten in de gewichtsruimte zijn, maar "symmetrie-gerelateerde representaties" die effectief gealigneerd en samengevoegd kunnen worden via eenvoudige lineaire paden. Door parameter-symmetrieën op te lossen via duale optimalisatie, stelt de methode betrouwbare model-merging en compositie mogelijk op schalen die voorheen als onhandelbaar voor linear mode connectivity werden beschouwd. Het werk overbrugt de kloof tussen de theoretische geometrie van het loss-landschap en praktische model-hergebruik, en biedt een fundament voor de schaalbare compositie van grote pre-trained modellen.