Physics-Aligned Self-Supervised Learning for Scientific Imaging

Official code for the GCPR 2026 paper Physics-Aligned Self-Supervised Learning for Scientific Imaging.

Data augmentations define the invariances learned by self-supervised learning (SSL), but standard augmentation pipelines were designed for natural images. This repository provides a principled, reproducible procedure for augmentation design in scientific SSL: physics-aligned augmentations are formalised as a union of measurement-consistent symmetries and acquisition-driven perturbations (T_phys = T_sym ∪ T_acq), instantiated for real-space electron microscopy and reciprocal-space 4D-STEM diffraction, and evaluated across five SSL paradigms.

• SSL methods: DINOv2, I-JEPA, MAE, SimCLR, VICRegL (shared ViT-B backbone, single-channel input)
• Pretraining data: CEM500K (real-space cellular EM) and simulated LiNiO2 4D-STEM diffraction patterns
• Downstream tasks: NFFA multi-class classification and 4D-STEM crystal-orientation (quaternion) regression
• Analyses: full finetuning, linear probing, low-label finetuning, robustness to acquisition variability, representation-geometry diagnostics, single-factor augmentation ablation

Each method is pretrained under two augmentation regimes that differ only in the augmentation pipeline:

Regime	Name in code	Description
T_orig	original	Natural-image augmentations (random crop, horizontal flip, blur, photometric perturbations)
T_phys	domain	Physics-aligned augmentations (measurement-consistent symmetries + acquisition perturbations: physically motivated noise, intensity variation, reciprocal-space scaling, diffraction tilt, EM artifacts)

Installation

git clone https://github.com/DL4EM/physics-aligned-ssl.git
cd physics-aligned-ssl

python -m venv .venv
source .venv/bin/activate
pip install --upgrade pip
pip install -e .

Optional extras:

pip install -e ".[hub]"     # Hugging Face Hub integration (pretrained weights)
pip install -e ".[wandb]"   # Weights & Biases logging

Pretrained weights

All 20 pretrained encoders (5 methods x 2 modalities x 2 augmentation regimes) are available on the Hugging Face Hub.

Load an encoder directly:

from em_ssl.hub import load_encoder

# <pretraining dataset>/<method>_<domain|original>
encoder = load_encoder("cem500k/dinov2_domain")

import torch
images = torch.randn(4, 1, 128, 128)   # grayscale images in [0, 1]
output = encoder(images)

Or download checkpoints into the training-output layout expected by the evaluation scripts:

# everything
python scripts/download_weights.py

# a subset
python scripts/download_weights.py --models 'cem500k/*' 4dstem/mae_domain

Checkpoints are self-describing and can also be used without this codebase:

import torch
from huggingface_hub import hf_hub_download

path = hf_hub_download("DL4EM/physics-aligned-ssl", "cem500k/dinov2_domain/encoder.pt")
ckpt = torch.load(path, map_location="cpu", weights_only=True)
state_dict = ckpt["encoder_state_dict"]

Data preparation

Three public datasets are used. The scripts expect the following layout (paths are relative to the repository root; adjust data.data_path in the configs or --data-root on the CLI to point elsewhere):

data/
├── cem500k/                  # CEM500K pretraining subset
│   ├── train/                # 10,000 .tiff images
│   └── val/                  # 2,000 .tiff images
├── 4dstem/
│   ├── pretrain/             # 4D-STEM pretraining subset
│   │   ├── train/            # 10,000 .png diffraction patterns
│   │   └── val/              # 2,000 .png diffraction patterns
│   └── downstream/           # orientation-regression splits
│       ├── train/            # 232,531 patterns
│       ├── val/              # 29,066 patterns
│       └── test/             # 29,067 patterns
└── nffa_eu_100/              # NFFA-EUROPE SEM dataset (10 class folders)
    ├── Biological/
    ├── Fibres/
    └── ...

Sources

• CEM500K — Conrad & Narayan, eLife 2021; download via EMPIAR-10592.
• NFFA-EUROPE 100% SEM dataset — Aversa et al.. Place the ten class folders under data/nffa_eu_100/. The 75/10/15 train/val/test split is generated deterministically at run time (seed 42).
• 4D-STEM LiNiO2 diffraction patterns — Scheunert et al. (simulated with the Bloch-wave algorithm in py4DSTEM). Orientation labels are parsed from the filenames.

Reproducing the exact subsets. manifests/ contains gzipped file lists for every split used in the paper (see manifests/README.md). To materialise a subset from a full dataset download:

mkdir -p data/cem500k/train
zcat manifests/cem500k_pretrain_train.txt.gz | \
  xargs -I{} cp /path/to/cem500k_full/{} data/cem500k/train/

Pretraining

Pretrain all methods under both augmentation regimes:

bash scripts/run_pretrain_cem500k.sh   # real-space EM
bash scripts/run_pretrain_4dstem.sh    # 4D-STEM diffraction

Or a single configuration:

python scripts/train_benchmark.py --config configs/benchmark/cem500k/dinov2_domain.yaml

Configs follow the naming configs/benchmark/<pretrain_dataset>/<method>_<domain|original>.yaml; the two regimes differ only in the data.augmentation / data.domain_augmentations sections. Checkpoints and logs are written to outputs/benchmark/<pretrain_dataset>/<method>_<regime>/.

Downstream evaluation

All downstream runs use three seeds (42, 43, 44) and report mean ± std, matching the paper.

NFFA classification (CEM500K-pretrained encoders)

# full-label finetuning (Table 2)
bash scripts/run_downstream_nffa_finetune.sh

# frozen-encoder linear probe (Table 3, left)
bash scripts/run_downstream_nffa_linear_probe.sh

# finetuning with 25% of the training labels (Table 3, right)
bash scripts/run_downstream_nffa_lowlabel.sh

4D-STEM orientation regression (4D-STEM-pretrained encoders)

# quaternion-regression finetuning (Table 4)
bash scripts/run_downstream_4dstem_quat_finetune.sh

# robustness to detector gain and blur at test time (Section 4.4)
bash scripts/run_robustness_4dstem.sh

Single-model example:

python scripts/eval_downstream.py \
  --checkpoint outputs/benchmark/cem500k/dinov2_domain/encoder_final.pt \
  --config configs/benchmark/cem500k/dinov2_domain.yaml \
  --dataset nffa \
  --mode finetune \
  --seeds 42,43,44

Representation-geometry diagnostics (Section 4.3)

Computes effective rank, uniformity, collapse ratio, and kNN accuracy from frozen features — the label-free validation step of the augmentation-design procedure:

python scripts/eval_representations.py \
  --config configs/benchmark/cem500k/dinov2_domain.yaml \
  --checkpoint outputs/benchmark/cem500k/dinov2_domain/encoder_final.pt

# batch mode over all checkpoints of one modality
python scripts/eval_representations.py --eval-all --dataset cem500k
python scripts/eval_representations.py --eval-all --dataset stem4d

Single-factor augmentation ablation (Section 4.5)

Pretrains DINOv2 on CEM500K once per ablation configuration (adding or removing one transform at a time) and finetunes each on NFFA:

bash scripts/run_ablation_cem500k.sh

The corresponding 4D-STEM ablation configs are under configs/benchmark/4dstem_ablation/.

Results

NFFA top-1 accuracy (%) after CEM500K pretraining, finetuned, mean ± std over 3 seeds (scratch baseline: 66.59 ± 2.76):

Method	T_orig	T_phys	Δ
DINOv2	66.70 ± 0.83	76.67 ± 0.39	+9.97
I-JEPA	72.27 ± 1.15	73.41 ± 0.75	+1.14
MAE	67.29 ± 12.72	79.10 ± 3.07	+11.81
SimCLR	89.68 ± 0.37	89.39 ± 0.16	−0.29
VICRegL	69.83 ± 2.68	73.42 ± 3.00	+3.59

4D-STEM quaternion regression, mean geodesic error in degrees (scratch baseline: 10.01 ± 1.89):

Method	T_orig	T_phys	Δ
DINOv2	9.85 ± 0.96	5.60 ± 0.58	+4.25
I-JEPA	11.49 ± 0.51	10.72 ± 0.51	+0.77
MAE	1.86 ± 0.01	1.75 ± 0.02	+0.11
SimCLR	11.06 ± 0.32	9.35 ± 0.93	+1.71
VICRegL	9.53 ± 1.97	7.05 ± 1.09	+2.48

Repository structure

em_ssl/                 # core package
├── config/             # experiment config schema
├── data/               # datasets, transforms (T_orig / T_phys pipelines), loaders
│   └── transforms/     # spatial, intensity, and domain (EM/4D-STEM) augmentations
├── eval/               # downstream trainers, linear probe, kNN, representation metrics
├── models/             # ViT/ResNet encoders, heads, input adapters
├── ssl/methods/        # DINOv2, I-JEPA, MAE, SimCLR, VICRegL
├── tasks/              # classification / regression task wrappers
├── train/              # SSL trainer, optimizers, schedulers
└── hub.py              # Hugging Face Hub loading utilities

configs/benchmark/      # pretraining configs (per dataset, method, regime) + ablations
manifests/              # exact file lists for all dataset splits
scripts/                # training, evaluation, and Hub upload/download entry points

Uploading weights to the Hub (for maintainers)

pip install -e ".[hub]"
hf auth login
python scripts/upload_to_hub.py --repo-id <user>/physics-aligned-ssl --dry-run   # stage only
python scripts/upload_to_hub.py --repo-id <user>/physics-aligned-ssl            # stage + upload

Citation

@inproceedings{kazimi2026physicsaligned,
  title     = {Physics-Aligned Self-Supervised Learning for Scientific Imaging},
  author    = {Kazimi, Bashir and Sandfeld, Stefan},
  booktitle = {DAGM German Conference on Pattern Recognition (GCPR)},
  year      = {2026}
}

License

MIT — see LICENSE.

About

Developed by the Deep Learning for Electron Microscopy (DL4EM) group at the Institute for Materials Data Science and Informatics (IAS-9), Forschungszentrum Jülich.