APRIL: Annotations for Policy evaluation with Reliable Inference from LLMs

TL;DR

This work tackles off-policy evaluation (OPE) for clinical contextual bandits by addressing dataset coverage gaps with scalable counterfactual annotations generated by large language models (LLMs). The authors demonstrate that LLMs can predict downstream lab values relevant to potassium and sodium repletion from MIMIC-IV data, and they show that incorporating LLM-generated counterfactuals into a direct-method OPE estimator ($\mathrm{DM}^+$) substantially improves RMSE under distribution shifts, with diminishing returns beyond a certain annotation count. An entropy-based stopping criterion $H(A)$ is proposed to identify when additional annotations cease to be useful. The results suggest a practical, scalable pathway to safer deployment of clinical decision policies by enriching behavior data with high-quality, model-generated counterfactuals while monitoring coverage via predictive entropy.

Abstract

Off-policy evaluation (OPE) estimates the value of a contextual bandit policy prior to deployment. As such, OPE plays a critical role in ensuring safety in high-stakes domains such as healthcare. However, standard OPE approaches are limited by the size and coverage of the behavior dataset. While previous work has explored using expert-labeled counterfactual annotations to enhance dataset coverage, obtaining such annotations is expensive, limiting the scalability of prior approaches. We propose leveraging large language models (LLMs) to generate counterfactual annotations for OPE in medical domains. Our method uses domain knowledge to guide LLMs in predicting how key clinical features evolve under alternate treatments. These predicted features can then be transformed using known reward functions to create counterfactual annotations. We first evaluate the ability of several LLMs to predict clinical features across two patient subsets in MIMIC-IV, finding that state-of-the-art LLMs achieve comparable performance. Building on this capacity to predict clinical features, we generate LLM-based counterfactual annotations and incorporate them into an OPE estimator. Our empirical results analyze the benefits of counterfactual annotations under varying degrees of shift between the behavior and target policies. We find that in most cases, the LLM-based counterfactual annotations significantly improve OPE estimates up to a point. We provide an entropy-based metric to identify when additional annotations cease to be useful. Our results demonstrate that LLM-based counterfactual annotations offer a scalable approach for addressing coverage limitations in healthcare datasets, enabling safer deployment of decision-making policies in clinical settings.

Figures (9)

• Figure 1: Our work improves OPE estimates using LLM-generated counterfactual annotations. We first query counterfactual annotations using domain knowledge guided prediction. We calculate the annotations using a known reward function $R$. Finally, we incorporate the counterfactual annotations and offline behavior dataset to learn an OPE estimate $\hat{v}(\pi_e)$.
• Figure 2: LLM-generated counterfactual annotations improve OPE estimates in settings with high divergence between actions observed in behavior and target policies. We report results for the potassium repletion task. Our baseline is a direct method estimator (blue) that does not use counterfactual annotations. The performance of $DM^+$ with annotations from each LLM is reported in the corresponding colors. Error bars represent standard error across $500$ bootstrapped datasets sampled with replacement. Since RMSE is non-negative, the lower bound of the error bars is truncated at 0 where necessary. Figures above each plot demonstrate the difference in distribution of actions observed in the behavior and target policies.
• Figure 3: Combining annotation sources yields limited returns. (Top) We compare $DM$ to $DM^+$ with annotations from the best-performing LLMs for potassium repletion in the dosage cohort, with two aggregation methods: pooling predictions and averaging annotations. Error bars show standard error over $500$ bootstrapped datasets, truncated at 0. (Bottom) Marginal entropy over the action space $H(A)$ when adding counterfactual annotations to the behavior dataset for the potassium cohort. The dashed line marks the maximum possible entropy.
• Figure 4: Reward functions for both decision-making tasks are a function of the corresponding reference range. Reward is bounded in the range $[0, 1]$, attaining its maximum when the lab value falls within the corresponding clinical reference range ($3.5-4.5$ mEq/L for serum potassium, and $135-145$ mEq/L for serum sodium). As the lab value deviates from this range, the reward decreases according to a Gaussian decay, with the lowest rewards assigned to critically low or high values.
• Figure 5: LLMs can be prompted to construct downstream lab value predictions. The prompt contains separate components that first describe the patient's clinical state four hours prior to receiving treatment, and then contains instructions to perform the lab value prediction. The prompt includes relevant information from UpToDate, a clinical resource, to help an LLM identify which features in the medical record are most predictive.