RAG 对比学习优化:用负样本提升检索精度的训练方法

发布时间:2026/7/17 16:04:09
RAG 对比学习优化:用负样本提升检索精度的训练方法 RAG 对比学习优化用负样本提升检索精度的训练方法一、RAG 返回了相关但不准确的结果怎么办标准的 RAG 流程用户提问 → Embedding 检索 → 拼接上下文 → LLM 生成答案。问题往往出在第二步——检索回来的 Top-K 文档看起来相关但内容对问题没有帮助。比如用户问Go 1.24 的 range over func 怎么用检索系统返回了一篇Go 1.18 泛型入门。关键词匹配度高Go、func、range但对用户毫无价值。根因在于 Embedding 模型的训练目标它学的是语义相似度而不是对当前问题是否有用的任务相关性。对比学习Contrastive Learning可以解决这个问题。核心思想让模型学会区分正样本对问题有用的文档和负样本看起来相关但无用的文档。二、对比学习的原理与负样本挖掘策略InfoNCE 损失函数的直观理解Loss -log( exp(sim(q, p)) / (exp(sim(q, p)) sum(exp(sim(q, n_i))) )sim(q, p)Query 和正样本的相似度sim(q, n_i)Query 和第 i 个负样本的相似度优化目标正样本相似度最大化负样本相似度最小化三、Python 实现负样本挖掘与对比训练负样本挖掘器import numpy as np from typing import List, Tuple, Optional from dataclasses import dataclass from collections import defaultdict dataclass class TrainingSample: 训练样本结构 query: str positive: str # 正样本对 query 有用的文档 negative_hard: str # 难负样本相似但无用的文档 negative_random: str # 随机负样本不相关的文档 class NegativeSampler: 负样本挖掘器——关键是挖掘难负样本 def __init__( self, embedding_model, # 已有 Embedding 模型SentenceTransformer corpus: List[str], # 全量文档集合 top_k_hard: int 20, # 初筛 top_k 个候选难负样本 similarity_threshold: float 0.6, # BM25 和 Embedding 相似度差异阈值 ): self.model embedding_model self.corpus corpus self.top_k_hard top_k_hard self.similarity_threshold similarity_threshold # 预计算全量语料的 Embedding离线做 self.corpus_embeddings: Optional[np.ndarray] None def precompute_embeddings(self): 离线预计算所有文档的 Embedding print(f正在编码 {len(self.corpus)} 篇文档...) self.corpus_embeddings self.model.encode( self.corpus, show_progress_barTrue, normalize_embeddingsTrue # 归一化方便余弦相似度计算 ) print(编码完成) def find_hard_negatives(self, query: str, positive_idx: int) - List[str]: 挖掘难负样本和 query 语义相似度高但不是正样本的文档 策略 1. 用 BM25词汇匹配初筛 top_k 候选 2. 用 Embedding语义匹配计算相似度 3. 取 BM25 高分但 Embedding 低分的文档作为难负样本 if self.corpus_embeddings is None: raise RuntimeError(请先调用 precompute_embeddings()) query_embedding self.model.encode( [query], normalize_embeddingsTrue )[0] # 计算 query 与所有文档的语义相似度 similarities np.dot(self.corpus_embeddings, query_embedding) # 排除正样本自身 similarities[positive_idx] -1.0 # 取相似度最高的 top_k 个最难区分的负样本 hard_indices np.argsort(similarities)[-self.top_k_hard:][::-1] return [self.corpus[i] for i in hard_indices] def sample_random_negatives(self, exclude_idx: int, count: int 3) - List[str]: 随机采样负样本——用于增加训练多样性 import random indices list(range(len(self.corpus))) indices.remove(exclude_idx) sampled random.sample(indices, min(count, len(indices))) return [self.corpus[i] for i in sampled]对比学习训练器import torch import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader from sentence_transformers import SentenceTransformer, losses, InputExample class ContrastiveRAGTrainer: RAG 对比学习训练器 def __init__( self, model_name: str BAAI/bge-base-zh-v1.5, temperature: float 0.05, # InfoNCE 温度参数 learning_rate: float 2e-5, ): self.model SentenceTransformer(model_name) self.temperature temperature self.learning_rate learning_rate def prepare_training_data(self, samples: List[TrainingSample]) - Dataset: 构造训练数据 - 每个样本包含 (query, positive, negative_1, ..., negative_n) - 使用 MultipleNegativesRankingLoss 自动构造 batch 内负样本 examples [] for sample in samples: # 核心正样本对 examples.append(InputExample( texts[sample.query, sample.positive] )) # 额外负样本模型会同时拉远这些样本的距离 examples.append(InputExample( texts[sample.query, sample.negative_hard] )) # 三元组损失query 应该更接近 positive 而非 negative examples.append(InputExample( texts[sample.query, sample.positive, sample.negative_hard] )) return examples def train( self, train_samples: List[TrainingSample], epochs: int 3, batch_size: int 16, output_path: str ./fine_tuned_embedding, ): 对比训练主流程 使用 MultipleNegativesRankingLoss - batch 内除正样本外的所有其他样本都作为负样本 - 天然扩大 batch_size 更多负样本 train_examples self.prepare_training_data(train_samples) # DataLoader 构造 train_dataloader DataLoader( train_examples, shuffleTrue, batch_sizebatch_size, ) # 损失函数batch 内对比学习 train_loss losses.MultipleNegativesRankingLoss(self.model) # 训练 print(f开始训练epochs{epochs}, samples{len(train_samples)}) self.model.fit( train_objectives[(train_dataloader, train_loss)], epochsepochs, warmup_stepsint(len(train_dataloader) * 0.1), output_pathoutput_path, show_progress_barTrue, ) print(f训练完成模型保存至: {output_path}) def evaluate( self, test_samples: List[TrainingSample], k_values: List[int] [1, 3, 5, 10], ) - dict: 评估检索精度 - RecallK: 正样本是否在 Top-K 中 - MRR: 第一个正样本的平均倒数排名 metrics defaultdict(float) total len(test_samples) for sample in test_samples: query_emb self.model.encode( [sample.query], normalize_embeddingsTrue )[0] # 候选池正样本 所有负样本 candidates [sample.positive, sample.negative_hard] candidate_embs self.model.encode( candidates, normalize_embeddingsTrue ) # 计算相似度并排序 scores np.dot(candidate_embs, query_emb) ranked_indices np.argsort(scores)[::-1] # 正样本的理想排名是 0排在第一位 pos_rank np.where(ranked_indices 0)[0][0] 1 # 1-based for k in k_values: if pos_rank k: metrics[fRecall{k}] 1 metrics[MRR] 1.0 / pos_rank # 取平均 for key in metrics: metrics[key] / total return dict(metrics)四、边界分析与 Trade-offs难负样本的双刃剑难负样本训练效果好但挖掘成本高如果难负样本中包含假阴性实际上是正样本但标注错了会严重损害模型性能建议先跑 3 轮人工标注 规则过滤确认难负样本质量后再训练温度参数的调节温度越低如 0.01模型越关注难负样本梯度越陡温度越高如 0.1模型对难易负样本的区分度降低训练更稳定生产建议从 0.05 开始观察 loss 曲线是否震荡评估数据泄漏确保测试集的负样本不包含训练集使用的负样本样本。否则评估指标会虚高。计算成本全量语料的 Embedding 预计算需要 GPU 时间。10 万文档 × 768 维用 V100 约 5 分钟。五、总结提升 RAG 检索精度不仅要优化 Embedding 模型本身更要优化训练数据正样本决定检索的上限——高质量标注是关键难负样本决定检索的下限——模型学会区分形似神不似的文档InfoNCE Loss天然适合 batch 内对比学习实操路径先用 LLM 辅助标注一批正负样本 → 跑对比学习微调 → 用标注集评估 → 部署上线 → 收集用户反馈持续迭代。不要一上来就改模型架构数据质量提升往往比模型改进更显著。