Rabtriever: Distilled Rationale-Aware Retriever

• Rabtriever is a distilled rationale-aware retriever that integrates JEPA-based latent matching to capture underlying reasoning and strategy in query-document pairs.
• It uses on-policy distillation from a generative reranker teacher to enable offline document encoding and efficient, independent query processing.
• Empirical results show competitive performance on rationale-based tasks with reduced online computational complexity compared to full cross-encoder methods.

Rabtriever is a distilled rationale-aware retriever introduced in “Efficient Rationale-based Retrieval: On-policy Distillation from Generative Rerankers based on JEPA” (Chen et al., 25 Apr 2026). It targets rationale-based retrieval, a setting in which relevance depends not only on explicit facts or semantic similarity but on retrieving the reasoning, strategy, procedural logic, or latent expert rationale appropriate to a query. The paper contrasts this with conventional dense retrieval through a knowledge-stratification view spanning Level-1: Explicit facts, Level-2: Implicit facts, Level-3: Interpretable rationales, and Level-4: Hidden rationales, and positions Rabtriever as a system that preserves much of a generative cross-encoder reranker’s cross query-document comprehension while restoring independent encoding of queries and documents for efficient retrieval (Chen et al., 25 Apr 2026).

1. Problem formulation and retrieval setting

Rabtriever is motivated by tasks in which the relevant document is not merely a passage that states matching facts, but one that contains the rationale needed to act or respond correctly. The paper gives two canonical examples: empathetic conversations, where retrieval concerns the best support strategy for a dialogue context, and robotic manipulation, where retrieval concerns the most appropriate low-level skill for a high-level instruction under environmental constraints (Chen et al., 25 Apr 2026).

This distinction matters because a strong solution for rationale-based retrieval is often a cross-encoder reranker that concatenates query and document and lets an LLM reason over both jointly. The paper identifies three costs of that design. Every query-document pair must be encoded online; the self-attention cost is quadratic in sequence length; and documents cannot be pre-indexed in the same way as in a bi-encoder. The resulting online complexity is written as

$n\,\mathcal{O}(L_D + L_Q)^2,$

where $n$ is the number of candidates, $L_D$ the document length, and $L_Q$ the query length (Chen et al., 25 Apr 2026).

Rabtriever addresses this efficiency-comprehension tradeoff. Its stated objective is to preserve the cross query-document comprehension needed for rationale-based retrieval while enabling independent encoding of documents and queries, offline document computation, and retriever-style indexing (Chen et al., 25 Apr 2026).

2. Teacher–student architecture

Rabtriever is trained from an LLM-based generative reranker that serves as a teacher. The teacher is prompted in a document-first causal format:

Document: $\{D\}$ Query: $\{Q\}$ Can Query be appropriately replied with Document? If the answer is true, choose <T>; otherwise, choose <F>.

The document-first order is central. Because the teacher is autoregressive, the paper states

$\text{LLM}(D, Q)[:L_D] = \text{LLM}(D),$

so the document portion is unchanged by later query tokens. This makes document-side representations reusable and is the basis for offline document encoding (Chen et al., 25 Apr 2026).

Rabtriever itself is a bi-encoder-style student initialized from the teacher: $$\mathcal{M} \leftarrow \mathcal{M}^{\text{ce}}, \qquad \mathcal{H} \leftarrow \mathcal{H}^{\text{ce}}.$$ The backbone encoder $\mathcal{M}$ and head $\mathcal{H}$ are then frozen. The student’s general retrieval form is

$n$0

where $n$1 and $n$2 are independently encoded document and query embeddings, and $n$3 is a lightweight predictor (Chen et al., 25 Apr 2026).

The key architectural move is to reconstruct a document-conditioned query representation that approximates what the cross-encoder teacher would have produced. The paper writes

$n$4

with $n$5 denoting concatenation. The predictor is implemented as

$n$6

where $n$7 is elementwise multiplication and $n$8 is the last document embedding. The paper’s ablations later show that this multiplication operator is not incidental but structurally important (Chen et al., 25 Apr 2026).

3. JEPA adaptation and on-policy distillation

Rabtriever frames its distillation procedure through the Joint-Embedding Predictive Architecture (JEPA). The generic JEPA objective is introduced as

$n$9

with input $L_D$0, target $L_D$1, latent control $L_D$2, and nonnegative discrepancy $L_D$3 (Chen et al., 25 Apr 2026).

In Rabtriever, the mapping is explicit. The query embedding $L_D$4 is the input, the teacher’s contextual representation is the target, and the document embedding $L_D$5 is the latent control. The adapted JEPA objective is

$L_D$6

This means the student does not simply regress a scalar score. It attempts to reconstruct the teacher’s document-conditioned latent geometry (Chen et al., 25 Apr 2026).

The implemented training objective combines latent-space MSE with reverse-KL on the teacher’s output distribution: $L_D$7 The paper states that the reverse-KL term is added to strengthen the sampling efficiency of on-policy distillation and to reshape the student’s logit distribution toward the teacher’s dominant modes, especially around the binary decision tokens $L_D$8 and $L_D$9 (Chen et al., 25 Apr 2026).

The distillation is described as on-policy because the teacher is frozen while the student is actively updating, and teacher embeddings and logits are computed on the student-experienced samples rather than treated as a static offline table. This suggests a training process aligned to the student’s evolving distribution rather than a purely fixed supervision set. A plausible implication is that the reverse-KL term is intended not merely as auxiliary calibration, but as a stabilizer for the student’s changing output policy (Chen et al., 25 Apr 2026).

4. Retrieval scoring and computational properties

At inference time, Rabtriever separates the heavy document computation from the online query path. Documents are encoded once: $L_Q$0 queries are encoded online: $L_Q$1 and candidate scoring then uses the lightweight predictor and frozen head (Chen et al., 25 Apr 2026).

The paper states that the student relevance score is computed from $L_Q$2 in the same way as the teacher, using the relative confidence of $L_Q$3 and $L_Q$4. This preserves the teacher’s binary relevance semantics while avoiding full cross-encoding at retrieval time (Chen et al., 25 Apr 2026).

The paper’s complexity comparison is central:

• Bi-encoder baseline:

$L_Q$5

• Cross-encoder reranker:

$L_Q$6

• Rabtriever:

$L_Q$7

The paper summarizes this as optimizing the teacher’s quadratic complexity on document length to linear in the online path. More precisely, the expensive online interaction over document tokens is removed from the main LLM computation and replaced by a lightweight predictor-based interaction (Chen et al., 25 Apr 2026).

Empirically, on the rationale-based tasks the reported average response times are 93 ms for LLaMA2Vec, 95 ms for Rabtriever, and 980 ms for the LaHoRe teacher. The paper also reports approximately linear response-time growth for Rabtriever as document length increases, versus an approximately quadratic trend for the teacher (Chen et al., 25 Apr 2026).

5. Empirical results and ablation evidence

The main experiments cover rationale-based retrieval tasks—ESConv, PsyQA, and SayCan—and conventional retrieval benchmarks including MS MARCO and BEIR (Chen et al., 25 Apr 2026).

The rationale-based results establish the paper’s central empirical claim: Rabtriever is close to the teacher and substantially stronger than conventional retriever baselines on tasks where relevance depends on strategy or procedural fit.

Dataset	Rabtriever	Reference point
ESConv	R@1 32.8, R@3 66.4, MRR 54.7	LaHoRe: 34.4, 67.1, 54.8
PsyQA	R@1 32.8, R@3 76.1, MRR 55.1	LaHoRe: 35.6, 77.6, 58.2
SayCan	R@1 67.3, R@3 76.9, R@5 78.8, MRR 73.7	LaHoRe: 71.2, 88.5, 94.2, 79.8

On ESConv, Rabtriever exceeds the strongest retriever baselines such as LLaMA2Vec and OneGen, while nearly matching the teacher’s MRR. On PsyQA, it remains close to the teacher and competitive with the best retriever baseline. On SayCan, it substantially outperforms retriever baselines and also exceeds RankLLaMA by a wide margin, which is the clearest rationale-based retrieval gain reported in the paper (Chen et al., 25 Apr 2026).

The conventional retrieval results are more conservative. On MS MARCO, the teacher reaches MRR 52.4 and Rabtriever reaches MRR 49.30. On BEIR, average nDCG@10 is 65.8 for the teacher, 59.2 for GritLM, 59.0 for E5-Mistral, and 58.2 for Rabtriever. The paper states that Rabtriever is best on 2/13 subtasks and second best on 4/13 subtasks, supporting the claim that a rationale-oriented student can remain competitive on standard retrieval (Chen et al., 25 Apr 2026).

The ablations are unusually decisive. Removing the multiplication operator causes collapse on PsyQA and SayCan, with R@1 2.0 and 9.6 respectively. Removing the latent MSE term yields R@1 0.0 on PsyQA and 23.1 on SayCan. Removing reverse KL drops PsyQA R@1 from 32.8 to 24.1 and SayCan R@1 from 67.3 to 21.2. Replacing reverse KL with forward KL partially recovers performance, but remains markedly weaker than the full model (Chen et al., 25 Apr 2026).

Two additional results are notable. On ESConv, the paper reports $L_Q$8 against LLaMA2Vec for Rabtriever R@1 = 32.8 and teacher R@1 = 34.4. In an end-to-end ESConv pipeline with Qwen2-70B-Instruct as generator, Rabtriever-based retrieval yields an 80.12% win rate against original responses in pairwise GPT-4 judging (Chen et al., 25 Apr 2026).

6. Position within retrieval research and stated limitations

Rabtriever belongs to a broader shift from generic semantic retrieval toward retrieval models specialized for reasoning, rationale, or dynamic orchestration. Within the papers provided, several neighboring lines are especially relevant.

RaDeR trains dense retrievers from retrieval-augmented mathematical reasoning trajectories and self-reflective relevance evaluation, using query types such as partial CoT, LLM-generated reasoning queries, and original questions (Das et al., 23 May 2025). MoR instead addresses query diversity through a zero-shot weighted combination of heterogeneous sparse, dense, and human retrievers, emphasizing retriever orchestration rather than distillation from a cross-encoder teacher (Kalra et al., 18 Jun 2025). RouterRetriever routes queries over domain-specific LoRA expert embedding models using a training-free pilot-embedding router and hard top-1 expert selection (Lee et al., 2024). These systems occupy adjacent parts of the design space: reasoning-aware supervision, mixture-based orchestration, and expert routing, whereas Rabtriever’s distinctive move is JEPA-based latent reconstruction from a generative reranker (Das et al., 23 May 2025, Kalra et al., 18 Jun 2025, Lee et al., 2024).

The paper is explicit about Rabtriever’s limitations. Its upper bound depends on teacher quality; it is not as cheap as a pure interaction-free bi-encoder because the predictor adds candidate-time cost; and scalability to broader massive-scale retrieval remains underexplored. The method also assumes access to a white-box teacher, since hidden states and logits are both distilled. The document-before-query causal property is essential to the reuse of document representations; this is a structural assumption of the framework, not an incidental prompt choice (Chen et al., 25 Apr 2026).

Taken together, these properties define Rabtriever as a distilled rationale-aware retriever: a system that preserves retriever-style indexability while importing reranker-style cross-input understanding. Its main significance lies not in replacing all dense retrieval, but in showing that rationale-based retrieval can be made practical through a frozen teacher-initialized backbone, a lightweight document-conditioned predictor, JEPA-style latent matching, and reverse-KL on-policy distillation (Chen et al., 25 Apr 2026).