用上下文标记提升因果型大模型的嵌入能力,无需改架构也不增算力。
Causal2Vec: Improving Decoder-only LLMs as Embedding Models through a Contextual Token
- 引入轻量编码器生成上下文标记,前置输入序列实现上下文感知。
- 在MTEB上超越现有纯公开数据训练模型,达到新SOTA性能。
- 适合需要高效高质文本嵌入的应用场景,如信息检索与相似度匹配。
解码器仅有的大型语言模型(LLMs)被广泛用于构建各类任务的嵌入模型。为克服因果注意力在表征学习中的固有限制,现有方法常修改注意力机制为双向,可能损害预训练中获得的语义信息。而主流单向方法多依赖额外输入文本生成上下文嵌入,导致计算成本上升。本文提出Causal2Vec,一种通用嵌入模型,可增强解码器仅有的LLMs性能,不改变其原始架构且无显著计算开销。首先,使用轻量级BERT风格模型将输入文本预编码为单一上下文标记,并将其前置至LLM输入序列,使每个标记即使不关注未来也能捕获上下文信息。此外,为缓解最后标记池化带来的近期偏差,采用上下文标记与结束标记(EOS)最后隐藏状态拼接作为最终文本嵌入。实验表明,Causal2Vec在仅使用公开检索数据集训练的模型中,在MTEB基准上取得新SOTA表现。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。