Skip to content

Repository files navigation

DistillPath

Paper Model weights Python 3.10+ CI

Official code for the paper DistillPath: An Efficient 22M Distilled Pathology Encoder Approaching Large Foundation Model Performance, accepted at the ECCV 2026 Workshop on Medical Foundation Models and Benchmarks (MedFM-Bench).

Pretrained DistillPath weights are available in the Hugging Face collection.

DistillPath turns a large released pathology foundation model into a compact 22M kaiko ViT-S/16 tile encoder using knowledge distillation. The recipe reads only the teacher's final class and patch tokens, so it does not need the teacher's DINO or iBOT pretraining heads. It trains on 6,000 public TCGA slides at low cost and applies to any released encoder that exposes backbone tokens.

Overview

DistillPath method overview

Given the same input tile, a frozen teacher and the trainable ViT-S/16 student both produce a token sequence. The student is supervised with three terms:

  1. A pointwise cosine loss on the class token.
  2. A relational RKD loss on the class token (pairwise distances and triplet angles across the batch).
  3. A pointwise cosine loss on the patch tokens, after the teacher grid is resized to the student grid.

A DINO-style projector maps the student features to the teacher dimension for the cosine terms. The teacher and projector are discarded after training, so inference uses the student backbone alone. See the paper for details.

We distill four teachers into the same student. The best teacher is task dependent: teacher size and teacher score do not predict student score.

Student (22M ViT-S/16) EVA (7-task) HEST PLISM
kaiko baseline (no distillation) 0.764 0.349 0.307
DistillPath-KS16-Virchow2 0.795 0.371 0.447
DistillPath-KS16-UNI2h 0.772 0.375 0.484
DistillPath-KS16-HOpt0 0.769 0.376 0.480
DistillPath-KS16-H0mini 0.771 0.387 0.495

Installation

git clone https://github.com/RamonKaspar/DistillPath.git
cd DistillPath
python -m venv .venv && source .venv/bin/activate
python -m pip install -e .

This installs DistillPath and all runtime dependencies.

Data setup

Training reads TCGA H&E whole-slide images directly, sampling tiles online during training rather than pre-extracting them. The exact 6,000-slide training set used in the paper is specified in data/tcga and totals 6.68 TB.

python scripts/download_tcga.py --output-dir /path/to/tcga/slides

See the TCGA data documentation for GDC client installation, resuming interrupted transfers, integrity checks, and downloading a single slide. The downloader creates the directory layout expected by the training code.

Online tile sampling and preprocessing are implemented with wsistream using its TiffSlide backend. See the wsistream documentation for its components and usage.

Copy .env.example to .env and set the paths:

cp .env.example .env
DATA_DIR=/path/to/tcga/slides       # training slides
VAL_DATA_DIR=/path/to/val/slides    # optional; leave empty to disable validation
OUTPUT_DIR=outputs
WEIGHTS_DIR=/path/to/weights        # where teacher and student weights live

Teacher and student weights

The student is initialized from public pretrained weights, and each teacher is frozen. Weights are looked up as $WEIGHTS_DIR/<model_name>.pth, where <model_name> is the registry name below. Each .pth is a plain timm backbone state dict (backbone.load_state_dict(...)).

Role <model_name> Backbone
Student kaiko_vits16 kaiko ViT-S/16
Student (generality) vits16_in21k ImageNet-21K ViT-S/16
Teacher virchow2 Virchow2 ViT-H/14
Teacher uni2_h UNI2-h ViT-H/14
Teacher h_optimus_0 H-optimus-0 ViT-g/14
Teacher h0_mini H0-mini ViT-B/14

Download each model from its official release (see the references in the paper) and save its backbone state dict to $WEIGHTS_DIR/<model_name>.pth.

Training

Training is driven by Hydra. Each experiment configuration selects a teacher, a student, and the calibrated loss coefficients. Run one of the eight paper experiments with:

# kaiko-initialized student (main results)
python -m distillpath.train experiment=ks16_virchow2
python -m distillpath.train experiment=ks16_uni2h
python -m distillpath.train experiment=ks16_hoptimus0
python -m distillpath.train experiment=ks16_h0mini

# ImageNet-21K-initialized student (generality experiment)
python -m distillpath.train experiment=is16_virchow2
python -m distillpath.train experiment=is16_uni2h
python -m distillpath.train experiment=is16_hoptimus0
python -m distillpath.train experiment=is16_h0mini

Any field can be overridden on the command line, for example:

python -m distillpath.train experiment=ks16_virchow2 max_steps=50000 batch_size=256

Each run trains for 50,000 steps at batch size 256 in bfloat16 (about 24 to 29 GPU-hours on one NVIDIA RTX 4090). Training checkpoints are written under OUTPUT_DIR. Exported student backbones are written under WEIGHTS_DIR when it is set, and under OUTPUT_DIR otherwise. slurm/submit.sh is an example cluster launcher.

Using a trained encoder

Pretrained checkpoints can be loaded directly from the Hugging Face collection with timm:

import timm

model = timm.create_model(
    "hf_hub:RamonK/DistillPath-KS16-HOpt0",
    pretrained=True,
    num_classes=0,
)
model.eval()

Alternatively, an exported .pth checkpoint is a standard timm ViT-S/16 state dict, so it loads without this package:

import timm
import torch

model = timm.create_model("vit_small_patch16_224", pretrained=False, num_classes=0)
model.load_state_dict(torch.load("distillpath_ks16_virchow2.pth", map_location="cpu"))
model.eval()

# Normalize tiles with mean=std=0.5 (the kaiko convention), then:
# tokens = model.forward_features(x)  # [B, 1 + num_patches, 384]
# cls = tokens[:, 0]                  # 384-dim tile embedding

Evaluation

Downstream results in the paper use the standard external protocols: EVA for tile-level tissue tasks, HEST for gene-expression prediction, and the PLISM robustness benchmark. Export a student with the training script above, then evaluate it with those frameworks.

Citation

@inproceedings{kaspar2026distillpath,
  title     = {DistillPath: An Efficient 22M Distilled Pathology Encoder Approaching Large Foundation Model Performance},
  author    = {Kaspar, Ramon and Ignatov, Andrey and Boeva, Valentina},
  booktitle = {Medical Foundation Models and Benchmarks (MedFM-Bench), ECCV 2026},
  year      = {2026}
}

About

Official implementation of DistillPath: backbone-token distillation from released pathology foundation models into efficient 22M ViT-S/16 encoders.

Topics

Resources

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages