arXiv:2604.23336cs.IRcs.CL2026-04被引 1

用生成式重排序器蒸馏出高效检索模型,速度提升且效果接近原模型。

Efficient Rationale-based Retrieval: On-policy Distillation from Generative Rerankers based on JEPA

论文配图:Efficient Rationale-based Retrieval: On-policy Distillation from Generative Rerankers based on JEPA
图 1 · 摘自论文原文
  • 基于生成式重排序器的在线蒸馏,让独立编码查询与文档的模型具备跨对理解能力。
  • 将文档长度复杂度从二次方降至线性,实验验证在多种任务上性能接近最优基线。
  • 适合追求高效推理的场景,尤其适用于需要快速响应的对话与机器人任务。

传统基于事实的检索通常依赖大语言模型对查询-文档对进行交叉编码,计算成本高。为解决此问题,我们提出Rabtriever,它独立编码查询与文档,却能实现与重排序器相当的跨对理解能力。首先训练一个基于LLM的生成式重排序器,将文档置于查询前,通过日志概率生成相关性评分。随后以该模型为教师,采用在线蒸馏框架,让Rabtriever作为学生重建教师的上下文感知查询嵌入。Rabtriever初始参数冻结自教师,引入联合嵌入预测架构(JEPA),在LLM层间加入轻量级可训练预测器,将查询嵌入投影至新隐藏空间,以文档嵌入为潜在向量,最小化投影嵌入与教师嵌入的分布差异。为提升在线蒸馏采样效率,还添加了反向KL辅助损失,重塑学生日志概率分布。理论与实证表明,Rabtriever将文档长度复杂度由二次方降为线性。实验显示,Rabtriever在情感对话、机器人操作等多样化的理由型检索任务中超越多个基线检索器,仅轻微损失准确率;同时在传统检索基准如MS MARCO和BEIR上表现良好,性能与最佳基线相当。

原文摘要 · Abstract (English)

Unlike traditional fact-based retrieval, rationale-based retrieval typically necessitates cross-encoding of query-document pairs using large language models, incurring substantial computational costs. To address this limitation, we propose Rabtriever, which independently encodes queries and documents, while providing comparable cross query-document comprehension capabilities to rerankers. We start from training a LLM-based generative reranker, which puts the document prior to the query and prompts the LLM to generate the relevance score by log probabilities. We then employ it as the teacher of an on-policy distillation framework, with Rabtriever as the student to reconstruct the teacher's contextual-aware query embedding. To achieve this effect, Rabtriever is first initialized from the teacher, with parameters frozen. The Joint-Embedding Predictive Architecture (JEPA) paradigm is then adopted, which integrates a lightweight, trainable predictor between LLM layers and heads, projecting the query embedding into a new hidden space, with the document embedding as the latent vector. JEPA then minimizes the distribution difference between this projected embedding and the teacher embedding. To strengthen the sampling efficiency of on-policy distillation, we also add an auxiliary loss on the reverse KL of LLM logits, to reshape the student's logit distribution. Rabtriever optimizes the teacher's quadratic complexity on the document length to linear, verified both theoretically and empirically. Experiments show that Rabtriever outperforms different retriever baselines across diverse rationale-based tasks, including empathetic conversations and robotic manipulations, with minor accuracy degradation from the reranker. Rabtriever also generalizes well on traditional retrieval benchmarks such as MS MARCO and BEIR, with comparable performance to the best retriever baseline.

检索增强模型蒸馏JEPA高效推理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。