Sphere Retraction Normalizations

Technische Samenvatting: Sphere Retraction Normalizations

Probleemstelling
Het stabiel trainen van diepe neurale netwerken blijft een aanzienlijke uitdaging. Bestaande methodologieën, die primair vertrouwen op residuele verbindingen gecombineerd met normalisatie (bijv. Pre-LN en Post-LN in Transformers), vertonen inherente beperkingen. Pre-LN architecturen lijden onder de "vloek van de diepte", waarbij diepere lagen een afnemende bijdrage leveren aan het leren, en laten onbegrensde residuele stromen toe die massieve activaties accumuleren, wat leidt tot attention sinks. Daarentegen kampen Post-LN architecturen met grote gradiënten nabij de outputlaag tijdens initialisatie, wat zorgvuldige learning-rate warm-up schema's vereist die de optimalisatie vertragen. Hoewel recent werk zoals Geodesic Normalization (GeoNorm) probeert de training te stabiliseren door verborgen toestanden binnen een hypersfeer te beperken via Riemanniaanse meetkunde, vertrouwt dit uitsluitend op de exponentiële afbeelding voor updates. De auteurs stellen dat de exponentiële afbeelding slechts één lid is van een bredere familie van retractie-afbeeldingen, wat de vraag oproept of alternatieve retracties op de hypersfeer superieure prestaties of stabiliteit kunnen bieden.

Methodologie
Het artikel stelt een verenigd kader voor diepe verbindingen voor op basis van retractie-afbeeldingen op een hypersfeer, getiteld Sphere Retraction Normalizations (SpheretNorms).

1. Verenigd Kader: De auteurs herformuleren standaard Residuele Verbindingen (RC) en GeoNorm binnen een enkel computationeel skelet gedefinieerd op een manifold $\mathcal{M}$. De update-regel wordt geabstraheerd als:
   $$h_l = R_{h_{l-1}}\left[ \alpha_l \cdot \text{proj}_{T_{h_{l-1}}\mathcal{M}}(F_l(h_{l-1})) \right]$$
   waarbij $R$ een retractie-afbeelding is. Terwijl GeoNorm de hypersfeer ($S^{d-1}$) en de exponentiële afbeelding selecteert, staat dit framework andere retracties toe.

2. Specifieke Instanties:

   • Proj-SpheretNorm: Instantieert het framework met de metrische projectie-retractie (ook wel bekend als de gnomonische retractie). Dit houdt in dat de vector som van de huidige toestand en de update wordt genormaliseerd om op de hypersfeer te liggen.
   • Cay-SpheretNorm: Gebruikt de Cayley-retractie (stereografische retractie), afgeleid via de Cayley-transformatie van een anti-symmetrische matrix gegenereerd door de huidige toestand en de update-vector. Dit levert een algebraïsche update in gesloten vorm op.

3. Generalisatie (p-SpheretNorm): De auteurs observeren dat zowel Proj-SpheretNorm als Cay-SpheretNorm een planaire rotatiestructuur delen waarbij de update-magnitude wordt omgezet in een rotatiehoek. Ze introduceren een een-parametrische familie, p-SpheretNorm, gedefinieerd door een rotatiehoek $\theta_l = p \arctan(\beta_l / p)$, waarbij $\beta_l$ de genormaliseerde stapgrootte is.

   • $p=1$ herstelt Proj-SpheretNorm.
   • $p=2$ herstelt Cay-SpheretNorm.
   • $p \to \infty$ convergeert naar GeoNorm (exponentiële afbeelding).
   • $p \to 0^+$ convergeert naar de identiteitsafbeelding.
   • Cruciaal is dat voor een eindige $p$, de rotatiehoek verzadigt, wat ongebreidelde groei voorkomt, in tegen tegenstelling tot de exponentiële afbeelding.

4. Implementatiedetails: Om stabiliteit te garanderen, gebruiken de auteurs een vervactorfactor $\lambda_l = 1/\sqrt{l}$ voor de leerbare stapgrootte $\alpha_l$ en passen ze een gemodificeerde ScaleNorm toe in de eerste laag om de straal van de verborgen toestand te initialiseren.

Kernbijdragen

1. Theoretische Unificatie: Het artikel verenigt Euclidische residuele verbindingen en GeoNorm als retractie-gebaseerde updates, en demonstreert dat GeoNorm een specifiek geval is van een Riemanniaanse residuele verbinding gedreven door de exponentiële afbeelding.
2. Nieuwe Normalisatielagen: Het introduceert twee nieuwe, norm-behoudende normalisatielagen (Proj-SpheretNorm en Cay-SpheretNorm) die uitsluitend vertrouwen op algebraïsche operaties, waardoor de computationele kosten van de trigonometrische functies die vereist zijn voor de exponentiële afbeelding worden vermeden.
3. Parametrische Familie: Het generaliseert deze methoden tot p-SpheretNorm, een een-parametrische familie die de kloof overbrugt tussen de identiteitsafbeelding, algebraïsche retracties en de exponentiële afbeelding. De auteurs bewijzen dat alle leden van deze familie tweede-orde retracties zijn met identieke eerste-orde tangentiële dynamiek.
4. Empirische Superieuriteit: Uit uitgebreide experimenten op nanoGPT (tot 36 lagen, 6.55B tokens) blijkt dat 1-SpheretNorm (overeenkomend met $p=1$) de laagste pre-training en validatie verliezen en de hoogste zero-shot nauwkeurigheid bereikt over de OpenWebText en FineWeb-Edu datasets, waarbij het Pre-LN, Pre-DyT, Peri-LN, Keel en GeoNorm overtreft. Opvallend genoeg convergeert 1-SpheretNorm in settings waar Pre-LN en Keel divergeren.

Resultaten

• Pre-training: Op medium (24 lagen) en grote (36 lagen) modellen behaalde 1-SpheretNorm de beste validatie perplexiteit. Bijvoorbeeld, op FineWeb-Edu met 24 lagen reduceerde het de validatie perplexiteit van 19.09 (Peri-LN) naar 16.81.
• Stabiliteit: In tegen tegenstelling tot Pre-LN (die divergeerde op OpenWebText bij medium schaal) en Keel (die divergeerde op FineWeb-Edu bij kleine schaal), convergeerden alle SpheretNorm varianten in alle geteste settings.
• Downstream Taken: Op de grote modellen getraind op FineWeb-Edu behaalden SpheretNorm varianten de beste score op 11 van de 12 metrieken; op OpenWebText behaalden ze de beste score op 10 van de 12 metrieken. De winst was het meest uitgesproken in perplexiteit-metrieken (bijv. LAMBADA en WikiText-2), waar 1-SpheretNorm significant beter presteerde dan GeoNorm en andere baselines.
• Parametergevoeligheid: De resultaten wijzen erop dat de optimale prestatie wordt bereikt bij een eindige $p$ (specifiek $p=1$), wat suggereert dat de exponentiële afbeelding (de limiet als $p \to \infty$) niet de geprefereerde retractie is voor sferische residuele stromen, maar slechts één uiteinde van een spectrum.

Betekenis en Claims
Het artikel claimt dat de keuze van de retractie-afbeelding een kritische ontwerpkeuze is voor diepe verbindingen op hypersferen. Door verder te gaan dan de exponentiële afbeelding naar algebraïsche retracties (specifiek metrische projectie), demonstreren de auteurs dat het mogelijk is om norm-behoud en trainingsstabiliteit te bereiken met superieure prestaties. Het werk suggereert dat de "exponentiële afbeelding" geen fundamentele noodzaak is voor stabiel diep leren op manifolds, maar eerder een specifiek punt in een breder spectrum van geldige geometrische updates. De auteurs concluderen dat het ontwerpen van diepe verbinding-mechanismen een wiskundig probleem is op het snijvlak van grafentheorie, optimalisatie en dynamische systemen, en dat het identificeren van de optimale afweging tussen informatiepropagatie en Jacobian-conditiering een openstaande richting is voor toekomstig werk. Ze erkennen beperkingen met betrekking tot gespecialiseerde architecturen (bijv. Equivariant Neural Networks) en de noodzaak van evaluatie op grotere schaal modellen buiten de huidige scope.