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

arXiv: https://arxiv.org/abs/2604.23336 · HTML: https://mlsanity.com/paper/2604.23336/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Knowledge distillation for LLMs, one of 12 directions in Knowledge distillation.
- Builds on 25 of its 31 references that are in the map.
- First posted on arXiv 25 Apr 2026.

## Placed in

- Knowledge distillation / [Knowledge distillation for LLMs](https://mlsanity.com/distill/knowledge-distillation-for-llms/)

## Builds on (most cited first, 25 of 25)

- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Do As I Can, Not As I Say: Grounding Language in Robotic Affordances](https://mlsanity.com/paper/2204.01691/) (2022)
- [MS MARCO: A Human Generated MAchine Reading COmprehension Dataset](https://mlsanity.com/paper/1611.09268/) (2016)
- [LlamaFactory: Unified Efficient Fine-Tuning of 100+ Language Models](https://mlsanity.com/paper/2403.13372/) (2024)
- [BEIR: A Heterogenous Benchmark for Zero-shot Evaluation of Information Retrieval Models](https://mlsanity.com/paper/2104.08663/) (2021)
- [Inner Monologue: Embodied Reasoning through Planning with Language Models](https://mlsanity.com/paper/2207.05608/) (2022)
- [A Survey on RAG Meeting LLMs: Towards Retrieval-Augmented Large Language Models](https://mlsanity.com/paper/2405.06211/) (2024)
- [Is ChatGPT Good at Search? Investigating Large Language Models as Re-Ranking Agents](https://mlsanity.com/paper/2304.09542/) (2023)
- [Retrieval-Augmented Generation for AI-Generated Content: A Survey](https://mlsanity.com/paper/2402.19473/) (2024)
- [One Embedder, Any Task: Instruction-Finetuned Text Embeddings](https://mlsanity.com/paper/2212.09741/) (2022)
- [Robots That Ask For Help: Uncertainty Alignment for Large Language Model Planners](https://mlsanity.com/paper/2307.01928/) (2023)
- [Improving Text Embeddings with Large Language Models](https://mlsanity.com/paper/2401.00368/) (2023)
- [Towards Emotional Support Dialog Systems](https://mlsanity.com/paper/2106.01144/) (2021)
- [Fine-Tuning LLaMA for Multi-Stage Text Retrieval](https://mlsanity.com/paper/2310.08319/) (2023)
- [RocketQAv2: A Joint Training Method for Dense Passage Retrieval and Passage Re-ranking](https://mlsanity.com/paper/2110.07367/) (2021)
- [Task-aware Retrieval with Instructions](https://mlsanity.com/paper/2211.09260/) (2022)
- [PsyQA: A Chinese Dataset for Generating Long Counseling Text for Mental Health Support](https://mlsanity.com/paper/2106.01702/) (2021)
- [Adversarial Retriever-Ranker for dense text retrieval](https://mlsanity.com/paper/2110.03611/) (2021)
- [Llama2Vec: Unsupervised Adaptation of Large Language Models for Dense Retrieval](https://mlsanity.com/paper/2312.15503/) (2023)
- [Retrieval Augmented Generation (RAG) and Beyond: A Comprehensive Survey on How to Make your LLMs use External Data More Wisely](https://mlsanity.com/paper/2409.14924/) (2024)
- [Introduction to Latent Variable Energy-Based Models: A Path Towards Autonomous Machine Intelligence](https://mlsanity.com/paper/2306.02572/) (2023)
- [A Comprehensive Survey on Knowledge Distillation](https://mlsanity.com/paper/2503.12067/) (2025)
- [Promptriever: Instruction-Trained Retrievers Can Be Prompted Like Language Models](https://mlsanity.com/paper/2409.11136/) (2024)
- [OneGen: Efficient One-Pass Unified Generation and Retrieval for LLMs](https://mlsanity.com/paper/2409.05152/) (2024)
- [Large Language Model Can Be a Foundation for Hidden Rationale-Based Retrieval](https://mlsanity.com/paper/2412.16615/) (2024)

## Abstract (arXiv)

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.