STAMP: Predicting Out-of-Distribution Generalization without Target Data
Paper β’ 2609.32672 β’ Published
NIH ChestX-ray14: 40-Model Architecture & Pretraining Comparison Suite
This repository contains 40 models trained and evaluated on the NIH ChestX-ray14 dataset for 14-class multi-label chest radiograph classification. The collection was built to study how backbone architecture, attention mechanisms, and pretraining paradigm (supervised vs. self-supervised / vision-language) affect downstream performance and representation geometry on the same fixed medical imaging task.
This is a research artifact, not a clinical or diagnostic tool. See Intended Use & Limitations below.
The 40 models split into two groups with different training protocols β this distinction is load-bearing and is preserved throughout this card rather than averaged away.
- Group A (31 models): fully fine-tuned, ImageNet-initialized backbones.
- Group B (9 models): frozen self-supervised / vision-language backbones with only a linear probe head trained on top. These are ablation models, not directly comparable to Group A.
Model Description
Task: Multi-label binary classification over 14 thoracic pathology findings from a single frontal chest radiograph.
Input: 224Γ224 RGB image, ImageNet-normalized (mean
[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]).Output: 14 raw logits; apply
sigmoidfor per-class probability. Class order:Atelectasis, Cardiomegaly, Effusion, Infiltration, Mass, Nodule, Pneumonia, Pneumothorax, Consolidation, Edema, Emphysema, Fibrosis, Pleural_Thickening, HerniaModel count: 40 (31 Group A + 9 Group B). (Note: this repository does not include the 4 external TorchXRayVision baseline models that were used for comparison during evaluation β those are third-party pretrained models, not part of this release.)
Group A β architectures (31, fully fine-tuned)
ResNet-18/34/50/101/152 Β· DenseNet-121/169/201 Β· EfficientNet-B0/B4 Β· EfficientNetV2-S Β· ConvNeXt-T/S/B Β· ResNet50-CBAM Β· ResNet50-SE Β· ResNet50-BAM Β· ResNet101-CBAM Β· ViT-Tiny/16 Β· ViT-Small/16 Β· ViT-Base/16 Β· DeiT-Base/16 Β· Swin-T/S/B Β· CoAtNet-0 Β· MaxViT-T Β· MViTv2-T Β· ConvFormer-S18 Β· CaFormer-S18 Β· MLP-Mixer-B/16
Group B β architectures (9, frozen backbone + linear probe)
| Model | Pretrained source |
|---|---|
| DINOv2 ViT-B/14 | facebookresearch/dinov2 (torch.hub) |
| MAE ViT-B/16 | timm β vit_base_patch16_224.mae |
| DINO ViT-B/16 | timm β vit_base_patch16_224.dino |
| SAM ViT-B/16 | timm β vit_base_patch16_224.sam_in1k |
| CLIP ViT-B/16 | timm β vit_base_patch16_clip_224.openai |
| BiomedCLIP ViT-B/16 | microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224 (open_clip) |
| MoCov3 ViT-B/16 | custom MoCo v3 pretrained checkpoint |
| PubMedCLIP ViT-B/32 | flaviagiammarino/pubmed-clip-vit-base-patch32 |
| RAD-DINO | microsoft/rad-dino |
Training Data
- Dataset: NIH ChestX-ray14 (112,120 frontal chest radiographs, 14 pathology labels mined from radiology reports).
- Split: the dataset's official patient-level split was used:
train_val_list.txtfor training/validation,test_list.txtfor testing. From the train/val pool, patients were further split 85% train / 15% validation at the patient level (seed=42), so that no patient's images appear in more than one of train, validation, or test. This matters specifically for medical imaging: splitting by image rather than by patient can leak patient-specific visual cues across the split boundary and inflate reported performance. - Test set: 25,596 images from the official held-out test list, entirely untouched during training, validation, or model selection.
- Label characteristics: ChestX-ray14 labels are extracted via NLP from radiology report text rather than expert-adjudicated per image, and are known in the literature to contain a non-trivial noise rate. Reported AUCs should be read with this in mind.
Preprocessing
train: Resize(256) β RandomRotation(10) β RandomCrop(224) β
ColorJitter(brightness=0.2, contrast=0.2) β ToTensor() β
Normalize(IMAGENET_MEAN, IMAGENET_STD)
eval: Resize(256) β CenterCrop(224) β ToTensor() β
Normalize(IMAGENET_MEAN, IMAGENET_STD)
Training Procedure
Group A β full fine-tuning (31 models)
ImageNet-pretrained backbones, fine-tuned in two stages:
| Stage 1 (backbone frozen) | Stage 2 (full fine-tune) | |
|---|---|---|
| Epochs | 5 | 25 |
| Optimizer | AdamW | AdamW |
| Learning rate | 3e-4 (head only) | 5e-6 backbone / 5e-5 head (differential) |
| Weight decay | 1e-4 | 1e-4 |
| LR schedule | Cosine annealing (Ξ·_min=1e-7) | Cosine annealing (Ξ·_min=1e-7) |
Shared across both stages: mixed-precision training (AMP), gradient clipping
at global norm 1.0, batch size 32, early stopping on validation loss
(patience = 7 epochs). Loss function: label-smoothed (smoothing = 0.1)
binary cross-entropy with per-class pos_weight to address class imbalance
across the 14 findings.
Group B β frozen linear probe (9 models)
The pretrained backbone is loaded from its original public source and
entirely frozen (requires_grad=False). Only a Dropout(0.3) β Linear
head is trained on top of the frozen embedding, using the same data split
described above. These models isolate what a frozen self-supervised /
vision-language representation can achieve on this task without any
domain adaptation of the backbone, and are intended as ablations against
Group A rather than as competitive standalone classifiers.
Evaluation
Macro-averaged AUC across all 14 classes, computed on the untouched official NIH test split (25,596 images, 27,206 positive labels total).
Group A β fully fine-tuned
| Model | Macro AUC |
|---|---|
| ResNet-18 | 0.7777 |
| ResNet-34 | 0.7721 |
| ResNet-50 | 0.7802 |
| ResNet-101 | 0.7754 |
| ResNet-152 | 0.7805 |
| DenseNet-121 | 0.7885 |
| DenseNet-169 | 0.7904 |
| DenseNet-201 | 0.7917 |
| EfficientNet-B0 | 0.7754 |
| EfficientNet-B4 | 0.7347 |
| EfficientNetV2-S | 0.7973 |
| ConvNeXt-T | 0.7998 |
| ConvNeXt-S | 0.8118 |
| ConvNeXt-B | 0.8107 |
| ResNet50-CBAM | 0.7827 |
| ResNet50-SE | 0.7734 |
| ResNet50-BAM | 0.7771 |
| ResNet101-CBAM | 0.7661 |
| ViT-Tiny/16 | 0.7868 |
| ViT-Small/16 | 0.7922 |
| ViT-Base/16 | 0.7921 |
| DeiT-Base/16 | 0.7901 |
| Swin-T | 0.7979 |
| Swin-S | 0.8024 |
| Swin-B | 0.8030 |
| CoAtNet-0 | 0.7886 |
| MaxViT-T | 0.7904 |
| MViTv2-T | 0.7897 |
| ConvFormer-S18 | 0.7922 |
| CaFormer-S18 | 0.7936 |
| MLP-Mixer-B/16 | 0.7754 |
Best in Group A: ConvNeXt-S (0.8118).
Group B β frozen linear probe (ablation only)
| Model | Macro AUC |
|---|---|
| DINOv2-ViT-B/14 | 0.7217 |
| MAE-ViT-B/16 | 0.7033 |
| DINO-ViT-B/16 | 0.6968 |
| SAM-ViT-B/16 | 0.7325 |
| CLIP-ViT-B/16 | 0.6511 |
| BiomedCLIP-ViT-B/16 | 0.6905 |
| MoCov3-ViT-B/16 | 0.7365 |
| PubMedCLIP-ViT-B/32 | 0.6995 |
| RAD-DINO | 0.8008 |
Note the spread within Group B (0.6511β0.8008) is large and should not be read as "frozen representations don't work" β RAD-DINO, a chest-X-ray-specific self-supervised backbone, is competitive with Group A's mid-tier models even fully frozen, while general-purpose vision-language backbones (CLIP, PubMedCLIP, BiomedCLIP) trail noticeably, consistent with the domain shift between their pretraining data and chest radiographs.
How to Load a Model
Checkpoint formats differ between the two groups β see below.
Group A checkpoints store a full state_dict (under a model,
state_dict, or model_state_dict key) and load directly into the matching
architecture wrapper.
Group B checkpoints store only the linear probe head (head key).
The frozen backbone is not bundled in the checkpoint and must be reloaded
from its original public source at inference time, then the head weights
loaded on top.
See loading_utils.py in this repository for complete, runnable loading
code for every architecture in both groups.
Intended Use & Limitations
- Intended use: research into architecture/pretraining effects on medical image classification, transfer learning, self-supervised representation quality, and related methodology (e.g., memorization/privacy analysis, out-of-distribution generalization studies).
- Not intended for: clinical diagnosis, triage, or any use affecting patient care. These models are not FDA-cleared or approved for any clinical use, have not undergone clinical validation, and were trained on a single public dataset with known label noise and limited demographic documentation.
- Known limitations:
- ChestX-ray14 labels are NLP-mined from report text, not adjudicated by radiologists per image; label noise affects reported AUCs for all 40 models equally but should not be forgotten when interpreting them.
- No external validation on other institutions' data is included in this release; performance on other patient populations, scanners, or acquisition protocols is unknown and may differ substantially.
- Group B models reflect a frozen-backbone ablation setting and are not representative of these architectures' best achievable performance with full fine-tuning.
Ethical Considerations
Chest X-ray datasets, including ChestX-ray14, have documented demographic and label-quality limitations discussed in the broader literature; this release does not include a demographic subgroup performance breakdown. Users applying these models to any real-world screening or triage context should conduct their own validation, bias auditing, and regulatory review before doing so.
License
[Pending β specify your chosen license here, e.g. MIT, Apache-2.0, or a custom research-only license, before publishing.]
Citation
@article{mahbub2026stamp,
title={STAMP: Predicting Out-of-Distribution Generalization without Target Data},
author={Mahbub, Md Kawsher and Biswas, Milon},
journal={arXiv preprint arXiv:2609.32672},
year={2026}
}
(Update once your paper is public β placeholder above.)