Learning to Remove: Towards Isotropic Pre-trained BERT Embedding

TL;DR

This paper identifies the non-isotropic nature of BERT embeddings and proposes a simple method to make them more isotropic, resulting in improved performance on various NLP tasks.

Contribution

The paper introduces a novel, effective technique to enhance BERT embeddings by removing dominant directions, improving their isotropy and downstream task performance.

Findings

Processed embeddings are more isotropic after applying the method.

Significant performance improvements on word similarity, analogy, and textual similarity tasks.

Method outperforms baseline approaches and is robust to hyperparameter changes.

Abstract

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.