Beyond Patient Invariance: Learning Cardiac Dynamics via Action-Conditioned JEPAs

Technical Summary: Beyond Patient Invariance: Learning Cardiac Dynamics via Action-Conditioned JEPA

Paper Title: Beyond Patient Invariance: Learning Cardiac Dynamics via Action-Conditioned JEPA
Venue: ICLR 2026, 2nd Workshop on World Models
Authors: Jose Geraldo Fernandes, Luiz Facury, Pedro Robles Dutenhefner & Wagner Meira Jr. (UFMG, Brazil)

---

1. The Problem: The "Invariance Paradox" in Medical AI

The authors identify a fundamental misalignment between current Self-Supervised Learning (SSL) paradigms and clinical reality.

• The Invariance Flaw: Most medical SSL (e.g., PCLR, MICLe) relies on Patient-Invariance, training models to maximize similarity between different recordings of the same patient. While this helps identify stable anatomy, it mathematically forces the model to treat transient pathological changes (like the onset of an arrhythmia) as "noise" to be suppressed.
• Augmentation Risks: Standard data augmentations (cropping, jittering) used in SSL can be semantically destructive to 1D physiological signals like ECG, potentially removing critical clinical markers (e.g., the P-wave) or mimicking pathologies.
• The Goal: Clinical monitoring requires the detection of change, not the preservation of stability.

2. Methodology: Action-Conditioned Cardiac World Models

The authors propose shifting from static invariance to dynamic causality by treating disease onset as a "translational action" within a latent space.

• Core Framework (LeJEPA): Instead of generative models (which risk "hallucinating" unrealistic waveforms), they adapt the Joint-Embedding Predictive Architecture (JEPA). The model predicts future latent states directly in a geometrically regularized embedding space.
• Action-Conditioned Dynamics: They redefine the clinical label not as a static category, but as a transition vector ($a_t = y_{t+1} - y_t$). This vector represents the "action" that moves a patient from a healthy state to a pathological one.
• Architecture Components:
  • Encoder ($\text{Enc}_\theta$): An xResNet1d50 that maps raw 12-lead ECG signals to a latent representation ($h_t$).
  • Action Projector ($\text{Proj}_\omega$): An MLP that transforms the sparse transition vector into a continuous action embedding.
  • Dynamics Predictor ($\text{Dyn}_\phi$): A residual MLP that predicts the future latent state: $\hat{h}_{t+1} = \text{Dyn}_\phi(h_t, \text{Proj}_\omega(a_t))$.
• Regularization (SIGReg): To prevent "representation collapse" (where all inputs map to the same vector), they use Sketched Isotropic Gaussian Regularization. This ensures the latent space follows a standard Gaussian distribution, making it rotationally symmetric and ideal for linear downstream tasks.

3. Key Contributions

1. New Paradigm: Recontextualizes disease onset as a translational action variable in a world model.
2. Disentanglement: Demonstrates that the model can separate stable Patient Identity (anatomical substrate) from dynamic Disease Physics (pathological forces).
3. Superior Supervision: Proves that predicting latent dynamics provides a "denser" and more robust supervision signal than simple categorical classification.

4. Experimental Results

The model was evaluated on the MIMIC-IV-ECG dataset, comparing it against fully supervised baselines and naive patient-invariance SSL.

• Generalization (Triage vs. Monitoring):
  • On the Triage task (predicting diagnosis from the first ECG), the proposed model slightly outperformed the supervised baseline (0.742 vs 0.735 AUROC).
  • On the Monitoring task (all ECGs), it achieved parity with supervised models.
• Data Efficiency (The "Low-Resource" Win):
  • In a 10% data regime, the supervised model's performance collapsed (0.74 $\to$ 0.64 AUROC), whereas the World Model remained robust (0.74 $\to$ 0.69 AUROC), outperforming the baseline by over 0.05 AUROC.
• Failure Mode Analysis: At the extreme 1% data regime, the model's "Linear Probe" (frozen encoder) still maintained a signal (0.56 AUROC), while "Fine-tuning" collapsed. This suggests the representations are learned correctly, but the adaptation process is too noisy when labels are scarce.

5. Significance and Future Work

The paper moves medical AI away from "identifying patients" toward "simulating disease."

Significance: By modeling the physics of how a heart changes, the model becomes more sample-efficient and better at detecting acute events that traditional invariant models would ignore.

Future Directions:

• Continuous Action Embeddings: Moving from binary $\{ -1, 0, 1 \}$ vectors to continuous embeddings to capture the magnitude and speed of disease progression.
• Neural ODEs: Integrating the model with Neural Ordinary Differential Equations to handle the "arrow of time" and the difference between acute events (sudden jumps) and chronic drift (slow changes).
• Counterfactual Simulation: Using the model to answer "What-if" questions (e.g., "How would this specific patient's ECG change if they developed Atrial Fibrillation?").