AAAI 2026 Fair Image Classification · Diffusion Models

Harnessing Diffusion-Generated Synthetic Images for Fair Image Classification

Fine-tuning Stable Diffusion directly on each biased training group — via LoRA, DreamBooth, and a new Clustered DreamBooth — generates group-balanced synthetic data that trains fairer classifiers than prompt-only generation. Clustered DreamBooth reaches 79.1% worst-group accuracy on average across three benchmarks, beating Group-DRO by up to 60.7 points on Waterbirds once dataset bias becomes severe.

Abhipsa Basu¹ Aviral Gupta² Abhijnya Bhat³ R. Venkatesh Babu¹

¹ Indian Institute of Science, Bangalore ² BITS Pilani ³ Stanford University

arXiv 2511.08711 PDF Paper </> Code {} BibTeX

Indian Institute of Science BITS Pilani Stanford University

01 / The problem: generative models don't know your dataset

Four generation pipelines: Vanilla SD, LoRA-finetuned, DreamBooth, and Clustered DreamBooth, each producing a Blond Male image from a text prompt or training images. — Prompting vanilla Stable Diffusion with a group description (e.g. “photo of a blond male person”) drifts from the true data distribution. Fine-tuning the diffusion model directly on each training group — via LoRA, DreamBooth, or per-cluster DreamBooth — keeps generations in-distribution while still representing the target group.

02 / Abstract

Image classification systems often inherit biases from uneven group representation in training data — e.g. in face datasets for hair-color classification, blond hair may be disproportionately associated with females, reinforcing stereotypes. A recent approach, FFR, leverages Stable Diffusion (SD) to generate balanced training data, but such prompt-only models often struggle to preserve the original data distribution.

We explore multiple diffusion-finetuning techniques — LoRA and DreamBooth — to generate images that more accurately represent each training group by learning directly from its samples. To prevent a single DreamBooth model from being overwhelmed by excessive intra-group variation, we introduce Clustered DreamBooth: clustering images within each group and training a separate DreamBooth model per cluster.

These models generate group-balanced data for pretraining a classifier, which is then fine-tuned on real data. Across three fairness benchmarks, the studied fine-tuning approaches outperform vanilla Stable Diffusion on average, and achieve results comparable to SOTA debiasing techniques like Group-DRO — while surpassing them as dataset bias severity increases.

Two-stage pipeline

Pretrain on group-balanced synthetic images, then fine-tune only the softmax layer on the real, biased dataset — no full-network hyperparameter tuning.

Clustered DreamBooth

A pool of per-cluster DreamBooth models captures intra-group diversity that a single DreamBooth model overwrites.

Robust under severe bias

Traditional debiasing (Group-DRO, SELF) collapses at a 0.999 bias ratio; generative pretraining degrades far more gracefully.

03 / Method

Two-stage pipeline: Stage 1 pretrains a classifier on group-balanced generated data; Stage 2 finetunes only the linear classifier on the original biased data. — Stage 1 generates images uniformly across every group (e.g. non-blond female, non-blond male, blond female, blond male) and pretrains a classifier with CE + SupCon losses. Stage 2 freezes the feature encoder and fine-tunes only the linear classifier on the original biased dataset.

Four generation strategies

Vanilla SD generates from class + bias labels alone — independent of the training data, so it can drift out-of-distribution. LoRA-finetuning fine-tunes SD on each training group separately, making the generator aware of the data. DreamBooth goes further, binding a group to a unique identifier token (“[V]”) for stronger resemblance to real samples.

Clustered DreamBooth addresses a failure mode of plain DreamBooth: a group like “Blond Male” contains many individuals sharing only one trait. We cluster each group's CLIP embeddings into k subsets and train a separate DreamBooth model per cluster, so no single model is overwhelmed by intra-group variation.

$$\text{CLIP-Score}(I,p_c,\bar z_g) = \alpha \cdot \text{CLIP-Label}(I,p_c) + (1-\alpha)\cdot \text{CLIP-Centroid}(I,\bar z_g)$$

CLIP-Label scores image–text similarity to the class prompt; CLIP-Centroid scores similarity to the group's CLIP-embedding centroid. The top 75% of generated images by this combined score are kept for training.

Two-stage training objective

Both the pretraining stage (on generated data) and the fine-tuning stage (on real data) optimize a weighted combination of Cross-Entropy and Supervised Contrastive loss, which sharpens class separation in the learned features.

$$\mathcal{L}_{CE} = -\frac{1}{|X|}\sum_{i=1}^{|X|}\sum_{j=1}^{|Y|} p_{ij}\log \hat p_{ij}$$

$$\mathcal{L} = \beta \cdot \mathcal{L}_{CE} + (1-\beta)\cdot \mathcal{L}_{\text{sup-con}}, \quad \beta = 0.5$$

Stage 2 freezes the pretrained feature encoder e^pre and fine-tunes only the linear classifier c^pre — full-network fine-tuning risks reintroducing the original dataset bias.

Evaluated on Waterbirds, CelebA (Blond Hair), and UTKFace (Gender vs. Age), reporting worst-group (WGA) and average-group (AGA) accuracy following prior work.

04 / Generating bias-conflicting samples under severe bias

A DreamBooth model trained on Blond Females generates Blond Males by swapping the gender token in the prompt and removing the learned [V] identifier token, with a negative prompt suppressing the female attribute. — At a 0.999 bias ratio, a training group like Blond Male has almost no real samples. Removing the learned “[V]” identifier from the prompt of a model trained on the bias-aligned group (Blond Females) — while swapping the attribute word and adding a negative prompt — transfers the group's visual style onto the target group.

When we push the bias ratio to 0.999 (99.9% of training images from bias-aligned groups), diffusion models must generate the minority groups largely from scratch. For Vanilla SD and FFR this just means prompting with the bias and class labels. For LoRA, DreamBooth, and Clustered DreamBooth — which rely on training-group images — bias-conflicting samples (e.g. Blond Males in CelebA) are instead generated from the model trained on the corresponding bias-aligned group (Blond Females).

On manual inspection, dropping the learned identifier token during this cross-group generation produces images that visually follow the source group's distribution while depicting the target group — a simple, training-free way to reach severely under-represented groups.

05 / Accuracy vs. time complexity

Averaged across all three datasets and both bias ratios: each fine-tuning step buys measurable worst-group accuracy, at the cost of training more per-group diffusion models. Clustered DreamBooth's cluster count k_D can be dialed down to trade fidelity for compute — k_D=1 recovers plain DreamBooth.

Method

Time complexity

WGA↑

AGA↑

Vanilla SD

O(1)

71.3

82.6

LoRA-finetuning

O(|G_D|)

72.8

85.4

DreamBooth

O(|G_D|)

76.2

84.8

Clustered DreamBooth

O(|G_D|·k_D)

79.1

86.2

06 / Quantitative results

Worst-group (WGA) and average-group (AGA) accuracy on the original dataset and its 0.999 bias-ratio variant. † uses the original authors' codebase; generative methods use no bias annotations besides group balancing.

Method

Orig. WGA

Orig. AGA

0.999 WGA

0.999 AGA

Clustered DreamBooth achieves the highest worst-group accuracy averaged across all datasets (79.1%). At the 0.999 bias ratio, GDRO and SELF collapse (e.g. 23.5% and 25.5% WGA on Waterbirds) while generative pretraining stays within ∼7% of its original-dataset performance.

07 / Qualitative analysis

Grid of real training images (red border) and Clustered DreamBooth generations (yellow border) for every group across Waterbirds, CelebA, and UTKFace. — Real training images (red border) alongside Clustered DreamBooth generations (yellow border) for every group across all three datasets. Generated images closely track both the group attribute and the domain-specific visual style of each dataset.

DreamBooth generations of UTKFace female children are visually narrow; Clustered DreamBooth generations span a wider range of demographics, backgrounds, and settings. — DreamBooth vs. Clustered DreamBooth for UTKFace “Female Children”. A single DreamBooth model narrows toward one demographic and setting; clustering first lets each sub-model specialize, so the pool covers far more of the group's real diversity.

This diversity gap explains why plain DreamBooth underperforms on UTKFace (Table 1): with a wide within-group age and background range, a single model averages toward a narrow mode instead of representing the full group.

Clustered DreamBooth's advantage is largest exactly where intra-group variation is highest — facial datasets with varied demographics, lighting, and pose — and smallest on Waterbirds, where DreamBooth already attains the top worst-group accuracy on the original dataset.

08 / Design choices & ablations

CLIP-Score weighting α

α=1 uses only image–text label similarity; α=0 uses only similarity to the group centroid. α=0.5 (used in the main results) balances both and performs best on most datasets, beating unscored random sampling.

Selection

WB WGA

CelebA WGA

α = 1

87.0

83.8

α = 0

84.7

83.3

α = 0.5 (used)

88.1

84.1

Random sampling

84.9

81.1

Generative pretraining + GDRO fine-tuning

At a 0.999 bias ratio, fine-tuning a Clustered-DreamBooth-pretrained model with the GDRO loss (instead of plain last-layer retraining) lifts worst-group accuracy by +58.3 points on Waterbirds and +37.2 on CelebA over GDRO alone — showing the two approaches are complementary.

Method

Waterbirds WGA

CelebA WGA

GDRO alone

23.5

27.2

Clustered DreamBooth + GDRO

81.8

64.4

09 / Conclusion

We investigated Stable Diffusion fine-tuning — LoRA and DreamBooth — for generating representative, group-balanced training images, and introduced Clustered DreamBooth to address intra-group diversity. Pretraining a classifier on this synthetic data, then fine-tuning only its softmax layer on real data, consistently outperforms vanilla SD and FFR, and matches or exceeds SOTA debiasing methods like Group-DRO — especially as dataset bias becomes severe. Optimizing the cluster count and combining generative pretraining with debiasing losses like GDRO remain open directions.

10 / BibTeX

@inproceedings{basu2026harnessing,
  title     = {Harnessing Diffusion-Generated Synthetic Images for Fair Image Classification},
  author    = {Basu, Abhipsa and Gupta, Aviral and Bhat, Abhijnya and Babu, R. Venkatesh},
  booktitle = {Proceedings of the AAAI Conference on Artificial Intelligence},
  year      = {2026},
  note      = {arXiv:2511.08711}
}