Treinamento de Sistema RAG Agente de Ponta a Ponta para Raciocínio Diagnóstico Rastreável

End-to-End Agentic RAG System Training for Traceable Diagnostic Reasoning

August 21, 2025

Autores: Qiaoyu Zheng, Yuze Sun, Chaoyi Wu, Weike Zhao, Pengcheng Qiu, Yongguo Yu, Kun Sun, Yanfeng Wang, Ya Zhang, Weidi Xie

cs.AI

Resumo

O diagnóstico preciso com modelos de linguagem de grande escala (LLMs) na área médica é prejudicado por lacunas de conhecimento e alucinações. Métodos de recuperação e aumento com ferramentas ajudam, mas seu impacto é limitado pelo uso fraco de conhecimento externo e pela baixa rastreabilidade do raciocínio com feedback. Para enfrentar esses desafios, introduzimos o Deep-DxSearch, um sistema RAG (Retrieval-Augmented Generation) agente treinado de ponta a ponta com aprendizado por reforço (RL) que permite um raciocínio aumentado por recuperação rastreável para diagnóstico médico. No Deep-DxSearch, primeiro construímos um corpus de recuperação médica em larga escala, composto por registros de pacientes e fontes confiáveis de conhecimento médico, para apoiar o raciocínio consciente da recuperação em diversos cenários diagnósticos. Mais crucialmente, enquadramos o LLM como o agente central e o corpus de recuperação como seu ambiente, utilizando recompensas personalizadas em formato, recuperação, estrutura de raciocínio e precisão diagnóstica, evoluindo assim a política RAG agente a partir de dados em larga escala por meio do RL. Experimentos demonstram que nossa estrutura de treinamento agente RL de ponta a ponta supera consistentemente abordagens de engenharia de prompt e RAG sem treinamento em múltiplos centros de dados. Após o treinamento, o Deep-DxSearch alcança ganhos substanciais em precisão diagnóstica, superando bases fortes de diagnóstico como GPT-4o, DeepSeek-R1 e outras estruturas específicas para medicina, tanto para diagnósticos de doenças comuns quanto raras, em cenários de distribuição interna e externa. Além disso, estudos de ablação sobre o design de recompensas e componentes do corpus de recuperação confirmam seus papéis críticos, destacando a singularidade e eficácia de nossa abordagem em comparação com implementações tradicionais. Por fim, estudos de caso e análises de interpretabilidade destacam melhorias na política diagnóstica do Deep-DxSearch, fornecendo insights mais profundos sobre seus ganhos de desempenho e apoiando clínicos na entrega de diagnósticos preliminares mais confiáveis e precisos. Consulte https://github.com/MAGIC-AI4Med/Deep-DxSearch.

English

Accurate diagnosis with medical large language models is hindered by knowledge gaps and hallucinations. Retrieval and tool-augmented methods help, but their impact is limited by weak use of external knowledge and poor feedback-reasoning traceability. To address these challenges, We introduce Deep-DxSearch, an agentic RAG system trained end-to-end with reinforcement learning (RL) that enables steer tracebale retrieval-augmented reasoning for medical diagnosis. In Deep-DxSearch, we first construct a large-scale medical retrieval corpus comprising patient records and reliable medical knowledge sources to support retrieval-aware reasoning across diagnostic scenarios. More crutially, we frame the LLM as the core agent and the retrieval corpus as its environment, using tailored rewards on format, retrieval, reasoning structure, and diagnostic accuracy, thereby evolving the agentic RAG policy from large-scale data through RL. Experiments demonstrate that our end-to-end agentic RL training framework consistently outperforms prompt-engineering and training-free RAG approaches across multiple data centers. After training, Deep-DxSearch achieves substantial gains in diagnostic accuracy, surpassing strong diagnostic baselines such as GPT-4o, DeepSeek-R1, and other medical-specific frameworks for both common and rare disease diagnosis under in-distribution and out-of-distribution settings. Moreover, ablation studies on reward design and retrieval corpus components confirm their critical roles, underscoring the uniqueness and effectiveness of our approach compared with traditional implementations. Finally, case studies and interpretability analyses highlight improvements in Deep-DxSearch's diagnostic policy, providing deeper insight into its performance gains and supporting clinicians in delivering more reliable and precise preliminary diagnoses. See https://github.com/MAGIC-AI4Med/Deep-DxSearch.