GRIT: Out-of-Distribution Robustness from Finite and Noisy Invariant Pairs

Ruqi Bai; Easton Currie; Yao Ji; Mingyu Kim; Zeyu Zhou; David I. Inouye

cs.LGMay 30, 2025

GRIT: Out-of-Distribution Robustness from Finite and Noisy Invariant Pairs

Authors: Ruqi Bai, Easton Currie, Yao Ji, Mingyu Kim, Zeyu Zhou, David I. Inouye

Organizations: Purdue University · Georgia Institute of Technology

Abstract

Predictors often exploit spurious correlations that fail under distribution shift, undermining out-of-distribution (OOD) robustness. Invariant pairs, such as the same object photographed against different backgrounds, directly indicate which variation should not affect prediction. Yet real pair collections are finite and noisy, so they can miss relevant variation or lead a learner to discard predictive information. To learn from such collections, we propose Geometric Robustness Invariant Training (GRIT), which estimates an unstable subspace from the covariance of pair differences in fixed pretrained features and projects it out before fitting a predictor. Because pair processing is separate from predictor fitting, the same projection can accompany ERM or other robustness objectives. Unlike quadratic prediction consistency, which penalizes variation along every direction with nonzero pair covariance, including variation introduced by pair noise, GRIT enforces exact invariance only to the estimated subspace. We bound GRIT's target risk in terms of pair count, noise relative to the weakest clean signal, and selected rank, assuming that the label distribution given a linear projection of the features is invariant across environments. Empirically, GRIT improves the mean OOD accuracy of each of eight learning objectives on ColoredMNIST-INV and Waterbirds-INV and has higher mean accuracy than a prediction-consistency penalty using the same pairs in 15 of 16 settings. On UrbanCars, combining pairs that vary the background with pairs that vary a co-occurring object raises worst-group accuracy above that of either pair type alone.

Figures & tables

Panel of Figure 1 — Figure 1: GRIT uses noisy invariant pairs to approach the ideal invariant classifier. (a) Eight pairs connect two training domains (shapes); colors denote classes. (b) Their difference covariance estimates an unstable direction; the red arc marks estimation error. (c) Projecting before fitting improves on ERM here, but GRIT retains some sensitivity from estimation error (shading). Test points ( $+$ ) shift along the true unstable direction, reversing its correlation with the label. The ideal and GRIT boundaries are parallel to the true and estimated directions in (b), respectively, making each classifier invariant to shifts along its corresponding direction. Arrows show classifier normals: ${\bm{w}}_{\mathcal{S}}^{*}$ for the ideal classifier and ${\bm{w}}_{1}$ for rank-one GRIT; shading marks disagreement with the ideal boundary.

Panel of Table 1 — Table 1: Mean OOD accuracy (%, ten seeds): ColoredMNIST-INV test accuracy and Waterbirds-INV worst-group accuracy. Within each objective and dataset, bold / italic mark the best/second-best mean. Solid/dotted boxes mark the best/second-best mean among all entries for each dataset. PC: prediction consistency; +PC and +GRIT add the corresponding pair-based method to the row objective. Appendix F.3 reports sample standard deviations.

Figure 2: Stronger signal in the weakest direction and less noise improve recovery and target accuracy; more pairs yield near-

k^{-1/2}

subspace error and a GRIT advantage over PC at larger budgets. On CueMNIST, the top row shows GRIT’s subspace error; the bottom compares target accuracy. Defaults are

k=256

\sigma_{\xi}=0.1

, and clean spectrum

(1,1)

. Top-panel legends vary a second setting; bottom panels use the defaults. Bands show sample SD across 10 replicates; the dotted line is a shape guide. Y-axis ranges differ by panel.

Panel of Table 2 — Table 2: Combining two pair types improves UrbanCars robustness. $n_{B},n_{O}$ count pairs that vary the background and object, respectively. B/O shift changes the indicated factor’s label alignment, holding the other aligned; worst group is the minimum across eight label–background–object groups. Accuracy (%); bold marks each row’s best method per metric. Solid/dotted boxes mark the highest/second-highest distinct mean across all allocations and methods within each metric. +PC/+GRIT add prediction consistency/GRIT to ERM.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

ColoredMNIST-INV (CLIP feat.)
	Vanilla		+PC		+GRIT
Objective	Mean	s.d.	Mean	s.d.	Mean	s.d.
ERM	17.02	2.68	57.16	0.47	67.89	0.55
SD	17.02	2.68	56.68	0.31	67.77	0.62
RDM	39.95	3.47	59.55	0.89	67.60	0.58
IRMv1	59.89	2.71	61.28	1.19	70.27	0.41

Appendix

Table 3: Expanded core comparison: means and sample standard deviations (s.d.) over ten final training seeds. +PC and +GRIT add prediction consistency and GRIT to each objective.

$n_{B}$	$n_{O}$	Method	B shift	O shift	Both shift	Worst group	Train-weighted
256	0	ERM	45.4 $\pm$ 5.4	76.0 $\pm$ 11.7	28.5 $\pm$ 12.4	20.4 $\pm$ 9.5	83.1 $\pm$ 14.9
		PC	85.7 $\pm$ 1.4	65.1 $\pm$ 0.9	55.1 $\pm$ 1.0	52.9 $\pm$ 1.6	90.2 $\pm$ 0.5
		GRIT	88.5 $\pm$ 0.6	68.4 $\pm$ 0.8	60.3 $\pm$ 1.9	56.7 $\pm$ 1.5	92.3 $\pm$ 0.8
128	128	ERM	45.4 $\pm$ 5.4	76.0 $\pm$ 11.7	28.5 $\pm$ 12.4	20.4 $\pm$ 9.5	83.1 $\pm$ 14.9
		PC	76.0 $\pm$ 0.8	84.4 $\pm$ 0.5	71.7 $\pm$ 1.3	68.3 $\pm$ 1.3	87.5 $\pm$ 0.6
		GRIT	79.1 $\pm$ 0.4	85.1 $\pm$ 0.8	74.8 $\pm$ 0.9	71.9 $\pm$ 0.9	89.3 $\pm$ 0.9

Appendix

Table 4: Mixed pair types improve worst-group accuracy under worst-group validation. Test accuracy (%), mean

\pm

sample SD, for all five allocations, sorted by

n_{O}

. PC denotes prediction consistency. “Train-weighted” reweights test groups by training frequencies; it is a test metric, distinct from the validation selector. ERM uses no pairs and repeats for comparison. Bold marks the highest mean for each allocation and metric.

$n_{B}$	$n_{O}$	Method	B shift	O shift	Both shift	Worst group	Train-weighted
256	0	ERM	46.4 $\pm$ 2.9	86.5 $\pm$ 1.7	16.4 $\pm$ 1.8	15.3 $\pm$ 1.5	94.6 $\pm$ 0.7
		PC	70.8 $\pm$ 1.4	81.7 $\pm$ 0.6	32.0 $\pm$ 1.8	30.0 $\pm$ 2.9	97.2 $\pm$ 0.1
		GRIT	87.3 $\pm$ 0.6	69.0 $\pm$ 0.7	51.8 $\pm$ 0.4	49.3 $\pm$ 1.3	95.6 $\pm$ 0.2
128	128	ERM	46.4 $\pm$ 2.9	86.5 $\pm$ 1.7	16.4 $\pm$ 1.8	15.3 $\pm$ 1.5	94.6 $\pm$ 0.7
		PC	70.7 $\pm$ 1.4	81.8 $\pm$ 0.7	32.0 $\pm$ 1.9	29.9 $\pm$ 2.9	97.2 $\pm$ 0.1
		GRIT	80.8 $\pm$ 0.6	78.8 $\pm$ 0.2	53.2 $\pm$ 0.5	51.7 $\pm$ 0.8	95.5 $\pm$ 0.2

Appendix

Table 5: Training-weighted validation reduces the benefit of mixing pair types and changes the collection ordering. Test accuracy (%), mean

\pm

sample SD, for all five allocations, sorted by

n_{O}

Validation	Method	B shift	O shift	Both shift	Worst group	Train-weighted
WG	ERM	45.4 $\pm$ 5.4	76.0 $\pm$ 11.7	28.5 $\pm$ 12.4	20.4 $\pm$ 9.5	83.1 $\pm$ 14.9
	ERM, matched	66.0 $\pm$ 4.0	82.5 $\pm$ 0.8	25.8 $\pm$ 4.3	20.6 $\pm$ 6.6	96.8 $\pm$ 0.5
	Reference	84.0 $\pm$ 1.2	83.8 $\pm$ 0.7	82.0 $\pm$ 0.4	77.4 $\pm$ 0.9	88.9 $\pm$ 0.9
TW	ERM	46.4 $\pm$ 2.9	86.5 $\pm$ 1.7	16.4 $\pm$ 1.8	15.3 $\pm$ 1.5	94.6 $\pm$ 0.7
	ERM, matched	69.4 $\pm$ 1.1	81.5 $\pm$ 0.3	31.3 $\pm$ 0.8	27.4 $\pm$ 1.9	97.1 $\pm$ 0.0
	Reference	93.7 $\pm$ 0.3	66.4 $\pm$ 0.6	62.5 $\pm$ 0.6	61.4 $\pm$ 1.1	95.0 $\pm$ 0.2

Appendix

Table 6: Optimizer and pair-budget controls on UrbanCars. Accuracy (%), mean

\pm

sample SD over ten final seeds. WG and TW identify worst-group and training-weighted validation selection; the final column is training-weighted test accuracy. Controls use no allocation-specific pair banks.

Figure 3: Stronger signal in the weakest direction and less noise improve recovery and target accuracy; more pairs yield near-

k^{-1/2}

subspace error and a GRIT advantage over PC at larger budgets. On the Gaussian synthetic task, the top row shows GRIT’s subspace error; the bottom compares target accuracy. Defaults are

k=256

\sigma_{\xi}=0.1

, and clean spectrum

(1,1)

. Top-panel legends vary a second setting; bottom panels use the defaults. Bands show sample SD across 20 replicates; the dotted line is a shape guide. Y-axis ranges differ by panel.

Figure 4: GRIT’s fitted instability follows the recovery trends and decreases faster than subspace error as pair count grows. Mirror of the three sweeps on CueMNIST, adding normalized predictor sensitivity to the subspace-error curves. Defaults are

k=256

\sigma_{\xi}=0.1

, and clean spectrum

(1,1)

; bands show sample SD. The dotted

k^{-1/2}

line describes subspace recovery, not a predictor-rate guarantee.

Figure 5: GRIT’s fitted instability follows the recovery trends and decreases faster than subspace error as pair count grows. Mirror of the three sweeps on the Gaussian synthetic task, adding normalized predictor sensitivity to the subspace-error curves. Defaults are

k=256

\sigma_{\xi}=0.1

, and clean spectrum

(1,1)

; bands show sample SD. The dotted

k^{-1/2}

line describes subspace recovery, not a predictor-rate guarantee.

OOD Generalization

Out-of-distribution generalization studies how machine-learning models perform when test data, domains, tasks, or environments differ from training.

See all 395 papers

Feb 26, 2025

Invariance Pair Guidance: Robustness to Spurious Correlations via Corrective Gradients

Machine learning models are inherently bound to the distribution of the training data, often exploiting non-causal shortcuts. As a result, achieving robustness to spurious correlations remains a challenge. While existing approaches rely on data manipulation or re-weighting strategies to achieve robustness, they typically require dense group labels, multiple training domains, or specialized pre-processing. We propose Invariance Pair Guidance (IPG), a method to mitigate reliance on spurious correlations using a sparse set of counterfactual pairs. Unlike other methods demanding extensive supervision, IPG utilizes a novel dual-update mechanism to dynamically correct the optimization trajectory. We generate input pairs that isolate the spurious attribute to define the invariance, a characteristic that should not affect the outcome of the model. Based on these pairs, we define a corrective gradient that complements the traditional gradient descent approach. The correction adapts via a predefined invariance condition. Experiments on ColoredMNIST, Waterbirds-100, and CelebA datasets demonstrate the effectiveness of our approach and its robustness to group shifts, supported by a theoretical convergence analysis. IPG offers a data-efficient and theoretically grounded path to robustness.

Martin Surner, Abdelmajid Khelil, Ludwig Bothmann

Institute for Data and Process Science, Landshut University of Applied Sciences, Landshut, Germany. · Computer Science Department, Landshut University of Applied Sciences, Landshut, Germany. · Department of Statistics, LMU Munich, Munich, Germany. +1

Spurious Correlation Robustness Invariant Representation Learning

Ask PDF

Sep 25, 2026

Can Circuit Alignment Predict OOD Generalization?

Can out-of-distribution (OOD) generalization be predicted from a trained model's weights alone, without any target-domain data? Existing representational similarity metrics (CKA, SVCCA, RSA) compare activations rather than forecast generalization. We show they are provably insensitive to structural rerouting in the computational graph, the very change distribution shift induces. We close this gap with the Circuit Alignment Score (CAS), which compares class-specific circuits across domains via graph kernels, decomposed into same-class coherence and cross-class confusion. Casting CAS as a Lebesgue integral over the domain distribution, we prove its Monte Carlo estimate recovers the ground-truth ranking of learners by OOD accuracy, with pairwise inversion error vanishing at rate

O(1/M)

, where

M

is the number of sampled domains. Across

48

learners on PACS, CAS attains

0.88

rank correlation with OOD accuracy, versus

0.58

(CKA),

0.23

(SVCCA), and

0.14

(RSA), with similar trends on other benchmarks and even against data-dependent methods, making it the first provably consistent predictor of distributional robustness requiring neither target-domain data nor labels. The code is available at: https://github.com/ayanban011/ACE

Ayan Banerjee, Abhra Chaudhuri, Josep Llados +2

Computer Vision Center, Universitat Autònoma de Barcelona · Fujitsu Research of Europe · Indian Statistical Institute, Kolkata +1

Distribution Shift Robustness OOD Generalization Neural Network Generalization Representation Alignment

Ask PDF

Apr 9, 2026

Adversarial Label Invariant Graph Data Augmentations for Out-of-Distribution Generalization

Out-of-distribution (OoD) generalization occurs when representation learning encounters a distribution shift. This occurs frequently in practice when training and testing data come from different environments. Covariate shift is a type of distribution shift that occurs only in the input data, while the concept distribution stays invariant. We propose RIA - Regularization for Invariance with Adversarial training, a new method for OoD generalization under convariate shift. Motivated by an analogy to

Q

-learning, it performs an adversarial exploration for counterfactual data environments. These new environments are induced by adversarial label invariant data augmentations that prevent a collapse to an in-distribution trained learner. It works with many existing OoD generalization methods for covariate shift that can be formulated as constrained optimization problems. We develop an alternating gradient descent-ascent algorithm to solve the problem in the context of causally generated graph data, and perform extensive experiments on OoD graph classification for various kinds of synthetic and natural distribution shifts. We demonstrate that our method can achieve high accuracy compared with OoD baselines.

Simon Zhang, Ryan P. DeMilt, Kun Jin +1

Department of Computer Science, Purdue University, West Lafayette, USA · Department of Computer Science and Engineering, The Ohio State University, Columbus Ohio, USA · Department of Industrial and Systems Engineering, The Ohio State University, Columbus Ohio, USA

OOD Generalization Covariate Shift Graph Data Augmentation Data Augmentation Counterfactual Data Augmentation Adversarial Training

Ask PDF

Abstract

OOD Generalization

Explore similar work

Invariance Pair Guidance: Robustness to Spurious Correlations via Corrective Gradients

Can Circuit Alignment Predict OOD Generalization?

Adversarial Label Invariant Graph Data Augmentations for Out-of-Distribution Generalization