rmxjck/ltaf-ecg-beats-classifier-htf

LTAF ECG Beat Classifier — HTF (Time + Frequency + History)

A 3-class beat classifier (Normal / Atrial premature / Ventricular premature) trained on PhysioNet's Long-Term Atrial Fibrillation (LTAF) database.

Inspired by the History-Time-Frequency (HTF) ensemble in alberto-rota/PAC-PVC-Beat-Classifier-for-ECGs: three parallel streams (raw waveform, FFT log-magnitude, RR-interval history with previous beat labels) fused before a small MLP head.

Metric	Value
Test accuracy	0.954
Test balanced accuracy	0.952
Test macro F1	0.936

vs. frozen Chronos-2 + MLP baseline on the same beats: F1 = 0.906 (+3 pp with 4× the head capacity but no LLM forward at inference).

Per-class F1 (test, n = 57 693 beats from 9 records):

• N (Normal): 0.976
• A (Atrial premature, PAC / APC): 0.930
• V (Ventricular premature, PVC / VE): 0.902

Test confusion matrix (rows = true, cols = pred, order N/A/V):

N: [36778,  406, 1278]   (recall 95.6 %)
A: [   98, 8539,  677]   (recall 91.7 %)
V: [   55,  112, 9750]   (recall 98.3 %)

Architecture

EcgBeatHTFClassifier(num_classes=3, n_channels=2, window_samples=256, history_k=5, history_use_labels=True, time_base_channels=32, freq_base_channels=32, head_hidden=128):

• Time stream — 5 Conv1d→BN→ReLU→MaxPool blocks on a (B, 2, 256) raw R-peak-centered window. Adaptive avg pool → 256-dim feature.
• Frequency stream — 4 conv blocks on log|rFFT| of the time signal (B, 2, 129). Adaptive avg pool → 256-dim feature.
• History stream — K=5 preceding RR intervals (sec) + one-hot of preceding K beat labels → 64-dim MLP feature.
• Head — concat (576-dim) → Linear(128) → ReLU → Dropout → Linear(3).
• Total parameters: 1,143,939.

Quickstart

pip install torch huggingface_hub numpy

import numpy as np
import torch
from huggingface_hub import hf_hub_download
from model import EcgBeatHTFClassifier, BEAT_CLASS_NAMES

ckpt = hf_hub_download("rmxjck/ltaf-ecg-beats-classifier-htf", "best_classifier.pt")
model = EcgBeatHTFClassifier.load(ckpt, device="cuda")
model.eval()

# Inputs:
# x_time: (B, 2, 256) — 2 s @ 128 Hz, R-peak centered, z-scored.
# rr_history: (B, 5) — RR intervals (seconds) to preceding K=5 beats.
# label_history: (B, 5) int64 — preceding K=5 beat labels (0=N, 1=A, 2=V); -1 = missing.
x_time = torch.randn(1, 2, 256).cuda()
rr = torch.tensor([[0.85, 0.83, 0.87, 0.82, 0.85]]).cuda()
lbl = torch.tensor([[0, 0, 0, 0, 0]]).cuda()  # all N
with torch.no_grad():
    pred = model(x_time, rr, lbl).argmax(-1).item()
print(BEAT_CLASS_NAMES[pred])  # "N", "A", or "V"

For a full beat-sequence example (autoregressive history), see inference.py.

Input format

• Time signal (B, 2, 256) float32: 2 s @ 128 Hz, 2-lead, centered on the R-peak sample, per-channel z-scored.
• RR history (B, 5) float32: seconds. rr[k] is the gap between the (k+1)th-prior beat and the kth-prior beat. Use 0.0 for unknown.
• Label history (B, 5) int64: preceding 5 beat labels (0=N, 1=A, 2=V), most-recent first. Use -1 for unknown / record start.

At inference time, label history can be filled by:

• The model's own previous predictions (autoregressive — see predict_beat_sequence in inference.py).
• All -1 if you don't have any (still works, slight recall loss).

Training recipe

.venv/bin/python scripts/train_ecg_beat_htf.py \
    --epochs 15 --batch-size 256 --history-k 5 \
    --output-dir results/ecg_classifier/beats_htf

• Dataset: LTAF beat timelines from nicozumarraga/ltaf-haystack. 67 train / 8 val / 9 test records, deterministic seed 42.
• Loss: weighted cross-entropy with sqrt-dampened inverse-frequency class weights (cap 10), label smoothing 0.05.
• N (normal) beats subsampled per epoch to 2 × n_nonN to balance the ~97 % / 1.7 % / 1.5 % class skew.
• Cosine LR 1e-3 → 0 over 15 epochs. AdamW (wd 1e-4).
• Best checkpoint by val macro F1 (epoch 9 in our run).
• Training time on a single H100 80GB: ~22 min.

Source repo: scripts/train_ecg_beat_htf.py and src/models/ts_llm/ecg_beat_htf.py in rmxjck/TSLM-Arena.

What was tried and didn't help

• Raw Chronos-2 frozen encoder + MLP head: F1 = 0.906 (-3 pp). The HTF ensemble's combination of morphology + spectral + R-R-context beats the foundation-model encoder for this single-beat task.
• Various history-K values (3, 7, 10): K=5 was best.
• No-history HTF (no label_history): -1.5 pp F1.
• No-frequency HTF: -1 pp F1.

Not for clinical use

Research artifact only. Not FDA-cleared. Not suitable for triage, diagnosis, or any patient-facing application.

Citation

@misc{petrutiu2008ltafdb,
  title         = {Abrupt Changes in Fibrillatory Wave Characteristics at the Termination of Paroxysmal Atrial Fibrillation in Humans},
  author        = {Petrutiu, Simona and Sahakian, Alan V. and Swiryn, Steven},
  year          = {2008},
  howpublished  = {PhysioNet},
  url           = {https://physionet.org/content/ltafdb/}
}

The HTF architecture is inspired by: alberto-rota/PAC-PVC-Beat-Classifier-for-ECGs.