Class-Conditional Distribution Balancing for Group Robust Classification

TL;DR

This work reframes spurious correlations as imbalances in class-conditional distributions and introduces Class-Conditional Distribution Balancing (CCDB), a bias-annotation-free method that maximizes the uncertainty of the label given spurious factors. Through a three-stage, class-aware sample-reweighting pipeline, CCDB balances class-conditional statistics via $W_2$ distances and per-class weights, effectively diminishing spurious cues. Extensive experiments on synthetic multi-class tasks and real-world benchmarks show CCDB achieving state-of-the-art or competitive performance without bias supervision, often rivaling bias-guided approaches. The approach leverages causal-inference notions of covariate balance and demonstrates practical utility across diverse domains, with detailed ablations and analysis of learned sample weights.

Abstract

Spurious correlations that lead models to correct predictions for the wrong reasons pose a critical challenge for robust real-world generalization. Existing research attributes this issue to group imbalance and addresses it by maximizing group-balanced or worst-group accuracy, which heavily relies on expensive bias annotations. A compromise approach involves predicting bias information using extensively pretrained foundation models, which requires large-scale data and becomes impractical for resource-limited rare domains. To address these challenges, we offer a novel perspective by reframing the spurious correlations as imbalances or mismatches in class-conditional distributions, and propose a simple yet effective robust learning method that eliminates the need for both bias annotations and predictions. With the goal of maximizing the conditional entropy (uncertainty) of the label given spurious factors, our method leverages a sample reweighting strategy to achieve class-conditional distribution balancing, which automatically highlights minority groups and classes, effectively dismantling spurious correlations and producing a debiased data distribution for classification. Extensive experiments and analysis demonstrate that our approach consistently delivers state-of-the-art performance, rivaling methods that rely on bias supervision.

Figures (10)

• Figure 1: Left: Causal graph. Right: Analyzing the effectiveness of our method under three extreme configurations where $\boldsymbol{z} _{\text{s} }$ and $y$ are: (a) strongly correlated with well-disentangled predefined features; (b) completely independent; and (c) strongly correlated with predefined features that are independent but semantically entangled. Each case consists of two classes (marked in blue and red). Data points in the yellow regions are assigned with high weights, while those in the blue regions are assigned with low weights.
• Figure 2: t-SNE visualization of the latent space learned in Stage1. Colored by bias annotations. As $\gamma$ decreases, the feature space exhibits an increasingly pronounced clustering effect aligned with bias annotations.
• Figure 3: Distribution of the learned sample weights within each group. The distributions are normalized within each group.
• Figure 4: Left: Ablation study on the split ratio $\gamma$ and threshold $T$. Right: Boxplot of the feature correlations with bias and class information, before and after CCDB sample reweighting. Refer to appendix for more results.
• Figure 5: Visualization of ERM and CCDB attention on Waterbirds and CelebA samples.