Isotropic Representation Can Improve Dense Retrieval

Consider a set of representation vectors given by $W$, where each row of $W$ corresponds to a representation vector. Mu et al. (2017) suggested $I\left(W\right)$ for measuring isotropy where a partition function $q\left(a\right)={\sum }_{i=1}^N\exp \left(w_i{a}^T\right)$ introduced in (Arora et al., 2016) is utilized. Then, $I\left(W\right)$ is approximated as

where $E$ represents the set of eigenvector of $W^{T}W$. The range of $I\left(W\right)$ is $[0,1]\subset R$. If the representations follow the standard Gaussian distribution, $I\left(W\right)$ becomes close to one. In previous studies, $I\left(W\right)$ was used to measure the isotropy of representations (Yu et al., 2022; Wang et al., 2019).

Suppose that the representation vectors are isotropically spread around the origin. In this case, the average cosine similarity between two randomly selected representation vectors should be close to zero. Based on this simple heuristics, the average cosine similarity $avgcos\left(W\right)$ is adopted as the second metric of isotropy. It can be expressed as

where $w_i$ represents the $i^{th}$ row of the representation matrix $W$. Ethayarajh (2019) also used the average cosine similarity of representations to measure the degree of isotropy.

Normalizing Flow transforms a simple and tractable distribution, such as a standard Gaussian, into the target data distribution by applying a series of invertible and (almost everywhere) differentiable mappings (Tabak and Vanden-Eijnden, 2010; Tabak and Turner, 2013; Kobyzev et al., 2020). By doing this, flow-based models can easily infer the target density and perform sampling from it. Formally speaking, let $Z$ be a random variable of the simple distribution, and $X$ be a random variable of the target distribution. Let $g$ be an invertible and differentiable mapping satisfying $X=g\left(Z\right)$, and $f$ be the inverse of $g$. The probability density of the target distribution can then be computed using the probability density of the simple distribution and volume correction:

where $p_X$ and $p_Z$ are the densities of $X$ and $Z$, respectively, and $Df\left(x\right)=\frac{\partial f}{\partial x}$ is the Jacobian matrix of $f$. Here, the determinant term represents the volume correction induced by the change of variables. This approach is referred to as density estimation. On the other hand, sampling data can be described as the following:

Normalizing Flow can be used for either density estimation or sampling. Between the two purposes, we use Normalizing Flow for density estimation, i.e., to construct flow function $f$ that maps BERT representations into an isotropic Gaussian distribution as Li et al. (2020) did. Among the various flow-based models, we use NICE and Glow in our experiments. More details on each model can be found in the original papers (Dinh et al., 2014; Kingma and Dhariwal, 2018).

As illustrated in Figure 2(b)-left, we aim to transform the BERT representations of query and document into the standard Gaussian distribution using the Normalizing Flow’s density estimation. To do so, in Eq. (5), we set $p_Z$ as the density of the standard Gaussian distribution. Then, we maximize the likelihood of BERT representations, that is, $p_X$ in Eq. (5). Consequently, the flow function $f$ is trained to transform the BERT representations $X$ into the standard Gaussian distribution. Because there is no need for any annotation information, the training is unsupervised. For the inference shown in Figure 2(b)-right, we simply cascade an aggregator that performs the cosine similarity calculation.

For DR, we propose two different implementations of Normalizing Flow. In token-wise implementation, Normalizing Flow is applied to each individual output token of BERT. In this case, each row of $X$ is a token representation of BERT. In sequence-wise implementation, Normalizing Flow is applied to the entire sequence of BERT output. In this case, each row of $X$ is a sequence representation of BERT. Both are applicable to RepBERT that is a single-vector model, but only token-wise is applicable to ColBERT because it is a multi-vector model and isotropy needs to be enforced at the token level and not at the sequence level.

Whitening is a linear transformation that renders the data distribution spherical. That is, it eliminates the structures of location, scale, and correlations in the distribution (Friedman, 1987). Let $\{x_i{\}}_{i=1}^N$ be a set of BERT representations where each $x_i$ is a row vector in $R^D$. To apply whitening, a mean vector $\mu$ and a covariance matrix $\Sigma$ of $\{x_i{\}}_{i=1}^N$ are pre-computed as

where $\Sigma$ denotes the unbiased covariance. The covariance matrix $\Sigma$ has SVD decomposition

where $\Lambda \in R^{D\times D}$ is a diagonal matrix of positive elements, and $U\in R^{D\times D}$ is an orthogonal matrix, i.e., $U^{T}U=I$. Then, the whitened representation vector $z_i$ can be represented as

where each $z_i$ follows a distribution of zero mean and identity covariance (Su et al., 2021).

For DR, we first pre-compute the mean and covariance using a train set of the source data as shown in Figure 2(c)-left. For the inference shown in Figure 2(c)-right, we obtain the whitened output $z_i$ in Eq. (10) using the pre-computed mean and covariance and cascade an aggregator that performs the cosine similarity calculation. As in the Normalizing Flow, we consider both token-wise and sequence-wise implementations of whitening.

As explained earlier, ColBERT (Khattab and Zaharia, 2020) and RepBERT (Zhan et al., 2020) are investigated. In case of ColBERT, we removed the linear layer just before the cosine similarity to clearly examine the effects of Normalizing Flow and whitening. As for the datasets, we examined three popular document re-ranking datasets, Robust04 (Voorhees and others, 2004), WebTrack 2009 (ClueWeb09b) (Callan et al., 2009), and MS-MARCO (Nguyen et al., 2016) as in (MacAvaney et al., 2019). For Robust04, we used the document collections from TREC Disks 4 and 5²²2520K documents, 7.5K triplet data samples, https://trec.nist.gov/data-disks.html. For ClueWeb09b, the document collections from ClueWeb09b³³350M web pages, 4.5K triplet data samples, https://lemurproject.org/clueweb09/ of WebTrack 2009 were used. We also used the large document collections ⁴⁴422G documents, 372K triplet data samples, https://microsoft.github.io/msmarco/TREC-Deep-Learning-2019 from MS-MARCO. To evaluate the ranking models, we used P@20 and NDCG@10 as the performance metrics.

Following Huston and Croft (2014), we divided each of Robust04 and ClueWeb09b into five folds: three folds for training, one for validation, and the remaining one for the test. For MS-MARCO, we divided the dataset into training, validation, and test data. Three different random seeds were used in each experiment. Using the three performance values for the three random seeds, we conducted one-tailed $t$-test under the assumption of homoscedasticity.

For fine-tuning the BERT weights as shown in Figure 2(a), we used a learning rate of 1e-4, batch size of 16, and an Adam optimizer. We set the maximum epoch to be 30 for Robust04 and ClueWeb09b and 10 for MS-MARCO, following the advice of Zhang et al. (2020). We then selected the model with the highest validation performance among the checkpoints of all epochs.

To train Normalizing Flow as shown in Figure 2(b), we stack a Normalizing Flow network at the top of the fine-tuned BERT and train it in an unsupervised manner for either ten epochs (Robust04 and ClueWeb09b) or three epochs (MS-MARCO) with a learning rate of 1e-4, while keeping the fine-tuned BERT weights frozen. For the NICE network, we used a five-layer network with 1000 units in each layer. For the Glow network, we used two levels and depth of three following (Li et al., 2020). For both NICE and Glow, we used a learning rate of 1e-4.

For whitening, we pre-compute mean and covariance as shown in Figure 2(c). Specifically, we collect training data for ten epochs (Robust04 and ClueWeb09b) or three epochs (MS-MARCO) and use the collected representation vectors for the pre-computation of mean and covariance.

For all the experiments, we tuned the hyper-parameters only lightly.

Our experiments are implemented using Python 3.8, Torch 1.10, and huggingface-hub 0.4.0. We fine-tuned the pre-trained BERT-base-uncased model provided by the huggingface transformers. We used RTX3090 GPUs, each of which has 25.6G memory.

In Table 1, we compare the re-ranking performance of fine-tuned ColBERT before and after applying Normalizing Flow or whitening to the BERT representations. Because ColBERT computes the cosine similarity between a query and a document’s multiple token representations, we perform the token-wise post-processing. The results demonstrate that the degree of isotropy is enhanced and the re-ranking performance is improved after the post-processing. For example, when we whiten the BERT representations of ColBERT that is fine-tuned on MS-MARCO, $I\left(W\right)$ increases from 0.6105 to 0.9178, and $avgcos\left(W\right)$ drops from 0.2229 to 0.0088. At the same time, P@20 and NDCG@10 improves by 3.33% and 6.77%, respectively. Across all three datasets, all of whitening, NICE, and Glow methods improve the re-ranking performance of the fine-tuned ColBERT. It can be confirmed that transforming BERT token representations to follow an isotropic distribution consistently improves the performance of ColBERT.

Table 2 presents the results for RepBERT. We compare the performance improvement of both token-wise and sequence-wise post-processing methods. Although RepBERT computes the cosine similarity between sequence representations of queries and documents, we also examine token-wise methods in addition to the sequence-wise methods because sequence-wise transformation might not be powerful enough if the query and document sequences have sufficiently different characteristics. As for the overall improvement, the results in Table 2 show that the performance is improved for all the cases except for one when the token-wise method is used, but not so consistently when the sequence-wise method is used. The peak improvement is as large as 22.81%. Between token-wise and sequence-wise methods, it can be observed that token-wise performs better for Robust04 and ClueWeb09b and sequence-wise performs better for MS-MARCO. We provide a discussion on this issue in Section 6.3.

Among the three post-processing methods, Glow is the most robust method for RepBERT in the sense that the performance is always improved regardless of the choice of dataset or the choice of token-wise or sequence-wise method. While not as consistent, whitening almost always achieves the best performance for each dataset as long as token-wise is used for Robust04 and ClueWeb09b and sequence-wise is used for MS-MARCO. Overall, whitening improves the performance of the fine-tuned model by 6.88% $\sim$ 22.81% on NDCG@10 over the three datasets.

We show the effectiveness of isotropic representations for OOD generalizability in Table 3. In the results, it can be observed that the baseline OOD generalizability of a fine-tuned model is always improved by applying a post-processing method that enforces isotropy to the BERT representations. Performance of P@20 is improved by 4.17% $\sim$ 27.50%, and NDCG@10 by 5.27% $\sim$ 24.98%, across the three datasets. The overall result indicates that isotropically distributed BERT representations are robust in the sense that they can perform better for an OOD dataset that has been unseen during the training.

In particular, when ColBERT is trained with MS-MARCO as the source data and used for inference with the target data of either Robust04 or ClueWeb09b, the OOD performance with isotropy even surpasses the ID performance. This is visualized in Figure 3 for Robust04. In the figure, the ID performance without enforcing isotropy is shown as a red line (the ID performance values can be found in Table 1). The baseline OOD performance with fine-tuning only results in a performance drop of 18.44% and 14.63% on P@20 and NDCG@10, respectively, when compared to the red line. After the fine-tuned OOD model undergoes a post-processing (whitening, NICE, or Glow), however, the OOD model mostly performs better than the baseline ID performance shown in the red line. This result may be interpreted to be surprising because the OOD models were trained without the target data and because the OOD models outperform the ID models that were trained with the target data. The conclusion remains the same even if we choose the red line to be the best ID performance with isotropy that can be found in Table 1. Although we did not include the results for RepBERT, similar results can be obtained for RepBERT as well.