ColoredMNIST: OOD Benchmark for Robustness

• ColoredMNIST is a synthetic benchmark that probes spurious correlations by associating digit shapes with colors in controlled environments.
• It supports both binary and multiclass settings by systematically altering the relationship between invariant digit features and spurious color cues.
• The benchmark enables rigorous evaluation of OOD generalization by comparing standard ERM with invariance-enhancing algorithms under reversed spurious associations.

ColoredMNIST is a synthetic benchmark derived from MNIST, explicitly designed to probe the tendency of machine learning models—especially deep neural networks—to rely on spurious correlations. It is a canonical testbed for evaluating the robustness and out-of-distribution (OOD) generalization of algorithms, offering precise control over the association between invariant and spurious features. The core idea is to correlate digit class labels (“shape”) with color in the training data, then reverse or modify this association in the test data to assess whether models have learned robust, invariant representations or have overfit to the shortcut (color) (Salaudeen et al., 31 Mar 2025, Choe et al., 2020, Robey et al., 2021).

1. Formal Construction and Generative Process

ColoredMNIST is constructed by modifying the binary or multiclass MNIST dataset as follows:

• Let $Y$ denote the label: $Y=1$ if digit $\geq5$, $Y=0$ otherwise (binary), or $Y\in\{0,\dots,9\}$ (multiclass).
• Apply random label noise (e.g., 25% chance of flipping) to decorrelate shape and true class, bounding attainable accuracy.
• Assign a color, $C\in\{\text{red},\text{green}\}$, to each image according to a distribution $p_e(c|y)$ parameterized by the “environment” $e$. Typically, $P_e(C=y|y)=p_e$.
• Generate environments by varying $p_e$:
  • In training, $Y=1$0 is high (e.g., 0.8, 0.9), so color is highly but spuriously predictive of the label.
  • In OOD test, $Y=1$1 is set low (e.g., 0.1), reversing or weakening the spurious correlation (Salaudeen et al., 31 Mar 2025, Robey et al., 2021, Choe et al., 2020).

Formally, the observed feature space $Y=1$2 decomposes as $Y=1$3, where $Y=1$4 is digit shape (domain-general) and $Y=1$5 is color (spurious). ColoredMNIST enables explicit control of spurious correlation strength and supports both binary and multiclass variants.

2. Spurious and Invariant Feature Formalism

ColoredMNIST operationalizes the distinction between:

• Domain-general (invariant) features ($Y=1$6): features whose conditional distribution $Y=1$7 is identical across environments.
• Spurious (domain-specific) features ($Y=1$8): features for which $Y=1$9 changes across environments, and whose predictive association with $\geq5$0 can be manipulated via $\geq5$1.

The platform supports a formal test of model robustness: a well-specified OOD benchmark occurs when the optimal ERM predictor using both $\geq5$2 and $\geq5$3 fails on the OOD environment, while the optimal predictor using only $\geq5$4 (digit) generalizes—this requires a shift or reversal in spurious correlation (Salaudeen et al., 31 Mar 2025).

3. Benchmark Protocols and OOD Evaluation

A typical ColoredMNIST experiment involves:

• Training on several environments with $\geq5$5, causing standard ERM to exploit color-label shortcuts.
• Evaluation on OOD environments with $\geq5$6 (spurious reversal) or $\geq5$7 (spurious persists).
• Generation of large pools of predictors (varying hyperparameters) to analyze the relationship between in-domain and OOD accuracy (“accuracy on the line”): perfect positive line ($\geq5$8) indicates misspecification, while negative correlation ($\geq5$9) indicates a well-specified, spurious-exposing benchmark (Salaudeen et al., 31 Mar 2025).

Quantitative metrics such as accuracy, OOD accuracy, and the “worst-group” or “worst-domain” scores are reported. The theoretical maximum OOD accuracy is determined by shape–label correlation (e.g., $Y=0$0 under 25% label noise).

4. Algorithmic Approaches and Results

ColoredMNIST is used to evaluate diverse domain generalization and OOD-robustness algorithms, including:

• Empirical Risk Minimization (ERM): Catastrophically fails in OOD tests, with accuracies often below chance when spurious correlation is reversed (e.g., 16.9% (Le et al., 22 Apr 2025), 36.7% (Robey et al., 2021)).
• Invariant Risk Minimization (IRM/IRMv1): Penalizes predictors whose risks vary across environments; achieves up to 70.6% OOD accuracy given substantial spurious-correlation diversity (Choe et al., 2020).
• Model-Based Domain Generalization (MBDG): Formulates domain generalization as a semi-infinite constrained optimization enforcing feature invariance under known color transformation; achieves 68.5% OOD accuracy, a ~30pp improvement over ERM (Robey et al., 2021).
• Rich Feature Construction (Bonsai): Iterative feature discovery and distillation across “hard” subpopulations yields significant improvements for both ERM and OOD-robust methods (e.g., IRM+Bonsai-cf up to 69.9%) (Zhang et al., 2022).
• Invariance Pair-Guided Learning (IPG): Enforces output invariance via explicit red–green digit pairs; reaches 72.7% OOD accuracy, within 2.3pp of the label-noise ceiling and above prior baselines (Surner et al., 26 Feb 2025).
• Language-Guided Attention Alignment: Uses pseudo-attention maps generated by large vision-LLMs (e.g., CLIP) to regularize model saliency toward shape rather than color; 64.9% OOD test accuracy achieved for zero-supervision attention guidance (Yang et al., 25 Sep 2025).

A summary of representative OOD test results (accuracy on reversed correlation) is tabulated below:

Method	OOD Test Accuracy (%)	Reference
ERM	16.1–36.7	(Le et al., 22 Apr 2025, Robey et al., 2021)
IRM	66.9–70.6	(Choe et al., 2020, Zhang et al., 2022)
MBDG (full PD)	68.5	(Robey et al., 2021)
Bonsai (IRM+Bonsai-cf)	69.9	(Zhang et al., 2022)
IPG (no aug)	72.7	(Surner et al., 26 Feb 2025)
Attention-align (VL)	64.9	(Yang et al., 25 Sep 2025)
Oracle (gray ERM)	72.7–73.1	(Le et al., 22 Apr 2025, Surner et al., 26 Feb 2025)

5. Critique, Pathologies, and Benchmark Specification

A key insight established is that many ColoredMNIST (and similar) benchmarks are frequently misspecified—i.e., they fail to reliably penalize reliance on spurious features:

• If $Y=0$1 continues the spurious correlation direction present in training ($Y=0$2 when $Y=0$3), both ID and OOD accuracy are tightly correlated across models (“accuracy on the line”), and ERM does not fail OOD. Such splits do not test for robust invariance (Salaudeen et al., 31 Mar 2025).
• Well-specified splits require a reversal of spurious correlation between train and test: only then does the domain-invariant predictor outperform the best ID ERM model, yielding negative (or non-positive) correlation between in-domain and OOD accuracy.

Beyond the correlation metrics, benchmark construction must prioritize:

• Explicit inclusion of spurious reversal environments.
• Screening splits using large-scale ID/OOD accuracy correlation analysis.
• Focusing on worst-case OOD error rather than average-case.

Colloquially, benchmarks that do not induce failure when models exploit the spurious signal (e.g., color) do not serve their intended purpose in evaluating invariance.

6. Methods for Environment Construction and Annotation-Free Learning

Classic ColoredMNIST assumes access to environment labels during training to enable methods like IRM or GroupDRO. Recent research demonstrates:

• Annotation-free environment inference: By clustering the penultimate-layer features from an ERM-trained CNN, environments dominated by color (spurious) are naturally discovered (purity >99%), and minority/conflict samples can be rebalanced or paired to yield datasets with effectively inverted or nulled spurious association (Le et al., 22 Apr 2025).
• Training IRM on these inferred environments yields OOD robustness on par with supervised-environment methods (e.g., 68.0% vs. 66.9% for IRM with hand-crafted environments) (Le et al., 22 Apr 2025).
• This procedure enables fully automated robust learning pipelines on ColoredMNIST without any manual group label annotation.

7. Impact on Model Comparison, Theory, and Future Benchmark Design

ColoredMNIST has influenced both empirical practice and theoretical understanding of OOD generalization:

• It demonstrates the brittleness of deep models to spurious cues and motivates penalty, invariance, and data-driven solutions.
• Theoretical results formalize when OOD benchmarks genuinely reveal spurious reliance, situating ColoredMNIST as both a cautionary example (susceptible to misspecification) and a testbed for novel algorithms.
• Best practices recommend: always ensuring spurious correlation reversal in test environments, evaluating correlation between ID and OOD accuracy to reveal benchmark flaws, and constructing synthetic benchmarks that allow direct control over the strength and direction of spurious features (Salaudeen et al., 31 Mar 2025, Robey et al., 2021).

ColoredMNIST continues to be used in the evaluation of new methodologies in OOD generalization, robust learning, and causal representation, owing to its controllable and analyzable structure. Its limitations have driven advances in both benchmark design and the practical deployment of domain-robust algorithms.