2025-11-12T01:28:29.133817

Stability of Transformers under Layer Normalization

Kan, Li, Zhang et al.
Despite their widespread use, training deep Transformers can be unstable. Layer normalization, a standard component, improves training stability, but its placement has often been ad-hoc. In this paper, we conduct a principled study on the forward (hidden states) and backward (gradient) stability of Transformers under different layer normalization placements. Our theory provides key insights into the training dynamics: whether training drives Transformers toward regular solutions or pathological behaviors. For forward stability, we derive explicit bounds on the growth of hidden states in trained Transformers. For backward stability, we analyze how layer normalization affects the backpropagation of gradients, thereby explaining the training dynamics of each layer normalization placement. Our analysis also guides the scaling of residual steps in Transformer blocks, where appropriate choices can further improve stability and performance. Our numerical results corroborate our theoretical findings. Beyond these results, our framework provides a principled way to sanity-check the stability of Transformers under new architectural modifications, offering guidance for future designs.
academic

Stabilität von Transformern unter Layer Normalization

Grundlegende Informationen

  • Paper-ID: 2510.09904
  • Titel: Stability of Transformers under Layer Normalization
  • Autoren: Kelvin Kan (UCLA), Xingjian Li (UT Austin), Benjamin J. Zhang (UNC Chapel Hill), Tuhin Sahai (SRI International), Stanley Osher (UCLA), Krishna Kumar (UT Austin), Markos A. Katsoulakis (UMass Amherst)
  • Klassifizierung: cs.LG, cs.AI, math.OC
  • Veröffentlichungsdatum: 10. Oktober 2025
  • Paper-Link: https://arxiv.org/abs/2510.09904

Zusammenfassung

Obwohl Transformer weit verbreitet sind, kann das Training tiefer Transformer instabil sein. Layer Normalization (Schichtennormalisierung) als Standardkomponente verbessert die Trainingsstabilität, aber die Positionswahl ist oft ad hoc. Dieses Papier führt eine prinzipielle Untersuchung der Vorwärtsstabilität (versteckte Zustände) und Rückwärtsstabilität (Gradienten) von Transformern unter verschiedenen Layer-Normalization-Positionen durch. Die theoretische Analyse offenbart wichtige Erkenntnisse über die Trainingsdynamik: ob das Training den Transformer zu regulären Lösungen oder pathologischem Verhalten führt. Bei der Vorwärtsstabilität werden explizite Grenzen für das Wachstum versteckter Zustände in trainierten Transformern hergeleitet. Bei der Rückwärtsstabilität wird analysiert, wie Layer Normalization die Gradienten-Rückpropagation beeinflusst und damit die Trainingsdynamik jeder Layer-Normalization-Position erklärt. Die Analyse leitet auch die Skalierung von Residuenschrittweiten in Transformer-Blöcken an, wobei angemessene Auswahl die Stabilität und Leistung weiter verbessern kann.

Forschungshintergrund und Motivation

Problemdefinition

Das Kernproblem dieser Forschung ist: Der Mechanismus der Auswirkung verschiedener Layer-Normalization-Positionen auf die Trainingsstabilität von Transformern. Dies umfasst konkret:

  1. Vorwärtsstabilitätsproblem: Kontrolle des Wachstums versteckter Zustände in tiefen Netzwerken
  2. Rückwärtsstabilitätsproblem: Stabilität von Gradienten während der Rückpropagation
  3. Architektur-Designanleitung: Wie man theoretische Anleitung für neue Transformer-Varianten bereitstellt

Bedeutungsanalyse

  1. Praktischer Wert: Transformer sind die Grundlage des modernen Deep Learning, und ihre Trainingsstabilität beeinflusst direkt die Modellleistung und Trainingseffizienz
  2. Theoretische Lücke: Die aktuelle Auswahl von Layer-Normalization-Positionen basiert hauptsächlich auf Erfahrung und entbehrt theoretischer Grundlagen
  3. Industrielle Anforderung: Mit zunehmender Modellgröße werden Trainingsstabilitätsprobleme immer drängender

Einschränkungen bestehender Methoden

  1. Post-LN: Erfordert feine Optimierungsplanung, Leistung ist oft suboptimal
  2. Pre-LN: Verbessert zwar die frühe Trainingsstabilität, erzeugt aber übermäßig große versteckte Zustände, was zu numerischer Instabilität führt
  3. Peri-LN: Zeigt zwar gute praktische Leistung, aber theoretische Eigenschaften sind nicht ausreichend verstanden

Forschungsmotivation

Die Autoren verwenden eine neue Perspektive der kontinuierlichen Zeitdynamik und optimalen Kontrolltheorie, modellieren das Transformer-Trainingsproblem als Mean-Field-Kontrollproblem und können damit die Eigenschaften des Modells nach Trainingskonvergenz analysieren, anstatt sich nur auf das Verhalten bei der Initialisierung zu konzentrieren.

Kernbeiträge

  1. Theoretischer Rahmen-Innovation: Vorschlag eines neuen Rahmens basierend auf optimaler Kontrolltheorie zur systematischen Analyse der Transformer-Stabilität unter verschiedenen Layer-Normalization-Positionen
  2. Vorwärtsstabilitätsanalyse: Herleitung expliziter Grenzen für das Wachstum versteckter Zustände, Beweis dass Pre-LN zu unbegrenztem Wachstum führt während Peri-LN kontrolliertes Wachstum aufrechterhält
  3. Rückwärtsstabilitätsanalyse: Offenlegung des Mechanismus, wie Layer Normalization die Gradienten-Rückpropagation beeinflusst
  4. Residuenschrittweiten-Skalierung: Vorschlag einer Residuenschrittweiten-Skalierungsmethode zur Verbesserung von Stabilität und Leistung
  5. Experimentelle Validierung: Validierung theoretischer Erkenntnisse an GPT-2-Modellserien

Methodische Details

Aufgabendefinition

Untersuchung der Stabilität von Transformern unter verschiedenen Layer-Normalization-Positionen, einschließlich:

  • Eingabe: Sequenz X0∈Rd×nX_0 \in \mathbb{R}^{d \times n} nach Einbettung und Positionskodierung
  • Ausgabe: Versteckte Zustände XDX_D nach D Transformer-Blöcken
  • Ziel: Analyse der Stabilität von Vorwärts- und Rückwärtspropagation

Kontinuierliche Zeitmodellierung

Kontinuierliche Zeitdarstellung des Transformers

Die Sprungverbindungsstruktur des Standard-Transformer-Blocks wird als Euler-Diskretisierung einer kontinuierlichen Zeitdynamik interpretiert:

f_{attn}(X(t), t; \theta_{attn}(t)), & t \in [t_i, t_i + \Delta t) \\ f_{ffn}(X(t), t; \theta_{ffn}(t)), & t \in [t_i + \Delta t, t_{i+1}) \end{cases}$$ wobei $\Delta t = \frac{T}{2D}$, $t_i = 2i\Delta t$. #### Modellierung als Mean-Field-Kontrollproblem Das Trainingsproblem wird als kontinuierliches Mean-Field-Kontrollproblem formuliert: $$\min_\theta \mathbb{E}_{(X_0,y)} G(X(T), y)$$ $$\text{s.t. } \frac{dX(t)}{dt} = f(X(t), t; \theta(t))$$ wobei $f \in \{f_{Pre}, f_{Peri}\}$ verschiedenen Layer-Normalization-Positionen entspricht. ### Geometrische Eigenschaften der Layer Normalization **Schlüsselles Lemma 1**: Die Layer-Normalization-Ausgabe liegt auf einer Ellipsoidfläche: $$\mathcal{E} = \{z \in \mathbb{R}^d : (z - \beta)^T\Gamma^{-2}(z - \beta) = d\}$$ wobei $\Gamma = \text{diag}(\gamma)$. ### Vorwärtsstabilitätsanalyse #### Unbegrenztheit von Pre-LN **Theorem 2**: Das optimale Lösungsproblem des Pre-LN-Trainings ist in der Magnitude unbegrenzt. **Beweisansatz**: Durch Analyse der Hamilton-Jacobi-Bellman (HJB) partiellen Differentialgleichung wird bewiesen, dass der entsprechende Hamiltonsche nicht existiert, was zum Degenerieren des Trainingsproblems führt. **Theorem 3**: Selbst mit Gewichtsabfall zeigen Pre-LN Transformer versteckte Zustände mit exponentiellem Wachstum: $$MA(X_D) \leq (1 + C(\lambda))^D \frac{\|X_0\|_F}{\sqrt{nd}} = O(e^D)$$ #### Kontrolliertes Wachstum von Peri-LN **Theorem 4**: Die versteckten Zustände von Peri-LN Transformern zeigen lineares Wachstum: $$MA(X_D) \leq \frac{\|X_0\|_F}{\sqrt{nd}} + 2D(\gamma_{max} + \beta_{max}) = O(D)$$ Die Varianz zeigt quadratisches Wachstum: $$\text{Var}(X_D) \leq \frac{(\|X_0\|_F + 2D\sqrt{nd}(\gamma_{max} + \beta_{max}))^2}{nd - 1} = O(D^2)$$ ### Rückwärtsstabilitätsanalyse Gradienten-Berechnungsformel: $$\nabla_{\theta_i} G(X_D) = \nabla_{\theta_i} X_{i+1} \cdot J_{i:D} \cdot \nabla_{X_D} G(X_D)$$ wobei die Jacobi-Matrix: $$J_{i:D} = \prod_{j=i+1}^D (I + \nabla_{X_{j-1}} f(X_{j-1}; \theta_{j-1}))$$ **Proposition 7**: Unter Pre-LN wächst die Sensitivität $\nabla_{X_{j-1}} f_{Pre}$ proportional zu Aktivierungswerten. **Proposition 8**: Unter Peri-LN ist die Sensitivität $\nabla_{X_{j-1}} f_{Peri}$ invariant gegenüber der Aktivierungswertmagnitude. ## Experimentelle Einrichtung ### Datensätze - **OpenWebText-Datensatz**: Etwa 9 Milliarden Trainings-Token, 4 Millionen Validierungs-Token - Vortraining mit GPT-2-Serienarchitektur ### Modellkonfiguration - **GPT-2** (124M Parameter) - **GPT-2 Large** (774M Parameter) - **GPT-2 XL** (1,5B Parameter) ### Bewertungsmetriken - **Verwirrung (Perplexity)** - **ROUGE-Scores** (Rouge1, Rouge2, RougeL) - **BERT-Scores** (BertP, BertR, BertF1) - **Trainingsstabilität**: Statistik der divergierenden Läufe ### Implementierungsdetails - Verwendung von für Pre-LN optimierten Hyperparametern, keine separate Optimierung für Peri-LN - Residuenschrittweiten-Skalierung: $\Delta t \in \{0.1, 1\}$ - Hardware: NVIDIA H200 GPU ## Experimentelle Ergebnisse ### Vergleich der Trainingsstabilität | Layer-Normalization-Einstellung | Gewichtsabfall aktiviert | Gewichtsabfall deaktiviert | |----------------------------------|--------------------------|---------------------------| | Pre-LN | 1/5 divergiert | 3/5 divergiert | | Peri-LN | 0/5 divergiert | 0/5 divergiert | | Keine LN | 5/5 divergiert | — | ### Leistungsvergleichsergebnisse **GPT-2 (124M) Modellergebnisse**: - Pre-LN ($\Delta t=1$): Validierungsverlust 5,43, Verwirrung 247,52 - Pre-LN ($\Delta t=0.1$): Validierungsverlust 3,13, Verwirrung 24,43 - Peri-LN ($\Delta t=1$): Validierungsverlust 3,12, Verwirrung 24,17 - Peri-LN ($\Delta t=0.1$): Validierungsverlust 3,10, Verwirrung 23,63 ### Analyse des Wachstums versteckter Zustände Experimente validieren theoretische Vorhersagen: - Pre-LN zeigt bei größerem $\Delta t$ schnelles Wachstum - Peri-LN behält regelmäßigeres lineares Wachstum bei - Residuenschrittweiten-Skalierung kontrolliert effektiv die Wachstumsrate ### Effekt der Residuenschrittweiten-Skalierung 1. **Leistungsverbesserung**: Peri-LN + $\Delta t=0.1$ zeigt beste Leistung bei allen Metriken 2. **Stabilitätsverbesserung**: Pre-LN wird bei $\Delta t=0.1$ von instabil zu stabil 3. **Wachstumskontrolle**: Effektive Reduzierung der Wachstumsrate von Mittelwert und Varianz versteckter Zustände ## Verwandte Arbeiten ### Layer-Normalization-Forschung - **Post-LN**: Frühestes Transformer-Design, erfordert feine Planung - **Pre-LN**: Verbessert Trainingsstabilität, erzeugt aber große Aktivierungswerte - **Peri-LN**: Kürzlich in großen Modellen wie Gemma2, OLMo2 übernommen ### Theoretische Analysemethoden - Bestehende Arbeiten konzentrieren sich hauptsächlich auf Verhalten bei Initialisierung oder basieren auf empirischen Beobachtungen - Dieses Papier analysiert innovativ die Modelleigenschaften nach Trainingskonvergenz - Die kontinuierliche Zeitperspektive bietet neue Werkzeuge für Architekturanalyse ## Schlussfolgerungen und Diskussion ### Hauptschlussfolgerungen 1. **Theoretische Mängel von Pre-LN**: Das Trainingsproblem ist von Natur aus schlecht gestellt, was zu unbegrenzten Lösungen führt 2. **Vorteile von Peri-LN**: Bietet gut definiertes Optimierungsproblem und kontrolliertes Wachstum versteckter Zustände 3. **Wert der Residualskalierung**: Einfache und effektive Methode zur Stabilitätsverbesserung ### Einschränkungen 1. **Vereinfachte Annahmen**: Theoretische Analyse basiert auf kontinuierlicher Zeitapproximation 2. **Hyperparameter-Abhängigkeit**: Experimente verwenden für Pre-LN optimierte Hyperparameter 3. **Skalierungsbeschränkung**: Experimente hauptsächlich auf mittleren Modellgrößen durchgeführt ### Zukünftige Richtungen 1. **Architektur-Screening-Framework**: Theoretisches Screening-Kriterium für neue Architekturmodifikationen 2. **Validierung in größerem Maßstab**: Validierung theoretischer Erkenntnisse an größeren Modellen 3. **Andere Normalisierungsmethoden**: Erweiterung der Analyse auf RMSNorm und andere Varianten ## Tiefgreifende Bewertung ### Stärken 1. **Starke theoretische Innovation**: Erste Anwendung der optimalen Kontrolltheorie auf das Layer-Normalization-Positionsproblem 2. **Mathematische Strenge**: Vollständige theoretische Herleitung und Beweise 3. **Hoher praktischer Wert**: Residuenschrittweiten-Skalierungsmethode ist einfach und effektiv 4. **Angemessenes Experimentdesign**: Validierung theoretischer Erkenntnisse über mehrere Modellgrößen ### Mängel 1. **Theorie-Praxis-Lücke**: Kontinuierliche Zeitannahme und tatsächliche diskrete Implementierung weichen ab 2. **Begrenzte Experimentreichweite**: Validierung hauptsächlich auf GPT-2-Serie, fehlende Validierung über mehr Architekturen 3. **Hyperparameter-Fairness**: Keine spezialisierte Hyperparameter-Optimierung für Peri-LN durchgeführt ### Einflussanalyse 1. **Akademischer Beitrag**: Bietet neuen theoretischen Rahmen für Transformer-Stabilitätsanalyse 2. **Praktischer Wert**: Leitet tatsächliches Modelldesign und Trainingsstrategien an 3. **Reproduzierbarkeit**: Zusage zur Veröffentlichung von Code und Modellen ### Anwendungsszenarien 1. **Tiefes Transformer-Training**: Besonders geeignet für großskalige tiefe Modelle 2. **Architektur-Designanleitung**: Bietet theoretische Grundlagen für neue Architekturmodifikationen 3. **Trainingsstabilitätsverbesserung**: Verbessert Trainingsstabilität durch Residualskalierung ## Referenzen Das Papier zitiert mehrere wichtige Arbeiten, einschließlich: - Ba et al. (2016): Originalpapier zu Layer Normalization - Xiong et al. (2020): Vergleichsstudie Pre-LN vs. Post-LN - Kim et al. (2025): Empirische Forschung zu Peri-LN - He et al. (2016): Bahnbrechendes Werk zu Residualverbindungen --- **Gesamtbewertung**: Dies ist ein hochqualitatives Papier mit guter Kombination von Theorie und Praxis, das einen neuen mathematischen Rahmen für die Transformer-Stabilitätsanalyse bietet und wichtige akademische und praktische Bedeutung hat. Die theoretische Analyse ist streng und tiefgreifend, die experimentelle Validierung ist umfassend und bietet wertvolle Anleitung für das Design tiefer Lernarchitekturen.