💻 computer science

Improving Rare Medication Recommendation with Counterfactual Data Augmentation and Large Language Models

The paper introduces GenRxR, a novel framework that leverages large language models to generate counterfactual medical data and model co-recommended medication relationships, thereby significantly improving the predictive performance of rare medication recommendations compared to existing baselines.

Original authors: Shinhwan Kang, Soo Yong Lee, Jaewon Kim, Kijung Shin, Buru Chang

Published 2026-07-29

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Shinhwan Kang, Soo Yong Lee, Jaewon Kim, Kijung Shin, Buru Chang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a world where a super-smart robot doctor helps real doctors decide which medicines to give patients. This isn't science fiction; it's a growing field called AI-based medication recommendation. These systems look at a patient's history—their past illnesses, surgeries, and what they've taken before—to guess the best new medicine. Think of it like a very advanced recipe app that knows your taste buds and allergies better than you do. But here's the catch: just like a chef who only cooks with common ingredients like salt and sugar, these AI robots have a hard time with the "rare" ingredients. In the medical world, these are medicines given to very few people, often for strange or severe conditions. Because the AI has never seen these rare medicines very often in its training data, it tends to ignore them completely, which can be dangerous for the patients who desperately need them.

This is the problem a team of researchers from KAIST and Korea University tackled in their new paper. They noticed that while AI is great at suggesting common drugs, it often fails miserably at suggesting the rare ones. To fix this, they built a new system called GenRxR. Instead of just waiting for more real-world data (which is hard to get because of privacy rules), they used a "time-traveling" trick with Large Language Models (LLMs)—the same kind of super-intelligent AI that can write stories or chat with you. They asked the AI to imagine "what if" scenarios: "What if this patient had been given this rare medicine? How would their story change?" By generating these fake but medically logical stories, they created a massive new library of examples for the AI to learn from. They also taught the AI to pay closer attention to how different medicines work together, like a conductor ensuring every instrument in an orchestra plays in harmony.

The result? Their new system, GenRxR, became a much better guesser for those tricky, rare medicines. In their tests, it didn't just do okay; it crushed the competition. It outperformed 14 other methods, including other AI systems that use similar big-brain models. Most impressively, when it came to the rare medicines that other systems often missed entirely, GenRxR improved its prediction accuracy by up to 30.9%. They didn't just guess; they ran the numbers and showed that by using these "what-if" stories and teaching the AI to reason like a doctor, they could make medication recommendations safer and more accurate for everyone, even the ones with the rarest conditions.

The Problem: The AI's "Common Sense" Bias

Imagine you are trying to learn to play a video game, but you only get to play the first level a million times. You become a master of that level, but if someone suddenly asks you to play the secret, hidden boss level that only appears once in a blue moon, you have no idea what to do. This is exactly what happens with current AI medication systems.

The researchers found that these systems suffer from a "long-tail" problem. In the real world, a few medicines are prescribed constantly (the "common" drugs), while thousands of others are prescribed very rarely (the "rare-meds"). Because the AI learns from historical hospital records, it sees the common drugs over and over again. It gets really good at them. But the rare drugs? They are like ghosts in the data. There are so few examples that the AI basically forgets they exist.

This isn't just a small glitch. Many of these rare medicines are "orphan drugs," designed to treat very specific, severe, or uncommon diseases. If the AI ignores them, a patient in need might be left without a treatment plan. The researchers also noticed that existing AI systems often fail to understand how medicines work together. Prescribing a rare drug usually means it's part of a specific cocktail of other treatments. If the AI doesn't understand the relationships between these drugs, it can't make a safe recommendation.

The Solution: Time-Traveling with "What-If" Stories

To fix this, the team created GenRxR, a framework that uses Large Language Models (LLMs) in two clever ways.

1. The "What-If" Machine (Counterfactual Data Augmentation)
The first big hurdle was the lack of data. You can't teach a student to drive a Ferrari if you only have one Ferrari in the world. So, the researchers used an LLM to create more Ferraris. They didn't just make up random numbers; they used the AI's medical knowledge to write "counterfactual" stories.

Here is how it works:

The team picks a patient from a real hospital record who didn't get a specific rare medicine.
They ask the LLM: "Imagine this patient did get this rare medicine. What else would have happened in their medical history to make that logical?"
The LLM acts like a creative writer with a medical degree. It might say, "Okay, if they got this rare drug, they probably also had a specific type of infection, underwent a certain surgery, and were taking a different set of supporting drugs."
The AI then generates a brand new, synthetic patient record that includes this rare medicine and all the logical changes that would go with it.

This is like giving the AI a time machine to visit alternate realities where rare medicines were used. By training on these thousands of new, realistic "what-if" stories, the AI finally learns what rare medicines look like and when to use them.

2. The "Team Player" (Instruction Tuning)
The second part of the puzzle was making sure the AI understood the context. A rare medicine is rarely a solo act; it's part of a team. To help the AI understand this, the researchers used a technique called instruction tuning.

They taught the LLM to summarize medical records in a way that highlights the connections between a patient's condition and the medicines they need. They gave the AI specific instructions to act like a clinical reasoning expert. Instead of just looking at a list of drugs, the AI learned to look at the whole picture: "This patient has these symptoms, so this specific combination of drugs makes sense." This step ensures that when the AI recommends a rare drug, it's not just guessing; it's reasoning through the clinical context, just like a human doctor would.

The Results: A Big Leap Forward

The team tested GenRxR against 14 other methods, including five other systems that also used Large Language Models. The results were clear: GenRxR was the champion.

Overall Performance: It beat almost every other method in predicting the right medicines.
The Rare Medal: The most exciting finding was for the rare medicines. While other systems struggled to even guess the right rare drug, GenRxR improved its performance by up to 30.9% compared to the strongest baseline.
Safety First: Crucially, getting better at rare medicines didn't hurt its ability to recommend common ones. It became a better all-around doctor without losing its touch on the basics.

The researchers didn't just claim this worked; they measured it using real hospital data (from the MIMIC-III and MIMIC-IV datasets). They showed that by using "what-if" stories to fill in the gaps and teaching the AI to reason through clinical contexts, they could solve a major blind spot in medical AI.

Why This Matters

This paper doesn't claim to have solved every problem in medicine, but it offers a powerful new tool. It shows that we don't always need more real-world data to train AI; sometimes, we just need to be smarter about how we use the data we have. By letting AI imagine "what if" scenarios, we can prepare it for the rare, difficult cases that happen in the real world. For patients with rare conditions, this means a future where AI doctors are just as ready to help them as they are for the common cold. The code and the data they generated are now open for others to use, inviting the whole community to build on this "what-if" approach to make healthcare safer for everyone.

Technical Summary: GenRxR for Rare Medication Recommendation

Problem Definition

The paper addresses a critical limitation in existing AI-based medication recommendation systems: their significantly reduced predictive performance on rarely prescribed medications (rare-meds). While these systems generally aim to enhance patient safety and therapeutic outcomes, the authors observe that they exhibit strong biases toward frequent medications, often failing to predict rare-meds entirely.

The authors attribute this failure to two intrinsic limitations:

Data Scarcity: Rare-meds suffer from a severe lack of training instances due to the long-tail distribution of medication occurrences in Electronic Health Records (EHRs) and strict privacy regulations limiting data availability.
Limited Contextual Modeling: Existing methods often fail to adequately model the complex relationships among co-recommended medications, which provide crucial clinical context for prescribing rare treatments.

The inability to recommend rare-meds is clinically significant because many such medications (e.g., orphan drugs) are essential for treating uncommon or severe conditions with limited therapeutic options. Failure to prescribe them can lead to severe or life-threatening outcomes.

Methodology: The GenRxR Framework

To address these limitations, the authors propose GenRxR, a novel framework leveraging Large Language Models (LLMs) through two primary mechanisms: counterfactual data augmentation and instruction-tuned clinical reasoning.

1. Counterfactual Data Augmentation

To mitigate data scarcity, GenRxR employs an LLM (denoted as $LLM_{CF}$ ) to generate synthetic, counterfactual medical data. This process involves two steps:

Counterfactual Scenario Construction: The system defines hypothetical "what-if" scenarios where a target rare medication ( $m^*$ ) is hypothetically prescribed to a patient who historically did not receive it.
Counterfactual Inference: The LLM is prompted to reason about the plausible clinical changes required to maintain medical consistency. Specifically, given a set of candidate patient records and clinical elements associated with the target medication (selected based on statistical relative risk), the LLM:
1. Summarizes the clinical context of candidate patients.
2. Selects the most plausible candidate for the hypothetical prescription.
3. Generates modified diagnoses, procedures, and co-prescribed medications that would logically accompany the new prescription.

This process enriches the training dataset with diverse clinical patterns, specifically increasing the representation of rare-meds without compromising the integrity of the original data.

2. Instruction Tuning for Clinical Reasoning

To better align the LLM's capabilities with the specific demands of medication recommendation, the authors introduce an instruction tuning phase.

Medical Record Summarization: Using the summaries generated during the counterfactual step, a second LLM (with stronger medical interpretation capabilities) is fine-tuned to perform medical record summarization relevant to specific medications.
Integration: This instruction-tuned LLM is then integrated into the recommendation process. This allows the model to better capture patient clinical contexts, including the nuanced relationships between co-recommended medications, thereby improving its ability to handle rare-meds cases.

Key Contributions

The paper outlines three primary contributions:

Identification of a Critical Gap: The authors reveal and analyze the overlooked limitation of existing AI methods in handling rare-meds, distinguishing this from standard item cold-start problems due to the added constraints of clinical consistency and safety.
First Counterfactual Augmentation Strategy for Medication: To the best of the authors' knowledge, this is the first work to propose a counterfactual data augmentation strategy specifically for medication recommendation, leveraging LLMs to generate plausible synthetic patient data.
Effective LLM Integration: The proposed framework, GenRxR, utilizes instruction-tuned LLMs to significantly outperform existing baselines. The authors release the code and generated counterfactual data to ensure reproducibility.

Experimental Results

The authors evaluated GenRxR against 14 baseline methods, including 5 LLM-based approaches and 9 state-of-the-art medication recommender systems, using the MIMIC-III and MIMIC-IV datasets.

Overall Performance: GenRxR outperformed all baselines in most cases.
Rare-Meds Performance: The framework achieved up to 30.9% higher predictive performance for rare-meds compared to the strongest baseline.
Distribution Mitigation: The counterfactual data augmentation effectively mitigated the long-tail distribution issue, increasing the coverage of rare medications in the training data (e.g., increasing the proportion of top 20% most frequent medications coverage from 63.7% to 71.3% in MIMIC-IV, implying better handling of the tail).
Consistency: The improvements in rare-meds prediction were achieved without compromising performance on frequently prescribed medications.

Significance and Claims

The paper claims that GenRxR offers a novel solution to the data scarcity and contextual modeling challenges inherent in rare medication recommendation. By leveraging the medical knowledge and clinical reasoning capabilities of LLMs, the framework not only generates high-quality synthetic data but also aligns the model's reasoning with clinical realities. The authors position this work as a crucial step toward improving patient safety and therapeutic outcomes for conditions requiring rare or specialized treatments, moving beyond the limitations of purely data-driven models that struggle with long-tail distributions.