direct-rag-learning
通过 E-step 采样 + M-step 加权优化，实现检索选择器与 LLM 生成器的端到端联合训
加载项目详情…
本应用为开源项目，仅供学习研究，请遵守其开源协议。
微信扫码咨询客服
POMOAI — 让开源 AI 即刻运行
通过 E-step 采样 + M-step 加权优化，实现检索选择器与 LLM 生成器的端到端联合训
加载项目详情…
本应用为开源项目，仅供学习研究，请遵守其开源协议。
当你向一个大模型提问「为什么加州总是闹火灾」，一个优秀的 RAG 系统会先检索相关文档，再基于检索到的知识生成答案。听起来很美好，但长期以来，检索器（Retriever）和生成器（Generator）之间存在一个根本性的矛盾——它们各自为战，互不了解。检索器不知道生成器需要什么样的知识，生成器也不理解检索器为什么挑了这些文档。这种「鸡同鸭讲」的困境，严重制约了 RAG 系统的上限。
图1：DRO 框架 E-step 和 M-step 交替优化示意图
你有没有注意到，同样是检索到的文档，只要调整一下排列顺序，模型的回答质量可能就天差地别？这是因为 LLM 是一个对上下文顺序极为敏感的模型——相关文档放前面还是放后面，直接影响它「看到」什么信息。
但现有方法存在两个致命缺陷：
第一种：独立优化，只调检索器或只调生成器。这相当于训练一个短跑运动员和一个跳远运动员，但从不让他们一起配合跑接力。
第二种：枚举训练，把所有可能的文档排列都拿来训练。但文档数量稍有增加，排列组合数就爆炸了——Top-5 文档就有 120 种排列，根本算不过来。
TOIS 2026 收录的这篇论文提出了一个优雅的解决方案：Direct Retrieval-augmented Optimization（DRO）——把文档排列当作一个「隐变量」，用期望最大化（EM）的思想，让检索选择器和答案生成器在同一个训练循环里互相学习、共同进化。
DRO 把训练过程拆成两个交替进行的步骤：
E-step（估计步）：给一个问答对 (x, y)，由选择模型（Selection Model）根据问题 x 生成候选文档的排列分布，然后通过重要性采样给每个排列打一个权重——权重由生成器给出，高权重意味着这个排列能帮生成器答对这道题。这个步骤不需要遍历所有排列，只要采样若干个就够了。
M-step（最大化步）：拿到 E-step 的采样结果和权重后，用加权的方式同时更新选择模型和生成器。选择模型学会「什么样的文档排列更适合这道题」，生成器学会「如何更好地利用已有的文档来回答」。两边互相给反馈，形成正循环。
这个过程的数学基础很扎实：从 log 似然出发，推导出 importance-weighted ELBO，可以证明训练过程中答案正确的概率只会单调递增、不会下降。整个训练过程是稳定的，不会出现 RL 训练中常见的「reward hacking」或者崩溃问题。
mangopy/direct-rag-learning 是一个标准的 Python 研究代码库，代码结构清晰，按功能分为三大模块：
src/prepare/ — 数据预处理模块
prepare.py：统一多个问答数据集（NQ、HotpotQA、MuSiQue、2WikiMultihopQA、WoW）的格式sampling.py：E-step 的重要性采样核心逻辑rank.py：列表级排序模型（RankLLM）的调用封装entropy.py / weight.py：M-step 的权重计算和重加权src/tuning/ — 模型训练模块
train.py + trainer.py：基于 DeepSpeed 的分布式训练入口，支持 ZeRO-3 优化model/：LoRA / 全量微调两种模式hparams/：超参数配置src/utilize/ — 推理与评测模块
utilize.py：调用 vLLM 进行批量推理metrics.py：pytrec_eval 计算 NDCG、MAP 等检索指标_evaluate.py：vLLM LLM 推理封装（独立进程启动）整个项目的技术栈非常前沿：PyTorch 2.3 + Transformers 4.43 + DeepSpeed 0.14 + vLLM 0.5（推理加速），依赖中还有 xformers（Flash Attention）和 peft（LoRA 微调）。代码质量在学术项目中属于上乘，有完整的参数解析（argparse + fire）、日志记录（logging）和随机种子固定。
论文在 5 个主流问答数据集上做了实验，DRO-Llama-3-8B 在 NQ 数据集上达到了 45.76% EM / 55.42% F1，相比最强基线提升了 12.32% EM；在 MuSiQue 多跳推理数据集上提升最大，达到 16.54% EM。
这个提升幅度在 RAG 领域相当可观。更值得注意的是文档选择精度的提升——Recall@K 提升了 17.78%，说明 DRO 学到的选择模型确实知道什么文档「配得上」这道题。
训练收敛速度也令人满意：主要收益在前 3 个 epoch 就能获得，第 4 个 epoch 开始进入精调阶段。对于一个需要多卡训练的大模型来说，这种快速收敛特性节省了大量计算成本。
没有任何技术是完美的，DRO 也有它的局限性：
GPU 门槛极高：训练需要 8 卡 A100/H100（每卡 24GB+），最低也需要 8x RTX 3090。这对于普通研究者来说几乎是不可及的。推理虽然可以用单卡，但前置的 ColBERT 检索模型也需要 GPU。
依赖外部检索器：DRO 的选择模型建立在 ColBERT 排序结果之上，ColBERT 本身的检索质量就是上限。如果 ColBERT 漏掉了关键文档，后面的优化也无济于事。
学术代码，未必适合生产：requirements.txt 中有 fire 这种非主流依赖，多卡训练依赖 DeepSpeed 配置文件（script/ds_z3_config.json），没有 Docker、没有 CI/CD，对于想快速部署到生产环境的团队来说，移植成本不小。
DRO 的核心贡献不在于某一次 SOTA，而在于提出了一种联合优化的新思路——不是让检索器和生成器各自为战，而是让它们在同一个训练目标下相互适应。这种思路对未来的 RAG 架构设计有重要的启示意义。
从更大的视角看，DRO 本质上是把强化学习的思想引入到了 RAG 的训练过程中（E-step 的重要性采样可以类比 policy gradient），但避免了 RL 的不稳定性。这是近年来 RAG 研究从「怎么检索」走向「怎么训练」趋势的一个典型代表。
TOIS（ACM Transactions on Information Systems）是信息检索领域的顶刊之一，能被 TOIS 2026 收录，说明学术界对这种「知识选择与答案生成联合优化」的范式是认可的。对于正在做 RAG 优化的工程师来说，这是一个值得关注的方向——它提供了一套有理论保障的训练方法，而不是靠玄学调参。