IryoNLP at MEDIQA-CORR 2024: Tackling the Medical Error Detection & Correction Task On the Shoulders of Medical Agents
Jean-Philippe Corbeil


Deep-learning Optical Flow Outperforms PIV in
Obtaining Velocity Fields from Active Nematics
    Correspondence email addresses:
^∗hongpeng@brandeis.edu, ^hagan@brandeis.edu 

    April 29, 2024
============================================================================================




In natural language processing applied to the clinical domain, utilizing large language models has emerged as a promising avenue for error detection and correction on clinical notes, a knowledge-intensive task for which annotated data is scarce. This paper presents MedReAct'N'MedReFlex, which leverages a suite of four LLM-based medical agents. The MedReAct agent initiates the process by observing, analyzing, and taking action, generating trajectories to guide the search to target a potential error in the clinical notes. Subsequently, the MedEval agent employs five evaluators to assess the targeted error and the proposed correction. In cases where MedReAct's actions prove insufficient, the MedReFlex agent intervenes, engaging in reflective analysis and proposing alternative strategies. Finally, the MedFinalParser agent formats the final output, preserving the original style while ensuring the integrity of the error correction process. One core component of our method is our RAG pipeline based on our ClinicalCorp corpora. Among other well-known sources containing clinical guidelines and information, we preprocess and release the open-source MedWiki dataset for clinical RAG application. Our results demonstrate the central role of our RAG approach with ClinicalCorp leveraged through the MedReAct'N'MedReFlex framework. It achieved the ninth rank on the MEDIQA-CORR 2024 final leaderboard. 




§ INTRODUCTION




In natural language processing applied to the clinical domain, the accurate detection and correction of medical errors are paramount tasks with profound implications for patient care and safety. This paper introduces the multi-agent framework MedReAct'N'MedReFlex, meticulously engineered to tackle medical error detection and correction, as delineated in the MEDIQA-CORR 2024 competition.

Our framework integrates four distinct types of medical agents: MedReAct, MedReFlex, MedEval, and MedFinalParser, each playing a specialized role in the error identification and rectification process. Drawing inspiration from existing frameworks like ReAct <cit.> and Reflexion <cit.>, our framework orchestrates a structured approach to error handling.

Leveraging a Retrieval-Augmented Generation (RAG) framework <cit.> based on MedRAG <cit.> and MedCPT <cit.>, our approach operates over ClinicalCorp, an extensive corpora curated to encompass crucial clinical guidelines. Additionally, we introduce MedWiki, a collection of medical articles from Wikipedia. By integrating these resources, our approach seeks to advance state-of-the-art clinical NLP by offering a comprehensive solution tailored to the intricate nuances of medical error handling. Furthermore, this paper documents the construction and release of MedWiki, a substantial repository comprising over 1.3 million article chunks. Additionally, we detail the assembly of the ClinicalCorp, a comprehensive corpus comprising MedWiki along with other clinical guideline datasets, such as parts of the MedCorp corpora <cit.> and parts of the guidelines <cit.>.

Our main contributions are:


    
  * We designed a multi-agent framework named MedReAct'N'MedReFlex to solve the medical error detection & correction task (MEDIQA-CORR 2024) based on four types of medical agents: MedReAct, MedReFlex, MedEval and MedFinalParser. We deployed this framework on ClinicalCorp using a retrieval-augmented generation approach.
    
  * We released the open-source MedWiki[<hf.co/datasets/jpcorb20/medical_wikipedia>], a version of Wikipedia 2022-12-22 focused solely on medical articles. This RAG-ready dataset contains about 1.3M chunks from more than 150K articles, which represents about 3% of the original corpus.
    
  * We provided the recipe to assemble our large corpora ClinicalCorp for RAG applications in the clinical domain, containing more than 2.3M chunks.
    
  * We released a RAG-ready version[<hf.co/datasets/jpcorb20/rag_epfl_guidelines>] of the open-source guidelines used to pre-train Meditron <cit.>, containing more than 710K chunks across eight open-source datasets.
    
  * We released our codebase on GitHub[<github.com/microsoft/iryonlp-mediqa-corr-2024>].




§ RELATED WORK




 §.§ Medical Large Language Models


Since the emergence of ChatGPT by OpenAI in December 2022, the landscape of large language models (LLMs) has witnessed a proliferation of both private and public initiatives, leading to the development of increasingly sophisticated models. OpenAI's journey from the GPT3.5-turbo architecture, as reported by <cit.> and <cit.>, culminated in the release of GPT-4 and its turbo variant <cit.>. Similarly, Google introduced Gemini, available in Nano, Pro, and Ultra configurations <cit.>, alongside its open-source Gemma model <cit.>. Anthropic contributed to this landscape with Claude3, offered in three sizes, ranging from Haiku to Opus. Other notable LLMs include Mistral and Mixtral <cit.>, as well as Llama 2 <cit.> and Yi <cit.>. These general-purpose LLMs, such as GPT-4, have demonstrated solid in-context learning capabilities in the medical field <cit.>.

Researchers have developed various open-source LLMs with diverse capabilities in the medical NLP domain. Examples include ClinicalCamel <cit.>, Med42 <cit.>, PMC-Llama <cit.>, BioMedGPT <cit.>, Meditron <cit.>, Apollo <cit.>, OpenMedLM <cit.>, and BioMistral <cit.>. Google also contributed Med-PaLM 2, a specialized LLM tailored for medical tasks <cit.>.

In this study, we employed OpenAI's GPT-4, specifically version turbo 0125, due to its proven state-of-the-art performances in various domains, its functional capabilities, and its large context window of 128K tokens. These attributes make it an ideal foundation for our approach. For instance, <cit.> demonstrated that utilizing in-context learning with GPT-4 — relying on prompt engineering (i.e. few-shot learning <cit.>, chain-of-thought <cit.>, self-consistency <cit.> and shuffling multiple choice <cit.>) — achieves state-of-the-art performances on medical question-answering tasks, surpassing specialized models like Med-PaLM 2. We relied on a similar approach as our early baseline for medical error detection and correction, discarding the self-consistency and the shuffling techniques since both do not apply to generative tasks. Nonetheless, we have observed low results from which we hypothesized that this approach using only parametric knowledge is lacking reliable knowledge <cit.>, which we addressed by applying agentic methods in a retrieval-augmented generation framework.



 §.§ Agentic Methods


Researchers have devised several agentic methods to enhance LLMs' responses and reasoning capabilities, such as ReAct <cit.>, Reflexion <cit.>, DSPy <cit.> and self-discovery <cit.>. Additionally, multi-agent paradigms <cit.> have found application in the medical domain <cit.>. Our approach draws inspiration from the Reflexion framework <cit.>, which we adapted into our MedReFlex agent. Specifically, we implemented a MedReAct agent  — inspired by the ReAct approach <cit.> — to generate trajectories in our environment. However, this agent realizes its sequence of actions in a different order (i.e., observation, thought, and action), enabling streamlined execution.

Given the reliance of the Reflexion framework on feedback mechanisms, we incorporated an LLM-based metric into our MedEval medical agents. Evaluation metrics based on prompting strong LLMs, such as GPT-4 <cit.>, have demonstrated a high correlation with human judgment. Similar findings have been reported in the medical NLP literature <cit.>. Our evaluation protocol involves prompting five GPT-4 reviewers with task-specific criteria: validity, preciseness, confidence, relevance, and completeness. The average and minimum of their scores are both utilized as success criteria, capturing an unbiased final score and the evaluators' confidence, respectively.



 §.§ Retrieval-Augmented Generation


Before the advent of LLMs, authors have proposed the retrieval-augmented generation (RAG) framework as a mechanism to incorporate non-parametric memory for knowledge-intensive tasks. This framework, as elucidated by Lewis et al. <cit.>, leverages both sparse <cit.> and dense <cit.> retrieval methods. In the medical NLP domain, MedCPT <cit.> serves as a prominent retrieval approach, augmented by a reranking stage based on a cross-encoder model. Notably, Xiong et al. <cit.> conducted a comprehensive study on RAG applications in the medical domain, culminating in developing the MedRAG framework and the MedCorp corpora. Our approach builds upon these foundations, employing the MedCPT retrieval techniques and two corpora from MedCorp.

A pivotal aspect of RAG is its search engine's collection of indexed documents. The guidelines corpora, part of the GAP-replay corpora, was curated to train Meditron <cit.>. This corpus comprises web pages describing medical guidelines from reputable healthcare websites like the World Health Organization. The StatPearls and Textbooks datasets, included in the MedCorp corpora used in MedRAG <cit.>, encompass documents from clinical decision support tools and medical textbooks <cit.>. While Wikipedia and PubMed datasets within MedCorp offer extensive data (i.e. more than 55M documents), we opted for efficiency by focusing on the smaller PubMed subset in the guidelines corpora and our MedWiki corpus.



§ METHODOLOGY




 §.§ MEDIQA-CORR Task


The goal of the medical error detection and correction task <cit.> from the clinical note is threefold: detect the presence of an error, locate the sentence containing the error and generate a corrected version of that sentence. The input of the dataset <cit.> is a clinical note of several sentences containing a medical description of a patient's condition, test results, diagnosis, treatment and other aspects. There are two parts for the validation and test sets: MS from Microsoft and UW from the University of Washington. As a primary evaluation metric, the organizers asked to utilize the aggregation score defined by <cit.> over Rouge-1, BertScore and BLEURT, demonstrating a higher correlation with human judgement.







 §.§ ClinicalCorp Corpora


Our corpus is detailed in Table <ref>.



  
guidelines We aggregated 13 datasets — which are open-source or closed-source — from the guidelines corpora. We adapted and ran the scrappers from the Meditron GitHub repository to gather the closed-source datasets. Then, we chunked the resulting documents using LangChain's recursive-character text splitter <cit.> with a chunk size of 1,000 characters and an overlap of 200 characters, as used for StatPearls (see next section).



  
MedCorp We gathered two of the four datasets contained in MedCorp from MedRAG <cit.>: StatPearls and Textbooks. The former was downloaded, cleaned and chunked using MedRAG GitHub repository, while the latter was readily available on the HuggingFaceHub[<hf.co/datasets/MedRAG/textbooks>].



  
MedWiki We filtered the 2022-12-22 Wikipedia dump[<hf.co/datasets/Cohere/wikipedia-22-12>] pre-processed into chunks by Cohere for medical articles only. To select the medical articles, we leveraged an available fine-tuned BerTopic[<hf.co/MaartenGr/BERTopic_Wikipedia>] <cit.>, trained on the same Wikipedia dump. We associated its 2,3K topics to the medical domain based on the topics' word representations — e.g. topic 1850 is related to the medical field, and it corresponds to the word representations: shingles, herpesvirus, chickenpox, herpes, smallpox, zoster, immunity, infectivity, inflammation, and viral. We made these predictions by prompting GPT3.5-turbo 0613 with a temperature of 1.0 followed by a majority vote over five predictions. If at least four were positive, we declared the topic medically relevant. In the manual verification of about 50 diverse medical terms on the resulting collection, we observed a near-perfect coverage of Wikipedia's articles related to diseases, treatments, bacteria, or drugs. Only two topics were missing[Index 509 related to biological taxonomy and 806 related to yeasts.], corresponding to one single example from the manual test. Given that our goal is to reduce the size of this dataset and use it in an RAG application, we added these topics manually. We obtained a corpus of 150K articles and nearly 1.4M chunks.



 §.§ Semantic Search


We followed the MedCPT approach <cit.> in two stages (see step B in Figure <ref>), which is composed of a fast bi-encoder retrieving stage followed by a cross-encoder reranking stage.

We implemented the first stage on a ChromaDB instance, in which we loaded ClinicalCorp. This stage aims to find relevant documents while maintaining a good accuracy/latency trade-off. This vector database embeds documents using a fast bi-encoder model <cit.>. Then, we provide a query to fetch the closest documents under a given distance, computed with the hierarchical navigable small world approximation (HNSW, by <cit.>). We experimented with three bi-encoders from the HuggingFaceHub: sentence-transformers/all-MiniLM-L6-v2 (default), NeuML/pubmedbert-base-embeddings-matryoshka and MedCPT original Query/Article encoders. According to our initial experiments, we discarded all-MiniLM-L6-v2 because we noticed a critical lack of knowledge about medical terminology hindering its accuracy despite a very low latency. NeuML's model and MedCPT's are Bert-based models of 768 hidden dimensions and 12 layers, a slow architecture to generate sentence embeddings. However, NeuML fine-tuned a recent model using the Matryoshka Representation Learning technique <cit.>, allowing to truncate dimensions down to 256 dimensions of the 768 embeddings, which significantly accelerated the computations. Our experiments employ this MRL encoder with truncation at 256 dimensions as a trade-off between accuracy and latency.

We implemented the reranking stage following the cross-encoder approach from MedCPT <cit.>. Our early experimentation demonstrated the superiority of this model compared to NeuML's MRL bi-encoder with all 768 dimensions as a reranker.



§ MEDREACT'N'MEDREFLEX FRAMEWORK


Unlike previous multi-agent frameworks <cit.>, our approach diverges from a free conversation format to adopt a fixed design schema, as illustrated in Figure <ref>. Within this structured framework, each medical agent intervenes at a specific step, facilitating a systematic and coordinated approach to address the error detection and correction task. Central to our methodology are four distinct medical agents: MedReAct, MedReFlex, MedEval, and MedFinalParser.



 §.§ MedReAct Agent

The MedReAct agent (see step A in Figure <ref>), inspired by the ReAct framework <cit.>, operates cyclically, beginning with an observation of the current context, followed by a thoughtful analysis, and concluding with an action (search or final_mistake). This agent generates a trajectory of up to N steps if the action is always a search with different queries. 

We also experimented with adding two other actions (get_doc_by_id and next_results_from_query), but MedReAct systematically ignored them.



 §.§ MedEval Agent

Upon MedReAct's selection of the final_mistake action, the MedEval agents (see step D in Figure <ref>), akin to the GPT-Eval approach <cit.>, evaluate the proposed solution. Five GPT-4-based evaluators assess the answer based on criteria such as validity, preciseness, confidence, relevance, and completeness. The ensemble of evaluators ensures comprehensive and unbiased feedback, contributing to robust error detection and correction. We leverage the average final score as well as the minimum review score. We added this condition on the minimum score to capture the confidence of the evaluation. If one reviewer gave a much lower score than the others, we experimentally observed that it was often a signal of lower confidence in the final answer.



 §.§ MedReFlex Agent

In scenarios where MedReAct's actions fail to yield satisfactory outcomes, the MedReFlex agent (see step C in Figure <ref>), drawing from the Reflexion framework <cit.>, intervenes. This agent engages in reflective analysis to reassess the situation. By considering contextual cues, past interactions and all five reviews, MedReFlex proposes alternative strategies to address the identified challenges. This iterative process allows for adaptive decision-making and fosters resilience in error detection and correction tasks.



 §.§ MedFinalParser Agent

Suppose the average score provided by the MedEval agents exceeds or equals 4, and the minimum score surpasses or equals 3. In that case, the MedFinalParser agent (see step D in Figure <ref>) proceeds to format the final answer into a JSON object. This agent also ensures the conservation of the original style of the clinical note, which the MedReAct agent tends to disrupt by copying the writing style of the search documents. Conversely, if the answer falls short of the predetermined thresholds, MedReFlex is triggered for further refinement. If MedReFlex's interventions prove ineffective after the M^th turn, the MedFinalParser agent concludes that no errors exist, ensuring the integrity of the error correction process.



§ RESULTS




 §.§ Results for the Competition


MedReAct'N'MedReFlex achieved the 9^th rank during the MEDIQA-CORR 2024 official testing period, corresponding to an aggregation score of 0.581. Nonetheless, we thoroughly optimize our method in the following sections. To complete these experiments in a reasonable amount of time, we randomly sample 50 examples from the MS validation set.



 §.§ Agentic Method Comparison


In Table <ref>, we compared the MedReAct agent only against using our proposed method MedReAct'N'MedReFlex. Our approach achieves more than a few absolute percent across metrics. We also experimented with a baseline inspired from <cit.> (i.e. "MedPrompt") with in-context learning prompting alone, but the results were drastically lower.





 §.§ Semantic Search Optimization


After the end of the MEDIQA-CORR 2024 shared task, we carried out a thorough analysis of our semantic search engine. The main parameters to tune are retrieval top-k, reranking top-k and the source included in ClinicalCorp.



  §.§.§ Retrieval Top-K


In Figure <ref>, we illustrate the performances across many retrieval top-k values employing a fixed reranking top-k of 20. For the official ranking of the MEDIQA-CORR 2024, we set this value to 300. However, we observe here that this setting is sub-optimal. A retrieval top-k of 50 improves the final performances by a few absolute percent. We interpret this observation as indicative of a misalignment between our task and the fine-tuning of the MedCPT reranker. The more documents we provide to the reranking model (e.g. 200 or 300), the more low-relevance documents are put in the top 20 by the reranker output.

Nonetheless, a reranking without surplus documents — i.e. retrieval top-k of 20 with a reranking top-k of 20 — remains sub-optimal, mainly in contrast to using 50 documents. In Figure <ref>, we provide the associated average latency for one react step in seconds. We notice that the latency seems to scale with the order of magnitude of the retrieval top-k, with a value of 20 and 50 having 17 seconds on average, while 100, 200 and 300 are around 20 seconds. We expected that the reranking of 300 examples against 100, for instance, would lead to noticeable latency, but it is negligible in contrast to the retrieval from ChromaDB over our 2.3M chunks.







We also show the average amount of MedReFlex turns and the average of the sum of all MedReAct turns in Figure <ref>. Overall, the trends are similar, with 4.8 total ReAct turns on average, but there is a slight increase in the average and variance for top-k values of 200 and 300. Therefore, these settings are underperforming and slower regarding latency and the number of turns needed to reach an answer.

Overall, the retrieval top-k of 50 leads to higher performances across all metrics and reduced latency and number of turns required by our algorithm.



  §.§.§ Reranker Top-K


In Figure <ref>, we fix the retrieval top-k at 300 and compute the performances across three reranking top-k values: 5, 10 and 20. Since the context window of the LLM limits us, we constraint the maximum of K to 20, given that these K documents are injected in the prompt up to N times for each MedReAct step. According to Figure <ref>, we observe that the more documents we provide in the prompt, the more we increase the performances — the aggregate score gains close to 10% absolute when augmenting from 5 to 20 documents.





  §.§.§ Sources in ClinicalCorp


We measure the impact of each source in ClinicalCorp in Figure <ref>. First, we observe that MedWiki is the lowest-performing source of documents with an aggregation score of nearly 0.47. guidelines and Textbooks provide a similar accuracy at about 0.51 in aggregate score. Finally, StatPearls leads to the highest score close to the full ClinicalCorp. Given our small validation set of 50 examples, we consider it a better practice to keep all ClinicalCorp for our task since more edge cases might appear at test time.



We show in Figure <ref> the distributions of sources from ClinicalCorp in general in comparison to the distributions of sources' chunks used by one run of MedReAct'N'MedReFlex. We observe, in general, a much larger utilization of the StatPearls' chunks in contrast to the MedWiki's chunks, while we remark similar distributions for the other two datasets. These results align with the previous analysis demonstrating a higher performance from using only StatPearls.





 §.§ MedEval Aggregation Thresholds


In Figure <ref>, we show the impact of applying different thresholds to the average and minimum review scores on the performance. For the minimum score criterion, we choose the integer values of 2.0, 3.0 and 4.0. We select values for the average score criterion: 3.0, 3.2, 3.5, 3.8, 4.0 and 4.2. We do not compute the performances for combinations where the minimum threshold is higher than the average threshold for mathematical consistency. We observe an optimal setting for a minimum evaluation score of 3.0 with a range of average evaluation scores in [3.5, 3.8].





§ CONCLUSION


In this paper, we introduced MedReAct'N'MedReFlex, a multi-agent framework developed for the MEDIQA-CORR 2024 competition aimed at medical error detection and correction in clinical notes. The framework incorporates four specialized medical agents: MedReAct, MedReFlex, MedEval, and MedFinalParser, leveraging the RAG framework and our ClinicalCorp. We detail the construction of our ClinicalCorp, including diverse clinical datasets such as guidelines, Textbooks, and StatPearls. Additionally, we released MedWiki, a corpus comprising Wikipedia medical articles. Our framework achieved the ninth rank in the competition with an aggregation score of 0.581. Through optimization experiments, we identified sub-optimal settings at the time, demonstrating substantial performance improvements with a retrieval top-k of 50, a reranking top-k of 20, an average review threshold of 3.8, and a minimum review threshold of 3. As future work, we envision refining the chunking strategy on the ClinicalCorp, applying further prompt engineering of the medical agents, and conducting a deeper analysis of the interactions between the MedReAct'N'MedReFlex's agents.


§ ACKNOWLEDGMENTS


We would like to thank Jean-Michel Attendu, François Beaulieu, and Paul Vozila from Microsoft Health & Life Sciences team for their support, as well as the ClinicalNLP workshop's organizers, reviewers and other participants. 



