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

Abstract: Unlike traditional fact-based retrieval, rationale-based retrieval typically necessitates cross-encoding of query-document pairs using LLMs, 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.

• The paper introduces Rabtriever, which efficiently distills cross-encoder rationale reasoning into a bi-encoder using JEPA-based on-policy distillation.
• It achieves nearly optimal recall and MRR on benchmarks like ESConv and SayCan while reducing online computation to linear time.
• The method employs joint MSE and reverse KL losses to align student predictions with the teacher, enabling scalable deployment in complex retrieval tasks.

Efficient Rationale-based Retrieval via On-Policy Distillation from Generative Rerankers based on JEPA

Motivation and Problem Setting

The paper introduces Rabtriever, an efficient LLM-based retriever architecture designed for rationale-based retrieval scenarios, where the relationship between queries and documents is grounded in sophisticated reasoning, expert logic, or procedural knowledge. Conventional dense retrieval methods, which encode queries and documents independently and compute similarity in a shared latent space, are insufficient for tasks requiring context-dependent rationale comprehension. LLM-based rerankers, which jointly encode query-document pairs, achieve superior reasoning but suffer from prohibitive computational complexity due to costly online cross-encoding for each pair. Rabtriever aims to bridge this gap, preserving reranker's nuanced relevance comprehension while matching the runtime efficiency of bi-encoder retrievers.

Architecture and Methodology

Rabtriever is constructed by initializing its backbone from a trained LLM-based cross-encoder reranker (LaHoRe-style teacher) and then performing distillation through the Joint-Embedding Predictive Architecture (JEPA) paradigm. The teacher reranker processes concatenated query-document inputs and outputs binary relevance scores based on generative LLM log-probabilities. Document embeddings, due to LLM causal attention, can be precomputed offline for both teacher and student.

The distillation phase leverages JEPA to embed query-document interaction into Rabtriever's latent space. Specifically, Rabtriever inserts a lightweight predictor (2-layer MLP + elementwise multiplication) between the frozen LLM backbone and scoring head. This predictor transforms the query embedding, implicitly grounding it in the offline document embedding. The objective combines MSE loss on reconstructed contextual query embeddings and an auxiliary reverse KL divergence loss between teacher and student logits, encouraging the student to precisely match the teacher’s prediction modes and suppress suboptimal behaviors.

Figure 1: Rabtriever employs JEPA to project query embedding into a latent space guided by the document embedding, minimizing the energy difference to the teacher’s cross-encoded embedding.

Figure 2: The Rabtriever pipeline: I) Teacher fine-tuning and cross-encoding; II) Student initialization with frozen backbone and JEPA predictor insertion; III) Predictor-conditioned embedding transformation and scoring.

This architecture and loss design allow Rabtriever to independently encode queries and documents while reconstructing the cross-encoding comprehension achieved by full rerankers.

Computational Efficiency and Complexity Analysis

Rabtriever's inference is highly efficient. Document embeddings are precomputed offline, and online computation involves only linear operations (MLP and vector multiplication) after encoding the query. Formal complexity analysis demonstrates that Rabtriever reduces the quadratic dependency on document length (characteristic of cross-encoder rerankers) to a linear dependency, providing orders-of-magnitude faster response times in practical retrieval pipelines.

Figure 3: Rabtriever achieves recall comparable to rerankers with runtime comparable to bi-encoder retrievers; reverse KL distillation targets mode-seeking behavior.

Figure 4: Empirical response times exhibit Rabtriever's linear scaling with document length, confirming the theoretical efficiency claims.

Experimental Evaluation

Rabtriever was evaluated on demanding rationale-based retrieval tasks (empathetic conversation and robotic skill selection) as well as large-scale fact-based benchmarks (MS MARCO, BEIR). Benchmarks included ESConv, PsyQA, and SayCan for rationale-based tasks, and MS MARCO/BEIR for generalized retrieval.

Rabtriever consistently outperformed all bi-encoder retriever baselines (TART-Contriever, RepLLaMA, GritLM, LlaMA2Vec, etc.), nearly matching its teacher cross-encoder reranker’s strong recall and MRR scores, with only marginal accuracy degradation. On ESConv, Rabtriever achieved R@1 of 32.8% (vs. teacher's 34.4%), with average response time only ~95ms, tenfold faster than cross-encoder-based rerankers. On SayCan, Rabtriever obtained R@1 of 67.3%, significantly surpassing bi-encoder competitors. Ablations confirmed that both JEPA MSE loss and reverse KL are necessary for distillation efficacy, and that simple predictor designs suffice for high performance.

Figure 5: MRR comparison on MS MARCO reveals Rabtriever’s near-optimal ranking performance among retrievers.

Generalization and Scalability

Rabtriever demonstrated strong zero-shot generalization on BEIR, achieving nDCG@10 competitive with SoTA retrievers and approaching cross-encoder rerankers. Method scalability was further corroborated: Rabtriever maintained competitive retrieval accuracy using smaller backbone LLMs (0.5B, 1.5B), implying flexible tradeoffs between speed and performance.

Figure 6: Rabtriever’s recall performance sustains across model sizes, supporting efficient deployment.

Methodological Implications

The on-policy distillation strategy employed in Rabtriever, using JEPA with latent energy minimization and reverse KL, represents a robust knowledge transfer approach from cross-encoder-based teachers to bi-encoder students. This paradigm enables retrievers to capture document-conditioned rationale comprehension efficiently, supporting practical deployment in complex, reasoning-driven retrieval scenarios where computational resources are a constraint.

The reverse KL loss specifically biases the student toward the teacher's most confident behaviors, critical for binary-choice generation tasks, and provides sharper mode-seeking than conventional forward KL objectives.

Figure 7: Recall curves under varying loss weights and MLP predictor layers indicate optimal hyperparameter selection for student distillation.

Theoretical Impact and Future Directions

Rabtriever advances the state of rationale-based retrieval by reconciling cross-encoder semantic comprehension with bi-encoder efficiency. Its JEPA-based distillation framework may catalyze broader adoption of energy-based joint embedding architectures for knowledge transfer in retrieval, multimodal alignment, and grounded reasoning. Extensions could include richer predictors for fine-grained rationale modeling, integration with multimodal JEPA (cf. M3-JEPA [pmlr-v267-lei25b]), or adaptation to open-domain retrieval pipelines requiring long-document handling and massive-scale inference.

The demonstrated efficiency-scaling opens avenues for deploying retrieval-augmented generation (RAG) systems in domains previously prohibitive due to cross-encoder bottlenecks. Further research is warranted to explore Rabtriever’s scalability for large document pools, extremely long documents, and integration into end-to-end generative LLM pipelines.

Conclusion

Rabtriever presents an effective solution for rationale-based retrieval challenges, distilling cross-encoder reranker capabilities into an efficient bi-encoder retriever using JEPA and reverse KL objectives. It achieves nearly optimal retrieval performance across several reasoning-intensive tasks and benchmarks, with practical inference efficiency suitable for real-world deployment. The work encourages future research into joint embedding predictive distillation, scalable rationale-based retrieval, and deeper integration of energy-based models within LLM grounding tasks.