Text Embeddings by Weakly-Supervised Contrastive Pre-training

Introduction

Text embeddings are low-dimensional vector representations for arbitrary-length texts and play key roles in many NLP tasks such as large-scale retrieval. Compared to the high-dimensional and sparse representations like TF-IDF, text embeddings have the potential to overcome the lexical mismatch issue and facilitate efficient retrieval and matching between texts. It also offers a versatile interface easily consumable by downstream applications.

While pre-trained language models such as BERT and GPT can produce transferrable text representations, they are not ideal for tasks such as retrieval and text matching where a single-vector embedding of texts is more desired due to its efficiency and versatility. To obtain better text embeddings, contrastive learning is often the go-to framework to enhance the sequence-level representations from text pairs. Along this line of research, some works are geared towards learning task-specific embeddings. For example, GTR and Sentence-T5 fine-tune pre-trained models with supervised datasets to learn embeddings customized for passage retrieval and semantic textual similarity, respectively. Other works learn unsupervised embeddings from automatically constructed text pairs. Typical methods to construct text pairs include Inverse Close Task (ICT) , random cropping and neighboring text spans , etc. While such synthetic data are of unlimited quantity, they are often poor in quality and the resulted embeddings fail to match the performance of the classic BM25 baseline without further fine-tuning .

In this work, we learn a high-quality general-purpose text embedding termed E5, EmbEddings from bidirEctional Encoder rEpresentations. E5 aims to provide strong off-the-shelf text embeddings suitable for any tasks requiring single-vector representations in both zero-shot or fine-tuned settings. To achieve this goal, instead of relying on limited labeled data or low-quality synthetic text pairs, we contrastively train E5 embeddings from CCPairs, a curated web-scale text pair dataset containing heterogeneous training signals. We construct the CCPairs dataset by combining various semi-structured data sources such as CommunityQA, Common Crawl and Scientific papers, and perform aggressive filtering with a consistency-based filter to improve data quality. We choose a simple contrastive learning recipe using in-batch negatives with a large batch-size to train our model. Extensive experiments on both BEIR and MTEB benchmarks demonstrate the effectiveness of the proposed method. On the BEIR zero-shot retrieval benchmark , E5 is the first model to outperform the strong BM25 baseline without using any labeled data. When fine-tuned on labeled datasets, the performance can be further improved. Results on $56$ datasets from the recently introduced MTEB benchmark show that our E5${}_{\text{base}}$ is competitive against GTR${}_{\text{xxl}}$ and Sentence-T5${}_{\text{xxl}}$, which have $40\times$ more parameters.

Related Work

There have been long-lasting interests in transforming texts into low-dimensional dense embeddings. Early works include Latent Semantic Indexing (LSA) and Latent Dirichlet Allocation (LDA) . LSA utilizes the decomposition of a word-document co-occurrence matrix to generate document embeddings, while LDA adopts probabilistic graphical models to learn topic distributions. Arora et al. show that a simple weighted average of word vectors can be a strong baseline for sentence embeddings.

With the development of pre-trained language models and large-scale labeled datasets such as SNLI and MS-MARCO , methods like Sentence-BERT , SimCSE , Sentence-T5 and SGPT directly fine-tune language models to output continuous embeddings. Most research focuses on short texts and thus uses the term "sentence embeddings". For long documents, it remains an open research question whether fixed-length embeddings can encode all the information. Contrastive loss popularized by SimCLR turns out to be more effective than classification-based losses for embeddings. LaBSE , LASER and CLIP further extend to multilingual and multi-modal scenarios using parallel sentences and image-text pairs.

Another direction is to design self-supervised pre-training tasks for text matching and retrieval. proposes the well-known inverse cloze task (ICT), where a random sentence within a passage is chosen as a pseudo-query and the rest is treated as a positive sample. However, Contriever shows that random cropping with data augmentation is more effective than ICT on a range of zero-shot information retrieval tasks. OpenAI text embeddings use neighboring texts as positives and scale up the model size to $175$B. Oguz et al. performs domain-matched pre-training to improve in-domain results. SPAR trains a dense retriever by treating BM25 as a teacher model. Although the aforementioned approaches can easily obtain abundant supervision signals, such synthetic data tend to be of low quality. Results on the BEIR benchmark show they struggle to match the performance of BM25 if not further fine-tuned on labeled datasets.

Evaluation and interpretation of text embeddings are also non-trivial. Most benchmarks measure the embedding quality through downstream task performances. For example, SentEval uses linear probing and a collection of semantic textual similarity (STS) datasets, while the BEIR benchmark focuses on zero-shot information retrieval scenarios. The recently introduced MTEB benchmark combines $56$ datasets spanning across $8$ tasks and $112$ languages. Experiments show no model can achieve state-of-the-art results on all embedding tasks yet. In this paper, we do not use the SentEval toolkit since its linear probing setup depends on the optimization hyperparameters.

Most closely related to our work is a series of community efforts by sentence-transformers https://github.com/UKPLab/sentence-transformers to train embeddings with a collection of labeled and automatically collected datasets. In this paper, we show that it is possible to train high-quality embeddings using self-supervised pre-training only. In terms of benchmark results, our model can achieve superior performance when fine-tuned on less labeled data.

CCPairs: A Large Collection of Text Pair Dataset

The quality and diversity of the data is crucial for training general-purpose text embeddings. In this work, we mine and assemble CCPairs, a large high-quality text pair dataset from web sources which provide diverse training signals transferring well to a wide range of tasks.

Harvesting semi-structured data sources Large-scale high-quality datasets like C4 and CCMatrix are vital for the success of language model pre-training and machine translation. For learning text embeddings, existing works either utilize small-scale human-annotated data such as NLI and MS-MARCO or adopt heuristics such as random cropping to obtain large-scale but very noisy supervision signals.

Instead, we curate a text pair dataset CCPairs (Colossal Clean text Pairs) by harvesting heterogeneous semi-structured data sources. Let ($q$, $p$) denote a text pair consisting of a query $q$ and a passage $p$. Here we use “passage” to denote word sequences of arbitrary length, which can be a short sentence, a paragraph, or a long document. Our dataset includes (post, comment) pairs from Reddit https://files.pushshift.io/reddit/, (question, upvoted answer) pairs from Stackexchange https://archive.org/details/stackexchange, (entity name + section title, passage) pairs from English Wikipedia, (title, abstract) and citation pairs from Scientific papers , and (title, passage) pairs from Common Crawl https://commoncrawl.org/ web pages and various News sources.

We only include data sources that can be automatically mined, and some subsets are directly reused from existing datasets. Simple heuristic rules are applied to filter data from Reddit and Common Crawl. For example, we remove Reddit comments that are either too long ($>4096$ characters) or receive score less than $1$, and remove passages from web pages with high perplexity . After preliminary filtering, we end up with $\sim 1.3$ billion text pairs, most of which come from Reddit and Common Crawl. For more details and examples, please refer to Appendix A.

Consistency-based filter To further improve data quality and make training costs manageable, we propose a consistency-based data filtering technique: a model is first trained on the $1.3$B noisy text pairs, and then used to rank each pair against a pool of $1$ million random passages. A text pair is kept only if it falls in the top-$k$ ranked lists. In other words, the model’s prediction should be consistent with the training labels. Here we set $k=2$ based on manual inspection of data quality. After this step, we end up with $\sim 270$M text pairs for contrastive pre-training.

The intuition for this technique comes from the memorization behaviors of neural networks : when trained on noisy datasets, neural networks tend to memorize the clean labels first and then gradually overfit the noisy labels. Similar techniques have been widely used for removing dataset noises. It is also possible to apply this filter iteratively, we will leave it for future work.

Method

Our embeddings can be trained with only unlabeled text pairs from CCPairs with contrastive pre-training. A second-stage fine-tuning on small, high-quality labeled datasets can be performed to further boost the quality of the resulted embeddings. See Figure 1 for an overview.

Contrastive pre-training aims to distinguish the relevant text pairs from other irrelevant or negative pairs. Given a collection of text pairs $\{(q_{i},p_{i})\}_{i=1}^{n}$, we assign a list of negative passages $\{p_{ij}^{-}\}_{j=1}^{m}$ for the $i$-th example. Then the InfoNCE contrastive loss is as follows:

where $s_{\boldsymbol{\theta}}(q,p)$ is a scoring function between query $q$ and passage $p$ parameterized by $\boldsymbol{\theta}$. Following the popular biencoder architecture, we use a pre-trained Transformer encoder and average pooling over the output layer to get fixed-size text embeddings $\mathbf{E}_{q}$ and $\mathbf{E}_{p}$. The score is the cosine similarity scaled by a temperature hyperparameter $\tau$:

Where $\tau$ is set to $0.01$ in our experiments by default. We use a shared encoder for all input texts and break the symmetry by adding two prefix identifiers “query:” and “passage:” to $q$ and $d$ respectively. For some data sources such as citation pairs, it is not obvious which side should be the query, we randomly choose one for simplicity. Such an asymmetric design turns out to be important for some retrieval tasks where there exist paraphrases of the query in the target corpus.

Another critical issue for contrastive training is how to select the negative samples. Here we choose to use the in-batch negatives , where the passages from other pairs in a batch serve as negative samples. We find that this simple strategy enables more stable training and outperforms methods such as MoCo when the batch size is sufficiently large.

2 Fine-tuning with Labeled Data

While contrastive pre-training on the CCPairs provides a solid foundation for general-purpose embeddings, further training on labeled data can inject human knowledge into the model to boost the performance. Although these datasets are small, existing works have shown that supervised fine-tuning leads to consistent performance gains. In this paper, we choose to further train with a combination of 3 datasets: NLI The version released by SimCSE. (Natural Language Inference), MS-MARCO passage ranking dataset , and NQ (Natural Questions) dataset . Empirically, tasks like STS (Semantic Textual Similarity) and linear probing benefit from NLI data, while MS-MARCO and NQ datasets transfer well to retrieval tasks.

Building on the practices of training state-of-the-art dense retrievers , we use mined hard negatives and knowledge distillation from a cross-encoder (CE) teacher model for the MS-MARCO and NQ datasets. For the NLI dataset, contradiction sentences are regarded as hard negatives. The loss function is a linear interpolation between contrastive loss $L_{\text{cont}}$ for hard labels and KL divergence $D_{\text{KL}}$ for distilling soft labels from the teacher model.

Where $p_{\text{ce}}$ and $p_{\text{stu}}$ are the probabilities from the cross-encoder teacher model and our student model. $\alpha$ is a hyperparameter to balance the two loss functions. $L_{\text{cont}}$ is the same as in Equation 1.

3 Applications to Text Embedding Tasks

After the above two steps, we obtain high-quality text embeddings transferring well to a wide range of tasks without fine-tuning the model parameters. Combined with techniques like approximate nearest neighbor search, embeddings provide a scalable and efficient solution for applications like web search. Here we briefly illustrate several use cases of our text embeddings.

Zero-shot Retrieval First, the passage embeddings for the target corpus are computed and indexed offline. Then for each query, we compute its query embedding and return the top-$k$ ranked lists from the corpus based on cosine similarity.

Few-shot Text Classification A linear classifier is trained on top of the frozen embeddings with a few labeled examples. Different tasks only need to train and save the parameters of the classification heads. It can be seen as a particular form of parameter-efficient learning .

Zero-shot Text Classification The input and label texts are converted to sentences based on manually written prompt templates. The predicted label is the one closest to the input text in the embedding space. Take the sentiment classification of movie reviews as an example, with the original input “I enjoy watching it”, the label text is “it is an example of terrible/great movie review” and the input text becomes “movie review: I enjoy watching it”.

Semantic Textual Similarity Given two text embeddings, we use the cosine function to measure their semantic similarity. Since the absolute similarity scores do not enable an easy interpretation, the evaluation is usually based on rank correlation coefficients.

Text Clustering Standard clustering algorithms such as k-means can be applied straightforwardly. Texts belonging to the same category are expected to be close in the embedding space.

For tasks other than zero-shot text classification and retrieval, we use the query embeddings by default.

Experiments

Pre-training We pre-train on our proposed text pair dataset for three model sizes: E5${}_{\text{small}}$, E5${}_{\text{base}}$ and E5${}_{\text{large}}$ initialized from MiniLM , bert-base-uncased, and bert-large-uncased-whole-word-masking respectively. The batch size is set to a large value of $32,768$ to increase the number of negatives. The learning rate is {$3,2,1$}$\times 10^{-4}$ for the {small, base, large} models, with linear decay and the first $1,000$ steps for warmup. We pre-train for $20k$ steps in total with AdamW optimizer, which is approximately $2.5$ epochs over the dataset. It takes {$16,32,64$} V100 GPUs and {$1,1,2$} days for the {small, base, large} models. To improve training efficiency and reduce GPU memory usage, we adopt mixed precision training and gradient checkpointing.

Fine-tuning is performed on the concatenation of $3$ datasets: MS-MARCO passage ranking , NQ , and NLI datasets. We reuse the mined hard negatives and re-ranker scores from SimLM for the first two datasets. Models are fine-tuned for $3$ epochs with batch size $256$ on $8$ GPUs. Learning rate is {$3,2,1$}$\times 10^{-5}$ for the {small, base, large} models with $400$ steps warmup. For each example, we use $7$ hard negatives. Since the NLI dataset only has $1$ hard negative for each example, $6$ sentences are randomly sampled from the entire corpus.

We use E5-PT to denote models with contrastive pre-training only. More implementation details can be found in Appendix B.

2 Evaluation Datasets

BEIR Benchmark is a collection of $19$ information retrieval datasets, ranging across ad-hoc web search, question answering, fact verification and duplicate question retrieval, etc. We evaluate the $15$ datasets that provide public downloads. The main metric is nDCG@$10$.

MTEB Benchmark is recently proposed for benchmarking massive text embedding tasks. Though MTEB is multilingual due to the inclusion of bitext mining datasets, most datasets are still only available in English. In this paper, we evaluate the English subsets, which have $56$ datasets spanning across $6$ categories: Classification (Class.), Clustering (Clust.), Pair Classification (PairClass.), Rerank, Retrieval (Retr.), STS, and Summarization (Summ.). The evaluation metrics are accuracy, v-measure, average precision, MAP, nDCG@10, and Spearman coefficients, respectively. Please refer to the MTEB paper for details.

3 Results on BEIR benchmark

Results with Unsupervised Methods In Table 1, we show model results that do not use any labeled data. When averaged over all $15$ datasets, E5-PT${}_{\text{base}}$ outperforms the classic BM25 algorithm by $1.2$ points. To the best of our knowledge, this is the first reported result that an unsupervised model can beat BM25 on the BEIR benchmark. When scaling up to E5-PT${}_{\text{large}}$, we see further benefits from $42.9$ to $44.2$.

In terms of pre-training tasks, Contriever adopts random cropping, while LaPraDor combines ICT and dropout-as-positive-instance from SimCSE. The methods can easily obtain large-scale training data, while our approach requires more effort in dataset curation. Such efforts pay off with better results. Recent studies also show that improving data quality is a vital step for training large language models.

Results with Supervised Fine-tuning In Table 2, we fine-tune our models on supervised datasets and then transfer them to the BEIR benchmark. Since our fine-tuning datasets include MS-MARCO and NQ, the corresponding numbers are in-domain results. For other datasets, these are zero-shot transfer results. Our E5${}_{\text{base}}$ model achieves an average nDCG@$10$ of $48.7$, already surpassing existing methods with more parameters such as GTR${}_{\text{large}}$ . Most datasets benefit from supervised fine-tuning, but there are also a few exceptions such as FiQA, Scidocs, and Fever, etc. This is likely due to the lack of enough domain diversity for the fine-tuning datasets.

4 Results on MTEB benchmark

In Table 3, E5 models not only substantially outperform existing ones with similar sizes, but also match the results of much larger models. The top-$2$ models on MTEB leaderboard https://huggingface.co/spaces/mteb/leaderboard, as of November 22, 2022 GTR${}_{\text{xxl}}$ and Sentence-T5${}_{\text{xxl}}$ have $4.8$B parameters, while our E5${}_{\text{large}}$ model is more than $10\times$ smaller with $300$M parameters. We expect that our model will benefit from continual scaling up.

Since the difference between BERT-FT${}_{\text{base}}$ and E5${}_{\text{base}}$ is that BERT-FT${}_{\text{base}}$ only has fine-tuning stage, their performance gap demonstrates the usefulness of contrastive pre-training on our proposed CCPairs dataset. For most task categories except Clustering, performance improves after supervised fine-tuning. Consistent with prior works , this once again demonstrates the importance of incorporating human knowledge for learning better text embeddings. It remains an open question whether state-of-the-art embeddings can be obtained in a purely self-supervised manner.

Table 4 shows the zero-shot text classification results on the dev set of the SST-2 dataset . By formulating text classification as embedding matching between input and label texts, our model can be much better than the “majority” baseline in a zero-shot setting. We use the prompt template from Section 4.3.

5 Analysis

In this section, we conduct a series of analyses to examine various design choices. All the numbers in this section are from base-size models. For the BEIR benchmark, we choose $6$ datasets with more stable results across different runs. Some negative results are also listed in Appendix C.

Impacts of Batch Size Since we use in-batch negatives for contrastive pre-training, larger batch size will provide more negatives and therefore improve the quality of the learned text embeddings. In Table 5, increasing batch size from $1$K to $32$K leads to consistent gains across all $6$ datasets. It is also possible to train with smaller batch sizes by adding hard negatives . However, the engineering efforts of mining hard negatives for large datasets (>$100$M) are non-trivial.

Fine-tuning Datasets GTR models are fine-tuned with “MS-MARCO + NQ”, while Sentence-T5 models use NLI instead. In Table 6, we can see that the “MS-MARCO + NQ” setting performs best on retrieval tasks, and the NLI data is beneficial for STS and linear probing classification. Similar observations are also made by Muennighoff et al. . Combining all of them leads to the best overall scores on the MTEB benchmark. This also illustrates the importance of dataset diversity for learning text embeddings.

Data Filtering One crucial step in our dataset curation pipeline is filtering out low-quality text pairs. In Table 7, when training with $1$M pairs, using filtered data has a nearly $6$ points advantage. When all the text pairs are used, the “w/o filter” setting has about $4\times$ more data but is still behind by $1.6$ points. Though recent studies show that deep learning models are quite robust to dataset noises, data filtering still has benefits in improving training efficiency and model quality.

Negative Sampling We explore two alternative methods to enlarge the number of negatives: Pre-batch negatives reuse embeddings from previous batches as additional negatives, while MoCo introduces a momentum encoder and uses a FIFO queue to store negatives. For both approaches, the negative size can be easily scaled up without incurring much GPU memory overhead. The downside is that most negatives are produced by an older version of model parameters. In Table 8, in-batch negatives still perform favorably. Empirically, we find that MoCo is more sensitive to certain hyperparameters such as temperature, better results are possible with more tuning.

BM25 vs Dense Retrieval With the rapid development of dense retrieval models, can we replace the long-standing BM25 algorithm from now on? The answer is likely “not yet”. BM25 still holds obvious advantages in terms of simplicity, efficiency, and interpretability. For long-tail domains such as Trec-Covid and retrieval tasks that involve long documents (Touche-2020) or rely heavily on exact lexical match (Fever) , further research efforts are still necessary to improve current dense retrievers.

Conclusion

In this work, we train a general-purpose text embedding model E5 from weak supervision signals. We adopt a simple contrastive training framework with in-batch negatives and learn from a large-scale text pair dataset we harvest from heterogeneous data sources across the web. E5 offers strong off-the-shelf performance for a wide range of tasks requiring single-vector text representations such as retrieval, semantic textual similarity, and text matching. When further customized for downstream tasks, E5 achieves superior fine-tuned performance compared to existing embedding models with $40\times$ more parameters on the large, 56-task MTEB benchmark datasets.

References

Appendix A Dataset Details

For Common Crawl, we download the 2022-33 snapshot and cc_net https://github.com/facebookresearch/cc_net is used for preprocessing including language identification, de-duplication, language model filtering, etc. Web pages from the MS-MARCO document ranking corpus are also included. For the data filtering step, we examine each pair of passages within a web page instead of just using the title as a query. For Wikipedia, we use the version released by Petroni et al. . To avoid possible data contamination, we remove text pairs that occur in the evaluation datasets based on exact string match.

Reddit data is collected from the year 2018 to August 2022. For the S2ORC data, we use a sample weight of $0.3$ during training to avoid over-fitting the scientific domains.

For the BEIR benchmark, we use the 15 datasets that provide public downloads: MS MARCO , Trec-Covid , NFCorpus , NQ , HotpotQA , FiQA , ArguAna , Touche-2020 , CQADupStack , Quora, DBPedia , Scidocs , Fever , Climate-Fever , and Scifact .

Appendix B Implementation Details

We list the hyperparameters in Table 11. Since some evaluation datasets have long texts, we freeze the position embeddings during both pre-training and fine-tuning and set the maximum text length to $512$ for evaluation.

For the Quora duplicate retrieval task in the BEIR benchmark, we add prefix “query: ” to all the questions. For other retrieval tasks, we use “query: ” and “passage: ” prefixes correspondingly.

The MS-MARCO results in Table 12 use document titles provided by RocketQA . This evaluation setup is consistent with most state-of-the-art dense retrievers. However, the MS-MARCO data from the BEIR benchmark does not have titles, so the results are expected to be lower.

In-domain Evaluation We report results for in-domain datasets in Table 12. These results can help illustrate the benefits brought by contrastive pre-training when abundant in-domain labeled data are available. For MS-MARCO passage ranking, MRR@10 and Recall@1k are reported. For the NQ dataset, Recall@20 and Recall@100 are the main metrics.

Appendix C Negative Results

Here are some attempts that we eventually give up on:

Adding BM25 hard negatives Similar to DPR , we add one BM25 hard negative for each positive pair during training. When using $15$M data, this strategy improves the overall results by $\sim 0.5$ points on the BEIR benchmark. However, running the BM25 algorithm over a $250$M+ dataset is too time-consuming even with multi-node and multi-process parallelism.

Using RoBERTa instead of BERT for initialization Though RoBERTa shows consistent gains on many NLP tasks, we empirically find that RoBERTa performs worse than BERT initialization on most of the BEIR benchmark datasets.

Auxiliary MLM objective We add a masked language modeling loss for $25$% of the training text pairs. The numbers are on par with removing this auxiliary objective, but the training cost goes up.