Technical Summary: ECG-WM: A Physiology-Informed ECG World Model for Clinical Intervention Simulation

Problem Statement

Current electrocardiogram (ECG) foundation models excel at static diagnostic tasks, such as arrhythmia classification and disease prediction, but fundamentally lack the capacity to simulate counterfactual future states under external interventions. In clinical practice, treatment is a Partially Observable Markov Decision Process (POMDP) requiring "what-if" reasoning to anticipate how a patient's cardiac state evolves under specific pharmacological actions. Existing approaches suffer from a "Clinical Imagination Gap":

1. Static Prediction: Most models estimate risk from static observations without explicitly modeling the influence of treatment actions.
2. Open-Loop Hallucinations: Generative models often focus on signal reconstruction without physiological constraints, leading to clinically implausible waveforms (hallucinations) when simulating drug responses.
3. Lack of Uncertainty Quantification: Deterministic point estimates fail to capture the inherent stochasticity of patient responses, which is critical for safety-critical decision-making.

The core challenge is to develop a framework that integrates candidate pharmacological actions, simulates action-conditioned cardiac dynamics, and evaluates downstream clinical risk signals while adhering to biophysical laws.

Methodology

The authors propose ECG-WM, an end-to-end ECG World Model designed for action-conditioned predictive simulation. The framework moves beyond disjoint pipelines by unifying treatment generation, physiological simulation, and risk evaluation.

1. Architecture Overview

The system operates on a pre-dose ECG signal $x_k$ and patient metadata $d$. It proceeds through three main stages:

• Action Proposal: A Vision-Language Model (VLM) or a trainable multimodal model proposes candidate pharmacological interventions (drugs and dosages).
• Physiology-Informed World Model: The core component simulates the post-dose ECG trajectory $x_{k+1}$.
• Uncertainty-Aware Risk Evaluation: A downstream ECG Foundation Model evaluates the simulated trajectories to estimate clinical risk and its variability.

2. Physiology-Informed Latent Diffusion

To model cardiac dynamics, the framework encodes raw 12-lead ECGs into a low-dimensional latent space using a pre-trained Variational Autoencoder (VAE). The evolution of the latent state $z$ under an action $a$ is modeled via a Denoising Diffusion Probabilistic Model (DDPM) enhanced with Energy Regularization.

• External Physiological Knowledge (EPK): Instead of learning dynamics purely from data, the model incorporates a mechanistic prior derived from the McSharry dynamical system, a continuous-time Ordinary Differential Equation (ODE) that describes ECG waveforms (P, Q, R, S, T waves) as a function of cardiac phase.
• Energy-Regularized Training: The diffusion process is guided by an energy function $E_\phi$ that penalizes deviations from the physiological prior. The training objective combines the standard diffusion reconstruction loss with a physiological regularization term:
  $$\mathcal{L} = \mathbb{E} \left[ \| \epsilon_\theta(z_\tau, \tau) - \epsilon \|^2_2 + \omega_\tau E_\phi(\hat{z}_0) \right]$$
  Here, $\hat{z}_0$ is the predicted clean latent state, and $E_\phi(\hat{z}_0) = \| \hat{z}_0 - \Pi_\phi(e_{EPK}) \|^2_2$, where $e_{EPK}$ is the discretized trajectory from the ODE solver and $\Pi_\phi$ projects it into the latent space. This ensures generated trajectories respect biophysical constraints (e.g., wave morphology and timing) without requiring expensive decoding during training.

3. Uncertainty-Aware Risk Calibration

Recognizing that clinical outcomes are stochastic, the framework does not rely on a single deterministic prediction. Instead, it leverages the stochasticity of the diffusion process:

• Multi-Sample Generation: For a given intervention $a$, the model samples $K$ independent future trajectories in the latent space.
• Risk Distribution: Each trajectory is decoded and evaluated by a frozen ECG Foundation Model to predict a distribution of clinical risks (e.g., probability of arrhythmia).
• Scoring Function: Interventions are ranked using a mean-variance objective:
  $$S(a) = \mu_R(a) + \lambda \cdot \sigma_R(a)$$
  where $\mu_R$ is the expected risk and $\sigma_R$ is the predictive uncertainty. This allows clinicians to prefer interventions with lower expected risk and lower variability.

Key Contributions

1. Physiology-Informed World Model: The first integration of non-linear cardiac ODE priors into a latent diffusion process via energy regularization. This ensures generated ECG trajectories adhere to biophysical constraints, capturing both underlying cardiac dynamics and complex waveform representations.
2. Closed-Loop Counterfactual Simulation: An end-to-end differentiable framework that unifies candidate treatment generation, post-intervention ECG simulation, and downstream risk evaluation. It enables intervention-aware analysis by explicitly modeling the physiological effects of different medications.
3. Uncertainty-Aware Risk Calibration: A novel evaluation strategy that characterizes both the expected clinical risk and its variability across candidate interventions, providing a more reliable basis for comparative assessment than deterministic point estimates.
4. Comprehensive Evaluation: The method demonstrates robust performance across diverse datasets, including controlled drug-response scenarios and real-world clinical records, outperforming state-of-the-art baselines in waveform fidelity and risk prediction.

Experimental Results

The authors evaluated ECG-WM on three datasets: MIMIC-IV-ECG & Clinical (real-world EHRs), Drug-Response ECG Datasets (ECGRDVQ and ECGDMMLD), and an In-house OOD Dataset (Verapamil-treated ICU patients).

• Simulation Fidelity: On the MIMIC-IV-ECG benchmark, ECG-WM achieved significantly lower Mean Squared Error (MSE) and Mean Absolute Error (MAE) in both latent and signal spaces compared to zero-shot large language models (e.g., GPT-5 mini, GPT-4o) and other diffusion baselines.
• Physiological Consistency: In terms of preserving key cardiac biomarkers (QTc, PR, and $T_{peak}-T_{end}$ intervals) under drug-induced variations, ECG-WM achieved 90.55% accuracy and 96.26% recall for QTc intervals, outperforming all GAN and diffusion baselines (e.g., DADM, Diffusets).
• Out-of-Distribution (OOD) Generalization: When tested on unseen pathological patterns (ICU patients with Verapamil), ECG-WM maintained superior performance (Latent MSE: 0.068) compared to diffusion baselines (0.080–0.096), demonstrating robustness to distribution shifts.
• Risk Modeling: The framework showed strong alignment with expert preferences, achieving a 76% Sign Agreement on the direction of drug-induced risk changes ($\Delta$Risk), compared to 58% for the DADM baseline. Pearson and Spearman correlations for risk prediction were also significantly higher (0.620 vs 0.266).
• Robustness: The model demonstrated graceful degradation under missing lead scenarios and increasing prior mismatch, maintaining stability where unconstrained models failed or produced hallucinations.

Significance and Claims

The paper positions ECG-WM as a robust foundation for safe and intervention-aware clinical decision support. By bridging the gap between generative simulation and downstream clinical analysis, the framework extends ECG analysis from a purely diagnostic setting to a simulation-driven paradigm.

The authors claim that their approach addresses the "Clinical Imagination Gap" by enabling AI systems to simulate counterfactual futures. Unlike previous open-loop models that risk generating clinically implausible hallucinations, ECG-WM uses physiological priors to ensure structural robustness. The integration of uncertainty quantification allows for a more nuanced evaluation of treatment options, prioritizing safety in high-stakes medical domains.

Limitations Acknowledged by the Authors:

• The model relies on discrete intervention tokens due to a lack of high-resolution pharmacokinetic time-series in public datasets.
• The ODE priors are simplified and may struggle with highly irregular patterns (e.g., severe ventricular fibrillation).
• The current evaluation focuses on aggregate stability and risk metrics rather than hard clinical endpoints.
• Comprehensive clinical validation is required before real-world deployment; the system is intended as a decision-support tool requiring expert validation, not a standalone clinical tool.