Contrastive Learning with Uncertainty Filtering on NeurIPS Multimodal Single-Cell Data

This repository trains a supervised contrastive (InfoNCE / SupCon) encoder on the NeurIPS 2021 multimodal single-cell integration benchmark, evaluates cell-type classification from the learned representations, and shows that MC-dropout uncertainty filtering improves held-out accuracy.

Dataset

The data come from the NeurIPS 2021 multimodal single-cell data integration challenge:

Multimodal single cell data integration challenge: Results and lessons learned
Lance et al., PMLR 176:162–176, 2022.

A central property of this dataset is batch effects. Cells were profiled across different laboratory sites and different donors. Technical and donor-driven variation can dominate gene-expression space, so models that ignore batch structure often learn site/donor signatures instead of cell-type biology.

We encode each site×donor combination as a batch ID (e.g. s1d1, s4d8) and use an explicit batch holdout split:

Split	Batches
Train	s1d1, s1d2, s1d3, s2d1, s2d4
Test	s3d6, s3d7, s4d8, s4d9

Train and test therefore come from disjoint sites and donors, which makes cell-type prediction a realistic generalization problem under batch shift.

Place the raw AnnData file as neurips.h5ad in this directory, then run:

python data_neurips_split.py

This writes data/neurips_batch_train.h5ad and data/neurips_batch_test.h5ad (HVG selection on train, shared gene set on test).

Method

1. Contrastive representation learning

We train an MLP encoder (InfoNCEEncoder in contrastive.py) on log-normalized counts. For each minibatch, embeddings of cells with the same cell-type label are pulled together and other cells are pushed apart via a supervised contrastive (SupCon / InfoNCE-style) loss. Dropout is enabled in the trunk so the network can later be queried in train mode for uncertainty.

python train_contrastive_neurips.py

Checkpoint: checkpoints/contrastive_neurips_infonce.pt.

2. Classification from representations

At test time we keep the encoder frozen, embed train and test cells, and classify test cells with a k-NN classifier in embedding space (cosine distance, default k=5). This measures how well the contrastive features transfer across unseen labs/donors.

python evaluate_contrastive_neurips.py

Learned embeddings organize by cell type while still reflecting the site×donor batch structure of the holdout setting:

Train vs held-out test batches in the same embedding space:

3. Uncertainty filtering with MC dropout

Even with strong representations, some test cells are ambiguous under batch shift. We estimate per-cell epistemic uncertainty by Monte Carlo dropout sampling:

1. Keep dropout active and draw N stochastic embeddings for each test cell.
2. Measure pairwise cosine similarity across samples; uncertainty is 1 - mean similarity (high when embeddings disagree).
3. Rank test cells by uncertainty and remove the most uncertain fraction.
4. Recompute k-NN accuracy / macro-F1 on the remaining cells.

We compare against removing the same fraction of cells at random. Preferentially dropping uncertain cells improves accuracy and macro-F1 more than random removal, which indicates that the dropout samples identify hard or batch-confounded cells rather than discarding signal arbitrarily.

python evaluate_infonce_uncertainty_neurips.py

Figures are also written under figures/ when you re-run evaluation.

Repository layout

File	Role
contrastive.py	Encoder + SupCon loss + MC-dropout uncertainty
data_neurips_split.py	Site×donor train/test split and HVG prep
train_contrastive_neurips.py	Train InfoNCE encoder
evaluate_contrastive_neurips.py	k-NN cell-type eval + UMAPs
evaluate_infonce_uncertainty_neurips.py	Uncertainty filtering curves
assets/	README figures from NeurIPS runs
requirements.txt	Python dependencies

Setup

pip install -r requirements.txt

GPU is used automatically when available.