Attention-Based Offline Reinforcement Learning and Clustering for Interpretable Sepsis Treatment

AI summary

Overview

• Research area: Machine learning for healthcare, specifically offline reinforcement learning, unsupervised clustering, generative synthetic-data augmentation, and large language model based explainability applied to sepsis treatment in intensive care units.
• Technical level: Intermediate. The paper assumes familiarity with reinforcement learning concepts (Q-functions, value functions, advantage weighting, offline RL) and density-based clustering, but its four components are presented modularly and are individually describable without deep math.
• Scope: The paper describes and evaluates an end-to-end decision-support pipeline that stratifies ICU patients by risk, augments sparse treatment data with synthetic trajectories, learns a treatment policy with attention-based Advantage-Weighted Regression ensembled with XGBoost and TabNet, and generates natural-language justifications with a locally deployed multi-modal LLM, evaluated on MIMIC-III and eICU.

What This Paper Is About

Sepsis is a life-threatening response to infection, and in the ICU the timing and choice of interventions such as fluids and vasopressors strongly affect survival. Learning treatment policies from historical ICU records is difficult because reinforcement learning cannot explore in a live hospital, rare but critical interventions are underrepresented in the data, and clinicians cannot act on recommendations they cannot understand. The paper's goal is to build an offline RL treatment recommender that is simultaneously accurate on rare treatment classes, generalizable to newly admitted patients with little or no ICU history, and able to explain each recommendation in plain clinical language.

Key Contributions

1. An interpretable offline RL pipeline that combines Advantage-Weighted Regression (AWR) with a simplified, lightweight attention mechanism, whose output is blended with ensemble predictions from XGBoost and TabNet to improve learning stability and treatment accuracy.
2. A clustering-driven stratification module using HDBSCAN (with UMAP for dimensionality reduction) that groups patients into low-, intermediate-, and high-risk categories and handles the cold-start problem for patients with limited or no ICU history by assigning them to similar historical trajectories.
3. A synthetic data augmentation pipeline that generates trajectories with a diffusion model and a conditional variational autoencoder (VAE) to counter class imbalance and data sparsity in critical interventions such as vasopressors and fluids.
4. A rationale generation module that integrates a multi-modal large language model into the inference pipeline, producing patient-specific natural-language justifications grounded in the current vitals, retrieved clinical knowledge, and the RL agent's selected action.

Main Findings

• Ensemble accuracy: The ensemble model (XGBoost + AWR with attention + TabNet) achieves the best overall accuracy at 83%, compared with 80% for AWR + Attention, 74% for BCQ + Attention, and 60% for BCQ. Average reward is not reported for the ensemble (listed as N/A).
• AWR beats the BCQ baselines: AWR + Attention reaches 80% accuracy and an average reward of -0.33, versus 74% and -0.47 for BCQ + Attention, and 60% and -0.60 for plain BCQ.
• Minority treatment classes improve: On per-class precision/recall, the ensemble reaches precision 0.50 and recall 0.60 on action A1 (fluids) and 0.65/0.60 on action A2 (vasopressors), while BCQ scores only 0.01/0.03 and 0.30/0.25 respectively. The ensemble also reaches 0.93/0.92 on A0 (no treatment) and 0.81/0.70 on A3 (combined treatments).
• Clustering produced clinically distinct risk groups: After aligning MIMIC-III and eICU on icustay_id, the merged dataset contained 874,108 time-stamped records across 27,799 ICU stays. Restricting to patients with between 1 and 80 time points preserved approximately 75% of the original data and retained all 27,799 unique patients, yielding a final feature matrix of 25,605 × 320. HDBSCAN with minimum cluster size 30, minimum samples 30, and epsilon 0.01 produced 124 distinct clusters, grouped into three risk categories. High-risk clusters showed 100.0% mortality, intermediate-risk clusters 32.6–62.4%, and low-risk clusters 0.0–4.2%.
• Noise handling: Noisy (unclustered) data comprised 5.3% of training data and 11.3% of test data, which the authors describe as acceptable given clinical data variability.
• Feature importance validates preprocessing: In the XGBoost gain-based ranking, SpO2 scored 910.93, platelets 732.08, MAP 239.74, hours since ICU admission 179.62, GCS total 138.09, ethnicity 116.90, systolic BP 104.86, and bicarbonate 70.51, while a deliberately injected random noise feature scored lowest at 0.93.
• Attention is clinically plausible: Figure 6 visualizes a patient trajectory where the model places the greatest attention weight on MAP (with secondary, thinner weights on other vitals and labs) during hypotensive episodes.
• LLM rationales are clinically coherent: The authors selected the multi-modal LLaMA3.2-Vision model for offline, healthcare-compatible deployment, using top-K sampling of 100, repeat penalty of 1.1, and temperature of 4.7. Sample outputs include "Vasopressor therapy was initiated due to persistent hypotension (MAP < 65 mmHg) and elevated lactate, suggesting ongoing hypoperfusion" and "No immediate action was taken as vital signs are stable and lactate levels are normal." No quantitative evaluation metrics for rationale quality are reported.

Methodology in Plain English

The pipeline has four stages that feed into one another.

Risk stratification first. Because many newly admitted patients have little or no ICU history, the authors cannot immediately run longitudinal models on them. Instead, they preprocess ICU records (temporal filtering, L2 normalization, UMAP dimensionality reduction) and apply HDBSCAN to group patients by similarity. Statistical validation using mortality variance and a chi-square test (with significance set at p < 0.001) confirms the clusters are meaningful. Clusters are then mapped to risk labels using mortality thresholds: 0–40% mortality is low risk, above 40% up to 75% is intermediate risk, and above 75% is high risk. Low-risk patients need no intervention, intermediate-risk patients are the target population for RL recommendations alongside clinical judgment, and high-risk patients require immediate intervention.

Then synthetic data generation. To correct the underrepresentation of fluid and vasopressor interventions, the authors generate two kinds of synthetic data. A lightweight diffusion model adds Gaussian noise to normalized observation windows through a noise schedule, trains a network to predict the added noise, and then denoises from pure noise back to a realistic sample that is unscaled, clipped, and added to the training set. A conditional VAE handles discrete, action-conditioned transitions: an encoder maps the state-action pair to a latent Gaussian distribution, a latent sample is drawn via the reparameterization trick, and a decoder reconstructs the next state conditioned on the latent variable plus action, reward, and terminal indicator. Training minimizes reconstruction error plus a KL divergence term, and post-processing clips outliers and filters implausible transitions.

Then the treatment agent. Each patient state is a 30-dimensional vector of vital signs, labs, and treatment indicators, standardized with z-scores and imputed by carrying the last observation forward (falling back to the feature median). Actions are discrete: no treatment, fluids, vasopressors, or combined treatments. The reward is composite, ranging from -1 for mortality within 48 hours to +0.8 when MAP is above 65, SpO2 is above 94, and lactate is below 2, each weighted (0.3, 0.3, and 0.2 respectively). AWR with feature attention encodes states and next states into latent embeddings, computes a one-step value target, trains the value head with expectile regression (weighting over- and under-estimates differently), trains the Q-head to match the same target, converts the advantage into an exponential weight, and trains the policy with that weight. Targets are updated softly with a rate alpha. The final recommendation comes from an ensemble: TabNet and XGBoost are consulted first, and if either predicts a treatment probability above a threshold omega, that action overrides the RL agent; otherwise the RL policy decides. This is intentionally conservative for sensitive interventions.

Finally, explanation. The patient state is mapped into the learned clustering space and merged with the selected action, then used as a query against a vector database via approximate nearest neighbor search over an expert-curated sepsis knowledge base that was pre-vectorized with a NOMIC encoder. The top-k retrieved tokens are inserted into a prompt template and passed to the LLM, which produces a natural-language rationale for the recommended action.

Why This Matters

• Research impact: The paper demonstrates that interpretability can be treated as a first-class design goal rather than an afterthought in clinical RL, by combining attention-based state encoding, cluster-based patient stratification, and LLM-generated rationales in one pipeline. It also shows a practical way to attack class imbalance in offline RL via generative augmentation with diffusion models and VAEs.
• Real-world applications:
  • ICU decision support at the bedside, where a clinician receives both a recommended action and a written justification for intermediate-risk patients.
  • Admission triage, using the clustering module to sort newly admitted patients into risk categories before any longitudinal history exists.
  • Clinical auditing and quality review, where natural-language rationales make policy behavior inspectable rather than opaque.
  • Education and training, where generated rationales illustrate the clinical reasoning behind fluid or vasopressor decisions.
• Industry relevance: The requirement for local, offline deployment (the reason LLaMA3.2-Vision was selected over hosted alternatives) reflects real hospital constraints around data privacy and security. The reliance on widely used tabular models (XGBoost, TabNet) and standard databases makes the approach compatible with existing clinical data infrastructure rather than requiring new instrumentation.

Future Directions

• Prospective or simulated validation: The reported results are derived offline from MIMIC-III and eICU; the paper does not report a prospective or simulator-based evaluation of whether the policies improve actual patient outcomes, which would be the natural next step.
• Quantitative evaluation of the rationales: The paper presents sample LLM outputs as evidence of clinical soundness but reports no metrics for faithfulness, factuality, or clinician agreement, leaving the quality of generated explanations an open question.
• Threshold and reward sensitivity: The ensemble relies on a confidence threshold omega, the AWR training relies on a temperature beta, an expectile tau, and a soft-update rate alpha, and the reward uses fixed weights of 0.3, 0.3, and 0.2. The paper does not report sensitivity analysis for these choices.
• Cold-start generalization: The paper claims clustering helps new patients with limited or no ICU history, but does not report a separate quantitative evaluation isolating that subgroup, which would clarify how well the stratification transfer actually works.

Target Audience

This paper is most useful to machine learning researchers working on offline reinforcement learning and healthcare applications, clinical informatics teams building ICU decision-support tools, and data scientists interested in combining generative augmentation with interpretable policy learning. Clinicians and hospital IT stakeholders evaluating whether an RL-based recommender could be deployed locally will find the deployment and explainability discussion relevant, though the evaluation details are more technical than clinical.

Authors’ abstract

Sepsis remains one of the leading causes of mortality in intensive care units, where timely and accurate treatment decisions can significantly impact patient outcomes. In this work, we propose an interpretable decision support framework. Our system integrates four core components: (1) a clustering-based stratification module that categorizes patients into low, intermediate, and high-risk groups upon ICU admission, using clustering with statistical validation; (2) a synthetic data augmentation pipeline leveraging variational autoencoders (VAE) and diffusion models to enrich underrepresented trajectories such as fluid or vasopressor administration; (3) an offline reinforcement learning (RL) agent trained using Advantage Weighted Regression (AWR) with a lightweight attention encoder and supported by an ensemble models for conservative, safety-aware treatment recommendations; and (4) a rationale generation module powered by a multi-modal large language model (LLM), which produces natural-language justifications grounded in clinical context and retrieved expert knowledge. Evaluated on the MIMIC-III and eICU datasets, our approach achieves high treatment accuracy while providing clinicians with interpretable and robust policy recommendations.