Mitigating Spurious Correlations in Image Datasets via Generative Counterfactual Image Augmentation

Abstract

Machine learning models often suffer from performance degradation and reduced generalizability when faced with out-of-distribution data and spurious correlations, as they tend to rely on biased shortcut features rather than learning the core characteristics. Although various solutions in the form of data manipulation, representation learning, and learning strategy have been proposed to address this challenge, a major limitation of many successful methods is their heavy reliance on group annotations for data balancing, a requirement that is costly and sometimes infeasible in real-world applications. To overcome this limitation, the present study proposes a novel method, termed Generative Counterfactual Image Augmentation (GCIA), which improves upon the DFR framework. By leveraging the capabilities of pre-trained image editing models, the proposed method generates counterfactual image data to implicitly neutralize spurious features at the dataset level. The key advantage of this approach is a significant reduction in the need for group annotations; reliance on these labels is eliminated during the balancing and training phases, and they are required solely for hyperparameter tuning. The performance of the proposed method was evaluated on standard datasets, including Waterbirds, CelebA, and UrbanCars. Experimental results demonstrate that, despite a severe reduction in access to group annotations, this method achieves highly competitive performance. It exhibits substantial superiority in improving Worst-Group Accuracy compared to baseline models like ERM (for instance, increasing this accuracy from 21.4% to 79.5% in the UrbanCars dataset and from 46.6% to 85.6% in CelebA). Ultimately, this research demonstrates that employing generative models for semantic data balancing is an efficient and cost-effective strategy to mitigate model bias, effectively enhancing the robustness of computer vision systems against spurious correlations.

Training

1. ERM Training:

dataset="urbancars"
seed=1
uv run main.py \
    --output_path logs \
    --experiment ERM \
    --dataset "$dataset" \
    --dataset_path "datasets/${dataset}" \
    --optimizer SGD \
    --learning_rate 1e-3 \
    --weight_decay 1e-4 \
    --epochs 300 \
    --pretrained_path imagenet \
    --batch_size 128 \
    --seed "$seed"

2. Edit images:

dataset="urbancars"
editing_model="flux2-klein-4b"
uv run edit_images.py \
    --dataset "${dataset}" \
    --dataset_path datasets/${dataset} \
    --output_dir "datasets/${dataset}_edited_${editing_model}" \
    --batch_size 16 \
    --edit_model "${editing_model}"

3. GCIA:

DATASET="celeba"
EDITING_MODEL="flux2_klein_4b"
SAMPLE_SIZE=128
BATCH_SIZE=128
LR=0.005
OPTIMIZER="adamW"
L1=1e-2
SEED=1
uv run main.py \
    --output_path logs \
    --dataset "$DATASET" \
    --dataset_path "datasets/${DATASET}" \
    --experiment GCIA \
    --sample_size "$SAMPLE_SIZE" \
    --batch_size "$BATCH_SIZE" \
    --learning_rate "$LR" \
    --pretrained_path "checkpoints/${DATASET}/erm_seed${SEED}/final_checkpoint.pt" \
    --l1 "$L1" \
    --epochs 100 \
    --optimizer "$OPTIMIZER" \
    --seed "$SEED" \
    --balanced_dataset_path "datasets/${DATASET}_edited_${EDITING_MODEL}"

Note: If the --feature_only flag is used, you should provide the pre-computed features of the specified dataset, which can be saved using the save_features.py file in the repository. If the flag is not specified, the raw image of the dataset should be provided. Here is an example script:

• For main datasets:

dataset="urbancars"
seed=1
uv run save_features.py \
    --dataset "$dataset" \
    --dataset_path "datasets/${dataset}" \
    --save_path "datasets_features/${dataset}/noaug_features_seed${seed}" \
    --pretrained_path "checkpoints/${dataset}/erm_seed${seed}/final_checkpoint.pt" \
    --batch_size 512 

• For Generative Counterfactual Image Augmented Datasets:

dataset="urbancars"
sample_size=128
editing_model="flux2_klein_4b"
seed=1
uv run save_features.py \
    --dataset "$dataset" \
    --dataset_path "datasets/${dataset}_edited_${editing_model}" \
    --save_path "datasets_features/${dataset}_edited_${editing_model}_samples${sample_size}/noaug_features_seed${seed}" \
    --pretrained_path "checkpoints/${dataset}/erm_seed${seed}/final_checkpoint.pt" \
    --batch_size 512 \
    --gcia True \
    --sample_size "$sample_size"

• GCIA with --feature_only used:

DATASET="celeba"
EDITING_MODEL="flux2_klein_4b"
SAMPLE_SIZE=128
BATCH_SIZE=128
LR=0.005
OPTIMIZER="adamW"
L1=1e-2
SEED=1
uv run main.py \
    --output_path logs \
    --dataset "$DATASET" \
    --dataset_path "datasets_features/${DATASET}/noaug_features_seed${SEED}" \
    --experiment GCIA \
    --sample_size "$SAMPLE_SIZE" \
    --batch_size "$BATCH_SIZE" \
    --learning_rate "$LR" \
    --pretrained_path "checkpoints/${DATASET}/erm_seed${SEED}/final_checkpoint.pt" \
    --l1 "$L1" \
    --epochs 100 \
    --optimizer "$OPTIMIZER" \
    --seed "$SEED" \
    --balanced_dataset_path "datasets_features/${DATASET}_edited_${EDITING_MODEL}_samples${SAMPLE_SIZE}/noaug_features_seed${SEED}" \
    --feature_only True

Results

Method	Group Info Train/Val	WaterBirds Worst (%)	WaterBirds Mean (%)	CelebA Worst (%)	CelebA Mean (%)	UrbanCars Worst (%)	UrbanCars Mean (%)
ERM	✗/✗✗	66.2±1.8	90.2±0.6	46.6±3.0	95.5±0.0	21.4±2.2	74.1±0.5
SUBG	✓/✓✓	89.1±1.1	-	85.6±2.3	-	-	-
GDRO	✓/✓✓	91.4	93.5	88.9	92.9	73.1	84.2
JTT	✗/✗✓	86.7	93.3	81.1	88	79.5	86.3
GDRO + EIIL	✗/✗✓	77.2±1.0	96.5±0.2	81.7±0.8	85.7±0.1	76.5±2.6	85.4±2.1
SELF	✗/✗✓	91.6±1.4	93.6±1.1	83.9±0.9	91.7±0.4	83.2±0.8	90.0±0.5
AFR	✗/✗✓	90.4±1.1	94.2±1.2	82.0±0.5	91.3±0.3	80.2±2.0	87.1±1.2
EVaLS-GL	✗/✗✓	89.4±0.3	95.1±0.3	84.6±1.6	91.1±0.6	83.5±1.7	88.3±0.9
GCIA (Ours)	✗/✗✓	90.3±0.4	95.1±0.3	85.6±2.1	90.8±0.7	79.5±1.9	87.5±0.4
DFR	✗/✓✓	92.9±0.2	94.2±0.4	88.3±1.1	91.3±0.3	79.6±2.2	87.5±0.6
SSA	✗/✓✓	89.0±0.6	92.2±0.9	89.8±1.3	92.8±0.1	-	-

XGrad-CAM

Waterbirds

CelebA

UrbanCars

Datasets

• CelebA
• Waterbirds
• UrbanCars

ERM Checkpoints

The ERM checkpoints for CelebA, Waterbirds, and UrbanCars are available here.

Generative Counterfactual Image Augmented Datasets

• CelebA
  • FLUX.2 klein 4B
  • FLUX.2 klein 9B
• Waterbirds
  • FLUX.2 klein 4B
  • FLUX.2 klein 9B
• UrbanCars
  • FLUX.2 klein 4B
  • FLUX.2 klein 9B

Acknowledgments

We would like to thank the authors of the following papers and repositories for their valuable contributions:

• EVaLS
• Deep Feature Reweighting
• Spurious Feature Learning
• WILDS
• DRO