Generative models improve fairness of medical classifiers under distribution shifts

In plain words

Medical image models often work well on the hospital they were trained in and worse on other hospitals, and they can also be less accurate for smaller patient groups. The usual fix is more labelled data, which is expensive and hard to get for rare conditions.

The paper's idea is to train an image generator on the images you already have, including unlabelled ones, and then use it to create extra training images. The generator is told which condition (and optionally which patient group or hospital) each image should show, so you can ask it for more of the groups you are short of. The synthetic images are mixed with the real labelled ones to train the classifier. The authors test this on three kinds of medical images and report better accuracy and fairness on data from new sites, but the generator is itself a large model and the gains depend on its quality.

TL;DR

• A conditional diffusion model is trained on labelled and unlabelled images; unlabelled images use an all-zero label vector app. E.
• Synthetic images are sampled with the training label distribution kept and sensitive attributes drawn uniformly §3.1 app. C.
• A classifier is trained on a weighted mix of real and synthetic images, with the mix ratio treated as a hyperparameter §3.1 app. C.
• Author-reported: out-of-distribution gains in histopathology, chest radiology and dermatology, plus smaller fairness gaps §3.2 §3.2.3.
• No official code was found.

Problem

Models trained on one population or hospital can fail on another, and performance can be uneven across subgroups §1. Prevalence differences, too little labelled data for some conditions, and evaluation sets that do not reflect real shifts all make this worse §1. The authors argue that robustness to shift and statistical fairness are rarely tackled together §1. They also note that generated data is not a substitute for collecting real data, and that synthetic images must be tagged and evaluated on real data §4.

Method

The compiled method is in the Supplementary Material, section C ("Method"), and in the overview of §3.1. The main-text methods file is commented out of the build (% \input{methods} in the source) and is not used here none; source: main.tex. The method has three steps §3.1.

1. Train a conditional generator. Train a DDPM-style diffusion model on labelled and unlabelled images. The condition is either the diagnostic label alone, or the label plus a property such as hospital id or a sensitive attribute §3.1. The network is a UNet with residual blocks and skip connections, following the architecture of app. E. Where a high-resolution output is needed, a second diffusion model upsamples the low-resolution output, and it uses the same conditioning vector §3.1 app. E. Histopathology is generated directly at 96×96; chest X-ray and dermatology use a low-resolution model plus an upsampler §3.2.1 §3.2.2 §3.2.3.
2. Sample synthetic data. Draw labels so that the diagnostic label distribution of the training set is preserved, and draw the sensitive attribute uniformly over its values. Repeated sampling gives several augmentations per condition §3.1 app. C.
3. Train the classifier on real plus synthetic images. Use a mixing ratio between real and synthetic batches, treated as a hyperparameter that depends on the application and modality §3.1. The objective weights the real-data loss by α and the synthetic-data loss by (1 − α) app. C.

Label-free data: the generator is trained with the same label slot, but unlabelled images get an all-zero conditioning vector. This also enables classifier-free guidance with label dropping app. E, in the sense of arXiv:2207.12598.

# Learned augmentation with a conditional diffusion model  [§3.1] [app. C]
# Input:  labelled set (image, label, attrs); unlabelled images
# Output: classifier parameters theta

# Step 1: train the generator on both sources  [§3.1] [app. E]
def train_generator(labelled, unlabelled, condition):
    # Unlabelled images get an all-zero conditioning vector  [app. E]
    for batch in mix(labelled, unlabelled):
        cond = condition_vector(batch, condition)   # zeros if unlabelled
        loss = diffusion_loss(unet, batch.images, cond)   # DDPM  [app. E]
        update(unet, loss)
    return unet

# Step 2: sample synthetic data with a fair strategy  [§3.1] [app. C]
def sample_synthetic(unet, label_prior, attr_values, count):
    # Keep the training label distribution; draw attributes uniformly
    labels = sample(label_prior, count)
    attrs = sample_uniform(attr_values, count)
    return [unet.sample(labels[i], attrs[i]) for i in range(count)]

# Step 3: train the classifier on real and synthetic data  [§3.1] [app. C]
def train_classifier(real, synthetic, alpha):
    # alpha: share of the loss from real data; a hyperparameter  [§3.1]
    for real_batch, synth_batch in zip(real, synthetic):
        loss_real = loss_fn(clf(real_batch.images), real_batch.targets)
        loss_synth = loss_fn(clf(synth_batch.images), synth_batch.targets)
        loss = alpha * loss_real + (1 - alpha) * loss_synth
        update(clf, loss)
    return clf

The toy experiments in app. D (Bernoulli and Gaussian mixture models) motivate the choice of mixing: with a finite training set, adding generated data can move the mixture closer to the target distribution app. D. They are not part of the training pipeline.

Dermatology labels are soft labels built from expert ratings: each rater gives up to three diagnoses with confidence, and the labels are aggregated with rank weighting §3.2.3 §3.3.

Heuristic baselines use fixed augmentation recipes listed in app. E (standard ImageNet-style colour jitter for histopathology; flips, crop, rotation, luminance and contrast for chest X-ray; flips, brightness, saturation, hue, contrast, rotation and blur for dermatology).

Hyperparameters

Source: app. E (model and classifier tables) unless noted. Values are as printed. "not stated" means the paper does not say.

Upsampler training preprocessing app. E, in order:

1. Bilinear upsample from 64×64 to the output size, with anti-aliasing with probability 0.5.
2. Add Gaussian noise with probability 0.2 and σ = 4.0 on the [0, 255] scale.
3. Gaussian blur with a 7×7 kernel, σ mean 0.0 and σ std 0.2 (as printed).
4. Quantise to 256 bins.
5. Normalise to [−1, 1].

Heuristic augmentation parameters app. E:

• Chest radiology: horizontal flip; random crop to 202×202; resize to 224×224 with bilinear interpolation and anti-aliasing; rotation of 15°; luminance shift uniform in [−0.1, 0.1]; contrast multiplier uniform in [0.8, 1.2], clipped to [0, 1].
• Dermatology: horizontal and vertical flips; brightness delta up to 0.1; saturation factor in [0.8, 1.2]; hue delta up to 0.02; contrast factor in [0.8, 1.2]; rotation in the range printed as [−150, 150]; Gaussian blur with σ sampled from {0.001, 0.01, 0.1, 1.0, 3.0, 5.0, 7.0}.
• Histopathology: brightness, contrast, saturation and hue jitter; magnitudes not stated.

Dataset sizes §3.2 app. A:

• Histopathology (CAMELYON17): labelled training 302,436; in-distribution validation 33,560; OOD validation 34,904; OOD test 85,054; unlabelled 1,799,247 app. A.
• Chest radiology: training 201,055 (CheXpert); in-distribution test 13,332; OOD test 17,723 (ChestX-ray14) §3.2.2 app. A.
• Dermatology: in-distribution 16,530 cases; OOD 1 6,639 images; OOD 2 3,900 images §3.2.3. The evaluation table in app. A lists different counts for the held-out slices (see Pitfalls).

Results

Author-reported. Setup: three medical tasks, ERM classifiers, comparison to heuristic augmentation and other baselines, five seeds, TPUs §3.1 app. F.1. Absolute values are mostly in figures, which this card does not reproduce.

The abstract gives 63.5% and the 7.5× fairness reduction; §3.2.3 adds the 27.3% in-distribution figure for the same split. Radiology gains are reported as averages over five conditions and two datasets, and the paper says the synthetic-data gain is largest out of distribution §3.2.2.

Paper vs code

No official code. Searched in order: the paper has no code link (abstract, comments, source); no project page is linked; GitHub search for "generative models fairness medical classifiers" under the google-deepmind organisation returned 0 results; "learned augmentations medical" returned xamyzhao/brainstorm (third-party, segmentation, MIT); "diffusion fairness medical" returned two third-party repositories (Hazrat-Ali9, mahmoudibrahim98/CompDiff). None of these is used.

Pitfalls

• The main method file is not in the build. methods.tex is commented out in main.tex, and it contains a one-hot label-plus-attribute encoding (concatenation of label vectors) that the compiled text does not state. The compiled text only says that the generator is conditioned on one or both of the label and a property §3.1 app. C. Treat the exact encoding as not stated.
• Sampling scheme changes for dermatology. The main text describes attribute-uniform sampling with the training label prior kept §3.1. The dermatology choice in the sampling-scheme analysis is different: synthetic samples of the four most common conditions are dropped by a filter, and the real-data importance weight is not used app. G. That analysis is run at 64×64, while the main dermatology results use 256×256 outputs §3.2.3 app. G.
• Mixing ratio and guidance are not stated. No numeric α is given, and no classifier-free guidance scale or sampler is given, although the method depends on guidance for label dropout app. E. Generator training steps and EMA are not stated.
• Zero label for unlabelled data. Unlabelled images use the all-zero conditioning vector. Classifier-free guidance then steers samples toward or away from that vector app. E. Implementers must choose the guidance weight themselves.
• Uncertain CheXpert labels. Uncertain labels are excluded from classifier training but used for the generator. The unmentioned label is mapped to negative §3.2.2 app. A.
• Dermatology count mismatches. The text gives OOD 2 as 3,900 images §3.2.3, but the evaluation table in app. A sums to 642 (396 female, 246 male). The text also says the minority subgroup is capped at 100 and 1,000 samples in the more and less skewed splits, but the training table shows 115 and 1,157 male samples §3.2.3 app. A.
• Rotation range. The dermatology heuristic rotation is printed as [−150, 150] degrees app. E. This is probably a typo, but the source does not say; check before copying.
• Fairness metric is underspecified. For dermatology the gap is a posterior estimate of the top-3 difference between a group and its outgroup, with a Beta(1, 1) prior and 100,000 posterior samples. The reported spread is the standard deviation of the MAP estimates, which the text calls the "central best estimate" app. F.2. It is not clear whether this is the reported statistic.
• Soft labels. Dermatology training and generation use aggregated soft labels from expert ratings, not one-hot labels §3.2.3 §3.3. The aggregation is described in §3.3. The loss used with soft targets is not stated.
• Hyperparameters were tuned on the baseline. They were selected on the baseline's in-distribution validation set and held fixed for the other methods app. E. Per-method tuning is not reported, so comparisons may favour the baseline's configuration.
• Upsampler dependency. Radiology and dermatology depend on a second diffusion model, and its training preprocessing (blur, noise, quantisation) is part of the pipeline app. E. Histopathology does not use an upsampler §3.2.1.
• Privacy is not addressed. The paper argues that synthetic data should not replace real data and that generative models need privacy guarantees before they can be used to anonymise data §2 §4.
• Expert study. The diagnosability numbers are from raters, and the real comparison uses a subset of real images with a leave-one-out aggregation §3.3.
• Compute. Hardware is TPUs; no compute budget or wall-clock cost is stated app. F.1.

Lineage

• Builds on: denoising diffusion probabilistic models §2 app. B; a UNet backbone app. E; classifier-free guidance, arXiv:2207.12598 app. E.
• Compared against: heuristic augmentation (RandAugment, ImageNet-style colour jitter, flips and crops), oversampling, focal loss, and ImageNet-21K or JFT pretraining app. E.2 §3.2.2 §3.2.3.
• Superseded by: not known from content on this wiki.