PaperlayerView exact revision on arXiv

Source-linked AI summary

Generative models improve fairness of medical classifiers under distribution shifts

Ira Ktena, Olivia Wiles, Isabela Albuquerque, Sylvestre-Alvise Rebuffi, Ryutaro Tanno, Abhijit Guha Roy, Shekoofeh Azizi, Danielle Belgrave, Pushmeet Kohli, Alan Karthikesalingam, Taylan Cemgil, Sven Gowal

arXiv:2304.09218v1cs.CV

TL;DR

Domain shifts, underrepresented groups, and limited labelled data challenge robust and fair medical classifiers. The paper learns label-conditioned augmentations with diffusion models from labelled and unlabelled data, then combines synthetic and real images for classifier training. Across histopathology, chest radiology, and dermatology, this improves robustness and fairness, with especially large reported gains under distribution shifts.

  • Problem

    Domain generalisation can amplify bias against underrepresented groups, while labelled medical data are costly or difficult to obtain.

  • Method

    The approach trains conditional diffusion models using labelled and available unlabelled data, generates specified synthetic images, and combines them with real labelled images for classifier training.

  • Results

    Synthetic data improve robustness and fairness across histopathology, chest radiology, and dermatology, including 5.2% higher average AUC OOD and 44.6% improvement in sex fairness gap for chest radiology.

  • Takeaways & Limitations

    Learned synthetic augmentations can improve statistical fairness, balanced accuracy, high-risk sensitivity, and in- and out-of-distribution robustness across three medical imaging settings.

  • Takeaways & Limitations

    Synthetic data should complement rather than replace real data, and realistic generated images may create overconfidence without continued data collection and caveat analysis.

Abstract

from arXiv · show

A ubiquitous challenge in machine learning is the problem of domain generalisation. This can exacerbate bias against groups or labels that are underrepresented in the datasets used for model development. Model bias can lead to unintended harms, especially in safety-critical applications like healthcare. Furthermore, the challenge is compounded by the difficulty of obtaining labelled data due to high cost or lack of readily available domain expertise. In our work, we show that learning realistic augmentations automatically from data is possible in a label-efficient manner using generative models. In particular, we leverage the higher abundance of unlabelled data to capture the underlying data distribution of different conditions and subgroups for an imaging modality. By conditioning generative models on appropriate labels, we can steer the distribution of synthetic examples according to specific requirements. We demonstrate that these learned augmentations can surpass heuristic ones by making models more robust and statistically fair in- and out-of-distribution. To evaluate the generality of our approach, we study 3 distinct medical imaging contexts of varying difficulty: (i) histopathology images from a publicly available generalisation benchmark, (ii) chest X-rays from publicly available clinical datasets, and (iii) dermatology images characterised by complex shifts and imaging conditions. Complementing real training samples with synthetic ones improves the robustness of models in all three medical tasks and increases fairness by improving the accuracy of diagnosis within underrepresented groups. This approach leads to stark improvements OOD across modalities: 7.7% prediction accuracy improvement in histopathology, 5.2% in chest radiology with 44.6% lower fairness gap and a striking 63.5% improvement in high-risk sensitivity for dermatology with a 7.5x reduction in fairness gap.

1. Introduction

Medical models can perform poorly across populations and distribution shifts, while labelled data are scarce and fairness and robustness are rarely addressed together. This work uses generative models and available unlabelled data to augment diagnostic training across modalities.

  • Only 24% of published studies evaluate algorithms on external cohorts or compare external performance with clinical experts.
  • Robust performance across populations is difficult because subgroup prevalence, data scarcity, evaluation-dataset acquisition, and realistic multi-axis shifts complicate generalisation.
  • Prior studies report unexpectedly poor performance on underrepresented populations in radiology, histopathology, and dermatology.
  • The proposed approach leverages generative models and potentially available unlabelled data to capture distributions and augment real samples for diagnostic models across three modalities.
  • The study reports improvements in top-level performance and fairness under distribution shifts, alongside high-quality diffusion-generated images across modalities.

2. Background

Medical imaging research has used generative models for augmentation, while diffusion models have recently shown strong image-generation performance. This paper situates conditional diffusion augmentation within concerns about fairness, clinical evaluation, and synthetic-data privacy.

  • Generative models, including GANs, have been used to improve medical imaging performance and address underrepresented conditions.
  • Diffusion models have been applied to medical image generation across domains, including 3D MR and CT images and text-conditioned chest X-rays.
  • The paper rejects treating synthetic images as replacement anonymized data without privacy guarantees for generative-model training.
  • The proposed workflow trains a diffusion model on labelled and available unlabelled data, samples specified synthetic images, and trains a classifier on real plus synthetic images.
  • Fairness studies have identified discrepancies across demographic and socioeconomic groups and their intersections in clinical machine-learning systems.

3. Results

The proposed approach uses diffusion models and available labelled and unlabelled data to generate targeted augmentations for improving robustness and fairness across medical imaging settings. Across histopathology, radiology, and dermatology experiments, synthetic data improves out-of-distribution performance, with especially strong gains in low-data and underrepresented-condition settings.

  • Overview of the proposed approach and experimental setting: Diffusion models are trained on labelled and unlabelled data, conditioned on diagnostic labels and optionally properties such as hospital identity or sensitive attributes.Synthetic samples are generated through fair sampling and combined with real training images for downstream classifiers.
  • Histopathology: 7.7% absolute improvement over the baseline model was achieved on the histopathology test hospital when generated samples were combined with color augmentation.Synthetic data also outperformed real-data baselines in-distribution and was most effective in the more skewed, low-data setting.
  • Chest radiology: 5.2% improvement on average AUC OOD and a 44.6% improvement in sex fairness gap were observed across five radiology conditions.Synthetic images improved average AUC both in-distribution and out-of-distribution, with particularly large gains for cardiomegaly.
  • Dermatology: Synthetic-data methods significantly improved dermatology high-risk sensitivity and fairness in the out-of-distribution setting, especially under greater sensitive-attribute skew.Label & property conditioning produced 27.3% better high-risk sensitivity than the baseline in the more skewed setting.
  • Dermatology: The conditional generative model captured characteristics across diverse dermatological conditions, including scarce conditions such as seborrheic dermatitis, alopecia areata, and hidradenitis.Expert rating and consensus-based soft labels were used to assess generated examples.
  • Distribution alignment: Learned augmentations yielded 18.6% lower MMD than heuristic augmentations, indicating improved alignment between domain representations.The analysis measured mismatches in the penultimate-layer representation space; a minor performance decline in the less skewed dermatology setting may reflect conflict with useful spurious correlations.

4. Discussion

The paper finds that conditional generative models can improve robustness and fairness across medical imaging settings, while synthetic data introduces important evaluation, bias, privacy, and transparency concerns.

  • Findings: Conditional diffusion models improved robustness, statistical fairness, balanced accuracy, and high-risk sensitivity across histopathology, radiology, and dermatology tasks.The approach also produced greater out-of-distribution than in-distribution benefits in some settings and required only a few labelled examples in histopathology and dermatology.
  • Findings: Synthetic data can outperform real-data baselines in histopathology and radiology, whereas real data remains more essential for downstream dermatology models.The authors attribute improvements partly to better domain alignment and more complex decision boundaries that reduce reliance on spurious correlations.
  • Practical implications: The same synthetic data may support augmentation across different models and potentially across tasks.This extends its utility beyond the specific downstream evaluations studied in the paper.
  • Practical implications: Synthetic images should complement, not replace, additional data collection and representation efforts for underrepresented conditions and populations.The authors recommend rigorous evaluation on real data, ideally outside the main source domain.
  • Risks and limitations: Poor-quality or biased generative models may reproduce or amplify training-data bias and cause downstream classifiers to make confident mistakes.This risk is especially important when the generator cannot represent certain label–attribute combinations or captures only part of the true distribution.
  • Risks and limitations: Conclusions are limited when evaluation datasets are not diverse, representative, or effective at capturing high-risk conditions and realistic out-of-distribution performance.The paper therefore emphasizes careful interpretation of what each evaluation setup measures.
  • Risks and limitations: Unobserved sensitive attributes, reductive categorical labels, and imperfect skin-tone measures constrain fairness assessment and future label-free approaches.Future work must also address memorization of rare attributes and conditions and provide privacy guarantees.
  • Practical implications: Synthetic images should be tagged when added to databases because they may perpetuate biases and affect reuse in other settings.The concern applies particularly when datasets are reused by different practitioners or for different purposes.

6. Author Contributions

The project involved shared contributions spanning method design, experiments, software engineering, and analyses across modalities and sampling schemes.

  • Project and method: O.W., S.G., and P.K. initiated the project, while O.W., I.K., and S.G. contributed to method and experiment design.O.W., S.G., and T.C. contributed to the method formulation.
  • Analysis and implementation: Contributors provided dataset pointers, software engineering, and analyses of distribution matching, spurious correlations, sampling schemes, and high-resolution image generation.These contributions supported the paper’s methodological and experimental work.
  • Experiments: I.K. and O.W. performed experiments across imaging modalities, with additional contributions to dermatology nearest-neighbour analysis and clinical problem formulation.The contributions covered multiple medical contexts and analyses.

A. Datasets

The study evaluates generative augmentation across three de-identified medical imaging datasets spanning hospital, demographic, label, and imaging-condition shifts.

  • Overview: Three de-identified datasets were used to train downstream classifiers and diffusion models across different medical modalities and contexts.The datasets support evaluation across multiple medical imaging settings.
  • Histopathology: CAMELYON17 provides labelled and unlabelled data from three training hospitals, plus in-distribution and out-of-distribution validation hospitals.Hospital differences arise from staining procedures, and the task detects breast cancer metastases in lymph-node histopathology patches.
  • Histopathology: CAMELYON17 hospital subsets are approximately evenly split between tumorous and non-tumorous images.This dataset statistic is reported for all hospitals.
  • Chest radiology: Figure A1 shows age histograms and normalized label distributions for five CheXpert conditions.The figure characterizes demographic and label distributions in the training data.
  • Chest radiology: CheXpert training used 201,055 samples from 119,352 male and 81,703 female individuals, with binary gender labels.Uncertain samples were used only for diffusion-model training, and the unmentioned label was mapped to negative, producing a highly imbalanced dataset.
  • Dermatology: The dermatology dataset was artificially skewed along sex, skin tone, and age to disentangle the effects of complex distribution shifts.The skewing procedure measures how increasingly severe shifts affect low-data regions and method performance.
  • Dermatology: Table A2 reports training-sample counts for sensitive-attribute labels across simulated splits.These counts characterize the data available in the skewed dermatology settings.

B.1. Learning augmentations with generative models in Health

Prior medical-imaging work used GANs and newer diffusion models for augmentation, but the literature also highlights limitations in diversity, quality, privacy, and bias.

  • Generative augmentation: GAN-based augmentation has been applied to medical diagnosis, including liver lesions, diabetic retinopathy, breast masses, dermoscopy, and chest X-rays.Some studies targeted balanced multiclass accuracy, high-risk underrepresented labels, or rare-disease classification.
  • Evaluation: Table A3 reports evaluation-sample counts annotated with sensitive-attribute labels across domains.These counts support assessment across domain-specific demographic distributions.
  • Generative augmentation: Conditional image generation has also been used to disentangle image content and style when conditioning information is unavailable.Applications included dermatoscopic images and lung CT scans.
  • Scope: GAN-based methods have additionally supported pixel-wise tasks such as vessel-contour and brain-lesion segmentation, outside this study’s whole-image scope.The paper refers readers to reviews for broader GAN-based augmentation work.
  • Diffusion models: Learned GAN augmentations have often been characterized by limited diversity and quality, motivating interest in denoising diffusion probabilistic models.Recent diffusion work includes large-scale generation, clinical 3D imaging, segmentation, and text-conditioned chest X-ray synthesis.
  • Ethics: Synthetic medical images raise ethical questions concerning privacy and data bias.These concerns remain relevant alongside their potential clinical utility.
  • Fairness: Machine-learning bias can arise from data acquisition, human decisions, missing features, and label scarcity throughout the development and deployment pipeline.Fairness analysis therefore considers disparities in predictive accuracy across demographic or socioeconomic traits.
  • Fairness: Training-set gender composition has been linked to predictive accuracy for the same gender group at test time in chest X-ray diagnosis.This prior finding motivates attention to demographic imbalance in medical classifier training.

C. Method

The method combines labelled training data with synthetic samples generated from labelled and unlabelled data to approximate a fair target distribution. Generative conditioning and a mixing parameter control how synthetic examples are sampled and incorporated into classifier training.

  • Motivation: Toy experiments motivate synthetic augmentation by testing whether classifiers can perform well when only a few samples represent the underlying distribution.The experiments explicitly target low-data settings and demonstrate utility even in these simplified cases.
  • Data and assumptions: The framework assumes labelled images, labels, attributes, and an additional dataset of unlabelled images, with attributes including demographics or hospital identity.The generative model is trained using the available data, while unlabelled images can be used as desired.
  • Data and assumptions: The target fair distribution may differ from training data through the joint distribution of attributes and labels, even when conditional image distributions match.This captures shifts such as overrepresentation of particular attributes or attribute-label combinations.
  • Synthetic-data integration: The approach combines real training samples with synthetic samples from a generative model to mimic the fair distribution and improve fairness.The combined dataset is constructed using a probability parameter α.
  • Synthetic-data integration: Synthetic and real data contribute separate loss expectations, weighted by α and 1 − α, when training the classifier.The mixing weight is treated as a hyperparameter, while the label distribution is preserved and attributes are sampled uniformly in the main settings.

D.1. A Bernoulli distribution over a single variable.

Toy experiments show that generated data can reduce mismatch between finite training data and a target distribution, including when hidden attributes cannot be directly resampled. Additional generated points also improve classifiers fitted from labelled and unlabelled data, although benefits are limited near-perfect or highly overlapping settings.

  • Bernoulli distribution: Combining real and generated data can perfectly match the target distribution in the specific toy example.The result is reported for the case where the generative distribution equals the target distribution.
  • Bernoulli distribution: As training-sample size increases, generated data becomes less helpful; it is most useful when its distribution resembles the true distribution, but remains helpful under some shift.The optimal synthetic-data weight decreases toward zero as sample size increases.
  • Labelled and hidden variables: When hidden attributes cannot be resampled directly, the method conditions on labels while retaining the conditional attribute distributions of the training and generative sources.The resulting distributions are optimized experimentally using KL divergence.
  • Labelled and hidden variables: Even for large sample sizes, additional generated data improves KL divergence; for small samples, the empirical data distribution can yield infinite divergence.This toy result supports generated data in low-data settings even when training and evaluation share the same underlying distribution.
  • GMM classification: A GMM fitted to unlabelled points is combined with labelled points to generate additional classifier-training examples across component and feature settings.The labelled points identify the classes associated with fitted mixture modes.
  • GMM classification: With 500 unlabelled points, the fitted GMM closely matches the true distribution, while additional generated data improves downstream performance across tested settings.Benefits are limited when labelled-data performance is already near perfect or mixture components overlap substantially.

F.1. Experimental setup

The experiments evaluate diffusion-based augmentation across medical imaging settings using performance and fairness metrics under in- and out-of-distribution conditions. The protocol includes domain-aware fairness definitions, repeated-seed evaluation, and comparisons with heuristic and baseline augmentation methods.

  • Experimental protocol: All model versions and baselines are evaluated with five initialization seeds, reporting averages and standard deviations across metrics.Experiments run on Tensor Processing Units.
  • Fairness metrics: Fairness is measured by modality-specific gaps or parity: hospital performance gaps in histopathology, AUC parity in radiology, and subgroup performance gaps in dermatology.The definitions vary with the classification task and sensitive attributes.
  • Distribution-shift estimation: Domain mismatches are estimated in penultimate-layer representations by projecting images into a 6144-dimensional space and computing MMD between datasets.The input images are represented from an initial 64 × 64 space.
  • Distribution-shift estimation: A cubic polynomial kernel is used for MMD to limit hyperparameter selection while capturing mismatches through third-order moments.Pairwise estimates use samples of size N = 300 and are tested with the Mann-Whitney U test at 95% significance.
  • Generated data: Class-conditioned diffusion models generate healthy and abnormal histopathology images, as well as healthy and condition-specific chest X-rays.The figures illustrate generated images used in the evaluation contexts.
  • Histopathology evaluation: In histopathology, varying labelled-sample counts tests the effect of synthetic augmentation on both in-distribution and out-of-distribution generalization.The dataset is balanced, so the experiment focuses on labelled-data availability rather than class imbalance.
  • Histopathology results: Generated samples combined with colour augmentation achieve about 10% absolute OOD improvement in high-data settings and about 24% in the one-label setting over baseline.The approach can match baseline performance using 1–10 labelled samples instead of 1,000 in the reported low-data comparison.

G.2.2. Results per condition

Chest-X-ray benefits from synthetic augmentation vary by condition and distribution. Cardiomegaly shows significant benefit, effusion benefits more out-of-distribution, and atelectasis shows only marginal OOD improvement.

  • Condition-specific outcomes: Cardiomegaly benefits significantly from synthetic data across the evaluated chest-X-ray conditions.The comparison is based on AUC values across methods and in- versus out-of-distribution datasets.
  • Condition-specific outcomes: Effusion benefits more from synthetic data out-of-distribution than in-distribution.The results are reported as condition-specific AUC comparisons.
  • Condition-specific outcomes: Atelectasis receives only marginal benefit from synthetic images in the out-of-distribution setting.The figure reports AUC per condition for in- and out-of-distribution datasets.

G.3.1. Multiple metrics across datasets

Across dermatology shifts, generated data generally maintains or improves accuracy, fairness, and high-risk sensitivity, including on the strongest out-of-distribution shifts. Representation analyses compare how learned, heuristic, and RandAugment augmentations affect domain alignment.

  • Accuracy: Generated data improves or maintains dermatology accuracy across distribution shifts and datasets, especially balanced accuracy under stronger prevalence shifts.RandAugment is also helpful, while oversampling and focal loss provide minimal improvements.
  • Fairness: Generated data consistently improves or maintains fairness for sex, skintone, and age, including out-of-distribution datasets with stronger shifts.Skin-tone fairness is reported only in-distribution because out-of-distribution skin-tone labels are unavailable.
  • High-risk sensitivity: Synthetic data maintains or improves high-risk sensitivity across shifts and datasets, with the strongest gains on the most out-of-distribution dataset.Its performance is consistently similar to or better than heuristic augmentation on this metric.
  • Representation analysis: The analysis estimates pairwise domain mismatches in penultimate-layer representations using multiple MMD estimates across domains.Models share experimental conditions to isolate the effect of augmentation on estimated distribution shifts; the table reports statistically significant comparisons.
  • Representation analysis: RandAugment produces more aligned representations than learned and heuristic augmentations across all considered domain pairs.The authors hypothesize that this removes domain-specific cues while potentially increasing reliance on spurious correlations.

G.3.3. Principal component analysis for spurious correlations

The paper examines how augmentation and sampling choices shape classifier representations, fairness, accuracy, and high-risk detection. It emphasizes filtering synthetic samples and comparing learned features across augmentation schemes under in- and out-of-distribution evaluation.

  • Principal component analysis for spurious correlations: The representation analysis uses projected penultimate-layer features to compare augmentation schemes and investigate whether learned augmentations promote better out-of-distribution generalization.The study explicitly accounts for simplicity bias and links it to reliance on spurious correlations.
  • Principal component analysis for spurious correlations: Figure G8 reports the average number of principal components needed to retain 0.60, 0.70, 0.80, 0.90, and 0.95 of projected-data variance across five model initializations.The comparison covers RandAugment, learned, and heuristic augmentations.
  • Individual-level analysis: Across training setups, some individuals are correctly classified by one augmentation scheme but not by the others.Each setup uses five runs, and consistent misclassification identifies underserved individuals for that setup.
  • Individual-level analysis: When training data is skewed toward females, generated data more often correctly classifies out-of-distribution males with high-risk conditions and reduces underserved individuals.The augmentation setups remain complementary because each correctly classifies some samples missed by the others.
  • Sampling schemes: The training objective combines real and synthetic examples, with w1 weighting real-example losses and w2 filtering generated samples; α controls the real-data proportion.The target fair distribution is approximated by the generative model.
  • Sampling schemes: The filtering function w2 excludes generated samples from the dominant four dermatology conditions and retains samples from less frequent conditions.It is defined as an indicator that outputs 0.0 for dominant conditions and 1.0 otherwise.
  • Sampling schemes: Without filtering, importance weighting at equality levels 0.5 and 1.0 increases overall top-3 accuracy and reduces the sex fairness gap in both evaluation settings.The authors offer reduced correlation between sex and some skin conditions as one possible explanation for improved generalisable features.
  • Sampling schemes: With filtering, importance weighting narrows the fairness gap but improves accuracy in-distribution while worsening it out-of-distribution.This indicates an interaction between importance weighting and synthetic-sample filtering.
Loading 2304.09218v1…