Causal2Vec: Improving Decoder-only LLMs as Versatile Embedding Models

📄 arXiv: 2507.23386v2 📥 PDF

作者: Ailiang Lin, Zhuoyun Li, Kotaro Funakoshi, Manabu Okumura

分类: cs.CL, cs.AI

发布日期: 2025-07-31 (更新: 2025-09-19)


💡 一句话要点

Causal2Vec:提升Decoder-only LLM作为通用嵌入模型的性能

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 文本嵌入 Decoder-only LLM 因果注意力 上下文表示 轻量级模型 MTEB 语义相似度

📋 核心要点

  1. 现有方法通过移除因果注意力掩码或增加额外输入,试图提升Decoder-only LLM的嵌入能力,但存在破坏语义信息或增加计算成本的问题。
  2. Causal2Vec通过预编码Contextual token并前置于LLM输入,以及连接Contextual和EOS token的隐藏状态,在不改变LLM架构的前提下提升性能。
  3. 实验表明,Causal2Vec在MTEB上取得了SOTA性能,同时显著降低了序列长度和推理时间,提升效率。

📝 摘要(中文)

Decoder-only的大型语言模型(LLM)越来越多地被用作嵌入模型,有效地将自然语言文本的语义信息编码成密集的向量表示,用于各种嵌入任务。然而,许多现有方法主要集中于移除LLM中的因果注意力掩码以实现双向注意力,这可能会削弱模型在预训练期间获得的语义信息提取能力。此外,领先的单向方法通常依赖于额外的输入文本来克服因果注意力的固有局限性,不可避免地增加了计算成本。本文提出了Causal2Vec,一种通用嵌入模型,旨在增强decoder-only LLM的性能,而无需改变其原始架构或引入显著的计算开销。具体来说,我们首先采用一个轻量级的BERT风格模型将输入文本预编码成一个单独的Contextual token,然后将其添加到LLM的输入序列中,使得每个token即使不关注未来的token也能捕获上下文信息。此外,为了减轻last-token pooling引入的近因偏差,并帮助LLM更好地利用Contextual token中编码的语义信息,我们将Contextual和EOS token的最后隐藏状态连接起来作为最终的文本嵌入。在实践中,Causal2Vec在Massive Text Embeddings Benchmark(MTEB)上,在仅使用公开可用的检索数据集训练的模型中,实现了最先进的性能,同时与性能最佳的方法相比,所需的序列长度最多减少了85%,推理时间最多减少了82%。

🔬 方法详解

问题定义:现有方法在利用decoder-only LLM进行文本嵌入时,为了克服单向注意力的限制,要么移除因果注意力掩码,导致模型丧失预训练获得的语义信息,要么引入额外的输入文本,增加计算开销。因此,如何在不改变decoder-only LLM原始架构的前提下,高效地提升其文本嵌入能力是一个关键问题。

核心思路:Causal2Vec的核心思路是利用一个轻量级的BERT风格模型预先提取输入文本的上下文信息,并将其编码为一个Contextual token。然后,将这个Contextual token添加到LLM的输入序列中,使得LLM的每个token都能感知到全局的上下文信息,从而在不使用双向注意力的情况下,提升模型对文本语义的理解能力。同时,通过连接Contextual和EOS token的隐藏状态,缓解last-token pooling带来的近因偏差。

技术框架:Causal2Vec的整体框架包含两个主要模块:1) Contextual Token生成器:使用一个轻量级的BERT风格模型(例如TinyBERT)将输入文本编码为一个Contextual token,该token包含了整个输入文本的上下文信息。2) Decoder-only LLM:将Contextual token前置于原始输入文本序列,输入到decoder-only LLM中进行处理。最后,将Contextual token和EOS token的最后隐藏状态连接起来,作为最终的文本嵌入。

关键创新:Causal2Vec的关键创新在于:1) 使用Contextual token来增强decoder-only LLM对上下文信息的感知能力,避免了直接修改LLM架构或引入额外输入文本。2) 通过连接Contextual和EOS token的隐藏状态,缓解了last-token pooling的近因偏差,使得模型能够更好地利用Contextual token中编码的语义信息。

关键设计:Contextual Token生成器可以使用预训练的TinyBERT模型,并对其进行微调以适应特定的嵌入任务。Contextual token的嵌入维度与decoder-only LLM的嵌入维度保持一致。损失函数可以使用对比学习损失或多任务学习损失,以优化Contextual token的生成和文本嵌入的质量。在推理阶段,只需要进行一次前向传播即可获得文本嵌入,计算效率高。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Causal2Vec在MTEB基准测试中取得了显著的性能提升,在仅使用公开可用的检索数据集训练的模型中,实现了最先进的性能。与最佳方法相比,Causal2Vec将所需的序列长度减少了高达85%,推理时间减少了高达82%。这些结果表明,Causal2Vec在保持高性能的同时,显著提高了计算效率。

🎯 应用场景

Causal2Vec作为一种高效的文本嵌入模型,可以广泛应用于各种自然语言处理任务,例如文本检索、文本分类、语义相似度计算、聚类等。其高效的推理速度和较低的计算成本使其特别适用于资源受限的场景,例如移动设备或边缘计算环境。未来,Causal2Vec可以进一步扩展到多语言场景,并与其他模态的信息进行融合,以实现更强大的跨模态表示学习。

📄 摘要(原文)

Decoder-only large language models (LLMs) are increasingly used to build embedding models that effectively encode the semantic information of natural language texts into dense vector representations for various embedding tasks. However, many existing methods primarily focus on removing the causal attention mask in LLMs to enable bidirectional attention, potentially undermining the model's ability to extract semantic information acquired during pretraining. Additionally, leading unidirectional approaches often rely on extra input text to overcome the inherent limitations of causal attention, 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 and help LLMs better leverage the semantic information encoded in the Contextual token, we concatenate the last hidden states of Contextual and EOS tokens as the final text embedding. In practice, Causal2Vec achieves state-of-the-art performance on the Massive Text Embeddings Benchmark (MTEB) among models trained solely on publicly available retrieval datasets, while reducing the required sequence length by up to 85% and inference time by up to 82% compared to best-performing methods.