When BERT-whitening Gets Hyperparameters: There's Always a Setting That Suits You

当BERT-whitening引入超参数：总有一款适合你 苏剑林 · 2022-05-18 · Machine Learning language model semantics semantic similarity

In "You Might Not Need BERT-flow: A Linear Transformation Rivals BERT-flow"], I proposed BERT-whitening, showing that a simple linear transformation could match the performance of BERT-flow, the SOTA method at the time. In addition, BERT-whitening can also be used to reduce the dimensionality of sentence embeddings, bringing lower memory usage and faster retrieval speed. However, in "Which Unsupervised Semantic Similarity Method Is Strongest? We Ran a Fairly Comprehensive Evaluation"], we also found that the whitening operation doesn't always bring improvements — some models are already well matched to the task (such as SimBERT, which has undergone supervised training), and in those cases the extra whitening operation often actually hurts performance.

To address this shortcoming, this post proposes introducing two hyperparameters into BERT-whitening. By tuning these two hyperparameters, we can almost always achieve results where "dimensionality is reduced without sacrificing performance." In other words, even for tasks where adding whitening used to degrade performance, there is now a chance to achieve comparable or even better results while still reducing dimensionality.

Method Overview

The current BERT-whitening pipeline is:

\begin{equation}\begin{aligned} \tilde{\boldsymbol{x}}_i =&\, (\boldsymbol{x}_i - \boldsymbol{\mu})\boldsymbol{U}\boldsymbol{\Lambda}^{-1/2} \\ \boldsymbol{\mu} =&\, \frac{1}{N}\sum\limits_{i=1}^N \boldsymbol{x}_i \\ \boldsymbol{\Sigma} =&\, \frac{1}{N}\sum\limits_{i=1}^N (\boldsymbol{x}_i - \boldsymbol{\mu})^{\top}(\boldsymbol{x}_i - \boldsymbol{\mu}) = \boldsymbol{U}\boldsymbol{\Lambda}\boldsymbol{U}^{\top} \,\,(\text{SVD decomposition}) \end{aligned}\end{equation}more

where $\boldsymbol{x}_i$ is the given sentence vector (unless otherwise stated, vectors are treated as row vectors by default), and $\tilde{\boldsymbol{x}}_i$ is the transformed vector. In the SVD decomposition results, $\boldsymbol{U}$ is an orthogonal matrix and $\boldsymbol{\Lambda}$ is a diagonal matrix whose diagonal elements are non-negative and arranged in decreasing order. As we can see, the current pipeline is entirely fixed — there are no tunable hyperparameters at all.

To create some room for adjustment, we can introduce two hyperparameters $\beta,\gamma$ (both scalars), turning it into

\begin{equation}\begin{aligned} \tilde{\boldsymbol{x}}_i =&\, (\boldsymbol{x}_i - {\color{red}\beta}\boldsymbol{\mu})\boldsymbol{U}\boldsymbol{\Lambda}^{-{\color{red}\gamma}/2} \\ \boldsymbol{\mu} =&\, \frac{1}{N}\sum\limits_{i=1}^N \boldsymbol{x}_i \\ \boldsymbol{\Sigma} =&\, \frac{1}{N}\sum\limits_{i=1}^N (\boldsymbol{x}_i - {\color{red}\beta}\boldsymbol{\mu})^{\top}(\boldsymbol{x}_i - {\color{red}\beta}\boldsymbol{\mu}) = \boldsymbol{U}\boldsymbol{\Lambda}\boldsymbol{U}^{\top} \,\,(\text{SVD decomposition}) \end{aligned}\end{equation}

Rationale

As we can see, when $\beta=\gamma=1$, this reduces to the original BERT-whitening; and when $\beta=\gamma=0$, the net transformation becomes

\begin{equation}\tilde{\boldsymbol{x}}_i =\boldsymbol{x}_i \boldsymbol{U}\end{equation}

Since $\boldsymbol{U}$ is an orthogonal matrix, it doesn't change the inner product result, i.e., $\tilde{\boldsymbol{x}}_i\tilde{\boldsymbol{x}}_i^{\top} = \boldsymbol{x}_i \boldsymbol{U} (\boldsymbol{x}_i \boldsymbol{U})^{\top} = \boldsymbol{x}_i\boldsymbol{x}_i^{\top}$. So when we use cosine similarity as the similarity measure, this transformation won't change the original result at all. In other words, by introducing this pair of hyperparameters, we get the possibility of "performance no worse than before the transformation" — and then, when we fine-tune these parameters, we might even achieve results better than before the transformation. This is precisely the design idea behind these two hyperparameters.

Moreover, under this modification, the original dimensionality-reduction capability is still preserved. We can split the transformation into two parts:

\begin{equation}\tilde{\boldsymbol{x}}_i = \color{red}{\underbrace{(\boldsymbol{x}_i - \beta\boldsymbol{\mu})\boldsymbol{U}}_{\text{part 1}}}\color{skyblue}{\underbrace{\boldsymbol{\Lambda}^{-\gamma/2}}_{\text{part 2}}}\end{equation}

The first part is essentially the orthogonal transformation $\boldsymbol{U}$, where $\boldsymbol{U}$ is the result of the SVD decomposition of the matrix $\boldsymbol{\Sigma}$. It transforms the vector $\boldsymbol{x}_i - \beta\boldsymbol{\mu}$ into a new vector whose components are as independent as possible, and the average fluctuation of each component of the new vector around 0 is precisely measured by the diagonal elements of $\boldsymbol{\Lambda}^{1/2}$. If the corresponding fluctuation is very close to 0, we can regard that component as effectively being zero, and dropping it won't affect the cosine similarity computation — this is the principle behind dimensionality reduction. Since the SVD decomposition has already sorted $\boldsymbol{\Lambda}$ in decreasing order in advance, we can simply keep the first $k$ dimensions via the operation $\tilde{\boldsymbol{x}}_i\text{[:}k\text{]}$ to achieve a reduction to $k$ dimensions.

As for the second part, $\boldsymbol{\Lambda}^{-\gamma/2}$, we can interpret it as the degree to which the current task depends on isotropy. If $\gamma=1$, then every component is treated equally — this can be regarded as an unsupervised prior result, but it isn't necessarily optimal for every task, so we can tune $\gamma$ to better adapt to the task at hand.

Experimental Results

The article "Which Unsupervised Semantic Similarity Method Is Strongest? We Ran a Fairly Comprehensive Evaluation"] has already shown that, on the ATEC, BQ, and LCQMC tasks, applying the default whitening operation (i.e., $\beta=\gamma=1$) to SimBERT causes a drop in performance. But if we instead take $\beta=\gamma=0$, the picture changes (here I show just two combinations at random; other combinations give similar results):

$$\small{\begin{array}{c} \text{BERT-P4 results table} \\ {\begin{array}{l|ccccc} \hline & \text{ATEC} & \text{BQ} & \text{LCQMC} & \text{PAWSX} & \text{STS-B} \\ \hline \beta=\gamma=1 & 24.51 / \color{green}{27.00} / \color{green}{27.91} & 38.81 / \color{red}{32.29} / \color{red}{37.67} & 64.75 / \color{green}{64.75} / \color{green}{65.65} & 15.12 / \color{green}{17.80} / \color{green}{15.34} & 61.66 / \color{green}{69.45} / \color{green}{69.37} \\ \beta=\gamma=0 & 24.51 / 24.51 / \color{green}{24.59} & 38.81 / 38.81 / \color{green}{38.99} & 64.75 / 64.75 / \color{red}{63.45} & 15.12 / 15.12 / \color{red}{14.59} & 61.66 / 61.66 / \color{green}{62.30} \\ \hline \end{array}} \\ \\ \text{SimBERT-P1 results table} \\ {\begin{array}{l|ccccc} \hline & \text{ATEC} & \text{BQ} & \text{LCQMC} & \text{PAWSX} & \text{STS-B} \\ \hline \beta=\gamma=1 & 38.50 / \color{red}{23.64} / \color{red}{30.79} & 48.54 / \color{red}{31.78} / \color{red}{40.01} & 76.23 / \color{red}{75.05} / \color{red}{74.50} & 15.10 / \color{green}{18.49} / \color{green}{15.64} & 74.14 / \color{red}{73.37} / \color{green}{75.29} \\ \beta=\gamma=0 & 38.50 / 38.50 / \color{green}{38.81} & 48.54 / 48.54 / \color{green}{48.66} & 76.23 / 76.23 / \color{red}{76.22} & 15.10 / 15.10 / \color{red}{14.88} & 74.14 / 74.14 / \color{green}{74.46} \\ \hline \end{array}} \end{array}}$$

As in the previous post, each entry in the table has the form $a / b / c$, representing the score for that task under that model "without whitening" as $a$, "with whitening" as $b$, and "with whitening reduced to 256 dimensions" as $c$. If $b > a$, then $b$ is shown in green, and in red otherwise; if $c > a$, then $c$ is shown in green, and in red otherwise. As noted earlier, without dimensionality reduction, the net transformation under $\beta=\gamma=0$ is $\boldsymbol{U}$, which doesn't change the cosine similarity result — so under $\beta=\gamma=0$, all values of $a,b$ are equal.

In this table, we're mainly interested in the third result in $a/b/c$, namely $c$, which is the result of reducing the vector from 768 dimensions down to 256 dimensions. As we can see, when $\beta=\gamma=0$, whether for the unsupervised BERT or the supervised SimBERT, this result is essentially very close to that of the original vector (i.e., $a$), and some results are even better. This means that the combination $\beta=\gamma=0,k=256$ can basically be considered a "free lunch" — performance is nearly unaffected while dimensionality reduction is achieved.

I also tried fine-tuning $\beta,\gamma$, and on some tasks this did indeed yield better results than the two combinations above; but fine-tuning requires labeled data, which could be more contentious, so I won't demonstrate it here. If the original sentence-embedding model was already obtained via supervised training, and BERT-whitening is being applied purely for the sake of dimensionality reduction, then it's entirely reasonable to use a validation set to fine-tune $\beta,\gamma$ and $k$ — in that scenario there's no controversy at all.

Summary

This post introduces two hyperparameters to give BERT-whitening some room for tuning, giving it the potential for "performance no worse than before the transformation" while still retaining its dimensionality-reduction capability. In other words, even for sentence-embedding models that have already been trained, we can now use this new version of BERT-whitening to reduce their dimensionality while keeping performance essentially unchanged — and sometimes even improving it.