TRACE is a 12-lead ECG autoencoder whose 32-dimensional clinical latent space is partitioned by a literature-derived diagnostic ontology (perfusion / structure / conduction), with each subspace routed by design to its own diagnostic head. On the paper's three cohorts it reaches 91.97 macro AUC / 72.59 macro F1 on PTB-XL, 87.31 / 67.40 on Georgia and 89.34 / 64.54 on CPSC2018, and it stores 3.93M parameters including the reconstruction pathway (1.87M in the path that produces a diagnosis). This repository is the reference tree that produced those numbers for the paper TRACE: Tractable Routing Autoencoder for Clinical ECG.
The TRACE architecture: residual 1D backbone, system inference network, ontology-partitioned clinical latent with its routed heads, and the reconstruction pathway.
The concrete layer stack behind each stage above, with detail panels for the residual block, the backbone, one diagnostic head and the upsampling stack.
| Cohort | Macro AUC | Macro F1 | Labels |
|---|---|---|---|
| PTB-XL | 91.97 ± 0.05 | 72.59 ± 0.13 | NORM, MI, STTC, CD, HYP |
| Georgia | 87.31 ± 0.09 | 67.40 ± 1.60 | NORM, MI, STTC, CD, HYP |
| CPSC2018 | 89.34 ± 0.43 | 64.54 ± 2.41 | NORM, AF, I-AVB, LBBB, RBBB, PAC, PVC, STD, STE |
Per-class AUC and macro F1 for TRACE across the three cohorts.
Macro AUC / macro F1 in percent, sorted by AUC.
| Model | Macro AUC | Macro F1 |
|---|---|---|
| TRACE (ours) | 91.97 | 72.59 |
| Inception1D | 91.81 | 71.89 |
| ECGFounder (frozen probe) | 91.43 | 70.06 |
| ResNet1D | 91.41 | 72.22 |
| FCN | 91.13 | 71.88 |
| ConceptBottleneck | 90.89 | 72.02 |
| Transformer1D | 90.20 | 70.25 |
| HeartLang (frozen probe) | 89.58 | 65.59 |
| 1D-MAE | 89.43 | 68.90 |
| 1D-SimCLR | 88.79 | 67.69 |
The two foundation models are also probed on the other cohorts: ECGFounder 86.38 / 65.42 on Georgia and 95.75 / 77.67 on CPSC2018; HeartLang 80.74 / 55.06 and 88.36 / 57.32.
| Component | Parameters |
|---|---|
| TRACE, as stored | 3.93M |
| — diagnostic path | 1.87M |
| — reconstruction path | 2.06M |
| ECGFounder | 30.8M |
| HeartLang (ST-ECGFormer) | 44.6M |
The reconstruction path does not change any prediction.
- Linear probes on the latent subspaces separate the diagonal (subspace to its own head) from the off-diagonal pairing.
- Targeted erasure drives the erased subspace's own head to chance (macro AUC 50.00) while the non-target heads hold their baseline values.
- Capacity-matched permutation control holds every head at 8 input dimensions, so routing and capacity are decoupled.
- Mechanism-layer zeroing of the depolarization and repolarization pathways, with R-peak-anchored QRS and ST-T windows, moves the reconstruction without moving the diagnosis.
Linear probes on the latent subspaces: diagonal pairings against their off-diagonal counterparts.
Ablations: unstructured latent, loss components, capacity-matched permutation control, latent dimensions, orthogonality weight and reconstruction fidelity.
Reconstruction overlays for two PTB-XL test recordings over the 10-second segment: original in blue and reconstruction in red for all 12 leads, two per row. Top: sample index 42; bottom: sample index 100.
Three public 12-lead ECG cohorts, all distributed by PhysioNet under CC BY 4.0. None of them is bundled with this repository, and data/ is git-ignored.
| Cohort | PhysioNet record | Version | Records | Role |
|---|---|---|---|---|
| PTB-XL | ptb-xl | 1.0.3 | 20,008 label-filtered | Main cohort, 5 superclasses |
| Georgia | challenge-2020 group georgia |
1.0.2 | 10,344 | Cross-cohort evaluation, same 5 superclasses |
| CPSC2018 | challenge-2020 group cpsc_2018 |
1.0.2 | 6,877 | Rhythm-dominant 9-label cohort |
Split sizes, PTB-XL superclass prevalence, and Georgia and CPSC2018 label prevalence.
Expected layout. --data-root defaults to ./data; each loader reads its own subdirectory.
data/
├── ptb-xl/ # dataset ptbxl
│ ├── ptbxl_database.csv # metadata + scp_codes labels
│ └── records500/ # WFDB records named by the filename_hr column
├── georgia/ # dataset georgia; any directory holding the records
│ ├── <record>.hea # header carrying the "# Dx:" SNOMED-CT line
│ └── <record>.mat # 12-lead signal
└── CPSC2018/ # dataset cpsc
├── REFERENCE.csv # Recording, First_label, Second_label, Third_label
└── ecg_signals/ # <record>.mat or <record>.npy; falls back to the root
- Python 3.12, PyTorch 2.5.1, CUDA 12.4, Ubuntu 22.04, one NVIDIA RTX 4090 (24 GB). Seeds are 42 / 2025 / 2026 throughout.
pip install -r requirements.txtinstalls the runtime dependencies (torch, numpy, scipy, pandas, scikit-learn, scikit-multilearn, wfdb, matplotlib, tqdm).- A GPU is needed for training; evaluation, the analyses that read trained weights and every figure script run on CPU.
Every command runs from the repository root, and the steps are a chain: each one reads what the previous steps wrote. Steps 1 and 2 apply to every cohort; steps 3 to 9 only need the cohort whose numbers you want.
python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt# PTB-XL 1.0.3
wget -r -N -c -np https://physionet.org/files/ptb-xl/1.0.3/
# Challenge 2020 v1.0.2, which carries both the Georgia and the CPSC2018 groups
wget -r -N -c -np https://physionet.org/files/challenge-2020/1.0.2/Arrange them as described under Data above, then set --data-root accordingly
on every later command. Everything downstream writes under --output-root
(./outputs by default) and reads it back.
Every PTB-XL analysis and ablation below reads preprocessed signals from
<output-root>/ptbxl/ptbxl_cache/{train,val,test}.npz. That cache is a required
input: prepare_ptbxl_cache is its only producer, and no consumer decodes the
records on the fly when it is missing.
python src/prepare_ptbxl_cache.py --data-root ./data --output-root ./outputsIt decodes the official folds through trace.data.PTBXLDataset and writes X,
y and mask per split; expect a few minutes. Re-run with --force to rebuild
a split that already exists. Skip this step only if you run nothing that reads
PTB-XL test signals (steps 4, 5, 9 and 10).
python src/train.py --dataset ptbxl --data-root ./data --output-root ./outputs
python src/train.py --dataset georgia --data-root ./data --output-root ./outputs
python src/train.py --dataset cpsc --data-root ./data --output-root ./outputspython src/evaluate.py --dataset ptbxl --data-root ./data --output-root ./outputs
python src/evaluate.py --dataset georgia --data-root ./data --output-root ./outputs
python src/evaluate.py --dataset cpsc --data-root ./data --output-root ./outputspython src/baselines/train_baselines.py --dataset ptbxl --data-root ./data --output-root ./outputspython src/baselines/eval_baselines.py --dataset ptbxl --data-root ./data --output-root ./outputspython src/foundation/ecgfounder_probe.py --dataset ptbxl --data-root ./data \
--ecgfounder-root <path-to-ecgfounder> --ckpt <path-to-ecgfounder-checkpoint>
python src/foundation/heartlang_probe.py --dataset ptbxl --data-root ./data \
--heartlang-root <path-to-heartlang> \
--ckpt <path-to-heartlang>/checkpoints/pretrain/MIMIC-IV/checkpoint-200.pth# routing contract: linear probes and targeted subspace erasure
python src/analyses/probes.py --dataset ptbxl --data-root ./data --output-root ./outputs
python src/analyses/erasure.py --dataset ptbxl --data-root ./data --output-root ./outputs
# mechanism layer: zeroing the depolarization and repolarization pathways
python src/analyses/perturbation.py --dataset ptbxl --data-root ./data --output-root ./outputs
# robustness
python src/analyses/robustness.py --dataset ptbxl --data-root ./data --output-root ./outputs \
--models TRACE FCN ResNet1D Inception1D Transformer1D ConceptBottleneck 1D-MAE 1D-SimCLR
# robustness of the two frozen-feature foundation probes
python src/analyses/foundation_robustness.py --dataset ptbxl --data-root ./data \
--models ECGFounder --ecgfounder-root /path/to/ECGFounder-main \
--ecgfounder-ckpt /path/to/12_lead_ECGFounder.pth \
--split-seed 42 --seeds 42 2025 2026 --out foundation_rob_ptbxl.csv
# reconstruction fidelity and the overlay export
python src/analyses/reconstruction.py --dataset ptbxl --output-root ./outputs
# parameter counts for the model-size table, the only analysis that needs no trained weights
python src/analyses/count_params.py --dataset ptbxl --output-root ./outputs# unstructured latent
python src/ablations/train_unstructured.py --dataset ptbxl --data-root ./data --output-root ./outputs
python src/ablations/train_unstructured.py --dataset georgia --data-root ./data --output-root ./outputs
python src/ablations/eval_unstructured.py --dataset ptbxl --data-root ./data --output-root ./outputs
# loss components
python src/ablations/component_ablation.py --dataset ptbxl --data-root ./data --output-root ./outputs
# capacity-matched permutation control
python src/ablations/train_permutation.py --dataset ptbxl --data-root ./data --output-root ./outputs --assignment true
python src/ablations/train_permutation.py --dataset ptbxl --data-root ./data --output-root ./outputs --assignment A
python src/ablations/train_permutation.py --dataset ptbxl --data-root ./data --output-root ./outputs --assignment B
python src/ablations/train_permutation.py --dataset ptbxl --data-root ./data --output-root ./outputs --assignment C
python src/ablations/eval_permutation.py --dataset ptbxl --data-root ./data --output-root ./outputs
# all four assignments must be trained first; eval_permutation evaluates the four by
# default and silently skips any whose checkpoint is missing
# latent dimensions
python src/ablations/dim_sweep.py --dataset ptbxl --data-root ./data --output-root ./outputs
# orthogonality weight
python src/ablations/lambda_sweep.py --dataset ptbxl --data-root ./data --output-root ./outputspython figures/georgia_stats.py # needs the Georgia .hea headers; writes figures/source/georgia_stats.csv
python figures/make_dataset_figure.py # dataset figure
python figures/make_classwise_group.py # class-wise figure
python figures/make_probe_2x2.py # probing figure
python figures/make_ablation_group.py # ablation figure
python figures/plot_recon_overlay.py # reconstruction-overlay figure├── src/
│ ├── trace/ Library: models.py, baselines.py,
│ │ data.py, metrics.py, config.py
│ ├── prepare_ptbxl_cache.py The only producer of ptbxl_cache/; every
│ │ PTB-XL consumer below requires that cache
│ ├── train.py Train TRACE on PTB-XL, Georgia or CPSC2018
│ ├── evaluate.py Main-table row, per-class thresholds, class-wise table
│ ├── baselines/
│ │ ├── train_baselines.py FCN, ResNet1D, Inception1D, Transformer1D,
│ │ │ ConceptBottleneck, 1D-MAE, 1D-SimCLR
│ │ └── eval_baselines.py Baseline rows of the main and robustness tables
│ ├── ablations/
│ │ ├── train_unstructured.py Unstructured-latent ablation
│ │ ├── eval_unstructured.py Its evaluation
│ │ ├── component_ablation.py Loss-component ablation
│ │ ├── train_permutation.py Capacity-matched permutation control
│ │ ├── eval_permutation.py Fold-10 optimal-threshold evaluation of the control
│ │ ├── dim_sweep.py Latent-dimension sweep
│ │ └── lambda_sweep.py Orthogonality-weight sweep
│ ├── analyses/
│ │ ├── probes.py Linear probes on the latent subspaces
│ │ ├── erasure.py Targeted subspace erasure
│ │ ├── perturbation.py Mechanism-layer pathway zeroing
│ │ ├── robustness.py Four-condition robustness evaluation
│ │ ├── foundation_robustness.py The same four conditions for the frozen probes
│ │ ├── reconstruction.py Reconstruction fidelity and overlay export
│ │ └── count_params.py Parameter counts
│ └── foundation/
│ ├── ecgfounder_probe.py ECGFounder frozen-feature linear probe
│ └── heartlang_probe.py HeartLang frozen-feature linear probe
├── figures/
│ ├── architecture.png Architecture figure, drawn by the authors
│ ├── module_stack.png Implementation-stack figure, drawn by the authors
│ ├── georgia_stats.py make_dataset_figure.py make_classwise_group.py
│ ├── make_probe_2x2.py make_ablation_group.py plot_recon_overlay.py
│ ├── source/ Reference CSV tables the figure scripts read
│ └── output/ Figures generated by the scripts above
├── requirements.txt
├── LICENSE
└── README.md
MIT — see LICENSE. The datasets keep their own licences: PTB-XL and Challenge 2020 v1.0.2 each state CC BY 4.0, which requires attribution when the data are reused, so cite the dataset records alongside this repository. The third-party foundation-model weights keep the terms of their own releases.