Improving Group Robustness on Spurious Correlation via Evidential Alignment (Extended Abstract)

1 Introduction

Deep neural networks trained with empirical risk minimization (ERM) [Vapnik, 1999] often exhibit spurious biases, i.e., reliance on brittle, non-causal features such as background, texture, or secondary objects [Stock and Cisse, 2018; Beery et al., 2018; Sagawa et al., 2020; Ye et al., 2024]. For example, in a cow/camel classification task, the camel class can be spuriously correlated with desert backgrounds, leading ERM-trained models to fail on minority groups (e.g., camels on grass) [Beery et al., 2018]. Such failures are particularly problematic in high-stakes domains including medicine [Zech et al., 2018], fairness [Khani and Liang, 2021], and criminal justice [Barocas et al., 2023].

A model's robustness to such failures is typically measured by its worst-group accuracy (WGA) [Sagawa et al., 2020].

∗ Full version appears in the Proceedings of KDD 2025 [Ye et al., 2025], recipient of the Best Paper Award in Research.

Supervised approaches such as Group DRO [Sagawa et al., 2020], DFR [Kirichenko et al., 2023], and JTT [Liu et al., 2021] require explicit group labels that pair class labels with spurious attributes, but such annotations are costly and often miss subtle biases. Recent annotation-free methods [Nam et al., 2020; Li et al., 2024] sidestep this requirement by using deterministic auxiliary models to amplify or pseudo-label biased samples. However, deterministic auxiliaries provide only a narrow view of how a model exploits spurious features, limiting their effectiveness.

We propose Evidential Alignment, a framework that leverages epistemic uncertainty to provide a principled, annotation-free signal for identifying and mitigating spurious biases. Drawing on the Dempster-Shafer Theory of Evidence [Yager and Liu, 2010] and recent advances in evidential deep learning [Sensoy et al., 2018], we transform a biased ERM model from first-order risk minimization to secondorder risk minimization by learning a distribution over distributions via a Dirichlet parameterization. The resulting epistemic uncertainty is then used to reweight calibration samples and retrain the classifier, suppressing spurious correlations while preserving core features.

The main contributions of this paper are:

• We propose Evidential Alignment, a novel framework that integrates second-order risk minimization with evidential calibration to leverage model uncertainty for identifying and mitigating spurious correlations without group annotations.
• We provide theoretical insights showing how uncertainty quantification captures spurious biases in latent space.
• Extensive experiments across diverse architectures and modalities demonstrate that Evidential Alignment significantly enhances group robustness and generalization. 1

2 RelatedWork

Group robustness. Methods for improving group robustness against spurious correlations differ mainly in how much group supervision they assume. With full training-set group annotations, Group DRO [Sagawa et al., 2020] minimizes

1 Code: https://github.com/wenqian-ye/evidential alignment.

(a) Second-order Risk Minimization

(b) Evidential Calibration

the worst-case loss across predefined groups. With only validation-set annotations, DFR [Kirichenko et al., 2023] retrains the final layer on a group-balanced held-out set, and JTT [Liu et al., 2021] as well as AFR [Qiu et al., 2023] upweight samples that a biased ERM model misclassifies. Annotation-free approaches such as LfF [Nam et al., 2020] and BAM [Li et al., 2024] train a deliberately biased auxiliary model and use its errors as a proxy signal. SELF [LaBonte et al., 2024] retrains the last layer with held-out data but without group labels. All of these annotation-free strategies rely on deterministic bias estimators, which capture only a narrow slice of how a model exploits spurious features.

Uncertainty quantification. Predictive uncertainty is typically decomposed into aleatoric (data) and epistemic (model) components [Kendall and Gal, 2017; H¨ullermeier and Waegeman, 2021]. Bayesian methods, MC dropout, and deep ensembles provide such estimates but require sampling or multiple model passes. Evidential deep learning [Sensoy et al., 2018], grounded in the Dempster-Shafer theory of evidence [Yager and Liu, 2010], instead places a Dirichlet distribution over class probabilities in a single forward pass. Uncertainty has been used for calibration, active learning, and out-of-distribution detection, but its potential for diagnosing and correcting spurious bias has remained largely unexplored. Our method closes this gap.

3 Method

3.1 Problem Setup

We consider a classification problem with training set Dtrain = {(x, y)} partitioned into latent groups Dtr g , where each group g := (y, a) pairs a class y ∈ Y with a spurious attribute a ∈ A. Let fθ : X → Y be a model with parameters θ = (ϕθ, hθ) decomposed into a feature extractor ϕθ and a classifier hθ. Standard ERM minimizes the average training loss; group robustness instead targets the worst-group risk

Uncertainty decomposition. The posterior predictive distribution can be split into aleatoric (data) and epistemic (model) parts by marginalizing over θ:

For neural networks this integral is intractable. Following [Malinin and Gales, 2018; Ulmer et al., 2021; H¨ullermeier and Waegeman, 2021], we factorize it through an intermediate distribution π over class probabilities and use a point estimate ˆ θ, yielding p(y | x) = ∫ P(y | π) p(π | x, ˆ θ) dπ, which has a closed-form solution when p(π | x, ˆ θ) is a Dirichlet distribution. We therefore focus on the epistemic component, which can be efficiently captured by a Dirichlet over class probabilities [Sensoy et al., 2018; H¨ullermeier and Waegeman, 2021] and provides the signal we use to debias fθ without group labels.

3.2 Second-Order RiskMinimization

Conventional ERM optimizes a deterministic predictive distribution p(y | θ, x), conflating aleatoric and epistemic uncertainty. We instead model a distribution over distributions p(θ | π(x)), where π(x) parameterizes uncertainty about the prediction outcomes. We extend the ERM model to output non-negative evidence ek(x) ≥ 0 for each class k ∈ {1, . . . ,K}, defining the Dirichlet parameters

The Dirichlet distribution over class probabilities p(x) is

with expected class probability E[pk(x)] = αk(x)/S(x), where S(x) = ∑k αk(x).

The second-order risk minimization loss for a sample (x, y) combines a classification objective with an evidence regularizer:

where Dir(1) is a non-informative uniform Dirichlet prior. The classification objective is the negative log expected probability of the true class,

and the KL divergence between the predicted Dirichlet and the uniform prior has the closed form

where ψ(·) = d dz lnΓ(·) is the digamma function. The KL term penalizes overconfidence and promotes calibrated uncertainty. We adopt a dynamic annealing coefficient λt = min(t/η, 1) to stabilize training, where t is the current epoch and η is the annealing-step hyperparameter.

3.3 Evidential Calibration

After second-order training, the learned Dirichlet yields a persample epistemic uncertainty

On a held-out calibration set Dcalib, we upweight misclassified samples in proportion to their uncertainty,

and retrain only the last layer with a class-balanced reweighted cross-entropy plus a proximal regularizer:

Last-layer retraining suffices to learn invariant features, following [Kirichenko et al., 2023; LaBonte et al., 2024]. The two stages are illustrated in Figure 1.

3.4 Theoretical Analysis

We establish two guarantees; full statements and detailed proofs are in the original paper.

Theorem 1 (ELBO for Second-Order Risk). Let p(π | y) be the true posterior over class probabilities given y, and let p(π | x, θ) be the variational evidential distribution. Training with L1 optimizes the evidence lower bound (ELBO) on log p(y), so second-order risk minimization is equivalent to maximizing the log marginal likelihood of observed labels under a variational Dirichlet posterior.

Theorem 2 (PAC-Bayes Bound for Reweighted Empirical Risk). Let H have [0, 1]-bounded loss and G be a finite set of latent groups with ng i.i.d. samples each, nmin = ming ng. For any reweighting {wg} with α = ming wg > 0 and any prior P and posterior Q over H, the worst-group risk satisfies, with probability at least 1 - δ,

Together, Theorem 1 certifies that our second-order objective is principled, since it maximizes a likelihood by treating evidence as the parameters of a variational Dirichlet posterior. Theorem 2 certifies that any reweighting whose induced empirical risk is small bounds the worst-group risk; our uncertainty-based weighting is exactly such a scheme, since high-uncertainty (likely-minority) samples are upweighted, raising α = ming wg and tightening the bound.

4 Experiments

Datasets. We evaluate on four standard group-robustness benchmarks. Waterbirds [Sagawa et al., 2020; Welinder et al., 2010] contains waterbird/landbird images whose label is spuriously correlated with the water/land background; the minority groups are waterbirds on land and landbirds on water. CelebA [Liu et al., 2015] requires predicting hair color (blond vs. non-blond), which is spuriously correlated with gender, with blond males being a tiny minority. MultiNLI [Williams et al., 2017] is a sentence-pair entailment task whose contradiction label is spuriously correlated with the presence of negation words. CivilComments [Borkan et al., 2019; Koh et al., 2021] is a toxicityclassification task whose label spuriously correlates with eight demographic identity mentions; we use the WILDS split.

Setup. For images we use a ResNet-50 [He et al., 2016] backbone pretrained on ImageNet; for text we use BERTbase [Kenton and Toutanova, 2019]. We form Dcalib from one half of the validation split and reserve the other half for model selection; this matches the protocol of DFR† and SELF† , but unlike DFR we use no group labels at any stage. Following [Kirichenko et al., 2023], only the last layer is updated in both stages of our method. Baselines fall into two groups: methods that use validation-set group labels (JTT, AFR, DFR† , SELF† ) and methods that use no group labels (LfF, BAM, Ours† ). We report worst-group accuracy (WGA), average accuracy (Acc.), and their gap, averaged over three random seeds.

Main results. Tables 1 and 2 compare Evidential Alignment with state-of-the-art baselines. Among methods that use no group annotations (bottom blocks), our method achieves the best WGA on all four benchmarks, lifting Waterbirds from 72.6% (ERM) to 92.2% and CelebA from 47.2% to 84.4%. Concretely, on Waterbirds we improve over LfF by 14.2 points and BAM by 3.1 points; on CelebA, the gain over LfF is 7.2 points and over BAM 4.3 points. The trend persists in language: on MultiNLI we surpass the strongest annotation-free baseline (BAM) by 3.7 points, and on CivilComments we set a new annotation-free record at 80.2% WGA. We even match or exceed several methods that rely on group labels for validation: we outperform AFR, JTT, and SELF on most datasets, and remain competitive with DFR despite DFR's access to a group-balanced calibration set. Beyond accuracy, our method also yields a small accuracy gap (e.g. 2.9 on Waterbirds and 6.9 on CivilComments), indicating balanced performance across groups rather than skewed majority versus minority trade-offs. The full paper additionally reports a Colored MNIST synthetic study, results on the safety-critical CheXpert chest-radiograph dataset, and ViT-base experiments on Waterbirds, all confirming the same trend across architectures and modalities. Computationally, Evidential Alignment retrains only the last layer in both stages, so its overhead is on par with DFR and AFR rather than with JTT [Liu et al., 2021], which retrains a full auxiliary model from scratch.

Synthetic study. On a Colored MNIST setup [Vapnik, 1999] where 90% of training samples carry a class-color spurious correlation but only 10% of test samples do, ERM collapses to 3.74% WGA on the minority group while Evidential Alignment reaches 84.58%, with majority-group accuracy Upweighted Samples unchanged. t-SNE embeddings of the test set further show high-uncertainty samples clustering on the minority groups in feature space, confirming that epistemic uncertainty localizes the data the biased model gets wrong.

Ablations and qualitative analysis. The full paper ablates the annealing step η and proximal coefficient β, showing that the KL evidence regularization and class-balanced calibration sampling are individually beneficial and combine to give the highest WGA. As Figure 2 illustrates, GradCAM on Waterbirds reveals that the samples our method upweights, those with high epistemic uncertainty, concentrate on background regions tied to the spurious attribute, whereas low-uncertainty samples attend to the bird itself. Quantitative analysis across all four benchmarks shows non-trivial correlation between per-sample uncertainty and the latent group identity, supporting the central claim that epistemic uncertainty is a reliable, annotation-free proxy for spurious-bias-induced failure modes.

5 Conclusion

We introduced Evidential Alignment, an uncertainty-guided framework that mitigates spurious correlations without requiring group annotations. By transforming ERM into second-order risk minimization with a Dirichlet parameterization and applying an evidential calibration step that reweights samples by their epistemic uncertainty, our method captures and corrects spurious biases that deterministic auxiliaries miss. Theoretical analysis and extensive experiments across vision and language benchmarks confirm that Evidential Alignment delivers state-of-the-art worst-group accuracy with minimal computational overhead, providing a scalable and principled path toward robust deep learning.

Acknowledgments

This work is supported in part by the US National Science Foundation under grants 2217071, 2213700, 2106913, and 2008208. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the National Science Foundation.