PaperlayerView exact revision on arXiv

Source-linked AI summary

LLM2Vec: Large Language Models Are Secretly Powerful Text Encoders

Parishad BehnamGhader, Vaibhav Adlakha, Marius Mosbach, Dzmitry Bahdanau, Nicolas Chapados, Siva Reddy

arXiv:2404.05961v2cs.CLcs.AI

TL;DR

Decoder-only LLMs are underused for text embeddings because causal attention limits access to full-sequence context. LLM2Vec transforms them with bidirectional attention, masked next token prediction, and unsupervised contrastive learning, achieving strong word-level results and state-of-the-art MTEB performance. The approach requires little, data- and parameter-efficient adaptation, while the paper notes model-size challenges and leaves Mistral’s bidirectional behavior for future investigation.

  • Problem

    Causal attention limits decoder-only LLMs’ ability to capture information across entire input sequences, restricting their use for rich text embeddings.

  • Method

    LLM2Vec transforms pre-trained decoder-only LLMs into universal text encoders through bidirectional attention, masked next token prediction, and unsupervised contrastive learning.

  • Results

    LLM2Vec-transformed models outperform strong encoder-only models on word-level tasks, reach 56.8 on MTEB in the best unsupervised setting, and achieve state-of-the-art public-data performance with supervised contrastive training.

  • Takeaways & Limitations

    Decoder-only LLMs can produce universal text embeddings with little adaptation, without labeled data or expensive synthetic GPT-4-generated data.

  • Takeaways & Limitations

    The paper acknowledges challenges associated with the large size of decoder-only LLMs and leaves detailed investigation of Mistral’s bidirectional behavior for future work.

Abstract

from arXiv · show

Large decoder-only language models (LLMs) are the state-of-the-art models on most of today's NLP tasks and benchmarks. Yet, the community is only slowly adopting these models for text embedding tasks, which require rich contextualized representations. In this work, we introduce LLM2Vec, a simple unsupervised approach that can transform any decoder-only LLM into a strong text encoder. LLM2Vec consists of three simple steps: 1) enabling bidirectional attention, 2) masked next token prediction, and 3) unsupervised contrastive learning. We demonstrate the effectiveness of LLM2Vec by applying it to 4 popular LLMs ranging from 1.3B to 8B parameters and evaluate the transformed models on English word- and sequence-level tasks. We outperform encoder-only models by a large margin on word-level tasks and reach a new unsupervised state-of-the-art performance on the Massive Text Embeddings Benchmark (MTEB). Moreover, when combining LLM2Vec with supervised contrastive learning, we achieve state-of-the-art performance on MTEB among models that train only on publicly available data (as of May 24, 2024). Our strong empirical results and extensive analysis demonstrate that LLMs can be effectively transformed into universal text encoders in a parameter-efficient manner without the need for expensive adaptation or synthetic GPT-4 generated data.

1 Introduction

LLM2Vec addresses the limited use of decoder-only LLMs for text embeddings by transforming them into universal text encoders through a simple unsupervised, efficient pipeline. Across word- and sequence-level evaluations, the transformed models outperform encoder-only baselines and achieve state-of-the-art MTEB results.

  • Motivation: Decoder-only LLMs offer sample-efficiency, mature tooling, and strong capabilities, but causal attention restricts token representations to preceding context.This restriction prevents representations from capturing information across the entire input sequence, which is important for text embeddings.
  • Approach: LLM2Vec transforms any pre-trained decoder-only LLM into a universal text encoder using bidirectional attention, masked next token prediction, and unsupervised contrastive learning.The approach requires no labeled data and is described as data- and parameter-efficient.
  • Evaluation: 4 decoder-only LLMs ranging from 1.3B to 8B parameters are evaluated on word- and sequence-level tasks.The evaluated models include S-LLaMA-1.3B, LLaMA-2-7B, Mistral-7B, and Meta-LLaMA-3-8B.
  • Results: On chunking, named-entity recognition, and part-of-speech tagging, LLM2Vec-transformed models outperform strong encoder-only models by a large margin.The result is presented as evidence that the method produces rich contextualized token representations.
  • Conclusion: Decoder-only LLMs can produce universal text embeddings with only little adaptation, and the code and pre-trained models are publicly available.The paper also acknowledges challenges associated with the large size of these models.

2 LLM2Vec

LLM2Vec adapts decoder-only LLMs for embedding by enabling bidirectional attention, training with masked next token prediction, and applying unsupervised contrastive learning. The method adapts representations for token- and sequence-level tasks, while its evaluation compares transformed models against encoder-only baselines.

  • Enabling Bidirectional Attention: Bidirectional attention replaces the causal mask with an all-ones matrix, giving each token access to every other token in the sequence.Naively enabling bidirectional attention decreases embedding performance for most models, motivating further adaptation.
  • Masked Next Token Prediction: Masked next token prediction adapts models to bidirectional attention by masking input tokens and predicting them from past and future context.The loss uses logits from the representation at the previous position rather than the masked position itself.
  • Unsupervised Contrastive Learning: Unsupervised contrastive learning complements the first two steps by adapting representations for sequence-level tasks.The step addresses that decoder-only LLMs are not explicitly trained to capture the context of an entire sequence.
  • Evaluation: The experiments evaluate chunking, named-entity recognition, and part-of-speech tagging with linear probing against DeBERTa-v3-large.The tasks use CoNLL-2003, with frozen representations and task-specific linear classifiers.

3 LLM2Vec-transformed models are strong unsupervised text embedders

LLM2Vec produces strong unsupervised word- and sequence-level text embeddings by adapting decoder-only LLMs with bidirectional attention, MNTP, and SimCSE. Across MTEB, it improves over causal baselines and reaches state-of-the-art unsupervised performance.

  • Word-level tasks: LLM2Vec is evaluated on chunking, named-entity recognition, and part-of-speech tagging using frozen representations with task-specific linear classifiers.The comparison includes DeBERTa-v3-large as an encoder-only baseline.
  • Word-level tasks: MNTP improves word-level performance for every model and task, while adding SimCSE performs worse because it adapts representations for sequence-level tasks.Chunking improves by 5% for S-LLaMA-1.3B, 4% for LLaMA-2-7B, and 4% for Mistral-7B.
  • Sequence-level tasks: Mean pooling performs best among EOS, mean, and weighted mean pooling on the 15-task MTEB subset, where causal attention and EOS pooling are sub-optimal.LLM2Vec is compatible with all three pooling choices.
  • Sequence-level tasks: SimCSE further boosts S-LLaMA-1.3B, LLaMA-2-7B, and Mistral-7B on the MTEB subset by 49.8%, 23.2%, and 37.5% versus the best causal baseline.MNTP improves S-LLaMA-1.3B and Mistral-7B before the additional SimCSE gains.
  • Sequence-level tasks: 16.4% improvement for Mistral-7B follows bidirectional attention and MNTP on full MTEB, with considerable gains observed for all four evaluated models.The full MTEB evaluation uses models selected from ablations on a representative 15-task subset.
  • Sequence-level tasks: LLM2Vec outperforms Echo embeddings for three models and is almost on par for Mistral-7B, while avoiding Echo’s doubled sequence length during inference.The direct comparison uses only LLM2Vec’s first two steps because SimCSE learns sequence representations.
  • Sequence-level tasks: Adding SimCSE makes LLM2Vec Mistral-7B state of the art among unsupervised models on full MTEB, with a score of 56.80.Overall, LLM2Vec-transformed models outperform previous unsupervised approaches on the challenging MTEB leaderboard.

4 How does LLM2Vec affect a model?

LLM2Vec increases future-token sensitivity in transformed representations, while Mistral-7B already exhibits similar behavior under bidirectional attention without training. The analysis links these representation changes to separation of positive and negative sentence examples.

  • Attention analysis: Figure 5 compares causal- and bidirectional-attention representations across token positions and layers without training.The comparison measures cosine similarity between the two representation types.
  • Future-token information: For S-LLaMA-1.3B, bidirectional attention plus masked next token prediction separates positive from negative examples.The positive and negative examples share a prefix, so separation indicates use of information beyond that prefix.
  • Future-token information: For Mistral-7B, every tested setup assigns higher cosine similarity to positive than negative examples.The analysis includes causal, bidirectional, and trained configurations.
  • Attention analysis: Bidirectional attention substantially changes S-LLaMA-1.3B and LLaMA-2-7B representations, producing low cosine similarity across most layers and token positions.The effect appears without additional training.
  • Attention analysis: Mistral-7B representations remain highly similar under causal and bidirectional attention throughout the analyzed layers and positions.This behavior is reported without training.
  • Interpretation: The authors speculate that Mistral models may have received some bidirectional-attention pre-training, possibly through prefix language modeling, but leave investigation for future work.This interpretation is based on replicated analyses and Mistral-7B’s strong unsupervised results.

5 Combining LLM2Vec with supervised contrastive learning

Supervised contrastive learning improves LLM2Vec-transformed models on MTEB, while the unsupervised SimCSE stage can be unnecessary for final supervised performance. Nevertheless, retaining SimCSE improves sample efficiency during supervised training.

  • Supervised MTEB: LLM2Vec transformation improves every evaluated model over the strong Uni + weighted mean baseline in supervised MTEB training.The supervised setting uses a public E5 replication, hard negatives, in-batch negatives, and LoRA fine-tuning.
  • Supervised MTEB: Meta-LLaMA-3-8B + LLM2Vec without SimCSE achieves new state-of-the-art performance among models trained only on publicly available data.This is the best-performing model in the reported supervised comparison.
  • Sample efficiency: LLM2Vec-transformed models reach better performance earlier in supervised training across S-LLaMA-1.3B, LLaMA-2-7B, and Mistral-7B.The results come from evaluations on a 15-task MTEB subset during training.
  • Sample efficiency: For S-LLaMA-1.3B, masked next token prediction alone considerably improves sample efficiency.This is especially relevant when high-quality labeled data is difficult to acquire.

6 Related Work

Related work has adapted encoder-only models through supervised and unsupervised training, while decoder-only LLM adaptations have commonly relied on final-token representations or more computationally expensive procedures. LLM2Vec differs through parameter-efficient adaptation.

  • Supervised text encoders: Supervised text encoders evolved from NLI and sentence-similarity training toward retrieval and multi-stage learning pipelines.These approaches primarily adapt BERT-like models for sentence and retrieval representations.
  • Unsupervised text encoders: Unsupervised text encoders create multiple representations of the same sentence, using input perturbations or different model instances for contrastive learning.SimCSE is identified as an approach used in this work.
  • Decoder-only LLMs: Decoder-only LLM sentence embeddings have commonly used the final hidden state of the last token as the text representation.This literature has had limited impact on sentence representation learning relative to decoder-only LLM performance on language understanding tasks.
  • Decoder-only LLMs: Prior work has explored removing or selectively disabling causal masking during supervised adaptation for classification, NER, and chunking.These studies examine the causal mask as a limitation for decoder-only representations.
  • Concurrent approaches: Compared with concurrent approaches, LLM2Vec uses parameter-efficient fine-tuning and 1000 gradient steps, while avoiding duplicated input sequences at inference.The duplicated-sequence strategy increases inference cost, especially for longer documents.

7 Conclusion

LLM2Vec transforms decoder-only LLMs into universal text embedders with strong word- and sequence-level performance in unsupervised and supervised settings. Its simplicity, compute efficiency, and sample efficiency support use in low-resource and compute-constrained scenarios.

  • Conclusion: LLM2Vec transforms decoder-only LLMs into universal text embedders and performs effectively on word- and sequence-level tasks.The evaluation covers both unsupervised and supervised settings.
  • Conclusion: Mistral-7B reaches state-of-the-art MTEB performance among unsupervised approaches after applying LLM2Vec.The conclusion also reports supervised state-of-the-art performance for Meta-LLaMA-3-8B among publicly trained approaches.
  • Conclusion: Meta-LLaMA-3-8B achieves state-of-the-art performance among approaches trained only on publicly available data when LLM2Vec is combined with supervised contrastive fine-tuning.The comparison is dated as of May 24, 2024.
  • Conclusion: The analysis identifies an intriguing property of Mistral-7B that explains its strong out-of-the-box performance with bidirectional attention.The authors connect this observation to the underlying model behavior.
  • Conclusion: The approach’s simplicity, compute efficiency, and sample efficiency make it promising for low-resource and compute-constrained scenarios.The conclusion presents these settings as avenues supported by the empirical results.

A Limitations

The paper identifies model scale, possible pre-training contamination, and English-only evaluation as important scope boundaries. Its experiments also use a representative MTEB subset and parameter-efficient training configurations.

  • Model scale: Decoder-only models can be substantially larger and more computationally intensive than smaller bidirectional encoders.The paper notes model sizes up to 540B parameters and larger output dimensions, increasing training, inference, memory, and indexing costs.
  • Data contamination: Pre-training contamination cannot be ruled out because the complete pre-training data for LLaMA-2-7B and Mistral-7B are unavailable.The authors propose evaluating on newly designed benchmarks that are not part of the models’ pre-training data.
  • Language scope: The methodology and benchmarks in this work use only English text corpora and English benchmarks.The authors state that LLM2Vec is language-agnostic but leave evaluation on other languages to future work.
  • Evaluation scope: MTEB ablations use a representative 15-task subset rather than the full benchmark.The subset covers each category in approximately the same proportion as the full MTEB, while full evaluation can take more than 40 hours for Mistral-7B.
  • Training and evaluation: MNTP improves word-level performance after bidirectional attention is enabled, while adding SimCSE performs worse on word-level tasks.The authors attribute the difference to SimCSE adapting representations for sequence-level tasks.

D.2.2 Sentence-level task results

On the MTEB subset, pooling choice depends on the attention and training setup: weighted mean is best for causal models, whereas mean pooling is best for LLM2Vec.

  • Pooling methods: Weighted mean pooling performs best for causal Uni models, while mean pooling performs best for LLM2Vec models.The comparison is reported across the MTEB subset and different pooling methods.
  • Unsupervised results: The paper additionally reports a breakdown of unsupervised LLM2Vec performance on MTEB.This breakdown appears in Table 11.

E.1 Reproducibility

The reproducibility analysis implements Echo embeddings under the paper’s evaluation setup and obtains a nearly matching score to the original report. It also finds a substantial efficiency disadvantage for Echo embeddings.

  • Echo reproduction: 55.22 versus 55.07: the reproduced Echo embedding score nearly matches the reported score on the 26-task MTEB subset.The evaluation uses Mistral-7B-Instruct-v0.1 and the same task subset for direct comparison.
  • Instruction wording: The reproduction supports the observation that exact instruction wording or template has little effect on Echo embedding performance.The paper presents this as a validation of both its implementation and the earlier observation.
  • Efficiency: Echo embeddings take much longer to evaluate because they copy the input text, creating a problem for large retrieval corpora.The paper reports approximate evaluation times using 8x 80GB A100 GPUs.

F More analysis results

Additional analyses examine cosine similarities, attention representations, instructions, training setup, and leaderboard position. They extend the reported comparisons across models and provide context for interpreting the results.

  • Analysis materials: The analysis uses toy sentences and task instructions alongside a dataset description for the reported experiments.Table 7 contains toy data for the cosine-similarity analysis, while Table 8 lists instructions for E5 datasets.
  • Cosine similarity analysis: Figure 7 compares query–positive and query–negative cosine similarities for S-LLaMA-1.3B, LLaMA-2-7B, and Mistral-7B, including Echo embeddings.The accompanying analysis states that cosine similarity results are provided for all three models and LLM2Vec-transformed models.
  • Representation analysis: Figures 8 and 9 compare causal- and bidirectional-attention representations across token positions and layers without training.Figure 9 extends the representation analysis to Mistral model variants.
  • Representation analysis: The same behavior across Mistral variants is treated as additional evidence that the Mistral-7B base model was trained with some bidirectional attention.This is presented as an interpretation of the representation analysis.
  • Training setup: The training setup uses LoRA, bfloat16 quantization, gradient checkpointing, and flash attention 2 on 8 NVIDIA A100 GPUs.Models are trained for 1000 steps with effective batch size 512 and maximum sequence length 512 tokens.
  • Leaderboard position: LLM2Vec achieves sixth place overall and first place among models trained only with publicly available data on the cited MTEB leaderboard.The ranking is reported as of 2024-03-29.

G.3 Results

LLM2Vec models rank highly on MTEB, including the sixth overall leaderboard score and the top score among models trained only on public data. Supervised evaluation is detailed separately for the Bi + MNTP variants, which performed best after supervised fine-tuning.

  • 6th on the MTEB leaderboard and 1st among models trained only with publicly available data.
  • Supervised LLM2Vec performance is reported across the full MTEB in Table 12.
  • Bi + MNTP models are used for the supervised MTEB results because they performed best after supervised fine-tuning.
Loading 2404.05961v2…