RAG系统优化实战:从召回到重排序

这篇文章记录了我在优化公司知识库问答系统的三个月里,从召回率40%提升到78%的全过程。包含了大量踩坑经验、对比实验数据和代码实现细节。

背景:一个"能用但不好用"的RAG系统

去年9月,公司上线了一个基于RAG的内部知识库问答系统。技术栈很常规:OpenAI的text-embedding-ada-002做向量化,Pinecone做向量数据库,GPT-3.5-turbo做生成。刚上线时大家觉得挺新鲜,但用了一个月后,抱怨开始多起来:

• "明明文档里有这个信息,为什么回答说不知道?"
• "检索出来的内容完全不相关,纯纯是在瞎编"
• "产品名称搜不到,得换个说法才行"

我接手后测试了100个用户真实问题,发现召回率只有可怜的40%。也就是说,10个问题里有6个压根检索不到相关文档。这种情况下,LLM再强也无济于事——巧妇难为无米之炊。

第一步:Embedding模型选型

为什么要换模型?

text-embedding-ada-002是个好模型,但有两个问题:

1. 成本高:我们的知识库有3万份文档,平均每份5000字,embedding一次就要花$120,而且每次更新都得重新embedding
2. 中文支持一般:我们90%的文档是中文,ada-002在中文语义理解上不如专门训练的中文模型

对比实验

我在MTEB-Chinese基准上测试了几个候选模型,并在我们的200条标注问答对上做了评测:

BGE-large-zh-v1.5在我们的测试集上表现最好,召回率提升了6.5个百分点。更重要的是,它可以本地部署,推理速度也完全可以接受。

部署实现

from FlagEmbedding import FlagModel
import torch

# 加载模型
model = FlagModel(
    'BAAI/bge-large-zh-v1.5',
    query_instruction_for_retrieval="为这个句子生成表示以用于检索相关文章:",
    use_fp16=True  # 半精度推理,速度快一倍
)

def embed_texts(texts, batch_size=32):
    """批量embedding,避免OOM"""
    embeddings = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i+batch_size]
        batch_embeddings = model.encode(batch)
        embeddings.extend(batch_embeddings)
    return np.array(embeddings)

# Query和Document的embedding方式不同
def embed_query(query):
    # query会自动加上instruction
    return model.encode_queries([query])[0]

def embed_docs(docs):
    # document不加instruction
    return model.encode(docs)

第二步:Chunking策略优化

问题诊断

换了embedding模型后,召回率提升到了52%,但还是不够。我分析了几十个失败case,发现了一个严重问题:我们的文档平均5000字,直接用512 tokens分块,很多时候把一个完整的知识点切成了两半。

比如这样一段文档:

产品X的部署流程分为三步:
1. 环境准备
   - 安装Docker 20.10+
   - 配置网络...
[被切断]
2. 数据库初始化
   - 创建数据库...

用户问"产品X怎么部署",检索到的chunk可能只有"2. 数据库初始化"这部分,缺少上下文,LLM根本不知道在说什么。

改进方案:语义感知分块

import re
from typing import List, Tuple

def semantic_chunking(
    text: str,
    max_chunk_size: int = 800,
    min_chunk_size: int = 200,
    overlap: int = 100
) -> List[Tuple[str, dict]]:
    """
    基于文档结构的语义分块
    返回: [(chunk_text, metadata), ...]
    """
    chunks = []
    
    # 1. 按标题分段
    # 匹配 ## 标题 或 **标题** 或 1. 2. 3.
    sections = re.split(
        r'(\n#{1,3}\s+.+\n|\n\*\*.+\*\*\n|\n\d+\.\s+.+\n)',
        text
    )
    
    current_chunk = ""
    current_title = ""
    
    for i, section in enumerate(sections):
        # 判断是否是标题
        is_title = bool(re.match(
            r'\n#{1,3}\s+.+\n|\n\*\*.+\*\*\n|\n\d+\.\s+.+\n',
            section
        ))
        
        if is_title:
            current_title = section.strip()
            
        # 如果当前chunk已经够大,或遇到新标题,就输出
        if len(current_chunk) >= min_chunk_size and is_title:
            if current_chunk:
                chunks.append((
                    current_chunk.strip(),
                    {"title": current_title, "position": len(chunks)}
                ))
            current_chunk = section
        else:
            current_chunk += section
            
        # 强制分块(超过max_size)
        if len(current_chunk) >= max_chunk_size:
            chunks.append((
                current_chunk.strip(),
                {"title": current_title, "position": len(chunks)}
            ))
            # 保留overlap的内容作为下一个chunk的开头
            overlap_text = current_chunk[-overlap:] if overlap > 0 else ""
            current_chunk = overlap_text
    
    # 处理最后一块
    if current_chunk.strip():
        chunks.append((
            current_chunk.strip(),
            {"title": current_title, "position": len(chunks)}
        ))
    
    return chunks

# 实际使用
doc = load_document("产品手册.md")
chunks = semantic_chunking(doc, max_chunk_size=800, overlap=100)

# 带metadata存储,方便后续过滤和调试
for chunk_text, metadata in chunks:
    chunk_id = store_to_vector_db(
        text=chunk_text,
        embedding=embed_docs([chunk_text])[0],
        metadata={
            **metadata,
            "doc_id": doc.id,
            "doc_title": doc.title
        }
    )

改用语义分块后,召回率又提升了8个点,达到60%。而且检索结果的可读性明显更好,LLM生成的答案质量也跟着提升。

第三步:混合检索

纯向量检索的局限

即使有了更好的embedding和chunking,纯向量检索在某些场景下还是不行。典型的问题是专有名词。

举个例子,用户问"AgentX的配置在哪",向量检索可能返回一堆关于"Agent配置"的文档,但就是找不到AgentX。因为"AgentX"这个词在embedding空间里的表示和"Agent"太接近了。

BM25救场

BM25是经典的TF-IDF改进算法,对精确关键词匹配特别好。我们用BM25补充向量检索的不足:

from rank_bm25 import BM25Okapi
import jieba

class HybridRetriever:
    def __init__(self, vector_db, documents):
        self.vector_db = vector_db
        self.documents = documents
        
        # 构建BM25索引
        tokenized_docs = [list(jieba.cut(doc)) for doc in documents]
        self.bm25 = BM25Okapi(tokenized_docs)
        
    def search(self, query, top_k=10, alpha=0.5):
        """
        alpha: 向量检索的权重 (1-alpha为BM25权重)
        """
        # 1. 向量检索
        vector_results = self.vector_db.search(
            query_vector=embed_query(query),
            top_k=top_k * 2
        )
        
        # 2. BM25检索
        tokenized_query = list(jieba.cut(query))
        bm25_scores = self.bm25.get_scores(tokenized_query)
        bm25_results = sorted(
            enumerate(bm25_scores),
            key=lambda x: x[1],
            reverse=True
        )[:top_k * 2]
        
        # 3. RRF融合 (Reciprocal Rank Fusion)
        # RRF比加权求和更稳定,不受分数scale影响
        scores = {}
        k = 60  # RRF参数
        
        # 向量检索的贡献
        for rank, (doc_id, _) in enumerate(vector_results):
            scores[doc_id] = scores.get(doc_id, 0) + alpha / (k + rank + 1)
        
        # BM25的贡献
        for rank, (doc_id, _) in enumerate(bm25_results):
            scores[doc_id] = scores.get(doc_id, 0) + (1-alpha) / (k + rank + 1)
        
        # 按融合分数排序
        final_results = sorted(
            scores.items(),
            key=lambda x: x[1],
            reverse=True
        )[:top_k]
        
        return [self.documents[doc_id] for doc_id, _ in final_results]

参数调优

我在验证集上测试了不同的alpha值:

• alpha=1.0 (纯向量): 召回率 60.2%
• alpha=0.7: 召回率 68.5%
• alpha=0.5: 召回率 70.1%
• alpha=0.3: 召回率 66.8%
• alpha=0.0 (纯BM25): 召回率 52.3%

最终选择alpha=0.5,向量和BM25各占一半。这样既能捕捉语义相似,又能精确匹配关键词。

第四步:重排序(Reranking)

为什么需要重排序?

混合检索已经把召回率提升到了70%,但还有个问题:检索出来的top-10里,真正相关的可能只有前3个,后面7个都是噪音。直接把这10个都扔给LLM,会导致:

1. Context太长,浪费token
2. 噪音信息干扰LLM判断
3. 生成速度变慢

Cross-Encoder精排

Bi-encoder(我们的embedding模型)把query和document分别编码,然后算相似度。而Cross-encoder是把query和document拼在一起,直接判断相关性,精度更高。

from sentence_transformers import CrossEncoder

class Reranker:
    def __init__(self, model_name='BAAI/bge-reranker-large'):
        self.model = CrossEncoder(model_name)
        
    def rerank(self, query, documents, top_k=5):
        """
        对检索结果重排序
        """
        if len(documents) == 0:
            return []
        
        # 构造query-document pairs
        pairs = [[query, doc.text] for doc in documents]
        
        # 批量打分
        scores = self.model.predict(pairs)
        
        # 按分数排序
        doc_scores = list(zip(documents, scores))
        doc_scores.sort(key=lambda x: x[1], reverse=True)
        
        # 只返回top-k个高分文档
        return [doc for doc, score in doc_scores[:top_k] if score > 0.3]

# 完整的检索pipeline
def retrieve_and_rerank(query, top_k_retrieve=20, top_k_final=5):
    # 1. 混合检索,多召回一些候选
    candidates = hybrid_retriever.search(query, top_k=top_k_retrieve)
    
    # 2. 重排序,精选最相关的
    final_docs = reranker.rerank(query, candidates, top_k=top_k_final)
    
    return final_docs

效果对比

加入重排序后,在标注的100个问题上:

虽然延迟增加了200ms,但召回率提升了8个点,而且LLM生成质量明显更好。用户满意度从之前的60分提升到了82分。

性能优化:实际生产中的坑

1. 批量embedding加速

# ❌ 慢:逐个embedding
for doc in documents:
    emb = model.encode([doc])
    store(emb)

# ✅ 快:批量embedding
batch_size = 64
for i in range(0, len(documents), batch_size):
    batch = documents[i:i+batch_size]
    embeddings = model.encode(batch)
    store_batch(embeddings)

批量处理后,embedding 3万份文档的时间从4小时降到了25分钟。

2. 异步BM25索引更新

每次文档更新都重建BM25索引太慢了。我们改成增量更新:

import redis
import pickle

class IncrementalBM25:
    def __init__(self, redis_client):
        self.redis = redis_client
        self.load_index()
        
    def load_index(self):
        data = self.redis.get('bm25_index')
        if data:
            self.bm25 = pickle.loads(data)
        else:
            self.bm25 = BM25Okapi([])
    
    def add_documents(self, new_docs):
        """增量添加文档"""
        tokenized = [list(jieba.cut(doc)) for doc in new_docs]
        
        # 更新BM25内部状态
        self.bm25.corpus_size += len(tokenized)
        self.bm25.doc_len.extend([len(doc) for doc in tokenized])
        # ... 更新其他统计信息
        
        # 持久化
        self.redis.set('bm25_index', pickle.dumps(self.bm25))
    
    def rebuild(self, all_docs):
        """定期全量重建(每周一次)"""
        tokenized = [list(jieba.cut(doc)) for doc in all_docs]
        self.bm25 = BM25Okapi(tokenized)
        self.redis.set('bm25_index', pickle.dumps(self.bm25))

3. Reranker推理优化

重排序延迟主要来自模型推理。我们做了几个优化:

• 模型量化:FP16推理,速度提升1.8x,精度几乎无损
• 批量推理:一次处理20个pairs,而不是逐个处理
• 缓存热门query:Redis缓存最近1小时的query结果

优化后,P99延迟从650ms降到了420ms。

教训与反思

1. 一定要建评测集

最开始我没有评测集,全靠人工看case。后来花了一周时间,让标注团队整理了500个真实用户问题,标注了正确答案和相关文档。有了这个评测集,每次改动都能量化效果,迭代速度快了很多。

2. 不要盲目追求最新技术

我曾经尝试过HyDE、Query Expansion、Multi-Query等"高级"技术,结果效果提升有限,复杂度却增加了不少。最后发现,把基础做扎实(好的embedding、合理的chunking、混合检索)比什么都重要。

3. 监控很重要

上线后,我们在Grafana上加了这些监控:

• 召回率@5、@10(基于用户点击反馈估算)
• 各阶段延迟(embedding/retrieval/rerank/generation)
• 失败case采样(自动记录rerank分数低的query)

这些监控帮我们及时发现了几次索引损坏、模型OOM等问题。

最终效果与成本

虽然延迟增加了0.6秒,但用户明显感觉答案质量更好了。成本方面,embedding省了一大笔,但多了重排序的GPU开销,整体反而更低。

下一步优化方向

1. Fine-tune Embedding模型:在我们的领域数据上微调BGE,可能还能提升5-10个点
2. 多路召回融合:除了向量+BM25,还可以加入关键词抽取、实体识别等其他召回路径
3. 个性化检索:根据用户历史行为调整检索权重
4. 主动学习:让用户标注bad case,持续优化检索效果

总结

RAG系统的优化是个系统工程,没有银弹。我的经验是:

1. 先建立评测体系,才能量化迭代
2. Embedding选型要针对自己的场景测试
3. Chunking策略比你想象的更重要
4. 混合检索是必选项,不是可选项
5. Reranking效果显著,但要控制延迟
6. 监控和Bad Case分析要持续做

希望这篇文章对正在做RAG的同学有帮助。如果有问题欢迎留言讨论!

更新记录:
2024-11-08: 初版发布
2024-11-15: 补充了性能优化部分
2024-11-22: 更新了最终效果数据