[Paper Review] IsoBN: Fine-Tuning BERT with Isotropic Batch Normalization

Wenxuan Zhou, Bill Yuchen Lin+1 more|arXiv (Cornell University)|May 2, 2020

Topic Modeling4 citations

TL;DR

This paper proposes Isotropic Batch Normalization (IsoBN), a novel regularization method that enhances BERT and RoBERTa fine-tuning by dynamically penalizing dominant principal components to produce more isotropic [CLS] embeddings. By modeling correlations as block-diagonal instead of ignoring them, IsoBN improves performance by approximately 1.0 absolute F1/Spearman score point across seven GLUE tasks, outperforming standard batch normalization.

ABSTRACT

Fine-tuning pre-trained language models (PTLMs), such as BERT and its better variant RoBERTa, has been a common practice for advancing performance in natural language understanding (NLU) tasks. Recent advance in representation learning shows that isotropic (i.e., unit-variance and uncorrelated) embeddings can significantly improve performance on downstream tasks with faster convergence and better generalization. The isotropy of the pre-trained embeddings in PTLMs, however, is relatively under-explored. In this paper, we analyze the isotropy of the pre-trained [CLS] embeddings of PTLMs with straightforward visualization, and point out two major issues: high variance in their standard deviation, and high correlation between different dimensions. We also propose a new network regularization method, isotropic batch normalization (IsoBN) to address the issues, towards learning more isotropic representations in fine-tuning by dynamically penalizing dominating principal components. This simple yet effective fine-tuning method yields about 1.0 absolute increment on the average of seven NLU tasks.

Motivation & Objective

• To investigate the isotropy of pre-trained [CLS] embeddings in BERT and RoBERTa, particularly their variance and correlation structure.
• To identify that high variance in standard deviation and strong correlations between dimensions hinder fine-tuning performance.
• To develop a practical regularization method that improves isotropy during fine-tuning without incurring high computational cost.
• To demonstrate that reducing the dominance of top principal components leads to faster convergence and better generalization in downstream NLU tasks.

Proposed method

• Propose isotropic batch normalization (IsoBN), which normalizes [CLS] embeddings using both standard deviation and pairwise correlation coefficients.
• Model the absolute correlation coefficient matrix as block-diagonal, assuming groups of highly correlated dimensions, rather than assuming diagonal (as in standard batch norm).
• Apply dynamic scaling to suppress the contribution of dominant principal components by penalizing high-variance and highly correlated dimensions.
• Integrate IsoBN into the fine-tuning pipeline of BERT and RoBERTa, replacing standard batch normalization layers.
• Use explained variance (EV_k) to quantitatively measure isotropy improvement, comparing pre-normalization, batch norm, and IsoBN.
• Train models with half-precision and evaluate on GLUE benchmark tasks to ensure compatibility and efficiency.

Experimental results

Research questions

• RQ1 How isotropic are the pre-trained [CLS] embeddings in BERT and RoBERTa, in terms of variance and correlation across dimensions?
• RQ2 To what extent do dominant principal components in [CLS] embeddings limit fine-tuning performance?
• RQ3 Can a modified batch normalization method that accounts for inter-dimensional correlations improve isotropy and downstream performance?
• RQ4 Does IsoBN outperform standard batch normalization in reducing explained variance of top principal components?

Key findings

• Pre-trained [CLS] embeddings in BERT and RoBERTa exhibit high variance in standard deviation and strong correlations between dimensions, indicating poor isotropy.
• The top three principal components explain approximately 86% of the variance on average, with some datasets (e.g., STS-B) exceeding 90%.
• Standard batch normalization reduces explained variance (EV₃) by only 0.025 (BERT) and 0.148 (RoBERTa), due to ignoring correlations.
• IsoBN reduces EV₃ to 0.123 (BERT) and 0.425 (RoBERTa), significantly improving isotropy compared to standard batch norm.
• IsoBN yields an average absolute improvement of 1.0 F1/Spearman score point across seven GLUE benchmark tasks.
• The method is stable, efficient, and compatible with half-precision training, making it practical for real-world fine-tuning.