Technische Samenvatting: Het Mechanisme van Weak-to-Strong Generalisatie

Probleemstelling

Weak-to-Strong (W2S) generalisatie beschrijft het fenomeen waarbij een sterk, voorgetraind model, wanneer het wordt gefinetuned op de output van een zwakker, taakgespecialiseerd model, de prestaties van het zwakke model op die specifieke taak kan overtreffen terwijl het zijn bredere capaciteiten behoudt. Hoewel empirische studies (bijvoorbeeld Burns et al., 2023) hebben aangetoond dat volledige finetuning (feature learning) superieur is aan lineaire probing (vaste representaties), zijn bestaande theoretische analyses grotendeels beperkt tot abstracte kaders of lineaire/willekeurige feature-modellen die feature learning niet afleiden uit het optimalisatieproces zelf.

De centrale open vraag die door dit artikel wordt aangepakt, is of multi-stap Stochastic Gradient Descent (SGD) bewezen succesvol kan zijn in W2S feature learning – specifiek het leren van de doel-featuresrichting terwijl diverse voorgetrainde capaciteiten behouden blijven – in een setting die meer lijkt op voorgetrainde Large Language Models (LLM's).

Methodologie en Probleemsetting

De auteurs formaliseren W2S in de context van reward-model learning met behulp van twee-laags neurale netwerken. De opzet omvat:

1. Additief Reward Model: De ware reward $r^*(x)$ wordt gemodelleerd als een som van $K$ single-index componenten, $r^*(x) = \sum_{k=1}^K \pi_k \sigma^*_k(\theta_k^\top x)$, waarbij elke taak $k$ overeenkomt met een specifieke feature-richting $\theta_k$.
2. Voorgetrainde Subruimtes: In overeenstemming met empirische bevindingen dat LLM's concepten coderen in laag-dimensionale subruimtes, wordt aangenomen dat het sterke model voorgetrainde kennis heeft van taak-specifieke subruimtes $V_k$ met dimensie $s \ll d$. Het sterke model kent de subruimte $V_k$ die $\theta_k$ bevat, maar niet de exacte richting $\theta_k$ zelf.
3. Zwak Model: Een klein, twee-laags neurale netwerk dat gespecialiseerd is in een doeltaak $\kappa$. Het is getraind om de reward-component $r^*_\kappa$ te benaderen en bezit feature-alignment met $\theta_\kappa$.
4. Sterk Model: Een groter twee-laags netwerk met neuronen opgesplitst in groepen die overeenkomen met de $K$ taken. Het wordt geïnitieerd met neuronen die gelokaliseerd zijn in hun respectievelijke subruimtes.
5. Trainingsdynamiek:
   • W2S Algoritme: Het sterke model wordt getraind via online SGD op inputs getrokken uit zijn eigen generatieve distributie (geconcentreerd op de doel-subruimte $V_\kappa$) en gesuperviseerd door de output van het zwakke model. Cruciaal ondergaat de output van het zwakke model een niet-lineaire transformatie ($z \mapsto z \exp(z)$) voordat deze wordt gebruikt als label.
   • Standaard SFT (Basislijn): Ter vergelijking wordt het sterke model getraind via online SGD op isotrope inputs ($x \sim N(0, I_d)$) met directe labels van de ware reward-functie $r^*_\kappa$.

Belangrijkste Bijdragen en Theoretische Resultaten

1. Feature Elicitatie via Multi-stap SGD (Stelling 1)

Het artikel biedt de eerste multi-stap SGD convergentie-garantie voor W2S in het feature-learning regime.

• Mechanisme: De niet-lineaire transformatie van het leraarsignaal verlaagt de effectieve "generatieve exponent" van het leerverhaal. Hoewel de doel-linkfunctie $\sigma^*_\kappa$ een hoge informatieve exponent $p > 2$ kan hebben, stelt de transformatie het sterke model in staat om de feature-richting $\theta_\kappa$ te leren met een steekproefcomplexiteit die afhankelijk is van de generatieve exponent (die 2 wordt voor even polynomen).
• Steekproefcomplexiteit: Het sterke model bereikt alignment met de doel-feature $\theta_\kappa$ met een steekproefcomplexiteit van $\tilde{O}(s^{3/2})$, waarbij $s$ de dimensie van de taak-subruimte is. Dit is aanzienlijk lager dan de $\tilde{O}(d^{p/2} s^{(p-2)/2})$ complexiteit die nodig is voor het zwakke model om dezelfde feature vanaf nul te leren, en lager dan de $\tilde{O}(d^{p-1})$ complexiteit van standaard SFT onder isotrope sampling.
• Subruimte Prior: Het resultaat is afhankelijk van het feit dat de input-distributie van het sterke model geconcentreerd is op de doel-subruimte $V_\kappa$ (geparametriseerd door $\lambda_\kappa \approx 1$), wat effectief de leerdimensie verlaagt van $d$ naar $s$.

2. Feature Behoud (Propositie 2)

Een cruciale bijdrage is het bewijs dat W2S voorgetrainde features behoudt voor niet-doel-taken ($k \neq \kappa$).

• Impliciete Regularisatie: De combinatie van de generatieve distributie van het sterke model (geconcentreerd op $V_\kappa$) en de niet-lineaire leraar-transformatie fungeert als een impliciete regularisator.
• Voorwaarde voor Behoud: Niet-doel-features worden behouden mits de niet-doel-taken meer neuronen toegewezen krijgen dan de doeltaak (of kleinere reward-gewichten dragen), gekwantificeerd door de localisatie-parameter $\chi_k = N_\kappa \pi_k / (N_k \pi_\kappa)$.
• Resultaat: De alignment van niet-doel-neuronen met hun oorspronkelijke richtingen $\theta_k$ wordt gehandhaafd, waarbij de daling in alignment begrensd wordt door termen gerelateerd aan overlap bij initialisatie en SGD-ruis, in plaats van catastrofale vergetelheid.

3. Catastrofaal Vergeten bij Standaard SFT (Stelling 2)

In tegenstelling tot W2S, bewijst het artikel dat standaard Supervised Fine-Tuning (SFT) leidt tot catastrofale vergetelheid van niet-doel-features wanneer die features richting-afhankelijkheid delen met het doel.

• Mechanisme: Onder isotrope sampling verspreidt het gradiëntsignaal zich naar alle neuronen. Als de initiële alignment tussen een niet-doel-neuron en de doel-richting $\theta_\kappa$ niet-verwaarloosbaar is (specifiek, als $\theta_k^\top \theta_\kappa \gtrsim s^{-1/2}$), convergeren de niet-doel-neuronen sneller naar $\theta_\kappa$ dan de doel-neuronen convergeren naar $\theta_\kappa$.
• Uitkomst: De voorgetrainde feature $\theta_k$ gaat verloren voordat de doeltaak volledig is geleerd, wat resulteert in de vernietiging van de algemene capaciteiten van het model.

Numerieke Experimenten

De auteurs valideren hun theoretische bevindingen met synthetische data met input-dimensie $d=1024$ en subruimte-dimensie $s=128$.

• W2S Prestaties: Experimenten bevestigen dat het sterke model de doel-feature-richting ($\theta_1$) leert terwijl het alignment behoudt met de niet-doel-feature ($\theta_2$), in overeenstemming met Stelling 1 en Propositie 2.
• SFT Falen: In de SFT-setting aligneren de niet-doel-neuronen met de doel-richting $\theta_1$ voordat de doel-neuronen volledig convergeren, wat de theoretische voorspelling van catastrofale vergetelheid bevestigt (Stelling 2).

Betekenis en Claims

Het artikel claimt een theoretische fundering te leggen voor W2S generalisatie in het feature-learning regime, verdergaand dan eerdere analyses die vaste representaties of lineaire modellen aannamen.

• Mechanisme van Elicitatie: Het toont aan dat W2S niet alleen nieuwe capaciteiten "installeert", maar effectief latente kennis "eliciteert" door gebruik te maken van voorgetrainde subruimtes en niet-lineaire transformaties om de steekproefcomplexiteit te verlagen.
• Voordeel ten opzichte van SFT: Het werk biedt een theoretische verklaring voor waarom W2S superieur is aan standaard SFT in het behoud van algemene capaciteiten. Dit wordt toegeschreven aan het algoritmische verschil: W2S vraagt het zwakke model om de generatieve distributie van het sterke model (geconcentreerd op de taak-subruimte), terwijl SFT vertrouwen stelt op isotrope sampling wat interferentie veroorzaakt tussen gecorreleerde features.
• Beperkingen: De auteurs merken op dat hun analyse zich richt op feature-alignment in de eerste laag en formeel geen analyse uitvoert van leerproces in de tweede laag (gewichtsupdates voor de outputlaag). Zij erkennen ook dat de vermindering van de generatieve exponent vereist dat de doel-linkfunctie een even polynoom is.

Kortom, het artikel betoogt dat W2S slaagt omdat de voorgetrainde structuur van het sterke model (subruimte-localisatie) in combinatie met het specifieke trainingsprotocol (generatieve distributie + niet-lineaire transformatie) het in staat stelt om de doel-feature efficiënt te leren zonder de orthogonale of zwak gecorreleerde features te verstoren die zijn bredere kennis coderen.