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

TL;DR

This work tackles sepsis treatment as a sequential decision problem by building an interpretable pipeline that combines risk stratification via clustering, synthetic data augmentation with diffusion models and conditional VAEs, offline reinforcement learning using Advantage Weighted Regression with a lightweight attention encoder, and rationale generation through a multimodal LLM. The ensemble RL approach integrates AWR with TabNet and XGBoost, applying a conservative rule to prioritize safety, while a clustering module addresses cold-start and data sparsity at ICU admission. Evaluations on MIMIC-III and eICU show the method achieves high treatment accuracy (83%), with particularly strong performance on underrepresented interventions, and the LLM-generated rationales provide natural language explanations grounded in clinical context. The framework advances interpretable, robust decision support for sepsis care, enabling better generalization across patients and auditable reasoning for clinicians.

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.

Figures (6)

• Figure 1: Overview of the interpretable sepsis treatment pipeline. (A) Patients without prior ICU history are stratified into low-, intermediate-, or high-risk groups using clustering. (B) To address data sparsity, synthetic transitions $(s, a, r, s', d)$ are generated using a VAE and a diffusion model, then added to the RL training set. (C) For intermediate-risk or historical patients, a feature-attention encoder produces a latent state $z = \mathrm{Attn}(s)$, used by an AWR policy $\pi_{\phi}(a \mid z)$, a Q-network, and a value function. The final recommendation $a' = \arg\max_{a} [\text{blend}(\pi_{\phi}, \pi_{\text{XGB}})]$ combines outputs from AWR and a clinician-trained XGBoost policy. (D) To enhance interpretability, a local LLM generates a natural-language rationale $r = f(s, a', \text{context})$ using the patient state, selected action, and retrieved clinical context.
• Figure 2: Clustering-based risk stratification pipeline. After preprocessing ICU data, including feature selection, temporal sequence construction, and UMAP Dimensionality reduction, HDBSCAN is used to group patients by similarity. The resulting clusters are validated using risk trends to ensure clinical relevance.
• Figure 3: Minority-class augmentation pipeline. Cleaned MIMIC-III + eICU records are filtered to classes 1 & 2, then two parallel generators create new samples: (1) a VAE that decodes latent draws, and (2) a conditional diffusion process that iteratively denoises scaled states with a time-stepped loop.
• Figure 4: AWR with a custom feature attention pipeline. The $30‑$D state goes through a custom attention block and a shared encoder. Policy ($\pi$), Q, and V heads are trained together. We weight the policy loss using the advantage value and use a target-value network to smooth the learning updates.
• Figure 5: Overview of the LLM-based rationale generation process. A query, combining the patient's cluster state and the RL agent's action, is used to retrieve relevant context from a vectorized sepsis knowledge base. The retrieved information is then passed to an LLM to generate a final, knowledge-grounded explanation for the action.