Technische Samenvatting: Leren in een Gekromde Gewichtsruimte (SymExpLin)

1. Probleemstelling

Standaard training van neurale netwerken vertrouwt op additieve updates in een lineaire parameterruimte. Veel interne netwerkoperaties zijn echter inherent multiplicatief (bijv. het schalen van feature-amplitudes, normalisatie-gains of singulariteitswaarden). In een lineaire ruimte vereisen relatieve veranderingen van gelijke grootte (bijv. verdubbelen versus halveren van een gewicht) ongelijke optimalisatieafstanden. Zo vereist het bewegen van een gewicht van 1,0 naar 2,0 een stap van grootte 1,0, terwijl het bewegen van 1,0 naar 0,5 een stap van 0,5 vereist.

Deze mismatch creëert optimalisatieproblemen:

• Heavy-Tailed Verdelingen: Pretrained modellen vertonen vaak gewichtsverdelingen met zware staarten (heavy-tailed), waarbij een klein deel van de parameters aanzienlijk groter wordt dan hun initialisatieschaal (bijv. 30× groter), terwijl de meeste nabij de initialisatie blijven.
• Optimizer Invariantie: Adaptieve optimizers zoals Adam normaliseren updates per coördinaat, waardoor stapgroottes invariant zijn voor de parametergrootte. Dit resulteert in een situatie waarin een kleine parameter (bijv. 0,004) en een grote parameter (bijv. 0,4) vergelijkbare absolute updates ontvangen. Dit leidt tot een 100× verschil in relatieve perturbatie, wat kan resulteren in het over-perturberen van kleine parameters terwijl grote, kritieke parameters onder-perturbeerd worden.
• Statische vs. Dynamische Schaling: Bestaande oplossingen zoals Maximal Update Parameterization (µP) adresseren schaal-mismatches statisch op basis van de tensorbreedte bij initialisatie, maar passen zich niet aan de dynamische groei van individuele parameters tijdens de training aan.

2. Methodologie: SymExpLin (SEL)

De auteurs stellen SymExpLin (SEL) voor, een niet-lineaire gewichtsreparameterisatie die effectieve gewichten ontleedt in een symmetrisch-exponentieel pad en een lineair pad. Dit creëert een gekromde parameterruimte-geometrie waarbij additieve updates in de ruwe ruimte mappen naar magnitude-proportionele veranderingen in de effectieve gewichtsruimte.

2.1 De Transformatie

Het effectieve gewicht $w_{eff}$ wordt berekend uit een ruw gewicht $w$ met behulp van de volgende componenten:

1. Symmetrisch-Exponentieel Pad ($E_\theta$): Een leerbaar, teken-behoudend functie gebaseerd op de symmetrische exponent (symexp). Voor kleine $|w|$ gedraagt het zich lineair; voor grote $|w|$ groeit het exponentieel. Dit zorgt ervoor dat grote ruwe stappen grotere effectieve verplaatsingen produceren.
   • $E_\theta(|w|) = e_w (\exp(\beta |w| \cdot m - n) - \exp(-n))$
   • Parameters: $\beta$ (kromming), $m$ (interne multiplier), $n$ (offset voor helling bij de oorsprong), $e_w$ (output schaal).
2. Lineair Pad ($L_\theta$): Een identiteit-achtige pathway $L_\theta(w) = w \cdot l_w$ die een directe gradiëntroute biedt, wat naar hypothese de optimalisatie stabiliseert.
3. Sign Wrapping: Het teken van het ruwe gewicht wordt toegepast op de som van de paden: $w_{eff} = \text{sign}(w) \cdot (E_\theta(|w|) + L_\theta(w)) + \Delta_\phi$.
4. Residu ($\Delta_\phi$): Een optionele low-rank (LoRA-stijl) additieve correctie.

2.2 Belangrijke Ontwerpkeuzes

• Mismatched Initialisatie: De auteurs identificeren een gunstige asymmetrie waarbij ruwe gewichten worden geïnitialiseerd door een "congruente" combinatie regel te inverteren (waar lineaire en exponentiële delen buiten de sign wrap worden opgeteld), maar de forward pass een "mismatch" regel gebruikt (waar beide binnen de sign wrap worden opgeteld). Dit resulteert in negatieve gewichten die bij een kleinere magnitude worden geïnitialiseerd dan positieve gewichten, wat fungeert als een vroege symmetrie-doorbrekende mechanisme.
• Gestructureerde Schaalparameters: Leerbare schalen ($e_w, l_w, m, n$) zijn geparametriseerd met behulp van rij-kolom multiplicatieve patronen (rank-1 vectoren) om expressiviteit en computationele kosten in balans te houden.
• Learning Rate Annealing: Een gericht schema wordt gebruikt voor de schaalparameters. De externe schaal $e_w$ begint met een hoge learning rate om snel een teken te committeren, terwijl de interne multiplier $m$ laag begint en opwaarts annelt om de kromming te verfijnen zodra de schaal stabiliseert.
• Kromming Calibratie: De krommingparameter $\beta$ wordt heuristisch ingesteld op basis van modelbreedte om te garanderen dat het exponentiële regime actief wordt nabij typische initialisatie-magnitudes (die krimpen naarmate de breedte toeneemt).

3. Belangrijkste Bijdragen

• SymExpLin (SEL): Een gewichtsreparameterisatie toepasbaar op elke leerbare matrix of bias, die gewichten ontleedt in symmetrisch-exponentiële en lineaire componenten met onafhankelijk leerbare gestructureerde schalen.
• Mismatched Initialisatie: De ontdekking dat het gebruik van verschillende combinatie regels voor initialisatie-inversie versus de forward pass een gunstige asymmetrische bias creëert die vroege optimalisatie verbetert.
• Gestructureerde Schaal & Annealing: Introductie van rij-kolom multiplicatieve schaalparameters en een specifiek learning rate annealing schema om redundante exponentiële controles te beheren.
• Praktische Toepasbaarheid: De methode is implementatievriendelijk. De trainings-transformatie is element-wise, en post-training worden de effectieve gewichten vooraf berekend en samengevoegd (folded) in standaard lineaire lagen, wat resulteert in nul inferentiekosten.
• Complementariteit: SEL opereert op een dynamische, per-parameter as, en complementeert statische, per-tensor methoden zoals µP.

4. Experimentele Resultaten

De methode werd geëvalueerd op autoregressieve taalmodellering met OpenWebText over negen configuraties van breedte $\{1024, 2048, 3072\}$ en diepte $\{12, 24, 36\}$.

• Loss Descent: SEL bereikte consistent een lagere validatie loss dan de baseline (Adam met standaard parameterisatie) over alle configuraties.
• Stap Efficiëntie: SEL bereikte een gematchte validatie loss in 1,32× tot 1,49× minder trainingsstappen. De winst was het meest uitgesproken in grotere breedte-modellen (bijv. 1,49× versnelling bij 3072×24).
• Wall-Clock Speedup: Ondanks een per-stap overhead van 1,7% (klein model) tot 5,5% (groot model) door de transformatie, resulteerde de reductie in totale stappen in een geschatte 1,29× tot 1,37× wall-clock speedup.
• Ablaties:
  • Lineair Pad: Het verwijderen van het lineaire pad resulteerde in een "failed-to-learn" geval, wat de noodzaak ervan bevestigt voor stabiliteit.
  • Mismatch vs. Congruent: Mismatched initialisatie presteerde consistent beter dan congruente initialisatie, met name in kleinere modellen.
  • Annealing: Het geannelleerde learning rate schema voor schaalparameters bood consistente, zij het kleine, verbeteringen ten opzichte van vaste multipliers.

5. Betekenis en Claims

Het artikel claimt dat SEL de fundamentele mismatch tussen de multiplicatieve aard van veel neurale berekeningen en de additieve aard van standaard optimalisatie aanpakt. Door de parameterruimte te krommen, stelt SEL de optimizer in staat om stappen te nemen die van nature proportioneel zijn aan de magnitude van de gewichten, wat effectief relatieve schaling biedt zonder dat er per-parameter learning rate tuning nodig is.

De auteurs benadrukken dat SEL een praktische, drop-in vervanging is voor standaard lineaire lagen die:

1. Convergentie versnelt, met name voor modellen met heavy-tailed gewichtsverdelingen.
2. Geen architecturale wijzigingen vereist tijdens inferentie.
3. Kan worden geïnverteerd, waardoor standaard pretrained gewichten kunnen worden omgezet naar de SEL-ruimte voor voortgezette training of fine-tuning.

Het werk suggereert dat de "soft log-space" parameterisatie hels helpt om gewichten de grote-magnitude targets te laten bereiken die vereist zijn voor de computatie van het netwerk, welke vaak benadeeld worden door lineaire additieve updates.