FairDiffusion

Despite the strong performance of these generative models, it remains an open question whether the quality of image generation is consistent across different demographic subgroups. To tackle these biases, we introduce FairDiffusion, an equity-aware latent diffusion model that enhances fairness in medical image generation via Fair Bayesian Perturbation.

Requirements

To install the prerequisites, run: conda env create -f environment.yml or pip install -r requirements.txt

FairGenMed Dataset

We present FairGenMed, the first dataset for studying fairness of medical generative models, providing detailed quantitative measurements of multiple clinical conditions to investigate the semantic correlation between text prompts and anatomical regions across various demographic subgroups.

The dataset can be accessed via this link. This dataset can only be used for non-commercial research purposes. At no time, the dataset shall be used for clinical decisions or patient care. The data use license is CC BY-NC-ND 4.0.

Our dataset includes 10,000 subjects for glaucoma detection with comprehensive demographic identity attributes including age, gender, race, ethnicity, preferred language, and marital status. Each subject has one Scanning Laser Ophthalmoscopy (SLO) fundus photo and one npz file. The size of SLO fundus photos is 512 x 664.

The NPZ files have the following attributes glaucoma: the label of glaucoma disease, 0 - non-glaucoma, 1 - glaucoma oct_bscans: images of OCT B-scans race: 0 - Asian, 1 - Black, 2 - White male: 0 - Female, 1 - Male hispanic: 0 - Non-Hispanic, 1 - Hispanic maritalstatus: 0 - Married or Partnered, 1 - Single, 2 - Divorced, 3 - Widowed, 4 - Legally Separated, and -1 - Unknown language: 0 - English, 1 - Spanish, 2 - Others

We have compiled all clinical measurements related to the 10,000 samples into a meta CSV file named data_summary.csv. Specifically, the cup-disc ratio, severity of vision loss, and status of spherical equivalent are denoted by the column names 'cdrstatus', 'mdseverity', and 'sestatus', respectively, in the datasummary.csv file.

Experiments

Train Stable Diffusion cd examples/text_to_image export MODEL_NAME="stabilityai/stable-diffusion-2-1" accelerate launch --main_process_port 29601 --multi_gpu --num_processes 2 --mixed_precision="fp16" train_text_to_image.py --text_encoder_type clip --dataset_dir <DATASET_DIR> --checkpointing_steps 5000 --pretrained_model_name_or_path=$MODEL_NAME --train_data_dir tmpp --datasets glaucoma --race_prompt --gender_prompt --ethnicity_prompt --use_ema --resolution=512 --center_crop --random_flip --train_batch_size=16 --gradient_accumulation_steps=1 --gradient_checkpointing --max_train_steps=100000 --learning_rate=1e-05 --max_grad_norm=1 --lr_scheduler="constant" --lr_warmup_steps=0 --output_dir=<OUTPUT_DIR> --cache_dir <CACHE_DIR>

Train FairDiffusion cd examples/text_to_image export MODEL_NAME="stabilityai/fair-diffusion-2-1" export TIME_WINDOW=30 export EXPLOITATION=0.95 accelerate launch --main_process_port 29601 --multi_gpu --num_processes 2 --mixed_precision="fp16" train_text_to_image.py --text_encoder_type clip --dataset_dir <DATASET_DIR> --checkpointing_steps 5000 --pretrained_model_name_or_path=$MODEL_NAME --train_data_dir tmpp --datasets glaucoma --race_prompt --gender_prompt --ethnicity_prompt --use_ema --resolution=512 --center_crop --random_flip --train_batch_size=16 --gradient_accumulation_steps=1 --gradient_checkpointing --max_train_steps=100000 --learning_rate=1e-05 --max_grad_norm=1 --lr_scheduler="constant" --lr_warmup_steps=0 --output_dir=<OUTPUT_DIR> --cache_dir <CACHE_DIR> --fair_time_window ${TIME_WINDOW} --fair_exploitation_rate ${EXPLOITATION}

Visualize Diffusion Model (Generation) python visualize_fairdiffusion.py --dataset_dir <DATASET_DIR> --initial_model stabilityai/stable-diffusion-2-1 --model_path <OUTPUT_DIR>/checkpoint-<TBD>/unet --datasets glaucoma --race_prompt --gender_prompt --ethnicity_prompt --vis_dir <VIS_DIR> --prompts prompts.txt --repeat_prompt 5

Evaluate Fairness of Diffusion Model (Generation) python evaluate_fairdiffusion.py --metrics_calculation_idx 1 --dataset_dir <DATASET_DIR> --initial_model stabilityai/stable-diffusion-2-1 --model_path <OUTPUT_DIR>/checkpoint-<TBD>/unet --datasets glaucoma --race_prompt --gender_prompt --ethnicity_prompt > OUTPUT.txt

Evaluate Fairness of Diffusion Model (Classification) - ViT-B ``` cd classificationcodebase DATASETDIR= RESULTDIR= MODELTYPE=ViT-B MODALITYTYPE=slofundus

VITWEIGHTS=imagenet BATCHSIZE=64 BLR=5e-4 WD=0.01 LD=0.55 DP=0.1

Baselines

EXPNAME=BASELINEGLAUCOMA PERFFILE=${MODELTYPE}${MODALITYTYPE}${EXPNAME}.csv python scripts/trainglaucomafair.py --task cls --epochs 50 --batchsize ${BATCHSIZE} --blr ${BLR} --minlr 1e-6 --warmupepochs 5 --weightdecay ${WD} --layerdecay ${LD} --droppath ${DP} --datadir ${DATASETDIR}/ --resultdir ${RESULTDIR}/${MODELTYPE}${MODALITYTYPE}${EXPNAME} --modeltype ${MODELTYPE} --modalitytypes ${MODALITYTYPE} --perffile ${PERFFILE} --vitweights ${VITWEIGHTS}

EXPNAME=BASELINEMDSEVERITY PERFFILE=${MODELTYPE}${MODALITYTYPE}${EXPNAME}.csv python scripts/trainglaucomafair.py --task mdseverity --epochs 50 --batchsize ${BATCHSIZE} --blr ${BLR} --minlr 1e-6 --warmupepochs 5 --weightdecay ${WD} --layerdecay ${LD} --droppath ${DP} --datadir ${DATASETDIR}/ --resultdir ${RESULTDIR}/${MODELTYPE}${MODALITYTYPE}${EXPNAME} --modeltype ${MODELTYPE} --modalitytypes ${MODALITYTYPE} --perffile ${PERFFILE} --vitweights ${VITWEIGHTS}

EXPNAME=BASELINECDRSTATUS PERFFILE=${MODELTYPE}${MODALITYTYPE}${EXPNAME}.csv python scripts/trainglaucomafair.py --task cdrstatus --epochs 50 --batchsize ${BATCHSIZE} --blr ${BLR} --minlr 1e-6 --warmupepochs 5 --weightdecay ${WD} --layerdecay ${LD} --droppath ${DP} --datadir ${DATASETDIR}/ --resultdir ${RESULTDIR}/${MODELTYPE}${MODALITYTYPE}${EXPNAME} --modeltype ${MODELTYPE} --modalitytypes ${MODALITYTYPE} --perffile ${PERFFILE} --vitweights ${VITWEIGHTS}

EXPNAME=BASELINESESTATUS PERFFILE=${MODELTYPE}${MODALITYTYPE}${EXPNAME}.csv python scripts/trainglaucomafair.py --task sestatus --epochs 50 --batchsize ${BATCHSIZE} --blr ${BLR} --minlr 1e-6 --warmupepochs 5 --weightdecay ${WD} --layerdecay ${LD} --droppath ${DP} --datadir ${DATASETDIR}/ --resultdir ${RESULTDIR}/${MODELTYPE}${MODALITYTYPE}${EXPNAME} --modeltype ${MODELTYPE} --modalitytypes ${MODALITYTYPE} --perffile ${PERFFILE} --vitweights ${VITWEIGHTS}

Generated Dataset

EXPNAME=FAIRDIFFUSIONGLAUCOMA PERFFILE=${MODELTYPE}${MODALITYTYPE}${EXPNAME}.csv python scripts/trainglaucomafair.py --task cls --fairdiffusion --initialmodel stabilityai/stable-diffusion-2-1 --modelpath /checkpoint-/unet --epochs 50 --batchsize ${BATCHSIZE} --blr ${BLR} --minlr 1e-6 --warmupepochs 5 --weightdecay ${WD} --layerdecay ${LD} --droppath ${DP} --datadir ${DATASETDIR}/ --resultdir ${RESULTDIR}/${MODELTYPE}${MODALITYTYPE}${EXPNAME} --modeltype ${MODELTYPE} --modalitytypes ${MODALITYTYPE} --perffile ${PERFFILE} --vitweights ${VITWEIGHTS}

EXPNAME=FAIRDIFFUSIONMDSEVERITY PERFFILE=${MODELTYPE}${MODALITYTYPE}${EXPNAME}.csv python scripts/trainglaucomafair.py --task mdseverity --fairdiffusion --initialmodel stabilityai/stable-diffusion-2-1 --modelpath /checkpoint-/unet --epochs 50 --batchsize ${BATCHSIZE} --blr ${BLR} --minlr 1e-6 --warmupepochs 5 --weightdecay ${WD} --layerdecay ${LD} --droppath ${DP} --datadir ${DATASETDIR}/ --resultdir ${RESULTDIR}/${MODELTYPE}${MODALITYTYPE}${EXPNAME} --modeltype ${MODELTYPE} --modalitytypes ${MODALITYTYPE} --perffile ${PERFFILE} --vitweights ${VITWEIGHTS}

EXPNAME=FAIRDIFFUSIONCDRSTATUS PERFFILE=${MODELTYPE}${MODALITYTYPE}${EXPNAME}.csv python scripts/trainglaucomafair.py --task cdrstatus --fairdiffusion --initialmodel stabilityai/stable-diffusion-2-1 --modelpath /checkpoint-/unet --epochs 50 --batchsize ${BATCHSIZE} --blr ${BLR} --minlr 1e-6 --warmupepochs 5 --weightdecay ${WD} --layerdecay ${LD} --droppath ${DP} --datadir ${DATASETDIR}/ --resultdir ${RESULTDIR}/${MODELTYPE}${MODALITYTYPE}${EXPNAME} --modeltype ${MODELTYPE} --modalitytypes ${MODALITYTYPE} --perffile ${PERFFILE} --vitweights ${VITWEIGHTS}

EXPNAME=FAIRDIFFUSIONSESTATUS PERFFILE=${MODELTYPE}${MODALITYTYPE}${EXPNAME}.csv python scripts/trainglaucomafair.py --task sestatus --fairdiffusion --initialmodel stabilityai/stable-diffusion-2-1 --modelpath /checkpoint-/unet --epochs 50 --batchsize ${BATCHSIZE} --blr ${BLR} --minlr 1e-6 --warmupepochs 5 --weightdecay ${WD} --layerdecay ${LD} --droppath ${DP} --datadir ${DATASETDIR}/ --resultdir ${RESULTDIR}/${MODELTYPE}${MODALITYTYPE}${EXPNAME} --modeltype ${MODELTYPE} --modalitytypes ${MODALITYTYPE} --perffile ${PERFFILE} --vitweights ${VITWEIGHTS}

```

Evaluate Fairness of Diffusion Model (Classification) - EfficientNet ``` cd classification_codebase

DATASETDIR= RESULTDIR= MODELTYPE=( efficientnet ) NUMEPOCH=10 MODALITYTYPE='slofundus' ATTRIBUTE_TYPE=race

OPTIMIZER='adamw' OPTIMIZERARGUMENTS='{"lr": 0.001, "weightdecay": 0.01}'

SCHEDULER='steplr' SCHEDULERARGUMENTS='{"step_size": 30, "gamma": 0.1}'

LR=1e-3

Baselines

BATCHSIZE=6 EXPNAME=BASELINEGLAUCOMA PERFFILE=${MODELTYPE}${MODALITYTYPE}${EXPNAME}.csv python scripts/trainglaucomafair.py --task cls --datadir ${DATASETDIR}/ --resultdir ${RESULTDIR}/${MODELTYPE}${MODALITYTYPE}${EXPNAME} --modeltype ${MODELTYPE} --imagesize 200 --lr ${LR} --weight-decay 0. --momentum 0.1 --batchsize ${BATCHSIZE} --epochs ${NUMEPOCH} --modalitytypes ${MODALITYTYPE} --perffile ${PERFFILE}

BATCHSIZE=6 EXPNAME=BASELINEMDSEVERITY PERFFILE=${MODELTYPE}${MODALITYTYPE}${EXPNAME}.csv python scripts/trainglaucomafair.py --task mdseverity --datadir ${DATASETDIR}/ --resultdir ${RESULTDIR}/${MODELTYPE}${MODALITYTYPE}${EXPNAME} --modeltype ${MODELTYPE} --imagesize 200 --lr ${LR} --weight-decay 0. --momentum 0.1 --batchsize ${BATCHSIZE} --epochs ${NUMEPOCH} --modalitytypes ${MODALITYTYPE} --perffile ${PERFFILE}

BATCHSIZE=16 EXPNAME=BASELINECDRSTATUS PERFFILE=${MODELTYPE}${MODALITYTYPE}${EXPNAME}.csv python scripts/trainglaucomafair.py --task cdrstatus --datadir ${DATASETDIR}/ --resultdir ${RESULTDIR}/${MODELTYPE}${MODALITYTYPE}${EXPNAME} --modeltype ${MODELTYPE} --imagesize 200 --lr ${LR} --weight-decay 0. --momentum 0.1 --batchsize ${BATCHSIZE} --epochs ${NUMEPOCH} --modalitytypes ${MODALITYTYPE} --perffile ${PERFFILE}

BATCHSIZE=16 EXPNAME=BASELINESESTATUS PERFFILE=${MODELTYPE}${MODALITYTYPE}${EXPNAME}.csv python scripts/trainglaucomafair.py --task sestatus --datadir ${DATASETDIR}/ --resultdir ${RESULTDIR}/${MODELTYPE}${MODALITYTYPE}${EXPNAME} --modeltype ${MODELTYPE} --imagesize 200 --lr ${LR} --weight-decay 0. --momentum 0.1 --batchsize ${BATCHSIZE} --epochs ${NUMEPOCH} --modalitytypes ${MODALITYTYPE} --perffile ${PERFFILE}

Generated Dataset

BATCHSIZE=6 EXPNAME=FAIRDIFFUSIONGLAUCOMA PERFFILE=${MODELTYPE}${MODALITYTYPE}${EXPNAME}.csv python scripts/trainglaucomafair.py --task cls --fairdiffusion --initialmodel stabilityai/stable-diffusion-2-1 --modelpath <OUTPUTDIR>/checkpoint-/unet --datadir ${DATASETDIR}/ --resultdir ${RESULTDIR}/${MODELTYPE}${MODALITYTYPE}${EXPNAME} --modeltype ${MODELTYPE} --imagesize 200 --lr ${LR} --weight-decay 0. --momentum 0.1 --batchsize ${BATCHSIZE} --epochs ${NUMEPOCH} --modalitytypes ${MODALITYTYPE} --perffile ${PERF_FILE}

BATCHSIZE=6 EXPNAME=FAIRDIFFUSIONMDSEVERITY PERFFILE=${MODELTYPE}${MODALITYTYPE}${EXPNAME}.csv python scripts/trainglaucomafair.py --task mdseverity --fairdiffusion --initialmodel stabilityai/stable-diffusion-2-1 --modelpath <OUTPUTDIR>/checkpoint-/unet --datadir ${DATASETDIR}/ --resultdir ${RESULTDIR}/${MODELTYPE}${MODALITYTYPE}${EXPNAME} --modeltype ${MODELTYPE} --imagesize 200 --lr ${LR} --weight-decay 0. --momentum 0.1 --batchsize ${BATCHSIZE} --epochs ${NUMEPOCH} --modalitytypes ${MODALITYTYPE} --perffile ${PERF_FILE}

BATCHSIZE=16 EXPNAME=FAIRDIFFUSIONCDRSTATUS PERFFILE=${MODELTYPE}${MODALITYTYPE}${EXPNAME}.csv python scripts/trainglaucomafair.py --task cdrstatus --fairdiffusion --initialmodel stabilityai/stable-diffusion-2-1 --modelpath <OUTPUTDIR>/checkpoint-/unet --datadir ${DATASETDIR}/ --resultdir ${RESULTDIR}/${MODELTYPE}${MODALITYTYPE}${EXPNAME} --modeltype ${MODELTYPE} --imagesize 200 --lr ${LR} --weight-decay 0. --momentum 0.1 --batchsize ${BATCHSIZE} --epochs ${NUMEPOCH} --modalitytypes ${MODALITYTYPE} --perffile ${PERF_FILE}

BATCHSIZE=16 EXPNAME=FAIRDIFFUSIONSESTATUS PERFFILE=${MODELTYPE}${MODALITYTYPE}${EXPNAME}.csv python scripts/trainglaucomafair.py --task sestatus --fairdiffusion --initialmodel stabilityai/stable-diffusion-2-1 --modelpath <OUTPUTDIR>/checkpoint-/unet --datadir ${DATASETDIR}/ --resultdir ${RESULTDIR}/${MODELTYPE}${MODALITYTYPE}${EXPNAME} --modeltype ${MODELTYPE} --imagesize 200 --lr ${LR} --weight-decay 0. --momentum 0.1 --batchsize ${BATCHSIZE} --epochs ${NUMEPOCH} --modalitytypes ${MODALITYTYPE} --perffile ${PERF_FILE}

```