Contrastive Learning in Biomedical AI

📖 What is Contrastive Learning?

Contrastive learning trains representations by pulling semantically related views together in an embedding space while pushing unrelated views apart. Unlike reconstruction objectives (VAEs, autoencoders) or token prediction (masked language modeling), contrastive losses operate directly on geometry of the embedding — defining what is similar by what the practitioner pairs together.

Why It Matters for Biology

• Scalability without labels: Single-cell atlases routinely exceed 10M cells, but most lack consistent annotations. Contrastive objectives turn the dataset itself into supervision — each cell becomes its own anchor.
• Cross-modal alignment: CLIP-style objectives align modalities that share no feature space (H&E images ↔ transcriptomes, CT scans ↔ radiology reports, cells ↔ text) into a single embedding where retrieval, zero-shot classification, and chat-based query become possible.
• Spatial neighborhood structure: Graph-contrastive variants (DGI, GCL) exploit spatial adjacency as a natural notion of positive pair, encoding tissue architecture into cell embeddings.
• Discriminating near-identical inputs: Cancer-vs-wild-type peptides differ by one amino acid; supervised contrastive losses can force those embeddings apart where cross-entropy struggles.
• Atlas-scale reference mapping: Once a contrastive embedding is trained, new query cells project into it in seconds — Concerto maps 10K query cells onto a 10M cell reference in 8s.

The Four Contrastive Families You'll See

The same paper often combines multiple families — e.g., MedMPT uses both intra-modal (CT slice ↔ CT slice) and inter-modal (CT ↔ report) contrastive losses alongside masked reconstruction.

🔍 Visual Glossary: Hard Concepts at a Glance

A pictorial cheat sheet for the trickier ideas behind the methods below — InfoNCE geometry, EMA teacher–student, graph corruption, similarity matrices, the modality gap. Each diagram is the intuition, not the math.

📅 Evolution Timeline: From One-Modality Cells to Multimodal Foundation Models

🔬 Method Taxonomy: Four Contrastive Families

Choosing a Contrastive Recipe

⚖️ Side-by-Side: Self-Augmentation vs. Cross-Modal Pairing

Concerto (2022) and OmiCLIP (2025) both minimise InfoNCE on the unit hypersphere, but they form positives — and avoid collapse — in completely opposite ways: one cell viewed twice through dropout vs. a matched H&E tile and transcriptome profile.

📄 Landmark Papers (2022–2025)

🔵 Instance Discrimination on Single Cells

📊 Graph Contrastive Learning

🔗 Cross-Modal CLIP-Style Models

🎯 Foundation Models: Contrastive vs. Masked-LM Pretraining

🧭 Putting Contrastive Learning into Practice

Designing Positive & Negative Pairs

The most important decision is what you call a positive pair — it defines the geometry of the resulting space. Get this wrong and the embedding will collapse onto the wrong invariance.

Common Pitfalls

• Batch shortcut: If samples within a batch share technical artifacts, the model learns to discriminate batches rather than biology. Mitigate with source-aware batch norm or domain adaptation (Concerto's approach), or careful batch construction.
• Modality gap: CLIP-style models often produce embeddings that cluster by modality before clustering by semantics. Watch for this in retrieval failure modes — even when same-modality retrieval works, cross-modal retrieval may not.
• Caption leakage: When using LLM-generated captions for cell or image data, the captions can encode information (cell type names, disease labels) that the contrastive loss then trivially recovers. Audit captions before celebrating zero-shot performance.
• Augmentation that erases biology: Gene dropout / masking is a popular augmentation but can remove marker genes for rare cell types. Test that augmentations preserve the variability of interest.
• Temperature tuning: The InfoNCE temperature τ controls how concentrated representations become on the hypersphere. Too low → collapse; too high → uniform. Concerto uses τ=0.1; MedMPT uses τ=0.07. Tune on a validation downstream task.
• Evaluating on the training distribution: Contrastive models often look great on in-distribution retrieval but degrade on truly novel cell types or tissues. Always hold out a tissue / disease / technology, not just cells.

When NOT to use contrastive learning

Software & Reproducibility

• Concerto: TensorFlow reference; sparse index-value gene encoding; multi-GPU distributed training.
• GraphST: PyTorch + PyG; clean baseline for graph-contrastive on spatial data.
• Nicheformer: Pretrained weights and 110M-cell corpus on Hugging Face; HF Datasets integration.
• Loki / OmiCLIP: End-to-end platform for histology ↔ spatial-transcriptomics alignment, retrieval, and prediction.
• CellWhisperer: Geneformer + BioBERT encoders, Mistral-7B chat layer, hosted demo at cellwhisperer.bocklab.org.
• SpatialEx: Hypergraph + UNI pathology foundation model; tutorials at spatialex-tutorials.readthedocs.io.

🛠️ Hands-On Practice

The sections below give you a concrete path from a blank environment to a working contrastive training loop for single-cell data — with the packages, data structures, a runnable snippet, and the pitfalls you'll actually hit.

Environment & packages

Install into a fresh conda or venv environment. lightly wraps common SSL objectives (NT-Xent, MoCo, BYOL) and works with any PyTorch encoder. info-nce-pytorch is a minimal InfoNCE reference implementation. pytorch-lightning handles the training loop boilerplate.

# Python 3.10+
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118
pip install pytorch-lightning lightly info-nce-pytorch
pip install scikit-learn anndata scanpy
# optional: sentence-transformers for text-side of CLIP-style objectives
pip install sentence-transformers

Hardware. NT-Xent/InfoNCE requires large effective batch sizes (256–2048) to supply enough in-batch negatives; a single A100 (80 GB) handles batch 512 with a 128-D encoder. For CPU-only prototyping, reduce batch to 64 and use a small synthetic dataset — expect slow convergence but correct gradients.

Data structures & formats

• Positive/negative pairs: for single-cell, a "view" is a stochastic augmentation of a count vector — typically random gene masking (set 10–30% of gene counts to 0) or Gaussian noise on log-normalised expression. Two augmented views of the same cell form the positive pair; all other cells in the batch are negatives.
• Batched tensors: the encoder receives a (N, G) float32 tensor (N = batch size, G = number of genes) and produces (N, D) L2-normalised embeddings. The loss operates on the (N, D) pair (z1, z2).
• AnnData / .h5ad: the standard container for single-cell data. After QC and log-normalisation, store the expression matrix in adata.X and load rows as training samples via a PyTorch Dataset.
• Augmentation pipeline: unlike images, single-cell augmentations should be biologically plausible — random dropout of genes (simulates technical dropout), small Gaussian noise, or subsetting to a random gene panel. Avoid augmentations that erase marker genes for rare populations.
• Memory bank (optional): for small batches, a FIFO queue of past embeddings (MoCo-style) extends the effective negative set without increasing batch size.

Minimal code walkthrough

Define an MLP encoder, build two augmented views of each cell in a batch, compute NT-Xent loss via info_nce, and run one gradient step — a self-contained 25-line training loop.

import torch
import torch.nn as nn
import torch.nn.functional as F
from info_nce import InfoNCE

# --- Encoder: gene counts (G,) -> unit-sphere embedding (D,) ---
class CellEncoder(nn.Module):
    def __init__(self, n_genes: int = 2000, d_emb: int = 128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_genes, 512), nn.BatchNorm1d(512), nn.ReLU(),
            nn.Linear(512, 256),    nn.BatchNorm1d(256), nn.ReLU(),
            nn.Linear(256, d_emb),
        )
    def forward(self, x):
        return F.normalize(self.net(x), dim=-1)  # L2-normalise onto unit sphere

# --- Augmentation: random gene dropout (simulates technical dropout) ---
def augment(x: torch.Tensor, drop_prob: float = 0.2) -> torch.Tensor:
    mask = torch.bernoulli(torch.full_like(x, 1.0 - drop_prob))
    return x * mask

# --- Training step ---
encoder = CellEncoder(n_genes=2000, d_emb=128)
loss_fn  = InfoNCE(temperature=0.1, reduction="mean", negative_mode="unpaired")
optimizer = torch.optim.Adam(encoder.parameters(), lr=1e-3)

# x: (N, G) float32 log-normalised expression for one batch
x = torch.randn(256, 2000).clamp(min=0)   # synthetic placeholder
view1 = augment(x)
view2 = augment(x)
z1 = encoder(view1)   # (N, D)
z2 = encoder(view2)   # (N, D)

loss = loss_fn(z1, z2)  # NT-Xent: pull (z1_i, z2_i) together, push apart from all j != i
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"InfoNCE loss: {loss.item():.4f}")   # ~log(N) at init, decreases as pairs align

Common pitfalls & tips

• Large batch / sufficient negatives is non-negotiable. With N=32, the loss has only 31 negatives per anchor — insufficient signal and the model learns trivially. Use N≥256, or add a MoCo-style memory bank for GPU-constrained settings.
• Augmentation design is domain-specific — cells ≠ images. Random horizontal flip or colour jitter corrupts gene identity. Use biologically motivated perturbations: gene dropout, Gaussian noise on log-counts, or random subsampling to lower sequencing depth.
• Temperature τ needs tuning. Start at 0.1 (Concerto) or 0.07 (MedMPT). Too small → representations collapse to a single point; too large → uniform distribution, no contrastive signal. Tune by tracking linear-probe accuracy on a held-out cell-type annotation task.
• Representation collapse. If training loss drops to near-zero in the first epoch, check that BatchNorm is active (it provides implicit negative correlation) and that the two views are truly different augmentations. Asymmetric architectures (teacher–student, as in Concerto) are more collapse-resistant.
• Evaluate embeddings, not just loss. InfoNCE loss correlates poorly with downstream utility. After each epoch, run a k-NN (k=15) or linear probe on a labelled subset: classify cell types and track accuracy. UMAP visualisation can reveal whether batch effects dominate the geometry before you waste GPU time.
• Augmentations must not erase marker genes. Test augmentations by checking that rare cell-type markers (e.g. FOXP3 for Tregs) survive dropout at the chosen rate. If a marker gene appears in <5% of cells it may be reliably masked out, destroying the signal for that population.