Source-linked AI summary
IsoBN: Fine-Tuning BERT with Isotropic Batch Normalization
Wenxuan Zhou, Bill Yuchen Lin, Xiang Ren
TL;DR
The isotropy of pre-trained [CLS] embeddings is under-explored, despite reported downstream benefits of isotropic representations. The paper analyzes their variance and correlations and proposes IsoBN, reporting about a 1.0-point average increment across seven NLU tasks.
Problem
The isotropy of pre-trained [CLS] embeddings is under-explored, while the paper identifies unequal variance and correlated dimensions in those embeddings.
Source: §1 · p. 1Method
IsoBN regularizes fine-tuning by modeling correlated embedding dimensions in groups and dynamically penalizing dominating principal components.
Source: §1 · p. 2, §3.2 · p. 3, §4.2 · p. 5Results
Takeaways & Limitations
IsoBN produces more isotropic embeddings than conventional batch normalization and reduces the computation effort of calculating the inverse.
Source: §1 · p. 2Whitening is numerically unstable, computationally expensive, and incompatible with half-precision training because PTLM embedding covariance matrices are ill-conditioned.
Source: §4.1 · p. 4
Abstract
from arXiv · showhide
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.
1 Introduction
The introduction frames isotropy of pre-trained [CLS] embeddings as an open question, identifies unequal variance and correlated dimensions, and proposes IsoBN to regularize fine-tuning. It reports improved isotropy and about a 1.0 absolute average increment across seven GLUE tasks.
1 Introduction
- Prior representation-learning work reports downstream benefits from isotropic embeddings, whose dimensions have equal variance and are uncorrelated.The introduction also links conventional feature normalization to faster convergence and better generalization.Source: §1 · p. 1
- The paper treats the isotropy of pre-trained [CLS] embeddings as an open question and studies how to improve it for downstream fine-tuning.It describes this work as the first study of isotropy in pre-trained [CLS] embeddings.Source: §1 · pp. 1–2
- Pre-trained [CLS] embeddings show very different variance and clusters of highly correlated dimensions, indicating they are far from isotropic.The paper identifies unit variance and uncorrelatedness as the two aspects of isotropic embeddings.Source: §1 · p. 1
- Whitening requires an ill-conditioned covariance inverse, making it numerically unstable, computationally expensive, and incompatible with half-precision training; standard batch normalization ignores inter-dimension correlation.These constraints motivate a different fine-tuning regularizer.Source: §1 · p. 2
- IsoBN models correlated dimensions as groups, assuming absolute correlation 1 within a group and 0 between groups.This replaces the diagonal-only correlation assumption used by conventional batch normalization.Source: §1 · p. 2
- About 1.0 absolute increment on average across seven GLUE tasks is reported for IsoBN, which also produces more isotropic embeddings than conventional batch normalization.The introduction reports improvements for both BERT and RoBERTa.Source: §1 · p. 2
2 Why isotropic [CLS] embeddings?
The section argues that classifier fine-tuning barely changes its random initialization and that leading principal components dominate logits. This motivates dynamically penalizing those components so fine-tuning can use knowledge in other components.
2 Why isotropic [CLS] embeddings?
- With random weight initialization, the first few principal components usually contribute most of the prediction results.This finding is presented as part of the paper's analysis of [CLS] embeddings.Source: §2 · p. 2
- The fine-tuned classifier weight remains nearly at its initialization; on RoBERTa CoLA, their cosine similarity is 0.9997.The comparison uses cosine similarity and L2 distance after projecting weights onto the top 10 eigenvectors of Cov(h).Source: §2 · p. 2, Figure 2 · p. 2
- The analysis compares each principal component's logit variance under the fine-tuned classifier with that under an optimal classifier fitted on all training data.The optimal classifier is logistic regression; the paper treats variance along informative components as desirable and variance along irrelevant ones as undesirable.Source: §2 · p. 2
- The first principal component accounts for over 90% of logit variance throughout training, while the top components are consistently exaggerated.Figure 3 summarizes average logit-variance percentages for the top 1, 5, and 30 principal components.Source: §2 · p. 2, Figure 3 · p. 3
- The top-component dominance motivates dynamically penalizing those components during fine-tuning to avoid losing knowledge from the PTLM.The paper says the exaggeration is especially pronounced at the beginning of training.Source: §2 · pp. 2–3
3 How isotropic are [CLS] embeddings?
The analysis evaluates pretrained [CLS] embedding isotropy on four NLU tasks using standard deviations and Pearson correlations. Both measures show that BERT and RoBERTa embeddings are far from isotropic.
3 How isotropic are [CLS] embeddings?
- The analysis covers pretrained [CLS] embeddings from MRPC, RTE, CoLA, and STS-b, examining standard deviation and Pearson correlation.The authors visualize the standard-deviation distributions and cluster dimensions by pairwise correlation for the correlation matrices.Source: §3 · p. 3, §3.1 · p. 3, §3.2 · p. 3
3.1 Analysis of Standard Deviation
- Both BERT and RoBERTa show widely varying [CLS] standard deviations; RoBERTa is better than BERT on this measure but remains far from isotropic.BERT’s standard deviations range from 10^-10 to 1 on MRPC.Source: Figure 4 · p. 3, §3.1 · p. 3
3.2 Analysis of Correlation Coefficient
- Both models usually show high inter-dimension correlations, with BERT forming several large clusters and RoBERTa tending toward one extreme cluster.The correlation pattern is less severe for BERT on CoLA and RoBERTa on RTE.Source: §3.2 · p. 3, Figure 5 · p. 4
- The authors argue that these variance and correlation patterns indicate non-isotropy, whose prior bias may contribute to suboptimal fine-tuning performance.They connect high correlations to classifier bias toward correlated features, which may require more training data or epochs for downstream convergence and generalization.Source: §3.2 · pp. 3–4
4 Approach
The paper introduces IsoBN to improve the isotropy of [CLS] representations during PTLM fine-tuning. It motivates the method by contrasting whitening and batch normalization with a correlation-aware, grouped-dimension approach.
4 Approach
- IsoBN is proposed to learn more isotropic [CLS] representations during PTLM fine-tuning by explicitly accounting for correlations between feature dimensions.Its key difference from conventional batch normalization is the introduction of γ.Source: §4 · p. 4, §4.2 · p. 5
4.1 Whitening and Batch Normalization
- Whitening requires inverting an ill-conditioned covariance matrix, making it numerically unstable, computationally expensive, and incompatible with half-precision training.Batch normalization simplifies the inverse by assuming a diagonal covariance matrix, but ignores embedding correlations that matter for [CLS] representations.Source: §4.1 · p. 4
4.2 Isotropic Batch Normalization
- IsoBN assumes the absolute correlation matrix is block-diagonal, with highly correlated dimensions grouped together and dimensions across groups uncorrelated.The paper presents this as a more realistic approximation of correlation structure than batch normalization’s diagonal assumption.Source: §4.2 · pp. 4–5, Figure 6 · p. 4
- Because embedding dimensions do not form hard groups, IsoBN uses the correlation matrix to estimate group size and introduces γ to account for inter-dimension correlations.Under the stated assumption, the estimate matches the group size; the method then applies a correlation-aware scaling transformation.Source: §4.2 · p. 5
- The implemented IsoBN is applied before the final classifier, removes the unstable mean, rescales to preserve total variance, and uses β to control normalization strength.It updates moving covariance and standard-deviation statistics during training because estimates from a single batch are inaccurate.Source: §4.2 · p. 5
5 Evaluation
The evaluation tests IsoBN on seven GLUE tasks and measures embedding isotropy with explained variance (EVk).
5.1 Experiment Setup
- Five random seeds are used; early stopping and hyperparameter selection rely on task-specific development-set metrics.The setup also specifies AdamW, a learning-rate schedule, and IsoBN hyperparameter ranges.Source: §5.1 · p. 5
5.2 Experimental Results
- IsoBN improves evaluation metrics on all seven GLUE datasets, increasing average scores by 1% for BERT-base and 0.8% for RoBERTa-large.Table 1 reports dev-set results from five random seeds as medians and standard deviations.Source: §5.2 · p. 5, Table 1 · p. 6
- IsoBN improves average performance more on small than large datasets: 1.6% versus 0.15% for BERT and 1.3% versus 0.25% for RoBERTa.Small datasets are MRPC, RTE, CoLA, and STS-B; large datasets are MNLI, QNLI, and SST-2.Source: §5.2 · p. 5
5.3 Experiments of EVk Metric
- EVk measures variance differences across embedding directions; smaller values indicate more evenly distributed variation and greater angular symmetry.The metric is used to assess embedding isotropy.Source: §5.3 · p. 6
- For EV3, IsoBN reduces explained variance to 0.123 for BERT and 0.425 for RoBERTa, versus batch-normalization reductions of 0.025 and 0.148.Before normalization, average EV3 is around 0.86 for both models; on some datasets, the top three components explain over 90% of variance.Source: §5.3 · p. 6, Table 2 · p. 6
6 Related Work
The related work contrasts normalization approaches, prior fine-tuning strategies, and studies of representation spectra with this paper’s focus on [CLS] embeddings.
Input Normalization.
- Batch normalization ignores correlations between input dimensions, while IsoBN uses a block-diagonal assumption to simplify whitening and scale dimensions against dominating components.Figure 7 compares EVk for no normalization, batch normalization, and IsoBN; the passage reports lower EVk with IsoBN than with batch normalization.Source: Input Normalization · pp. 6–7, Figure 7 · p. 7
Fine-tuning Language Models.
- Prior fine-tuning strategies include adversarial training, gradual unfreezing, and multitasking; this work studies isotropy in fine-tuning language models.The paper identifies its study of text-representation isotropy in fine-tuning as the first of its kind to the authors’ knowledge.Source: Fine-tuning Language Models · p. 7
Spectral Control.
- Prior spectrum studies address word embeddings, contextual knowledge, and text-generation degradation; this work studies [CLS] spectra and controls top components to improve fine-tuning.The cited prior work includes Mu, Bhat, and Viswanath (2018), Ethayarajh (2019), and Gao et al. (2019) and Wang et al. (2020).Source: Spectral Control · p. 7
7 Conclusion
The paper studies isotropy in pretrained [CLS] embeddings and proposes IsoBN, reporting improved fine-tuning results on BERT and RoBERTa.
7 Conclusion
- The paper analyzes pretrained [CLS] embedding isotropy through visualizations of standard deviation and correlation coefficients.The authors describe the analysis as a contribution of the work.Source: §7 · p. 7
- IsoBN stably improves fine-tuning of BERT and RoBERTa toward more isotropic representations, with an absolute increment around 1.0 point on seven NLU tasks.The authors identify more isotropic pretraining and fine-tuning as directions for future research.Source: §7 · p. 7