You Might Not Need BERT-flow: A Linear Transformation Rivaling BERT-flow
BERT-flow comes from the paper On the Sentence Embeddings from Pre-trained Language Models, accepted at EMNLP 2020, which mainly uses a flow model to correct the distribution of sentence embeddings produced by BERT, making the resulting cosine similarities more reasonable. Given my habit of regularly browsing Arxiv, I actually came across this paper as soon as it was posted there, but wasn't particularly interested at the time. Unexpectedly, it became somewhat popular recently, and in a short span of time quite a few interpretations appeared on WeChat public accounts, Zhihu, and elsewhere—I believe readers have more or less had it flood their feeds by now.
Judging from the experimental results, BERT-flow indeed achieves a new SOTA. But my first reaction to this result was: something doesn't feel right! To be clear, I'm not saying there's a problem with the results, but rather that, based on my understanding, a flow model is unlikely to be playing such a decisive role. Following this hunch, I did some analysis, and sure enough, I found that although BERT-flow's approach is not flawed, a simple linear transformation can achieve comparable results—the flow model turns out not to be that critical.
The Assumption Behind Cosine Similarity
Generally speaking, when comparing or retrieving semantic similarity, we compute a sentence embedding for each sentence and then compare or rank them by the cosine of the angle between them. So have we ever considered this question: what assumption does cosine similarity impose on the input vectors? Or put differently, what condition must vectors satisfy for cosine similarity to work well as a comparison metric?
We know that the geometric meaning of the inner product of two vectors $\boldsymbol{x},\boldsymbol{y}$ is "the product of their respective magnitudes and the cosine of the angle between them." So cosine similarity is the inner product of two vectors divided by their respective magnitudes, and the corresponding coordinate-based formula is
\begin{equation}\cos(\boldsymbol{x},\boldsymbol{y}) = \frac{\sum\limits_{i=1}^d x_i y_i}{\sqrt{\sum\limits_{i=1}^d x_i^2} \sqrt{\sum\limits_{i=1}^d y_i^2}}\label{eq:cos}\end{equation}
However, let's not forget one thing: the equality above only holds under an "orthonormal basis." In other words, the "cosine of the angle" between vectors has a clear geometric meaning in itself, but the right-hand side of the formula above is merely an operation on coordinates—and coordinates depend on the chosen basis. Different bases give different coordinate formulas for the inner product, and hence different coordinate formulas for the cosine value.
Therefore, assuming BERT sentence embeddings already contain sufficient semantic information (e.g., the original sentence could in principle be reconstructed from them), if using formula $\eqref{eq:cos}$ to compute cosine values performs poorly for comparing sentence similarity, the reason might be that the coordinate system the sentence embeddings live in is not an orthonormal basis. So how do we know which specific basis it uses? In principle there's no way to know for certain, but we can make an educated guess. The basis for this guess is that when we choose a basis for a set of vectors, we tend to use each basis vector as evenly as possible. From a statistical standpoint, this manifests as each component being used independently and uniformly—if the basis is orthonormal, the corresponding set of vectors should exhibit "isotropy."
Of course, this isn't really a rigorous derivation, just a heuristic guide. It tells us that if a set of vectors is isotropic, we can consider it as having originated from an orthonormal basis, in which case it's reasonable to use formula $\eqref{eq:cos}$ to compute similarity. Conversely, if it is not isotropic, we can try to make it more isotropic and then apply formula $\eqref{eq:cos}$ to compute similarity—and this is precisely the approach BERT-flow took, using a "flow model" for this purpose.
Some Rambling Thoughts on Flow Models
In my view, flow models are truly a class of models that leave one with mixed feelings—there's enough to say about them to fill several pages, so let me try to keep it brief here. In mid-2018, OpenAI released the Glow model, and the results looked quite impressive, which drew me to study flow models further—I even reimplemented the Glow model myself. That work is documented in Trickling Flow, Part 1: NICE—Basic Concepts and Implementation of Flow Models and Trickling Flow, Part 2: RealNVP and Glow—Inheritance and Refinement of Flow Models. If you're not familiar with flow models, feel free to check out those two posts. In short, a flow model is a vector transformation model that can convert the distribution of input data into a standard normal distribution—and clearly a standard normal distribution is isotropic, which is why BERT-flow chose the flow model.
So does the flow model have any drawbacks? Actually, I already grumbled about this in an earlier post, Trickling Flow, Part 3: Invertible ResNets—The Aesthetics of Extreme Brute Force. Let me repeat it here:
(The flow model,) through fairly clever design, makes the inverse transformation at each layer relatively simple, and the Jacobian matrix is triangular, so the Jacobian determinant is easy to compute. Such a model is theoretically elegant and beautiful, but has one serious problem: because the inverse transformation must remain simple and the Jacobian determinant must remain tractable, the nonlinear transformation capacity of each layer is very weak. In fact, in models like Glow, only half the variables get transformed at each layer, so to ensure sufficient fitting capacity, the model has to be stacked extremely deep (for example, for generating 256-resolution face images, Glow stacks about 600 convolutional layers with 200 million parameters), which is computationally very expensive.
Reading this, you can now understand why my first reaction to BERT-flow was that "something doesn't feel right." The gripe above tells us that flow models are actually quite weak. So how large was the flow model used in BERT-flow? It was a Glow model with level=2 and depth=3—these two parameters might not mean much to most readers, but suffice it to say it's tiny, small enough that the overall model barely adds any computational overhead. So my "something's off" intuition was:
The flow model itself is weak, and the flow model used in BERT-flow is even weaker, so it seems unlikely that the flow model is playing a decisive role in BERT-flow. Turning this around: maybe we can find a simpler, more direct method that achieves the same effect as BERT-flow.
Standardizing the Covariance Matrix
After some exploration, I did indeed find such a method, and as the title of this post suggests, it's just a linear transformation.
The idea is actually quite simple. We know that a standard normal distribution has mean 0 and identity covariance matrix, so why not try transforming the sentence embeddings so that their mean becomes 0 and their covariance matrix becomes the identity? Suppose the (row) vector set is $\{\boldsymbol{x}_i\}_{i=1}^N$, and we apply the transformation
\begin{equation}\tilde{\boldsymbol{x}}_i = (\boldsymbol{x}_i - \boldsymbol{\mu})\boldsymbol{W} \end{equation}
such that $\{\tilde{\boldsymbol{x}}_i\}_{i=1}^N$ has mean 0 and identity covariance matrix. Readers familiar with traditional data mining might recognize that this is essentially equivalent to the whitening operation used there, which is why I call this method BERT-whitening.
Making the mean 0 is simple—just set $\boldsymbol{\mu}=\frac{1}{N}\sum\limits_{i=1}^N \boldsymbol{x}_i$. The slightly trickier part is solving for the matrix $\boldsymbol{W}$. Let's denote the covariance matrix of the original data as
\begin{equation}\boldsymbol{\Sigma}=\frac{1}{N}\sum\limits_{i=1}^N (\boldsymbol{x}_i - \boldsymbol{\mu})^{\top}(\boldsymbol{x}_i - \boldsymbol{\mu})=\left(\frac{1}{N}\sum\limits_{i=1}^N \boldsymbol{x}_i^{\top}\boldsymbol{x}_i\right) - \boldsymbol{\mu}^{\top}\boldsymbol{\mu}\end{equation}
Then it's not hard to show that the covariance matrix of the transformed data is $\tilde{\boldsymbol{\Sigma}}=\boldsymbol{W}^{\top}\boldsymbol{\Sigma}\boldsymbol{W}$, so what we actually need to solve is the equation
\begin{equation}\boldsymbol{W}^{\top}\boldsymbol{\Sigma}\boldsymbol{W}=\boldsymbol{I}\quad\Rightarrow \quad \boldsymbol{\Sigma} = \left(\boldsymbol{W}^{\top}\right)^{-1}\boldsymbol{W}^{-1} = \left(\boldsymbol{W}^{-1}\right)^{\top}\boldsymbol{W}^{-1}\end{equation}
We know that the covariance matrix $\boldsymbol{\Sigma}$ is a positive semi-definite symmetric matrix, and when there's enough data it's typically positive definite, admitting an SVD decomposition of the form
\begin{equation}\boldsymbol{\Sigma} = \boldsymbol{U}\boldsymbol{\Lambda}\boldsymbol{U}^{\top}\end{equation}
where $\boldsymbol{U}$ is an orthogonal matrix and $\boldsymbol{\Lambda}$ is a diagonal matrix with all positive diagonal entries. Therefore, simply setting $\boldsymbol{W}^{-1}=\sqrt{\boldsymbol{\Lambda}}\boldsymbol{U}^{\top}$ solves the equation:
\begin{equation}\boldsymbol{W} = \boldsymbol{U}\sqrt{\boldsymbol{\Lambda}^{-1}}\end{equation}
A reference implementation in NumPy:
def compute_kernel_bias(vecs):
"""计算kernel和bias
vecs.shape = [num_samples, embedding_size]，
最后的变换：y = (x + bias).dot(kernel)
"""
mu = vecs.mean(axis=0, keepdims=True)
cov = np.cov(vecs.T)
u, s, vh = np.linalg.svd(cov)
W = np.dot(u, np.diag(1 / np.sqrt(s)))
return W, -mu
Some might wonder how this handles large corpora. First, the algorithm above only requires knowing the mean vector $\boldsymbol{\mu}\in\mathbb{R}^{d}$ and covariance matrix $\boldsymbol{\Sigma}\in\mathbb{R}^{d\times d}$ ($d$ being the embedding dimension) of the full set of sentence embeddings, where $\boldsymbol{\mu}$ is the mean of all sentence embeddings $\boldsymbol{x}_i$, and the mean can be computed recursively:
\begin{equation}\boldsymbol{\mu}_{n+1} = \frac{n}{n+1}\boldsymbol{\mu}_{n} + \frac{1}{n+1}\boldsymbol{x}_{n+1}\end{equation}
Similarly, the covariance matrix $\boldsymbol{\Sigma}$ is just the mean of all $\boldsymbol{x}_i^{\top}\boldsymbol{x}_i$ minus $\boldsymbol{\mu}^{\top}\boldsymbol{\mu}$, which can also be computed recursively:
\begin{equation}\boldsymbol{\Sigma}_{n+1} = \frac{n}{n+1}\left(\boldsymbol{\Sigma}_{n}+\boldsymbol{\mu}_{n}^{\top}\boldsymbol{\mu}_{n}\right) + \frac{1}{n+1}\boldsymbol{x}_{n+1}^{\top}\boldsymbol{x}_{n+1}-\boldsymbol{\mu}_{n+1}^{\top}\boldsymbol{\mu}_{n+1}\end{equation}
Since these can be computed recursively, it means we can compute $\boldsymbol{\mu},\boldsymbol{\Sigma}$ under bounded memory, so BERT-whitening is not a problem even for large corpora.
Comparison with BERT-flow
Now we can test the effect of the BERT-whitening approach described above. To compare with BERT-flow, I used bert4keras to run tests on the STS-B task; the reference script is available at:
GitHub link: https://github.com/bojone/BERT-whitening
The comparison of results is as follows:
$$\begin{array}{l|c} \hline & \,\,\text{STS-B}\,\, \\ \hline \text{BERT}_{\text{base}}\text{-last2avg}\,(\text{paper results}) & 59.04 \\ \text{BERT}_{\text{base}}\text{-flow}\,(\text{target, paper results}) & 70.72 \\ \text{BERT}_{\text{base}}\text{-last2avg}\,(\text{personal reproduction}) & 59.04 \\ \text{BERT}_{\text{base}}\text{-whitening}\,(\text{target, personal implementation}) & 71.20 \\ \hline \text{BERT}_{\text{large}}\text{-last2avg}\,(\text{paper results}) & 59.56 \\ \text{BERT}_{\text{large}}\text{-flow}\,(\text{target, paper results}) & 72.26 \\ \text{BERT}_{\text{large}}\text{-last2avg}\,(\text{personal reproduction}) & 59.59 \\ \text{BERT}_{\text{large}}\text{-whitening}\,(\text{target, personal implementation}) & 71.98 \\ \hline \end{array}$$
As we can see, the simple BERT-whitening approach indeed achieves results comparable to BERT-flow. Beyond STS-B, my colleagues performed a similar comparison on internal Chinese business data, and the results consistently showed that the improvement brought by BERT-flow is similar to that of BERT-whitening. This suggests that introducing a flow model may not be all that necessary, since flow-model layers aren't standard layers—they require dedicated implementation and involve a nontrivial amount of engineering effort to train—whereas BERT-whitening is very simple to implement, being just a linear transformation, and can be readily applied on top of any sentence embedding model. (Of course, if one insists on arguing otherwise, one could say that whitening is a flow model implemented via a linear transformation...)
Note: I'll add here as an aside that the "last2avg" mentioned in the BERT-flow paper was originally meant to refer to the average of the outputs of the last two layers, but its actual code computes the average of the outputs of the "first layer + last layer." See the related discussion in this
issue
.
Dimensionality Reduction Can Do Even Better
Now we know that BERT-whitening's transformation matrix $\boldsymbol{W} = \boldsymbol{U}\sqrt{\boldsymbol{\Lambda}^{-1}}$ can transform the data's covariance matrix into the identity matrix. What if we ignore $\sqrt{\boldsymbol{\Lambda}^{-1}}$ and transform using only $\boldsymbol{U}$? It's not hard to show that if we transform using only $\boldsymbol{U}$, the covariance matrix of the data becomes $\boldsymbol{\Lambda}$, which is a diagonal matrix.
As mentioned earlier, $\boldsymbol{U}$ is an orthogonal matrix, which is equivalent to just rotating the overall data without changing the relative positions of the samples—in other words, it's a completely "faithful" transformation. Each diagonal entry of $\boldsymbol{\Lambda}$, on the other hand, measures the magnitude of variation along that particular dimension. If its value is very small, that means this dimension of the feature barely varies, close to a constant, which implies that the original sentence embeddings might actually live in a lower-dimensional space, and we could discard this dimension. This would let us reduce dimensionality while at the same time making the results of cosine similarity more sensible.
In fact, the diagonal matrix $\boldsymbol{\Lambda}$ resulting from SVD is already sorted from largest to smallest, so we only need to keep the leading few dimensions to achieve this dimensionality-reduction effect. Readers familiar with linear algebra will recognize that this operation is actually PCA! And the code only needs a single line changed:
def compute_kernel_bias(vecs, n_components=256):
"""计算kernel和bias
vecs.shape = [num_samples, embedding_size]，
最后的变换：y = (x + bias).dot(kernel)
"""
mu = vecs.mean(axis=0, keepdims=True)
cov = np.cov(vecs.T)
u, s, vh = np.linalg.svd(cov)
W = np.dot(u, np.diag(1 / np.sqrt(s)))
return W[:, :n_components], -mu
The results are as follows:
$$\begin{array}{l|c} \hline & \,\,\text{STS-B}\,\, \\ \hline \text{BERT}_{\text{base}}\text{-last2avg}\,(\text{paper results}) & 59.04 \\ \text{BERT}_{\text{base}}\text{-flow}\,(\text{target, paper results}) & 70.72 \\ \text{BERT}_{\text{base}}\text{-last2avg}\,(\text{personal reproduction}) & 59.04 \\ \text{BERT}_{\text{base}}\text{-whitening}\,(\text{target, personal implementation}) & 71.20 \\ \text{BERT}_{\text{base}}\text{-whitening-256}\,(\text{target, personal implementation}) & 71.42 \\ \hline \text{BERT}_{\text{large}}\text{-last2avg}\,(\text{paper results}) & 59.56 \\ \text{BERT}_{\text{large}}\text{-flow}\,(\text{target, paper results}) & 72.26 \\ \text{BERT}_{\text{large}}\text{-last2avg}\,(\text{personal reproduction}) & 59.59 \\ \text{BERT}_{\text{large}}\text{-whitening}\,(\text{target, personal implementation}) & 71.98 \\ \text{BERT}_{\text{large}}\text{-whitening-384}\,(\text{target, personal implementation}) & 72.66 \\ \hline \end{array}$$
From the table above, we can see that when we keep only the first 256 dimensions out of the base model's 768, the performance actually improves; and since we've reduced the dimensionality, vector retrieval speed is naturally much faster too. Similarly, keeping only the first 384 dimensions of the large model's 1024 also improves performance while reducing dimensionality. This result suggests that sentence embeddings trained without supervision are actually "general-purpose"—for applications in specific domains, many of the features they contain are redundant, and removing these redundant features often yields both a speedup and a performance boost.
By contrast, flow models are invertible and non-dimension-reducing, which is an advantage in some scenarios but a drawback in many others, since they cannot discard redundant dimensions, limiting performance. For example, GAN research has shown that a 256-dimensional Gaussian vector suffices to randomly generate $1024\times 1024$ face images, which suggests that these face images actually lie on a fairly low-dimensional manifold. But if one were to use a flow model for this, invertibility would force the use of a Gaussian vector with as many dimensions as $1024\times 1024\times 3$ to randomly generate images, greatly increasing computational cost while still failing to reach comparable performance.
(Note: for follow-up experimental results, see Which Unsupervised Semantic Similarity Method Is Best? A Fairly Comprehensive Evaluation.)
So the Final Conclusion Is
So, the current findings are as follows: my experiments show that a simple linear transformation operation (BERT-whitening) can basically match the performance of the BERT-flow model. This suggests that introducing a flow model into sentence embedding models may not be as critical as it seems—its correction of the distribution might only operate at a shallow level, and directly correcting the covariance matrix of sentence embeddings through a linear transformation can achieve comparable results. Moreover, BERT-whitening also supports dimensionality reduction, achieving both a speedup and a performance boost.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.