[Paper Review] Understanding and Improving Layer Normalization

Jingjing Xu, Xu Sun+3 more|arXiv (Cornell University)|Nov 16, 2019

Natural Language Processing Techniques175 citations

TL;DR

The paper shows that derivatives of mean and variance, not forward normalization, drive LayerNorm's effectiveness; it introduces AdaNorm, replacing bias/gain with an adaptive transformation, improving performance on most tasks.

ABSTRACT

Layer normalization (LayerNorm) is a technique to normalize the distributions of intermediate layers. It enables smoother gradients, faster training, and better generalization accuracy. However, it is still unclear where the effectiveness stems from. In this paper, our main contribution is to take a step further in understanding LayerNorm. Many of previous studies believe that the success of LayerNorm comes from forward normalization. Unlike them, we find that the derivatives of the mean and variance are more important than forward normalization by re-centering and re-scaling backward gradients. Furthermore, we find that the parameters of LayerNorm, including the bias and gain, increase the risk of over-fitting and do not work in most cases. Experiments show that a simple version of LayerNorm (LayerNorm-simple) without the bias and gain outperforms LayerNorm on four datasets. It obtains the state-of-the-art performance on En-Vi machine translation. To address the over-fitting problem, we propose a new normalization method, Adaptive Normalization (AdaNorm), by replacing the bias and gain with a new transformation function. Experiments show that AdaNorm demonstrates better results than LayerNorm on seven out of eight datasets.

Motivation & Objective

• Investigate what makes LayerNorm effective beyond forward input normalization.
• Assess the roles of bias and gain in LayerNorm and their impact on overfitting.
• Analyze how backward gradients are affected by the derivatives of mean and variance.
• Propose AdaNorm to replace bias/gain with an adaptive transformation and evaluate its performance.

Proposed method

• Reproduce LayerNorm variants and compare with a w/o Norm baseline across multiple tasks.
• Introduce DetachNorm to detach derivatives of mean and variance and measure impact on forward/backward effects.
• Theoretically analyze gradient re-centering and re-scaling in LayerNorm via mean/variance derivatives (Theorem 1).
• Propose AdaNorm with a transformation phi(y) applied to normalized features, ensuring differentiability and bounded average output (Theorem 2).
• Empirically compare LayerNorm, LayerNorm-simple, DetachNorm, and AdaNorm on eight datasets spanning MT, LM, classification, parsing, OCR, and NLP tasks.

Experimental results

Research questions

• RQ1 What components drive LayerNorm's effectiveness: forward normalization or gradient effects from mean/variance?
• RQ2 Do bias and gain contribute positively, or do they cause overfitting across tasks?
• RQ3 Can gradient normalization (via derivatives of mean/variance) explain LayerNorm's training behavior and performance?
• RQ4 Does an adaptive normalization (AdaNorm) outperform LayerNorm by replacing fixed affine parameters with input-dependent scaling?

Key findings

• Forward normalization has limited explanatory power for LayerNorm's success; derivatives of the mean and variance are more important.
• Detaching derivatives of mean/variance (DetachNorm) degrades performance, indicating gradient normalization aids optimization, especially in deeper models.
• LayerNorm-simple (no bias/gain) matches or exceeds LayerNorm on several datasets and can achieve state-of-the-art in En-Vi MT.
• Bias and gain tend to increase overfitting and are often ineffective across tasks.
• AdaNorm, which replaces bias/gain with an adaptive phi(y), outperforms LayerNorm on seven of eight datasets, showing better generalization.