Causal2Vec: Improving Decoder-only LLMs as Embedding Models through a Contextual Token

TL;DR

Causal2Vec enhances decoder-only LLMs as embedding models by prepending a contextual token encoded by a lightweight model, improving performance without altering architecture or increasing computational costs.

Contribution

It introduces a method to generate contextualized embeddings in decoder-only LLMs using a lightweight pre-encoding step, avoiding architecture modifications.

Findings

Achieves state-of-the-art on MTEB benchmark with publicly available data.

Improves embeddings without increasing computational overhead.

Utilizes a lightweight BERT-style model for context encoding.

Abstract

Decoder-only large language models (LLMs) have been increasingly adopted to build embedding models for diverse tasks. To overcome the inherent limitations of causal attention in representation learning, many existing methods modify the attention mechanism to be bidirectional, potentially undermining LLMs' ability to extract semantic information acquired during pre-training. Meanwhile, leading unidirectional approaches often rely on extra input text to generate contextualized embeddings, inevitably increasing computational costs. In this work, we propose Causal2Vec, a general-purpose embedding model tailored to enhance the performance of decoder-only LLMs without altering their original architectures or introducing significant computational overhead. Specifically, we first employ a lightweight BERT-style model to pre-encode the input text into a single Contextual token, which is then prepended to the LLM's input sequence, allowing each token to capture contextualized information even without attending to future tokens. Furthermore, to mitigate the recency bias introduced by last-token pooling, we concatenate the last hidden states of Contextual and EOS tokens as the final text embedding. In practice, Causal2Vec achieves a new state-of-the-art performance on the MTEB benchmark among models trained solely on publicly available retrieval datasets.