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.
翻译:不同于传统基于事实的检索,基于推理的检索通常需要利用大语言模型对查询-文档对进行跨编码,带来大量计算成本。为解决此问题,本文提出R abtriever,通过独立编码查询和文档,同时具备与重排序器相当的跨查询-文档理解能力。我们首先训练基于LLM的生成式重排序器,该模型将文档置于查询之前,并利用对数概率促使LLM生成相关性得分。随后将其用作在线策略蒸馏框架中的教师模型,以R abtriever为学生模型重构教师的环境感知查询嵌入。为实现该效果,R abtriever首先从教师模型初始化并冻结参数,进而采用联合嵌入预测架构(JEPA)范式,在LLM层与输出头之间集成轻量可训练预测器,将查询嵌入投影至新隐空间(以文档嵌入为潜变量)。JEPA通过最小化该投影嵌入与教师嵌入之间的分布差异。为增强在线策略蒸馏的采样效率,我们额外添加基于LLM对数概率反向KL散度的辅助损失函数,以重塑学生模型的对数概率分布。R abtriever将教师模型在文档长度上的二次复杂度优化为线性复杂度,该结论经理论与实验双重验证。实验表明,在包括共情对话和机器人操控等多样化基于推理的任务中,R abtriever均优于各类检索基线,相较于重排序器仅有微小精度下降。该模型在MS MARCO和BEIR等传统检索基准上亦展现出良好泛化性,性能与最优检索基线持平。