Llama2Vec: Unsupervised Adaptation of Large Language Models for Dense Retrieval

In plain words

A search engine that finds documents by meaning first turns each query and each document into a list of numbers, then returns the documents whose lists look most like the query's. Large language models understand text well, but they were trained to guess the next word, so the numbers they produce for a whole text are not tuned for this comparison.

The authors teach the model, before any search training, to build a summary of a text that is good enough to rebuild the text and to guess the text that comes after it. The model does this by reading the passage followed by two short instructions, and it reads the answer from a small marker placed after each instruction. The extra training uses only ordinary Wikipedia text, with no labelled search data. After that, the usual search training is run on top, and the authors report better results than earlier models on the standard benchmarks.

TL;DR

• Llama2Vec is an unsupervised step that adapts LLaMA-2-7B to produce one embedding per text for dense retrieval. It is trained on Wikipedia with two pretext tasks, EBAE (rebuild the input) and EBAR (predict the next text) §1 §3.2.
• The two prompts share one forward pass. A modified attention mask keeps them from seeing each other, which the authors say saves almost half the compute of two separate passes §3.2 fig. 2.
• Search fine-tuning uses contrastive learning with LoRA on MS MARCO §4.1.
• Author-reported: MS MARCO passage MRR@10 43.1 (RepLLaMA 41.2) tab. 1; document MRR@100 47.9 (RepLLaMA 45.6) tab. 2; BEIR zero-shot average NDCG@10 56.4 (RepLLaMA 53.9) tab. 3.
• The released code does not match the paper in several places: the prompts, the embedding read-out, an extra next-token loss, the loss form, and position ids. Read "Paper vs code" before implementing.

Problem

Dense retrieval represents queries and documents as vectors in one space and ranks documents by similarity §3.1. Earlier work reached good results with BERT-sized encoders §2. Decoder-only LLMs carry more knowledge, but they are trained to predict the next token, so their hidden state at the last position mainly describes what comes next, not the meaning of the whole text. The authors call this a mismatch between language modelling and text embedding §1 §2.

Method

Embedding read-out (baseline). A decoder-only model sees the whole input only at its last token, so a single text embedding is usually read from a special end-of-text token. For LLaMA the paper writes this as the state of that token after the input §3.1 eq. 1.

Joint prompt. Llama2Vec feeds the input text T followed by two prompt blocks: SELF, with the words "The input sentence is:", and NEXT, with the words "The next sentence is:". Each block ends in a special token. The embedding e^α is the state at the end of the SELF block, and e^β is the state at the end of the NEXT block §3.2 fig. 1 eq. 2 eq. 3.

Attention mask. The SELF and NEXT blocks cannot attend to each other, but both can attend to T. Because T is most of the sequence, the authors say one joint pass saves almost 50 percent of the cost of two passes §3.2 fig. 2.

Pretext objective. For each embedding, a linear projection to the vocabulary predicts the tokens of its target: the input itself for e^α (EBAE) and the next sentence for e^β (EBAR). The loss is the average negative log-likelihood of those target tokens §3.2 eq. 4. The paper does not say how the two terms are weighted.

# Llama2Vec adaptation loss, as the paper describes it  [§3.2] [eq. 2-4]
# Input:  text ids T, the same text as summary target, next-text ids
# Output: scalar loss

def adaptation_loss(text_ids, next_ids, lm_head):
    # Joint prompt: T, then the SELF block, then the NEXT block  [§3.2]
    prompt = text_ids + SELF_ids + NEXT_ids
    # Blocks may not attend to each other  [§3.2] [fig. 2]
    hidden = llama(prompt, mask=joint_prompt_mask(prompt))  # (L, d)
    # Inductive embedding: state at the end of the SELF block  [eq. 2]
    e_alpha = hidden[end_of_self_block]  # (d,)
    # Deductive embedding: state at the end of the NEXT block  [eq. 3]
    e_beta = hidden[end_of_next_block]  # (d,)
    # Predict each target token from its embedding  [§3.2] [eq. 4]
    loss_alpha = mean(-log_softmax(lm_head(e_alpha))[text_ids])
    loss_beta = mean(-log_softmax(lm_head(e_beta))[next_ids])
    # Paper does not state how the two terms are combined
    return loss_alpha + loss_beta

Fine-tuning. The adapted model is trained with contrastive learning on (query, answer) pairs: the positive is scored against the other passages in the batch and against hard negatives §3.2 eq. 5 §4.1. Queries use the NEXT prompt and answers use the SELF prompt in the default QA setting. For paraphrase tasks the authors use SELF for long documents and NEXT for short texts §3.2. Fine-tuning uses LoRA on the LLM §4.1.

# Contrastive fine-tuning step  [§3.2] [eq. 5] [§4.1]
# Input:  queries, positives, hard negatives; temperature
# Output: scalar loss

def finetune_loss(queries, positives, negatives, temperature):
    # Query embedding with the NEXT prompt  [§3.2]
    query_emb = normalize(embed(queries, block="NEXT"))  # (B, d)
    # Candidates (positives and negatives) with the SELF prompt  [§3.2]
    cand_emb = normalize(embed(candidates, block="SELF"))  # (B*(K+1), d)
    # Similarity of every query with every candidate in the batch
    # [code: research/LLARA/finetune/modeling.py#L107-L110]
    scores = query_emb @ cand_emb.T / temperature  # (B, B*(K+1))
    # Each query's positive sits at its own index; the rest are negatives
    # [code: research/LLARA/finetune/modeling.py#L128-L130]
    return cross_entropy(scores, positive_index)

Key equations, in the paper's notation:

\boldsymbol{e}_t \leftarrow \mathrm{LLaMA}(T)[\langle\backslash s\rangle] eq. 1

\boldsymbol{e}^{\alpha}_t \leftarrow \mathrm{LLaMA}(T,\mathrm{SELF},\langle\backslash s\rangle)[-1] eq. 2

\boldsymbol{e}^{\beta}_t \leftarrow \mathrm{LLaMA}(T,\mathrm{NEXT},\langle\backslash s\rangle)[-1] eq. 3

\min -\frac{1}{|\mathcal{T}|}\sum_{t\in\mathcal{T}} \log\frac{\exp(\boldsymbol{e}^{\top}\boldsymbol{W}_{t})} {\sum_{v\in V}\exp(\boldsymbol{e}^{\top}\boldsymbol{W}_{v})} eq. 4

\min \sum_{q} -\log\frac{\exp(\langle \boldsymbol{e}^{\alpha}_{q},\boldsymbol{e}^{\beta}_{a}\rangle)} {\sum_{a'\in\mathcal{A}'}\exp(\langle \boldsymbol{e}^{\alpha}_{q},\boldsymbol{e}^{\beta}_{a'}\rangle)} eq. 5

Eq. 5 is printed with the query in α and the answer in β, which contradicts the text that follows it (see Pitfalls).

Hyperparameters

Adaptation and fine-tuning settings. "not stated" means the paper gives no value.

Results

Author-reported. Setup: LLaMA-2-7B adapted on Wikipedia, then fine-tuned on MS MARCO passage or document data; evaluated on MS MARCO dev and DL'19 and DL'20 for retrieval, and on BEIR (zero-shot) and a LlamaIndex set §4.2 §4.3. The paper reports no variance or seed information.

Passage retrieval on MS MARCO tab. 1. Columns: dev MRR@10, dev Recall@1000, DL'19 NDCG@10, DL'20 NDCG@10.

Document retrieval on MS MARCO tab. 2. Columns: dev MRR@100, dev Recall@100, DL'19 NDCG@10, DL'20 NDCG@10.

Zero-shot retrieval on BEIR, NDCG@10 tab. 3. Columns are the models; the paper gives no per-model size for BM25 or Ada-2.

I recomputed the Llama2Vec column: the 14 values average 56.4, which matches the printed average. Llama2Vec beats BM25 on 13 of 14 datasets (not on Touche) and RepLLaMA on 12 of 14 (not on FEVER or SciFact).

Llama Index zero-shot retrieval tab. 4. The dataset is not described in the paper; the metric names are MRR and Hit Rate.

Prompt usage on BEIR, NDCG@10 tab. 5. N2S = NEXT for the query and SELF for the document, S2S = SELF for both, N2N = NEXT for both, None = no prompt, Ada* = the per-dataset choice the paper argues for.

The Ada* values per dataset are the same as the Llama2Vec column of tab. 3. The printed Ada* average (57.4) does not match its own 14 values, which average 56.4. I also get 55.0 for the N2S row, against 56.1 printed.

Lexical similarity of top-N vocabulary predictions, measured with BM25 tab. 6. Units are not stated; the paper says larger means more lexical overlap.

Embedding dimension tab. 7. The metric is not named in the caption. DimRed means a jointly learned projection during fine-tuning; DimRed* means the projection is learned by distillation on a frozen retriever §4.4; Sparse keeps the top-N entries §4.4.

The technical-factor studies are the prompt study on BEIR tab. 5, the lexical study on MS MARCO tab. 6 and the dimension study tab. 7. The lexical study supports the claim that adaptation raises BM25 similarity between the query's and the answer's predicted vocabulary §4.4.

Paper vs code

Official repository: https://github.com/FlagOpen/FlagEmbedding at commit fd1a2bdf69488ffebe0327999d4400d8c8058a0b, MIT licence, copyright held by Shitao Xiao, an author. The paper's abstract links it [abstract]. The method is under research/LLARA/ (the paper's earlier name is LLARA). The repository contains the adaptation code (pretrain/), the fine-tuning code (finetune/), a toy dataset and a README, but not the MS MARCO or BEIR evaluation scripts and not the Wikipedia data, which the README links on Hugging Face README.md.

Differences between paper and code:

• Prompt text. The paper gives SELF as "The input sentence is:" and NEXT as "The next sentence is:" §3.2. The code wraps the input in quotes and adds a suffix that asks to summarise or predict the passage in eight words: "summarize the above passage within eight words" and "predict the following passage within eight words" data.py data.py.
• Read-out tokens. The paper uses one special token per block, with the state at that token as the embedding §3.2 eq. 2 eq. 3. The code adds sixteen new tokens <s1> to <s16> and resizes the embedding matrix; each block uses eight of them run.py modeling.py.
• Embedding at inference and in fine-tuning. The paper reads the single token's state eq. 2 eq. 3. The code takes the mean of the last eight hidden states, for queries and passages alike modeling.py README.md. The README's usage example, which loads the released checkpoints, takes this mean.
• Adaptation loss. The paper has only the multi-class term of eq. 4. The code adds next-token cross-entropy on the input passage tokens modeling.py, which the paper does not mention.
• Bag-of-words term. The code does not use a per-token average over the target sentence. It max-pools the logits over the eight positions of each block, applies log-softmax, averages the negative log-probability over the unique target ids above 2 (which drops padding and special tokens), and divides by 10 modeling.py modeling.py. The two terms are then averaged modeling.py.
• Attention mask. The code blocks the predict block from the summary block columns and leaves the rest of the causal mask as is modeling.py. Under a causal mask, this is the same as the paper's "mutually invisible" rule.
• Position ids. The code sets the predict block's position ids to those of the summary block modeling.py. The paper does not describe position ids.
• Attention implementation. The adaptation code refuses FlashAttention-2 modeling.py, so the README disables it README.md.
• Stop words. The code has an option to remove stop words and punctuation from the bag-of-words targets. It is off by default arguments.py data.py, and the README turns it on README.md. The paper does not mention it.
• Adaptation data. The paper says "next sentence" §3.2. The toy data has whole passages: the input, the same input as the summary target, and a following passage toy_pretrain_data.jsonl. How the real corpus was split is not stated.
• Fine-tuning temperature. Not in the paper. Code default 0.02 arguments.py; README 0.01 README.md.
• Hard negatives. The paper uses ANN hard negatives §4.1. The code reads a neg field from each training record and samples from it data.py. Where those negatives come from is not in the repository. The README says the released BEIR model uses BM25 negatives and the passage model uses dense-retriever negatives README.md README.md.
• Batch size and steps. The paper uses batch 256 for 10,000 steps §4.1. The README example uses batch 1 per GPU on 8 GPUs for one epoch on toy data README.md, and says to use the paper's hyperparameters for the full data.
• Title. The README cites an earlier title, "Making Large Language Models A Better Foundation For Dense Retrieval" README.md.

No discrepancies were found for the learning rate: 1e-5 in both §4.1 and README.md.

Pitfalls

• Read-out choice. The paper's embedding is one special-token state. The code's embedding is the mean of eight special-token states, and the README's usage example, which loads the released models, takes that mean. Implement the code's read-out if you want the reported numbers to be reachable, and state that you did so.
• Eq. 5 notation. eq. 5 puts the query embedding in α (SELF) and the answer embedding in β (NEXT). The text right after says the query uses NEXT and the answer uses SELF §3.2, and the code does the same data.py README.md. Follow the text.
• Eq. 2 and eq. 3 use index −1. That is the last position of the separate prompt. In the joint prompt, the paper reads the SELF and NEXT states from their own blocks, not from the last position. The index notation does not describe the joint form.
• Joint prompt only during adaptation. The custom mask exists only in the adaptation code. Fine-tuning and inference encode each block as its own prompt with the standard model modeling.py README.md. An implementation that uses the joint mask at inference would differ from the released model.
• Multi-term loss. The paper does not say how the two pretext losses are combined. The code uses one next-token loss plus two bag-of-words terms with a factor of 1/10 modeling.py.
• Max-pooling in the bag-of-words term. The code takes the max over eight positions before the softmax. This is not the per-token formula of eq. 4, and it is not in the paper.
• Adaptation is full training in the README, not LoRA. The paper does not say which. The released command sets use_lora False README.md, while the code default is LoRA arguments.py.
• Sequence length. The paper uses 1024 tokens §4.1. The code default is 512 and the README uses 128 arguments.py README.md.
• Stop-word removal is not in the paper. The README enables it. Whether the reported results used it is not stated.
• Optimizer, warmup and precision are not stated in the paper. Code uses the Hugging Face Trainer defaults with a DeepSpeed stage-1 config, a warmup of 0.1, and bf16 in the README stage1.json README.md README.md.
• Which MS MARCO model feeds BEIR. The paper evaluates a model fine-tuned on MS MARCO passages §4.3, but it does not say which negative sampler it used. The released passage and BEIR models differ in negatives README.md README.md.
• Table 5 average. The printed Ada* average (57.4) and N2S average (56.1) do not match their rows (56.4 and 55.0 by my recount). Use the per-dataset values; do not rely on the printed averages.
• Table 7 metric and sparse size. The metric is not named, and N for the sparse variant is not stated §4.4.
• Llama Index set. The dataset, its size and the metric definitions are not described tab. 4. The comparison models may use different settings.
• Lexical study. The BM25 parameters and the units of tab. 6 are not stated.
• Cost claim. The 50 percent saving is the authors' statement from the shared-prompt design §3.2. No timing measurement is reported.
• Single runs. The paper reports one number per cell with no seeds or significance tests tab. 1 tab. 2 tab. 3. Differences of a point or two on small BEIR sets may not hold up.
• Ethical limits. The authors do not recommend the model for sensitive retrieval, because LLaMA-2 carries its biases into the embeddings §7.

Lineage

• Builds on: LLaMA-2-7B (Touvron et al. 2023) as the backbone §4.1; the fine-tuning recipe of RepLLaMA (Ma et al. 2023), which this paper follows and uses as its closest baseline §4.1 §4.2; LoRA (Hu et al. 2021) §4.1; ANN hard negatives (Xiong et al. 2020) §4.1; the Wikipedia corpus curated by DPR (Karpukhin et al. 2020) §4.1. None of these has a card in this wiki yet.
• Compared against: BM25, ANCE, ADORE, Condenser, coCondenser, TAS-B, RocketQAv2, AR2+SimANS, GTR-XXL, SimLM, RetroMAE, RetroMAEv2, SGPT, OpenAI-Ada-002, bge-m3 and OpenAI text-embedding-3 tab. 1 tab. 2 tab. 3 tab. 4. The strongest baseline is RepLLaMA §4.2.
• Superseded by: not stated in the wiki.