Technische Samenvatting: GeoNorm: Verenig Pre-Norm en Post-Norm met Geodetische Optimalisatie

Probleemstelling

De plaatsing van normalisatielagen binnen Transformer-architecturen—specifiek de keuze tussen Pre-Norm en Post-Norm—blijft een cruciale ontwerpbeslissing die de trainingsstabiliteit, modelprestaties en gradiëntstroom beïnvloedt. Hoewel Pre-Norm de dominante keuze is geworden in moderne Large Language Models (LLM's) vanwege de stabiliteit, introduceert het ongebalanceerde gradiënten waarbij lagere lagen grotere updates ontvangen dan hogere lagen. Daarentegen is Post-Norm, oorspronkelijk ontworpen om de training van diepe netwerken te faciliteren, gevoelig voor trainingsinstabiliteit en verliespieken (loss spikes). Bestaande oplossingen zoals DeepNorm en SandwichNorm verbeteren de stabiliteit door middel van parameterinitialisatie of aanvullende normalisatielagen, maar missen een verenigd theoretisch kader om deze mechanismen te analyseren.

Methodologie

De auteurs stellen GeoNorm voor, een nieuwe normalisatiemethode die geworteld is in manifold-optimalisatie. Het kerninzicht van de theorie is het herinterpreteren van de standaard Transformer-updateoperaties ($x + \text{FFN}(x)$ en $x + \text{Attention}(x)$) als optimalisatiestappen op een sfeer. In dit perspectief:

1. Token-embeddings fungeren als de initiële staat op een manifold.
2. Attention- en FFN-modules genereren update-richtingen (geschatte gradiënten).
3. Normalisatie wordt gezien als een projectie terug naar een restrictieset (de sfeer).

Standaard normalisatie (Post-Norm) voert een onbeperkte update uit gevolgd door een Euclidische projectie terug naar de sfeer. De auteurs stellen dat deze projectie geometrische informatie weggooit en de update-richtingen kan verstoren. In plaats daarvan vervangt GeoNorm de projectiestap door een geodetische update met behulp van de exponentiële afbeelding (exponential map).

Voor een punt $x$ op een sfeer en een tangentiële vector $v$ (de geprojecteerde update-richting), wordt de update uitgevoerd intrinsiek langs de geodeet:
$$\text{exp}_x(v) = \cos\left(\frac{\|v\|}{\|x\|}\right)x + \|x\|\sin\left(\frac{\|v\|}{\|x\|}\right)\frac{v}{\|v\|}$$

Deze formulering zorgt ervoor dat de representatie soepel beweegt langs de manifold, waardoor de intrinsieke geometrie van de feature-ruimte behouden blijft. Verder introduceren de auteurs een laag-specifieke update-decay mechanisme voor de stapgrootte ($\alpha_k$), analoog aan learning rate schedules. Ze evalueren drie decay-strategieën:

• Harmonisch: $\alpha_k = \alpha / k$
• Vierkantswortel: $\alpha_k = \alpha / \sqrt{k}$
• Lineair: $\alpha_k = \alpha (T - k) / T$

Het paper demonstreert ook theoretisch dat Pre-Norm geïnterpreteerd kan worden als een speciaal geval van GeoNorm met niet-adaptieve, diepte-afhankelijke hoeken, waardoor de twee benaderingen onder één enkel geometrisch kader worden verenigd.

Belangrijkste Bijdragen

1. Theoretisch Kader: De auteurs bieden een verenigd theoretisch perspectief waarin Transformer-lagen worden geïnterpreteerd als iteratieve optimalisatiestappen op een sferische manifold, waarmee de kloof tussen Pre-Norm en Post-Norm wordt overbrugd.
2. GeoNorm Algoritme: Zij introduceren GeoNorm, dat de standaard projectie-gebaseerde normalisatie vervangt door geodetische updates via de exponentiële afbeelding. Deze methode vereist geen extra trainbare parameters en voegt verwaarloosbare computationele overhead toe.
3. Laag-specifieke Decay: Een nieuw voorstel voor laag-specifieke update-decay voor attention- en FFN-componenten om de trainingsdynamiek te stabiliseren.
4. Empirische Validatie: Uitgebreide experimenten over variërende modelgroottes (125M tot 1.3B), datasets (Arxiv, Books3, FinWeb-Edu) en trainingslengtes demonstreren de effectiviteit van de methode.

Experimentele Resultaten

Het paper rapporteert dat GeoNorm consistent beter presteert dan bestaande normalisatiemethoden (Pre-Norm, Post-Norm, DeepNorm, SandwichNorm) in alle geteste configuraties:

• Prestaties: Op de Arxiv- en Books3-datasets met een 122M-model behaalde GeoNorm het laagste validatieverlies (1.8792 op Arxiv; 3.4040 op Books3) vergeleken met alle baselines.
• Schaalbaarheid: Het voordeel in prestaties blijft bestaan en wordt groter naarmate de modelgrootte toeneemt. Voor een 1.3B-model behaalde GeoNorm een verlies van 3.1759, waarmee het de beste baseline (Pre-Norm op 3.2067) versloeg.
• Stabiliteit: GeoNorm vertoonde superieure trainingsstabiliteit, met name bij lange trainingsruns (tot 4096 sequentielengte). Terwijl Post-Norm verliespieken en trainingsinstabiliteit vertoonde bij langere lengtes, behield GeoNorm een stabiele convergentie.
• Decay Mechanisme: Onder de decay-strategieën leverde de harmonische decay ($\alpha/k$) de beste prestaties op, waarmee nieuwe benchmarks werd gezet op de geteste datasets.

Betekenis en Claims

Het paper claimt dat GeoNorm een naadloze integratie biedt in standaard Transformer-architecturen, waarbij geen aanpassingen aan de kernblokken (Attention of FFN) nodig zijn en er geen extra trainbare parameters worden geïntroduceerd. De auteurs stellen dat door gebruik te maken van geodetische optimalisatie, GeoNorm representaties in staat stelt om coherenter te bewegen langs de sferische manifold, wat leidt tot verbeterde informatiepropagatie en trainingsstabiliteit.

Dit werk positioneert zich als een fundamentele stap in het begrijpen van normalisatiemechanismen door de lens van Riemanniaanse geometrie. Het suggereert dat het empirische succes van Pre-Norm verklaard kan worden door de impliciete stapgrootte-decay, een mechanisme dat GeoNorm expliciet en adaptief maakt. De auteurs concluderen dat dit verenigde perspectief een robuuste basis biedt voor toekomstige innovaties in het ontwerp van Transformer-architecturen, met name voor het schalen van modellen naar grotere dieptes en complexiteiten.