CustoDistillPath

Official implementation of the MICCAI 2026 early-accepted paper:

Harnessing Adversarial Distillation to Customise Debiased, Disease-Specific Pathology Foundation Models for Breast Cancer

CustoDistillPath customises pathology foundation models for breast cancer analysis through a two-stage adversarial distillation pipeline. It first learns an adversarial model from metadata-supervised signals, then trains a compact disease-specific student model with foundation-model distillation, adversarial alignment, and optional self-supervised regularisation.

Repository Structure

.
|-- CustoDistillPath/
|   |-- configs/                 # Distillation training configs
|   |-- data/                    # Dataset and metadata utilities
|   |-- losses/                  # Distillation, adversarial, and SSL losses
|   |-- metadata/                # Example TCGA-BRCA metadata file
|   |-- models/                  # Adversarial and student model definitions
|   |-- save_coords_images/      # Utilities for patch extraction from coordinates
|   |-- utils/                   # Training helpers, metrics, logging
|   |-- requirements.txt
|   `-- train.py                 # Main two-stage training entry point
|-- evaluate_n_runs_bracs/
|   |-- config/                  # MIL evaluation configs
|   |-- mil_model.py             # Attention-based MIL model
|   |-- mil_train.py             # Multi-run BRACS evaluation
|   `-- utils.py
`-- README.md

Installation

Create a Python environment and install the project dependencies:

cd CustoDistillPath
pip install -r requirements.txt

The code is designed for PyTorch-based GPU training. For distributed runs, launch the training script with torchrun.

Data Preparation

Before training, prepare the following inputs:

1. Patch images extracted from whole-slide images.
2. Precomputed patch-level teacher features from the selected pathology foundation models.
3. Metadata used by the adversarial supervision tasks.
4. Downstream slide labels and patch features for MIL evaluation.

You can use TRIDENT or an equivalent WSI preprocessing pipeline to extract patches and foundation-model features.

Training

Run commands from the CustoDistillPath/ directory.

Full Pipeline

This runs phase 1 followed by phase 2:

python train.py --config configs/config_tinyvit_21m.yaml

Phase 1: Adversarial Model Pretraining

Phase 1 trains the adversarial model with supervised contrastive learning and cross-entropy losses over metadata-defined tasks.

python train.py --config configs/config_tinyvit_21m.yaml --phase 1

Phase 2: Student Distillation

Phase 2 trains the disease-specific student model using teacher feature distillation, adversarial loss, and optional DINO-style SSL loss.

python train.py \
  --config configs/config_tinyvit_21m.yaml \
  --phase 2 \
  --adversarial_checkpoint path/to/best_adversarial_model.pth

Distributed Training

For multi-GPU training:

torchrun --nproc_per_node=4 train.py --config configs/config_tinyvit_21m.yaml

Feature Extraction for Downstream Tasks

After training the student model, extract patch-level features for the downstream dataset, such as BRACS. The MIL evaluation expects one .h5 feature file per slide in the configured feature directory.

Downstream MIL Evaluation

Run Attention-based MIL evaluation from the evaluate_n_runs_bracs/ directory:

cd evaluate_n_runs_bracs
python mil_train.py --config config/tinyvit_21m/config_tinyvit_21m_loss_all.yaml

The script performs multiple runs with different random seeds and reports:

• Balanced Accuracy (BACC)
• AUC
• Weighted F1

Acknowledgements

This repository builds on open-source tools and libraries including timm, TRIDENT, and vit-pytorch.

Citation

If this repository or paper is useful for your research, please cite our work. The BibTeX entry will be updated after the official proceedings metadata is available.

coming soon