PaperlayerView exact revision on arXiv

Source-linked AI summary

SAGE: Subpopulation-Aware Generative Enhancement for Mitigating Spurious Correlations

Yiming Luo, Rongqiang Zhao, Jie Liu

arXiv:2609.01051v1cs.LGcs.CV

TL;DR

Spurious correlations can cause ERM classifiers to rely on majority shortcuts, especially when spurious attributes and balanced validation data are unavailable. SAGE addresses this gap with clustered sub-labels, targeted generative augmentation, and synthetic validation for last-layer reweighting, achieving strong worst-group accuracy across Waterbirds, CelebA, and MetaShift.

  • Problem

    Spurious correlations and dataset imbalance can make ERM models rely on majority spurious attributes, while group labels and balanced validation data may be unavailable.

  • Method

    SAGE clusters semantic features into sub-labels, fine-tunes a conditional generator with class labels, targets underrepresented regions through inverse-density sampling, and builds balanced synthetic validation data.

  • Results

    89.5%, 85.7%, and 79.1% worst-group accuracy are achieved on Waterbirds, CelebA, and MetaShift, outperforming prior methods without group labels or curated validation data.

  • Takeaways & Limitations

    SAGE provides a data-centric generative framework for mitigating spurious correlations without explicit spurious-group annotations, external discovery tools, or curated validation sets.

  • Takeaways & Limitations

    Inverse-density sampling targets empirically rare clusters rather than verified worst-case groups, and some low-density clusters may reflect non-spurious variation such as pose or lighting.

Abstract

from arXiv · show

Spurious correlations pose a significant challenge to the robustness of modern machine learning. The inherent imbalance in dataset distributions often leads traditional Empirical Risk Minimization (ERM) models to rely on majority spurious attributes for classification, resulting in poor performance on minority groups. This problem becomes particularly challenging when the spurious attributes are unavailable. Existing group-label-free methods often upsample minority groups or misclassified real training examples; repeating the same instances can reduce effective diversity and encourage overfitting. To mitigate these spurious correlations from a data-centric perspective in the absence of prior knowledge, we introduce Subpopulation-Aware Generative Enhancement (SAGE), a two-stage generative augmentation framework. Using cluster-derived sub-labels and class labels, we fine-tune a conditional generative model and text encoder, generating targeted synthetic data to fill underrepresented regions in the training set and construct a balanced validation set for last-layer reweighting. We experimentally show that SAGE achieves 89.5%, 85.7%, and 79.1% worst-group accuracy on Waterbirds, CelebA, and MetaShift, respectively, outperforming the best group-label-free baselines by up to 7.7 percentage points.

Introduction

SAGE addresses spurious-correlation mitigation when group labels and balanced validation data are unavailable by deriving subpopulation proxies from feature clusters and using generative augmentation. It combines targeted training synthesis with balanced synthetic validation, achieving strong worst-group accuracy across three benchmarks.

  • Motivation: Spurious correlations can make classifiers rely on shortcut features, causing catastrophic test-performance drops when those correlations break.Waterbirds illustrate this failure when background rather than the bird becomes the primary classification basis.
  • Motivation: Existing approaches either use spurious group information or infer proxy signals from model behavior, input regions, or discovered concepts.These categories modify objectives or reweighting strategies, or avoid group labels through inferred signals.
  • Research gap: Both spurious group labels and balanced validation data are often unavailable, while existing methods in this stricter setting remain less competitive in worst-group accuracy.Generative alternatives may additionally require external discovery tools, explicit annotations, or manually curated validation sets.
  • Approach: SAGE treats feature-cluster assignments as proxies for bias-relevant subpopulations rather than ground-truth spurious groups.The approach relies on the observation that samples sharing shortcut features tend to cluster in pretrained-encoder feature space.
  • Approach: SAGE uses inverse-density sampling for targeted synthetic training data and uniform sampling for a sub-label-balanced validation set used in last-layer reweighting.Cluster-derived sub-labels and class labels fine-tune a conditional generative model before these two sampling strategies are applied.
  • Results: 89.5%, 85.7%, and 79.1% worst-group accuracy are achieved on Waterbirds, CelebA, and MetaShift, respectively, outperforming prior methods without group labels or curated validation data.The contributions also include analyses of complementary training and validation components and subpopulation-specific visual factors.

Related Work

Prior work mitigates spurious correlations through group-aware training or model selection, group-label-free proxy signals, and synthetic data generation. These approaches differ in how they obtain robustness-relevant information and construct augmented data.

  • Robust Learning for Spurious Correlations: Group-aware methods modify training objectives or model selection using spurious group information.Examples include gDRO, which upweights high-loss groups, and DFR, which retrains the last linear layer on a group-balanced validation set.
  • Robust Learning for Spurious Correlations: Group-label-free methods use proxy signals from model behavior, input regions, or discovered concepts to reduce bias.Examples include JTT, LfF, CnC, and DaC, which identify hard examples, exploit biased-example signals, use contrastive learning, or decompose training.
  • Generative Data Augmentation for Spurious Correlations: Generative augmentation methods construct synthetic samples for robust classification, including approaches that balance synthetic data before real-data training.Diffusion models support controllable image generation and have motivated synthetic-data methods for spurious-correlation mitigation.

Method

SAGE is a two-stage pipeline that discovers subpopulation structure through feature clustering, embeds it in a conditional diffusion generator, and samples synthetic training and validation data. Its sampling design merges redundant tokens to preserve diversity while targeting rare regions and balancing validation coverage.

  • Stage 1: Semantic Clustering and Conditional Generator Fine-Tuning: SAGE first clusters training images in pretrained-encoder feature space using Affinity Propagation without explicit spurious-attribute labels.Affinity Propagation identifies representative exemplars without requiring a predefined number of clusters.
  • Stage 1: Semantic Clustering and Conditional Generator Fine-Tuning: LoRA updates the diffusion model’s UNet while the remaining text-encoder and UNet weights stay frozen.The new token embeddings and LoRA parameters are jointly optimized during conditional generation fine-tuning.
  • Stage 1: Semantic Clustering and Conditional Generator Fine-Tuning: Each cluster receives a learnable sub-label token that is combined with the class label to fine-tune a conditional diffusion generator.The sub-label token is optimized to capture fine-grained subpopulation characteristics while the class label provides the global semantic anchor.
  • Stage 2: Targeted Data Generation and Downstream Model Training: Inverse-density sampling generates auxiliary training data for underrepresented condition regions after the conditional generator is optimized.The sampled synthetic set is combined with the original dataset to form the mixed training set used for downstream classifier training.
  • Stage 2: Targeted Data Generation and Downstream Model Training: Merging nearly redundant sub-label tokens prevents naive inverse-density sampling from reducing generation diversity.Token pairs with sufficiently high cosine similarity are merged transitively within each class, while training probabilities are normalized globally over all sub-labels.
  • Stage 2: Targeted Data Generation and Downstream Model Training: Inverse-density sampling upweights empirically rare clusters rather than verified worst-case class-attribute groups.Some low-density clusters may instead reflect non-spurious variation such as pose or lighting, and global normalization accounts partly for long-tailed structure.
  • Stage 2: Targeted Data Generation and Downstream Model Training: Uniform sampling creates a deterministic, globally balanced synthetic validation set across merged groups for Deep Feature Reweighting.The downstream classifier is trained with ERM, after which DFR retrains only the last classification layer using the synthetic validation set.

Experiment

Experiments evaluate SAGE across three benchmarks and compare its group-label-free performance, generative augmentation, targeted synthesis, and synthetic validation components.

  • Results: 0.8, 7.7, and 5.6 percentage points separate SAGE from the best no-group-label baseline on Waterbirds, CelebA, and MetaShift, respectively.The comparison uses worst-group accuracy under the setting without group labels in training or validation.
  • Generative comparisons: SAGE outperforms ASPIRE+ERM on Waterbirds and CelebA and surpasses ASPIRE+DFR without using a manually curated validation set.It slightly trails DDB, which relies on an external segmentation model.
  • Ablations and visualization: Changing sub-label tokens alters generated images while class labels remain fixed, indicating control over subpopulation-specific visual variation.Rare conflicted Waterbirds clusters with 25 training samples each are upweighted by inverse-density sampling.
  • Ablations and visualization: Targeted Dmix augmentation yields substantially larger gains than generic synthetic augmentation, which provides limited benefit or can be harmful.The ablation attributes improvement to targeted augmentation of underrepresented subpopulations rather than generation alone.
  • Ablations and visualization: Synthetic validation is competitive with manual validation, but performance depends on benchmark structure: manual validation is strongest on Waterbirds, while synthetic validation can match or exceed it on CelebA and MetaShift.Reported comparisons include 89.7% versus at most 88.6% on Waterbirds, up to 82.6% versus 62.8% on CelebA, and 73.4% versus 69.5% on MetaShift.

Conclusion

SAGE is a two-stage, data-centric method that uses clustered sub-labels to expand minority samples and synthesize a balanced validation set. Experiments across three benchmarks demonstrate its efficacy and indicate potential for datasets with subpopulation shifts.

  • Conclusion: SAGE uses clustered sub-labels to fine-tune conditional generative models and guide generation in a two-stage pipeline.The method expands minority samples and synthesizes a balanced validation set.
  • Conclusion: Experiments across three distinct benchmark datasets demonstrate the efficacy of the SAGE framework.The conclusion describes SAGE as independent of explicit prior knowledge and specific dataset distribution characteristics.
  • Conclusion: SAGE may generalize to datasets with subpopulation shifts and broader general-purpose datasets because it does not rely on explicit prior knowledge or specific distribution characteristics.This is presented as potential rather than an experimentally established result.

Appendix A: Dataset Statistics

Appendix A documents benchmark split sizes, group distributions, and the separation of training, validation, and test data used in the experiments.

  • Dataset statistics: 4,795, 162,770, and 1,024 images are used in the Waterbirds, CelebA, and MetaShift training splits, respectively.Clustering and generative fine-tuning use training data only; benchmark validation splits are not used in the method pipeline.
  • Dataset statistics: Tables 5–7 report the number of images in each (y, a) group across training, validation, and test splits.The appendix defines these as group distributions by split.
  • Dataset statistics: Waterbirds has 4,795 training, 1,199 validation, and 5,794 test images, with approximately group-balanced validation and test splits.These counts are reported in Table 5.
  • Dataset statistics: CelebA has 162,770 training, 19,867 validation, and 19,962 test images.These counts are reported in Table 6.
  • Dataset statistics: MetaShift has 1,024 training, 81 validation, and 460 test images, with shelf background appearing only in validation and test.The split-specific counts and background scope are reported in Table 7.

Appendix B: Implementation Details

Appendix B specifies the shared experimental configuration, generator fine-tuning setup, token-merging procedure, validation synthesis, and hardware allocation used across the three datasets.

  • Overview: The appendix provides reproducibility-oriented implementation details, using the same configuration across Waterbirds, CelebA, and MetaShift unless otherwise stated.It covers dataset-related appendix details, clustering hyperparameters, and post-hoc purity analysis.
  • Conditional generator fine-tuning: SAGE fine-tunes Stable Diffusion v1.5 with LoRA while jointly optimizing newly added sub-label token embeddings.Each cluster receives a token initialized from “photo,” and training uses class-label and sub-label prompts.
  • Conditional generator fine-tuning: The implementation uses LoRA rank r = 128, α = 128, learning rate 1e-5, AdamW, batch size 4 per GPU, 100 epochs, and 512-pixel images.Trainable modules are UNet attention-layer LoRA parameters and sub-label token embeddings.
  • Token similarity merging: Redundant sub-label tokens are merged using within-class pairwise cosine similarities and a dataset-specific multiplier β.The multipliers are β = 2.5 for Waterbirds and β = 4.0 for CelebA and MetaShift.
  • Validation-set synthesis: Validation synthesis replaces stochastic inverse-density sampling with a deterministic samples-per-token parameter to reduce distributional variance.This choice applies when constructing a small-scale validation set.
  • Hardware: Stable Diffusion fine-tuning uses two NVIDIA RTX 4090 GPUs, while the remaining pipeline stages run on one RTX 4090 GPU.The remaining stages include feature extraction, clustering, image generation, ResNet-50 training, and DFR.

Appendix C: Clustering Analysis

Appendix C analyzes cluster construction, purity, and scalability. It reports strong alignment with binary spurious attributes on Waterbirds and CelebA, weaker alignment on MetaShift, and a scalable CelebA procedure.

  • Clustering strategy: Affinity Propagation clusters CLIP embeddings independently within each class, using damping and preference hyperparameters rather than a pre-specified cluster count.The preference values are set from mean pairwise similarities scaled by αpref.
  • Purity analysis: Ground-truth spurious labels are used only for post-hoc purity analysis, not during SAGE training or generation.Purity measures how closely discovered clusters align with spurious attributes.
  • Purity analysis: The overall purity reported in Table 8 is a sample-weighted micro average over clusters.Table 8 reports sample counts, resulting clusters, and clustering purity for each dataset.
  • Purity results: Purity exceeds 90% on Waterbirds and CelebA, while MetaShift reaches 72.49%.The appendix attributes the weaker MetaShift alignment to four spurious attributes and stronger interclass visual differences beyond background cues.
  • CelebA scalability: CelebA clustering reduces the infeasible N × N similarity matrix by selecting K = 10,000 diverse samples, clustering them, and assigning all data to nearest exemplars.The resulting memory cost is O(N × D + K^2), with K ≪ N.

Appendix D: Effectiveness Analysis of DFR

Appendix D visualizes DFR using a shared two-dimensional PCA representation of generated validation and held-out test features. The analysis shows how last-layer reweighting shifts the classifier boundary toward rarer class-attribute combinations.

  • Visualization setup: Figure 3 compares generated reweighting samples with held-out test samples under the same frozen feature representation.The left column shows reweighting samples, and the right column shows held-out test samples.
  • Feature-space construction: PCA is fit on the union of frozen features from the visualization test subset and generated DFR validation samples.Additional generated reweighting samples are projected afterward using the fixed mean and basis.
  • Visualization setup: For each dataset, the visualization uses a balanced held-out subset with 180 images from each (y, a) group.This preserves class-attribute structure while keeping the visualization readable.
  • Decision boundaries: The ERM and DFR boundaries are compared after projecting the binary final linear layer onto the PCA plane.The PCA approximation substitutes the feature representation into the original decision-boundary equation.
  • Downstream classification analysis: DFR updates only the last linear layer, shifting the boundary toward a classifier less dependent on majority-dominated correlations.The generated validation set provides coverage for rare or under-represented class-attribute combinations.
  • Benchmark encoding: Figure 3 identifies Waterbirds minority groups G2 and G3, CelebA minority group G4, and the relevant target and spurious features for each benchmark.The figure uses blue and orange points for target classes and dashed green and red lines for DFR and ERM boundaries.
Loading 2609.01051v1…