HASSL: Hierarchy-Aware Self-Supervised Learning Framework for Single Cell Microscopy

Figure 1: Hierarchy in cell imaging (left) and the embedding objective with HASSL (right). HASSL learns a more hierarchical embedding space, leading to improved morphological representation and tighter subclusters.

[ 📜 Paper] [ 📖 BibTeX]

Biological cell images present a distinct challenge for self-supervised learning: coarse imaging factors (acquisition modality, staining protocol) systematically dominate the learned representation, overwhelming the fine morphological signals that distinguish biologically distinct subtypes. The result is a latent space where semantically different cells appear identical and hierarchically related subtypes collapse into the same cluster.

We propose HASSL, a hierarchy-aware self-supervised training framework built on top of DINOv3 that directly counteracts this tendency via two tightly integrated components:

1. Double-Teacher Distillation — A segmentation teacher is incorporated alongside the standard EMA image teacher. Pre-computed zero-shot segmentation masks (CellposeSAM) provide structure-aware supervision that biases the student away from modality cues and toward cell morphology, initiating the emergence of morphology-based subclusters.

2. Hierarchy-Aware Contrastive Loss — At each training step, HDBSCAN is run on the current batch embeddings to obtain a condensed cluster tree. For each anchor, stability-weighted (λ) positive and negative prototypes are mined at every level of the hierarchy. The resulting hinge-contrastive loss pulls each cell toward its ancestor cluster centroids while repelling cells from different subtypes at each hierarchical level.

Together, these two components push the embedding space toward a structure that is simultaneously morphologically grounded and hierarchically consistent. We train and evaluate on a curated corpus of 2.3M single cells aggregated from 20 microscopy datasets covering 208 cell classes, achieving +2.8% average top-K accuracy, +6.3% top-9 retrieval on the deep-hierarchy benchmark, and +7.8% F1-score on drug perturbation classification over state-of-the-art SSL baselines.

Figure 2: Overview of the HASSL training pipeline. (1) Zero-shot segmentation maps are generated per cell. (2) A student ViT is trained via double-teacher distillation from a global image teacher and a segmentation teacher. (3) HDBSCAN derives a cluster hierarchy from batch embeddings; stability-λ-weighted prototypes define positive and negative pairs. (4) The resulting latent space organises cells into superclusters with clearer morphology-driven subclusters.

---

Dataset

The dataset can be downloaded here

---

Installation

The training and evaluation code requires PyTorch and a Weights & Biases account for experiment tracking. Clone the repository and create the conda environment:

conda env create -f conda.yaml
conda activate dinov3

---

Training

Step 1 — Configure Weights & Biases

Set the following environment variables before launching training. The trainer reads these at runtime — no credentials are stored in the codebase.

Variable	Description	Required
WANDB_ENTITY	Your W&B username or team name	Yes
WANDB_PROJECT	Project name for this run	No (defaults to dinov3-cell)
WANDB_TAGS	Comma-separated list of run tags	No

export WANDB_ENTITY="<your-wandb-entity>"
export WANDB_PROJECT="<your-project-name>"
export WANDB_TAGS="dinov3,hdbscan,finetuning"   # optional

Step 2 — Prepare the Training Manifest

Option A — Use our dataset. We provide pre-built manifests for our training and evaluation sets. Download the dataset and manifests:

File	Description	Download
Cell image dataset	All .npy cell images used for training and evaluation	link
manifest_train_fixed.csv.gz	Training manifest (included in repo)	—
manifest_test_fixed.csv.gz	Evaluation manifest (included in repo)	—

After downloading the dataset, update the img_path and mask_dir columns in the manifests to reflect the location of the dataset on your machine:

import gzip, csv, io

OLD_PREFIX = "path_to_dataset/"   # placeholder in the provided manifests
NEW_PREFIX = "/your/local/path/to/dataset/"

for fname in ["manifest_train_fixed.csv.gz", "manifest_test_fixed.csv.gz"]:
    with gzip.open(fname, "rt", newline="") as f:
        reader = csv.DictReader(f)
        fieldnames = reader.fieldnames
        rows = list(reader)
    for row in rows:
        for col in ("img_path", "mask_dir"):
            if col in row:
                row[col] = row[col].replace(OLD_PREFIX, NEW_PREFIX)
    buf = io.StringIO()
    writer = csv.DictWriter(buf, fieldnames=fieldnames)
    writer.writeheader()
    writer.writerows(rows)
    with gzip.open(fname, "wt", newline="") as f:
        f.write(buf.getvalue())

Option B — Bring your own dataset. If you are training on a custom cell image collection, generate the manifest using build_manifest_csv() provided in dinov3/data/datasets/n_cells.py:

from dinov3.data.datasets.n_cells import build_manifest_csv

build_manifest_csv(
    root="/path/to/your/dataset/root",
    split="train",
    out_csv_gz="manifest_train.csv.gz",
)

The dataset root must follow this directory structure:

<root>/
  <origin>/          # e.g. N_PanNuke, N_MoNuSeg, …
    <split>/         # train / val / test
      [<label>/]     # optional label sub-directories
        original/    # cell images as .npy files  (H × W × 3, uint8 or float32)
        mask/        # binary instance masks as .npy files

The manifest is a gzip-compressed CSV with the following schema:

Column	Description
img_path	Absolute path to the cell image (.npy)
origin	Source dataset name (e.g. N_PanNuke)
label	Cell type or class label
mask_dir	Directory containing the corresponding mask file
has_empty	1 if an empty-mask variant exists, else 0
stem	Filename stem (no extension)
h	Height of the image in pixels
w	Width of the image in pixels
area	Area of the image in pixels (h × w)

Step 3 — Register Cell Classes

Add your cell origin names from the manifest file to:

dinov3/data/datasets/n_cells.py  (line 27)

Step 4 — Launch Training

Run DINOv3 pre-training with the double-teacher distillation and hierarchy-aware contrastive loss on a single node:

PYTHONPATH=${PWD} python -m dinov3.run.submit dinov3/train/train.py \
  --nodes 1 \
  --config-file dinov3/configs/train/vitl_im1k_lin834.yaml \
  --output-dir <PATH/TO/OUTPUT/DIR> \
  train.dataset_path=NCells:root=/<PATH/TO/CSV.GZ>:split=TRAIN \
  finetune.path='' \
  triplet.enable=true:weight_scaling=global \
  checkpointing.checkpointing_goal_epoch=40

Leave finetune.path='' to train from scratch.

Fine-tuning from a Checkpoint

To resume from or fine-tune an existing checkpoint — for example, to apply the HDBSCAN contrastive loss on top of the Double Teacher checkpoint — set finetune.path to the checkpoint path:

PYTHONPATH=${PWD} python -m dinov3.run.submit dinov3/train/train.py \
  --nodes 1 \
  --config-file dinov3/configs/train/vitl_im1k_lin834.yaml \
  --output-dir <PATH/TO/OUTPUT/DIR> \
  train.dataset_path=NCells:root=/<PATH/TO/CSV.GZ>:split=TRAIN \
  finetune.path='<PATH/TO/CHECKPOINT>' \
  triplet.enable=true:weight_scaling=global \
  checkpointing.checkpointing_goal_epoch=40

---

License

This project is released under the DINOv3 License. See LICENSE.md for full terms.

Contributing

We welcome contributions. See CONTRIBUTING.md for guidelines.

Code of Conduct

This project follows the Contributor Covenant. By participating, you agree to uphold these standards.

---