Whitening Not Recommended for Classification Tasks in LLMs

Sentence embedding plays a fundamental role in NLP Le and Mikolov (2014). Despite the widespread success of Large Language Models (LLMs) in generative tasks, embeddings obtained from pre-trained models are not impressive Li and Li (2023). Sometimes, they are not even competitive with traditional word2vec-based approaches on machine learning tasks such as classification and Semantic Text Similarity (STS). Consequently, there has been a flurry of research aimed at improving the quality of embeddings from pre-trained models Gao et al. (2021); Jiang et al. (2022); Li and Li (2023).

Among this group of work, whitening has been shown to be an effective post-processing method for improving embeddings obtained from LLMs Zhuo et al. (2023); Su et al. (2021); Huang et al. (2021). We find that the efficacy of whitening is both model-dependent and task-dependent. Although we reproduced the result that whitening does work for some models on STS tasks, it does not work for other models. More importantly, the effectiveness of the whitening operation is restricted to STS tasks. For classification tasks, whitening degrades embedding quality consistently and sometimes with a large margin. The result is supported consistently for all the evaluated models and all the datasets in SentEval Conneau and Kiela (2018). To further consolidate the surprising results, we explored a variety of whitening operations, including Principal Component Analysis (PCA) Friedman (1987), Cholesky matrix decomposition Siarohin et al. (2018), and Zero-Phase Component Analysis (ZCA) Bell and Sejnowski (1997). Although some variants of whitening induce different performances, the overall conclusion remains unchanged.

A by-product of our research is an embedding evaluation platform for LLMs, which we call SentEval⁺, to streamline the evaluation of embedding quality. LLMs are big and costly to run. SentEval Conneau and Kiela (2018) provides a platform for embedding evaluation on a variety of models, tasks, and datasets. It works well on smaller models such as BERT. To facilitate the evaluation of LLMs on commodity machines, we provide the embeddings for all sentences in our evaluation datasets.

There is not much detailed comparison of the performance of embeddings from OpenAI, maybe partially due to the cost for API calls. We observe that embeddings from OpenAI are on par with LLaMA overall. Another interesting observation is that LLaMA and LLaMA2 are very close in terms of embedding performance.

Our work is important for both practitioners and researchers in LLMs. For LLM providers such as openAI, various post-processing are commonly applied to the embeddings they serve. They may want to serve different types of embeddings for different tasks, with the understanding of our result. For researchers in the area, running on a variety of LLMs is prohibitive computationally. Our SentEval⁺ makes experiments feasible on commodity machines.

LLM embeddings have the isotropy problem Timkey and van Schijndel (2021); Kovaleva et al. (2021); Rudman et al. (2022). Whitening is a post-processing technique that converts spatially correlated, anisotropic feature representations into uncorrelated, isotropic ones Sasaki et al. (2023); Rudman and Eickhoff (2024). For this purpose, whitening transforms the feature representations such that the mean is centred at the origin, covariances are eliminated, and the variance is normalized to an identity matrix.

Given $N$ number of sentence embeddings $x_{1},x_{2},\dots,x_{N}$. Let $X=(x_{1},x_{2},\dots,x_{N})^{T}\in\mathbb{R}^{N\times d}$, where $d$ is the dimension of the embeddings. The covariance matrix for $X$ is $\Sigma=(X-\mu)(X-\mu)^{T}$, where $\mu$ is the mean of $\{x_{i}\}_{i=1}^{N}$. Whitening transformation is achieved using a matrix $W$resulting in unit diagonal “white” covariance $var(Z)=I$:

$W$ in Equation 1 varies as in Equation 2. The most commonly used whitening operation is called PCA-whitening, which is also the one used in the first a few papers on the performance gain of whitening on LLMs. Since our initial result on PCA-whitening shows the opposite for classification tasks, and Wang and Wu (2023) reported different behaviour of ZCA-whitening, we exhaustively investigate all variations of whitening operations.

In Equation 2, $\Lambda$ is the eigenvectors, and $U$ is the eigenvalues of the covariance matrix, i.e., $\Sigma=U\Lambda U^{T}$. The matrix $L$ corresponds to the Cholesky decomposition of the inverse of $\Sigma$, such that $LL^{T}=\Sigma^{-1}$. The matrices $V$ and $\Theta$ result from the eigen decomposition of the correlation matrix $P$, expressed as $P=V\Theta^{-\frac{1}{2}}V^{T}$, where $V$ is the eigenvector matrix and $\Theta$ contains the corresponding eigenvalues.

We experimented with 8 models on classification and STS tasks. The embeddings are extracted from the last layer of the BERT and LLaMA models, following the practice described in Reimers and Gurevych (2019). We also explored other pooling strategies and observed similar pattern. Embeddings of SBert, AnglE, and SimCSE are generated using their provided frameworks. While AnglE and SimCSE typically use the CLS pooling method to extract embeddings, which involves using the output of the ’CLS’ token from the model to represent the entire input sequence, SimCSE employs the mean pooling method instead. For all mentioned models, we used the original tokenizers. For generating ChatGPT embeddings, we choose the recent text-small-3-embeddings.

Next, we employ the SentEval setting to evaluate the embeddings. The classification setup involves using an MLP (Multi-Layer Perceptron) classifier with no hidden layers, utilizing the RMSprop optimizer. We also experimented with other classifiers including logistic regression, SVM, and Random Forests. Although the accuracy of the classification varies, the overall conclusion remains the same. Following the practice in SentEval, we report accuracy instead of F1 because the datasets are balanced.

We show that the performance of whitening is model-dependent and task-dependent. For classification tasks, we do not recommend to apply whitening. For STS tasks, the performance varies from model to model. We conjecture that it works only for LLMs before fine-tuning. Also, the technical details of ChatGPT remain to be a mystery. Based on its reaction to the whitening operation, we can infer that it may be fine-tuned, probably using NLI data. Another contribution of our work is an embedding evaluation platform for LLMs.