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

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Chen, Teng, Xu, Sheng, Guo, Feixiang, Wang, Xiaoyu, Gu, Qingqing, Li, Hongyan, Ji, Luo
Formato: Preprint
Publicado: 2026
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866909040541433856
author Chen, Teng
Xu, Sheng
Guo, Feixiang
Wang, Xiaoyu
Gu, Qingqing
Li, Hongyan
Ji, Luo
author_facet Chen, Teng
Xu, Sheng
Guo, Feixiang
Wang, Xiaoyu
Gu, Qingqing
Li, Hongyan
Ji, Luo
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2604_23336
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Efficient Rationale-based Retrieval: On-policy Distillation from Generative Rerankers based on JEPA
Chen, Teng
Xu, Sheng
Guo, Feixiang
Wang, Xiaoyu
Gu, Qingqing
Li, Hongyan
Ji, Luo
Information Retrieval
Computation and Language
Machine Learning
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.
title Efficient Rationale-based Retrieval: On-policy Distillation from Generative Rerankers based on JEPA
topic Information Retrieval
Computation and Language
Machine Learning
url https://arxiv.org/abs/2604.23336