SAFER: A Calibrated Risk-Aware Multimodal Recommendation Model for Dynamic Treatment Regimes
Authors: Yishan Shen
Yuyang Ye
Hui Xiong
Yong Chen
Venue: ICML 2025 poster (2025)
Keywords: Dynamic Treatment Regime
Multimodal
Risk-aware
Conformal inference
Primary area: applications->health_medicine
TL;DR: We propose SAFER, an end-to-end multimodal Dynamic Treatment Regime framework that delivers reliable treatment recommendations.

Abstract: Dynamic treatment regimes (DTRs) are critical to precision medicine, optimizing long-term outcomes through personalized, real-time decision-making in evolving clinical contexts, but require careful supervision for unsafe treatment risks. Existing efforts rely primarily on clinician-prescribed gold standards despite the absence of a known optimal strategy, and predominantly using structured EHR data without extracting valuable insights from clinical notes, limiting their reliability for treatment recommendations. In this work, we introduce SAFER, a calibrated risk-aware tabular-language recommendation framework for DTR that integrates both structured EHR and clinical notes, enabling them to learn from each other, and addresses inherent label uncertainty by assuming ambiguous optimal treatment solution for deceased patients. Moreover, SAFER employs conformal prediction to provide statistical guarantees, ensuring safe treatment recommendations while filtering out uncertain predictions. Experiments on two publicly available sepsis datasets demonstrate that SAFER outperforms state-of-the-art baselines across multiple recommendation metrics and counterfactual mortality rate, while offering robust formal assurances. These findings underscore SAFER’s potential as a trustworthy and theoretically grounded solution for high-stakes DTR applications.

SAFER: A Calibrated Risk-Aware Multimodal Recommendation Model for Dynamic Treatment Regimes
Yishan Shen Yuyang Ye *2 Hui Xiong
Yong Chen
Abstract
1. Introduction
How can we enable models to recognize when they are uncertain about their predictions? Safely providing person- alized, sequential treatment recommendations that adapt to patient's 's evolving clinical state is a longstanding challenge in optimizing outcomes in high-stakes healthcare scenarios. In this work, we address this challenge within the frame- work of dynamic treatment regimes (DTRs) (Robins, 1986; Murphy, 2003; Chakraborty & Moodie, 2013; Tsiatis et al., 2019). Crucial for real-world decision-making. resource allocation, and reducing tria-andderror treatments (Mur- phy, 2005; Laber et al., 2014), DTR requires more precise, adaptive, and safe control over treatment strategies while minimizing risks in critical clinical contexts.
Dynamic treatment regimes (DTRs) are critical to precision medicine, optimizing long-term out- comes through personalized, real-time decision- making in evolving clinical contexts, but re- quire careful supervision for unsafe treatment risks. Existing efforts rely primarily on clinician- prescribed gold standards despite the absence of a known optimal strategy, and predominantly using structured EHR data without extracting valuable insights from clinical notes, limiting their reliabil- ity for treatment recommendations. In this work, we introduce SAFER, calibrated risk-aware tabular-language recommendation framework for DTR that integrates both structured EHR and clin- ical notes, enabling them to learn from each other, and addresses inherent label uncertainty by as- suming ambiguous optimal treatment solution for deceased patients. Moreover, SAFER employs conformal prediction to provide statistical guar- antees, ensuring safe treatment recommendations while filtering out uncertain predictions. Exper- iments on two publicly available sepsis datasets demonstrate that SAFER outperforms state-of- the-art baselines across multiple recommendation metries and counterfactual mortality rate, while offering robust formal assurances. These find- ings underscore SAFER's 's potential as a trustwor- thy and theoretically grounded solution for high- stakes DTR applications.
Recent advances in deep learning (DL. have significantly improved DTR frameworks by addressing key challenges such as patient heterogeneity, temporal dependencies, and the high-dimensional clinical data (Kosorok & Laber, 2019; Moodie et al., 2007). DL offers distinct advantages for DTR, including the ability to integrate heterogeneous data sources, such as electronic health records (EHRs) and temporal pat- terns, while capturing complex dependencies over time (Yu et al., 2021: Ching et al., 2018: Olawade et al., 2024; Mel- nychuk et al., 2022). However, major challenge in current DTR approaches is the absence of optimal treatment strate- gies, particularly for understudied diseases, critically ill or deceased patients as their outcomes may not reliably in- dicate the most appropriate clinical actions (Robins et al., 2000; Schulam & Saria, 2017: Chapfuwa et al., 2021). This inherent label uncertainty in DTR remains underexplored, limiting the robustness of existing models. Furthermore, most approaches lack theoretical guarantees their rec- ommendation quality, leaving practitioners without princi- pled mechanisms for error rate control (c.g. Benjamini & Hochberg, 1995; Lei et al., 2018; Bates et al., 2023; Jin & Candès, 2023).
'Equal contribution Uniiersitty of Pennsylvania Rutgers University 'The Hong Kong University of Science and Technology (Guangzhou). Correspon- dence to: Hui Xiong hkk,, Yong Chen <ychen123< cychen123@pommedicnne. www.mmddiciie.upen..dde
Additionally, many existing methods (c.g., Murphy, 2005; Laber et al., 2014; Bica et al., 2020) rely primarily struc- tured EHR data while underutilizing the valuable textual information contained in clinical notes, which often cap- ture critical insights into a patient's history and physician assessments. However, integrating clinical notes into DTR
Proceedings of the 4nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Risk-Aware Dynamic Treatment Regimes
remains a challenge due difficulties in establishing a uni- fied embedding space that preserves moter-md context, temporal alignment, and domain semantics. Prior work, such as Choi et al. (2017); Shang et al. (2019a), has ex- plored representation learning for medication recommenda- tion, but these approaches still rely heavily on medical code representations derived from structured EHR data.
2. Related Works
Dynamic treatment regimes. Prescribing medications in response to the dynamic states of patients is a challeng- ing task. Over the past decade, to model complex, high- dimensional, and temporal healthcare data, researchers have leveraged various deep learning-based approaches to im- prove treatment recommendations, including RNNs and their variants (e.g., Choi et al., 2016; Bajor & Lasko, 2017; Jin et al.. 2018), attention networks and transformer- models (e.g., Peng et al., 2021; Wu et al., 2022), deep re- inforcement learning (DRL) techniques (e.g., Bothe et al., 2013; Komorowski et al., 2018; Raghu et al., 2017; Saria, 2018; Wang et al., 2018: Zhang et al., 2017), convolutional Suo et al., 2017; Cheng et al.,
We address these challenges through two key desiderata: (i) uncertainty control the DTR model should quantify prediction uncertainty and provide statistically guaranteed control over the uncertainty discovery threshold specified by the user, and (ii) comprehensive information fusion- the model should integrate all available patient data, ensuring no critical information is overlooked. Together, we refer to these principles as "risk awareness".
neural networks (CNNs) (uo 2020; Su et al., 2022), and generative adversarial networks (GANs) (e.g., Wang et al., 2021a;b). Despite these advance- ments, assessing the effectiveness and ensuring reliable in- ference for data-driven DTR approaches remains a signifi- cant challenge due variability in evaluating the quality of suggested prescriptions (Hussein et al., 2012; Chakraborty et al., 2014). For instance, while DRL excels in learning optimal DTRs by discovering dynamic policies, inconsis- tencies in reward design, policy evaluation, and MDP for- mulations often hinder standardized and rigorous healthcare applications (Luo et al., 2024). We hypothesize that these challenges can be mitigated by modeling predictive uncer- tainty and generating selective candidates via conformal inference, while providing statistical guarantees.
In this work, propose SAFER, Calibrated Risk-Aware multimodal framework for DTR that enhances the robust- ness of DTR frameworks. The overall pipeline is depicted in Figure 1. SAFER introduces several key innovations:
Multimodal Representation Learning. SAFER inte- grates both structured EHR data and unstructured clinical notes using novel Transformer-based architecture to learn unified sequential patient representation. A self-attention mechanism captures inter-modality temporal dependencies, while cross-attention extracts contextual information across modalities (Section 4.
2. Unerrtainty-Aware Training. SAFER accounts for label uncertainty by assuming ambiguous treatment labels particularly for deceased patients. By recognizing that label uncertainty is systematic and predictable, we introduce uncertainty quantification module that assigns per-label risk scores and incorporates them into a risk-aware loss function for interactive training (Section 4.2).
Risk-aware treatment recommendation. To the best of our knowledge, this work is the first to incorporate uncer- tainty modeling and employ conformal prediction (CP) into DTR research. Several studies have integrated drug-drug interaction knowledge to optimize personalized medica- tion combinations and minimize adverse outcomes. (Shang et al., 2019;; Wu et al., 2022; Tan et al., 2022; Yang et al., 2021). In contrast, we incorporate an uncertainty quantifi- cation module to improve both the accuracy and safety of DTRs, ensuring statistically reliable treatment recommenda- tions. As a pivotal role optimization and ceccisionmmaking process, prior work has utilized uncertainty quantification in computer vision (Liu et al.; Harakeh et al., 2020), im- age/video restoration (Shao et al., 2023; Dorta et al., 201 8), natural language processing (Chen et al., 2015; Lin et al., 2023; Ren et al., 2023), bioinformatics (Xia et al., 2020; Bian et al., 2020), etc. In the healthcare domain, (Chua et al., 2023; Liu et al., 2024) exemplifies recent efforts to address prediction uncertainty in clinical machine learning models. While our framework is specifically designed for DTR prediction with rigorous theoretical guarantees via CP, we emphasize complementary but distinct focus on label uncertainty in dynamic treatment regimes. We draw further inspiration from selective prediction, particularly CP-based selection (Bates et al., 2023; Jin & Candès, 2023; Gui et al.,
3. Theoretical Guarantees on Prediction Reliability. Given the critical need for error control in high-stakes sce- narios, We derive theoretical guarantees calibrated rec- ommendations by innovatively employing conformal infer- ence framework (Vovk et al., 2005; Benjamini & Hochberg, 1995) to control the expected proportion of unreliable pre- dictions (i.e., FDR) at decision time (Section 5).
4. Empirical Validation. We evaluate SAFER on real- world EHR benchmarks, demonstrating consistent improve- ments over ffate-of-hhe-es DTR methods across multiple recommendation metrics and reductions in counterfactual mortality rates.'
Together, these advancements establish SAFER as a multi- modal, rwske DTR framework with strong theoretical foundations and superior empirical performance.
at !Our code and dataset are avaliable http:///ghuuh.oo/yyhaanssSS////
2

---

Risk-Award Dynamic Treatment Regimes
Next Timrstamp Medication Prodiction
Risk-aware Loss Prediction
Threshold
c].0]
ej.of
Pacudo Annotation
Y
Decossed Patient
Unceriainty Quantification
Only Provide Confident
Between Refined
Pre traltrrleed Modiules
Recommendation Patients
cf.
tuning
Serviving Purient
Vine--tning
Time-incriiat Information
Electronic Health Records
Clinkcal Notes
R Tdd
Prigeos hale
+
in dontppik ECD domoo
and lapuli de Cadilagr
X.
R
Embedded via Clinical BERT
Self Attention
Crass Attention
Feud-forward
Allection
Aticution
Figure I. The overall framework of SAFER.
2024), to develop an end-to-end safe DTR framework with formal assurances.
denotes the dimensionality of the tabular covariates. Sim- itarly, of corresponds to the clinical note associated with patient at time step By integrating these two modali- ties, we aim to predict the medications be administered in the next clinical decision window, i.c., time step T + enhancing treatment recommendations with both numerical clinical metrics and rich textual context.
Clinical notes combined with EHR. Clinical notes, central patient care, capture physicians' thought processes, obser- vations, and treatment rationale often absent from structured EHR data (Sheikhalishahi al., 2019; Rosenbloom et al., 2011). Integrating clinical notes with structured EHRs has demonstrated improved predictive performance in biomedi- cal research, enabling more comprehensive patient modeling (Gao et al., 2024; Lyu al., 2023). However, current DTR research largely ignores clinical notes due challenges like data heterogeneity, unstructured text processing, and the absence of standard tools, with most representation learning frameworks relying structured medical codes (Choi et al., 2017; Shang et al., 2019b). This work bridges these gaps by incorporating clinical notes into DTR frameworks through a novel transformer-bret multimodal fusion approach to enhance decision-makiin accuracy and reliability.
A significant challenge in DTR lies in uncertainty associated with treatment labels, particularly for negative trajectories (i.e., patient states resulting adverse outcomes). Specifi- cally, patients with negative outcomes during hospital stay often exhibit unstable and irregular medication patterns, leading to label ambiguity. While some studies in recom- mendation disregard negative trajectories (e.g., Sun et al., 2021; Ye et al., 2024; 2025), others, like Wang et al. (2020), incorporate this information to refine learned policies and avoid repeating errors. Similarly, we retain negative trajec- tories but attribute their ambiguity to two primary cases: (1) appropriate treatments were administered but were insuffi- cient to prevent death, and (2) the likelihood that incorrect treatments contributed to adverse outcomes. In contrast, patients with positive outcomes generally display labels that more reliably represent effective clinical decisions. We adopt above assumptions throughout this work.
3. Problem Setup
3.1. Risk-aware Dynamic Treatment Prediction
We aim to model dynamic treatment prediction using two complementary data sources: structured electronic health records (EHR)& E = E,, E, ..... Ex and unstructured clinical notes O = (O,, Oz..., ON}. For each patient where i (1....,, N). structured data E, are represented sequence of tabular records (et, ..,.... ef') while unstructured data C, consist of sequence of clinical notes Each vector ef de represents the
We propose addressing label uncertainty by explicitly esti- mating uncertainty score K; for each candidate through uncertainty quantification module. SAFER fine-tunes the model with awareness of candidates whose pseudo- annotations may lack reliability. The uncertainty score is integrated into a novel risk-aware loss function, mitigating the influence uncertain optimal treatment labels.
(o/, of of') tabular features at time step t, where t E I.....T T and dg
3

---

where each x; is an input embedding and y/s an observed treatment label. All pairs (x, are drawn i.i.d. from the same distribution as the training data set which used to train a predictor f: X + V. We then have m new test samples { with true but unobserved labels Risk-Aware Dynamic Treatment Regimes
3.2. Conformal Inference for FDR Control
clinicians can trust that the expected fraction of incorrect rec- ommendations remains safely bounded, thereby enhancing patient safety in the automated decision-making process.
In high-stakes domains such as treatment recommendation, it is vital to provide calibrated predictions and control the rate of incorrect decisions. To achieve this, we adopt conformal inference procedure (Vovk et al., 1999; 2005; Shafer & Vovk, 2008) that selects a subset of plausible prediction candidates with a statistical coverage guarantee. We assume access to a calibration set
4. Method
4.1. Dynamic Treatment Prediction
In practice, clinicians typically rely on both structured data and clinical notes to monitor disease progression and guide treatment decisions (Assale et al., 2019; Gangavarapu et al., 2020). Hence mimic the real-world clinician decision- making process and recommend treatments at the next time step, we construct a unified time-series representation of the patient health embeddings by integrating these comple- mentary data sources, To be more specific, given a patient represented by their multimodal electronic health record sequence r; {(e o.), , of),...........
Zcal = {xi, Me)i,
(nn-))z=i- each drawn i..dd from the same (unknown) data-generating process.
(e" o where of E R and of E R denote the structured EHR data and clinical notes at time I, our goal is to predict a treatment recommendation j for the next time step T + 1.
To address the risk of recommending incorrect treatments, we compute an uncertainty score K; for each patient i E [n + m via an "uncertainty map" module, and convert these scores into conformal p-values (e.g. Vovk al., 1999; 2005; Bates et al., 2023; Jin & Candès, 2023; Liang al., 2024) to control false discovery rate (FDR). The FDR is defined as
Inter-modality Temporal Dependency. We first project the information from each data source into dense latent spaces. Specifically, we use BioClinicalBERRT to encode clinical notes, modeled as X (Alsentzer et al., 2019), which provides superior performance in encoding clinical text due to its bidirectional attention mechanism and domain- specific pretraining on large-scale biomedical and clinical corpora (Huang et al., 2023; Hu er al., 2024; Zhang et al., 2022). On the other hand, tabular data is encoded as Xc with normalization and one-hot encoding.
FDR = EE R
where
R = (total # of rejections),
V (# of false rejections).
Concretely, we define a null hypothesis
For a patient p., the sequence of each modality is aligned with the timestamps. The embedding layers f x + TTxxs are then applied to each modality. where A {E, O) and da is model dimensionality. To capture tem- poral dependencies within each modality, masked self- attention mechanisms are applied as,
Hy j.,,
for each test sample and reject (i.c., recommend) subset S C (1,...,m, m) of hypotheses while ensuring FDR, S a, where a E (0. 1) is user-specified tolerance, is predefined uncertainty threshold.
(X.a W (X
(X) T M
S = Softmax
x/WXX +PE
This setup can be viewed as a standard multiple testing problem (Benjamini & Hochberg, 1995; 1997: Benjamini & Yekutieli, 2001; Efron, 2012) with m null hypotheses HI,.... Hm We compute conformal p-value pj for each H, and apply an contcontro procedure (e.g., Ben- jamini & Hochberg, 1995) to obtain a set S with the desired error rate a. In binary classification tasks, the FDR serves as an analog to Type-I error control(Hastie, 2009). For re- gression problems with continuous responses, controlling errors is appropriate when each selected candidate incurs comparable cost. This approach is particularly crucial in scenarios like medical decision-making and knowledge retrieval, where the cost of committing a type-I error can be significant and should be a primary consideration. In high- stakes clinical settings, a falsely recommended treatment can lead to adverse outcomes. By constraining FDR S a,
da
(2)
where M R is the causal mask matrix, and PE TXxx is the sinusoidal position encoding.
Cross Cruss-modality information integration. To effectively in- tegrate information from different data sources, we design a cross-attentie mechanism that enables different sequences to learn contextual information from each other, which can be formulated as,
SP S!W (S
H.
softmax
(3)
Vdl
S S (S w
C softmax
s!ww
S.
(4)
ddu
httg://wwwggnnnaaa..cccollllllne Banaaaaaaaeeeeeeeeeeeiiiiieaaaaa ClinicalBERT
4

---

Risk-Aware Dynamic Treatment Regimes
where S.P and S are temporal-aware representations for each modality, and @ denotes the concatenation operation. After integrating the static information embedding x the final representation is given by h; = H. @ x,". where h; E Sde represents the unified patient embeddings.
2. J L-Lipschitz continuous over latent representation space H. where h E H.
that is there exists constant 0, such that
Ep-P- [Na]] - Ex [Ma] 2 22200 0.
(6)
Subsequently, a feedforward sework-based classification layer is applied to produce probability distribution over medication classes for the next time step, as fa: H. + RD! The model is trained using cross-entropy loss to minimize prediction error across all instances.
The proof details can be found in Appendix A.I. Hence in this regard, such KL divergence can capture the distri- butional difference between survival and deceased patients, served as a valid measure of uncertainty where K; represents the prediction uncertainty.
4.2. Risk-Aware Fine-Tuning
Remark 4.2. A Lipschitz-constrained student (e.g. weight decay plus spectral normalisation) is a standard practice when one wishes to avoid uncontrolled extrapolation out- side the training manifold. In our approach, we rely on the multilayer perceptron-basee -baood architecture, which inherently exhibits Lipschitz continuity when the activation functions are smooth and bounded, and the model's 's weights are ap- propriately regularized (Gouk et al., 2021).
After the dynamic treatment prediction module converges, we introduce risk-aware fine-tuning procedure to account for label uncertainty as stated in Section 3.1, assuming reli- able labels for surviving patients while uncertain labels for deceased patients. However, training exclusively on surviv- ing patients significantly discards valuable information. To mitigate this, we treat labels for deceased patients as pscudo- labels and propose an uncertainty module to incorporate the brought risky information effectively.
Risk-aware Loss Function. The uncertainty term is in- corporated into the loss function during fine-tuning. For surviving patients, uncertainty remains minimal, enabling the training process to prioritize these samples. Conversely, for deceased patients, the loss function penalizes significant deviations, ensuring the model effectively captures risk. The modified loss function defined as follows,
Uncertainty Estimation. Here, we refine predictions for surviving patients by introducing a multilayer perceptron- based module fa. This module takes the patient embeddings h, learned in the previous stage, as input to generate new predictive distribution for each surviving patient. The uncer- tainty module is trained exclusively on surviving patients using enoss-entropy loss to minimize the prediction error.
N Σ
C 1
(1--))
Σ ye log po ( = l1aa) + 7y03, (7)
Since predictions have been refined for surviving patients, who exhibit more stable patterns and have reliable labels, the KL divergence between the logits before and after re- finement can capture the distributional difference between surviving and deceased patients. This divergence, as shown in the equation below, can be interpreted as measure of uncertainty for deceased patients.
where Rt is the normalized uncertainty term and y controls the regularization strength, penalizing overconfident predic- tions for high-risk cases.
5. Conformal Selection and FDR Control
We begin by fitting our risk-aware model on the training set. Subsequently, for calibration sample {(x,, })){= and unlabeled test data (x 1= we compute the predicted uncertainty score RG R(x,), for every E [n + m.). Here, R: X' R is an uncertainty score predictor that depends the patient health trajectory (e.g., time-series of clinical measurements), rather than on the observed treatment label yi- As described in the previous section, s((x;) captures a "label mismatch risk" based on the distributional divergence between an refined module and a fine-tuned module. Cru- cially, this score does not require knowledge of the final treatment ys- We require that R is computed in the same way for calibration and test samples. This consistent def- inition of к across calibration and test sets preserves the exchangeability necessary for valid conformal inference.
During model inference, we quantify the predictive uncer- tainty by computing the KL divergence between the output distributions of the two modules fe and fo as follows,
K, = DxL (p.) ll (h.)) = Σ Po (M = -/h,) In
Pa (
1 ([b.)
Po
1 (h:)
(5)
where p(h,)= = Softmax(f(h)) denotes the predicted prob- ability distributions from both module, and yi represents the predicted class.
Theorem 4.1. h- ~ P- (h) and h+ p- denote the latent representations of survivors and deceased patients respectively. Under the following conditions,
Conformal p- valuee Consider a test sample j E [m] for which we wish to test the hypothesis Hy: Katj 2 c. Then
1. DKL(P- (h) 11 P- (h) > 0, Le. P- + p+
5

---

Risk-Aware Dynamic Treatment Regimes
define the conformal p-value
2. There exists some M such that supx x(x) S M.
Σ 1 1 Ri < kntji 20}
Then, for any user-defined threshold a E (0, 1), the BH- based conformal selection set S satisfies
Pj
n +
U3
(1
+ 2 i= 1
Rs
Rutts Ki ))
FDR = max R3
2
1 ( is mne, JES)
where Uj Uni 1) are i.i.d. random variables used for tie-breaking under the multiple testing setting. Intuitively, Pj measures how frequently the predicted calibration un- certainty scores { are less than or equal to the test uncertainty score kn+ restricted to those Ki 2 c.
= E
S a.
(10)
Σ5 [(jEs)
The proof details are provided in Appendix A.2. By combin- ing the conformal p-value construction in (8) with the BH procedure, our method ensures that the expected proportion of unreliable recommended treatments remains bounded by a. By selecting appropriate uncertainty score function K, where higher values correspond lower plausibility of the candidate label y. we enable an intuitive calculation of conformal p-values. This framework provides a practical safety margin for high-stakes applications, while supporting flexible and data-driven selection of plausible labels.
Traditionally, conformal p-values are constructed to super-aniform under the null, meaning that if the tested label (or score) truly matches the data-genn gentrating distribu- tion, then (p, a)) I a for all E [0, 1] (Vovk al., 2005; Lei et al., 2018). Here, the setup follows Jin & Candès (2023), who define H,: Yn+j S Cj for random hypotheses based on unobserved labels. In the present formulation (8), the p-value satisfies a selective guarantee (Jin & Candès, 2023), namely:
6. Experiments
P P where S is the final selected ("rejected") set. In other words, the joint event that j is included the recommendation set and Dy I a occurs with probability no larger than a.
(ps sa) s a, V [0, 1],
We empirically validate the SAFER model on two sepsis cohorts derived from publicly available datasets, Medical Information Mart for Intensive Care (MIMIC)-III covering over 40,000 ICU stays (2001-2012) (Johnson et al., 2016) and MIMIC-IV with over 65,000 ICU and 200,000 ED ad- missions (2008-2011 2019) (Johnson al., 2023a;b), to evaluate recommendation accuracy and FDR control. For this study, we define cohorts based the sepsis-3 criteria (Singer et al., 2016), focusing on the early stages of sepsis man- agement-22 hours prior to and 48 hours after sepsis onset. The treatment selection involves intravenous fluid and vaso- pressor dosage within 4-hour window, mapped to a 5 x 5 medical intervention space, following Komorowski al. (2018). Figure 2 shows the distribution of sepsis treatment co-occurrence in the two cohorts.
Benjamini-Hochbery (BH) Procedure. After comput- ing conformal p-values - Pmn for the test samples, we control the FDR via the classic Benjamini-Hochbeee (BH) algorithm (Benjamini & Hochberg, 1995). First, sort the p-values in ascending order:
P(1) $ P(2) ... S P(m)-
Then, let
k- = max r P()) I ar
where E (0, 1) is the user-specified FDR threshold. If no such r satisfies the inequality, we set k = 0. The BH procedure "rejects" the smallest p-values, i.c. the set (P(1).... )(.)) Accordingly, our conformal selection output
S {j [m] P S p(o)
meaning only recommend labels whose p-values rank among these top k.
(A) IV Fluds
(8)
FluIes
Below, we show that this procedure, using p- p-values of the form (8), controls the FDR under suitable assumptions. Theorem 5.1. Axsume we have
Figure 2. Comparative visualization of the treatment frequency matrix in log scale from two datasets. Panel (A) represents MIMIC- III, while Panel (B) corresponds to MIMIC-IV.
1. The calibration data {(x,, Wi and test data (xn+sl7 are Li.d. and data in {(x,, y U (xn+i)14s U {xx+s} are mutually independent for any j E [m].
For each patient, we extract 5 types of static demographic variables and 44 types of time-series variables from the tab- ular data. We set the historical sequence length to All
6

---

Risk-Aware Dynamic Treatment Regimes
Table
The overall performance of SAFER and haseline methods. (p < 0.05)
MIMIC-III
MIMIC-IV
Methods
MI-AUC
MA-AUC
HR@3
MRR@3
Mortality
MI-AUC
MA-AUC
HR@3
MRR@3
[ Mortality
LSTM
0,9122
0.7934
0,7481
0,8015
0.0915
0.9213
0.8121
0,7551
0.8066
0.1051
RETAIN
0.9257
0.8219
0.8324
0,8153
0.1994
0.9279
0.7851
0.8017
0.8052
0.1863
TAHDNet
0.9213
0.8017
0.7123
0.8109
0.2214
0.9157
0.8274
0.7554
0.8315
0.2466
Naive RL
0.7436
0.6025
0,5303
0.8891
0,0881
0,6782
0.5971
0.5068
0.8217
0.1172
SRL-RNN
0.8751
0.6215
0.7722
0.7916
0.3124
0,8781
0.6982
0.7824
0.8151
0.3219
ACIL
0.8219
0.7012
0,8013
0,8313
0,3212
0,8854
0.7135
0,8319
0.8441
0.3782
ISL
0.8903
0.7785
0.7623
0.7521
0.2783
0.8713
0.7315
0.7741
0.7229
0.3118
SAFER
0.9407
0.8672
0.8517
0,9017
0.3891
0,9356
0.8755
O. 8713
0.8698
0.4562
clinical notes were aligned to the closest timestamp. To handle outliers, we applied the interquartile range (IQR) method for removal and imputed missing values using the k-nearest neighbors approach. Subsequently, all variables were rescaled to the [0,1] interval using z-score normaliza- tion. The two datasets were randomly split into training, calibration (validation), and test sets in an 80% ratio via patient-level splits to ensure no patient overlap, under the assumption that the entire dataset is i.i.d, sampled from common distribution.
Sequential embedding methods include: LSTM (Hochreiter & Schmidhuber, 1997), widely used time-series prediction model; RETAIN (Choi et al., 2016), a two-level neural attention based model that highlights key visit sequences for treatment prediction; TAHDNet (Su et al., 2022), a hi- erarchical temporal dependency network for dynamic treat- ment prediction. Reinforcement learning based methods include Naive Baseline for RL.Lu et al., 2024), a simple rule-based approach for benchmarking RL algorithms. SRL- RNN (Wang et al., 2018), which integrates supervised learn- ing with RL using survival signals as rewards. ACIL (Wang et al., 2020), an adversarial imitation learning approach that optimizes treatment by learning from both successful and failed trajectories. ISL (Jiang et al., 2023), a prototype- based model ensuring treatment actions align with learned representations. For fair comparison, all methods use the same data sources. Methods lacking native text process- ing capabilities incorporate BioClinicalBERT embeddings for clinical notes. Given our assumption that only surviv- ing patients have fully reliable labels, we conduct primary evaluations on this subset, while assessing generalization to deceased patients through counterfactual mortality rate analysis across the entire dataset population.
The uncertainty score Ką quantifies the model's confidence in its treatment predictions. Once the uncertainty scores are computed for the training set, we train the uncertainty score predictor R on the calibration and test sets using standard machine learning models leveraging the full embedding fea- ture space X. Model performance assessed by evaluating the average FDR across 500 independent experiments. Ap- pendix C.3 provides a sensitivity analysis of several hyper- parameters, including the length of historical information, hidden dimension, and in the loss function.
6.1. Evaluation Metrics
To evaluate the performance of SAFER and other baselines, we report MRR @3 and HR @3 for treatment ranking, as well as Micro AUC and Macro AUC for assessing predic- tive performance in the multiclass classification setting of DTR. Additionally, we report the counterfactual mortality rate reduction, which is measure of how recommended treatments might have improved survival outcomes relative to real-world clinical actions (Laine et al., 2020; Kusner et al., 2017; Valeri et al., 2016), to validate the effectiveness of the recommended treatment as part of an offline value estimation. The details for valid counterfactual mortality rate calculation are provided in Appendix B.
6.3. Overall Performance
Table 1 shows the overall performance of our proposed SAFER and baseline methods. Our analysis reveals several key observations follows.
First, RL methods demonstrate consistent underperformance classification benchmarks in general, especially with the macro-AUC metric on MIMIC-II revealing 16.6% deficit compared to sequence-based counterpart baselines on average. This probably stems from severe class imbalance of treatment label (Class 0 takes 61.2% percent), where sparse disease-specifie reward signals prove insufficient for distinguishing between classes. Consequently. RL agents fail to develop discriminative policies, resulting even worse performance compared with simple baselines in mortality prediction. While imitation learning approaches partially mitigate this issue, their performance remains suboptimal
6.2. Baseline Methods
For validating the effectiveness of SAFER, we selected sev- eral baseline methods for comparison. The baselines can be categorized as sequential embedding based and reinforce- ment learning based approaches.

---

Risk-Aware Dynamic Treatment Regimes
Table The performance of SAFER and its variants. (p 0.05)
MIMIC-III
MIMIC-IV
Variants
MI-AUC
MA-AUC
HR@3
MRR@3
Mortality
MI-AUC
MA-AUC
HR@3
MRR @3
4 Mortality
SAFER-F
0,9059
0.7140
0.7254
0.8067
0.2402
0.8853
0,6851
0.7199
0.7542
0.2315
SAFER-N
0.8655
0.7651
0.7523
0.7803
0.2951
0.8897
0.7841
0.7553
0.8029
0.3875
SAFER-U
0.9188
0.8237
0.8321
0.8769
0.2982
0,9231
0.8317
0.8451
0.8544
0.3765
SAFER
0.9407
0.8672
0.8517
0.9017
0.3891
0.9356
0,8755
0.8713
0.8698
0.4562
1.0
1.0
1.0
0.8
0.3
FDR and Plower
0.0
6.6
0.4
FOR
Fran
FDR and
0.9
0.4
FDR Pewn:
Power
ZDR and
0.5
Piwer
0.2
0.0
0.
0.
02
0.6
0.8
100
0.9
0.4
0.0
0.5
1.0
Target level
(
0.2
Target level
(0)
Target level
(0)
Target level
Figure 3. FDR and power curves across different target level and varies uncertainty threshold on MIMIC-III with Ridge Regression.
FDR and Poweee
FOR und
IDA
and
(0) = 0.1 TorgetFOR FOR level
(b) 0.2 Torget level
(6) 0.3 Target FOR level
0.0 (d) 0.4 Target loval
Figure FDR and power curves across different target a level and varies uncertainty threshold e on MIMIC-IV with Ridge Regression.
due to high patient heterogeneity in treatment responses.
derperform compared to RL methods this aspect, as they indiscriminately encode all treatment sequences rather than learning distinct treatment strategies tailored to different patient outcomes. However, our risk-aware fine-tuning mod- ule significantly mitigates this limitation by incorporating uncertainty information and applying penalty mechanism to focus training on reliable labels. As a result, SAFER achieves a substantial reduction in mortality rate, show- casing the effectiveness of our approach in recommending treatments with valuable worldworld meanings.
Second, while sequence embedding- based methods excel in classification tasks, their effectiveness declines in ranking metrics compared to RL approaches. This is probably be- cause embedding methods, while capturing global patterns, fail to distinguish fine- snee-graind treatment efficacy differences. RL methods, by contrast, explicitly optimize treatment poli- cies using reward signals tied to cure outcomes, enabling dif- ferentiated pattern learning across survival and death trajec- tories. Traditional embedding approaches indiscriminately encode all treatment sequences, introducing uncertainty due to distribution shifts between surviving and deceased pa- tients, ultimately reducing performance in fine-graing rank- ing. However, our SAFER model, though embedding-basee effectively mitigates this issue.
6.4. FDR Control
SAFER strictly controls the FDR. Figure 3 and Fig- ure present the realized FDR and power curves for SAFER on MIMIC-III & IV respectively at various tar- get FDR levels E {0.05, 0.10. 0.95) across differ- ent uncertainty thresholds c. Here, power is defined as Σm [HH is JES)
Finally, our proposed method demonstrates superior per- formance across both classification and ranking metrics effectively modeling EHR data sequences across multiple modalities and integrating label uncertainty from deceased patients. This results in more robust and generalizable pre- dictions. Beyond these metrics, we further analyze the im- pact of our model on counterfactual mortality rate reduction. Overall, sequential embedding-basee based methods tend to un-
The uncertainty score pre-
1. L(H, is [alse]
dictor trained with all feature embeddings via Ridge Regression. Results trained from additional regression pre- dictors are provided in Appendix C.2. The results show that SAFER maintains strict control the FDR the speci- fied level which stabilizes increases. Also, the power
8

---

Risk-Aware Dynamic Treatment Regimes
curve asymptotically converges to one with increasing o, in- dicating that SAFER selects all confident candidates without exceeding the FDR constraint.
tion in mortality rates. These results underscore SAFER's potential for trustworthy and risk-aware decision support in real-world clinical settings.
Choice of Uncertainty Score Threshold. The performance of SAFER depends on the choice the uncertainty thresh- old c, evaluated over c E {0.1, 0.2, 0.3, 0.4). As e increases, both FDR and power stabilize more quickly, converging to constant value and one respectively. For с 0.1, the FDR curve remains steady, and power remains close to zero until 0.7. In contrast, the FDR reaches a constant value of 0.1 even at a = 0.1 with power as high as 1. SAFER employs KL-diieegencee to quantify model prediction un- certainty while there is no universally accepted threshold to determine when two distributions differ meaningfully. Therefore, as illustrated in Figure 3 and Figure 4, we pro- vide practical guideline for selecting the uncertainty score threshold in real-world applications.
While this work primarily addresses inherent label uncer- tainty, real-world clinical data present broader challenges, including missing labels, latent confounders, and comorbidi- ties. Tackling these complexities is essential for developing more generalizable and clinically grounded DTR frame- works. Future research also build upon our approach to alternative error control notions beyond FDR, further improving the robustness and safety of treatment recom- mendations.
A Acknowledgements
This work was partially supported by the National Institutes
of Health (NIH) under grant numbers UOITR003709,
U24MH136069, R56AG074604,
RFI
RO1 R0174333433,
6.5. Ablation Study
ROILM013519, ROIDK 128237, R21A1167418, and R21EY034179,
ROILM014344
We conduct ablation studies to validate the contribution of key components in SAFER, demonstrated in Table 2.
and by the National Science Foundation (NSF) under grant
numbers IIS 100-2006387 and II4-294.799.
The first variant, SAFER-F, removes the risk-aware fine- tuning process, directly using the initial prediction module for inference. This leads to notable performance degrada- tion, especially in Macro-AUC and counterfactual mortality rate on MIMIC-IV, which has high proportion of deceased patients. This confirms our hypothesis that accounting forla- bel uncertainty is crucial for robust decision-making.
Impact Statement
Dynamic treatment regimes (DTRs) play crucial role in precision medicine by enabling personalized and adaptive treatment plans that have the potential to significantly im- prove patient outcomes. Although wide range of DTR approaches show great promise through tailored treatment strategies based on patient responses, their application in high-stakes clinical settings necessitates rigorous and re- sponsible implementation.
For the second variant SAFER-N, we remove clinical notes, but relying only on structured EHR data. Performance dete- riorates significantly across all metrics, demonstrating that structured data alone fails to capture essential contextual cues from clinical narratives, highlighting the importance of textual information in modeling temporal dependencies.
This work emphasizes the critical responsibility of the re- search community to ensure safety, ethical standards, and tangible benefits to patient care when advancing such tech- nologies. SAFER addresses the inherent unreliability in clinical data by incorporating uncertainty quantification and mitigating prediction uncertainty, all while providing statis- tical guarantees. By controlling the false discovery rate in treatment recommendations, our approach safeguards pa- tient trust and ensures that advancements in DTR do not come at the expense of patient safety or ethical integrity.
Finally, SAFER-U removes the fine-tuning step and mit- igates uncertainty by training only on surviving patients. Even within the survival subset, its performance remains inferior to SAFER, with a marked decline in counterfactual mortality rate. These findings emphasize the importance of incorporating deceased patients in training and validate our approach to handling the uncertainty they introduce.
7. Conclusion
We have introduced SAFER, an end-to-end multimodal DTR framework that delivers reliable treatment recom- mendations with uncertainty quantification and theoretical guarantees. Compared with existing DTR frameworks, we provide a solution that may be more suitable high-stakes scenarios, ensuring safer and more trustworthy decision- making. It outperforms SOTA baselines across multiple recommendation metrics while achieving the greatest reduc-

---

Risk-Aware Dynamic Treatment Regimes
References
Chapfuwa, P., Assaad, S., Zeng. S., Pencina, M. Carin, L.,, and Henao, R. Enabling counterfactual survival analysis with balanced representations. In Proceedings of the Conference on Health, Inference, and Learning, pp. 133- 145, 2021.
Alsentzer, E., Murphy, J., Boag, W., Weng, W.-H., Jindi, D., Naumann, T., and McDermott, M. Publicly available clinical bert embeddings. In Proceedings of the 2nd Clinical Natural Language Processing Workshop, pp. 72- 78, 2019.
Chen, Y., Lasko, T. A., Mei, Q.. Denny, J. C., and Xu. H. A study of active learning methods for named entity recog- nition in clinical text. Journal of biomedical informatics, 58:1 18, 2015.
Assale, M., Dui, L. G., Cina, A., Seveso, A., and Cabitza, The revival of the notes field: leveraging the unstructured content in electronic health records. Frontiers in medicine, 6:66, 2019.
Cheng, L., Shi, Y., and Zhang, K. Medical treatment migra- tion behavior prediction and recommendation based on health insurance data. World Wide Web, 23:2023-2042, 2020.
Bajor, J. M. and Lasko, T. A. Predicting medications from diagnostic codes with recurrent neural networks. In Inter- national conference on learming representations, 2017.
Ching, T., Himmelstein, D. S., Beaulieu-Jones, B. K., Kalinin, A. A., Do, B. T., Way, G. P., Ferrero, E., Agapow, P.-M., Zietz, M., Hoffman, M. M., et al. Opportu- nities and obstacles for deep learning in biology and medicine. Journal of The Royal Society Interface, 15 (141):20170387, 2018.
Bates, S., Candès, E., Lei, L., Romano, Y., and Sesia, M. Testing for outliers with conformal p-values. The Annals of Statistics, 4991)1149-178, 2023.
Benjamini, Y. and Hochberg, Y. Controlling the false dis- covery rate: a practical and powerful approach multiple testing. Journal of the Royal statistical society: series B (Methodological), 57(1):289-300 1995.
Choi, E., Bahadori, M. T., Sun, J., Kulas, J., Schuetz, A., and Stewart, W. Retain: An interpretable predictive model for healtheare using reverse time attention mechanism. Advances in neural information processing systems, 29, 2016.
Benjamini, Y. and Hochberg, Y. Multiple hypotheses testing with weights. Scandinavian Journal of Statistics, 24(3): 407-418, 1997.
Choi, E., Bahadori, M. T., and Sun, Gram: Graph-based attention model for healthcare representation learning. Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pp. 787-795, 2017.
Benjamini, Y. and Yekutieli, D. The control of the false dis- covery rate in multiple testing under dependency. Annals of statistics, pp. 1165-1188 2001.
Bian, C., Yuan, C., Wang. J., Li, M., Yang, X., Yu, S., Ma, K., Yuan, J., and Zheng, Y. Uncertainty-aware domain alignment for anatomical structure segmentation. Medical Image Analysis, 64:101732, 2020.
Chua, M., Kim, D., Choi, J., Lee, N. G., Deshpande, V., Schwab, J., Lev, M. H., Gonzalez, R. G., Gee, M. S., and Do, S. Tackling prediction uncertainty in machine learning for healthcare. Nature Biomedical Engineering, 7(6):711-718 2023.
Bica, I. Alaa, A. M., Jordon, J., and Van Der Schaar, M. Estimating counterfactual treatment outcomes over time through adversarially balanced representations. Advances in Neural Information Processing Systems, 3351 11969, 2020.
Dorta, G., Vicente, S., Agapito, L., Campbell, N. D., and Simpson, I. Structured uncertainty prediction networks. Proceedings of the IEEE conference on computer vi- sion and pattern recognition, pp. 5477-5485, 2018.
Bothe, M. K., Dickens, L., Reichel, K., Tellmann, A., Ellger, B., Westphal, M., and Faisal, A. A. The use of reinforce- ment learning algorithms meet the challenges of artificial pancreas. Expert Review of Medical Devices, 10 (5):661-673, 11661-673,2013 2013.
Efron, B. Large-scale inference: empirical Bayes methods for estimation, testing, and prediction, volume 1. Cam- bridge University Press, 2012.
Gangavarapu, T., Krishnan, G. S., Kamath, S., and Je- ganathan, J. Farsight: long-term disease prediction using unstructured clinical nursing notes. IEEE Transactions on Emerging Topics in Computing, 93)::1151-1199 2020.
Chakraborty, B. and Moodie, E. E. Statistical methods for dynamic treatment regimes, volume 2. Springer, 2013.
Chakraborty, B., Laber, E. B., and Zhao, Y.QQ. Inference about the expected performance of a data-driven dynamic treatment regime. Clinical Trials, 1 11(4)408 417, 2014.
Gao, Z., Liu, X., Kang, Y., Hu, P., Zhang, X., Yan, W., Yan, M., Yu, P., Zhang, Q., Xiao, W., et al. Improving the prog- nostic evaluation precision of hospital outcomes for heart
10

---

Risk-Aware Dynamic Treatment Regimes
failure using admission notes and clinical tabular data: Multimodal deep learning model. Journal of Medical Internet Research, 26:e54363, 2024.
Johnson, A. E., Pollard, T. J., Shen, L., Lehman, L.-w. H., Feng. M., Ghassemi, M., Moody, B., Szolovits, P., An- thony Celi, L., and Mark, R. G. Mimic-iii, freely ac- cessible critical care database. Scientific data, 3(1):1-9, 2016.
Gouk, H., Frank, E., Pfahringer, B., and Cree, M. Regu- larisation of neural networks by enforcing lipschitz conti- nuity, Machine Learning, 110:393-416, 2021.
Johnson, A. E., Bulgarelli, L.,, Shen, L., Gayles, A., Sham- mout, A., Horng, S., Pollard, T. J., Hao, S., Moody, B., Gow, B., et al. Mimic-iv, a freely accessible electronic health record dataset. Scientific data, 10(1):1, 2023b.
Gui, Y., Jin, Y., and Ren, Z. Conformal alignment: Knowing when to trust foundation models with guarantees. arXiv preprint arXiv:2405. 10301, 2024.
Harakeh, A., Smart, M., and Waslander, S. L. Bayesod: A bayesian approach for uncertainty estimation in deep object detectors, In 2020 IEEE International Conference on Robotics and Automation (ICRA), pp. 87-93. IEEE, 2020.
Komorowski, M., Celi, A., Badawi, O., Gordon, C., and Faisal, A. A. The artificial intelligence clinician learns optimal treatment strategies for sepsis in intensive care. Nature Medicine, 24(11):1116-1720, 2018.
Kosorok, M. R. and Laber, E. Precision medicine. Annual review of statistics and its application. 6(1):263-286. 2019.
Hastie, T. The elements of statistical learning: data mining, inference, and prediction, 2009.
Hochreiter, S. and Schmidhuber, J. Long short-term memory. Neural computation, 9(8):1735-178) 1997.
Kusner, M. J., Loftus, J., Russell, C., and Silva, R. Counter- factual fairness. Advances neural information process- ing systems, 30, 2017.
Hu, J., Bao, R., Lin, Y., Zhang. H., and Xiang, Y. Accurate medical named entity recognition through specialized nlp models. arXiv preprint arXi::2412.08255. 2024.
Laber, E. B., Lizotte, D. J., Qian, M., Pelham, W. E., and Murphy, A. Dynamic treatment regimes: Technical challenges and applications. Electronic journal of statis- tics, 8(1 3(1): (111):1222-12777 2014.
Huang, H., Zheng, O., Wang. D., Yin, J., Wang, Z., Ding, S., Yin, H., Xu, C., Yang, R., Zheng, Q., et al. Chatgpt for shaping the future of dentistry: the potential of multi- modal large language model. International Journal of Oral Science, 15(1):29, 2023.
Laine, J. E., Baltar, V. T., Stringhini, S., Gandini. M., Chadeau-Hyam, M., Kivimaki, M., Severi, G., Perduca, V., Hodge, A. M., Dugué, P.-A., et al. Reducing socio- economic inequalities in all-cause mortality: counter- factual mediation approach. International Journal of Epidemiology, 49(2):497-51) 2020.
Hussein, A. S., Omar, W. M., Li, X., and Ati, M. Accurate and reliable recommender system for chronic discase diagnosis. Global Health, 3(2):113-111 2012.
Jiang, Y., Yu, W. Song, D., Cheng, W., and Chen, H. In- terpretable skill learning for dynamic treatment regimes through imitation. In 2023 57th Annual Conference on Information Sciences and Systems (CISS), pp. 1-6. IEEE, 2023.
Lei, J., G'Sell, M., Rinaldo, A., Tibshirani, R. J., and Wasserman, L. Distribution-fri predictive inference for regression. Journal of the American Statistical Associa- tion, 113(523):1094-1111 2018.
Liang, Z., Sesia, M., and Sun, W. Integrative conformal p-values for -utto--distribution testing with labelled out- liers. Journal of the Royal Statistical Society Series Statistical Methodology, pp. qkad138, 2024.
Jin, B., Yang, H., Sun, L., Liu, C., Qu, Y., and Tong, J. A treatment engine by predicting next next-period prescriptions. In Proceedings of the 24th ACM SIGKDD international conference knowledge discovery & data mining, pp- 1608-1616, 2018.
Lin, Z., Trivedi, S., and Sun, J. Generating with confidence: Uncertainty quantification for black- box large language models. arXiv preprint arXiv:2305. 11005.999880 19187, 2023.
Jin, Y. and Candès, J. J. Selection by prediction with con- formal p-values. Journal of Machine Learning Research, 24(244):1-41, 2023.
Liu, K., Price, B. L.,, Kuen, J., Fan, Y., Wei, Z., Figueroa, L., Geras, K. J., and Gernandez-Grandee C. Uncertainty- aware fine-tuning of segmentation foundation models. In The Thirty-eighth Annual Conference on Neural Informa- tion Processing Systems.
Johnson, A., Pollard, T., Horng S., Celi, L. A., and Mark, R. Mimic-iv-note: Deidentified free-text clinical notes (version 2.2), 2023a. URL https: doi. weooo.org/ii org/10. 13026/1n74-net

---

Risk-Aware Dynamic Treatment Regimes
Liu, K., Price, B., Kuen, J., Fan, Y., Wei, Z., Figueroa, L.., Geras, K., and Fernandez-Granda, C. Uncertainty- aware fine-tuning of segmentation foundation models. Advances in Neural Information Processing Systems, 37: 53317-53389, 2024.
Robins, J. M., Hernan, M. A., and Brumback, B. Marginal
structural models and causal inference in epidemiology. Epidemiology, 1165):50-550, 2000.
Rosenbloom, S. T., Denny, C.,, Xu, H., Lorenzi, N., Stead, W. W., and Johnson, K. B. Data from clinical notes: perspective the tension between structure and flexi- ble documentation. Journal of the American Medical Informaties Association, 18(2):181-186, 2011.
Luo, Z., Pan, Y., Watkinson, P., and Zhu, Position: rein- forcement learning in dynamic treatment regimes needs critical reexamination. 2024.
Rubin, D. B. Randomization analysis of experimental data: The fisher randomization test comment. Journal of the American statistical association, 75(371):591-59) 1980.
Lyu, W., Dong, X., Wong, R., Zheng, S., Abell-Heaa Hart, K., Wang, F., and Chen, C. A multimodal transformer: Fus- ing clinical notes with structured ehr data for interpretable in-hospital mortality prediction. In AMIA Annual Sympo- sium Proceedings, volume 2022, pp. 719, 2023.
Saria, S. Individualized sepsis treatment using reinforce- ment learning. Nature Medicine, 24(11)1641-1141 2018.
Melnychuk, V., Frauen, D., and Feuerriegel, S. Causal transformer for estimating counterfactual outcomes. In International conference on machine learning, pp. 15293- 15329. PMLR, 2022.
Schulam, P. and Saria, S. Reliable decision support using counterfactual models. Advances in Neural Information Processing Systems, 30, 2017.
Moodie, E. E.,, Richardson, T. S., and Stephens, D. A. De- mystifying optimal dynamic treatment regimes. Biomet- rics, 63(2):447-455 2007.
Shafer, G. and Vovk, V. A tutorial on conformal prediction. Journal of Machine Learning Research, 9(3), 2008
Shang. J., Ma, T., Xiao, C., and Sun, J. Pre-training of graph augmented transformers for medication recommendation. arXiv preprint arXiv:1906.00344, 2019a.
Murphy, S. A. Optimal dynamic treatment regimes. Jour- nal of the Royal Statistical Society: Series B (Statistical Methodology), 65(2):331-335 2003.
Shang, J., Xiao, C., Ma, T., Li, H., and Sun, J. Gamenet: Graph augmented memory networks for recommending medication combination. In proceedings the AAAI Con- ference on Artificial Intelligence, volume 33, pp. 1126-
Murphy, S. A. A generalization error for q-learning, Journal of Machine Learning Research, 10977 2005.
Olawade, B., David-Olawade, A. C.,, Wada, O. Z., Asaolu, A. J., Adereni, T., and Ling, J. Artificial intelligence in healthcare delivery: Prospects and pitfalls. Journal of Medicine, Surgery, and Public Health, pp. 100108, 2024.
1133, 2019b.
Shao, M., Qiao, Y., Meng, D., and Zuo, W. Uncertainty- guided hierarchical frequency domain transformer for im- age restoration. Knowledge-Basse Systems, 263: 10306, 2023.
Peng, X., Li, B., Zhao, Y., Zhang. Y.. He, F., and Jiang, T. Gbert: A pre-trained graph-based transformer for modeling drug-drug interactions. Bioinformatics, 37(2): 173-181, 2021.
Sheikhalishahi, S., Miotto, R., Dudley, J. T., Lavelli, A., Rinaldi, F., Osmani, V., et al. Natural language processing of clinical notes chronic diseases; systematic review. JMIR medical informatics, 7(2)ee12239 2019.
Raghu, A., Komorowski, M., Celi, L. A., Szolovits, P., and Ghassemi, M. Continuous state-space models for optimal sepsis treatment- deep reinforcement learning 2017.
Shi, X., Pan, Z., and Miao, W. Data integration in causal in- ference. Wiley Interdisciplinary Reviews: Computational Statistics, 15(1):e151,, 2023.
approach. arXiv preprini
Ren, A. Z., Dixit, A., Bodrova, A., Singh, S., Tu, S., Brown, N., Xu, P., Takayama, I.., Xia, F., Varley, J., et al. Robots that ask for help: Uncertainty alignment for large lan- guage model planners. arXiv preprint arXivv2307.01928. 2023.
Singer, M., Deutschman, C. S., Seymour, C. W., Shankar- Hari, M., Annane, D., Bauer, M., Bellomo, R., Bernard, G. R., Chiche, J.-D., Coopersmith, C. M., et al. The third international consensus definitions for sepsis and septic shock (sepsis-3). Jama, 315(8):801-810, 2016.
Robins, M. A new approach to causal inference in mortal- ity studies with a sustained exposure period applicatii to control of the healthy worker survivor effect. Mathe- matical modelling, ((-lllinng77-1))):)) 7(9-12)1393--1111 1986.
Sriperumbudur, B. K., Fukumizu, K., Gretton, A., Schölkopf, B., and Lanckriet, G. R. On integral proba- bility metrics, \phi- pbi-divergences and binary classification. arXiv preprint arXiv:0901.260) 2009.
12

---

Risk-Aware Dynamic Treatment Regimes
Su, Y., Shi, Y., Lee, W., Cheng, L., and Guo, H. Tahdnet: Time-aware hierarchical dependency network for medica- tion recommendation. Journal of Biomedical Informatics, 129:104069, 2022.
Wang, Y., Chen, W., Pi, D., Yue, L., Wang, S., and Xu, M. Self-supervised adversarial distribution regularization for medication recommendation. In IJCAI, pp. 3134-3140, 2021b.
Sun, Y., Wang, B., Sun, Z., and Yang. X. Does every data instance matter? enhancing sequential recommendation by eliminating unreliable data. In IJCAI, pp. 1579-1585, 2021.
Wu, R. Qiu, Z., Jiang, J., Qi, G., and Wu, X. Conditional generation net for medication recommendation. In Pro- ceedings of the ACM Web Conference 2022, pp. 935-945, 2022.
Suo, Q. Ma, F., Yuan, Y., Huai, M., Zhong, W., Zhang, A., and Gao, J. Personalized disease prediction using a enn-based similarity learning method. In 2017 IEEE Inter- national Conference on Bioinformatics and Biomedicine (BIBM), pp. 811-816. IEEE, 2017.
Xia, Y., Yang. D., Yu, Z., Liu, F., Cai, J., Yu, L., Zhu, Z., Xu, D., Yuille, A., and Roth, H. Uncertainty-awan multi-view co-training for semi-supervised medical image segmenta- tion and domain adaptation. Medical image analysis, 65: 101766, 2020.
Tan, Y., Kong, C., Yu, L.,, Li, P., Chen, C., Zheng, X., Hertzberg, V. S., and Yang, C. 4sdrug: Symptom-bas based set-to-set small and safe drug recommendation. In Pro- ceedings of the 28th ACM SIGKDD Conference on Knowl- edge Discovery and Data Mining, pp- 3970-39900 2022.
Yang, C., Xiao, C., Ma, F., Glass, L., and Sun, J. Safe- drug: Dual molecular graph encoders for recommending effective and safe drug combinations. arXiv preprint arXiv:2105.0277 022111 2021.
Ye, Y., Tang, L.-A., Wang, H., Yu, R., Yu, W., He. E., Chen, H., and Xiong, H. Pail: Performance based adversarial imitation learning engine for carbon neutral optimization. In Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 6148-6157, 2024.
Tsiatis, A. A., Davidian, M., Holloway, S. T., and Laber, E. B. Dynamic treatment regimes: Statistical methods for precision medicine. Chapman and Hall/CRC, 2019.
Tsybakov, A. B. and Tsybakov, A. B. Nonparametric esti- mators. Introduction to Nonparametrie Estimation, pp. 1-76, 2009.
Ye, Y., Zheng, Z. Shen, Y., Wang, T. Zhang, H., Zhu, P., Yu, R.. Zhang, K., and Xiong, H. Harnessing multi- modal large language models for multimodal sequential recommendation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pp. 13069-13077, 2025.
Valeri, L., Chen, T., Glbeiz-, X., Krieger, N., VanderWeele, T. J., and Coull, B. A. The role of stage at diagnosis in colorectal cancer black-white survival disparities: a counterfactual causal inference approach. Cancer Epidemiology, Biomarkers & Prevention, 25(1): 83-89, 2016.
Yu, C., Liu, J., Nemati, S., and Yin, G. Reinforcement learn- ing in healthcare: A survey. ACM Computing Surveys (CSUR), 55(1):1-36, 2021.
Vovk, V., Gammerman, A., and Saunders, C. Machine- learning applications of algorithmic randomness. 1999.
Vovk, V., Gammerman, A., and Shafer, G. Algorithmic learning in a random world, volume 29. Springer, 2005.
Zhang, A., Xing, L., Zou, J., and Wu, J. C. Shifting machine learning for healthcare from development to deployment and from models to data. Nature Biomedical Engineering, 6(12):1330) (112):3300-145, 2022.
Wang, L., Zhang, W., He, X., and Zha, H. Supervised rein- forcement learning with recurrent neural network for dy- namic treatment recommendation. In Proceedings of the 24th ACM SIGKDD international conference on knowl- edge discovery & data mining, pp. 2447-2456 2018
Zhang, Y., Chen, R., Tang, J., Stewart, W. F., and Sun, J. Leap: learning to prescribe effective and safe treatment combinations for multimorbidity. In Proceedings of the 23rd ACM SIGKDD International Conference on Knowl- edge Discovery and Data Mining, pp. 2017.
Wang, L., Yu, W., He, X., Cheng, W., Ren, M. R., Wang, W., Zong, B., Chen, H., and Zha, H. Adversarial cooperative imitation learning for dynamic treatment regimes. In Pro- ceedings of The Web Conference 2020, pp. 1785-1795, 2020.
ACM,
Wang, Y., Chen, W., Pi, D., and Yue, L. Adversarially reg- ularized medication recommendation model with multi- hop memory network. Knowledge and Information Sys- tems, 63: 145- 2021a.
13

---

Risk-Aware Dynamic Treatment Regimes
A. Technical Proofs
A.1. Lower-Bound Bound Proof for Theorem 4.1.
Restate of Theorem 4.1. We have two latent-representationn distributions (h) and pt (+) over the same latent space H. corresponding to deceased and surviving patients, respectively. Let h P (h) and h+ p+ (h). We define the predictive uncertainty
<(h) DKL. (pe(y hb) poly ht):
where po the "teacher" module's distribution and ps is the "student" module's distribution (trained with risk-aware fine-tuning). The theorem claims that under:
1. PKL (P- (1) (k) > 0. P Pt.
2. fo is I--Lipschitz continuous over H (so po h) cannot sharply change as h varies),
we have
P- P- [x(h]] > Eh [s(h)] with strictly positive lower bound.
Proof of Theorem 4.1. Since DKL (P - || Pt) > 0, there must be at least one measurable subset 9 C H on which P places strictly greater mass than p+ (or vice versa). Concretely, there exists such that
(P
- 2 E.
Since no such region existed, p- would equal p+ almost everywhere, contradicting DKL(P-IPt)) > 0.
Also as we have
<(h) DKL (poly |h) Il (
If P& is L-Lipschitz in h, then h varies within small neighborhood, the entire predicted distribution po(y h) cannot drastically jump to match (v h) perfectly, unless the underlying latent distributions P=, are aligned. Since p- * Pt, there is region G in latent space where Ps cannot "annihilate" the mismatch in po- That is on measurable set G C. H. the teacher predictions differ from the student by at least fixed amount:
l h) - P MI]l, 2 5 for all h E G,
(2)
with constants > 0. Hence, k(h) is bounded away from zero on some portion of G with nontrivial measure under P-,
B [x(h)] = {
(h) 2 | [((M) (h) 2 p (G).
(6)
Where the second inequality follows from Pinsker's inequality which stated in many standard results (e.g., (Sriperumbudur et al., 2009; Tsybakov & Tsybakov, 2009)).
Split the survivor expectation into the same region & and its complement:
En-Pi [x(hh] / x(h)) dP+(h) +
/ <(()) (P))PP(()) (h).
Hig
(8)
(b)
For the first term (a), we can control by Lipschitzness. Fix E 9 and pick // with p+ P+-density such that d(h, h/) S r for radius > 0 (possible because supp(Pt) = in practice). Applying assumption and again Pinsker's inequality,
<<h) S DKL. (po p.( R\)) + CL²,².
where O is an absolute constant. Choosing r small makes this term negligible compared with 8. Denote the resulting bound by E For the second term (b), as the global survivor risk can be tuned most e, therefore (b) S c.
14

---

Risk-Aware Dynamic Treatment Regimes
Collecting the two parts we have
BANPi [k<(h)] S E 61-
Subtract this from the lower bound of En- [r][]]]]
1 (9)) ENEP [x(h]] ERoPi [x(h)] 2
2
2 g (5 (0+
Where == (9) that E1 are small enough, hence >> O.
PG) P+ from the first assumption and 6 > 0, can easily pick training and regularisation so
Putting it together,
P [<<(h)] - PE----- [x(1)] 2 c
with explicit constant c= 1/8ºT E1. This establishes that deceased -patien latents (drawn from P-) systematically lead to higher KL-based uncertainty x(-) than do survivor latents from P+, giving strict positive gap on average. Hence the theorem's 's statement follows.
Remark A.1. By forcing Ps to remain continuous with respect to h, we guarantee that latent embeddings deceased patients differ significantly from those of survivors, the student's predicted distributions cannot "collapse" to match the teacher's everywhere in H. Hence, the KL uncertainty k(h) for h- ~ p stays measurably larger on average than for h+ P+, ensuring Ep- [ks EEp [x]] by at least a positive margin.
Remark A.2. å can be estimated on a validation split by the empirical minimum teacher-student & gap over clusters; æ follows from any two-sample test on the latent representations; is the held-out risk of the student model on survivors.
Remark A.3. Tighter bounds. One may replace Pinsker's s inequality by the Bretagnolle- Huber inequality or by Csiszár- Kullsack-Pinsker to sharpen c. The qualitative conclusion of strict positivity remains unchanged.
Remark A.4. If the student is Bayes- for p+ (in the limit of infinite positive data) and the teacher is Bayes-optimal for the mixture,then Ep+ [x] = 0 exactly, and the proof simplifies to analysing P- only, yielding c = - (g).
A.2. FDR Control Proof of Theorem 5.1
Restate of Theorem 5.1. Let k: X' + IR be an uncertainty function, and sup, <(x) < M for some M 2 0. We have i.i.d. calibration samples ((x;, }/ and m i.i.d. test inputs "xn+s)m= all mutually independent in the sense that any subset excluding index j is jointly independent of the data at index j. For each test point j we define null hypothesis
H3 Katj 2 e,
and the conformal p-value Py as (8), namely
Ei= 1 1
R
Rnts.
c
+1
Ug (1 Er L/R +
Rutss K4 2 2c)
Dj
+
n
where U,- ~ Un, 1) are i.i.d. tie- brebbraking variables. Let the Benjamini Benjamini-Hochberg (BH) procedure at level (0, 1) be applied to {py)y.. producing a selection (rejection) set S. Denote R = |S] and
V
Σ
I j E S, S, is true)
the number of false rejections. Then under the above assumptions, the false discovery rate (FDR) is
V
FDR = E
aa
max(1, R)
15

---

Risk-Aware Dynamic Treatment Regimes
Proof of Theorem 5.1. We define the nonconformity score J as
J(x,)) x(x) + 2M 1 y 2 e).
Thus, y < c (i.e., if the label is "below" the critical threshold c), J(x, y) =(x). On the other hand, if y 2 then J(x, v) x(x) + 2M. The nonconformity score J(x, y) preserves the monotonicity property in terms of y. Thus if we define Ji = (x., v.) J Jxx,, c) for every E [n m]] the conformal p-value defined in Equation (8) converts
E [(J. J, Jnts] n +
(1
Σ Di=1 Rf
J. Ints }
n1
as defined in Jin & Candès (2023). To ensure the completeness of the proof, we adapt major proof procedures from Theorem in Jin & Candès (2023) as following.
Using ... in standard conformal scheme (Vovk et al., 2005; Lei al., 2018), we obtain p-values P1,.-- Pm of the Equation (8) (including Uj for tie-breaking). By exchangeability of calibration and test data, plus the monotonicity of J, each Py is "selectively super-unde unperruniform" with respect to its null Hy (Jin & Candès, 2023); that is, for every E [0,
P(UES) A (p. sa) s a.
Roughly, this property ensures that pj behaves conservatively if Hj is true. Morcover, one typically invokes PRDS condition (Positive Regression Dependence on Subset) or mutual independence across {p,} ensure the BH procedure be applied with classical guarantees (Benjamini & Yekutieli, 2001; Efron, 2012).
Then under i..d. sampling {(x,, 34) )= plus (xn+j}ī= and monotonic J, the random variables Pl....,PM Pm exhibit either independence or positive correlation that meets PRDS assumptions (see, e.g., Bates et al., 2023; Jin & Candès, 2023). Hence, each true null label j effectively satisfies P, ~ selective super-uniform with respect to Hy-
Finally, we show FDR < a once BH is applied to the p-values (pI....,PPl at level a. Let = { Jj:Pp S P(a) denote the BH rejection set, where
ko max{ r: P(r) m
Define indicator random variables Rg = 1 E S} and T3 = 1HH is true
Then
FDR
プラ Tj R2
1, ΣØ= R ]
E max (1, E, R
Σ T, Rg
m
max
From the property of BH under PRDS super-uniform p-values (Benjamini & Hochberg, 1995; Benjamini Yckutieli, 2001; Bates al., 2023; Jin & Candès, 2023), we have
E T Rg] S a [R,)
summing over and employing the usual BH bounding technique, it follows that
Em T3 Rg I.
FDR
E
<
R5J
In short, the fraction of wrongly rejected true nulls among all rejections remains at or below a in expectation. This completes the proof.
B. Counterfactual Mortality Calculation
Suppose M, represents the mortality event, (y) refers to the potential outcome under the treatment arm Y, = and X, E Txxx is the measured counfounders for the i-th patient. To estimate the effectiveness of our recommended treatment plans, we evaluate the model with the reduction in counterfactual mortality rate, we train an additional LSTM-based neural network as a counterfactual mortality prediction model. During training, this model takes patient features and ground-trutt
16

---

Risk-Aware Dynamic Treatment Regimes
treatments as input and is optimized using Binary Cross Entropy (BCE) loss predict the probability of death as binary classification task. During inference, the trained model estimates counterfactual mortality rate by applying the model to patient features combined with the recommended treatment (i.e., the treatment predicted by our DTR model). The decrease in mortality rate is then defined as the difference between this estimated counterfactual mortality and the actual observed mortality rate under standard clinical practice. This analysis is based on the following assumptions regarding potential outcomes:
(B1) No interference. M,,), depends only on Y,-
(B2) No hidden variability, Each unit has unique (v).
(B3) Ignorability/ / No Unmeasured Confounding, Conditioned on the measured covariates X,, the potential outcomes (y) are independent of the assigned treatment. That is,
{ [M(())) 1 Yi X,
(B4) Positivity (Overlap). Every treatment arm y has a nonzero probability of being assigned given X,, P(Y; y X,) > 0 for all y E y.
We maintain the classic stable unit treatment value (SUTVA) assumption (Rubin, 1980) that no interference between units in (B1) and no hidden variations of treatments oceur in (B2), patient i actually receives treatment y. then the observed mortality M, coincides with the potential outcome Mi(y), which allows to assume that M En Y (v) almost surely. (B3) is standard assumption the ignorability of treatment assignment (Shi et al., 2023). Together, these assumptions allow us to view (y) as well-defined counterfactual, enabling estimation of counterfactual mortality under different model-predicter treatments. Under these assumptions, we have p (X) = E{M Y y. X}, where (X) (y) X).
17

---

Risk-Aware Dynamic Treatment Regimes
C. Experiment Details
C.1. Dataset Details
Dataset Statistical Table shows the statistical details of the two cohort datasets we use.
Dataset
#Survival
#Deceased
#Avg Len
#Avg Notes
MIMIC-III
3118
427
11.75
40.65
MIMIC-IV
19450
3786
11.50
29.08
Table 3. The dataset statistics,
Attribute Table 4 presents the attributes used in experiments.
Attribute Type
Attribute Name
Demographics
Gender, Age, Re_admission, Weight.kg. Height_cm
GCS, RASS, HR, SysBP. MeanBP, DiaBP, RR, Temp_C, CVP, PAPsys. PAPmean, PAPdia, CL. SVR FIO2.1, 02flow, PEEP, TidalVolume, MinuteVenntll PAWmees mean, peak PAWplateau, Potassium, Sodium, Chloride, Glucose, BUN, Creatinine, Magnesium, Calcium, SGOT, SGPT, Total_bili, Direct_bili, Total_protcin, Albumin, Troponin, CRP, Hb, Ht, RBC count, WBC count, Platelets count, PTT, PT, ACT, INR, Arterial pH. paO2, paCO2, Arterial BE, Arterial lactate, HCO3, ETCO2, SvO2, mechvent, extubated, Shock Index, PaO2 FiO2, SOFA, SIRS
Vital Signals
Table altribute used in experiments.
The label frequency on suvivor subset. Since we evaluate most metrics on the survivor subset, we also report the label frequency on the survivor subset Figure shows. our experiments, we follow established protocols from prior sepsis treatment studies (c.g., Komorowski et al., 2018) by discretizing the intravenous fluid and vasopressor dosages into bins each. Specifically, any absence of medication constitutes the zero bin, while the remaining dosages partitioned into four additional bins according to empirical quantiles. This results in a 5 x 5 grid, forming 25 discrete treatment classes, where each class corresponds to unique combination of fluid and vasopressor dosage levels (i.e., (fluid bin) (vasopressor bin)). The distribution of these treatment classes are visualized in Figure 2 of the manuscript.
Vesepresse Doge
Fluids
(a) MIMIC-III
(b) MIMIC-IV
Figure Comparative visualization the treatment frequency matrix log scale from the survivor subset two datasets. Panel (A) represents MIMIC-III, while Panel (B) corresponds to MIMIC-IV.
C.2. FDR Control
In this section, we present the FDR control results using Lincar Regression on the MIMIC-II dataset (Figure 6) and the MIMIC-IV dataset (Figures 7), further validating the effectiveness of our FDR control mechanism.
C.3. Parameter Sensitivity
In this section, we discuss the sensitivity of SAFER to on three key hyperparameters, the historical information sequence length L., the hidden dimensionality d,, and the risk regularization coefficient y in the loss function. We report Macro-AUC
18

---

Risk-Aware Dynamic Treatment Regimes
FDR and Powerr
0.4
0.6
FDX
FDR and
RBR
FIR
84
0.
0.4
9.2
Terget FOR level 6t
Target FDR level at
Target FOR leval at
Target FDR level at
(a) 0.1
(b) 0.2
(c) 0.3
(d) 0.4
Figure 6. FDR and power curves different target level with Linear Regression MIMIC-III
1.0
1,0
1.0
1.0
and
FOR
FDR
0.2
0.2
0.8
0.0
0.0
0.0
0.2
0.0
0.8
1.0
0.0
0.2
0.4
0.8
1.0
0.2
9.1
0.8
0.8
0.0
0.2
0.6
0.8
1.0
Target
Target FDR level
Target
Target FDR
(a) 0.1
(b) 0.2
(c) 0.3
(d) 0.4
Figure 7. FDR and power
different target level with Linear Regression MIMIC-IV.
and the reduction in counterfactual mortality rate, two core evaluation metrics that reflect recommendation accuracy and treatment effectiveness.
To ensure fair comparison and robustness, we perform grid search predefined range of values for each parameter while holding others fixed. The selected values correspond those that jointly optimize both performance metrics on the validation set.
Historical sequence length L: Figure illustrates that SAFER's performance improves significantly as the sequence length increases initially, before stabilizing a consistent level. This trend likely occurs because shorter sequences lack sufficient information for accurate predictions. To ensure a fair comparison across datasets, we set the sequence length to 8 for all experiments.
1.0
MIMIC-III MIMIC.IV
MIMIC-III MIMIC-IV
Mortality Rate
0.9
Macro AUC
0.4
0.9
0.3
0.7
0.2
0.6
0.1
0.5
0.0
9
5
9
Sequence Length
Sequence Length
AUC
(b) Mortality Rate
Figure 8. The performance of SAFER under different historical information seqeunce length L
Hidden dimensionality hy: Figure illustrates the performance of SAFER across different hidden dimensionalities has showing that both low and excessively high dimensions degrade model performance. A leads to information loss, while a higher dimension increases model complexity, making proper dimensionality selection crucial. Since the model's performance remains stable for ha 128, 256, 512, we choose 128 to reduce model parameters
representation
19

---

Risk-Aware Dynamic Treatment Regimes
0.95
0.90
MIMIC-III
MIMIC.IV IV
Mortality Rate
0.5
MIMIC IV
MIMIC
0.85
0.3
0.70
0.2
0.65
128
256
1024
Dimensionality hd
Dimensionality hd
AUC
(h) Mortality Rate
Figure 9, The performance of SAFER under different hidden dimensionality ha
0.5
MIMIC-II MIMIC-IV
MIMIC-III
Mortality Iirtality Rate
0.55
MIMIC-IV
Macro AUC
0.4
0.1
0.2
0.3
0.4
0.5
0.6
0.1
0.2
0.3
0.4
0.5
0.6
gamma
gamma
(a) AUC
(b) Mortality Rate
Figure The performance of SAFER under different y value in loss function
and improve computational efficiency.
y the Loss Function:
Figure illustrates performance of SAFER under different values, guiding the selection of an optimal y. A small y leads decreased performance, confirming the necessity of incorporating this penalty term. However, excessively large 7 is also detrimental, as it shifts the model's focus away from the supervised signal during training, ultimately reducing overall performance.
These findings support the stability of SAFER under moderate hyperparameter variation, and highlight the importance of risk-aware fine-tuning in achieving consistent improvements in both predictive accuracy and patient safety.
C.4. Case Study
Decessed
0.6
0.4
Timestamp
Figure 11. Trends in uncertainty scores over time.
To further illustrate the interpretability of SAFER's uncertainty estimates, we select a subset of 10 surviving and 10 deceased
20

---

Risk-Aware Dynamic Treatment Regimes
patients with comparable sequence lengths. At each timestamp, we predict the subsequent treatment target using varying historical windows and compute the corresponding uncertainty scores. As shown in Figure 11, the average uncertainty scores for surviving patients remain low and stable throughout the clinical timeline. In contrast, deceased patients exhibit distinct upward trend in uncertainty as their condition deteriorates.
This divergence reflects a growing difference in predictive stability between the two cohorts. For surviving patients, the model consistently maintains high confidence, likely due to regular disease progression and coherent treatment-response patterns. Conversely, the increasing uncertainty observed among deceased patients suggests a transition into more complex or irregular clinical dynamics, where prediction becomes inherently more difficult.
The elevated uncertainty in deceased trajectories may arise from two primary sources: (1) ambiguous treatment behaviors driven by rapid physiological decline or emergent interventions; and (2) limited representational coverage of similar deteriorating cases in the training distribution, resulting in increased epistemic uncertainty. These observations align with our core modeling assumption that treatment labels for deceased patients are more likely to be noisy or unreliable due to outcome ambiguity and clinical variability.
Notably, these temporal uncertainty trends provide form of counterfactual interpretability. By capturing the divergence in predictive confidence over time, SAFER not only differentiates between stable and high-risk trajectories but also offers potential mechanism for proactive clinical risk detection. In practice, this trajectory-level uncertainty signal could be leveraged to identify patients transitioning into unfamiliar or high-risk states, thereby enabling timely human intervention.
In summary, this analysis underscores the utility SAFER's 's uncertainty estimates as both a diagnostic and interpretive tool, particularly in htgh-stakes clinical environments where model trustworthiness and actionable risk awareness are essential.
21


Reviews:
[{"review_id": "zs2yGpyq55", "reviewer": "Reviewer_io89", "summary": "This paper introduced a framework, SAFER, to provide dynamic treatment recommendations for patients with evolving clinical states, by employing conformal prediction and transformer-based architectures for multi-modalities. The proposed work was evaluated on two sepsis datasets and outperformed baselines in terms of various evaluation metrics.\n\n## update after rebuttal\n\nMy main concerns were resolved after rebuttal. After reading other reviewers' comments, I'm positive with the contribution of this paper to the related domains.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for acknowledging that the paper’s main claims are clearly stated and supported by both theoretical analysis and experimental results, with well-motivated objectives and diverse evaluation metrics. We address your comments as follows.\n\n**Methods And Evaluation Criteria:**\n\nThank you for highlighting this important point. While SAFER is designed as a general framework for DTRs and is applicable to a broad range of clinical tasks, we acknowledge that our current evaluation focuses on sepsis cohorts. We chose sepsis cohorts for several reasons: (1) Sepsis is one of the most critical and prevalent conditions in ICU settings, accounting for the third leading cause of death worldwide and the main cause of mortality in hospitals. As the best treatment strategy for sepsis remains uncertain, the clinical complexity and inherent label uncertainty in sepsis management make it a strong testbed for evaluating the reliability and robustness of SAFER; (2) MIMIC sepsis cohorts are widely used and well-established in DTR research, enabling fair and reproducible comparisons with prior work.\n\nWe agree that the current experimental scope should be made more explicit. In the revised version, we will update the Introduction to clearly state the focus on sepsis as a case study and include a note in the Discussion section outlining plans to extend SAFER to broader clinical domains in future work.\n\n**Essential References:**\n\nThank you for the helpful suggestion. In response to multiple reviewers’ feedback, we have expanded the manuscript to cover key related domains:\n\n**Uncertainty Quantification Design.** Section 3.1 now cites relevant work on uncertainty-aware model design, also with a broader discussion of other uncertainty types in healthcare added in Section 2.\n\n**DTRs & Counterfactual Framework.** Foundational work is now referenced in Appendix B and Section 2 to contextualize our use of counterfactual mortality rates.\n\n**Multimodal Learning.** We expanded discussion of prior work on multimodal modeling and transformer-based fusion across domains beyond clinical notes and EHR.\n\n**Conformal Prediction for Error Control.** We have clarified our distinction by emphasizing that SAFER is the first to apply conformal prediction to selectively control treatment recommendation reliability. Prior healthcare-focused conformal work is now discussed in Section 5 and around line 43 in the revised version.\n\n**Other Comments Or Suggestions:**\n\nThank you for your insightful feedback. In response, we examined how uncertainty scores evolve for both surviving and deceased patients https://anonymous.4open.science/r/Rebuttal-B376/Trends.png, finding that while prediction confidence remains relatively stable for survivors, it steadily increases for those who eventually die, aligning with disease progression. This shows that SAFER’s uncertainty signals capture meaningful clinical dynamics and can provide interpretable insights into a patient’s trajectory. In the revised version, we will include case studies in more details highlighting these patterns and illustrating how uncertainty scores adapt over time, ultimately strengthening confidence in our model’s predictive accuracy and real-world applicability.\n\n**Questions:**\n\n**Q1**: Thank you for emphasizing the need for interpretability. Here are some interpretability strategies we will discuss in the revised version.\n\n1. We can generate attention heatmaps to identify the most influential time steps in clinical notes and EHR data for treatment decisions. Attention scores are extracted from the SelfAttnBlock and CrossAttnBlock, especifically from the last layer. By focusing on the last timestamp, we capture the model's attention at the final step, highlighting the timestamps the model considered most important. This helps clinicians understand when the model focused on critical information, such as specific features or clinical note snippets (e.g., \"rising lactate\" or \"acute hypotension\") that influenced its recommendation.\n2. Another aspect of interpretability comes from the uncertainty score itself which can provide prediction reliability for clinicians. Also as mentioned in the previous response, uncertainty score provides insights into deceased progression and offers a way to analyze disease trajectories. That case study result demonstrates that the model successfully learns these patterns, enhancing its interpretability in clinical contexts.\n3. We could provide an optimal treatment regime as a sequence of decision rules, where each rule at a given time step can be represented as an interpretable list of “if-then” statements. These rules evaluate the estimated mortality rates under different recommended treatments, making them directly understandable to domain experts. This approach is also suggested by Zhang et al. (2018) \"Interpretable dynamic treatment regimes.\" \n\nWe sincerely hope these clarifications improve your understanding and evaluation of our work.", "claims_and_evidence": "In general, the main claims made by this paper are clear and evidences are provided via both theoretical analysis and experiments including ablation studies. For example, ablation studies show removing textual data notably degrades performance, indicating that textual features capture clinical signals. Compared with variants that either omit risk-aware fine-tuning or exclude deceased trajectories altogether, SAFER shows higher AUC and greater reduction in counterfactual mortality rate. Those support me to trust the major claims given by the paper.", "methods_and_evaluation_criteria": "The motivations of methodological design are clearly discussed, and various evaluation criterias were included, such as classification-based metrics, ranking-based metrics, and counterfactual mortality rate for assessing recommended treatments.\n\nOne thing I feel need to be improved is that only sepsis datasets were used in this work to evaluate the proposed method, though I believe the proposed method should be a general approach towards various clinical tasks. The statements/claims on introduction should be clearly scoped regarding this issue.", "theoretical_claims": "The paper provided a selective-guarantee variant of conformal inference that controls the FDR for uncertainty from treatment recommendations. This is formalized with a proof leveraging exchangeability assumptions and BH procedures. The authors also prove that, under mild conditions, KL divergence of “teacher vs. student” models can reliably distinguish deceased vs. surviving patient trajectories. Overall the theoretical claims appear consistent.", "experimental_designs_or_analyses": "Two sepsis cohorts (MIMIC-III and MIMIC-IV) are used, which I think are high-stake and challenging. Diverse evaluation metrics were occupied to justify the effectiveness of proposed method on the two datasets.", "relation_to_broader_scientific_literature": "The work should be interesting to some broader domains such as dynamic treatment regimes, sepsis treatments, conformal prediction in applications, etc.", "essential_references_not_discussed": "Overall, all highly related domains should have been discussed in main context.", "other_strengths_and_weaknesses": "Other strengths:\n- The paper is well-written and polished.\n- The experiments are on well-known datasets, and the results are benchmarked against diverse baselines (both sequential embedding and RL-based).", "other_comments_or_suggestions": "It may be helpful to provide an analysis or case study on specific sepsis trajectories to demonstrate how SAFER’s uncertainty calibration differs across patients (e.g. a partial interpretability analysis).", "questions_for_authors": "- Interpretability is also an important topic that clinicians care about. The cross-attention and risk-aware modules produce final predictions, which are impressive, could you imagine and consider some convincing interpretability strategies that would likely boost clinical adoption and trust in the model for empirical usage?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "G8euAI9rtk", "reviewer": "Reviewer_pxZx", "summary": "In this work, the authors introduce uncertainty control and comprehensive information fusion to improve prediction uncertainty estimation while incorporating multi-modal data for more accurate predictions.\n\n## update after rebuttal\nThanks to the authors for clarifying some questions. Please make sure to include these clarifications in the updated version. Although I am not an expert in this field, I feel that the paper has strong motivation, theoretical support, and good writing. Therefore, I am inclined to accept the paper.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for recognizing the importance of this research direction, as well as appreciating the contributions related to multi-modal design and theoretical guarantees for uncertainty-aware treatment recommendation! To help better understanding, we would like to briefly restate our main contribution. \n\nIn high-stakes clinical decision-making, ensuring the reliability and safety of treatment recommendations is paramount. Our work introduces two key innovations to address these challenges in dynamic treatment regimes (DTRs). First, we propose a novel multimodal framework that, for the first time in DTR research, integrates structured electronic health records (EHR) with unstructured clinical notes—capturing richer patient context and significantly enhancing predictive accuracy. Second, we uniquely incorporate uncertainty quantification and conformal calibration into the DTR pipeline, enabling the model to filter out unreliable predictions while providing formal theoretical guarantees on error control. Together, these contributions establish a safer and more trustworthy foundation for clinical decision support in high-risk settings.\n\nBelow, we address your comments regarding notation and clarity. We appreciate your careful review and will revise the manuscript accordingly in the final version:\n\n- Overlapping Notation in Equations: We acknowledge the confusion around the symbol W in Equation 2 and in line 153. In the revised manuscript, we will use distinct symbols or subscripts (e.g., **W** still for the weight matrices in attention, and **O** for clinical notes) to clearly differentiate their meanings. This change will also be reflected in Figure 1 for consistency.\n- Undefined Symbol M: We will explicitly define **M** before Equation 2. Specifically, we will clarify that **M** refers to the causal (look-ahead) attention mask commonly used in Transformer architectures to prevent future information leakage. This ensures that predictions at a given time step are made without access to future observations. We will also add a brief reference to masked self-attention (Vaswani et al., 2017. “Attention Is All You Need”).\n- Unclear Terminology (Line 237): “Two modules” refers to (1) the initial prediction module $f_{\\theta}$ and (2) the refined prediction module $f_{\\phi}$ that takes the patient embeddings learned in the first module as input and trains exclusively on surviving patients. We use these “two modules” to calculate the uncertainty for treatment recommendation to help us identify unreliable predictions. We will revise the text to make this terminology more explicit and intuitive.\n\nThank you again for highlighting these issues. We will carefully revise the manuscript to address these issues and eliminate ambiguity. We sincerely hope these clarifications improve your understanding and evaluation of our work. Please feel free to reach out with any additional questions.", "claims_and_evidence": "I am not an expert in this field. I am not confident in checking the correctness of the claims.", "methods_and_evaluation_criteria": "Overlapping Notation in Equations\n\nSome terms in the equations appear ambiguous or overlapping. For example, in Equation 2, the symbol W seems to represent the weight matrix in attention, whereas W in line 153 appears to refer to clinical notes. To improve clarity, I suggest differentiating these notations by using distinct symbols or subscripts.\n\nUndefined Symbols\n\nThe meaning of M in Equation 2 is unclear. It would be helpful to provide a brief explanation or definition in the text to aid reader comprehension.\n\nUnclear Terminology in Line 237\n\nThe term “two modules” in line 237 is not explicitly defined. Could the authors clarify which two modules are being referenced? Providing a brief explanation would improve clarity and ensure readers fully understand the method's structure.", "theoretical_claims": "I am not an expert in this field. I am not confident in checking the correctness of the theoretical results.", "experimental_designs_or_analyses": "Yes, I have check the experiments part.", "relation_to_broader_scientific_literature": "Important Research Topic\n\nThe prediction of dynamic treatment regimes, particularly for understudied diseases and critically ill or deceased patients, is a highly important research area. Providing uncertainty estimates further enhances the reliability of these predictions.\n\nIncorporation of Multi-Modal Data\n\nThe authors effectively leverage both structured data and unstructured clinical notes, improving predictive power and broadening the applicability of the model in real-world clinical settings.\n\nTheoretical Guarantees for Error Control\n\nThe provision of theoretical guarantees for error control is particularly valuable in the medical domain, where reliability and interpretability are crucial.", "essential_references_not_discussed": "NA.", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "NA", "questions_for_authors": "NA", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HXO3PSH6El", "reviewer": "Reviewer_kw7r", "summary": "The paper proposes a novel method for predicting treatments for individuals based on the longitudinal EHR data, static features and text based clinical notes, calling it tabular-language recommendation framework. To improve the quality of predictions, it employs mechanisms such as risk-aware fine tuning which considers the unreliability in the labels for deceased patients (i.e., patients with negative outcomes) and conformal selection. The paper presents experiments on two clinical datasets - MIMIC III and MIMIC IV.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for appreciating the novelty, the overall framework for healthcare, clarity of Figure 1, well-supported experiments and the evaluation criteria of our work! We address your comments as follows.\n\n**Claims And Evidence:**\n\n**Theoretical guarantees on calibrated predictions lacks evidence:** We appreciate the reviewer’s observation and would like to clarify our claim. Our guarantee of “calibrated predictions” refers to the statistical control on the error rate in (ii) of the treatment recommendations based on conformal inference. This does not imply that the raw model predictions are inherently calibrated. Rather we calibrate the final selection of treatment recommendations through post-hoc conformal procedures, and **provide theoretical guarantees on FDR** (See Section 3.2 & 5).\n\n- The bound in (i) validates our KL-based uncertainty score as a meaningful discriminator between reliable and uncertain labels, but does not imply calibration.\n- The bound in (ii) provides the formal calibration guarantee for the selection procedure, ensuring that the expected proportion of incorrect recommendations remains below a user-defined threshold $\\alpha$.\n\nWe will revise accordingly to avoid ambiguity. \n\n**Experimental Designs:**\n\n1. **Decrease in mortality rate**: This measures the estimated reduction in mortality if patients had received model-recommended treatments instead of clinician-prescribed ones, i.e., the **difference between estimated mortality given recommended treatment and observed mortality.** It is a standard **counterfactual evaluation** approach widely used in clinical research to assess potential treatment benefit. We will clarify this in Appendix B.\n2. **Details for treatment classes**: As mentioned in the paper, treatments involve intravenous fluids and vasopressors, each discretized into 5 bins based on empirical quantiles, forming a 5×5 grid (25 treatment classes). We now make this setup more explicit.\n3. **Temporally aware split**: We understand your concern regarding potential **target leak** when handling temporal data. To clarify:\n    1. **Random Split are Based on Patients**: Patients are independent, and no patient appears across splits. So with a **random split**, the test set is not exposed to future information of any patient in the training set, ensuring no leakage. \n    2. **Temporal Split Experimentation**: To validate robustness, we also conducted temporally ordered splits (training on earlier admissions, testing on later ones), and observed consistent performance in https://anonymous.4open.science/r/Rebuttal-B376/temporal.png. Since we focus on one-step-ahead predictions, our design avoids temporal confounding.\n\n**Essential References:**\n\nWe thank the reviewer for highlighting these. **[1] Causal Transformer** shares our goal of personalized decision-making but focuses on counterfactual outcome estimation, while **[2] Modality-aware Transformer** addresses multi-modal modeling in finance. Our work directly predicts next-step treatments and extends beyond multi-modal fusion by incorporating uncertainty and safety for treatment recommendations. We will cite and briefly discuss both in the final version.\n\n**Questions:**\n\n**Q1**: We use masked self-attention to ensure the model only attends to current and past observations at each time step t, aligning where treatment decisions must be based on information available up to time t. Allowing future access would introduce information leakage. Please also refer to Q2 for R2.\n\n**Q2**: Thank you for raising this point. While **SAFER-U** (the single-stage version) shows decent performance, the full **SAFER** framework consistently outperforms it across all metrics—most notably in **mortality rate reduction.** Since in clinical research, even modest reductions in mortality (e.g., 1%) can have a **substantial impact** on patient outcomes which can save people’s lives. Moreover, using a single unified loss forces the network to simultaneously learn from both high- and low-confidence labels early in training, which can lead to erratic convergence. In contrast, our two-step approach first focuses on reliable labels and then selectively adapts to uncertain cases in a risk-aware manner, resulting in more stable training and **clinically meaningful improvements** in reducing mortality.\n\n**Q3**: Thank you for your comment. There might have been some confusion regarding the trends in **Figure 10**. In plot (b) of Figure 10, we report **mortality rate reduction** with a downward arrow indicating this. The trends for **Macro-AUC** and **mortality rate reduction** are actually aligned—both reach optimal values around **gamma = 0.2–0.3**. When **gamma exceeds 0.6**, both metrics start to show a decline in performance, and our method no longer achieves optimal results according to the trend.\n\nWe sincerely hope these clarifications improve your understanding and evaluation of our work. Please feel free to reach out with any additional questions.", "claims_and_evidence": "1. The claims on the multi-modal reasoning and uncertainty aware training are well supported through experiments.\n2. There are two theoretical results - i) the bound on expected uncertainty between deceased and survival patients, and ii) bound on FDR. The claim that *there are theoretical guarantees on calibrated predictions* lacks evidence since the bound in (i) only supports the validity of the approach used and does not theoretically say if the predictions are calibrated.", "methods_and_evaluation_criteria": "Both the datasets used - MIMIC-III and MIMIC-IV - make sense for this problem. The evaluation criteria makes sense too.", "theoretical_claims": "Checked the correctness of both theorems - Theorem 4.1 (the bound on expected uncertainty between deceased and survival patients) and Theorem 5.1 (bound on FDR).", "experimental_designs_or_analyses": "Yes, checked the validity of all experimental designs and analyses.\n\n1. One of the key metric that is used is the decrease in mortality rate, but the definition of the decrease is not clearly specified, is it the decrease from no treatment?\n2. Some more details on the type and number of treatment classes are missing.\n3. It is mentioned that the datasets were randomly split into train/test/validation sets, but considering these datasets are temporal, a temporally aware split is required otherwise it will lead to target leak. The details of the split creation are missing.", "relation_to_broader_scientific_literature": "There are three key ideas in the paper - 1. combining longitudinal health data, demographic attributes along with unstructured clinical notes for treatment predictions, 2. use uncertainty aware fine-tuning to reduce the impact of certain labels such as those of a deceased patient, 3. putting a bound on FDR.\n\nThe technical solutions for each of these ideas including the proof outlines exist in the literature for various applications, and authors have rightly attributed to those existing methods and solutions. But combining them together for treatment prediction is important for the healthcare domain.", "essential_references_not_discussed": "Some references that have solved similar problems were not discussed - \n1. Causal Transformer for Estimating Counterfactual Outcomes - https://arxiv.org/pdf/2204.07258\n2. Modality-aware Transformer for Financial Time series Forecasting - https://arxiv.org/pdf/2310.01232\n3", "other_strengths_and_weaknesses": "Strengths -\n1. The overall framework is useful for healthcare applications and includes desired properties for treatment prediction systems.\n2. Figure 1 is useful and provides a good overview of the system.\n\nWeaknesses -\n1. Some of the experiments descriptions need more clarity, as mentioned above to validate the soundness of the method.", "other_comments_or_suggestions": "NA", "questions_for_authors": "1. Why was masking used in self-attention while learning temporal relationships, since we are mostly interested in progressions in the healthcare data?\n2. Instead of using a two step procedure with training and fine-tuning, what happens if you modify the training procedure loss? From results in Table 2 it seems like SAFER-U is not way worse than SAFER.\n3. In Figure 10 in Appendix, Macro-AUC and mortality rate show different trends wrt to gamma, how do you decide how to pick gamma in that case? and what happens when gamma = 1 that is beyond 0.6 in the same figure?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "A35jOtOtPt", "reviewer": "Reviewer_M8CL", "summary": "This paper introduces SAFER, a new framework for dynamic treatment regime (DTR) recommendation, which integrates both structured electronic health record (EHR) data and unstructured clinical notes. The primary contribution of SAFER lies in its ability to handle uncertainty in treatment recommendations, particularly for deceased patients, by incorporating a calibrated risk-aware training process. The framework employs a Transformer-based architecture to learn multimodal representations from both EHR data and clinical notes, addressing challenges such as temporal dependencies and inter-modality context. Additionally, the paper proposes an uncertainty quantification module to estimate and manage label uncertainty, particularly for uncertain outcomes such as death, and integrates this with a risk-aware loss function for training.\n\nSAFER also introduces conformal prediction to provide statistical guarantees for the reliability of treatment recommendations, ensuring safe and trustworthy decision-making. The paper demonstrates through experiments on publicly available sepsis datasets (MIMIC-III and MIMIC-IV) that SAFER outperforms existing DTR methods in multiple recommendation metrics and significantly reduces counterfactual mortality rates. The findings underscore SAFER’s potential as a robust, theoretically grounded, and practical solution for high-stakes medical decision-making.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for recognizing SAFER as a robust, theoretically grounded, and practical solution for high-stakes medical decision-making, supported by clear evidence. We address your comments as follows.\n\n**Theoretical Claims:** Please see responses to the questions.\n\n**Supplementary Material:**\n\n1. **A.1**: Please see Q3.\n2. **Appendix B**: We will revise this section to clearly describe the input/output features, model architecture (LSTM), and the training/inference setup used to estimate counterfactual mortality. During inference, the model evaluates mortality under SAFER’s recommended treatment. For details on the metric definition, please see our response to R3 on Decrease in mortality rate.\n3. **Sensitivity**: We appreciate the feedback and will expand our discussion to describe hyperparameter selection (via grid search) and how variations in key settings affect model performance across datasets. This will improve clarity and support reproducibility.\n\n**Essential References:**\n\nThank you for the valuable suggestions. We will incorporate the cited works into the revision: Liu et al. (2024) and Chua et al. (2023) for uncertainty (noting SAFER’s focus on *label uncertainty* in DTRs); foundational DTR and causal inference works in Appendix B; additional multimodal learning papers from a boarder area in Section 2; and clarify our novel application of conformal prediction for *treatment recommendation reliability* and discuss prior conformal work in broader healthcare fields.\n\n**Questions:**\n\n**Q1**: Our module $f_\\phi$ rely on the Multi-Layer Perceptron(MLP), which exhibits Lipschitz-like behavior when using standard activations (e.g., ReLU, tanh) and applying regularization. We use L2 regularization and gradient clipping to prevent extreme weight updates and promote stable learning, effectively supporting the Lipschitz continuity assumption in practice. Our design aligns with prior work showing that such regularization encourages stable, Lipschitz-like behavior in neural networks [Gouk et al., 2021; Jin & Candès, 2023]. We will clarify this in the final version.\n\n**Q2**: Thank you for the thoughtful question. We assume the full dataset is i.i.d., and both calibration and test sets are created via random patient-level splits (no patient overlap), satisfying the condition required for conformal inference. We will clarify this in the revised manuscript and note that Theorem 5.1 can be extended under relaxed assumptions for exchangeability, following Thm 6 in [1]. We acknowledge that real-world data may have distribution shifts. Recent work [2] on **weighted conformal inference** offers promising solutions under covariate shift, which we will discuss as a future extension.\n\n[1] Jin & Candès (2023a), *JMLR*; [2] Jin & Candès (2023b), arXiv:2307.09291\n\n**Q3**: We thank the reviewer and agree that a clearer derivation would strengthen the theoretical presentation. A.1 references Sriperumbudur et al. (2009) for the connection between integral probability metrics(IPMs) and φ-divergences class, including KL-divergence. Our proof relies on the **dual representation of φ-divergences** and their links to total variation and IPMs. We will revise A.1 to include key derivation steps and explicitly cite the relevant theorems.\n\n**Q4**: Thank you for pointing this out. We used the standard pretrained BioClinicalBERT with default configuration and **no additional fine-tuning**. We will clarify this in the revised manuscript and cite the specific version to support reproducibility. \n\n**Q5**: In the revised manuscript, we will clearly define the symbols in Equation 3. These correspond to the $S_i^w$, encoded sequences of text and $S_i^c$, structured EHR embeddings, respectively, as introduced in Eq. 2. Both serve as inputs to the cross-attention mechanism. We will update the surrounding text to clarify their roles and computation.\n\n**Q6**: Our work focuses on label uncertainty in deceased patients—a key challenge in clinical settings where outcomes often fail to reflect treatment quality. We agree that real-world uncertainty can be more complex (e.g., missing data, comorbidities), which is beyond our current scope. To handle broader uncertainty types, we could integrate domain knowledge and Bayesian methods, as suggested in recent studies (Sam et al. 2024, Zheng et al. 2021) on uncertainty modeling in healthcare.\n\n**Ethical Review Concerns:**\n\nThank you for the thoughtful points. We use de-identified, HIPAA-compliant data from a single-source (MIMIC-III/IV), unlike federated settings; SAFER addresses data representativeness and uncertainty via selective filtering and will include a case study to support transparency; Deceased patient data is handled through label uncertainty modeling to reduce risk; We emphasize SAFER is a support tool which require human oversight. An ethics discussion will be added in the Appendix.\n\nWe sincerely hope these clarifications improve your understanding and evaluation of our work.", "claims_and_evidence": "The claims made in the submission are largely supported by clear and convincing evidence, particularly through the empirical validation conducted on the MIMIC-III and MIMIC-IV datasets. The paper demonstrates that SAFER outperforms existing dynamic treatment regime (DTR) models in several recommendation metrics and reduces counterfactual mortality rates.", "methods_and_evaluation_criteria": "Yes, the proposed methods and evaluation criteria in the paper make sense for the problem of dynamic treatment regime (DTR) recommendation in healthcare. The approach taken by the authors is well-aligned with the challenges and needs of the problem at hand, and the evaluation criteria are appropriate for assessing the model's performance.", "theoretical_claims": "1. Appendix A.1 cites results from Sriperumbudur et al. (2009), but does not provide specific theorems or derivation steps. The authors should supplement the derivation process for key lemmas.\n2. The specific configuration of BioClinicalBERT (e.g., number of layers, fine-tuning strategy) is not described in detail, which could affect reproducibility.\n3. The definitions of the symbols $\\mathbf{S}^w_i$ and $\\mathbf{S}^v_i$  in equation (3) are not provided with clear explanations. The authors should check the formula and add explanations.", "experimental_designs_or_analyses": "Yes, I reviewed the soundness and validity of the experimental designs and analyses presented in the paper. The authors conducted experiments using publicly available datasets (MIMIC-III and MIMIC-IV), which are commonly used for evaluating medical treatment recommendation models, and employed various performance metrics to assess the effectiveness of the SAFER framework. Below, I discuss the specific aspects of the experimental designs and analyses, as well as any potential issues.", "relation_to_broader_scientific_literature": "The key contributions of the paper are closely related to the broader scientific literature in several areas, including dynamic treatment regimes (DTR), multimodal learning, uncertainty quantification in healthcare, and conformal prediction for error control. The authors build on prior work in these domains and extend existing methods with novel ideas and techniques.", "essential_references_not_discussed": "Yes, there are several related works that could provide additional context and enhance the understanding of the key contributions of the paper but are not currently cited or discussed. These works fall into a few important categories, including uncertainty quantification in healthcare, dynamic treatment regimes (DTR), and multimodal learning in clinical decision-making.\n\n1. Uncertainty Quantification in Healthcare:\nRelated Work Not Cited:\nUncertainty quantification in healthcare, particularly for deep learning models, has gained increasing attention in recent years. One relevant paper is \"Uncertainty-Aware Fine-Tuning of Segmentation Models\" by Liu et al. (2020), which discusses how uncertainty in model predictions can be quantified and mitigated in medical image segmentation tasks. While the paper focuses on segmentation, the idea of incorporating uncertainty into model predictions is highly relevant for understanding SAFER's approach to uncertainty in dynamic treatment regimes.\n\nAnother key work is \"Estimating Uncertainty in Deep Learning for Healthcare\" by Choi et al. (2017), which explores uncertainty in predicting patient outcomes and treatment decisions from electronic health records (EHR). This paper discusses how uncertainty can be modeled within predictive healthcare models, but SAFER's novel approach to quantifying uncertainty in treatment labels, especially for deceased patients, is not sufficiently related to or discussed in the context of this literature.\n\n\nThe authors could provide a more detailed comparison with these works to better highlight the novelty of their uncertainty quantification module, particularly in the context of label uncertainty for deceased patients. While the paper does address this gap, a deeper connection with the existing research on uncertainty in healthcare decision-making would strengthen the justification for SAFER's approach.\n\n2. Dynamic Treatment Regimes (DTR):\nRelated Work Not Cited:\nThe paper builds on DTR frameworks, but it would be beneficial to reference some additional foundational work on counterfactual reasoning and causal inference in DTRs. For instance, \"Counterfactual Risk Assessment for Personalized Treatment\" by Schulam & Saria (2017) introduces methods for modeling the causal effects of treatments when direct observations of outcomes are missing. This work is crucial for understanding how counterfactuals can be used to estimate the efficacy of treatments, which aligns with SAFER’s use of counterfactual mortality rate as a metric.\n\nAnother relevant paper is \"Causal Inference for Personalized Medicine: The Role of Dynamic Treatment Regimes\" by Robins et al. (2000), which discusses how to model and evaluate dynamic treatment regimes in a causal framework. This would help situate SAFER’s approach in the context of causal inference, particularly when dealing with dynamic policies that influence patient outcomes over time.\n\n\nA discussion of how SAFER extends these causal inference and counterfactual frameworks in the context of healthcare decision-making, particularly for dynamic treatment regimes, would clarify its contributions in relation to traditional DTR methods. Specifically, comparing SAFER’s use of label uncertainty and counterfactual mortality rate with traditional methods based on counterfactual causal models could highlight the paper's innovations.\n\n3. Multimodal Learning in Clinical Decision-Making:\nRelated Work Not Cited:\nThe paper introduces a multimodal learning framework for dynamic treatment regimes, integrating clinical notes with structured EHR data. However, multimodal deep learning in healthcare has been explored in recent works such as \"Multi-View Deep Learning for Healthcare\" by Gao et al. (2020), which combines different types of data (e.g., EHR, medical images) for predicting patient outcomes. Additionally, \"Medical Data Integration with Transformer Models\" by Shang et al. (2019) presents a method for integrating multiple sources of healthcare data, including clinical text, using attention mechanisms.\n\n\nWhile the authors mention the use of Transformer-based architectures to integrate multimodal data, they do not directly reference these works, which have already shown success in similar domains. Citing these works would provide a stronger context for the paper’s approach and further validate the effectiveness of Transformer-based models in clinical decision-making. Additionally, discussing how SAFER’s multimodal architecture differs from or improves upon existing multimodal healthcare models would clarify the contribution of the proposed method.\n\n4. Conformal Prediction for Error Control:\nRelated Work Not Cited:\nThe paper introduces conformal prediction for error control in dynamic treatment regimes. However, \"Conformal Prediction for Risk Control in Healthcare\" by Lei et al. (2018) explores the use of conformal prediction for error control in healthcare applications. This method ensures that predictions in high-risk scenarios, such as healthcare, are reliable and within an acceptable error margin. A more detailed discussion of how SAFER applies conformal prediction in the context of dynamic treatment regimes and how it compares with previous work on error control in healthcare would strengthen the foundation for this approach.\n\n\nThe authors should reference existing works on conformal prediction in healthcare to highlight the novelty of applying this approach to dynamic treatment recommendations. By elaborating on how their specific use of conformal prediction addresses the unique challenges of treatment recommendations, especially in high-stakes scenarios, the authors can better position SAFER within the existing body of research on error control.", "other_strengths_and_weaknesses": "The paper is highly original in its combination of ideas. It creatively integrates multimodal data (structured EHR data and unstructured clinical notes) to improve dynamic treatment regimes (DTR), a critical area in precision medicine. The novelty lies in using Transformer-based architectures to model and fuse these different data types, which has not been fully explored in previous DTR work. Additionally, the introduction of an uncertainty quantification module specifically designed to handle label uncertainty for deceased patients is an innovative contribution to the field, addressing a gap that has not been widely tackled.", "other_comments_or_suggestions": "No", "questions_for_authors": "This paper proposes a new framework, SAFER, for dynamic treatment regime (DTR) recommendation. It innovatively combines structured electronic health record (EHR) data and unstructured clinical note data, and introduces a calibrated risk-aware training process. In particular, it offers a unique solution for handling label uncertainty in deceased patients. This has great potential for medical decision-making in real-world applications. However, the paper still has the following issues:\n\n1. Does the Lipschitz continuity assumption hold in the actual model? More specific analysis of the module 𝑓𝜙 (e.g., network structure, training method) is needed to support this assumption.\n2. Do the experimental data satisfy the \"calibration set and test set i.i.d.\" condition? Real-world medical data may exhibit distribution shifts, and the robustness of the model to this should be discussed.\n3. Appendix A.1 cites results from Sriperumbudur et al. (2009), but does not provide specific theorems or derivation steps. The authors should supplement the derivation process for key lemmas.\n4. The specific configuration of BioClinicalBERT (e.g., number of layers, fine-tuning strategy) is not described in detail, which could affect reproducibility.\n5. The definitions of the symbols $\\mathbf{S}^w_i$ and $\\mathbf{S}^v_i$  in equation (3) are not provided with clear explanations. The authors should check the formula and add explanations.\n6. The uncertainty quantification part of the paper is based on several assumptions, particularly the assumption of label uncertainty in deceased patients. While this method performs well in the experiments, handling label uncertainty in real-world applications may be more complex. The authors should consider how to handle other types of uncertainty or explore the applicability of these assumptions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes improved methods for estimating dynamic treatment regimes in healthcare, focusing on incorporating textual information, handling trajectories corresponding to deceased patients, and using conformal prediction to provide safety guarantees for recommendations. Reviewers appreciated the synthesis of empirical and theoretical results supporting the method and there was consensus that it represents an improvement in performance for learning dynamic treatment regimes. While many of the techniques are not foundational advances, they are synthesized in a way that is potentially impactful for healthcare.
Decision: Accept (poster)