[Paper Review] Improving BERT Fine-tuning with Embedding Normalization

Wenxuan Zhou, Junyi Du+1 more|arXiv (Cornell University)|Nov 10, 2019

Topic Modeling9 references4 citations

TL;DR

This paper identifies a biased embedding distribution in BERT's [CLS] token—where activations concentrate on a few dimensions and lack zero-centering—causing optimization instability during fine-tuning. To address this, it proposes simple embedding normalization techniques that re-center and normalize [CLS] embeddings, leading to improved performance across multiple text classification tasks without architectural changes.

ABSTRACT

Large pre-trained sentence encoders like BERT start a new chapter in natural language processing. A common practice to apply pre-trained BERT to sequence classification tasks (e.g., classification of sentences or sentence pairs) is by feeding the embedding of [CLS] token (in the last layer) to a task-specific classification layer, and then fine tune the model parameters of BERT and classifier jointly. In this paper, we conduct systematic analysis over several sequence classification datasets to examine the embedding values of [CLS] token before the fine tuning phase, and present the biased embedding distribution issue---i.e., embedding values of [CLS] concentrate on a few dimensions and are non-zero centered. Such biased embedding brings challenge to the optimization process during fine-tuning as gradients of [CLS] embedding may explode and result in degraded model performance. We further propose several simple yet effective normalization methods to modify the [CLS] embedding during the fine-tuning. Compared with the previous practice, neural classification model with the normalized embedding shows improvements on several text classification tasks, demonstrates the effectiveness of our method.

Motivation & Objective

• To investigate the distributional properties of BERT's [CLS] token embeddings before fine-tuning.
• To identify that biased, non-zero-centered [CLS] embeddings hinder optimization during fine-tuning.
• To propose normalization techniques that stabilize training and improve performance on sequence classification tasks.
• To evaluate the effectiveness of normalized [CLS] embeddings across multiple benchmark text classification datasets.

Proposed method

• Analyzing the [CLS] token embedding distribution across multiple sequence classification datasets before fine-tuning.
• Applying normalization techniques such as batch normalization and layer normalization to the [CLS] embedding vector during fine-tuning.
• Re-centering the [CLS] embedding by subtracting its mean across the batch to achieve zero-centering.
• Modifying the input to the classification head by normalizing the [CLS] representation before feeding it into the classifier.
• Integrating normalization as a lightweight post-processing step without changing BERT's architecture or pre-training.
• Evaluating the impact of normalization on downstream classification performance using standard benchmarks.

Experimental results

Research questions

• RQ1 How are the [CLS] token embeddings distributed before fine-tuning in BERT on sequence classification tasks?
• RQ2 To what extent does the non-zero-centered, concentrated distribution of [CLS] embeddings affect fine-tuning optimization?
• RQ3 Can simple normalization of [CLS] embeddings improve fine-tuning stability and model performance?
• RQ4 How does normalization compare to standard fine-tuning across multiple text classification datasets?

Key findings

• The [CLS] token embeddings exhibit a biased distribution with high concentration on a few dimensions and are not zero-centered before fine-tuning.
• This biased distribution leads to unstable gradients and degraded performance during fine-tuning.
• Normalization techniques such as batch re-centering and layer normalization effectively stabilize the optimization process.
• The proposed method achieves consistent improvements across multiple text classification benchmarks compared to standard fine-tuning.
• The gains are achieved without architectural changes or additional parameters, making the method lightweight and deployable.