SAFER: A Calibrated Risk-Aware Multimodal Recommendation Model for Dynamic Treatment Regimes

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.

Theoretical Claims: Please see responses to the questions.

Supplementary Material:

1. A.1: Please see Q3.
2. 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.
3. 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.

Essential References:

Thank 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.

Questions:

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.

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.

[1] Jin & Candès (2023a), JMLR; [2] Jin & Candès (2023b), arXiv:2307.09291

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.

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.

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.

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.

Ethical Review Concerns:

Thank 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.

We sincerely hope these clarifications improve your understanding and evaluation of our work.

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.

Claims And Evidence:

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).

• The bound in (i) validates our KL-based uncertainty score as a meaningful discriminator between reliable and uncertain labels, but does not imply calibration.
• 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$.

We will revise accordingly to avoid ambiguity.

Experimental Designs:

1. 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.
2. 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.
3. Temporally aware split: We understand your concern regarding potential target leak when handling temporal data. To clarify:
   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.
   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.

Essential References:

We 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.

Questions:

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.

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.

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.

We sincerely hope these clarifications improve your understanding and evaluation of our work. Please feel free to reach out with any additional questions.

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.

In 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.

Below, we address your comments regarding notation and clarity. We appreciate your careful review and will revise the manuscript accordingly in the final version:

• 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.
• 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”).
• 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.

Thank 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.

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.

Methods And Evaluation Criteria:

Thank 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.

We 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.

Essential References:

Thank you for the helpful suggestion. In response to multiple reviewers’ feedback, we have expanded the manuscript to cover key related domains:

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.

DTRs & Counterfactual Framework. Foundational work is now referenced in Appendix B and Section 2 to contextualize our use of counterfactual mortality rates.

Multimodal Learning. We expanded discussion of prior work on multimodal modeling and transformer-based fusion across domains beyond clinical notes and EHR.

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.

Other Comments Or Suggestions:

Thank 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.

Questions:

Q1: Thank you for emphasizing the need for interpretability. Here are some interpretability strategies we will discuss in the revised version.

1. 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.
2. 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.
3. 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."

We sincerely hope these clarifications improve your understanding and evaluation of our work.