Massive cleanup for paper submission

.github/CODEOWNERS

@@ -3,14 +3,7 @@
 .gitmodules @StevenSong
 .pre-commit-config.yaml @StevenSong
 .secrets.baseline @StevenSong
-configs/audio @StevenSong
-configs/multi-arm-heedb @StevenSong
-configs/pretrain-supervised.yaml @StevenSong
-configs/pretrain-unsupervised.yaml @StevenSong
-configs/target-blackbox.yaml @StevenSong
-configs/target-guided-14ppl.yaml @StevenSong
-configs/target-guided-5ppl.yaml @StevenSong
-configs/target-unguided.yaml @StevenSong
+configs/ @StevenSong
 data-preprocessing/ @StevenSong
 env.yaml @StevenSong
 external/PierreElias-IntroECG @StevenSong
@@ -41,7 +34,9 @@ protossl/lightning_utils.py @StevenSong
 protossl/models/__init__.py @StevenSong
 protossl/models/_base_classifier.py @StevenSong
 protossl/models/_blackbox_classifier.py @StevenSong
+protossl/models/_blackbox_contraster.py @StevenSong
 protossl/models/_pretrained_utils.py @StevenSong
+protossl/models/_protoecgnet.py @StevenSong
 protossl/models/_prototype_assigner.py @StevenSong
 protossl/models/_prototype_classifier.py @StevenSong
 protossl/models/_prototype_contraster.py @StevenSong
@@ -60,7 +55,9 @@ protossl/models/helpers/_prototype_ilp_assigner.py @sahilsethi0105
 protossl/models/layers/__init__.py @StevenSong
 protossl/models/layers/_multi_input_linear.py @StevenSong
 protossl/models/layers/_panns_backbones.py @StevenSong
+protossl/plotting.py @StevenSong
 protossl/trainer.py @StevenSong
+protossl/protoecgnet_trainer.py @StevenSong
 pyproject.toml @StevenSong
 README.md @StevenSong @sahilsethi0105
 requirements.txt @StevenSong
@@ -72,40 +69,33 @@ scripts/_slurm_wrapper.sh @StevenSong
 scripts/_submit_job.sh @StevenSong
 scripts/0-run-cache-data.sh @StevenSong
 scripts/1-run-blackbox-direct.sh @StevenSong
-scripts/2-run-labsup-proto-direct.sh @StevenSong
+scripts/2-run-supproto-direct.sh @StevenSong
 scripts/2-z-* @StevenSong
-scripts/3-run-protossl-heedb-pila.sh @StevenSong
-scripts/3-z-* @StevenSong
-scripts/4-run-labsup-proto-heedb-rila.sh @StevenSong
-scripts/4-y-* @StevenSong
-scripts/4-z-* @StevenSong
+scripts/3-* @StevenSong
+scripts/4-* @StevenSong
 scripts/5-1-run-ecgfounder-logreg.sh @StevenSong
 scripts/5-2-run-stmem-logreg.sh @StevenSong
+scripts/5-3-run-pclr-logreg.sh @StevenSong
 scripts/6-run-protossl-heedb-pia.sh @StevenSong
 scripts/6-z-* @StevenSong
-scripts/7-run-protossl-heedb-pit.sh @StevenSong
-scripts/8-run-protossl-heedb-pip.sh @StevenSong
+scripts/7-* @StevenSong
 scripts/9-0-run-ecgfounder-patches.sh @StevenSong
 scripts/9-1-run-ecgfounder-lap.sh @StevenSong
 scripts/9-2-run-ecgfounder-clustering.sh @StevenSong
 scripts/9-3-run-ecgfounder-random.sh @StevenSong
 scripts/audio @StevenSong
-scripts/ecg-fms/_compute_ecgfounder_embeddings.py @StevenSong
-scripts/ecg-fms/_compute_stmem_embeddings.py @StevenSong
+scripts/ecg-fms/ @StevenSong
 scripts/echonext/_tabular_logreg.py @StevenSong
 scripts/echonext/run-columbia-minimodel.sh @StevenSong
 scripts/echonext/run-tabular-logreg.sh @StevenSong
-scripts/pretrain/_cache_heedb.py @StevenSong
-scripts/pretrain/_compute_heedb_normalizations.py @StevenSong
-scripts/pretrain/run-heedb-normalizations.sh @StevenSong
-scripts/pretrain/run-pass-heedb-pretrain-no-attn.sh @StevenSong
-scripts/pretrain/run-pass-heedb-pretrain.sh @StevenSong
-scripts/pretrain/run-prosup-heedb-pretrain.sh @StevenSong
+scripts/pretrain/ @StevenSong
 scripts/prototypes-from-fms @StevenSong
 scripts/queue-experiments.sh @StevenSong
+scripts/queue-pclr-logreg.sh @StevenSong
 scripts/README.md @StevenSong @sahilsethi0105
 user-study/images @StevenSong
 user-study/metadata.csv @StevenSong
 user-study/prepare_samples.ipynb @StevenSong
 user-study/results.csv @StevenSong
 user-study/user_study_analysis.ipynb @sahilsethi0105
+large-user-study/ @StevenSong

.gitignore

@@ -219,3 +219,6 @@ OLD/
 slurm-logs/
 results/figs/
 ecgfounder-checkpoint/
+wandb/
+.devcontainer/
+user-study/images-pdf

README.md

@@ -1,37 +1,52 @@
-# ProtoSSL: Interpretable Prototype Learning from Unlabeled Time-Series Data
+# ProtoSSL: Self-Supervised Pretraining and Downstream Transfer for Projection-Based Prototype Models
 
-This repo contains the source code for **ProtoSSL**, a novel framework for label-free learning of interpretable, projection-based prototypes that are readily adaptable to downstream tasks. Our key innovation is to separate *motif discovery* from *label alignment*. ProtoSSL first learns a reusable prototype bank using a self-supervised objective applied directly to prototype activations, and then aligns these prototypes to downstream tasks through a novel and efficient assignment procedure. We study ProtoSSL and provide code for our experiments primarily over ECGs and additionally audio waveforms.
+This repo contains the source code for **ProtoSSL**, a novel framework for pretraining a foundational latent prototype bank on unlabeled data and transferring it to downstream tasks to create interpretable, projection-based prototype models. Our key innovation is to separate *motif discovery* from *label alignment*. ProtoSSL first learns a reusable latent prototype bank using a self-supervised objective applied directly to prototype activations, and then aligns these prototypes to downstream tasks through a novel and efficient assignment procedure. We study ProtoSSL and provide code for our experiments primarily over ECGs and additionally audio waveforms.
 
 ### Citation
 ```
-TODO: bibtex citation
+@article{song2026protossl,
+  title={ProtoSSL: Self-Supervised Pretraining and Downstream Transfer for Projection-Based Prototype Models},
+  author={Song, Steven and Sethi, Sahil and Beaulieu-Jones, Brett and Grossman, Robert L},
+  journal={arXiv preprint arXiv:2605.06943},
+  year={2026}
+}
 ```
 
-## Contributing
+## Repo Organization
+
+| Path | Contents |
+|------|----------|
+| `protossl/` | The `protossl` package: datasets, models, and the Lightning trainers (`protossl.trainer` for ProtoSSL and baselines, `protossl.protoecgnet_trainer` for multi-arm ProtoECGNet) |
+| `configs/` | Lightning CLI configs for pretraining (`pretrain-*.yaml`) and downstream tasks (`target-*.yaml`), with audio configs under `configs/audio/` |
+| `scripts/` | Run scripts for every experiment in the paper, plus the slurm queue scripts that submit them. See [`scripts/README.md`](scripts/README.md) |
+| `user-study/` | Sample preparation and analysis for the user study (`user-study/images` is a submodule of the de-identified study images) |
+| `results/` | Notebooks that produce the figures and tables in the paper from the experiment outputs |
+| `data-preprocessing/` | Notebooks that create the train/val/test splits and nested training subsets, plus precomputed HEEDB normalization statistics. See [`data-preprocessing/README.md`](data-preprocessing/README.md) |
+| `external/` | Submodules for baselines: [ST-MEM](https://github.com/vuno/ST-MEM) and our fork of [IntroECG](https://github.com/PierreElias/IntroECG) (EchoNext minimodel) |
 
-**Repo Organization:** The outline of this repo is centered around the package `protossl`, defined in the subdirectory of the same name. While the package definition enables experimentation, the `scripts` folder defines all of the actual experiment run scripts. Most of the experimental comparisons we run have standalone implementations also contained within the `scripts` folder.
+## Setup
 
-**Environment Setup:** We use pre-commit hooks to maintain coding style and to enforce data stewardship. Source data should never be committed to this repo, including exploratory notebooks which may accidentally leak source data. The pre-commit hooks do not necessarily prevent this and can be circumvented, but may help prevent obvious cases. To correctly set up your environment:
 ```bash
-# 1) clone repo
-git clone git@github.com:StevenSong/ProtoSSL.git
+# 1) clone repo, including submodules (only needed for the ST-MEM and EchoNext minimodel baselines, and the user study images)
+git clone --recurse-submodules git@github.com:StevenSong/ProtoSSL.git
 cd ProtoSSL
 
 # 2) create and activate environment
 # NOTE: if you don't use conda, make sure you're using the same python version, install from `requirements.txt`, and MAKE SURE YOU HAVE FFMPEG 5.* FOR TORCHCODEC (see below)
 conda env create -f env.yaml
 conda activate protossl
 
-# 3) enable pre-commit hooks
-pre-commit install
-
-# 4) install protossl from editable source
+# 3) install protossl from editable source
 pip install -e .
 
-# 5) dev away
+# 4) (contributors only) enable pre-commit hooks, see Contributing below
+pre-commit install
 ```
 
-**torchcodec:** torchcodec is a bit fragile with dependencies. We've pinned `torch==2.7.0` which is compatible with `torchcodec==0.4.0`, both compiled against CUDA 12.8 (which we use on our machines). This torchcodec version is only compatible with `datasets==4.0.0`. If you see errors relating to torchcodec (you can diagnose this by just importing torchcodec), make sure the dependencies are compatible not just relative to versioning, but also relating to the CUDA versions. We also use `ffmpeg=5.*` installed via conda. If you see an error relating to not being able to find `libnppicc.so.12`, it might be that the linker can't find the binaries (which we ensure are available by installing `nvidia-npp-cu12`). To fix this, you can try setting the `LD_LIBRARY_PATH` environment variable:
+<details>
+<summary><b>Troubleshooting torchcodec</b></summary>
+
+torchcodec is a bit fragile with dependencies. We've pinned `torch==2.7.0` which is compatible with `torchcodec==0.4.0`, both compiled against CUDA 12.8 (which we use on our machines). This torchcodec version is only compatible with `datasets==4.0.0`. If you see errors relating to torchcodec (you can diagnose this by just importing torchcodec), make sure the dependencies are compatible not just relative to versioning, but also relating to the CUDA versions. We also use `ffmpeg=5.*` installed via conda. If you see an error relating to not being able to find `libnppicc.so.12`, it might be that the linker can't find the binaries (which we ensure are available by installing `nvidia-npp-cu12`). To fix this, you can try setting the `LD_LIBRARY_PATH` environment variable:
 ```bash
 export LD_LIBRARY_PATH=$CONDA_PREFIX/lib/python3.10/site-packages/nvidia/npp/lib:$LD_LIBRARY_PATH
 # test by importing torchcodec in a python runtime
@@ -47,3 +62,49 @@ echo 'export LD_LIBRARY_PATH=$CONDA_PREFIX/lib/python3.10/site-packages/nvidia/n
 # Unset on deactivate
 echo 'export LD_LIBRARY_PATH=$(echo $LD_LIBRARY_PATH | sed "s|$CONDA_PREFIX/lib/python3.10/site-packages/nvidia/npp/lib:||g")' > $CONDA_PREFIX/etc/conda/deactivate.d/npp_lib.sh
 ```
+</details>
+
+## Data
+
+We do not redistribute any data; each dataset must be obtained from its original source under its own terms of use.
+
+* **ECG pretraining:** HEEDB
+* **ECG downstream:** EchoNext, PTB-XL, CinC 2020 (Georgia), MIMIC-IV-ECG, ZZU pECG, CODE-15%
+* **Audio pretraining:** AudioSet
+* **Audio downstream:** UrbanSound8K, IEMOCAP, VoxCeleb1
+
+The `data-preprocessing/` notebooks create the splits and nested training subsets (e.g. `ptb-xl-1k`) used for our data-efficiency experiments.
+
+## Reproducing Our Experiments
+
+1. Obtain the datasets above and create the splits and training subsets with the `data-preprocessing/` notebooks
+1. Find and replace the hardcoded paths in the run, queue, and pretrain scripts: `/opt/gpu_working/steven` (repo and outputs) and `/opt/gpudata` (source datasets)
+1. Pretrain: run the scripts in [`scripts/pretrain/`](scripts/pretrain/) for ECG and `scripts/audio/run-pretrain-audio-*.sh` for audio
+1. Submit the downstream experiments via slurm: [`scripts/queue-experiments.sh`](scripts/queue-experiments.sh) (ECG) and [`scripts/audio/queue-experiments.sh`](scripts/audio/queue-experiments.sh) (audio)
+1. Produce figures and tables with the `results/` notebooks
+
+See [`scripts/README.md`](scripts/README.md) for more detail. Each run script writes to `$RUN_DIR/<experiment-name>`, named `<model>-<pretraining data>-<transfer method>[-suffixes]`:
+
+* **Models:** `protossl`, `supproto` (supervised prototypes), `blackbox`, and linear probes over foundation model embeddings (`ecgfounder-logreg`, `stmem-logreg`, `pclr-logreg`)
+* **Transfer methods:** `pila` (project in linear assignment), `rila` (reproject in linear assignment), `pia` (project in assignment, using ProtoPool-based assignment), `pit` (project in transfer dataset without assignment or label supervision), `pip` (project in pretraining dataset without assignment or label supervision), `latent` (no projection), `direct` (trained only on the target dataset)
+* **Suffixes:** `-ft` is fine-tuned, `-Nppl` uses N prototypes per label, `heedb-150` pretrains over the 150-label HEEDB set
+
+### Environment Variables
+
+| Variable | Effect |
+|----------|--------|
+| `CACHE_DIR` | Where preprocessed datasets are cached (default `~/.cache/protossl_cache`) |
+| `INVALIDATE_CACHE` | Set to `"yes im sure"` to delete the cache on import |
+| `USE_HEEDB_150` | Use the 150-label HEEDB label set instead of the default (set by the `*-heedb-150-*` pretraining scripts) |
+| `HIGH_MEMORY` | Load all of HEEDB into memory (~500 GB of RAM) rather than streaming it from disk (set by the HEEDB pretraining scripts) |
+| `US8K_TEST_FOLD`, `IEMOCAP_TEST_FOLD` | Held-out fold for UrbanSound8K and IEMOCAP cross-validation (default `0`) |
+
+The ECGFounder baseline downloads its checkpoint from Hugging Face to `ecgfounder-checkpoint/` under the working directory (`scripts/` for the run scripts) on first use.
+
+## Contributing
+
+We use pre-commit hooks to maintain coding style and to enforce data stewardship. Source data should never be committed to this repo, including exploratory notebooks which may accidentally leak source data. The pre-commit hooks (which strip notebook outputs, scan for secrets, and block direct commits to `main`) do not necessarily prevent this and can be circumvented, but may help prevent obvious cases. Every tracked file must also be covered by `.github/CODEOWNERS`.
+
+## License
+
+Apache 2.0, see [`LICENSE`](LICENSE).

configs/audio/pretrain-supervised-v2.yaml

@@ -0,0 +1,60 @@
+# lightning.pytorch==2.6.0
+seed_everything: 42
+trainer:
+  devices: 1
+  precision: 32
+  logger:
+    class_path: protossl.lightning_utils.StrictWandbLogger
+    init_args:
+      project: ProtoSSL-Audio
+      # save_dir: /path/to/runs/
+      # name: exp-name
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.EarlyStopping
+      init_args:
+        monitor: val_loss
+        mode: min
+        patience: 10 # check reduce lr on plateau
+    - class_path: lightning.pytorch.callbacks.LearningRateMonitor
+      init_args:
+        logging_interval: epoch
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        monitor: val_loss
+        mode: min
+        save_last: True
+    - class_path: PredictionWriter
+  max_epochs: 1000
+  log_every_n_steps: 1
+optimizer:
+  class_path: torch.optim.AdamW
+  init_args:
+    lr: 0.001
+    weight_decay: 0.01
+lr_scheduler:
+  class_path: lightning.pytorch.cli.ReduceLROnPlateau
+  init_args:
+    monitor: val_loss
+    patience: 5 # check early stopping
+model:
+  class_path: LitModel
+  init_args:
+    backbone_type: Cnn14
+    conv_type: PANNS
+    input_channels: 1
+    prototype_type: partial
+    partial_len: 32000 # 1-second @ 32 kHz
+    partial_overlap: 0.5
+    prototype_h: 1
+    prototype_w: 1
+    n_prototypes_per_label: 12
+    # pretrained_weights: /path/to/weights
+data:
+  class_path: LitData
+  init_args:
+    # dataset_path: /path/to/dataset
+    sampling_rate: 32000
+    batch_size: 128
+    num_workers: 12
+    prefetch_factor: 2
+# pipeline_stage: learn-prototypes-supervised|project-prototypes-supervised|compute-embeddings|train-classifier

configs/audio/pretrain-supervised.yaml

@@ -24,7 +24,7 @@ trainer:
         mode: min
         save_last: True
     - class_path: PredictionWriter
-  max_epochs: 5
+  max_epochs: 1000
   log_every_n_steps: 1
 optimizer:
   class_path: torch.optim.AdamW

configs/audio/pretrain-unsupervised-v2.yaml

@@ -0,0 +1,63 @@
+# lightning.pytorch==2.6.0
+seed_everything: 42
+trainer:
+  devices: 1
+  precision: 32
+  logger:
+    class_path: protossl.lightning_utils.StrictWandbLogger
+    init_args:
+      project: ProtoSSL-Audio
+      # save_dir: /path/to/runs/
+      # name: exp-name
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.EarlyStopping
+      init_args:
+        monitor: val_loss
+        mode: min
+        patience: 10 # check reduce lr on plateau
+    - class_path: lightning.pytorch.callbacks.LearningRateMonitor
+      init_args:
+        logging_interval: epoch
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        monitor: val_loss
+        mode: min
+        save_last: True
+    - class_path: PredictionWriter
+  max_epochs: 1000
+  log_every_n_steps: 1
+optimizer:
+  class_path: torch.optim.AdamW
+  init_args:
+    lr: 0.001
+    weight_decay: 0.01
+lr_scheduler:
+  class_path: lightning.pytorch.cli.ReduceLROnPlateau
+  init_args:
+    monitor: val_loss
+    patience: 5 # check early stopping
+model:
+  class_path: LitModel
+  init_args:
+    backbone_type: Cnn14
+    conv_type: PANNS
+    input_channels: 1
+    prototype_type: partial
+    partial_len: 32000 # 1-second @ 32 kHz
+    partial_overlap: 0.5
+    prototype_h: 1
+    prototype_w: 1
+    n_prototypes: 6324 # 527 audioset labels, comparing to 12 prototypes per label for supervised pretraining and to accomodate 5 PPL in highest cardinality downstream task
+    model_kwargs: '{"cola_loss_weight": 1, "clar_loss_weight": 0, "koleo_loss_weight": 1}'
+    # pretrained_weights: /path/to/weights
+data:
+  class_path: LitData
+  init_args:
+    # dataset_path: /path/to/dataset
+    sampling_rate: 32000
+    batch_size: 128
+    num_workers: 12
+    prefetch_factor: 2
+    data_wrapper_kwargs: '{"cola_view_seconds": 2}'
+contrastive_pair_mode: cola
+# pipeline_stage: learn-prototypes|project-prototypes|compute-embeddings|train-classifier

configs/audio/pretrain-unsupervised.yaml

@@ -24,7 +24,7 @@ trainer:
         mode: min
         save_last: True
     - class_path: PredictionWriter
-  max_epochs: 5
+  max_epochs: 1000
   log_every_n_steps: 1
 optimizer:
   class_path: torch.optim.AdamW
@@ -58,6 +58,6 @@ data:
     batch_size: 128
     num_workers: 12
     prefetch_factor: 2
-    data_kwargs: '{"cola_view_seconds": 2}'
+    data_wrapper_kwargs: '{"cola_view_seconds": 2}'
 contrastive_pair_mode: cola+clar
 # pipeline_stage: learn-prototypes|project-prototypes|compute-embeddings|train-classifier