跳转到内容

多阶段 RAG:Bi-encoder + Cross-encoder 双阶段检索

多阶段 RAG:Bi-encoder + Cross-encoder 双阶段检索

学习目标

理解 Khoj 如何通过 bi-encoder 初筛 + cross-encoder 精排的双阶段检索架构实现高精度语义搜索,以及本地/远程 embedding 回退策略。

前置知识

本章涉及 RAG 流水线和分块策略的通用原理,建议先阅读:

下文假设你已理解上述概念,直接聚焦 Khoj 的双阶段检索具体实现。

项目实践

检索流水线概览

第一阶段:Bi-encoder 粗筛

Bi-encoder 将查询和文档片段编码为固定维度的向量,通过余弦相似度(归一化向量上的点积)快速筛选 Top-K 候选。

模型选择:默认使用 thenlper/gte-small,这是一个 33M 参数的小型模型,在速度与质量之间取得平衡。

三级 embedding 后端

  1. 本地SentenceTransformer 库直接加载模型
  2. HuggingFace Inference API:远程推理,适合没有 GPU 的部署
  3. OpenAI Embeddings API:最准确但有成本和延迟

这种设计允许自部署用户根据自身硬件条件选择:有 GPU 用本地模型,无 GPU 用远程 API。

磁盘缓存:嵌入结果保存为 .pt(PyTorch)文件,避免重复计算。文档变更通过 hashed_value 检测,仅重新计算变更部分的嵌入。

第二阶段:Cross-encoder 精排

Bi-encoder 返回的 Top-K 结果(默认 10 条)送入 cross-encoder 进行精排。

为什么需要重排? Bi-encoder 将查询和文档独立编码,丢失了交互信息。Cross-encoder 同时编码查询和文档,计算它们的交互分数,精度显著更高但计算成本也更高。

Khoj 使用 mixedbread-ai/mxbai-rerank-xsmall-v1,一个轻量级的 cross-encoder 模型,在精度和延迟之间取得平衡。

查询过滤器语法

Khoj 支持从查询字符串中直接解析过滤条件,类似 Lucene 语法:

过滤器示例作用
日期过滤+file:"2024-01"限定文档创建时间范围
文件过滤+file:notes.org限定文件名
关键词过滤+word:重要必须包含特定词汇

过滤器在 bi-encoder 编码前应用,先缩小候选范围再进行语义匹配,显著提升检索效率。

分块策略

Khoj 使用 LangChain 的 RecursiveCharacterTextSplitter,max tokens = 256,separator 优先级:

\n\n → \n → ! → ? → . → 空格 → \t → 空字符串

这个策略优先在段落边界分割,保留语义完整性。256 token 的限制确保了嵌入质量(短片段更精确),但可能丢失长文档的上下文——这是 Khoj 在精度与召回之间的取舍。

向量存储:pgvector

Entry 模型:
- corpus_id: UUID(文档级去重)
- hashed_value: SHA 哈希(变更检测)
- embedding: VectorField(pgvector)
- raw: 原始文本
- compiled: 分块后的文本
- file_type: 文件类型
- file_path: 文件路径

Khoj 使用 pgvector 的 VectorField,在归一化向量上使用点积操作符(等价于余弦相似度)。对于自建部署,pgvector 是最简单的向量存储方案——不需要额外的向量数据库,PostgreSQL 同时承担关系数据和向量检索两种角色。

问题与规避

问题影响规避策略
本地 embedding 模型内存占用大小内存服务器无法部署回退到 HuggingFace/ OpenAI 远程 API
256 token 分块过小长上下文被截断,丢失信息检索结果返回时可考虑合并相邻片段
Cross-encoder 延迟精排阶段增加 100-500ms使用 xsmall 变体;仅对 bi-encoder 的 Top-K 结果精排
pgvector 无 HNSW 索引大数据集检索变慢手动添加 IVFFlat 或 HNSW 索引;或限制文档总量

设计取舍

为什么选 pgvector 而非专用向量数据库(如 Milvus、Qdrant)?

Khoj 的目标用户是自部署的个人用户,增加一个专用向量数据库会显著提高部署复杂度。pgvector 作为 PostgreSQL 扩展,可以在同一个数据库中完成所有存储和检索,运维成本最低。代价是规模上限——当文档数量达到百万级时,pgvector 的性能不如专用向量库。但对于个人知识库(通常数万到数十万文档),pgvector 完全够用。

为什么 bi-encoder + cross-encoder 而不是只用其中一个?

只用 bi-encoder:速度快但精度有限,尤其对短查询。 只用 cross-encoder:精度高但需要对所有文档计算,不可行(O(n) 复杂度)。 双阶段:bi-encoder 将搜索空间从 O(n) 缩小到 O(K),cross-encoder 在 K 条结果上精排——兼顾速度和精度。

参考来源