Enhancing stability in cardiac risk stratification with equivariant neural fields

Standard deep learning approaches in medical imaging are sensitive to geometric transformations such as rotations and translations. These models, typically trained on image grids, lack inherent spatial invariance, leading to performance degradation under variations in input orientation. This sensitivity limits generalization in real-world clinical applications, especially in modalities like cardiac MRI (CMR) and CT.

This study investigates whether Equivariant Neural Fields (ENFs), which learn continuous latent point cloud representations grounded in spatial geometry, can provide improved robustness and generalization for predicting cardiac risk–related endpoints. To assess robustness, ENFs are compared against a ResNet50 baseline under rotated input conditions.

An ENF was trained using a meta-learning approach to reconstruct mid-slice, end-diastolic and end-systolic CMR images from 275 patients in the UK Biobank. The learned latent point cloud representations (illustrated in Figure 1) - obtained after reconstruction training - were used as input to a lightweight Transformer classifier for binary classification of left ventricular ejection fraction (LVEF ≤ 40% indicating high risk).

Although the UK Biobank includes tens of thousands of CMR samples, only a small fraction of patients exhibit reduced LVEF (≤ 40%). To ensure balanced model training, a stratified dataset was constructed with approximately 50/50 class balance between high-risk (LVEF ≤ 40) and low-risk (LVEF > 40) patients across the training (n = 275), validation (n = 60), and held-out test (n = 60) set.

To evaluate geometric robustness, the test set consisted of randomly rotated versions (90°, 180°, or 270°) of the original CMR images. The corresponding rotated latent point clouds were passed to the Transformer classifier. A ResNet50 model was trained directly on the original image slices using the same splits, training setup and class balance as the Transformer setup for fair comparison (Figure 2).

The ResNet model achieved an accuracy of 91.0% on the validation set but dropped to 83.0% on the rotated test set, indicating reduced generalization under geometric transformations. The ENF model reached a validation accuracy of 87.0% and maintained robust performance with 87.0% accuracy on the rotated test set. The latent point cloud representations were invariant to the tested transformations, resulting in consistent downstream performance without requiring rotation-specific augmentations during training.

ENFs provide a robust and geometry-aware alternative to conventional CNN-based models for cardiac risk stratification. Their inherent equivariance enables stable classification across input transformations, supporting broader applicability in clinical environments. Additionally, ENFs can offer modality- and resolution-agnostic representations, making them well-suited for multimodal and heterogeneous medical data.Figure 1:

Schematic Illustration of ENFs