S e RTS : Self-Rewarding Tree Search for Biomedical Retrieval-Augmented Generation (2024)

Abstract: Large Language Models (LLMs) have shown great potential in the biomedical domain with the advancement of retrieval-augmented generation (RAG). However, existing retrieval-augmented approaches face challenges in addressing diverse queries and documents, particularly for medical knowledge queries, resulting in sub-optimal performance. To address these limitations, we propose a novel plug-and-play LLM-based retrieval method called Self-Rewarding Tree Search (SeRTS) based on Monte Carlo Tree Search (MCTS) and a self-rewarding paradigm. By combining the reasoning capabilities of LLMs with the effectiveness of tree search, SeRTS boosts the zero-shot performance of retrieving high-quality and informative results for RAG. We further enhance retrieval performance by fine-tuning LLMs with Proximal Policy Optimization (PPO) objectives using the trajectories collected by SeRTS as feedback. Controlled experiments using the BioASQ-QA dataset with GPT-3.5-Turbo and LLama2-7b demonstrate that our method significantly improves the performance of the BM25 retriever and surpasses the strong baseline of self-reflection in both efficiency and scalability. Moreover, SeRTS generates higher-quality feedback for PPO training than self-reflection. Our proposed method effectively adapts LLMs to document retrieval tasks, enhancing their ability to retrieve highly relevant documents for RAG in the context of medical knowledge queries. This work presents a significant step forward in leveraging LLMs for accurate and comprehensive biomedical question answering.

## introduction large language models (llms) have made significant advancements and gained widespread adoption in various natural language tasks, such as dialogue systems @xcite and question-answering (kamalloo et al., * the first two authors have equal contributions. 2023; @xcite . however, their performance in specialized domains like medicine remains unsatisfactory due to the lack of publicly available training data, which limits the amount of domain-specific knowledge these llms can encode in their parameters @xcite . to address this limitation, recent studies have turned to retrieval-augmented generation (rag) frameworks @xcite @xcite . these frameworks can enhance llms' capabilities by retrieving relevant documents from large and private databases, thereby enabling the models to provide more accurate and informative answers. existing rag framework mainly consists of two modules: an external retriever to retrieve documents according to semantic similarity between documents and queries, and the reader to generate the final answer given retrieved documents and query @xcite . however, this kind of setting overlook the two significant challenges in medical domain: the complexity of the queries and the varying importance of different documents. for example, patients may describe the same disease in diverse ways, leading to variations in query formulation. additionally, documents from different medical disciplines should be weighted differently based on both of their relevance and importance. thus, effectively identifying the most pertinent and useful documents from a vast corpus of medical literature is crucial. this not only enhances performance but also improves the explainability and credibility of the system. to address these issues, we propose a plugand-play llm-based retrieval method, self-rewarding tree search (serts), combines the exceptional reasoning capabilities of llms with the effective searching capabilities of monte carlo tree search (mcts). specifically, we first formulate the task of document retrieval as a tree search problem, considering both the diversity of queries and different relationship between multiple documents. inspired by recent studies using mcts to search the optimal reasoning paths of llms, we alternativelly employ it to find the best query and most related documents from vast candidate spaces. to ensure the success of mcts, we carefully design each operations in mcts and provide detailed feedback from different levels of documents. furthermore, to validate the effectivenss and quality of our tree search retriever, we then collect these feedbacks with all intermediate reasoning results to finetune llms using ppo method. to sum up, our contributions can be summarized as follows: ## conclusion we propose self-rewarding tree search (serts), a plug-and-play retrieval method that combines llm reasoning with conventional retrievers to address limitations in biomedical retrieval-augmented approaches. serts introduces a reasoning process based on monte carlo tree search (mcts) and a self-rewarding paradigm, boosting zero-shot performance in retrieving high-quality, informative results for retrieval augmented generation (rag). we further enhance retrieval by fine-tuning llms with ppo objectives using serts retrieval trajectories as feedback. experiments using the bioasq-qa dataset with gpt-3.5-turbo and llama2-7b show that serts significantly improves bm25 retriever performance, surpassing the self-reflection baseline in efficiency and scalability. serts (ppo) effectively adapts llms to document retrieval, enhancing their ability to find highly relevant documents for rag when answering medical queries. our work presents an effective method for accurate, comprehensive medical knowledge retrieval using llms. while serts outperforms self-reflection, both require numerous simulations, limiting widespread application. future research should focus on developing more efficient reasoning strategies and leveraging the self-improvement capabilities of llms to further enhance retrieval efficiency and usability.