引言

RAG系统的检索质量很大程度上取决于嵌入模型的表示能力。通用嵌入模型(如text-embedding-ada-002、BGE)在开放域表现优秀,但在特定领域(医疗、法律、金融)往往力不从心。领域术语的特殊含义、专业表达方式的差异,都可能导致通用嵌入模型的语义匹配不准确。嵌入模型微调是解决这一问题的有效手段。本文分享嵌入模型微调的完整流程。

为什么需要微调嵌入模型

领域适配问题

以医疗领域为例,“阳性"在通用语境中是褒义的,但在医学语境中仅表示检测结果。通用嵌入模型可能将"阳性结果"与"好消息"的向量拉近,而医疗场景中"阳性结果"可能意味着确诊了某种疾病。

检索质量影响

嵌入模型不匹配会导致:相关文档检索不到(召回率低)、不相关文档排在前面(精确率低)、同义词无法匹配、专业缩写无法理解。

数据准备

训练数据格式

嵌入模型微调需要三元组数据:(query, positive, negative):

  • query:查询文本
  • positive:与query相关的正样本
  • negative:与query不相关的负样本

数据构建方法

方法一:利用已有标注数据

如果已有问答对或查询-文档配对数据,直接用作正样本对:

# 从FAQ库构建训练数据
def build_from_faq(faq_data):
    training_data = []
    for item in faq_data:
        query = item['question']
        positive = item['answer']
        # 随机采样负样本
        negatives = random.sample([other['answer'] for other in faq_data if other != item], 5)
        training_data.append({
            'query': query,
            'positive': positive,
            'negatives': negatives
        })
    return training_data

方法二:LLM辅助生成

使用LLM根据文档内容生成查询:

def generate_training_data(documents, llm):
    training_data = []
    for doc in documents:
        # 生成多种查询方式
        prompt = f"""
        根据以下文档内容,生成5个不同角度的查询问题。
        问题应该能通过这段文档内容回答。
        
        文档:{doc}
        
        生成的问题:
        """
        queries = llm.generate(prompt).split('\n')
        
        for query in queries:
            if query.strip():
                training_data.append({
                    'query': query.strip(),
                    'positive': doc,
                    'negatives': get_hard_negatives(query, doc, all_docs)
                })
    return training_data

方法三:挖掘难负样本

难负样本(Hard Negatives)对微调效果至关重要——它们是与query相关但不应该匹配的文档:

def get_hard_negatives(query, positive_doc, all_docs, top_k=5):
    """使用现有嵌入模型挖掘难负样本"""
    # 用现有模型检索
    results = vector_store.search(query, top_k=top_k * 3)
    
    hard_negatives = []
    for result in results:
        if result['doc'] != positive_doc and result['score'] > 0.5:
            hard_negatives.append(result['doc'])
        if len(hard_negatives) >= top_k:
            break
    
    return hard_negatives

数据质量检查

训练前进行数据质量检查:

  • 正样本一致性:确认query确实能从positive文档中找到答案
  • 负样本正确性:确认negative文档确实不包含答案
  • 多样性:确保训练数据覆盖不同查询类型和文档长度
  • 平衡性:每个文档作为正样本的次数大致均匀

微调方法

方法一:对比学习(Contrastive Learning)

使用InfoNCE损失函数,拉近正样本对、推远负样本对:

import torch
import torch.nn as nn
from sentence_transformers import SentenceTransformer

class ContrastiveTrainer:
    def __init__(self, model_name='BAAI/bge-large-zh-v1.5'):
        self.model = SentenceTransformer(model_name)
    
    def train(self, training_data, epochs=3, batch_size=32, lr=2e-5):
        optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr)
        
        for epoch in range(epochs):
            total_loss = 0
            for batch in self.get_batches(training_data, batch_size):
                queries = [item['query'] for item in batch]
                positives = [item['positive'] for item in batch]
                
                # 编码
                q_embeddings = self.model.encode(queries, convert_to_tensor=True)
                p_embeddings = self.model.encode(positives, convert_to_tensor=True)
                
                # InfoNCE损失
                scores = torch.matmul(q_embeddings, p_embeddings.T)
                labels = torch.arange(len(batch)).to(scores.device)
                loss = nn.CrossEntropyLoss()(scores / 0.05, labels)
                
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                
                total_loss += loss.item()
            
            print(f"Epoch {epoch+1}: loss = {total_loss / len(training_data):.4f}")
    
    def get_batches(self, data, batch_size):
        for i in range(0, len(data), batch_size):
            yield data[i:i+batch_size]

方法二:困难负样本挖掘迭代

交替进行训练和困难负样本挖掘:

def iterative_train(model, training_data, all_docs, iterations=3):
    for iteration in range(iterations):
        # 1. 用当前模型挖掘更难的负样本
        for item in training_data:
            item['negatives'] = get_hard_negatives(
                item['query'], item['positive'], all_docs,
                model=model, top_k=5
            )
        
        # 2. 用新数据训练
        trainer = ContrastiveTrainer(model)
        trainer.train(training_data, epochs=1)
        
        # 3. 评估
        metrics = evaluate(model, test_data)
        print(f"Iteration {iteration+1}: {metrics}")

方法三:LoRA微调

使用LoRA进行参数高效微调,降低计算成本:

from peft import LoraConfig, get_peft_model

def lora_finetune(model_name='BAAI/bge-large-zh-v1.5'):
    model = SentenceTransformer(model_name)
    
    # 对Transformer层添加LoRA
    lora_config = LoraConfig(
        r=8,
        lora_alpha=32,
        target_modules=["query", "key", "value"],
        lora_dropout=0.1
    )
    
    # 应用LoRA
    for module in model.modules():
        if hasattr(module, 'query'):
            get_peft_model(module, lora_config)
    
    return model

评估方法

检索质量评估

构建领域测试集,评估微调前后的检索质量:

def evaluate_retrieval(model, test_queries, ground_truth):
    """评估检索质量"""
    results = {'recall@5': 0, 'recall@10': 0, 'mrr': 0}
    
    for query, relevant_docs in zip(test_queries, ground_truth):
        # 检索
        query_emb = model.encode([query])
        scores = cosine_similarity(query_emb, doc_embeddings)
        ranked_indices = np.argsort(scores[0])[::-1]
        
        # 计算指标
        top5 = ranked_indices[:5]
        top10 = ranked_indices[:10]
        
        relevant_set = set(relevant_docs)
        results['recall@5'] += len(set(top5) & relevant_set) / len(relevant_set)
        results['recall@10'] += len(set(top10) & relevant_set) / len(relevant_set)
        
        # MRR
        for rank, idx in enumerate(ranked_indices, 1):
            if idx in relevant_set:
                results['mrr'] += 1 / rank
                break
    
    n = len(test_queries)
    return {k: v / n for k, v in results.items()}

对比基线

评估指标微调前微调后提升
Recall@562.3%81.7%+19.4%
Recall@1075.1%89.2%+14.1%
MRR0.480.71+0.23

实践建议

数据量要求

  • 最少需要500-1000个训练三元组才能看到明显效果
  • 2000-5000个高质量样本通常能达到良好效果
  • 超过10000个样本的边际收益递减

避免过拟合

  • 使用验证集监控训练,在验证指标下降时停止
  • 设置较小的学习率(1e-5到5e-5)
  • 训练轮数通常2-5轮即可
  • 使用Dropout和权重衰减

保持通用能力

微调可能导致模型在通用任务上表现下降(灾难性遗忘)。缓解策略:

  • 在训练数据中混入部分通用数据
  • 使用较小的学习率和较少的训练轮数
  • 使用LoRA而非全参数微调

结语

嵌入模型微调是提升领域RAG检索质量的有效手段。关键成功因素是高质量训练数据(特别是难负样本)和合理的训练策略。通过迭代式的困难负样本挖掘和模型评估,可以在保持通用能力的同时显著提升领域检索效果。建议在通用嵌入模型效果不达预期时,将微调作为首选优化方向。

加入讨论

这篇文章有姊妹讨论帖在硅基AGI论坛 — 全球首个碳基硅基认知交流平台。

碳基与硅基的智慧碰撞,认知差异创造无限可能。