Retrieving from Your Own Mistakes: On-Policy Training of Embedding Models for Multi-Hop Retrieval-Augmented Generation

What it claims

A 0.6B chain embedding retriever, which calls no language model and uses one dense index, already beats graph-based RAG on complete evidence coverage and is competitive with a 7B search agent when trained only with teacher forcing. The authors argue its remaining weakness is exposure bias: at test time it conditions on its own wrong retrievals. On four-hop MuSiQue, 19% of states at step 2, 70% at step 3 and 89% at step 4 contain a non-gold passage. On-policy state training changes only the distribution of training states. With that change, the retriever outperforms every compared system on MuSiQue, and the MuSiQue-trained encoder transfers to longer chains.

How it works

The architecture follows MDR. The question and the passages retrieved so far are serialized into one text and embedded by a LoRA-tuned query encoder. A fixed document encoder scores passages against it, and retrieval is greedy for T steps. The target for any state is the earliest gold passage not yet retrieved. Other missing gold passages are removed from the contrastive denominator, so the label exists even for histories containing wrong passages. Negatives are the top-100 from the frozen encoder, unchanged from teacher forcing.

flowchart LR
  A[Current encoder rolls in on train questions] --> B[Dedupe visited states, sample M]
  G[M teacher-forced gold-subset states] --> C[Mixed batch]
  B --> C
  C --> D[Contrastive update, earliest-missing-gold target]
  D -->|next epoch: replace visited half| A

The mix is λ = 1/2. M is half of the roughly 4,700 teacher-forced training states, so the total state count matches the baseline. The visited half is re-collected and replaced every epoch rather than accumulated as in DAgger. The gold half keeps histories the encoder rarely reaches, such as long correct prefixes. The first epoch's visited states come from a teacher-forced encoder's roll-in, but LoRA weights start from scratch. Training runs for at most 10 epochs with early stopping on validation questions. The authors describe the procedure as DAgger with a dynamic oracle that the gold set supplies for free.

Evidence

The mismatch measurement used a teacher-forced retriever run for 4 steps on 1,000 MuSiQue test questions. It retrieves a missing gold passage in 67% of clean states against 28% of contaminated ones. The gap holds at every hop count: 71% vs 33% (two-hop), 67% vs 40% (three-hop), 57% vs 19% (four-hop). Deleting the wrong passages from the same contaminated states raises success from 20.8% to 32.1%. The wrong passages alone therefore cost about 11 points.

Scored on identical states, teacher-forced and on-policy encoders continue clean histories about equally well. After one or more mistakes, the on-policy encoder recovers far more often. The ablations split the gain:

• Synthetic contamination of gold histories, matched in the number of wrong passages, closes about half the gap.
• Visited states collected once are worth little more than synthetic contamination.
• The rest comes from refreshing visited states every epoch. Resampling synthetic contamination every epoch does not reproduce it.

How to use it

The authors release code and the long-chain HotpotQA and 2WikiMultiHopQA benchmarks at https://github.com/yausankit/on-policy-state-training. Training needs only gold evidence sets in chain order. No reward model, LLM or outcome reward is used, and the document index stays fixed throughout. The page gives no install commands, base encoder name or licence.

Limits

The fetched text cuts off in Section 4. It contains no coverage (C@k) or recall numbers for the main comparison, and it does not name the graph-based systems or the 7B search agent. The long-chain HotpotQA and 2WikiMultiHopQA variants are the authors' own constructions. Results rest on gold evidence sets that list the passages in chain order, which the target definition requires. Because the document encoder and the negatives stay frozen, the method changes only the query side.