Retrieve-for-Train:用离线强化学习编译查询扇出,以扩散检索器绕过推理延迟瓶颈
核心概要
该工作提出 Retrieve-for-Train 框架,先用离线强化学习训练一个扇出语言模型(基于 Gemma3-4B 与 Qwen3-4B,每次生成 10 个子查询),以 groundedness、Vendi Score 多样性、对齐度三项复合奖励评估整组结果,再把该行为蒸馏进一个 53.9M 参数的扩散检索器,使其在连续嵌入空间中一次非自回归并行生成完整目标集合,从而在开放式抽象检索与弱监督组合检索两类任务上超过单查询检索、零样本扩展与 Best-of-N 基线,并相对自回归方法取得 12 至 20 倍加速。
深度剖析
提出“奖励到数据编译”的三步流水线:RL 训练扇出语言模型、离线合成(查询→目标集合)监督数据、训练扩散检索器。 把集合级属性的优化从推理时思考预算转移到一次性离线训练,且监督合成不需要人工标注。 文中给出完整三步流程描述,并在两个检索任务与两个领域(时尚文生图、音乐文生乐)上报告结果。
用 groundedness、Vendi Score 多样性、对齐度三项复合奖励作为互相制衡的锚点,避免奖励被钻空子。 传统监督训练按 learning to rank 逐项打分,无法度量多样性、互补性等不可分解的集合级属性;该工作以整组结果为单位评分。 消融观察显示,去掉多样性项后模型会退化为“line ending line ending”这类无意义字符串以数学上利用数据库向量坐标。
将学到的扇出行为蒸馏进 53.9M 参数扩散模型,在连续嵌入空间一次并行生成全部目标方向。 绕开基于文本的思维链推理 token,把自回归的逐 token 延迟下限替换为非自回归单次生成。 报告相对自回归方法 12 至 20 倍加速;大规模上下文批次下自回归扇出延迟线性增长至近 50 秒,而扩散版本保持在亚秒到数秒。
在开放式抽象检索与弱监督组合检索两类设定上,Retrieve-for-Train 在多样性、对齐度与召回上均优于标准检索与零样本基线。 零样本 LLM 易出现“释义坍缩”,生成近义重复查询;该方法生成如“boots”“lace”这类语义上真正分叉且仍锚定数据库流形的子查询。 文中以柱状图对比两类任务(OAR 与 WSCR)的多项指标,并称其持续优于 Best-of-N 基线。
启示与展望
该结果面向需要返回互补结果集合的搜索与推荐场景,尤其是专门或多模态领域;其设定为固定数据库,质量由多样性、对齐度与数据库 groundedness 等集合级属性衡量,实验覆盖时尚文生图与音乐文生乐两个领域,扇出模型每次固定生成 10 个子查询。
读者可能仍会关注:复合奖励三项权重的具体取值与敏感性、Vendi Score 作为反锚点在不同嵌入骨干下的稳定性、53.9M 扩散检索器在更大规模或跨领域数据库上的表现,以及本文未展开的图表数值细节。
