Improving Text Embeddings with Large Language Models

In plain words

A text embedding model turns a sentence or a passage into a list of numbers, so that texts with similar meaning get similar lists. Search engines and retrieval systems use these lists to find relevant documents. The usual recipe trains such a model in several stages on huge collections of loosely matched text pairs, then tunes it on a few carefully labelled datasets. That takes a lot of engineering and only covers the tasks and languages those datasets cover.

This paper asks whether a large language model can skip most of that work. A strong proprietary model is asked to invent search-style tasks, together with matching examples, in many languages. An open 7-billion-parameter language model is then fine-tuned on those examples with a standard contrastive loss. Each query is prefixed with a one-line description of the task, and documents are not. The embedding is read from a special end-of-text position.

The result is a model trained on generated data alone that performs well on English benchmarks. Adding labelled datasets to the training mix makes it stronger still. It is weaker on low-resource languages, and the paper also tests whether it can find a short passkey hidden in very long documents.

TL;DR

• Synthetic data: about 500k examples generated with GPT-4 and GPT-3.5-Turbo, covering 93 languages, from a two-step prompting scheme §3.1 §4.1.
• Training: Mistral-7B with LoRA rank 16, one epoch, InfoNCE with temperature 0.02, queries prefixed with a task instruction §3.2 §4.2 app. A.
• Author-reported: MTEB (56 English datasets) average 66.6 with full data, against 64.2 for BGE-large-en-v1.5; 63.1 with synthetic data only tab. 1.
• Author-reported: BEIR (15 retrieval datasets) average nDCG@10 56.9 tab. 3.
• Pitfall: the repository contains only evaluation code. The training recipe, data generation and passkey experiments cannot be checked against code.

Problem

Strong text embeddings have usually needed multi-stage training: weakly supervised pre-training on billions of pairs, then fine-tuning on labelled data. This is costly to build, and the labelled data is limited in task diversity and language coverage. Encoders are typically small BERT-style models, which leaves out recent LLM advances such as longer context §1.

Method

The method has two parts. A proprietary LLM generates synthetic training examples for many embedding tasks §3.1. A decoder-only LLM is then fine-tuned into an embedding model with a contrastive loss §3.2. The embedding of a text is the last-layer vector at a final [EOS] token §3.2.

Synthetic data §3.1:

1. Tasks are grouped as asymmetric (query and document related but not paraphrases) or symmetric (STS and bitext retrieval) §3.1.
2. Asymmetric tasks are split by length into four subgroups: short-long, long-short, short-short and long-long. Each subgroup has its own prompt templates §3.1 tab. 8 tab. 9 tab. 10 tab. 11.
3. For each asymmetric subgroup, the LLM first brainstorms about 20 one-sentence task definitions, then writes one JSON example per task §3.1 app. C. Symmetric tasks skip the brainstorming step §3.1.
4. Placeholders in the templates are sampled at runtime, which adds diversity §3.1 tab. 8 tab. 9 tab. 12 tab. 13. The language is sampled from the XLM-R language list with extra weight on high-resource languages §3.1 app. C.
5. Outputs that are not valid JSON are discarded, and exact duplicates are removed by string matching §3.1.
6. Generation uses Azure OpenAI with GPT-4 and GPT-3.5-Turbo, about 25% from GPT-3.5-Turbo, at temperature 1.0 and top-p 1.0 §4.1 app. C.

Training §3.2:

7. For a relevant pair (q^+, d^+), the query gets an instruction prefix (equation 1). Synthetic data uses its brainstormed task definition. Other datasets use manually written definitions, listed in tab. 14. Documents get no prefix §3.2.
8. Query and document are both passed through the LLM with [EOS] appended, and the last-layer [EOS] vector is the embedding §3.2.
9. The loss is InfoNCE over the positive, the in-batch negatives and the hard negatives (equation 2). Scores are cosine similarities divided by a fixed temperature \tau = 0.02 (equation 3) §3.2.
10. Fine-tuning uses Mistral-7B, one epoch, LoRA rank 16 on all linear layers, gradient checkpointing, mixed precision and DeepSpeed ZeRO-3 §4.2 app. A.
11. The full-data model mixes the synthetic data with 13 public datasets, about 1.8M examples after sampling §4.2 app. A. The default ablation setting uses synthetic data plus MS-MARCO only tab. 5.

Equations, as written in the paper:

q^{+}_{\mathrm{inst}} = \text{Instruct: } \{\text{task definition}\}\ \backslash n\ \text{Query: } \{q^{+}\} eq. 1

\mathbb{L} = -\log \frac{\phi(q^{+}_{\mathrm{inst}}, d^{+})}{\phi(q^{+}_{\mathrm{inst}}, d^{+}) + \sum_{n_i \in \mathbb{N}} \phi(q^{+}_{\mathrm{inst}}, n_i)} eq. 2

\phi(q, d) = \exp\left(\frac{1}{\tau} \cos(\mathbf{h}_q, \mathbf{h}_d)\right), \quad \tau = 0.02 eq. 3

Pseudocode, in our words. The query instruction and the embedding:

# Put the task instruction in front of a query  [§3.2] [eq. 1]
def instruct(task_definition, query):
    return "Instruct: " + task_definition + "\nQuery: " + query

# Encode a batch of texts, appending [EOS] to each  [§3.2]
def encode(texts):
    # Tokenize and append an [EOS] token to every text  [§3.2]
    token_ids = tokenize_append_eos(texts)           # (B, L)
    # Run the LLM with LoRA adapters active  [§4.2] [app. A]
    hidden = llm_with_lora(token_ids)                # (B, L, d)
    # Take the last-layer vector at the [EOS] position  [§3.2] [tab. 5]
    rows = arange(len(texts))
    return hidden[rows, eos_positions]               # (B, d)

The contrastive loss:

# Fixed temperature for all scores  [§3.2] [eq. 3]
TAU = 0.02

# Temperature-scaled cosine similarity matrix  [§3.2] [eq. 3]
def score(query_vecs, doc_vecs):
    cosines = normalize(query_vecs) @ normalize(doc_vecs).T
    return exp(cosines / TAU)                        # (Q, D)

# InfoNCE over positives, in-batch and hard negatives  [§3.2] [eq. 2]
def info_nce_loss(query_vecs, pos_vecs, hard_neg_vecs):
    # Candidates: every positive and hard negative in the batch  [§3.2]
    docs = concat(pos_vecs, hard_neg_vecs)           # (2B, d)
    scores = score(query_vecs, docs)                 # (B, 2B)
    # Query i's own positive is column i; all other columns are negatives
    rows = arange(len(query_vecs))
    positives = scores[rows, rows]                   # (B,)
    losses = -log(positives / scores.sum(axis=1))    # sum includes positive
    # Reduction over the batch is not stated in the paper
    return losses.mean()

Synthetic data generation, in our words. Asymmetric subgroups only:

# Generate examples for one asymmetric subgroup  [§3.1] [app. C]
def generate_asymmetric(subgroup):
    # Step one: ask for about 20 one-sentence task definitions  [app. C]
    tasks = llm(brainstorm_prompt(subgroup))
    examples = []
    for task in tasks:
        # Placeholders and language are sampled per call  [§3.1]
        values = sample_placeholders(subgroup)
        language = sample_language()
        # Step two: one JSON example for this task  [§3.1] [app. C]
        raw = llm(example_prompt(subgroup, task, values, language))
        example = parse_json(raw)
        # Discard anything that is not valid JSON  [§3.1]
        if example is not None:
            examples.append(example)
    # Remove exact string duplicates  [§3.1]
    return unique_by_exact_string(examples)

Hyperparameters

Placeholder values in the synthetic prompt templates tab. 8 to tab. 13:

The short-short and long-long templates list no placeholders beyond the task and language tab. 10 tab. 11.

Results

Author-reported throughout. MTEB here means the 56 English datasets. BEIR means the 15 public retrieval datasets in BEIR. Full-data models use the synthetic data plus labelled data §4.2 app. A. Baselines are as listed in each table, including commercial models whose training data is not public §4.3.

Headline comparisons:

Full tables follow, as printed in the paper.

Table 1, MTEB 56 English datasets, category averages tab. 1:

Table 2, MIRACL dev nDCG@10, high-resource and low-resource languages tab. 2:

Table 3, comparison with commercial and leaderboard models, as of 2023-12-22 tab. 3:

Table 4, bitext mining tab. 4:

Table 5, ablations on MTEB, all models trained on synthetic plus MS-MARCO data unless stated. The first row is the default: last-token pooling, LoRA rank 16, natural-language instructions tab. 5:

Table 6, MIRACL dev, all 16 languages, nDCG@10 and Recall@100 tab. 6:

Table 7, contrastive pre-training ablation tab. 7 fig. 3:

The paper pre-trains Mistral-7B with the mE5 recipe for 10k steps in the last row tab. 7.

Table 17, per-dataset MTEB scores, four settings tab. 17:

Other results in prose, with their figures and tables:

• Long context §5.2 fig. 4 fig. 5: the personalized passkey task has each document carry a person's name and a random passkey at a random position, with 100 candidates. Each context length uses 50 random queries. The 4k-window default reaches 100% within 4k tokens and degrades beyond it. Naively enlarging the window to 32k is worse. A RoPE base of 10^5 gives over 90% within 32k tokens, with a small loss on shorter contexts. The paper does not quantify that loss in the text.
• Training-data effect §5.1 fig. 3 tab. 7: contrastive pre-training helps XLM-R large by 8.2 retrieval points but has negligible effect on Mistral-7B.
• Instructions §5.3 tab. 5: removing the instruction costs 4.2 average points; the task-type prefix variant costs the same.
• Contamination app. B: a string-match check on MTEB test sets found no substantial contamination that would change the main findings.

Paper vs code

Searched: the paper's artifacts paragraph links github.com/microsoft/unilm under the e5 folder app. A. The repository is by the authors' organisation, so it was used without further search. It was cloned at depth 1 at commit 31c5b904ca1bf2afb4c234a6675c683a4e5fc7cd and its root LICENSE is MIT.

• Coverage: the e5/ folder contains evaluation scripts only: MTEB retrieval, MTEB non-retrieval, shared helpers, a model-config table and two shell scripts mteb_beir_eval.py README.md. No training loop, LoRA setup, loss, data-generation script or passkey experiment is in the repository, although the paper says the implementation is based on the HuggingFace PEFT library app. A. Every training detail above is therefore paper-only.
• Query instruction: the paper's equation 1 and the code's format match. The code builds Instruct: {task}\nQuery: and prepends it to each query utils.py mteb_beir_eval.py [code: e5/mteb_beir_eval.py#L111-L112].
• Document text: the paper says documents get no instruction prefix §3.2. The code does not add one for instruct models, but it builds each BEIR document as title plus space plus text, using the title when the corpus has one mteb_beir_eval.py. The paper does not mention titles.
• [EOS] token: the paper appends [EOS] to queries and documents §3.2. The evaluation code never adds an EOS token; it tokenises with the default tokenizer call utils.py. Whether the tokenizer appends one is not checked here.
• Pooling: the default for the Mistral model is last, which matches the paper's last-token default tab. 5 model_config.py. The pool function takes the final position when the batch is left-padded, and otherwise the last non-padding position utils.py. Behaviour therefore depends on the tokenizer's padding side.
• Normalisation: retrieval embeddings are L2-normalised [code: e5/mteb_beir_eval.py#L84], which makes the dot product equal to the cosine of equation 3. Non-retrieval tasks are normalised except for classification mteb_except_retrieval_eval.py. The paper does not state this.
• Non-retrieval protocol: the code prepends the task instruction to every input, including both sides of pair tasks such as STS mteb_except_retrieval_eval.py [code: e5/mteb_except_retrieval_eval.py#L103-L106]. The paper does not describe the non-retrieval protocol §4.2.
• Length: evaluation truncates to 512 tokens, matching §4.2 [code: e5/utils.py#L73-L82]. Training truncation at 512 app. A cannot be checked.
• Precision: evaluation loads weights in float16 and runs under autocast mteb_beir_eval.py mteb_beir_eval.py. The paper does not state an evaluation precision.
• Task instructions: the retrieval instruction strings I compared match the paper's evaluation instruction table, for example ArguAna, NQ and SciFact utils.py tab. 15. The paper's training instructions tab. 14 are not in the code.
• Released checkpoint: the README reports 56.9 BEIR for e5-mistral-7b-instruct README.md, which matches the paper's full-data BEIR number tab. 3. That match does not verify the training recipe, since the training code is absent.
• Released artefacts: the paper says it releases models and evaluation scripts at unilm/e5 app. A. The README links only the passkey dataset README.md. No synthetic training data is in the repository.

Pitfalls

• Training is underspecified. Not stated: optimizer, LoRA alpha and dropout, reduction of the loss over the batch, whether batches mix datasets, the precision used in training, and how one hard negative is chosen from the top 100 mined by mE5-base app. A.
• Equation 2 uses the positive and all negatives in its denominator. The paper does not say whether a hard negative of one query counts as a negative for every other query in the batch. The pseudocode above assumes it does.
• Not every generated example is a query-document pair. Long-short examples carry an input text, a label and a misleading label tab. 9. Symmetric examples are triples with graded scores tab. 12 tab. 13. Short-short and long-long examples have no negative document tab. 10 tab. 11. The paper does not say how these become positive and negative pairs, or how the scores are used.
• The text and the table disagree on one placeholder. Section 3.1 gives query lengths as "less than 5 words, 5-10 words, at least 10 words" §3.1. Table 8 gives "5 to 15 words" tab. 8. Follow the table if you are reproducing the templates, and note the conflict.
• Data quality: the paper says part of the GPT-3.5-Turbo output does not follow the template guidelines, and kept it §4.1. Filtering is limited to JSON parsing and exact-string duplicates §3.1.
• Contamination: the check in app. B is string matching. It found four DBPedia test questions in TriviaQA training data and notes overlapping Wikipedia corpora across DBPedia, NQ and TriviaQA. Several training sets are also MTEB evaluation sets: MS-MARCO, NQ, HotpotQA, FEVER and Quora app. A tab. 17. The paper does not report results with those removed.
• Long context: the paper calls the 4k setting a "sliding window size" without defining it further. It does not say whether the RoPE base change of 10^5 is applied during training or only at inference §5.2. The passkey data is on HuggingFace README.md. The passkey code is not released.
• Low-resource languages: the model is weaker there. In Table 2, sw scores 68.4 against mE5-large 74.9 tab. 2. On bn it is level with mE5-base, 70.3 against 70.2 tab. 6. Only about 1k synthetic examples per low-resource language were generated on average §4.1.
• The paper says a fair comparison with commercial models is not feasible, because their details are not public §4.3. The leaderboard comparison is dated 2023-12-22 tab. 3.
• Terms of the generator: the paper does not discuss the terms under which GPT-4 and GPT-3.5-Turbo outputs were used. Check the provider terms before you reproduce the data pipeline.
• Cost: about 180M generated tokens §4.1; about 18 hours on 32 V100 GPUs app. A; about 3 days of evaluation on 8 V100 GPUs §4.2; embeddings of 4096 dimensions [Limitations].
• Retrieval needs the instruction. The same instruction must be used at query time. Documents are encoded without any prefix, so the document index never has to be rebuilt when the task changes §3.2 §5.3.

Lineage

• Builds on: E5, the multi-stage baseline this paper argues against §1. Also the Mistral-7B base model, Mistral 7B §4.2.
• Uses, without a card here: the LoRA recipe follows RankLLaMA, and LoRA itself §4.2. The position-weighted mean pooling is taken from SGPT app. A.
• Compared against: mE5 (multilingual), BGE, GTE, GTR, Sentence-T5, Contriever, SimCSE, UAE-Large-V1, and commercial embedding APIs tab. 1 tab. 2 tab. 3.