Paper Explained 3: E5

How a simple architecture is transformed into a SOTA embedding model

Abstract

Have you ever wondered how search engines can interpret vague queries like amzn or ggle and still return the right results — yet sometimes completely miss the mark? The secret lies in a powerful technology shaping today’s world: text embeddings — a way to convert human language into numbers that computers can understand.

In this article, we explore an efficient yet powerful embedding model called E5 (EmbEddings from bidirEctional Encoder rEpresentations). While most tech giants are racing to build increasingly larger and complex language models, E5 takes an alternative approach and shows that with high-quality web-scale data and carefully designed training, simple model architectures can be transformed into state-of-the-art embedding models.

The E5 model family has been made public by Microsoft, and we will also provide a simple guide on how to use it for your own applications.

Table of contents for this article:

• Background: introduce fundamentals of text embedding, how contrastive learning works, and limitations of previous embedding models.
• E5: explain key innovations of E5 model, model variants and a practical guide on how to use it.
• Evaluation.
• Summary.
• References

Background

Text Embedding Fundamentals

In Machine Learning, Embedding refers to the process that maps high-dimensional, categorical, or sparse data into a continuous, lower-dimensional vector space, to make it more computationally efficient for machines to process while preserving the relationships and structure of the original data.

Following this idea, text embedding simply means transforming human language to dense, low-dimensional vectors. After embedding, vectors from semantically similar text pairs should be close to each other in the embedding space, so that we can retrieve these semantically similar texts using techniques like Approximate Nearest Neighbor (ANN) search.

Text embedding is a key element in today’s AI-driven world, and is widely used in many applications, including information retrieval, semantic search, recommendation systems, document classification, etc.

So what makes embedding so important? Before we have text embedding, texts can only be searched or matched using exact words, which cannot distinguish different meanings of the same word under different context (such as “bank” in “river bank” vs. “bank account”) nor understand different words with similar meanings (such as “dogs” vs. “puppies”).

In other words, text embedding enables us to search and match texts in the embedding space according to their semantic similarity rather than relying on the exact raw text, therefore leading to better performance.

But how does text embedding preserve the semantic similarity while mapping texts into the embedding space? The answer is through Contrastive Learning.

Contrastive Learning

Contrastive learning is a commonly used technique for embedding learning. It follows a self-supervised framework, hence doesn’t require explicit human labels, which enables it to leverage the vast amount of loosely annotated data or even synthesized data.

The core idea of contrastive learning is to project high-dimensional data like images or texts into a latent space where similar samples are pulled closer, and dissimilar ones are pushed apart, as shown in the figure below.

Figure 1. Contrastive learning. Image created by Author.

This is often achieved using a contrastive loss function such as InfoNCE (Noise Contrastive Estimation) [2], which encourages high similarity scores for positive pairs while minimizing similarity with negative pairs. This loss function can be formulated as below:

where

• q is the anchor sample, k_0 is the positive sample, and k_i (i from 1 to N) are all negative samples.
• hence, (q ⋅ k_i) represents the cosine similarity or inner product between the anchor point and each individual sample.
• τ is often called the temperature parameter, which controls the sharpness of the softmax distribution.

Therefore, by minimizing the above InfoNCE loss, the model will learn to maximize similarity between q and k_0 in the numerator and minimize the denominator.

Training with such a loss function requires a positive sample and a bunch of negative samples for each anchor point. In practice, the most common setting is to collect a dataset with just positive pairs (q, k_0), and use the rest samples in the same training batch as negative samples, which are called in-batch negatives. There are also techniques to mine hard negatives, which is more complex but can often lead to better embedding quality.

Limitations of Prior Embedding Models

Existing embedding models prior to E5 often suffer from two limitations:

• Models like BERT naturally generate a series of token-level embeddings rather than sequence-level embeddings, however in tasks such as retrieval and text matching, it is more useful to generate a single embedding vector for the entire input sequence.
• Most models are trained on either small-scale human annotated dataset or large-scale text pairs collected from website with lower quality, which often leads to poor embedding quality.

In the following section, we will explain how E5 mitigates both issues with a simple and efficient model architecture.

E5

The success of E5 comes from two major ingredients: a large scale high-quality dataset, plus carefully designed training strategies. In this section, we will take a closer look on both aspects.

CCPairs (Colossal Clean Text Pairs) Dataset

Imagine you are asked to collect a large scale yet high-quality dataset with text pairs, how would you ensure this dataset is both diverse and of high-quality?

Firstly, we need to make sure the text pairs are collected from as many sources as possible, so that they cover a broad range of topics. To achieve this, in [1] the authors collect text pairs from multiple sources, including (post, comment) pairs from Reddit, (question, upvoted answer) pairs from Stackexchange, citation pairs from Scientific papers, and (title, passage) pairs from Common Crawl web pages and various News sources.

The second issue is to improve data quality. For that, some simple data cleaning strategies are firstly applied in [1], including

• Removing Reddit comments that are either longer than 4096 characters or have scores less than 1.
• Removing web page passages with perplexity higher than 60, as high perplexity typically indicates lack of coherence.

After cleaning, the size of the dataset is reduced to roughly 1.3B pairs, but they didn’t stop there. Instead, they apply a consistency-based filter to further enhance the quality of this dataset, by training a model with the 1.3B noisy data first, and then use it to rank each pair against a pool of 1M random passages.

Intuitively, if a (query, passage) pair does present good relevance, the passage should be ranked before the random passages. Following this intuition, if a passage does not appear in the top 2 ranked list, it will be considered as a noisy pair and get filtered from the dataset. This way, the size of the dataset is further reduced to ~270M.

The figure below summarizes the collection process of CCPairs dataset.

Figure 2. CCPairs collection pipeline. Image from [1].

Contrastive Pre-training

As we mentioned before, contrastive learning is a commonly used strategy to learn embedding models, as it can pull embeddings from semantically similar pairs closer while pushing embeddings from dissimilar pairs further.

In E5, contrastive learning is applied in both pre-training and finetuning stage. Note that all the text pairs in the training dataset are just positive pairs, so we still need to generate negative pairs in order to formulate the contrastive loss. In [1], these negative pairs are selected using the in-batch negative strategy at pre-training, which is simple and effective when the batch size is sufficiently large.

In terms of implementation, E5 uses a shared Transformer encoder to encode both query and passage in the text pairs into embeddings. In order to distinguish between query and passage, two prefix identifiers “query:” and “passage:” are added to the input texts, as shown in the figure below.

Figure 3. E5 mdoel architecture. Image from [1].

Transformer by default will generate token-level embeddings, meaning that the output embedding will be of shape L x d, where L is the sequence length and d is the embedding dimension. In E5, average pooling is applied to these token-level embeddings to get sequence-level embedding for both query and passage, denoted by E_q and E_p, respectively.

Then, we can define the score between a certain text pair as below using cosine similarity scaled by a temperature hyperparameter τ:

where τ is set to 0.01.

With that, the training objective can be defined using the InfoNCE loss:

Fine-tuning with Labeled Data

Since embedding models can be used by many downstream tasks, finetuning on domain-specific labeled data can often lead to better performance on these downstream tasks.

Following this idea, E5 is further finetuned on a combination of three labeled datasets:

• NLI (Natural Language Inference): which determines the relationship between a premise and a hypothesis, by considering whether the hypothesis is entailment (true given the premise), contradiction (false given the premise), or neutral (neither). Finetuning on this dataset will benefit tasks like classification for reasoning and understanding entailment relations.
• MS-MARCO passage ranking: which retrieves the most relevant passages given a query, so finetuning on this dataset will benefit retrieval tasks.
• NQ (Natural Questions): which finds the best short answer span from a large document like Wikipedia. Finetuning on this dataset will benefit tasks like extractive question answering and document comprehension.

At this stage, the finetuning objective is a weighted combination of contrastive learning and knowledge distillation from a cross-encoder teacher model.

In case you are not familiar with these terminologies, let’s firstly explain these two concepts: cross-encoder and knowledge distillation.

Cross-Encoder Models

In cross-encoder models, both text pairs are concatenated before fed into the model, as shown in the figure below. This enables the model to capture interactions between tokens in both query and passage much earlier, and hence can often lead to better performance compared with bi-encoder.

Figure 4. Bi-encoder vs. Cross-encoder. Image from [3].

Knowledge Distillation

Knowledge distillation is a commonly used approach for model compression (other methods include pruning and quantization), where instead of learning a small model directly, we firstly train a larger teacher model, and then train the student model by mimicking the behavior of the teacher model. This has been proved to be an effective strategy to learn small models with better performance, since the teacher model provides finer-grained information rather than binary class labels that can help the student to learn faster.

Figure 5. Knowledge distillation. Image from [4].

A commonly used training objective for knowledge distillation is KL divergence, since it can measure the distance between two distributions, which in the teacher-student case, can be the output probability for teacher and student, respectively.

Putting everything together, the final training objective in finetuning stage of E5 is as below:

where p_ce and p_stu represent the probabilities generated by the teacher model and E5 student, respectively, and alpha is the weight to combine the KL divergence loss with the contrastive loss.

Model Variants and Usage

E5 presents three different variants, called E5-small, E5-base and E5-large, where

• E5-small is initialized from MiniLM.
• E5-base is initialized from bert-base-uncased.
• E5-large is initialized from bert-large-uncased-whole-word-masking.

All model variants are available in Huggingface. The following code example (provided by Huggingface)shows how to process texts with multilingual-e5-small:

from sentence_transformers import SentenceTransformer
model = SentenceTransformer('intfloat/multilingual-e5-small')
input_texts = [
    'query: how much protein should a female eat',
    'query: 南瓜的家常做法',
    "passage: As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 i     s 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or traini     ng for a marathon. Check out the chart below to see how much protein you should be eating each day.",
    "passage: 1.清炒南瓜丝 原料:嫩南瓜半个 调料:葱、盐、白糖、鸡精 做法: 1、南瓜用刀薄薄的削去表面一层皮     ,用勺子刮去瓤 2、擦成细丝(没有擦菜板就用刀慢慢切成细丝) 3、锅烧热放油,入葱花煸出香味 4、入南瓜丝快速翻炒一分钟左右,     放盐、一点白糖和鸡精调味出锅 2.香葱炒南瓜 原料:南瓜1只 调料:香葱、蒜末、橄榄油、盐 做法: 1、将南瓜去皮,切成片 2、油     锅8成热后,将蒜末放入爆香 3、爆香后,将南瓜片放入,翻炒 4、在翻炒的同时,可以不时地往锅里加水,但不要太多 5、放入盐,炒匀      6、南瓜差不多软和绵了之后,就可以关火 7、撒入香葱,即可出锅"
]
embeddings = model.encode(input_texts, normalize_embeddings=True)

Evaluation

E5 is evaluated on two popular benchmarks, BEIR and MTEB, where

• BEIR contains a bunch of tasks including web search, QA, fact verification and duplicated question retrieval.
• MTEB is a more recently proposed benchmark for text embedding.

On BEIR, E5 is evaluated on both pre-training and finetuned models, where it outperforms other baselines in both setting across most tasks:

On MTEB, E5 after finetuning can even outperform models much larger:

Summary

This article explores E5, an efficient yet powerful model family for text embedding.

More specifically, we start from a brief introduction of related concepts like text embedding fundamentals and contrastive learning, as well as limitations of existing text embedding models, to give more context for readers.

On top of that, we further explain the major innovations of E5, including the CCPairs dataset and the two-step pre-training/finetuning strategy, along with a brief introduction on its model variants and usage.

In Evaluation, we highlight some of the most important evaluation results to demonstrate the performance of E5.

My takeaways:

• Contrastive learning on top of a large-scale, diverse and high-quality dataset is crucial in learning text embedding models.
• Carefully designed training strategies (for example, contrastive pre-training with instruction prefix, hard-negative mining, knowledge distillation, etc.) can further boost model performance.

Thanks for reading!

Reference

• [1] Text Embeddings by Weakly-Supervised Contrastive Pre-training
• [2] SimCSE: Simple Contrastive Learning of Sentence Embeddings
• [3] Sentence-BERT
• [4] Knowledge Distillation: A Survey