Source-linked AI summary
Learning to Remove: Towards Isotropic Pre-trained BERT Embedding
Yuxin Liang, Rui Cao, Jie Zheng, Jie Ren, Ling Gao
TL;DR
Pre-trained BERT embeddings are anisotropic, with vectors clustered in narrow cones that can weaken their representational capacity. The paper learns weights to remove dominant directions and reports better performance than three baselines in most conditions across three evaluation tasks.
Problem
Pre-trained BERT embeddings are anisotropic, with vectors clustering in narrow cones and potentially reducing representation capacity and downstream-task accuracy.
Source: §1 · p. 2Method
The method learns weights on word-similarity tasks to control the removal of dominant directions in BERT embeddings.
Source: §1 · p. 2Results
The weighted-removal method outperforms three baselines in most conditions across word similarity, word analogy, and semantic textual similarity.
Source: §6 · p. 11Takeaways & Limitations
As the number of removed directions grows beyond a threshold, weighted removal generally stays steady or improves while ABTT performance drops.
Source: §5.3 · p. 11Further work is proposed on other influential factors, alternative training tasks, and combining the method with the Transformer architecture.
Source: §6 · p. 11
Abstract
from arXiv · showhide
Pre-trained language models such as BERT have become a more common choice of natural language processing (NLP) tasks. Research in word representation shows that isotropic embeddings can significantly improve performance on downstream tasks. However, we measure and analyze the geometry of pre-trained BERT embedding and find that it is far from isotropic. We find that the word vectors are not centered around the origin, and the average cosine similarity between two random words is much higher than zero, which indicates that the word vectors are distributed in a narrow cone and deteriorate the representation capacity of word embedding. We propose a simple, and yet effective method to fix this problem: remove several dominant directions of BERT embedding with a set of learnable weights. We train the weights on word similarity tasks and show that processed embedding is more isotropic. Our method is evaluated on three standardized tasks: word similarity, word analogy, and semantic textual similarity. In all tasks, the word embedding processed by our method consistently outperforms the original embedding (with average improvement of 13% on word analogy and 16% on semantic textual similarity) and two baseline methods. Our method is also proven to be more robust to changes of hyperparameter.
1 Introduction
Pre-trained BERT embeddings are anisotropic, clustering in narrow cones in ways that can weaken word representations. The paper analyzes this geometry and proposes learned weighted removal of dominant directions, evaluated against existing embeddings and post-processing baselines.
1 Introduction
- BERT embeddings are anisotropic: word vectors cluster in narrow cones, which can make unrelated words highly similar and affect downstream-task accuracy.The paper also links contextual embedding anisotropy to low-frequency words in the corpus.Source: §1 · p. 2
- ABTT removes a common vector and PCA-derived dominant directions, but its effectiveness decreases as more directions are removed, risking loss of useful linguistic information.The paper presents this as motivation for adjusting the degree of removal.Source: §1 · p. 2
- The paper measures BERT embedding geometry and reports a strong correlation between word-vector norm and average cosine similarity and word frequency.It also provides an intuitive explanation and theoretical analysis of these phenomena.Source: §1 · p. 2
- The proposed weighted removal method learns a weight for each dominant direction from word-similarity tasks to adjust how much is removed.The approach is presented as more flexible than removing directions completely.Source: §1 · p. 2
- Across word similarity, word analogy, and semantic textual similarity, the method outperforms the original BERT embedding, ABTT, and conceptor negation in most cases and remains relatively stable as hyperparameters change.The paper also reports that the method makes embeddings more isotropic and expressive.Source: §1 · p. 3
2 Related Work
Related work examines anisotropy in static and contextual embeddings and proposes geometric corrections. This paper distinguishes its approach by using word-similarity information to learn weighted removal, rather than relying only on geometric features.
2 Related Work
- BERT maps token embeddings together with position and segment information through Transformer encoders; this study focuses on the geometry of its pre-trained input embeddings.BERT is described as based on 12- to 24-layer Transformer encoders and trained on a hybrid corpus.Source: §2 · p. 3
- For static embeddings, Mu et al. address anisotropy by removing the common vector and dominant directions to capture stronger linguistic regularities; alternatives include conceptors, manifold learning, and isotropic-noise alignment.Source: §2 · p. 3
- Prior work finds anisotropy in contextual representations and reports frequency bias in contextual word embeddings.Gong et al. use adversarial training to learn a frequency-agnostic representation.Source: §2 · p. 3
- Other contextual-embedding studies describe degeneration into narrow cones and propose cosine-similarity regularization or spectrum control.The cited work also includes studies by Reif et al., Karve et al., and Zhou et al. [16].Source: §2 · p. 3
- Unlike post-processing based only on geometry, this method uses word similarity to make the correction semantically informed, without retraining the model.The paper describes its post-processing as more computationally efficient than methods that regularize word similarity during training.Source: §2 · pp. 3–4
- Compared with ABTT, the method learns weights for each dominant direction from word-similarity tasks and is evaluated against ABTT and conceptor negation.The paper identifies conceptor negation as a successor to ABTT.Source: §2 · p. 4
3 Observation
Pre-trained BERT embeddings are anisotropic: their mean is displaced, pairwise cosine similarity is well above zero, and singular values decay exponentially. The section relates these patterns to word frequency and to training losses that disproportionately affect rare words.
3.1 Anisotropic geometry of pre-trained BERT embedding
- The mean vector exceeds half the average word-vector norm, while average cosine similarity is well above zero, placing vectors in a narrow cone.The section also reports that word vectors are not centered around the origin.Source: §3.1 · p. 4
- BERT embedding singular values decay exponentially, indicating an imbalance across directions.The singular-value plot uses semi-logarithmic coordinates.Source: §3.1 · p. 4, Figure 1 · p. 5
- Norm and average cosine similarity correlate about −0.7 with log word count; the first PCA coefficient also encodes frequency at about −0.7.Figure 2 plots these frequency relationships and colors top-two-PCA projections by word frequency.Source: §3.1 · p. 5, Figure 2 · p. 5
3.2 Insights
- The analysis splits each token’s loss by whether its context contains that token; for most words, the non-containing-context term dominates, aligning vectors and raising mean cosine similarity.The cited analysis says the optimum is close to that term when its context probability is much larger.Source: §3.2 · p. 6
- Low-frequency words are affected more by the non-containing-context loss and are likely to cluster together farther from the origin.The proposed intuition is that these words encounter contexts without the token more often.Source: §3.2 · p. 6
4 Method
Weighted removal identifies dominant embedding directions with PCA, learns a weight for each, and removes them in proportion to those weights. The weights are trained on word similarity, avoiding the need to remove every selected direction completely.
4 Method
- The method uses learnable weights because completely removing dominant directions may discard semantic and syntactic information.The weights allow the removal ratio for each direction to be adjusted.Source: §4 · p. 6
- Weighted removal (WR) uses PCA to identify dominant directions, learns direction-specific weights, then subtracts weighted directions from word embeddings.The algorithm returns the processed embedding after removing the selected directions with their learned weights.Source: §4 · p. 7
- WR learns its weights on word similarity by minimizing the mean squared error between predicted and target similarities.The method uses normalized word vectors to predict similarity.Source: §4 · p. 7
- WR is compared with ABTT’s complete removal of top directions and CN’s suppression of high-variance features; it skips mean removal because its output is naturally centered.The comparison includes ABTT and conceptor negation (CN).Source: §4 · p. 7
5 Experiment
The experiments test WR against original BERT, ABTT, and CN on word similarity, analogy, and semantic textual similarity, alongside embedding-geometry analyses. WR improves the reported task outcomes and produces more isotropic embeddings; as d grows, its performance is steadier than ABTT’s.
5 Experiment
- Experiments used eight word-similarity datasets, tested d from 1 to 200, and compared WR with original BERT, ABTT, and CN; word similarity used a 70/30 split and Pearson correlation.Human annotations were scaled to [-1, 1].Source: §5 · pp. 7–8, §5.1 · p. 8
5.1 Towards Expressive
- WR improves all three evaluated task families: word-similarity performance beats original BERT and CN, analogy averages 13.04% improvement over original vectors, and STS averages 16.78% over original embeddings.Figure 3 covers the three tasks for BERT-base and BERT-large; Tables 3 and 4 report analogy and STS results, respectively.Source: §5.1 · pp. 8–10, Figure 3 · p. 9, Table 3 · p. 9
5.2 Towards Isotropy
- After WR processing, singular values become steadier across directions, vectors are centered around the origin and more uniformly distributed by frequency, and average cosine similarity drops near zero.Figure 4a shows singular values across d; Figure 4b projects embeddings processed with d=5 onto the top two PCA directions and colors words by frequency.Source: §5.2 · p. 10, Figure 4 · p. 11
5.3 Discussion
- For small d, WR and ABTT perform similarly and learned weights are near 1; at larger d, ABTT performance drops while WR generally stays steady or improves.The authors attribute WR’s relative stability to learning a weight for each direction to retain useful information.Source: §5.3 · p. 11
6 Conclusion
The paper addresses BERT embedding anisotropy with weighted removal of PCA-derived dominant directions, trained on word similarity tasks. The method outperforms three baselines in most conditions across three evaluation tasks.
6 Conclusion
- The weighted removal method outperforms three baselines in most conditions across word similarity, word analogy, and semantic textual similarity.The weights are trained on word similarity tasks.Source: §6 · p. 11
- The method removes several PCA-derived dominant directions using a learnable weight for each direction.Weights are trained on word similarity tasks.Source: §6 · p. 11
- Future work could examine other impactful embedding factors, alternative training tasks, and combining the method with transformer architecture.The paper presents these as possibilities for future investigation.Source: §6 · p. 11