Invariance Pair Guidance: Robustness to Spurious Correlations via Corrective Gradients
February 26, 2025
1 Introduction↩︎
The ability to learn representations from data makes neural networks highly applicable to a wide range of tasks. However, standard training assumes that training and test data are i.i.d., an assumption that rarely holds in real-world scenarios. For instance, in autonomous driving, models inevitably encounter unknown situations or shifted distributions [1]. To operate safely in such environments, models must generalize to unseen data by learning robust, causal features rather than relying on dataset-specific correlations [2].
A key barrier to robustness is the model’s tendency to internalize spurious correlations, features that correlate with the label in the training set but are not causally predictive. Examples are widespread: animal detection models often rely on background scenery (e.g., snow vs. grass) [3], [4], and medical diagnostic models have been shown to rely on hospital-specific tags on radiographs rather than pathology [5]. Similarly, substantial disparities in gender classification accuracy arise when models rely on skin color rather than facial features [6]. If these spurious characteristics are not explicitly suppressed, models fail when the spurious correlation breaks in the unseen environment [7].
To mitigate this risk, we aim to separate the defining characteristics of a class from spurious attributes. To this end, we introduce the concept of invariance pairs [8]. While conceptually related to contrastive learning [9], IPG differs fundamentally in efficiency and scope. Instead of structuring the entire latent space via dense, random augmentations, IPG utilizes a sparse set of targeted pairs (often \(<3\%\) of the training data) to efficiently correct specific spurious dependencies via a dynamic gradient update. An invariance pair consists of two samples that share the same semantic content (e.g., the same digit) but differ in a specific spurious attribute (e.g., color) (Fig. [fig:colored95mnist95invariance95pair]). We define invariance as the requirement that the model’s internal representation should remain identical for both elements of the pair. This creates a direct metric for robustness: if the model’s output varies within a pair, it has internalized a spurious dependency.
While identifying spurious attributes is non-trivial, recent advancements in generative modeling provide new possibilities in this process. In many high-stakes domains, experts already possess the domain knowledge of what should be invariant (e.g., a tumor classification should be invariant to the hospital scanner ID). Additionally, there are methods to identify such dependencies [10], [11]. We also utilize generative AI (e.g., instruction-based image editing) to synthesize counterfactual invariance pairs, such as gender-swapped portraits, which are otherwise difficult to obtain. In this work, we assume that a sparse set (e.g., \(3\%\) of the training data) of such pairs is available, either from domain knowledge or generative augmentation, and focus on the optimization challenge: how to efficiently utilize them to correct the model’s behavior.
We introduce Invariance Pair Guidance (IPG), a method that explicitly enforces these invariances during training. Unlike standard regularization, IPG extends stochastic gradient descent with a corrective gradient step, inspired by [12]. This correction minimizes the discrepancy between the learned representations of invariance pairs. Crucially, we use an adaptive scaling mechanism based on an invariance condition: the corrective gradient is dynamically prioritized only when the model demonstrates a violation of the invariance (i.e., high disagreement between pairs). This allows the model to internalize the invariance while optimizing on the primary task.
We evaluate IPG on three datasets representing synthetic, composite, and real-world shifts: ColoredMNIST, Waterbirds-100, and CelebA. Our key contribution is a novel dual-update optimization mechanism that utilizes invariance pairs to improve robustness against shifts in both perfect and partial spurious correlations, even under label noise. Figure 1 illustrates the approach and invariance pairs for each dataset. 1
The remainder of the paper is structured as follows. Section 2 discusses related work. In Section 3, we present IPG. Section 5 evaluates the performance, Section 6 concludes.
2 Related Work↩︎
We categorize approaches for robustness and out-of-distribution generalization into three streams: data manipulation, representation learning, and learning strategies. Our proposed IPG method sits at the intersection, integrating explicit representation guidance (invariance pairs) with an adaptive learning strategy (dual-update scheme).
2.0.0.1 Data Manipulation.
These methods aim to bridge the gap between training and target distributions via augmentation or generation. Standard techniques vary image style [13] or use synthetic data [14]. Perturbation-based strategies, such as adversarial augmentation [15], [16], generate worst-group perturbations to balance the data; for instance, DAIR [15] regularizes models against such variations. Other approaches enrich the data with concept-level information, utilizing pseudo-labels [17] or concept banks [18]. A distinct sub-category aims to subtract spurious information via pre-processing. Shock Graph [19] explicitly converts images into shape-based graph representations, thereby discarding texture and color information entirely. Similarly, Noisy Counterfactual Matching (NCM) [8], identifies a linear spurious subspace from counterfactual pairs and projects the input data orthogonally to remove these directions. Distinction: IPG differs from these approaches in three key ways: (i) unlike augmentation methods (DAIR) that require generating synthetic samples for the entire dataset, IPG guides training using a sparse set of pairs; (ii) unlike concept banks, it avoids dense annotation costs; and (iii) unlike feature removal (Shock Graph/NCM), IPG preserves full input information, teaching the model to dynamically ignore spurious correlations rather than statically deleting them.
2.0.0.2 Representation Learning.
The goal of these methods is to disentangle features or learn domain-invariant representations. Domain-invariant learning techniques, such as IRM [20], enforce an optimal classifier across varying training domains. A related subfield, feature disentanglement, aims to separate latent factors using generative models [21] or causal scaffolding [22]. Further, EPG [23], and GALS [24] offer explicit guidance. EPG penalizes feature attributions (calculated via B-cos or IxG) that fall outside a target bounding box, thereby effectively guiding the model spatially. GALS similarly aligns model attention with language-derived ground truths. DKT [25] generates counterfactual features by transferring semantic statistics between spuriously correlated groups to break dependencies. Distinction: While EPG and GALS require bounding boxes or language encoders to define spatial attention guidance, IPG defines semantic guidance via comparative pairs, making it applicable even when spatial localization is difficult. In contrast to IPG, DKT relies on group labels.
2.0.0.3 Learning Strategies.
Robustness can also be enforced through specialized optimization objectives. GroupDRO [26] minimizes the loss of the worst-performing group, providing strong robustness guarantees if group labels are known. In the absence of such labels, Identify-then-Mitigate approaches like JTT [27] and EvA-E [28] first utilize a biased model to identify misclassified samples, then up-weight them (JTT) or prune high-evidence activations (EvA-E). Similarly, LfF [29] employs a dual-model architecture to simultaneously amplify and counteract bias. Gradient-based methods also play a role; Fishr [30] enforces domain-invariance by matching the variances of gradients across domains, adapting the covariance-alignment principle of CORAL [31]. Distinction: Unlike these methods, IPG operates on a single domain without requiring group labels (vs. GroupDRO), multiple domains (vs. Fishr), or the presence of bias-conflicting samples in the training set (vs. JTT/LfF/EvA-E).
3 IPG: Invariance Pair Guidance↩︎
In this section, we describe our approach to formulating the invariance pairs, the corrective gradient, and the adaptive scaling by the invariance condition.
3.1 Preliminaries and Overview↩︎
We consider a classification task with input space \(\mathcal{X}\) and label space \(\mathcal{Y} = \{1, \dots, K\}\). A training set \(\mathcal{D}_\text{tr}\coloneq \{(\mathbf{x}_j, y_j)\}^n_{j=1}\) is drawn from a joint distribution \(P_{\mathcal{D}_\text{tr}}\). Associated with each instance \(\mathbf{x}_j\) is a spurious attribute \(a_j \in \mathcal{A}\). While \(a_j\) is causally irrelevant to the label \(y_j\), the training distribution typically exhibits a strong statistical correlation between specific values of \(a\) and \(y\). We define a group \(g\) as unique combination of the attribute and label: \(g \coloneq (a, y) \in \mathcal{A} \times \mathcal{Y}\). Due to dataset imbalance, standard models often rely on the spurious correlation between \(a\) and \(y\) rather than causal features. In extreme cases, certain groups may be entirely absent from \(\mathcal{D}_\text{tr}\) (perfect spurious correlation).
Our goal is to learn a model \(f\) that generalizes well to a test distribution \(P_{\mathcal{D}_\text{te}}\) where the correlation between \(a\) or \(y\) is reduced or inverted (\(P_{\mathcal{D}_\text{tr}} \neq P_{\mathcal{D}_\text{te}}\)). To measure robustness across all subpopulations, we target the worst-group accuracy [26]:
\[\text{Acc}_\text{wg}(f) \coloneq \min_{g \in \mathcal{A} \times \mathcal{Y}} \mathbb{E}_{(\mathbf{x}, y) \sim P_{\mathcal{D}_\text{te}^g }}[\mathbb{1}_{y = f(\mathbf{x})}].\] where \(P_{\mathcal{D}_\text{te}^g}\) denotes the distribution of test samples belonging to group \(g\).
3.1.0.1 Assumptions.
Our method relies on a set of invariance pairs \(\mathcal{I}\). We assume that for any pair \((\mathbf{x}_{\text{inv}}, \mathbf{x}_{\text{inv}}') \in \mathcal{I}\), both samples share the same semantic label \(y\) and features, but differ in their spurious attribute (\(a \neq a'\)). These pairs serve as explicit examples of the invariance the model must learn.
3.1.0.2 Method Overview.
We propose an extension to stochastic mini-batch gradient descent [32] that enforces group robustness through a dual-update mechanism (Fig. 2). First, utilizing batches sampled from \(\mathcal{I}\), we compute a corrective gradient \(\mathbf{g}_{C}\) that minimizes the discrepancy between the internal representations of invariant pairs. Second, we perform a task update using the standard loss gradient, but with a critical modification: we monitor a disagreement rate \(\delta\) between pair predictions. If \(\delta\) exceeds a threshold \(\tau\), indicating reliance on spurious features, we strictly scale the task gradient to prioritize the corrective signal. This adaptive scaling allows the model to learn invariance when necessary, while retaining optimization capacity for the task loss when the invariance condition is satisfied. We detail these components in the following section.
3.2 Rationale and Invariance Pair Definition↩︎
To compare learned characteristics and formulate the corrective gradient \(\mathbf{g}_{C}\), we utilize a latent representation known as the rationale matrix [33]. The rationale explicitly encodes the contribution of learned features to the final classification. We decompose the model \(f\) into two modules: a feature extractor \(\phi\), mapping inputs \(\mathbf{x}\) to features \(\mathbf{z}\coloneq \phi(\mathbf{x})\), and a classifier \(h\), mapping \(\mathbf{z}\) to logits \(\mathbf{o} \coloneq h(\mathbf{z})\). The rationale \(\mathbf{R}\) is defined as the element-wise product of the feature activations and the classifier weights \(\boldsymbol{\theta}_h\): \[\mathbf{R}\coloneq \text{diag}(\mathbf{z})\boldsymbol{\theta}_h \in \mathbb{R}^{D\times K},\] where \(D\) and \(K\) denote the feature dimension and number of classes, respectively. Unlike raw logits, \(\mathbf{R}\) retains fine-grained semantic information by incorporating the final layer’s weights, making it an ideal candidate for guiding the training process.
To identify and penalize reliance on spurious correlations, we contrast rationales derived from a set of invariance pairs \(\mathcal{I}\). We use a set of invariant pairs \((\mathbf{x}_{\text{inv}}, \mathbf{x}_{\text{inv}}') \in \mathcal{I}\) to define a specific invariance. Each pair \((\mathbf{x}_{\text{inv}}, \mathbf{x}_{\text{inv}}') \in \mathcal{I}\) contains two variations of the same semantic content that differ only in a specific spurious attribute. Ideally, since the semantic label \(y\) and class-relevant features are identical, the classification, and thus the rationale, should remain invariant to this change. For example, in the ColoredMNIST dataset, a pair consists of the same digit rendered in different colors (Fig. [fig:colored95mnist95invariance95pair]). These pairs provide the necessary ground truth to define the target model behavior independent of the spurious correlation.
3.3 Corrective Gradient and Invariance Condition↩︎
To guide the neural network towards robust representations, we use a sequential optimization strategy. During training, we sample a batch of invariance pairs \(\mathbf{I}\) from \(\mathcal{I}\) to compute the corrective updates. We establish a distance measure to quantify the internalization of the specified invariance. Given the batches of invariance pairs \((\mathbf{x}_{\text{inv}}, \mathbf{x}_{\text{inv}}') \in \mathbf{I}\), let \(\mathbf{R}, \mathbf{R}' \in \mathbb{R}^{K \times D}\) denote their corresponding rationale matrices, where \(R_{c,j}\) represents the contribution of feature \(j\) to class \(c\). The model predictions \(\hat{\boldsymbol{y}}_{\text{inv}}\) are equal to an aggregation of these rationales and an input-agnostic bias \(\boldsymbol{\theta}_{h,bias,c}\): \[\hat{\boldsymbol{y}}_{\text{inv}}= \text{softmax}\left(\sum_{j=1}^D R_{c,j} + \boldsymbol{\theta}_{h,bias,c}\right). \label{eq:rationale95prediction}\tag{1}\] Since the pairs are invariant by definition, their rationales, and consequently their predictions, should be aligned. We define an alignment loss based on the mean pairwise Frobenius distance: \[\mathcal{L}_A\coloneq \frac{1}{\vert\mathbf{I}\vert}\sum_i\Vert \mathbf{R}_i - \mathbf{R}_i'\Vert_F.\] To prevent model collapse, where the model maps all inputs to a trivial constant representation, we enforce a uniformity constraint. We define a hinge loss to ensure that rationales for class-conflicting elements differ by at least a fixed margin \(m\): \[\mathcal{L}_U\coloneq \frac{1}{2|\mathbf{I}|^2} \sum_{\mathbf{Q} \in \{\mathbf{R}, \mathbf{R}'\}} \sum_{i,j} \mathbb{1}(y_i \neq y_j) \cdot \max \left(0, m- \Vert \mathbf{Q}_i - \mathbf{Q}_j \Vert_F\right).\] The total corrective loss is the sum of these components: \(\mathcal{L}_C\coloneq \mathcal{L}_A+ \mathcal{L}_U\).
We extend standard mini-batch gradient descent by performing two distinct update steps per iteration: a correction step followed by a scaled task optimization step.
3.3.0.1 Step 1: Correction.
We first compute the corrective gradient \(\mathbf{g}_{C}\coloneq \nabla_{\boldsymbol{\theta}}\mathcal{L}_C\). This gradient directs the model parameters \(\boldsymbol{\theta}\) towards a region where the invariance constraints (alignment and uniformity) are satisfied. The optimizer performs the first update using \(\mathbf{g}_{C}\).
3.3.0.2 Step 2: Scaled Task Optimization.
Subsequently, we update the model to minimize the task loss \(\mathcal{L}_T\), defined as cross-entropy loss. However, we do not apply the raw gradient \(\nabla_{\boldsymbol{\theta}} \mathcal{L}_T\) directly. Instead, we compute a scaled task gradient \(\mathbf{g}_{T}\) that respects the invariance established in the previous step. We quantify this dependency using the disagreement rate \(\delta\), defined as the fraction of invariant pairs yielding inconsistent class predictions: \[\delta\coloneq \frac{1}{|\mathbf{I}|} \sum_{k} \mathbb{1}(\arg\max \hat{\boldsymbol{y}}_{\text{inv},k} \neq \arg\max \hat{\boldsymbol{y}}_{\text{inv},k}').\]
To formalize the switching behavior of our optimizer, we define an invariance condition: \(\delta< \tau\), based on a fixed predefined threshold \(\tau\in [0,1]\). When \(\delta\ge \tau\), the model violates the invariance condition, triggering a strict corrective mechanism. The mechanism rescales the raw task gradient \(\nabla_{\boldsymbol{\theta}}\mathcal{L}_T\) to a fixed fraction \(\alpha\in \left[0,1\right)\) of the magnitude of the corrective gradient \(\mathbf{g}_{C}\). Otherwise, gradient magnitudes are harmonized. We calculate the final scaled task gradient \(\mathbf{g}_{T}\) based on the invariance condition: \[\mathbf{g}_{T}\coloneq \begin{cases} \alpha\cdot \frac{\nabla_{\boldsymbol{\theta}} \mathcal{L}_T}{\|\nabla_{\boldsymbol{\theta}} \mathcal{L}_T\|_2} \|\mathbf{g}_{C}\|_2 & \text{if } \delta\ge \tau\\[2ex] \text{clip}\big(\nabla_{\boldsymbol{\theta}} \mathcal{L}_T, \, 2\cdot\|\mathbf{g}_{C}\|_2\big) & \text{else}, \end{cases} \label{eq:loss95grad}\tag{2}\] where \(\text{clip}(\mathbf{u}, c)\) rescales vector \(\mathbf{u}\) such that \(\|\mathbf{u}\|_2 \le c\), maintaining its direction. Finally, the optimizer performs the second update using \(\mathbf{g}_{T}\). We update the weights using an optimizer that can involve momentum-based techniques, allowing for a stabilizing effect. This adaptive scaling mechanism prioritizes invariant representation learning when spurious correlations are detected, while allowing the task loss to dominate when the invariance condition is satisfied.
An overview of the presented method is shown in Alg. 3 defining the extended update function for one gradient descent step for a training-data batch given by \((\mathbf{x}, \mathbf{y})\).
4 Theoretical Analysis↩︎
In this section, we analyze the theoretical properties of the approach, focusing on two aspects: the monotonic correction, its resulting convergence properties, and a comparative analysis of computational complexity.
4.1 Monotonic Correction and Convergence Properties↩︎
For this analysis, we show that IPG results in a monotonic descent on \(\mathcal{L}_{C,\boldsymbol{\theta}_{}}\) when the invariance condition is violated. Based on the theorem, we further derive convergence properties. We introduce the following simplifications: First, the parameters \(\boldsymbol{\theta}_k\) at the start of the IPG update step (Alg. 3) are fixed. Second, we treat the optimization process as standard gradient descent (ignoring momentum for the step-wise analysis). All norms applied are Euclidean. We rely on the following common assumptions:
Assumption 1 (Convexity). The function \(\mathcal{L}_{C,\boldsymbol{\theta}_{}}\) is convex over the domain \(\boldsymbol{\Theta}\subseteq \mathbb{R}^d\).
Assumption 2 (Smoothness). The loss function \(\mathcal{L}_{C,\boldsymbol{\theta}_{}}\) is \(L\)-smooth, meaning its gradient is Lipschitz continuous. Consequently, for all \(\boldsymbol{\theta}, \boldsymbol{\theta}' \in \mathbb{R}^d\): \[\mathcal{L}_C(\boldsymbol{\theta}') \leq \mathcal{L}_C(\boldsymbol{\theta}) + \langle \nabla \mathcal{L}_C(\boldsymbol{\theta}), \boldsymbol{\theta}' - \boldsymbol{\theta}\rangle + \frac{L}{2} \|\boldsymbol{\theta}' - \boldsymbol{\theta}\|^2.\]
Assumption 3 (Joint Feasibility). There exists a non-empty subset of parameters \(\boldsymbol{\Theta}^* \subset \boldsymbol{\Theta}\) such that for \(\boldsymbol{\theta}\in \boldsymbol{\Theta}^*\), \(\mathcal{L}_A(\boldsymbol{\theta}) = 0\) and \(\mathcal{L}_U(\boldsymbol{\theta}) = 0\). This implies the rationale dimension \(D\) is sufficient to satisfy margin constraints without violating invariance.
theoremmonotonicDescent Let Assumptions 1 and 2 hold. Suppose the invariance condition is violated (\(\delta\ge \tau\)) and \(\tau\in (0,1]\). Then, provided the task gradient scaling factor \(\alpha< 1\) and the learning rate \(\eta\) is sufficiently small, the algorithm ensures monotonic descent: \[\mathcal{L}_{C,\boldsymbol{\theta}_{k+1}} < \mathcal{L}_{C,\boldsymbol{\theta}_{k}}.\]
The full proof is provided in Appendix 9, showing a dominance of \(\mathbf{g}_{C}\) when \(\alpha< 1\) and thus the method results in a monotonic descent under given assumptions.
corollaryconvergenceProperties Let Assumptions 1, 2, and 3 hold. If the optimization follows the Monotonic Descent property (Theorem [thm:monotonic95desc]) whenever the invariance condition is violated (\(\delta> \tau\)), then, given Eq. 1 :
Finite Feasibility: For any relaxed threshold \(\tau > 0\), the sequence of parameters \(\{\boldsymbol{\theta}_k\}\) reaches the invariant region \(\mathcal{R}_\tau= \{ \boldsymbol{\theta}\mid \delta(\boldsymbol{\theta}) < \tau\}\) in a finite number of steps.
Finite Exact Convergence: For the strict case \(\tau= 0\), the discrete disagreement rate reaches the exact invariant state \(\delta(\boldsymbol{\theta}) = 0\) in a finite number of steps.
The proof is provided in Appendix 9, deriving both properties from the Theorem [thm:monotonic95desc].
4.2 Extension to Stochastic and Momentum-based Optimization↩︎
While the theoretical analysis above assumes full-batch gradient descent for clarity, the results generalize to practical training settings under standard assumptions.
4.2.0.1 Stochastic Gradient Descent (SGD).
In the stochastic setting, the monotonic descent condition (Theorem [thm:monotonic95desc]) translates to a descent in expectation: \(\mathbb{E}[\mathcal{L}_C(\boldsymbol{\theta}_{k+1})] \leq \mathcal{L}_C(\boldsymbol{\theta}_k)\). Provided that the learning rate follows a standard decay schedule and the stochastic gradients satisfy bounded variance assumptions, the convergence to the invariant region \(\mathcal{R}_\tau\) holds almost surely. Specifically, for \(\tau > 0\), the parameter trajectory enters the invariant region \(\mathcal{R}_\tau\) with probability 1. For \(\tau = 0\), the stability margin ensures that the parameters reach the exact invariant state \(\delta(\boldsymbol{\theta}_k) = 0\) in finite time almost surely.
4.2.0.2 Momentum.
For momentum-based optimizers, the update vector is a history-dependent aggregation of past gradients. In this context, the proposed gradient scaling (Eq. 2 ) ensures that all new information entering the momentum buffer satisfies the invariance constraint. While the momentum history may introduce a transient lag, the spurious components from previous steps decay exponentially (controlled by the momentum factor \(\beta < 1\)). Since the source term is continuously corrected, the momentum buffer effectively discards spurious alignment, guaranteeing that the update direction aligns with the invariant region or exact invariant state in finite time.
4.2.0.3 Summary.
The theoretical analysis demonstrates that the corrective gradient induces a monotonic descent in the correction error \(\mathcal{L}_C\) whenever the invariance condition is violated. For relaxed constraints (\(\tau > 0\)), this guarantees that the model satisfies the condition in a finite number of steps. For strict constraints (\(\tau = 0\)), the discrete classification becomes invariant in finite time once the correction error falls below the decision boundary’s stability margin.
4.3 Computational Complexity Analysis↩︎
The computational complexity of IPG during the training phase depends on the number of epochs \(N\), the dataset size \(D\), the task batch size \(B_D\), and the invariance pair batch size \(B_I\). We define the complexity in terms of the number of forward and backward passes through the network. Standard ERM requires one forward and one backward pass per sample, scaling linearly as \(\mathcal{O}(N \cdot D)\).
In contrast, IPG introduces an overhead due to the corrective mechanism. For every task batch, we sample a batch of invariant pairs (size \(B_I\)) and perform a corrective update. Since pairs consist of two images \((\mathbf{x}_{\text{inv}}, \mathbf{x}_{\text{inv}}')\), this adds \(2 \cdot B_I\) forward passes and a corresponding backward pass per iteration. The total number of iterations per epoch is \(\lceil D/B_D \rceil\). Thus, the total operations for IPG can be approximated as: \[C_{\text{IPG}} \approx N \cdot \left( D + \left\lceil\frac{D}{B_D}\right\rceil \cdot 2 \cdot B_I \right) \approx N \cdot D \left( 1 + \frac{2 B_I}{B_D} \right).\] While the asymptotic complexity remains linear, \(\mathcal{O}(N \cdot D)\), the computational complexity of IPG scales linearly as \(N \cdot D (1 + 2B_I / B_D)\). In our experiments, we set \(B_I = B_D\). However, this represents a conservative upper bound rather than a structural requirement.
We compare the computational characteristics of IPG with related approaches in Table 1. We categorize methods based on their estimated relative runtime factor compared to standard ERM.
| Approach | Mechanism | Relative Cost (\(\approx\)) |
|---|---|---|
| ERM | – | \(1\times\) |
| Pre-processing / Input Transformation | ||
| Shock Graph [19] | Input Transformation | \(1\times\) (+ Pre.) |
| NCM [8] | Linear Projection | \(1\times\) (+ Pre.) |
| Explanation-Guided | ||
| EPG [23] | Input Gradient Penalty | \(2\times\) |
| GALS [24] | Attention Alignment | \(2\times\) |
| Penalty-based Methods | ||
| IRM [20] | Gradient Norm Penalty | \(2\times\) |
| Fishr [30] | Gradient Variance Matching | \(2\times\) |
| GroupDRO [26] | Group-wise Loss | \(1\times\) |
| Two-Stage / Dual-Model | ||
| JTT [27] | Two-stage Training | \(2\times\) |
| LfF [29] | Bias-Amplified Training | \(2\times\) |
| Post-hoc / Feature Editing | ||
| EvA-E [28] | Feature Deletion | \(1\times\) (+ Post.)\(^\dagger\) |
| Proposed Method | ||
| IPG (Ours) | Dual-Update Mechanism | \(1 + \frac{2B_I}{B_D}\times\) |
5 Experiments↩︎
In this section, we present the experimental setup and results along with a detailed discussion.
5.1 Experimental Setup↩︎
We evaluate IPG on three datasets representing unique challenges in spurious correlation: label noise, perfect spurious correlation, and feature entanglement. For baseline comparisons, we report results directly from the original publications where the experimental protocol (architecture and data splits) is identical; otherwise, we reproduce baselines using official implementations. Full implementation details are provided in Appendix 8.
5.1.0.1 Label Noise (ColoredMNIST).
We use ColoredMNIST [20] to evaluate robustness under noisy supervision (\(25\%\) label noise). This dataset is structured as a domain generalization benchmark where the task is binary digit classification. The digit color (red or green) serves as a strongly correlated spurious attribute in the training set (\(p_\text{tr} \in \{0.1, 0.2\}\)) but is reversed in the test set (\(p_\text{te} = 0.9\)). To guide the model, we create the dataset with invariance pairs consisting of the same digit instance recolored with the opposite color (e.g., a red ‘1’ paired with a green ‘1’, see Fig. [fig:colored95mnist95invariance95pair]). We employ a standard CNN architecture following the DomainBed framework [34] and report the mean test accuracy on the out-of-distribution test set.
5.1.0.2 Perfect Spurious Correlation (Waterbirds-100).
To evaluate robustness without natural counter-examples, we use the Waterbirds-100 variant, following the protocol of Petryk et al. [24]. Unlike the standard benchmark, the training set is perfectly correlated (100% spurious), excluding any bias-conflicting samples. To overcome this lack of support, we generate invariance pairs by explicitly swapping the background of a given bird instance (e.g., pasting a waterbird onto a land background, see Fig. [fig:waterbird95invariance95pair]). This setting tests the model’s ability to extrapolate solely through this sparse guidance, as the test dataset consists of samples from all four groups. We also evaluate the reversed setting, where the target label is the background and the spurious attribute is the bird type. We train a ResNet-50 backbone [35] and report worst-group accuracy (\(\text{Acc}_\text{wg}\)) for both the standard and reversed settings.
5.1.0.3 Feature Entanglement (CelebA).
We use the CelebA dataset [26] to assess performance in a setting where spurious and semantic features are inherently entangled. The task is classifying hair color (blond vs. non-blond), which is spuriously correlated with gender. Unlike the previous datasets where attributes can be perfectly isolated (e.g., via segmentation), real-world attributes such as gender and hair style are semantically linked. We address this by generating invariance pairs via instruction-based editing using the FLUX.1 Kontext [dev] model [36], prompting it to alter the gender while preserving identity and hair color (Fig. [fig:celeba95invariance95pair]). Consequently, the generated invariance pairs inevitably contain feature entanglement and generation artifacts. See Appendix 10 for further details. This setting evaluates whether IPG can extract useful gradients from noisy, entangled pairs using a ResNet-50 model. We report worst-group accuracy (\(\text{Acc}_\text{wg}\)).
5.2 Accuracy Comparison↩︎
5.2.0.1 ColoredMNIST Results.
We evaluate IPG on the ColoredMNIST test set, where the spurious correlation (color-label association) is inverted relative to the training set. As shown in Table 2, IPG achieves a mean test accuracy of \(72.8\%\). This result effectively recovers the performance level of the grayscale oracle (\(73.1\%\)) and approaches the theoretical limit of \(75\%\) imposed by the dataset’s label noise. In comparison to state-of-the-art methods, IPG is statistically indistinguishable from DIAR-RA [15] (\(73.1\%\)), which relies on color-augmentation and thus breaks spurious correlation. Notably, IPG outperforms NCM [8] (\(69.3\%\)), a recent approach that also utilizes invariance pairs but uses a static linear subspace projection. This suggests that IPG’s dynamic gradient correction captures non-linear invariances more effectively than NCM’s linear removal. We further observe that IPG outperforms Shock Graph (\(71.6\%\)), which explicitly discards color information via preprocessing. Overall, IPG demonstrates superior performance compared to standard invariance baselines (IRM, Fishr) and competes with intensive augmentation strategies like DIAR, validating the effectiveness of sparse invariance guidance in extreme correlation shifts. Summarizing, IPG reaches oracle-level performance on ColoredMNIST, comparable to the state-of-the-art DIAR-RA.
| Approach | Mean Acc. ↑ | |||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| ERM | 16.1 ± 0.8 | |||||||||||||||||||||
| IRM [20] | 66.9 ± 2.5 | |||||||||||||||||||||
| Fishr [30] | 68.8 ± 1.4 | |||||||||||||||||||||
| NCM [8] | 69.3 | |||||||||||||||||||||
| meta-IRM [37] | 70.4 ± 0.9 | |||||||||||||||||||||
| Shock Graph [19] | 71.6 | |||||||||||||||||||||
| DIAR-RA [15] | 73.1 ± 0.1 | |||||||||||||||||||||
| IPG (ours) | 72.8 ± 0.4 | |||||||||||||||||||||
| Random | 50.0 ± 0.0 | |||||||||||||||||||||
| ERM grayscale (oracle) | 73.1 ± 0.4 | |||||||||||||||||||||
| Optimal | 75.0 ± 0.0 |
5.2.0.2 Waterbirds-100 Results.
We further evaluate the performance of IPG on the Waterbirds-100 dataset, which has a perfect spurious correlation in the training set: every waterbird appears against a water background, and every landbird against a land background, with no bias-conflicting samples available. This creates a challenging zero-shot generalization scenario for the worst-group (e.g., waterbirds on land). Following the protocol of [23] and [24], we also examine a reversed setting (Waterbirds-100-reverse), where labels and spurious attributes are swapped (predicting the background class instead of the bird species).
We compare our approach against standard ERM and two state-of-the-art guidance methods: EPG [23], which aggregates attributions from B-cos, \(\mathcal{X}\)-DNN, and IxG; and the language-guided GALS method [24]. Table 3 summarizes the results. In this data-scarce regime, IPG outperforms the baselines with 124 invariance pairs. In the standard setting, IPG achieves a worst-group accuracy of 86.3%, outperforming the best baseline (GALS, 56.7%) by a margin of +29.6 percentage points. Similarly, in the reverse setting, IPG improves upon GALS by +13.8 percentage points (86.7% vs. 72.9%). We conclude that IPG establishes a new state-of-the-art for this strict spurious correlation setting to the best of our knowledge. It effectively mitigates the reliance on background features even when the correlation is perfect, by using a sparse set of invariance pairs.
| Waterbirds-100 | Waterbirds-100-reverse | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| Model | Mean Acc. ↑ | \(\accWg\) ↑ | Mean Acc. ↑ | \(\accWg\) ↑ | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| ERM | 69.3 ± 1.1 | 37.9 ± 1.7 | 83.3 ± 1.8 | 62.7 ± 3.5 | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| EPG B-cos [23] | 71.1 ± 0.9 | 41.0 ± 2.1 | 83.6 ± 1.1 | 62.8 ± 2.1 | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| EPG \(\mathcal{X}\)-DNN [23] | 73.1 ± 3.4 | 47.0 ± 9.1 | 82.6 ± 2.0 | 63.9 ± 3.6 | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| EPG IxG [23] | 78.1 ± 2.6 | 56.1 ± 7.0 | 78.9 ± 1.9 | 56.5 ± 3.7 | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| GALS [24] | 79.7 | 56.7 | 86.8 | 72.9 | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| IPG (ours) | 89.5 ± 1.1 | 86.3 ± 1.9 | 93.0 ± 0.7 | 86.7 ± 1.7 | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
5.2.0.3 CelebA Results.
Finally, we evaluate IPG on CelebA, where the spurious correlation (gender bias) arises from complex, real-world data distributions. We compare IPG against the upper-bound baseline GroupDRO [26] and DKT [25] (which uses explicit group labels), and recent label-free approaches including LfF [29], JTT [27], and EvA-E [28]. Results are summarized in Table 4. IPG (\(74.7\%\)) achieves a substantial improvement over the standard ERM baseline (\(47.8\%\)) and outperforms the label-free LfF method (\(70.6\%\)). While IPG lags behind implicit methods like JTT (\(81.1\%\)) and feature-synthesis methods like DKT (\(92.3\%\)), it achieves this result using a single-stage training process with sparse input constraints. Notably, IPG maintains a high mean test accuracy (\(90.8\%\)), outperforming JTT (\(88.0\%\)) and LfF (\(86.0\%\)) in average performance. This indicates that IPG acts as a conservative regularizer: it successfully mitigates the worst-group failure modes (+26.9 percentage points gain) without over-correcting to the point of degrading the model’s core semantic capabilities.
| Model | Test Acc. ↑ | \(\accWg\) ↑ | ||||||||||||||||||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| ERM | 94.9 ± 0.2 | 47.8 ± 3.7 | ||||||||||||||||||||||||||||||||||||||||||
| LfF [29] | 86.0 | 70.6 | ||||||||||||||||||||||||||||||||||||||||||
| JTT [27] | 88.0 | 81.1 | ||||||||||||||||||||||||||||||||||||||||||
| GroupDRO [26] | 92.9 ± 0.2 | 88.9 ± 2.3 | ||||||||||||||||||||||||||||||||||||||||||
| EvA-E [28] | 88.7 | 82.7 | ||||||||||||||||||||||||||||||||||||||||||
| DKT [25] | 94.6 | 92.3 | ||||||||||||||||||||||||||||||||||||||||||
| IPG (ours) | 90.8 ± 1.6 | 74.7 ± 8.1 |
5.3 Effectiveness and Efficiency Analysis↩︎
5.3.0.1 Data Efficiency.
We analyze the sensitivity of IPG to the number of invariance pairs (Table 5). To this end, we evaluate four exponentially spaced sample sizes ranging from 40 to 1208 pairs. Remarkably, IPG achieves robust performance with as few as 124 pairs across all datasets. Increasing the number of pairs to 1208 pairs leads to diminishing improvements, confirming that IPG relies on a sparse signal rather than dense coverage. This validates our claim of data efficiency compared to contrastive methods requiring massive batch sizes. For this reason, we limit our method to \(<4\%\) when comparing.
5.3.0.2 Hyperparameter Stability.
We analyze the impact of the corrective step size \(\alpha\) and threshold \(\tau\) (detailed results in Appendix 7). We observe a dichotomy between synthetic and real-world shifts. On ColoredMNIST and Waterbirds-100, strict constraints (\(\tau=0.1, \alpha=0.1\)) perform best, effectively mitigating reliance of the model on the color or background bias. Conversely, on CelebA, looser constraints (\(\tau\approx 0.5\)) prevent the corrective gradient from impeding the learning of core semantic features.
| ColoredMNIST | Waterbirds-100 | CelebA | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| Number of Pairs | \(\accWg\) ↑ | % \(\trainDomain\) | \(\accWg\) ↑ | % \(\trainDomain\) | \(\accWg\) ↑ | % \(\trainDomain\) | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| 40 | 66.3 ± 4.3 | 0.1 | 83.9 ± 1.2 | 0.8 | 67.5 ± 8.5 | 0.02 | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| 124 | 69.3 ± 6.2 | 0.3 | 86.3 ± 1.9 | 2.6 | 66.8 ± 8.7 | 0.08 | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| 388 | 72.5 ± 0.5 | 1.0 | 86.9 ± 0.8 | 8.1 | 73.2 ± 6.2 | 0.2 | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| 1208 | 72.8 ± 0.4 | 3.2 | 87.0 ± 0.5 | 25.2 | 74.7 ± 8.1 | 0.7 | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
5.4 Discussion and Limitations↩︎
5.4.0.1 Oracle Recovery and Label Noise.
Our experiments on ColoredMNIST demonstrate that IPG effectively encodes group invariance, matching the performance of the grayscale oracle (\(72.8\%\) vs. \(73.1\%\)). Critically, because the invariance is defined by explicit pairs rather than dataset labels, the method is resilient to the inherent label noise. This suggests that for simple spurious correlations, explicit invariance guidance allows the model to discard spurious information as effectively as preprocessing steps that remove it.
5.4.0.2 Extrapolation and Semantic Agnosticism.
On Waterbirds-100, IPG achieves state-of-the-art results (86.3% worst-group accuracy), outperforming baselines that rely on other guiding mechanisms using standard additive regularization. This confirms that sparse invariance pairs can successfully guide the model to extrapolate to unseen groups (e.g., waterbirds on land) even when the training set provides zero support. Furthermore, the equally strong performance in the reversed setting demonstrates that IPG is agnostic to the semantic nature of the artifact. Whether the spurious feature is a background texture or a foreground object, the method mitigates reliance on the shortcut once it is defined by the pairs.
5.4.0.3 Real-World Entanglement.
On CelebA, IPG acts as a conservative regularizer: it improves worst-group accuracy by +26.9 percentage points over ERM while maintaining superior mean test accuracy compared to implicit methods like JTT. However, the remaining gap to dense supervision methods highlights the challenge of semantic entanglement.
Unlike Waterbirds, where the background is geometrically distinct from the object, real-world attributes like gender and hair style are coupled; altering one often necessitates changes in the other to maintain photorealism [38]. This violates the assumption that spurious attributes can be intervened upon independently to core features, resulting in generated pairs that contain inevitable artifacts.
5.4.0.4 Theoretical Integrity.
The efficacy of IPG is supported by our proof of monotonic descent toward the invariant region in Section 4. We establish that the corrective gradient mathematically forces the parameter trajectory to minimize alignment error whenever the invariance condition is violated. Consequently, the model is provably guided to a state that satisfies the invariance condition or exhibits zero disagreement. This provides a formal guarantee of robustness that distinguishes IPG from the heuristic-based updates of implicit methods. A remaining direction for future work is the formal analysis of the stability of this invariant region—specifically, determining the conditions under which the task gradient might displace the parameters from the region once the corrective force is reduced.
5.4.0.5 Complexity, Efficiency, and Limitations.
Although our experiments used a configuration (\(B_I=B_D\)) resulting in a \(3\times\) overhead, this represents a conservative performance upper bound. Our data-efficiency results prove that robust features are recovered with very few pairs, meaning the guidance ratio \(B_I/B_D\) can likely be reduced in practice to achieve an overhead of \(\approx 1.5\times\) ERM. However, certain limitations remain. First, performance is bounded by the quality of invariance pairs; poor generation can misguide the model, though improvements in generative AI will likely ease this bottleneck. Second, the current rationale formulation is specific to classification.
6 Conclusion↩︎
In this work, we presented Invariance Pair Guidance (IPG), a framework that decouples robust learning from the need for dense supervision or bias-conflicting data. By leveraging a sparse set of invariance pairs, IPG formulates a corrective gradient that penalizes model rationales only when they diverge from the desired invariance. Empirically, IPG demonstrates remarkable universality: it matches oracle performance on synthetic data (ColoredMNIST), achieves state-of-the-art extrapolation on data-scarce regimes (Waterbirds-100), and acts as a conservative regularizer on real-world data (CelebA). Our theoretical analysis supports these findings, guaranteeing monotonic descent toward the invariant region. Crucially, IPG offers a data-efficient path to robustness that eliminates the need for dataset preprocessing or multi-stage training pipelines. In future, we will focus on integrating IPG with latent-space generative models to synthesize high-fidelity, disentangled pairs for increasingly complex real-world attributes. In addition, we plan to generalize the framework to multi-objective settings.
7 Detailed Sensitivity Analysis↩︎
In this section, we provide the full empirical results regarding hyperparameter sensitivity and ablation studies, complementing the summary in Section 5.3.
7.1 Hyperparameter Sensitivity↩︎
We analyzed the impact of the corrective gradient step size \(\alpha\) and the disagreement threshold \(\tau\) across 10 independent trials. Table 6 and Table 7 visualize the trade-off between strict invariance enforcement and task learning.
7.1.0.1 Synthetic vs. Real-World Dynamics.
On ColoredMNIST (Table 6) and Waterbirds-100 (Table 7), the best performance is achieved with strict constraints (\(\tau=0.1, \alpha=0.1\)). Since the invariance pairs clearly define the spurious correlation, any deviation from the invariance condition indicates reliance on spurious correlation, requiring immediate correction. Conversely, on CelebA (Table 7), the model benefits from a looser threshold (\(\tau\approx 0.5\)). The value allows the optimization to suppress minor inconsistencies in the rationales of the invariance pairs, preventing the corrective gradient from overriding useful semantic feature learning.
| ColoredMNIST (\(\thresh\)) | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| \(\stepSize\) | 0.1 | 0.5 | 0.9 | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| 0.1 | 64.3 +- 8.6 | 34.4 +- 1.7 | 17.4 +- 1.4 | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| 0.5 | 55.6 +- 13.4 | 32.6 +- 1.6 | 14.6 +- 2.9 | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| 0.9 | 41.9 +- 7.7 | 29.1 +- 2.3 | 13.8 +- 0.3 | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| Waterbirds-100 (\(\thresh\)) | CelebA (\(\thresh\)) | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| \(\stepSize\) | 0.1 | 0.5 | 0.9 | 0.1 | 0.5 | 0.9 | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| 0.1 | 81.0 +- 1.9 | 78.2 +- 0.8 | 78.4 +- 1.1 | 68.9 +- 5.2 | 68.4 +- 4.0 | 71.7 +- 4.3 | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| 0.5 | 80.9 +- 1.9 | 78.2 +- 0.9 | 78.0 +- 1.0 | 69.8 +- 5.2 | 71.1 +- 8.4 | 71.3 +- 6.4 | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| 0.9 | 79.7 +- 1.3 | 78.3 +- 0.8 | 78.3 +- 0.8 | 70.9 +- 7.9 | 74.5 +- 5.7 | 70.9 +- 3.7 | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
7.2 Data Efficiency Visualization↩︎
In the main text, we reported the mean performance across varying numbers of pairs. Figure 4 provides the corresponding boxplot distribution, highlighting the variance reduction as the number of pairs increases.
7.3 Ablation Studies↩︎
To validate the specific design choices of IPG, we compare our method against two simplifications of the optimization mechanism (Table 8):
Joint Optimization (Sum of Losses): A standard multi-task formulation where the task loss \(\mathcal{L}_T\) and the alignment loss \(\mathcal{L}_A\) are applied consecutively. This removes the adaptive gradient scaling.
Continuous Correction (Full Control): A variant of IPG where the invariance threshold is set to an unreachable strict limit (\(\tau= 0\)). In this scenario, the invariance condition is effectively never satisfied, forcing the gradient correction to be applied at every single step. This tests whether the correction relaxation of our correction for the invariant region is beneficial.
As shown in Table 8, Joint Optimization fails to recover the Oracle performance on ColoredMNIST (\(10.2\%\) vs. \(72.8\%\)) or the performance on Waterbirds-100 (\(33.1\%\) vs. \(86.3\%\)). The difference indicates that simply adding up the losses is not sufficient to control for the spurious correlation. The model fails to prioritize the invariance over the strong spurious signal.
Continuous Correction achieves strong results on ColoredMNIST (\(72.8\%\)) and Waterbirds-100 (\(86.3\%\)), which is expected since the spurious correlation is perfect and requires constant suppression. However, on real-world datasets like CelebA, continuous correction degrades performance slightly compared to IPG (\(70.9\%\) vs. \(74.7\%\)). This indicates that the model benefits from an relaxed invariance condition allowing a certain rate of mismatches.
| Method | ColoredMNIST | Waterbirds-100 | CelebA |
|---|---|---|---|
| Joint Optimization | 10.2 ± 0.2 | 33.1 ± 4.0 | 72.2 ± 5.2 |
| Continuous Correction (\(\thresh = 0\)) | 72.8 ± 0.4 | 86.1 ± 2.0 | 70.9 ± 9.0 |
| IPG | 72.8 ± 0.4 | 86.3 ± 1.9 | 74.7 ± 8.8 |
8 Implementation Details↩︎
8.1 Training and Model Selection↩︎
We train the models with the following specifications. For ColoredMNIST, we employ a standard Convolutional Neural Network (CNN) following the architecture defined in the DomainBed framework [34]. For Waterbirds-100 and CelebA, we use a ResNet-50 backbone [35] pretrained on ImageNet. All models are optimized using Adam with a class-weighted cross-entropy loss. Accuracy is evaluated over \(10\) independent trials with different random seeds. Model selection is based on the highest validation accuracy for ColoredMNIST, and the maximum worst-group validation accuracy for Waterbirds-100 and CelebA.
8.2 Hyperparameters↩︎
The hyperparameters for our approach are listed in Table 9 and are selected using the Tree-structured Parzen Estimator with \(100\) trials per dataset [39]. We apply the same batch size for data and invariance pair batches.
| Dataset | Model | \(\stepSize\) | \(\thresh\) | \(\margin\) | \(|\invpairs|\) | \(\lr\) | Batch Size | Nr. Epochs | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| ColoredMNIST | IPG | 0.5 | 0.0 | 0.08 | 1208 | 1 × 10−3 | 128 | 18 | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| Waterbirds-100 | IPG | 0.1 | 0.0 | 0.1 | 124 | 1 × 10−4 | 128 | 10 | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| Waterbirds-100-reverse | IPG | 0.5 | 0.5 | 0.1 | 124 | 1 × 10−4 | 128 | 10 | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| CelebA | IPG | 0.5 | 1 × 10−3 | 1.5 | 1208 | 8 × 10−3 | 128 | 10 |
9 Proof of Monotonic Correction and Convergence Properties↩︎
We provide the full proof of the Theorem [thm:monotonic95desc] and Corollary [cor:convergence95properties].
Proof. Since \(\delta\ge \tau\), the update rule (Alg. 3) applies a correction. The task gradient is scaled such that: \[\|\mathbf{g}_{T}\| = \alpha\|\mathbf{g}_{C}\|.\label{eq:convergence95grad95correction}\tag{3}\] The parameter update is given by: \[\label{eq:update} \boldsymbol{\theta}_{k+1} = \boldsymbol{\theta}_k -\eta(\mathbf{g}_{C}+\mathbf{g}_{T}).\tag{4}\] Using the \(L\)-smoothness upper bound (Assumption 2):\[\label{eq:conv95smooth95inequality} \mathcal{L}_{C,\boldsymbol{\theta}_{k+1}} \leq \mathcal{L}_{C,\boldsymbol{\theta}_{k}} \underbrace{-\eta\nabla \mathcal{L}_{C,\boldsymbol{\theta}_{k}}^T(\mathbf{g}_{C}+\mathbf{g}_{T})}_{\text{Linear Term}} + \underbrace{\frac{L}{2}\eta^2\|\mathbf{g}_{C}+\mathbf{g}_{T}\|^2}_{\text{Quadratic Term}}.\tag{5}\] We analyze the linear term. In the worst-group scenario, the task gradient \(\mathbf{g}_{T}\) opposes the corrective gradient. Using the Cauchy-Schwarz inequality and Eq. 3 :\[\label{eq:linear95term} \begin{align} -\eta\mathbf{g}_{C}^T (\mathbf{g}_{C}+ \mathbf{g}_{T}) &= -\eta\|\mathbf{g}_{C}\|^2 - \eta\langle \mathbf{g}_{C}, \mathbf{g}_{T}\rangle \\ &\leq -\eta\|\mathbf{g}_{C}\|^2 + \eta\|\mathbf{g}_{C}\| \|\mathbf{g}_{T}\| \\ &= -\eta(1 - \alpha) \|\mathbf{g}_{C}\|^2. \end{align}\tag{6}\] This term is negative (providing descent) only if \(\alpha< 1\). Next, we bound the quadratic penalty term using the triangle inequality:\[\label{eq:quad95term} \begin{align} \|\mathbf{g}_{C}+ \mathbf{g}_{T}\|^2 &\leq (\|\mathbf{g}_{C}\| + \|\mathbf{g}_{T}\|)^2 \\ &= (1 + \alpha)^2 \|\mathbf{g}_{C}\|^2. \end{align}\tag{7}\] Substituting Eqs. 6 and 7 back into Eq. 5 :\[\label{eq:combined95inequality} \mathcal{L}_{C,\boldsymbol{\theta}_{k+1}} \leq \mathcal{L}_{C,\boldsymbol{\theta}_{k}} - \eta\left[ (1 - \alpha) - \frac{L\eta}{2}(1 + \alpha)^2 \right] \|\mathbf{g}_{C}\|^2.\tag{8}\] For monotonic descent (\(\mathcal{L}_{C,\boldsymbol{\theta}_{k+1}} < \mathcal{L}_{C,\boldsymbol{\theta}_{k}}\)), the term in the brackets must be positive. Since \(\alpha< 1\), the term \((1-\alpha)\) is strictly positive. Therefore, there exists a sufficiently small learning rate \(\eta\) such that: \[\frac{L\eta}{2}(1 + \alpha)^2 < 1 - \alpha.\] Under this condition, the loss strictly decreases. ◻
Proof. Let \(\gamma\) denote the effective stopping threshold. For the relaxed case, \(\gamma = \tau\). For the strict case (\(\tau=0\)), \(\gamma = \epsilon\), where \(\epsilon > 0\) is the stability margin of the discrete classifier. Since predictions are equal to aggregating rationales (Eq. 1 ), the local Lipschitz continuity ensures that small deviations in \(\mathbf{R}\) do not alter class labels provided the error is within this margin. Thus, in both cases, the target condition is satisfied if the continuous alignment error falls below \(\gamma\). We prove that the algorithm satisfies this condition in finite time by contradiction. Assume the optimization never satisfies the condition, meaning the alignment error remains above \(\gamma\) for all \(k\). Under this assumption, the corrective update is applied indefinitely. By Theorem [thm:monotonic95desc] (Monotonic Descent), the loss sequence \(\{\mathcal{L}_{C,\boldsymbol{\theta}_{k}}\}\) is strictly decreasing and bounded below by 0. By the Monotone Convergence Theorem and Assumption 3 (Feasibility), the sequence converges to the global minimum where the alignment error is zero: \[\lim_{k \to \infty} \mathcal{L}_A(\boldsymbol{\theta}_k) = 0.\] By the definition of the limit, for any \(\gamma > 0\), there exists a finite iteration \(K\) such that for all \(k \geq K\), \(\mathcal{L}_A(\boldsymbol{\theta}_k) < \gamma\). This contradicts the assumption that the error remains above \(\gamma\) indefinitely. Therefore, the process must reach the target state (entering \(\mathcal{R}_\tau\) for \(\tau > 0\), or achieving exact \(\delta= 0\) for \(\tau = 0\)) in a finite number of steps. ◻
10 Pair Generation↩︎
This section details the generation process for the invariance pairs, providing visual examples, the specific text prompts employed for the CelebA dataset, and an analysis of common failure modes.
10.1 Pair Illustrations↩︎
We provide visual representations of the invariance pairs generated for the three benchmark datasets: ColoredMNIST (Fig. 5), Waterbirds-100 (Figs. 6 and 7), and CelebA (Fig. 8). These examples demonstrate the isolation of spurious attributes while retaining semantic content.
10.2 Prompts for CelebA Pair Generation↩︎
To generate counterfactual pairs (Fig. 8)—specifically, gender-transformed images—we utilized Flux.1 [Context] dev [36], a rectified flow transformer capable of editing images based on text instructions.
For this process, we randomly sampled \(1,300\) instances from each group in the training set. For each source image, three candidate counterfactuals were generated using specific textual prompts. Below, we provide examples of the prompts used.
For instance, when transforming a source image of a dark-haired male, the following prompt was used:
Transform this male person into a female version, retaining the
unique characteristics, especially hair color, age, and position.
The result is photorealistic, with natural lighting.Conversely, for samples involving blond hair (e.g., transforming a blond female), the target hair color was explicitly specified to prevent attribute leakage:
Transform this female person into a male version with clearly
blond hair, retaining the unique characteristics, especially age
and position. The result is photorealistic, with natural
lighting.Prompts for other subgroups follow a similar template, with the gender-specific keywords (highlighted in blue) substituted accordingly.
10.3 CelebA Failure Cases↩︎
Figure 9 illustrates instances where the generation process failed to meet the invariance criteria. Common failure modes included: (i) unintentional alteration of the hair color, (ii) unsuccessful gender transformation, or (iii) significant modification of extraneous attributes such as accessories, skin tone, or head pose.
References↩︎
Code is available at https://github.com/msurner/ipg.↩︎