Impact of Layer Norm on Memorization and Generalization in Transformers

AI summary

Overview

Research area: Machine learning — transformer architecture analysis, specifically the role of Layer Normalization (LN) placement and parameters in learning (generalization) versus label memorization.

Technical level: Advanced. The empirical findings are stated plainly, but the paper supports them with gradient-norm theorems, backpropagation derivations, and upper-bound proofs in the appendix.

Scope (one sentence): The paper isolates how the learnable parameters of LayerNorm affect memorization and learning in Pre-LN versus Post-LN transformers, identifies early-layer LNs as the most influential, and explains the divergence through learning-versus-memorization gradient norms.

What This Paper Is About

Transformers come in two LayerNorm arrangements: Pre-LN, where normalization is applied before self-attention and feed-forward sub-layers, and Post-LN, where it is applied after the residual connection. Pre-LN models are now preferred because they stabilize gradient flow, yet the effect of LayerNorm on what a model learns versus what it merely memorizes was not clearly understood. The authors test what happens when the LayerNorm learnable parameters (the weight w and bias b) are removed while the normalization operation N(x) itself is kept intact, and they compare the outcomes across both architectures.

Key Contributions

1. A architecture-dependent role for LayerNorm. The paper identifies that LayerNorm is critical for learning in Pre-LN transformers, whereas in Post-LN transformers it governs memorization — a distinction the authors present as new relative to prior work, including Xu et al. (2019), which suggested LN contributes to overfitting in Pre-LN models.

2. Label recovery through LayerNorm parameter removal. In Post-LN models, removing LN learnable parameters suppresses memorization and restores genuine labels, while in Pre-LN models the same removal leaves memorization and random predictions persistent.

3. Early layers are the pivotal ones. Selectively removing LN parameters from early, middle, or later layers shows that early-layer LNs have the strongest effect — disrupting learning in Pre-LN models and mitigating memorization in Post-LN models.

4. A gradient-based explanation. The authors formalize three theorems: that the learning gradient norm is greater than or equal to the memorization gradient norm across all layers (Theorem 1), an upper bound on the gradient norm of the loss with respect to LN input for both configurations (Theorem 2), and that the upper bounds of early-layer LNs are higher than those of later layers (Theorem 3).

Main Findings

• Pre-LN models depend on LN for learning. Removing LN learnable parameters in Pre-LN transformers significantly disrupts learning. For Qwen2 on the News dataset, once learning is disrupted it does not recover through the end of training, which the authors describe as fundamental instability.

• Post-LN models are robust to LN parameter removal in learning. ELECTRA maintains stable learning throughout training with no signs of degradation. The paper reports similar results for other Post-LN models (BERT, DeBERTa, Longformer, RoBERTa) and Pre-LN models (GPT2, GPT-Neo, ViT-B, DeiT, ViT-S) in Appendix G.1.

• LN governs memorization in Post-LN models. In Post-LN models, LN parameter removal lowers memorization scores and raises recovery scores; for ELECTRA, memorization decreases over epochs while label recovery improves.

• LN parameter removal does not suppress memorization in Pre-LN models. For Qwen2, memorization persists throughout training and label recovery remains poor, with persistently high memorization and random prediction scores.

• Overfitting gap moves in opposite directions. LN parameter removal widens the train-test accuracy gap in Pre-LN models (denoted Δ_overfit^Pre), exacerbating overfitting, and narrows it in Post-LN models (Δ_overfit^Post) because memorization is mitigated.

• Early-layer LNs dominate. In DeiT (Pre-LN, UTK-Face), removing early LN parameters disrupts learning most; in DeBERTa (Post-LN, Emotions), removing early LN parameters mitigates memorization and improves recovery most. The paper formalizes this as Δ_overfit^Pre, early > Δ_overfit^Pre, middle > Δ_overfit^Pre, later, and Δ_overfit^Post, early < Δ_overfit^Post, middle < Δ_overfit^Post, later.

• Gradient norms explain the divergence. Learning gradient norms (‖g_x^learn‖₂) exceed memorization gradient norms (‖g_x^mem‖₂) across all layers in both configurations, and both are higher in earlier layers. The ratio ‖g_x^learn‖₂ / ‖g_x^mem‖₂ is much higher in Pre-LN models than in Post-LN models across all layers.

• Compensation differs by architecture. In Pre-LN models the gradient norms are predominantly high in the first layer with following layers having almost negligible norms, so early LN removal cannot be compensated. In Post-LN models both norms decrease gradually over layers, so later LNs can compensate for the absence of early ones, recovering learning while mitigating memorization.

• Empirical scope. The claims are validated through 13 models across 6 vision and language datasets, with all experiments run across 3 random seeds. The precise accuracy, memorization, recovery, and random prediction values are presented graphically in the figures rather than as numbers in the text.

Methodology in Plain English

The setup is deliberately simple and controlled. The authors take each transformer and train two versions: one with LayerNorm's learnable weight and bias intact, and one where those parameters are removed while the normalization step itself still runs. Because the normalization is kept, the comparison isolates the contribution of the learnable parameters.

To create something worth memorizing, they inject random label noise: for every dataset they change the labels of 1% of the training set, reassigning those samples to a randomly chosen class different from the original. Models are then trained until they reach 100% training accuracy, at which point the noisy samples are presumed memorized. Four metrics track behavior: learning (test) accuracy, memorization score, recovery score (noisy samples whose predictions return to the genuine label), and random prediction score (noisy samples predicted as an unrelated random label rather than the true or noisy label).

For the layer-wise study, the layers are grouped into early, middle, and later sets (defined in Appendix F.4), and LN parameters are removed from one group at a time. For the mechanistic explanation, the authors compute the gradient of the loss with respect to the LayerNorm input. Averaging its L2 norm over all test samples gives the learning gradient norm; averaging over only the injected noisy samples gives the memorization gradient norm. A layer whose LN has a high gradient norm is one that strongly influences the corresponding behavior.

Datasets span both modalities: CIFAR10, NICO++, UTK-Face, Emotions, News, and TweetTopic — 3 vision and 3 language classification datasets. Post-LN models are BERT, RoBERTa, DistilBERT, DeBERTa, ELECTRA, and Longformer (6). Pre-LN models are ViT-B, ViT-S, DeiT, GPT2, GPT-Neo, Qwen2, and RoBERTa-PreLayerNorm (7). The paper notes that only the language modality is available for the Post-LN architecture in practice and in the literature. Training configurations and full dataset details appear in Appendix F, and additional results for other models appear in Appendices G.1, G.2, and G.3.

Why This Matters

Impact on research. The work reframes LayerNorm from a generic training stabilizer into a component with a specific, architecture-dependent function: a learning enabler in Pre-LN models and a memorization enabler in Post-LN models. It also qualifies earlier claims about LN and overfitting in Pre-LN models, and it complements prior work on attention heads and feed-forward networks as locations of memorization by pointing to a different component. The gradient analysis supplies a quantitative reason for why deeper layers are less effective in Pre-LN models, connecting to prior observations about deeper layers.

Potential real-world applications (the paper does not enumerate deployed applications; these follow from its findings):

• Training on datasets with unavoidable label noise, where selectively removing Early LN parameters in Post-LN models could let a model fall back on genuine labels instead of fitting corrupted ones.
• Model quantization, given prior work cited in the paper showing that LayerNorm outlier neurons hinder transformer quantization.
• Architecture selection for large language and vision models, since the choice between Pre-LN and Post-LN placement changes whether LayerNorm removal is a memorization control or a learning hazard.
• Diagnosing overfitting, using the train-test accuracy gap across layer groups as a signal of where memorization is concentrated.

Industry relevance. Modern deployed architectures — GPT, Llama, and Vision Transformers, as the paper notes — are Pre-LN, precisely the configuration where the authors find that LN parameters are critical to learning and that their removal destabilizes training and exacerbates overfitting. For practitioners, this is a caution against treating LN parameters as removable overhead in Pre-LN stacks, while for Post-LN users it identifies a targeted lever for reducing memorization without sacrificing learning.

Future Directions

• Vision Post-LN models. The paper acknowledges that only the language modality is available for Post-LN architectures in practice and literature, so the Post-LN memorization findings for vision have not been tested on a comparable model set.

• Noise levels and training regimes. The study fixes label noise at 1% of the training set and trains to 100% training accuracy across 3 random seeds; whether the Pre-LN/Post-LN divergence holds at higher noise rates or shorter training budgets is not reported.

• Other normalization and modern variants. The experiments cover LayerNorm placements in the listed 13 models; the behavior of alternative normalization schemes and of the resulting parameter-removal effects is not reported.

• Broader impacts. The paper states that its broader impacts are discussed in Appendix K, leaving scope for follow-up work on the practical and ethical consequences of controlling memorization through normalization parameters.

Target Audience

This paper suits transformer architecture researchers and interpretability researchers who want a mechanistic account of where memorization lives in a network; machine learning engineers choosing between Pre-LN and Post-LN designs or tuning normalization layers; and graduate-level readers comfortable with gradient derivations, since the core evidence rests on gradient-norm theorems supported by empirical results across vision and language tasks.

Authors’ abstract

Layer Normalization (LayerNorm) is one of the fundamental components in transformers that stabilizes training and improves optimization. In recent times, Pre-LayerNorm transformers have become the preferred choice over Post-LayerNorm transformers due to their stable gradient flow. However, the impact of LayerNorm on learning and memorization across these architectures remains unclear. In this work, we investigate how LayerNorm influences memorization and learning for Pre- and Post-LayerNorm transformers. We identify that LayerNorm serves as a key factor for stable learning in Pre-LayerNorm transformers, while in Post-LayerNorm transformers, it impacts memorization. Our analysis reveals that eliminating LayerNorm parameters in Pre-LayerNorm models exacerbates memorization and destabilizes learning, while in Post-LayerNorm models, it effectively mitigates memorization by restoring genuine labels. We further precisely identify that early layers LayerNorm are the most critical over middle/later layers and their influence varies across Pre and Post LayerNorm models. We have validated it through 13 models across 6 Vision and Language datasets. These insights shed new light on the role of LayerNorm in shaping memorization and learning in transformers.