Triagerag improves domain adaptation for retrieval-augmented language models

TriageRAG: Failure-Aware Self-Training for Adapting Retrieval-Augmented Language Models to Specialized Domains

Computation and LanguageArtificial IntelligenceInformation RetrievalMachine Learning

Summary

Language models often use external information to answer difficult questions. The authors noticed that existing methods don’t check if the model already knows an answer or where it goes wrong when learning from its own generated examples. They created TriageRAG, a system that tests the model’s answers with and without extra information to figure out where it fails and then trains it on exactly those problems. This approach helps language models adapt better to specialized topics like medicine and outperforms previous methods in many tests.

What this means in practice

  • For medical ai developers: Adapt language models to medical texts using targeted self-training that improves accuracy on specialized queries.
  • For legal tech teams: Improve domain-specific question answering systems by identifying and addressing failures in retrieval and reading in specialized corpora.

Authors

Samarth Varshney, Ritu Yadav

Abstract

Retrieval-augmented generation (RAG) lets large language models (LLMs) answer knowledge-intensive questions from retrieved evidence, but adapting a general-purpose RAG system to a specialized domain such as medicine is difficult when only an unlabeled corpus is available. Self-training approaches such as SimRAG address this by letting the LLM write its own question-answer pairs from domain passages and fine-tuning on the pairs that survive a round-trip consistency filter. We observe that these approaches treat every retained example identically: they never ask whether the student already knows the answer, already extracts it correctly from the retrieved context, or fails - and if so, whether the failure lies with the reader or with the retriever. To tackle this, we propose TriageRAG, a failure-aware self-training approach that probes the student on each self-generated question with and without retrieved context, triages it into one of four categories (known, solved, context-failure, retrieval-failure), and routes each category to the training signal that addresses it, including a parametric route that teaches the reader to answer when retrieval misses. Probing is cheap and repeated every round, so the training mixture follows the student, and the triage histogram doubles as a label-free diagnostic of where the domain gap lies. Experiments on 11 datasets spanning three domains with a Qwen3-4B backbone, and on CS-Bench with Qwen3-8B, demonstrate that TriageRAG outperforms a budget-matched SimRAG by 1.2-6.4 points and is the only method that recovers the closed-book level in every case where retrieval hurts the base model. A single ablation identifies the mechanism: round-trip filtering removes every training context that lacks the evidence, and a reader that never sees such contexts learns to over-trust noisy retrieval.