Technische Samenvatting: Stabiliteit van Transformers onder Layer Normalization

Probleemstelling

Ondanks hun dominantie in moderne deep learning, blijft het trainen van diepe Transformers gevoelig voor instabiliteit, met name naarmate de modeldiepte toeneemt. Deze instabiliteiten manifesteren zich in twee kritieke gebieden: forward instabiliteit (ongecontroleerde groei van verborgen toestanden leidend tot numerieke problemen) en backward instabiliteit (exploderende gradiënten tijdens backpropagation).

Een primaire architecturale component die wordt gebruikt om deze problemen te mitigeren, is Layer Normalization (LN). De plaatsing van LN binnen Transformer-blokken is echter grotendeels geëvolueerd via empirische heuristieken in plaats van rigoureuze theoretische rechtvaardiging. Het document identificeert drie hoofdconfiguraties:

• Post-LN: LN toegepast na residuele verbindingen. Vereist vaak delicate optimalisatieschema's en vertoont suboptimale prestaties.
• Pre-LN: LN toegepast vóór aandacht- en feedforward-modules. Hoewel dit de standaard is voor verbeterde trainingsstabiliteit ten opzichte van Post-LN, wordt er empirisch waargenomen dat het excessief grote verborgen toestanden produceert, wat leidt tot potentiële numerieke instabiliteit en gradiëntexplosie in diepe netwerken.
• Peri-LN: Een nieuwere configuratie die LN toepast op zowel de input als de output van modules. Empirische studies suggereren dat het superieure stabiliteit en meer geregulariseerde verborgen toestanden biedt vergeleken met Pre-LN, maar de theoretische eigenschappen ervan blijven slecht begrepen.

Het kernprobleem dat wordt aangepakt, is het gebrek aan een principieel theoretisch kader om te verklaren waarom deze verschillende plaatsingen leiden tot verschillende trainingsdynamieken en stabiliteitsprofielen, en of training modellen naar reguliere oplossingen of pathologische gedragingen drijft.

Methodologie

De auteurs stellen een nieuw theoretisch kader voor, gebaseerd op discrete-time optimale controletheorie, om de trainingsdynamiek van Transformers te analyseren.

1. Formulering als Optimal Control: Het trainingstraject van de Transformer wordt gegoten in een discrete-time mean-field controleprobleem. De laagindex $i$ representeert discrete tijd, en de verborgen toestand $X_i$ representeert de systeemtoestand. De residuele verbindingen fungeren als de snelheid van de dynamiek.
2. Hamilton-Jacobi-Bellman (HJB) Analyse: De auteurs analyseren de optimaliteitsvoorwaarden van het trainingsprobleem met behulp van de HJB partiële differentiaalvergelijking (PDE) gekoppeld aan massatransportvergelijkingen.
   • Ze onderzoeken de Hamiltoniaan, die fungeert als een feedback-controlesysteem dat het optimale snelheidsveld voor verborgen toestanden definieert.
   • Voor Pre-LN demonstreren ze dat de Hamiltoniaan onbegrensd is, wat leidt tot slecht gedefinieerde optimaliteitsvoorwaarden waarbij verborgen toestanden willekeurig groot kunnen worden.
   • Voor Peri-LN beperkt de output-normalisatie het snelheidsveld tot een ellipsoïde (gebaseerd op Lemma 1), wat garandeert dat de Hamiltoniaan begrensd is en het probleem goed gedefinieerd is.
3. Gradiëntanalyse: Het paper analyseert backward stabiliteit door de propagatie van gradiënten te onderzoeken door het product van lokale sensitiviteiten (Jacobiaan-matrices) over lagen heen.
4. Residuele Schaling: Het framework wordt uitgebreid met een tijd-increment parameter $\Delta t < 1$, wat effectief de magnitude van residuele stappen schaalt om signaalpropagatie te controlen.

Belangrijkste Bijdragen

1. Theoretisch Kader voor Getrainde Modellen

In tegen tegenstelling tot eerder werk dat vaak focust op modelinitialisatie of vertrouwt op empirisch bewijs, analyseert dit kader getrainde modellen bij convergentie. Het biedt een systematische beoordeling van of trainingsdynamieken leiden tot reguliere oplossingen of pathologieën.

2. Rigoureuze Analyse van LN-plaatsingen

• Pre-LN Instabiliteit: Het paper bewijst dat voor Pre-LN de optimale oplossing van het trainingsprobleem onbegrensd is in magnitude. Zelfs met weight decay kunnen verborgen toestanden exponentieel groeien ($O(e^D)$), en gradiënten groeien proportioneel met activaties, wat leidt tot trainingsinstabiliteit.
• Peri-LN Stabiliteit: De auteurs leiden expliciete grenzen af die aantonen dat Peri-LN een lineaire groei garandeert voor de gemiddelde absolute waarde van verborgen toestanden ($O(D)$) en een kwadratische groei voor variantie ($O(D^2)$). Cruciaal is dat de lokale sensitiviteit van Peri-LN wordt getoond invariant te zijn voor de activatiemagnitude, wat gradiëntexplosie voorkomt, zelfs wanneer activaties groot zijn.

3. Onzekerheidskwantificering

De analyse stelt vast dat Peri-LN gecontroleerde propagatie van onzekerheid (gemeten via de Wasserstein-afstand) van input naar terminale verborgen toestanden biedt, terwijl Pre-LN de input-onzekerheid naar onbegrensde niveaus kan versterken.

4. Residuele Stap-schaling

Gestuurd door de stabiliteitsanalyse introduceren de auteurs een residuele stap-schalingsmechanisme ($\Delta t < 1$). Ze demonstreren theoretisch dat deze schaling:

• De grenzen op de groei van verborgen toestanden en output-onzekerheid verscherpt.
• Gradiëntexplosie mitigeert door de lokale sensitiviteit te schalen.
• Zowel forward als backward stabiliteit verbetert zonder extra computationele kosten.

Experimentele Resultaten

De auteurs valideren hun theoretische bevindingen met behulp van GPT-2 varianten (124M tot 1.5B parameters) getraind op de OpenWebText dataset.

• Trainingsstabiliteit: In experimenten met 5 trials divergeerden Pre-LN modellen in 1 run wanneer weight decay werd ingeschakeld en in 3 runs wanneer weight decay werd uitgeschakeld. In contrast bleven Peri-LN modellen stabiel in alle 5 de trials onder beide weight decay settings.
• Prestaties: Peri-LN bereikte prestaties vergelijkbaar met Pre-LN, zelfs wanneer hyperparameters werden afgestemd voor Pre-LN.
• Impact van Schaling: Het combineren van Peri-LN met residuele schaling ($\Delta t = 0.1$) leverde consistent de beste prestaties op over alle metrieken (Val Loss, Perplexity, ROUGE, BERT scores) en controleerde effectief de groei van de momenten van verborgen toestanden (gemiddelde absolute waarde en variantie) over de lagen, wat overeenkomt met de theoretische voorspellingen.

Betekenis en Claims

Het paper claimt de eerste rigoureuze theoretische rechtvaardiging te bieden voor de empirische observaties met betrekking tot de plaatsing van Layer Normalization in Transformers.

• Verklaring van Instabiliteit: Het legt uit waarom Pre-LN leidt tot onbegrensde verborgen toestanden en gradiëntexplosie via de slecht gedefinieerde aard van de onderliggende HJB PDE.
• Rechtvaardiging voor Peri-LN: Het biedt een theoretische basis voor de stabiliteit van Peri-LN, waarbij wordt aangetoond dat het gecontroleerde groei en invariante sensitiviteit behoudt.
• Ontwerprichtlijnen: Het framework biedt een principiële methode om nieuwe architecturale modificaties te "sanity-checken" voordat er dure empirische training plaatsvindt.
• Praktische Verbetering: De introductie van residuele stap-schaling wordt gepresenteerd als een eenvoudige, kosteloze modificatie die de stabiliteit en prestaties aanzienlijk verbetert, gevalideerd door experimenten die laten zien dat Peri-LN met schaling beter presteert dan standaard configuraties.

De auteurs benadrukken dat hun werk de kloof overbrugt tussen empirische heuristieken en theoretisch begrip, waarbij een continu-tijd perspectief wordt geboden dat zowel theoretische inzichten als praktische architecturale verbeteringen oplevert.