为RAG选择正确的嵌入模型

在这篇文章中,我想分享我现在知道的关于Embedding模型选择的知识——涵盖MTEB排行榜、SBERT模型,以及我如何权衡开源和专有选项之间的取舍。

为RAG选择正确的嵌入模型
微信 ezpoda免费咨询:AI编程 | AI模型微调| AI私有化部署
AI模型价格对比 | AI工具导航 | ONNX模型库 | Tripo 3D | Meshy AI | ElevenLabs | KlingAI | ArtSpace | Phot.AI | InVideo

去年,我花了三周时间调试一个RAG聊天机器人,它不断返回看起来非常自信但完全错误的答案。检索指标在纸面上看起来很好。LLM是GPT-4。向量数据库也很扎实。那问题出在哪里?

是Embedding模型。我从教程中随手抓了一个,插进去后就再也没质疑过它。结果证明它与领域完全不匹配。

那段经历迫使我真正理解我放在RAG流水线前端的是什么。在这篇文章中,我想分享我现在知道的关于Embedding模型选择的知识——涵盖MTEB排行榜、SBERT模型,以及我如何权衡开源和专有选项之间的取舍。

1、Embedding模型如何融入RAG流水线

在进行模型比较之前,让我确保思维模型是清晰的,因为我看到过很多困惑。

RAG分三个阶段工作:

阶段1——摄取:你将每个文档块通过Embedding模型处理,生成一个向量Embedding——该块语义内容的固定大小数值表示。这些向量存储在像Milvus这样的向量数据库中。

阶段2——检索:当用户提问时,你使用完全相同的模型对问题进行Embedding,然后运行最近邻搜索以找到语义最相似的块。

阶段3——生成:你将前K个检索到的块作为上下文注入LLM提示中,然后模型根据你的领域知识而不是其训练数据来回答问题。

这里的关键约束:你的查询Embedding和你的文档Embedding只有在由同一模型生成时才存在于同一向量空间中。在不重新索引另一个的情况下交换一个,你的检索质量会完全崩溃。

顺便说一句,有一个研究支持的原因让你保持top-K检索的聚焦。Lost in the Middle论文表明,当太多检索到的块被塞入上下文窗口时,LLM答案质量会下降。保持检索紧密和相关,而不是广泛和嘈杂。

2、SBERT:大多数Embedding模型背后的架构

你遇到的大多数Embedding模型都建立在SBERT(Sentence-BERT)之上。值得理解它与普通BERT有什么不同。

BERT设计用于理解上下文中的单个token。SBERT通过训练模型产生表示整个句子含义的单一固定大小向量来扩展这一点。它使用连体网络训练:两个句子分别编码,模型训练将语义相似的句子放在向量空间中的彼此附近。

实际结果:SBERT理解"the cat sat on the mat"和"the mat sat on the cat"含义不同。普通BERT不会可靠地捕捉到这一点。对于检索任务,你将问题与语义相关的段落匹配,这种句子级理解正是SBERT工作的原因。

像GPT-4这样的LLM建立在Transformer架构的解码器端——优化用于生成。Embedding模型建立在编码器端——优化用于表示。它们在同一流水线中是为不同目的服务的基本不同工具。

3、如何真正使用MTEB排行榜

HuggingFace MTEB排行榜是比较Embedding模型的行业标准,但大多数人都用错了。

MTEB在58个数据集上评估8个任务——检索、聚类、分类、语义文本相似性等。当你构建RAG流水线时,你关心一列:检索平均值,以NDCG@10(第十个结果的归一化折损累积增益)衡量。这个指标对排名较高的结果赋予更多权重,这与RAG实际工作的方式一致——前几个检索到的块承载大部分权重。

我选择模型时的工作流程:

  1. 按检索列降序排列MTEB排行榜
  2. 筛选符合我的内存和延迟预算的模型
  3. 选择达到可接受检索分数的最小模型
  4. 在真实领域查询的样本上运行我自己的评估——因为MTEB对某些模型有过拟合问题

最后一步是不可商榷的。我被那些在MTEB上得分很高但在技术或领域特定文本上表现差的模型烧过。排行榜是起点,不是最终答案。

4、我在生产中使用过的六个模型

让我来聊聊我实际用过的模型,附带关于每个模型优点和缺点的真实笔记。

创建者 模型 Embedding维度 上下文长度 开源 MTEB检索分数
BAAI bge-base-en-v1.5 768 512 tokens 53
BAAI bge-base-zh-v1.5 768 512 tokens 69
VoyageAI voyage-2 1024 4K tokens
VoyageAI voyage-code-2 1536 16K tokens
OpenAI text-embedding-3-small 512–1536 8K tokens 62
OpenAI text-embedding-3-large 256–3072 8K tokens 65

4.1 BAAI/bge-base-en-v1.5和bge-base-zh-v1.5

这些是我原型制作或需要将基础设施成本保持接近零时的首选模型。它们在HuggingFace上可用,在CPU上运行良好,没有API调用成本。

512token的上下文窗口是真正的约束。如果你的文档块自然低于这个限制——大多数段落级分块策略都是这样——你不会注意到它。但如果你处理的是难以干净分割的长技术段落,你会遇到截断问题,悄悄降低检索质量。

对于双语部署,bge-base-zh-v1.5是我为中文内容找到的最实用选项。在中文基准上的MTEB分数69确实很强。

4.2 VoyageAI的voyage-2和voyage-code-2

我在看到一篇案例研究后开始使用VoyageAI的模型,该研究显示在技术文档检索上比ada-002代有显著更好的NDCG@10。

voyage-2在对话和对话数据上训练,这意味着对于某些领域,它在问题到段落的匹配上比通用模型处理得更好。在我的客户支持RAG系统经验中,它在短意向重查询上明显优于bge。

voyage-code-2是有趣的地方。它专门在代码数据上训练,有16K上下文窗口——这组中最长的。对于代码搜索或文档RAG用例,那个上下文窗口意味着你可以嵌入整个函数或长文档字符串,而不会在尴尬的边界分块。VoyageAI报告代码检索任务召回率提高14%,在我自己的测试中那个数字感觉可信。

缺点:这些是专有的、仅API的模型。没有自托管,你依赖它们的可用性和定价。

4.3 OpenAI text-embedding-3-small和text-embedding-3-large

这些取代了ada-002,改进是有意义的——更高的MTEB检索分数、更好的多语言性能和更低的价格。

OpenAI在这里做的最有趣的工程决策是Matryoshka表示学习。不是在单一固定维度上训练,模型同时学习多尺度的表示。结果:你可以在查询时将Embedding截断到更小的维度,精度损失令人惊讶地很小。

对于text-embedding-3-large,从3072维到256维将MTEB检索分数从65降到62——存储和内存减少12倍,精度下降5%。对于存储数千万向量的高吞吐量应用,这种权衡通常是值得的。

在我自己对Milvus技术文档的RAG测试中,dim=256的text-embedding-3-small产生的答案与dim=1536在大多数查询是无法区分的。更高维度重要的边缘情况是细微的歧义问题——那种可能代表实际用户流量的5%。

5、代码:使用OpenAI Embedding和Milvus

这是我一直用的完整流水线。它连接到Zilliz Cloud(托管Milvus),但同样的代码适用于自托管Milvus实例。

步骤1——连接:

from pymilvus import connections, utility
import os
from dotenv import load_dotenv

load_dotenv()
TOKEN = os.getenv("ZILLIZ_API_KEY")
CLUSTER_ENDPOINT = "https://in03-xxxx.api.gcp-us-west1.zillizcloud.com:443"
connections.connect(
    alias='default',
    uri=CLUSTER_ENDPOINT,
    token=TOKEN,
)

步骤2——定义你的Embedding模型:

import openai
from openai import OpenAI

openai_client = OpenAI()
EMBEDDING_MODEL = "text-embedding-3-small"
EMBEDDING_DIM = 512  # 使用减少的维度——对大多数用例来说足够好

步骤3——创建集合:

from pymilvus import MilvusClient

COLLECTION_NAME = "my_rag_collection"
mc = MilvusClient(uri=CLUSTER_ENDPOINT, token=TOKEN)
mc.create_collection(
    COLLECTION_NAME,
    EMBEDDING_DIM,
    consistency_level="Eventually",
    auto_id=True,
    overwrite=True
)

步骤4——分块、嵌入和插入:

# 假设`chunks`是LangChain Document对象列表
chunk_list = []
for chunk in chunks:
    response = openai_client.embeddings.create(
        input=chunk.page_content,
        model=EMBEDDING_MODEL,
        dimensions=EMBEDDING_DIM
    )
    embeddings = response.data[0].embedding

    chunk_list.append({
        'vector': embeddings,
        'chunk': chunk.page_content,
        'source': chunk.metadata.get('source', ''),
    })
mc.insert(COLLECTION_NAME, data=chunk_list, progress_bar=True)
mc.flush(COLLECTION_NAME)

步骤5——查询和生成:

SAMPLE_QUESTION = "What do the parameters for HNSW mean?"

response = openai_client.embeddings.create(
    input=SAMPLE_QUESTION,
    model=EMBEDDING_MODEL,
    dimensions=EMBEDDING_DIM
)
query_embeddings = response.data[0].embedding
results = mc.search(
    COLLECTION_NAME,
    data=[query_embeddings],
    output_fields=["chunk", "source"],
    limit=3,
    consistency_level="Eventually"
)
context = [r['entity']['chunk'] for r in results[0]]
contexts_combined = ' '.join(context)
llm_response = openai_client.chat.completions.create(
    messages=[
        {"role": "system", "content": f"Answer using only the context below. Context: {contexts_combined}"},
        {"role": "user", "content": SAMPLE_QUESTION}
    ],
    model="gpt-3.5-turbo",
    temperature=0.1,
)
print(llm_response.choices[0].message.content)

完整的工作笔记本在GitHub的Milvus bootcamp上。

6、我的决策框架

在进行了很多这些实验后,我现在是这样决定的:

  • 原型制作/成本敏感/自托管:从bge-base-en-v1.5开始。它是免费的、快速的,足以验证你的流水线架构。
  • 生产英语/多语言聊天机器人:text-embedding-3-small在降低维度后在成本-质量曲线上难以击败。
  • 高风险领域RAG(法律、医学、技术文档):认真评估voyage-2。MTEB分数不能捕捉一切;领域特定的检索质量可能会让你惊讶。
  • 代码搜索/开发者工具:voyage-code-2及其16K上下文窗口值得API依赖。
  • 极端内存限制或数十亿向量:dim=256的text-embedding-3-large——Matryoshka方法确实是聪明的工程。

我会不断重复的一件事:总是在真实生产查询的切片上运行你自己的评估。MTEB是一个代理。你实际数据上真正的检索质量才是唯一重要的数字。


原文链接:Picking the Right Embedding Model for Your RAG Pipeline

汇智网翻译整理,转载请标明出处