pacl-rri-af

PACL-RRI-AF

Patient-Aware Contrastive Learning Preserves Per-Patient Structure in RR-Interval Representations

Official implementation and pretrained weights for Patient-Aware Contrastive Learning Preserves Per-Patient Structure in RR-Interval Representations — Accepted at GlobalSouthML: Bridging Gaps for Underrepresented Researchers in Machine Learning, ICML 2026 Workshop · July 9, 2026 · Seoul, South Korea.

📄 Paper: arXiv:2606.23570 · 💻 Code: github.com/EML-Labs/pacl-rri-af

---

Model Summary

This repository hosts a multi-scale 1D-CNN encoder trained with the Patient-Aware Contrastive Learning (PACL) objective to produce per-patient–consistent representations of RR-interval (RRI) windows derived from single-lead ECG. The frozen encoder + a lightweight logistic-regression probe distinguishes sinus rhythm (SR) from atrial fibrillation (AF) on a patient-independent test split.

• Input: 1D RR-interval window of length window_size = 200 (normalized), shape [B, 200].
• Output: L2-normalized embeddings [B, 128] (projection head) and L2-normalized latents [B, 128] (encoder). The latents are the recommended representation for downstream probes.
• Architecture: three parallel Conv1d branches (kernel sizes 3/5/7) → 1×1 channel mix → temporal attention pooling → linear projection to a 128-d latent, plus a 2-layer MLP projection head used only during contrastive training.
• Checkpoint: checkpoints/best_model.pth (PyTorch state_dict).
• Model type: self-supervised / supervised contrastive encoder for physiological time series.

---

Key Results

PACL achieves the highest per-patient SR cohesion and 2.6× lower seed variance than SupCon on a patient-independent test split (frozen encoder + logistic regression probe).

Per-Class Downstream Metrics (Proposed, 5 seeds)

---

The Core Idea

Standard supervised contrastive learning treats all same-class segments as positives, regardless of which patient they come from. When class differences partially overlap with between-patient differences, this collapses each subject's distinct sinus rhythm (SR) baseline into a shared class prototype, removing the very structure needed to generalize to unseen patients.

PACL restricts positives to same-patient, same-class pairs:

P(i) = { j ≠ i | patient_j == patient_i  AND  class_j == class_i }
N(i) = { j ≠ i | class_j ≠ class_i }

Same-class segments from different patients are excluded from P(i). This preserves each patient's individual SR structure while still pushing SR and AF apart.

---

How to Use

The model is a plain PyTorch nn.Module (not a transformers model). Clone the repository for the source, then load the checkpoint.

Installation

git clone https://github.com/EML-Labs/pacl-rri-af.git
cd pacl-rri-af
pip install -r requirements.txt

Requirements: Python 3.9+, PyTorch 2.x, numpy, scikit-learn, hydra. Exact pinned versions are in requirements.txt.

Load the encoder and extract embeddings

import torch
from src.Model import Model

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Architecture must match the training config (configs/model/model.yaml)
model = Model(latent_dim=128, projection_dim=128, dropout=0.12).to(device)

state_dict = torch.load("checkpoints/best_model.pth", map_location=device, weights_only=True)
model.load_state_dict(state_dict)
model.eval()

# RR-interval window: [B, window_size=200], normalized
rr_window = torch.randn(4, 200, device=device)

with torch.no_grad():
    embeddings, latents = model(rr_window)
    # embeddings: [B, 128] (L2-normalized projection head output)
    # latents:    [B, 128] (L2-normalized encoder output -> use this for probes)

Reproduce the downstream evaluation

The full evaluation (frozen encoder → logistic-regression probe → SR/AF metrics) is driven by Hydra:

python main.py

Environment: Single NVIDIA RTX 2080, ~3 minutes per run (5 seeds × 3 losses = 15 runs total). Seeds used: 1, 42, 102, 2026, 9999.

---

Training & Evaluation Data

This work uses the IRIDIA-AF dataset, comprising long-term single-lead ECG recordings from 167 patients with paroxysmal AF. Episodes are retained when AF duration ≥ 1 hr and the immediately preceding SR duration ≥ 4 hr. Splits are patient-level (119/24/24 train/val/test); after quality filtering, 154/27/24 episodes remain.

RR intervals are extracted per recording and segmented into windows of window_size = 200 with stride = 50. Two preprocessed HDF5 splits (Data/*_train.h5, Data/*_test.h5) ship with the repository for the test pipeline.

---

Comparison with Prior Work

Direct numerical comparison is limited by dataset and evaluation differences. † Cross-validation with within-patient data mixing inflates reported metrics. ‡ MIT-BIH AF contains no episodes satisfying our quality criteria (AF ≥ 1h, preceding SR ≥ 4h).

The proposed representations exceed Andersen et al.'s patient-holdout sensitivity by over 11 percentage points using only a frozen encoder and a logistic regression probe.

---

Intended Use & Limitations

Intended use. Research on patient-robust representation learning for subject-structured physiological signals, and SR-vs-AF discrimination from RR intervals as a downstream probe.

Out of scope. This model is a research artifact, not a medical device. It must not be used for clinical diagnosis, screening, or treatment decisions.

Limitations.

• Trained and evaluated on IRIDIA-AF (paroxysmal AF cohort); behavior on other populations, leads, or acquisition devices is unverified.
• Expects RR-interval windows preprocessed exactly as in this repository (length 200, matching normalization). Distribution shift in beat detection or normalization will degrade results.
• Only two rhythm classes (SR, AF) are modeled; other arrhythmias are not represented.

---

Generality

The PACL objective does not depend on the encoder or the signal type. It requires only subject identifiers in each mini-batch. The same positive-pair construction is expected to apply to other subject-structured physiological signals — EEG, EMG, PPG — wherever class differences partially overlap with between-subject differences.

---

Citation

If you find this work useful, please cite:

@inproceedings{niroshan2026pacl,
  title     = {Patient-Aware Contrastive Learning Preserves Per-Patient Structure
               in RR-Interval Representations},
  author    = {Niroshan, Yasantha and Wimalasiri, Wimalasiri and Hettiarachchi, Chathuranga},
  booktitle = {GlobalSouthML: Bridging Gaps for Underrepresented Researchers in Machine Learning,
               ICML 2026 Workshop},
  year      = {2026},
  eprint    = {2606.23570},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG}
}

License

MIT License. See LICENSE for details.