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

2026-04-25

AI Summary

This work addresses the high computational cost of traditional rationale-based retrieval, which relies on large language models to cross-encode query–document pairs. The authors propose Rabtriever, the first approach to adapt the Joint Embedding Predictive Architecture (JEPA) paradigm to retrieval distillation. Rabtriever learns context-aware query embeddings via policy distillation from a generative reranker, enabling independent encoding of queries and documents. To enhance sampling efficiency, it incorporates a reverse KL divergence auxiliary loss and reduces the document-length-dependent quadratic complexity to linear. Experiments demonstrate that Rabtriever significantly outperforms existing retrieval baselines on rationale-intensive tasks such as empathetic dialogue and robotic manipulation, achieving accuracy close to that of the teacher reranker, while also attaining state-of-the-art performance on standard benchmarks including MS MARCO and BEIR.

Technology Category

• Search and Optimization: Learning to Search
• Data Mining & Knowledge Management: Conversational Systems for Recommendation & Retrieval
• Knowledge Representation and Reasoning: Knowledge Acquisition

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Graph Algorithms and Modeling for the Web: Querying, indexing, and retrieval in Web-related graphs

Abstract

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.

Problem

• rationale-based retrieval
• computational cost
• cross-encoding
• large language models
• retrieval efficiency

Innovation

• rationale-based retrieval
• on-policy distillation
• Joint-Embedding Predictive Architecture
• generative reranker
• efficient retrieval