PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Chenjie Qiu, Zixuan Zeng, Shaobin Shi, Bingfeng Chen, Dongsheng Luo, Boyan Xu

arXiv:2610.40409v1cs.AI

TL;DR

Complex questions need evidence from multiple passages, yet chain retrievers train on gold histories and then condition on their own retrievals. The paper trains on states reached by the current encoder, labeling them from the remaining gold evidence. This improves complete evidence retrieval and recovery from mistakes, including transfer to longer-chain benchmarks.

  • Problem

    Chain retrievers train on gold histories but use their own retrievals at test time, so their test states can fall outside the histories represented in training.

    Source: §3.3 · p. 3

  • Method

    On-policy state training retrieves chains with the current encoder each epoch and trains on a mix of its visited states and gold states, labeling each visited state with the earliest missing gold passage.

    Source: §1 · pp. 1–2

  • Results

    On MuSiQue and both long-chain benchmarks, on-policy state training raises complete coverage over teacher forcing; on MuSiQue it reaches 50.5 C@5 versus 44.6 C@5 for teacher forcing.

    Source: §1 · p. 2, §5.2.1 · p. 5

  • Takeaways & Limitations

    The gain depends both on exposure to contaminated histories and on refreshing visited states as the encoder changes.

    Source: §1 · p. 2

Abstract

from arXiv · show

Complex questions in retrieval-augmented generation require evidence spread over several documents, and an answer is reliable only when the retriever finds all of it. Chain embedding retrievers collect such evidence step by step, re-encoding the question together with the passages retrieved so far, and we find that they are already strong: trained only on gold histories, a small chain embedding retriever that calls no language model outperforms graph-based systems and is competitive with a search agent. This training, however, leaves a gap. At test time the retriever conditions on its own retrievals, which soon contain mistakes: by the third step, about half or more of the states it visits hold a wrong passage, and these passages lower its chance of retrieving a missing gold passage well beyond what the difficulty of the states explains. We propose on-policy state training, which trains the encoder on the states it reaches itself. At every epoch the current encoder retrieves chains for the training questions; each state it visits, including those built on its own mistakes, is labeled with the earliest gold passage still missing while the other missing gold passages are masked, so that no correct continuation is penalized; and the encoder is updated contrastively on these states mixed with teacher-forced ones, so that its training distribution follows it as it learns. The resulting retriever retrieves complete evidence more often than teacher-forced training and the systems we compare, and it transfers without further training to harder long-chain variants of HotpotQA and 2WikiMultiHopQA that we construct. On the same states, it recovers from its own mistakes far more often, and controlled comparisons attribute the gain to training on the states the encoder currently reaches.

1 Introduction

Chain retrievers must find all evidence across multiple passages, but they train on gold histories and test on their own retrievals. The paper proposes on-policy state training and finds that it improves complete evidence retrieval, especially by helping the encoder recover from its mistakes.

1 Introduction

  • On four-hop MuSiQue questions, 70% of third-step states and 89% of fourth-step states contain a non-gold passage, while such histories reduce missing-gold retrieval success beyond state difficulty alone.Deleting wrong passages from the same states raises success by 11 points.Source: §1 · p. 1
  • On-policy state training retrieves chains with the current encoder each epoch, labels visited states with the earliest missing gold passage, and mixes them with gold states for training.Other missing gold passages are masked so that retrieving them is not penalized.Source: §1 · pp. 1–2
  • Even with teacher forcing, a 0.6B chain retriever that uses no language-model call retrieves complete evidence more often than graph-based systems and is competitive with a 7B search agent.The comparison is reported at a fraction of the retrieval cost.Source: §1 · p. 2
  • Changing only training states raises complete coverage over teacher forcing on MuSiQue and both long-chain benchmarks; the resulting retriever beats every compared system on MuSiQue and transfers to longer chains.The MuSiQue-trained encoder transfers to harder long-chain variants of HotpotQA and 2WikiMultiHopQA.Source: §1 · p. 2
  • On the same states, teacher-forced and on-policy encoders continue clean histories about equally well, but the on-policy encoder recovers far more often after mistakes.The recovery advantage accumulates along the chain.Source: §1 · p. 2
  • The gain comes partly from exposure to contaminated histories and partly from refreshing visited states as the encoder changes; re-sampling synthetic contamination each epoch does not reproduce it.Synthetic contamination matched in the number of wrong passages closes about half the gap to on-policy training.Source: §1 · p. 2

2 Related Work

Prior systems for complex-question retrieval use language models, graphs, or chain embedding retrievers. This work situates its approach as dense, per-step supervision on the query side of a chain retriever, addressing the mismatch between reference histories and model-generated histories.

2.1 Retrieval for Complex Questions

  • Unlike methods learning retrieval from sparse episode-level rewards, this method derives dense per-step targets from the gold evidence set for every visited state, without language-model calls in training or retrieval.Search-R1 and Q-RAG learn from sparse outcome rewards [20].Source: §2.1 · p. 3

2.2 Chain and Dense Retrievers

  • Chain retrievers condition each hop on earlier retrieved passages, while this method refreshes the query histories used for training and leaves the negatives unchanged.The fixed document index and chain architecture distinguish the training change from alterations to the candidate side of the contrastive objective.Source: §2.2 · p. 3

2.3 Exposure Bias and On-Policy Training

  • The procedure is DAgger with a dynamic oracle: the learner visits states and an oracle labels them, with the gold evidence set supplying targets even off the gold path.Unlike cost-to-go approaches such as AggreVaTe and LOLS [18], the oracle labels every state using gold evidence.Source: §2.3 · p. 3

3 Problem Setup

A chain retriever repeatedly encodes a question with retrieved passages, but teacher-forced training supplies only gold histories. The setup defines targets for states containing retrieval errors and measures how often such mismatched states occur.

3.1 Chain Dense Retrieval

  • At each step, the retriever encodes the question and previously retrieved passages, scores corpus passages against a fixed document index, and builds a ranking whose complete-coverage metric checks whether all gold evidence appears in the top k.The chain state consists of the question and passages retrieved so far; complete coverage measures whether a question's whole gold set is present.Source: §3.1 · p. 3

3.2 Supervision from the Gold Set

  • For any state with missing gold passages, the target is the earliest missing passage in chain order, and other missing gold passages are removed from the loss denominator because retrieving them is also progress.The target remains defined even when the history contains passages outside the gold set.Source: §3.2 · p. 3

3.3 Teacher Forcing and the State Mismatch

  • After a wrong retrieval, later states fall outside the support of gold-history training, since the wrong passage remains in the query; this is the state mismatch.At test time the retriever conditions on states visited by its own policy.Source: §3.3 · p. 3

3.4 How Large Is the Mismatch?

  • On four-hop questions, 19% of second-step states, 70% of third-step states, and 89% of fourth-step states contain a non-gold passage; deleting wrong passages raises success on the same states from 20.8% to 32.1%.The clean-versus-contaminated retrieval gap remains across question lengths, so it is not explained by question length alone.Source: §3.4 · p. 4

4 On-Policy State Training

On-policy state training changes the distribution of states used to train the query encoder while holding the target, loss, negatives, and document index fixed. The learner retrieves its own chains, and an oracle labels visited states, including those off the gold path.

4 On-Policy State Training

  • On-policy state training changes only the query encoder’s training-state distribution, using learner roll-ins and an oracle defined even off the gold path.The architecture, target, loss, negative candidates, and document index remain fixed.Source: §4 · p. 4

4.1 Objective

  • Each epoch mixes newly sampled states from the current retriever with teacher-forced states, retaining supervision on long correct prefixes it rarely reaches.The visited-state sample is refreshed each epoch; the same state count and training schedule are used in both regimes.Source: §4.1 · p. 4

4.2 Labels for Visited States

  • Visited states are labeled with the earliest missing gold passage, while other missing gold passages are masked so no correct continuation is penalized.This supplies dense per-step supervision on contaminated histories without a language model, reward model, or outcome reward.Source: §4.2 · p. 4

4.3 Fixed Index and Negative Rule

  • The comparison fixes document embeddings and uses negatives mined by the same rule and frozen encoder, isolating state distribution from changes to indexing or negative mining.Negative pools can still differ by state text, and non-gold passages in visited histories may appear among the negatives.Source: §4.3 · p. 4

5 Experiments

The experiments test whether on-policy state training improves multi-hop retrieval across MuSiQue and longer-chain benchmarks, and whether the gain reflects recovery from the retriever’s own mistakes. The results show improved complete evidence coverage, including on harder chains, with gains tied to training on states refreshed by the current encoder.

5.1.1 Datasets.

  • The evaluation covers MuSiQue and constructed HotpotQA-Long and 2Wiki-Long benchmarks, with C@k counting only questions for which every gold passage is retrieved.The long-chain benchmarks use chains of up to five hops and larger corpora; the MuSiQue-trained encoder is evaluated on them without further training.Source: §5.1.1 · pp. 4–5, §5.1.2 · p. 5

5.2.1 Overall Results.

  • On MuSiQue, on-policy state training reaches 50.5 C@5, compared with 44.6 for teacher forcing using the same encoder and recipe; teacher forcing also exceeds Search-R1’s 36.5 C@5.The trained chain embedding retriever calls no language model at retrieval time.Source: §5.2.1 · p. 5

5.2.2 Transfer to Longer Chains.

  • 29.6 C@10 on HotpotQA-Long and 32.8 C@10 on 2Wiki-Long, up from 1.0 and 2.0 for the frozen encoder, respectively.The MuSiQue-trained encoder transfers without further training to benchmarks with longer chains and larger corpora; Search-R1 leads on 2Wiki-Long C@10.Source: §5.2.2 · p. 5

5.2.3 Results by Question Type.

  • On MuSiQue three-hop questions, on-policy training completes 41.6% of chains within five passages versus 20.9% for the next system; it leads on every HotpotQA-Long type and 2Wiki-Long comparison questions.On four-hop MuSiQue questions, every retriever completes fewer than 7% of chains.Source: §5.2.3 · p. 5, Figure 2 · p. 6

5.2.4 Downstream Question Answering.

  • With on-policy retrieval passages, the fixed reader reaches 44.1 exact match, 8.6 points above Search-R1’s passages and 11.7 above Q-RAG’s.The same reader evaluates the retrievers’ top five passages on MuSiQue.Source: §5.2.4 · p. 5, Table 2 · p. 7

5.3 RQ2: Ablations

  • Removing visited states, while keeping the encoder, recipe, and number of states fixed, reduces C@5 by 5.9 points on MuSiQue and C@10 by 5.7 and 10.4 points on HotpotQA-Long and 2Wiki-Long, respectively.On HotpotQA-Long, fixed states yield 22.7 versus 22.3 C@10 with refreshed states; on 2Wiki-Long, they yield 28.3 versus 24.3. On MuSiQue, question-only fine-tuning gives no R@5 gain over the frozen encoder (54.6 versus 55.1).Source: §5.3 · p. 6

5.4.1 Recovery on Fixed States.

  • After one wrong passage, next-gold success rises from 25.0% with teacher forcing to 35.8% with on-policy training; after two or more, it rises from 14.5% to 29.2%.On contaminated states, removing wrong passages improves teacher-forced success from 20.8% to 32.1%, but on-policy success changes only from 33.1% to 35.3%.Source: §5.4.1 · p. 7, §5.4.2 · p. 7

5.4.3 Gain along the Chain.

  • Late-hop gains accumulate: next-gold retrieval rises from 27.7% to 38.7% at step three and from 15.2% to 23.1% at step four, while query trajectories move toward gold passages.Controlled comparisons find that synthetic contamination explains about half the C@5 gap, while refreshing states with the current encoder raises C@5 to 50.5; shared negative pools preserve a 3.8-point C@5 gain over teacher forcing.Source: §5.4.3 · p. 7, §5.4.4 · p. 7, §5.5.1 · p. 8

6 Conclusion

The study finds that training chain retrievers on states they visit addresses the mismatch between gold histories in training and retriever-generated histories at test time. The gains center on recovering from contaminated histories, require refreshing visited states, and extend to longer chains.

6 Conclusion

  • On-policy training improves recovery from contaminated histories, depends on refreshing visited states as the encoder learns, and transfers to chains longer than those seen in training.The authors report these implications within the setting studied.Source: §6 · p. 9
  • Within this setting, training-state distribution is a design choice for chain retrievers, and reporting test-time distribution mismatch would make results easier to interpret.The authors also suggest dense supervision on visited states as a cheap complement to outcome rewards in retrievers inside agents and reinforcement-learning pipelines.Source: §6 · p. 9
  • A trained chain embedding retriever that needs no language model at retrieval time is a strong baseline for graph-based and agentic retrieval of complex questions.This implication is stated within the setting studied.Source: §6 · p. 9

A Training Details

The training setup adapts a Qwen3-Embedding-0.6B query encoder with LoRA and uses serialized retrieval states with a fixed state budget. On-policy training refreshes visited states during training and combines them with teacher-forced states under a shared optimization recipe.

A.1 Encoder

  • The query encoder starts from Qwen3-Embedding-0.6B and is adapted with rank-16 LoRA on the attention projections, while document embeddings remain frozen.Queries use last-token pooling and ℓ2 normalization; document embeddings are computed once in half precision.Source: §A.1 · p. 10

A.2 States and Serialization

  • Teacher-forced states use random subsets of gold passages in chain order; on-policy training keeps about 2,350 such states and fills the rest of its roughly 4,700-state budget with deduplicated visited states.State text includes the retrieval instruction, question, and history passages in retrieval order; the passages are not shortened or truncated.Source: §A.2 · p. 10

A.3 Optimization

  • Training uses batches of 16 states, 15 sampled negatives per state, and all 256 batch candidates as negatives, while masking the state’s other missing gold passages.Optimization uses τ=0.05, AdamW with learning rate 3 × 10^-4, warmup and cosine decay, and runs retrieval for four steps on MuSiQue or eight on long-chain benchmarks.Source: §A.3 · p. 10
  • Visited states are collected by the teacher-forced encoder in epoch one and by the current encoder before each later epoch’s updates; targets are the earliest missing gold passage.The checkpoint is selected by validation C@5, with a patience of three epochs.Source: §A.3 · p. 10

B Shared Negative Pools

The shared-pool control gives teacher-forced and visited states the same question-only negative pool while leaving other settings unchanged. The supplied text describes the pool-overlap comparison but does not report its outcome.

B Shared Negative Pools

  • The control draws negatives for every state from the frozen encoder’s top 100 for the question alone, with gold passages removed; sampling and in-batch negatives remain unchanged.Next-step success is measured on fixed states by the number of wrong passages in the history.Source: B Shared Negative Pools · p. 10
  • With per-state pools, a contaminated state’s pool shares 43% of its passages with the question pool, compared with 57% for a clean history.Source: B Shared Negative Pools · p. 10
  • Table 5 covers per-state versus shared negative pools on MuSiQue, with all runs trained using the same seed.Source: Table 5 · p. 10

C Mixing Ratio

On MuSiQue, every tested visited-state share improves over teacher forcing, but the mixing-ratio curve is not monotone. The equal mixture performs best, while retaining gold states preserves supervision on histories the current encoder rarely reaches.

C Mixing Ratio

  • Every tested share of visited states improves on teacher forcing, and visited states alone recover about 80% of the best mixture’s gain.The curve is not monotone.Source: C Mixing Ratio · p. 11
  • The equal mixture is best, beating visited states alone by 1.2 points.The authors retain the gold half to preserve supervision on histories the current encoder rarely reaches.Source: C Mixing Ratio · p. 11
  • Figure 5 varies the gold-to-visited-state mixing ratio on MuSiQue; 0% denotes teacher forcing on gold states.Source: Figure 5 · p. 11

D HippoRAG 2 on Reduced Corpora

HippoRAG 2 and G-Retriever are evaluated on reduced corpora for the long-chain benchmarks, and Table 6 reports R@10 and C@10 results. HippoRAG 2 sometimes falls back to dense retrieval when its triple filter finds no facts.

D HippoRAG 2 on Reduced Corpora

  • The long-chain evaluation uses reduced corpora containing all gold passages and each question’s top 100 passages under a frozen encoder.The corpora contain 27,744 passages for HotpotQA-Long and 19,257 for 2Wiki-Long.Source: D HippoRAG 2 on Reduced Corpora · p. 11
  • HippoRAG 2’s triple filter keeps no fact for 87 HotpotQA-Long and 48 2Wiki-Long questions, prompting fallback to dense retrieval.On chains of three to five hops, a single round of fact retrieval rarely reaches later passages.Source: D HippoRAG 2 on Reduced Corpora · p. 11
  • Table 6 reports R@10 and C@10 on the reduced corpora, with the authors’ results given as the five-seed mean from Table 1.Source: Table 6 · p. 11

E Long-Chain Benchmarks

The paper constructs three-to-five-hop HotpotQA-Long and 2Wiki-Long benchmarks, then evaluates retrieval, question answering, cost, beam decoding, and recovery on them. On MuSiQue, on-policy state training improves question-answering scores over teacher forcing with both readers, while the long-chain probes compare recovery and decoding behavior.

E.1 Chains

  • 2Wiki-Long chains are three to five-hop Wikidata paths with functional relations, while HotpotQA-Long extends bridge-question pages through Wikipedia hyperlinks.2Wiki-Long excludes repeated entities and degenerate relation pairs; HotpotQA-Long limits hops to named-entity pages with in-degree at most 300 and sentences with fewer than five links.Source: §E.1 · p. 11

E.2 Questions and Filters

  • Questions are model-written and filtered for answerability, uniqueness, naturalness, and absence of intermediate entities or answers; a second model re-checks answerability and uniqueness.After re-checking, 62% of the 2Wiki-Long pool and 98% of the HotpotQA-Long pool pass; questions whose final gold passage appears in the top 10 of question-only BM25 are removed.Source: §E.2 · p. 11

E.3 Corpora

  • The benchmarks include 100 test questions per hop count for three-, four-, and five-hop HotpotQA-Long, and 170, 45, and 12 chains plus 73 comparison questions for 2Wiki-Long.Their corpora contain 554,812 HotpotQA-Long passages and 220,000 2Wiki-Long passages, with development splits of 100 questions per benchmark.Source: §E.2 · p. 11, §E.3 · p. 11

F.2 Cost

  • On-policy training takes 60 minutes after seven epochs, including 32 minutes for roll-ins, versus 37 minutes for ten teacher-forced epochs.Both regimes schedule 2,940 optimizer steps; the first on-policy roll-in uses a teacher-forced encoder.Source: §F.2 · pp. 11–12

F.1 Question Answering

  • 5.5 more EM and 5.7 more F1 with deepseek-flash, and 2.5 more EM and 3.5 more F1 with qwen-flash, are reported for on-policy state training over teacher forcing.The retriever ranking is the same under both readers.Source: §F.1 · p. 11, Table 8 · p. 12

F.2 Cost

  • The chain retriever adds 4.6M trained parameters to a 0.6B encoder, calls no language model, and takes 0.10 seconds per four-step question.The graph systems make 23.3k language-model calls with 14.6M tokens to build their MuSiQue graphs; Search-R1 runs a 7B language model for about four generation turns per question.Source: §F.2 · p. 12, Table 9 · p. 12

F.3 Beam Decoding

  • Greedy on-policy decoding outperforms teacher-forced decoding with a beam of four, which encodes three times as many queries.Beams improve teacher forcing by two to three points of complete coverage but leave a four-to-six-point gap.Source: §F.3 · p. 12, Table 10 · p. 12

F.4 Recovery on the Long-Chain Benchmarks

  • On teacher-forced states of the long-chain benchmarks, the three encoders are close with clean histories, while on-policy succeeds about twice as often with two or more wrong passages.The probe uses T=8.Source: §F.4 · p. 12

F.5 Error Propagation

On three- and four-hop MuSiQue questions, on-policy state training is associated with more recovery after a non-gold retrieval and more complete chains. On contaminated states, its query is more often closer to the missing target than to a wrong passage in the history.

F.5 Error Propagation

  • 33.7% versus 23.0% of steps after a non-gold retrieval find a missing gold passage under on-policy state training and teacher forcing, respectively.Across 482 questions, chains are complete after four steps for 87 on-policy questions and 68 teacher-forced questions; Figure 6’s dark bands mark recoveries from a miss to a hit.Source: §F.5 · pp. 12–13, Figure 6 · p. 13

F.6 Where Contaminated Queries Point

  • 30.6% versus 20.3% of 1,358 contaminated states place the query closer to the target than to a wrong passage in its history under on-policy state training and teacher forcing, respectively.Figure 7 compares query cosine similarity to the wrong history passage and the target.Source: Figure 7 · p. 13, §F.6 · p. 13

Related papers

Loading 2610.40409v1…