NV Embed

Analysis

Why This Paper Matters

Text embeddings are the backbone of modern retrieval and semantic search systems. While encoder-only models like BERT have dominated this space, decoder-only LLMs offer richer contextual understanding but struggle to produce high-quality embeddings due to their causal attention structure. NV-Embed addresses this gap by introducing two key innovations: a latent attention layer for better pooling and the removal of causal masking during contrastive training. The result is a generalist embedding model that achieves state-of-the-art results on the MTEB benchmark, demonstrating that LLMs can be effectively repurposed for embedding tasks without sacrificing simplicity or reproducibility.

How It Works

The core idea is to adapt a decoder-only LLM (Mistral 7B) for embedding tasks by making two architectural changes and employing a two-stage training recipe.

Bidirectional Attention: In standard decoder-only LLMs, the causal attention mask prevents tokens from attending to future positions, which is essential for autoregressive generation but limits representation power for embeddings. NV-Embed removes this mask during contrastive learning, allowing each token to attend to all others in the sequence. This bidirectional attention significantly improves the quality of the resulting embeddings.

Latent Attention Layer: Traditional pooling methods have drawbacks. Mean pooling dilutes important information from key phrases, while using the last <EOS> token suffers from recency bias. To overcome these, the authors propose a latent attention layer that takes the last layer hidden states as queries and uses a trainable latent array as keys and values. The output is computed as:

This is followed by an MLP with two linear layers and GELU activation, then mean pooling. The latent array acts as a learned dictionary that can capture more expressive representations of the sequence.

Two-Stage Instruction-Tuning: To create a generalist model, training is split into two stages. In the first stage, the model is trained on retrieval datasets (e.g., MS MARCO, HotpotQA) using contrastive learning with in-batch negatives and mined hard negatives. This stage focuses on the more challenging retrieval task. In the second stage, the model is trained on a mix of retrieval and non-retrieval datasets (classification, clustering, semantic similarity) without in-batch negatives, allowing it to adapt to diverse tasks and instructions.

Results

NV-Embed achieves a new record high score of 69.32 on the MTEB benchmark with 56 tasks and also attains the highest score of 59.36 on 15 retrieval tasks from the BEIR benchmark. Ablation studies confirm the effectiveness of each component:

• Bidirectional vs. Causal Attention: Bidirectional attention consistently outperforms causal attention across all pooling methods.
• Pooling Methods: Latent attention outperforms mean pooling and last-token pooling. Self-attention does not provide additional gains.
• Two-Stage Training: The second stage improves performance on non-retrieval tasks without harming retrieval performance.

Significance

NV-Embed provides a clean, reproducible recipe for transforming decoder-only LLMs into state-of-the-art embedding models. The latent attention layer is a particularly elegant solution to the pooling problem, offering a learnable way to extract task-relevant information from token sequences. By demonstrating that bidirectional attention and two-stage instruction-tuning can dramatically improve embedding quality, this work opens the door for using larger and more capable LLMs as generalist embedding models. The model is available on HuggingFace, making it accessible for practitioners. Future work could explore scaling to larger base models, different latent array sizes, or more efficient training techniques.