HSAT: Hierarchical Self-Supervised Adversarial Training for Robust Vision Models in Histopathology

Hashmat Shadab Malik, Shahina Kunhimon, Muzammal Naseer, Fahad Shahbaz Khan, and Salman Khan

Mohamed Bin Zayed University of Artificial Intelligence, Khalifa University, Linko ̈ping University, and Australian National University

Official PyTorch implementation

---

🔥 News

• (March 14, 2025)
  • HSAT paper is released on arxiv.
• (March 13, 2025)
  • We open source the code for model training and evaluation.

---

Abstract: Adversarial attacks pose significant challenges for vision models in critical fields like healthcare, where reliability is essential. Although adversarial training has been well studied in natural images, its application to biomedical and microscopy data remains limited. Existing self-supervised adversarial training methods overlook the hierarchical structure of histopathology images, where patient-slide-patch relationships provide valuable discriminative signals. To address this, we propose Hierarchical Self-Supervised Adversarial Training (HSAT), which exploits these properties to craft adversarial examples using multi-level contrastive learning and integrate it into adversarial training for enhanced robustness. We evaluate HSAT on multiclass histopathology cancer diagnosis dataset OpenSRH and the results show that HSAT outperforms existing methods from both biomedical and natural image domains. HSAT enhances robustness, achieving an average gain of 54.31% in the white-box setting and reducing performance drops to 3-4% in the black-box setting, compared to 25-30% for the baseline. These results set a new benchmark for adversarial training in this domain, paving the way for more robust models. Our code and pretrained models will be made publicly available.

---

We propose HSAT, a min-max optimization framework to learn robust representations via self-supervised contrastive learning. The inner maximization step, 𝓛_max, generates adversarial perturbations by crafting a hierarchy-wise attack. The outer minimization step, 𝓛_min, updates the encoder f_θ to minimize the hierarchical contrastive loss, encouraging robust feature representations across patch, slide, and patient levels.

---

Contents

1. Installation
2. Models and Data Preparation
3. Training
4. Robustness against White-Box Attacks
5. Robustness against Transfer-Based Black-Box Attacks
6. BibTeX
7. Contact
8. References

---

---

💿 Installation

conda create -n hsat
conda activate hsat
conda install pytorch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 pytorch-cuda=11.8 -c pytorch -c nvidia
pip install timm
pip install wandb
pip install tifffile
pip install dill
pip install hydra-core --upgrade
pip install matplotlib
pip install scikit-learn
pip install tqdm
pip install torchmetrics
pip install pandas
pip install opencv-python
pip install gdown
pip install ftfy regex tqdm
pip install transformers

🛠️ Models and Data Preparation

📦 Available Models

Model	Identifier
ResNet-50 (Scratch)	resnet50
ResNet-50 (ImageNet Pretrained)	resnet50_timm_pretrained
ResNet-50 (Adversarial ImageNet Pretrained)	resnet50_at
WResNet-50 (ImageNet Pretrained)	wresnet50_normal
WResNet-50 (Adversarial ImageNet Pretrained)	wresnet50_at
ResNet-101 (ImageNet Pretrained)	resnet101_normal
ResNet-101 (Adversarial ImageNet Pretrained)	resnet101_at

---

📊 OpenSRH Dataset

To access the OpenSRH dataset, please submit a data request here.

---

🚀 Training

📚 Baseline Non-Adversarial Training

To train different vision models on the OpenSRH dataset with the baseline non-adversarial training approach from HiDisc, run the following command:

torchrun --nproc_per_node=<NUM_GPUS> main.py \
    data.db_root=<DATA_PATH> \
    data.dynamic_aug=False \
    data.dynamic_aug_version=v0 \
    model.backbone=<MODEL_NAME> \
    training.batch_size=<BATCH_SIZE> \
    training.only_adv=False \
    training.attack.name=none \
    training.attack.eps=8 \
    training.attack.warmup_epochs=0 \
    training.attack.loss_type=p_s_pt \
    out_dir=Results/Baseline/<MODEL_NAME>_exp1 \
    wandb.exp_name=Baseline_backbone_<MODEL_NAME>_exp1 \
    wandb.use=True

• <NUM_GPUS>: Number of GPUs used for distributed training.
• <DATA_PATH>: Path to the OpenSRH dataset.
• <MODEL_NAME>: Name of the vision model backbone (for example resnet50, resnet50_at)
• <BATCH_SIZE>: Number of patients per batch during training.
• data.dynamic_aug=False: Disables dynamic augmentation, keeping the dataset unchanged for baseline training.
• training.only_adv=False: if True ensures that only adversarial samples are used during adversarial training.
• training.attack.name=none: Disables adversarial attacks for baseline training.
• out_dir=Results/Baseline/<MODEL_NAME>_exp1: Path to save the training results (exp1 for baseline training).
• wandb.exp_name=Baseline_backbone_<MODEL_NAME>_exp1: Sets the experiment name for tracking in Weights & Biases ( W&B).
• wandb.use=True: Enables logging with Weights & Biases for visualizing training metrics.

📚 Hierarchical Self-Supervised Adversarial Training (HSAT)

To train different vision models on the OpenSRH dataset with the proposed Hierarchical Self-Supervised Adversarial Training (HSAT) approach, run the following command:

torchrun --nproc_per_node=<NUM_GPUS> main.py \
    data.db_root=<DATA_PATH> \
    data.dynamic_aug=True \
    data.dynamic_aug_version=v0 \
    model.backbone=<MODEL_NAME> \
    training.batch_size=<BATCH_SIZE> \
    training.only_adv=True \
    training.attack.name=pgd \
    training.attack.eps=8 \
    training.attack.warmup_epochs=5000 \
    training.attack.loss_type=p_s_pt \
    out_dir=Results/Adv/<MODEL_NAME>_dynamicaug_true_epsilon_warmup_5000_only_adv_exp2 \
    wandb.exp_name=Adv_backbone_<MODEL_NAME>_dynamicaug_true_epsilon_warmup_5000_only_adv_exp2 \
    wandb.use=True

• training.attack.name=pgd: Adversarial attack used for HSAT training.

• training.attack.eps=8: Sets the maximum perturbation allowed for adversarial examples.

• training.attack.warmup_epochs=0: Number of warm-up epochs to linearly increase the adversarial attack strength.

• training.attack.loss_type=p_s_pt: Specifies the maximization loss type for crafting adversarial examples during training.

  • p_s_pt is a combination of the patch-level, slide-level, and patient-level contrastive losses.
  • pt is the patch-level contrastive loss.
  • s_pt is the combination of the slide-level and patch-level contrastive losses.

• training.only_adv: Set to False if you want to include clean samples as well during adversarial training.

(exp2 for HSAT training)

To train different vision models on the OpenSRH dataset with the proposed Hierarchical Self-Supervised Adversarial Training (HSAT) approach but using different levels of heirarchical contrastive losses during the training, run the following command:

# HAT-Patch Testing the effect of patch loss only (L_max = L_min) on the adversarial training. Increase the batch size accordingly to keep the same effective batch size.
torchrun --nproc_per_node=<NUM_GPUS> main.py \
    data.db_root=<DATA_PATH> \
    data.dynamic_aug=True \
    data.dynamic_aug_version=v0 \
    data.hidisc.num_slide_samples=1 \
    data.hidisc.num_patch_samples=1 \
    model.backbone=<MODEL_NAME> \
    training.batch_size=<BATCH_SIZE> \
    training.only_adv=True \
    training.attack.name=pgd \
    training.attack.eps=8 \
    training.attack.warmup_epochs=5000 \
    training.attack.loss_type=p_s_pt \
    out_dir=Results/Adv/<MODEL_NAME>_dynamicaug_true_epsilon_warmup_5000_only_adv_hat_patch_exp13  \
    wandb.exp_name=Adv_backbone_<MODEL_NAME>_dynamicaug_true_epsilon_warmup_5000_only_adv_hat_patch_exp13 \
    wandb.use=True

(exp13 for HSAT-Patch training)

#  HAT-Slide Testing the effect of patch-slide loss only (L_max = L_min) on the adversarial training. Increase the batch size accordingly to keep the same effective batch size.
torchrun --nproc_per_node=<NUM_GPUS> main.py \
    data.db_root=<DATA_PATH> \
    data.dynamic_aug=True \
    data.dynamic_aug_version=v0 \
    data.hidisc.num_patch_samples=1 \
    model.backbone=<MODEL_NAME> \
    training.batch_size=<BATCH_SIZE> \
    training.only_adv=True \
    training.attack.name=pgd \
    training.attack.eps=8 \
    training.attack.warmup_epochs=5000 \
    training.attack.loss_type=p_s_pt \
    out_dir=Results/Adv/<MODEL_NAME>_dynamicaug_true_epsilon_warmup_5000_only_adv_hat_slide_exp14  \
    wandb.exp_name=Adv_backbone_<MODEL_NAME>_dynamicaug_true_epsilon_warmup_5000_only_adv_hat_slide_exp14 \
    wandb.use=True

(exp14 for HSAT-Slide training)

We provide bash scripts to train different vision models on the OpenSRH dataset with the baseline and proposed Hierarchical Self-Supervised Adversarial Training (HSAT) approach.

bash scripts/training.sh <NUM_GPUS> <BATCH_SIZE> <EXP_NUM> <MODEL_NAME> <DATA_PATH>

# For example, training ResNet-50(scratch) on OpenSRH dataset with Baseline training using 4 GPU, batch size 16
bash scripts/training.sh 4 16 1 resnet50 path/to/dataset

# For example, training ResNet-50(scratch) on OpenSRH dataset with HSAT training using 4 GPU, batch size 16
bash scripts/training.sh 4 16 2 resnet50 path/to/dataset

Please check the scripts/training.sh file for more details. The logs and trained models will be saved in the Results folder.

🛡️ Robustness against White-Box Attacks

For evaluating the robustness of vision models against white-box adversarial attacks, we provide the following script:

python adv_eval_knn.py --data_db_root <DATA_PATH>  --model_backbone <MODEL_NAME>  --eval_ckpt_path <CKPT_PATH> --save_results_path <CKPT_DIR>/eval_knn_results /
--eps <EPSILON> --steps <STEPS> --eval_predict_batch_size <EVAL_BATCH_SIZE> --attack_name <ATTACK_NAME> 

• --eps: Perturbation budget for adversarial attacks
• --steps: Number of attack steps for iterative attacks
• --eval_predict_batch_size: Batch size for evaluation
• --attack_name: Name of the attack to be used
  • pgd_r for PGD attack with random initialization
  • bim_r for BIM attack with random initialization
  • mifgsn_r for MIFGSM attack with random initialization
• --eval_ckpt_path: Path to the model checkpoint
• --save_results_path: Path to save the evaluation results

Results are saved in the checkpoint directory under the eval_knn_results folder.

We provide bash scripts to evaluate different vision models trained on the OpenSRH dataset with the baseline and proposed Hierarchical Self-Supervised Adversarial Training (HSAT) approach.

bash scripts/eval_knn_whitebox_attack.sh <EXP_NUM> <MODEL_NAME> <EPSILON> <STEPs> <BATCH_SIZE>

# After training using the bash scripts provided above, for example, training ResNet-50(scratch) on OpenSRH dataset with Baseline training using 4 GPU, batch size 16
bash scripts/training.sh 4 16 1 resnet50 path/to/dataset
# you can evaluate the models using the following commands:
# To evaluate ResNet-50(scratch) trained on OpenSRH dataset with Baseline method against white-box attacks (PGD, BIM, MIFGSM) at epsilon=8/255, steps=10, batch size=16
bash scripts/eval_knn_whitebox_attack.sh 1 resnet50 8 10 16

# Similarly, training ResNet-50(scratch) on OpenSRH dataset with HSAT training using 4 GPU, batch size 16
bash scripts/training.sh 4 16 2 resnet50 path/to/dataset

# you can evaluate the models using the following commands:
# To evaluate ResNet-50(scratch) trained on OpenSRH dataset with HSAT method against white-box attacks (PGD, BIM, MIFGSM) at epsilon=8/255, steps=10, batch size=16
bash scripts/eval_knn_whitebox_attack.sh 2 resnet50 8 10 16

For more details, please check the scripts/eval_knn_whitebox_attack.sh and scripts/adv_knn.sh files.

🛡️ Robustness against Transfer-Based Black-Box Attacks

After generating adversarial examples using a surrogate model, the transferability of adversarial examples can be reported by evaluating them on unseen target models trained on the same dataset. To evaluate any target model on the adversarial examples, run the following script:

 python adv_eval_knn_transf.py --data_db_root <DATA_PATH>  --source_model_backbone <SOURCE_MODEL_NAME> --source_exp_no <SOURCE_EXP_NUMBER>   \
--source_ckpt_path <SOURCE_CKPT_PATH> --target_model_backbone <TARGET_MODEL_NAME> --target_exp_no <TARGET_EXP_NUMBER> \
--target_ckpt_path <TARGET_CKPT_PATH>  --save_results_path  <PATH_TO_SAVE_RESULTS> --eps <EPSILON> --steps <STEPS>  --load_source_from_ssl True

• <DATA_PATH>: Path to the dataset.
• <SOURCE_MODEL_NAME>: Name of the source model backbone.
• <SOURCE_EXP_NUMBER>: Experiment number of the source model
• <SOURCE_CKPT_PATH>: Path to the source model checkpoint.
• <TARGET_MODEL_NAME>: Name of the target model backbone.
• <TARGET_EXP_NUMBER>: Experiment number of the target model.
• <TARGET_CKPT_PATH>: Path to the target model checkpoint.
• <PATH_TO_SAVE_RESULTS>: Path to save the evaluation results.
• <EPSILON>: Perturbation budget for adversarial attacks.
• <STEPS>: Number of attack steps for iterative attacks.

Furthermore, bash scripts are provided to evaluate transferability of adversarial examples generated using a surrogate model on target models trained using the baseline and proposed Hierarchical Self-Supervised Adversarial Training (HSAT).

bash scripts/eval_knn_blackbox_attack.sh <SOURCE_MODEL_EXP_NUM>

Please check the scripts/eval_knn_blackbox_attack.sh and scripts/adv_knn_transf.sh files for more details. The scripts can be easily modified to evaluate different vision models on the OpenSRH dataset with the baseline and proposed Hierarchical Self-Supervised Adversarial Training (HSAT) approach.

Quantitative Evaluation 📊

We observe that vision models trained using our HSAT framework exhibit significantly higher adversarial robustness in white-box settings compared to baseline methods. Under the PGD attack at ϵ = 8/255 using a ResNet-50 backbone, HSAT achieves gains of 43.90%, 60.70%, and 58.33% in patch, slide, and patient classification, respectively, compared to non-adversarial hierarchical training methods. Additionally, HSAT outperforms instance-level adversarial training (HSAT-Patch) by 6.68%, 10.51%, and 10% across the same tasks. Despite a slight drop in clean accuracy compared to non-adversarial methods, HSAT maintains a superior clean performance with improvements of 15.68%, 17.12%, and 20% over HSAT-Patch.

We observe that target vision models trained using our HSAT framework are significantly more robust against transferability of adversarial examples crafted across different surrogate models. On average, HSAT trained target models show a performance drop of around 3-4%, while target models trained using Hidisc show a performance drop of around 25-30%.

Ablation studies for Hierarchical Adversarial Training (HSAT) in Tables 3 and 4 show that increasing hierarchical discrimination—progressing from patch-level (HSAT-Patch) to slide-level (HSAT-Slide) and patient-level (HSAT-Patient)—consistently improves adversarial robustness. These results highlight the effectiveness of multi-level adversarial training in aligning robust representations across hierarchical levels.

---

📚 BibTeX

@article{malik2025hierarchical,
  title={Hierarchical Self-Supervised Adversarial Training for Robust Vision Models in Histopathology},
  author={Malik, Hashmat Shadab and Kunhimon, Shahina and Naseer, Muzammal and Khan, Fahad Shahbaz and Khan, Salman},
  journal={arXiv preprint arXiv:2503.10629},
  year={2025}
}

---

📧 Contact

Should you have any question, please create an issue on this repository or contact at hashmat.malik@mbzuai.ac.ae

---

📚 References

Our code is based on hidisc. We thank them for open-sourcing their codebase.