How does the optimizer implicitly bias the model merging loss landscape?

Probleemstelling

Modelmerging (het combineren van onafhankelijk getrainde modellen tot één model) is een veelbelovende techniek om de inferentiekosten laag te houden terwijl de capaciteiten van verschillende modellen worden gefuseerd. Bestaande methoden, zoals lineaire interpolatie van gewichten en "task arithmetic", werken goed in de praktijk, maar het is fundamenteel onduidelijk waarom sommige modellen compatibel zijn voor merging terwijl andere falen, zelfs als ze vergelijkbare individuele prestaties behalen.

De bestaande literatuur heeft geïdentificeerd dat "mode connectivity" (het vermogen om minima met een lage loss te verbinden) cruciaal is. Echter, de rol van de optimalisatiedynamica (zoals leerfactor, batchgrootte, gewichtsdecay) in het vormgeven van het verlieslandschap tussen deze oplossingen is slecht begrepen. Er is een gebrek aan inzicht in welke hyperparameters de geometrie van het verlieslandschap zodanig beïnvloeden dat onafhankelijke oplossingen succesvol kunnen worden gemerged.

Methodologie

De auteurs onderzoeken hoe de implicit bias van de optimizer het verlieslandschap beïnvloedt. Ze introduceren en valideren het concept van de effectieve ruis-schaal (effective noise scale, $\tilde{S}$) als de unificerende factor die de succesvolheid van merging bepaalt.

De effectieve ruis-schaal wordt gedefinieerd als:
$$\tilde{S} \propto \frac{\eta}{B(1 - \mu)} \text{tr}(\Sigma_A)$$
Waarbij:

• $\eta$: Leerfactor (Learning rate)
• $B$: Batchgrootte
• $\mu$: Momentum
• $\Sigma_A$: Covariantie van de ruis door data-augmentatie.

De studie omvat uitgebreide experimenten over verschillende architecturen (MLP, ResNet18, DenseNet121, Transformers) en datasets (SVHN, CIFAR-10/100, TinyImageNet, TinyStories voor taalmodellen). Ze testen twee merging-methoden:

1. Lineaire Interpolatie: $\theta_{li} = (1-\alpha)\theta_A + \alpha\theta_B$.
2. Task Arithmetic: Mergen via task vectors ($\tau = \theta_{fine-tuned} - \theta_{base}$).

De auteurs variëren systematisch de hyperparameters om de impact op de "accuracy gain" (de verbetering van het gemerged model ten opzichte van de individuele modellen) te meten. Ze analyseren ook de geometrie van het landschap (vlakheid, barrières) en feature-alignment.

Belangrijkste Bijdragen en Resultaten

1. De Effectieve Ruis-Schaal als Unificerende Factor
De kernvinding is dat de effectieve ruis-schaal ($\tilde{S}$) de compatibiliteit van modellen voor merging bepaalt. Als men de leerfactor en batchgrootte onafhankelijk varieert, is er geen duidelijke trend. Echter, wanneer men de x-as reparametriseert naar $\tilde{S}$, komen alle curves samen in één niet-monotoon verband:

• Te kleine ruis: Leidt tot weinig merging-voordeel (modellen zitten in te scherpe minima).
• Te grote ruis: Leidt tot instabiliteit en falen van merging.
• Optimale (moderate) ruis: Bereikt een "kritiek punt" waar merging het meest effectief is.

2. Impact van Optimizer Componenten
Elk component van de optimizer beïnvloedt $\tilde{S}$ en volgt dezelfde kwalitatieve trend:

• Leerfactor ($\eta$): Een grotere leerfactor (tot een bepaald punt) identificeert oplossingen die beter te mergen zijn, zelfs als de individuele prestaties vergelijkbaar zijn met die van een kleinere leerfactor.
• Gewichtsdecay ($\lambda$): Voor schaal-invariante netwerken (met normalisatielagen) fungeert gewichtsdecay als een regelaar voor de effectieve leerfactor. Een grotere $\lambda$ verhoogt de effectieve ruis en verbetert de merging, vergelijkbaar met een grotere $\eta$.
• Batchgrootte ($B$): Een kleinere batchgrootte injecteert meer gradiëntruis, wat de merging-voordelen vergroot.
• Momentum ($\mu$) en Data Augmentatie: Hoger momentum en data-augmentatie moduleren eveneens de ruis en verbeteren de merging-voordelen.

3. Taalmodellen en Transfer Learning

• Taalmodellen (TinyStories): De bevindingen houden stand voor taalmodellen (GPT). Een grotere leerfactor en gewichtsdecay leiden tot een lagere loss-gain (betere merging).
• Transfer Learning: Bij het finetunen van een pre-getraind model (bijv. CLIP) leidt een grotere leerfactor tot een vlakkere loss-landschap en betere merging. Echter, er is een trade-off: de beste individuele prestaties worden vaak behaald met een gemiddelde leerfactor, terwijl de beste merging-voordelen worden gezien bij een iets hogere leerfactor.

4. Task Arithmetic en Differentiële Ruimtes
Bij task arithmetic (mergen van modellen getraind op verschillende taken) verandert de geometrie van het landschap afhankelijk van de initialisatie.

• Met een pre-getrainde initialisatie leidt een grotere leerfactor tot een robuustere (vlakker) landschap.
• Het mergen van modellen getraind op verschillende taken met een "moderate" ruis geeft de beste prestaties, maar vermindert de compatibiliteit met modellen die met een heel andere ruis zijn getraind.

5. Permutatie Symmetrie en Feature Alignement
De auteurs tonen aan dat een grotere effectieve ruis de "rebasin"-processen (het aligneren van permutatiesymmetrieën tussen modellen) vereenvoudigt. Interessant genoeg correleert een hogere ruis met minder feature-alignment (diversere features) tussen de takken, maar wel met hogere merging-voordelen. Dit suggereert dat merging profiteert van een zekere mate van feature-diversiteit die door de optimale ruis wordt gegenereerd.

Significantie en Conclusie

Dit paper verschuift het perspectief van modelmerging van een puur geometrisch probleem naar een dynamisch optimalisatieprobleem. De belangrijkste implicaties zijn:

1. Unificatie: Het concept van de "effectieve ruis-schaal" biedt een enkele metriek om te voorspellen of twee onafhankelijk getrainde modellen succesvol kunnen worden gemerged.
2. Stuurbaarheid: Het suggereert dat practitioners de training dynamica kunnen manipuleren (bijv. door de leerfactor of batchgrootte te kiezen) om modellen te produceren die niet alleen sterk zijn, maar ook "inherently mergeable".
3. Praktische Richtlijnen: Voor het beste merging-resultaat moeten modellen niet noodzakelijk met de kleinste mogelijke ruis worden getraind (wat vaak leidt tot scherpe minima), maar met een moderate, geoptimaliseerde ruis.
4. Toekomstig Onderzoek: De bevindingen openen de deur voor het systematisch afstemmen van hyperparameters en architecturale keuzes om "model soups" te creëren die zonder extra inferentiekosten superieure prestaties leveren.

Kortom, de paper toont aan dat de "ruis" in het trainingsproces geen bijproduct is dat moet worden geminimaliseerd, maar een cruciale eigenschap die de global loss landscape vormgeeft en bepaalt of onafhankelijke oplossingen kunnen worden samengevoegd.