Sparse Invariant Risk Minimization

Xiao ZhouYong LinWeizhong ZhangTong Zhang

article2022ICML91 citations

Proposes Sparse Invariant Risk Minimization, a method that enforces continuous global sparsity constraints during training to prevent overparameterized networks from learning spurious features under distribution shifts, outperforming existing techniques by up to 29% in accuracy.

Listen

Modern deep learning models frequently fail when deployed in real-world environments due to distributional shifts—situations where testing conditions differ from training data. While Invariant Risk Minimization was designed to solve this by learning stable, causal relationships across varying environments, standard deep neural networks suffer from severe overfitting when overparameterized. Modern architectures require massive parameter counts to ensure smooth training, yet this high capacity causes invariant learning techniques to mistakenly memorize deceptive, spurious correlations. As a result, increasing the size of conventional invariant models causes their real-world generalization performance to collapse.

The main objective of the article is to demonstrate that overparameterization systematically breaks traditional invariant learning methods and to introduce Sparse Invariant Risk Minimization, a framework that preserves robust generalization in large models by maintaining sparsity throughout training.

The authors evaluate this issue through theoretical proofs in linear systems and empirical benchmarks across multi-layer neural networks and deep architectures, such as ResNet-18. They test the methodology using four vision datasets—ColoredMNIST, FullColoredMNIST, ColoredObject, and CIFARMNIST—where spurious correlations deliberately reverse between training and evaluation phases. The study benchmarks the proposed technique against standard Empirical Risk Minimization, classical Invariant Risk Minimization variants, Bayesian approaches, and post-training pruning methods like MRM.

The evaluation reveals four primary findings. First, as standard invariant models grow in size, their test performance degrades significantly; for instance, expanding hidden layer dimensions decreased test accuracy by up to 27 percentage points due to spurious feature reliance. Second, pruning models after dense training fails because misleading features are already embedded and true invariant signals are discarded. Third, the proposed technique successfully resists overparameterization, improving test accuracy steadily as network capacity grows and outperforming prior state-of-the-art methods by up to 29 percentage points across complex vision tasks. Fourth, ablation studies confirm that the new approach prevents spurious information from entering the network representations entirely, yielding near-optimal feature filtering.

These results establish that enforcing a global parameter budget during the entire training cycle prevents the network from fitting abundant spurious correlations, forcing it to isolate true causal features instead. This insight provides a practical safeguard against catastrophic model failures in safety-critical, out-of-distribution deployments, such as medical diagnostics and automated driving. It bridges the fundamental conflict between model trainability and out-of-distribution generalization in modern artificial intelligence.

Organizations developing models under distribution shifts should transition from post-training compression routines to continuous sparse training frameworks. Development teams should adopt parameter budgets to naturally restrict feature memorization during optimization. Future work should focus on extending this paradigm to large-scale natural language processing, self-supervised learning, and vision-language foundation models.

A primary limitation of this study is that element-wise weight sparsity currently offers limited hardware acceleration on standard deep learning platforms such as PyTorch and TensorFlow, making training computationally intensive. However, given the consistent theoretical backing and experimental verification across varying model architectures, confidence in the method's ability to eliminate spurious features remains very high.

Google The Hong Kong University of Science and Technology deep learning feature selection genetic selection sparse invariant risk minimization global sparsity constraint invariant features

Zhou et al (2022).pdf

Paper: Invariant Risk Minimization, Martin Arjovsky et al. (2019). SparseIRM modifies IRM’s invariant-feature objective, so first learn how IRM defines invariance across environments and implements it with a penalty.

No sufficiently relevant recommendations were found.

Cover for Sparse Invariant Risk Minimization

Abstract

Invariant Risk Minimization (IRM) is an emerging invariant feature extracting technique to help generalization with distributional shift. However, we find that there exists a basic and intractable contradiction between the model trainability and generalization ability in IRM. On one hand, recent studies on deep learning theory indicate the importance of large-sized or even overparameterized neural networks to make the model easy to train. On the other hand, unlike empirical risk minimization that can be benefited from overparameterization, our empirical and theoretical analyses show that the generalization ability of IRM is much easier to be demolished by overfitting caused by overparameterization. In this paper, we propose a simple yet effective paradigm named Sparse Invariant Risk Minimization (SparseIRM) to address this contradiction. Our key idea is to employ a global sparsity constraint as a defense to prevent spurious features from leaking in during the whole IRM process. Compared with sparsify-after-training prototype by prior work which can discard invariant features, the global sparsity constraint limits the budget for feature selection and enforces SparseIRM to select the invariant features. We illustrate the benefit of SparseIRM through a theoretical analysis on a simple linear case. Empirically we demonstrate the power of SparseIRM through various datasets and models and surpass state-of-the-art methods with a gap up to 29%.

Topics

deep learning feature selection genetic selection sparse approximations sparse recovery scale invariant optimization

Keywords

sparse invariant risk minimization global sparsity constraint invariant features distributional shifts spurious features feature mask feature selection

1. Introduction
2. Related Work
2.1. Causality and invariance
2.2. Sparsity in Deep Neural Networks
2.3. Overparameterized Deep Neural Networks
3. Investigating the Effects of Overparameterization on IRM
3.1. Preliminaries
3.2. Analysis of Overparameterized IRM in a Linear Case
4. Sparse Invariant Risk Minimization
4.1. SparseIRM Framework
4.2. Understanding the Benefits of SparseIRM through Theoretical Analysis
5. Experiment
5.1. MLP on ColoredMNIST/FullColoredMNIST
5.2. ResNet18 on ColoredObject/CIFARMNIST
5.3. Ablation Studies
6. Limitations
7. Conclusion
Acknowledgements
References
Supplemental Material: Sparse Invariant Risk Minimization
A. Dataset Details
B. Proofs
B.1. More Discussions on Our Theoretical Analysis
B.2. The Relationship between Problem (6) and Problem (8)
B.3. Basics
B.4. Proof for Proposition 1
B.5. Proof of Corollary 1
B.6. Proof of Theorem 1
B.6.1. PRELIMINARIES
B.6.2. ASSUMPTIONS AND TECHNICAL LEMMAS
B.6.3. PROOF OF THE MAIN THEOREM
C. Algorithm
D. Experimental Configurations
E. Additional Experimental Results
E.1. Training and Testing Accuracy of MLP on FullColoredMNIST50000 with varying hidden dimensions
E.2. MLP on ColoredMNIST20000 and FullColoredMNIST20000 with Varying Hidden Dimensions
F. Future Directions

Knowls

Knowl 1 — SparseIRM imposes a global sparsity budget throughout IRM training
model/method
Sparse Invariant Risk Minimization (SparseIRM) trains an invariant-risk-minimization model while constraining how many weights or features may be active, rather than training a dense model and sparsifying it afterward. For training environments $e\in\mathcal E_{\mathrm{tr}}$ , let $R^e(w)$ be the empirical prediction loss in environment $e$ , and let $J(w)$ be an IRM penalty, such as the squared norm of environment-specific classifier gradients or the variance of environment risks. The objective is to minimize $\sum_{e\in\mathcal E_{\mathrm{tr}}}R^e(w)+\lambda J(w)$ subject to a binary mask $m\in\{0,1\}^{d_\Phi}$ satisfying $\lVert m\rVert_1\le K$ ; the masked model uses $m\circ\Phi$ , where $\Phi$ denotes the weights or features being selected and $K$ is the global active-unit budget. The mask is constrained during the full optimization, so the model must fit the training environments using a limited subnetwork. The proposed rationale is that when spurious and random features greatly outnumber invariant features, the budget prevents the model from retaining all such alternatives and encourages it to use invariant features.
Knowl 2 — SparseIRM optimizes a stochastic mask with projected Gumbel-Softmax updates
algorithm
SparseIRM relaxes its discrete mask optimization by assigning each mask coordinate an independent Bernoulli distribution with probability vector $s\in[0,1]^{d_\Phi}$ and enforcing $\mathbf 1^\top s\le K$ . It minimizes the expected IRM loss over masks sampled from this distribution, using projected stochastic gradient descent and a Gumbel-Softmax relaxation.
Input: Training environments, a dense network, target remaining ratio k_f = 0.5, training epochs T.
Initialize network weights; initialize mask probabilities s to 1; set K = k_f d_w and temperature τ = 1.
For each epoch from 1 to T:
    For each iteration:
        Form a minibatch by combining minibatches from the training environments.
        Sample independent Gumbel(0, 1) vectors g1 and g0.
        Form relaxed mask probabilities as sigmoid((logit(s) + g1 - g0) / τ).
        Compute the minibatch IRM loss using the relaxed mask.
        Take a gradient step on s and project it onto {s in [0,1]^dΦ : 1ᵀs ≤ K}.
        Take a gradient step on the network weights using the masked minibatch IRM loss.
    End iteration loop.
End epoch loop.
Sample a binary mask m from the Bernoulli distribution parameterized by s.
Output the masked network.
Here $d_\Phi$ is the number of mask coordinates and $d_w$ is the network-parameter dimension used in the paper's budget definition. The paper's reported training configurations use Adam for mask probabilities, a probability learning rate of $6\times10^{-3}$ , an IRM penalty weight of $10000$ , and cosine learning-rate scheduling. For the network weights, ColoredMNIST and FullColoredMNIST use Adam with learning rate $0.0004$ ; CIFARMNIST and ColoredObject use SGD with learning rate $0.01$ and momentum $0.9$ .
Knowl 3 — Overparameterized linear IRM can prefer all features or only spurious and random features
theoretical result
The paper analyzes a finite-sample linear setting with two training environments. Each input contains invariant features $x_{\mathrm{inv}}\in\mathbb R^{d_{\mathrm{inv}}}$ , spurious features $x_s\in\mathbb R^{d_s}$ , and random features $x_r\in\mathbb R^{d_r}$ . The label is generated as $y=\gamma^\top x_{\mathrm{inv}}+\epsilon_{\mathrm{inv}}$ , the spurious features as $x_s=y\mathbf 1+\alpha^e\circ\epsilon_s$ with environment-dependent coefficients $\alpha^e$ , and the random features as $x_r=\epsilon_r$ . The noises are independent, zero-mean sub-Gaussian variables with bounded variance. A binary feature selector $\Phi$ defines the linear predictor $f(x)=(\Phi\circ x)^\top v+b$ ; let $\widehat L(\Phi)$ denote the paper's empirical IRM objective after optimizing the linear predictor parameters.

If the total dimension $d=d_{\mathrm{inv}}+d_s+d_r$ exceeds the combined sample count $n_{e_1}+n_{e_2}$ , the all-feature selector $\Phi_{\mathrm{all}}$ satisfies $\widehat L(\Phi_{\mathrm{all}})=0\le\widehat L(\Phi_{\mathrm{inv}})$ , where $\Phi_{\mathrm{inv}}$ selects only invariant features. If the spurious-plus-random dimension $d_s+d_r$ itself exceeds the combined sample count, the selector $\Phi_{sr}$ that excludes all invariant features also satisfies $\widehat L(\Phi_{sr})=0\le\widehat L(\Phi_{\mathrm{inv}})$ . Thus, in this setting, empirical IRM minimization need not identify invariant features under overparameterization; sufficiently many non-invariant features can interpolate the training data even without invariant signal.
Knowl 4 — A feature budget yields invariant-feature selection in the analyzed linear setting
theoretical result
For the paper's linear data-generation setting, SparseIRM can identify the invariant feature mask with high probability when the active-feature budget equals the number of invariant features. Assume the same sample size $n$ in each of two training environments, and set $K=d_{\mathrm{inv}}$ . The guarantee uses the paper's technical conditions: bounded concentration of normalized features and positive bounded design-matrix eigenvalues; sub-Gaussian invariant-label noise; non-negligible signal and variance in every invariant feature; environment-dependent variation of each spurious-feature coefficient by at least $\Delta>0$ ; and nonzero label correlation for each spurious feature in at least one environment.

Under these assumptions, there are positive distribution-dependent constants $Q_1,Q_2$ such that, for failure probability $\delta\in(0,1)$ , if $n>Q_1+Q_2\ln(d/\delta)$ , then with probability at least $1-\delta$ the invariant-only selector has strictly smaller empirical objective than every other selector with at most $K$ active features: $\widehat L(\Phi_{\mathrm{inv}})<\widehat L(\Phi)$ for all $\Phi\ne\Phi_{\mathrm{inv}}$ with $\lVert\Phi\rVert_1\le K$ . The result is logarithmic in the total number of invariant, spurious, and random features, with constants depending on the stated distributional and signal conditions; it is a guarantee for the analyzed linear case, not a general deep-network theorem.
Knowl 5 — Sparsifying during training avoids a failure mode of prune-after-training IRM
model/method
The paper distinguishes SparseIRM from a sparsify-after-training strategy. In an overparameterized dense IRM model, training can first fit the labels using spurious or random features while assigning little useful weight to invariant features. A later pruning rule based on weight importance can then remove invariant features, leaving a subnetwork that cannot recover them. SparseIRM instead limits the active subnetwork throughout training, which is intended to prevent the dense model from first absorbing a large collection of shortcuts and noise features. The schematic comparison on page 5 illustrates this contrast: the SparseIRM path applies a mask during training, whereas the post-training path trains densely before pruning.
Knowl 6 — Benchmark experiments use spurious correlations that reverse at test time
experimental setup
The experiments evaluate generalization to a held-out environment in four benchmarks. Each dataset has two training environments and one test environment; the spurious-feature correlation with the label is strong in training and reversed in testing. ColoredMNIST uses digit as the invariant feature and color as the spurious feature, with correlation tuple $(0.9,0.8,0.1)$ and 25% label noise. FullColoredMNIST extends the task to 10 classes and uses $(0.999,0.7,0.1)$ with 20% label noise. ColoredObject uses object identity and background, with $(0.999,0.7,0.1)$ and 5% label noise. CIFARMNIST concatenates CIFAR10 images as invariant content with MNIST images as spurious content, with $(0.999,0.7,0.1)$ and 10% label noise. The tuples list the two training correlations followed by the testing correlation.

The MLP experiments use hidden dimensions from 64 to 640 and training-set sizes of 20,000 or 50,000. ResNet-18 experiments test more highly overparameterized settings on ColoredObject and CIFARMNIST. Comparisons include ERM, SparseERM, IRMv1, REx, the prune-after-training method MRM, BayesianIRM, and an Oracle ERM trained without spurious features. The MLP runs use 1,500 epochs; ResNet-18 runs use 75 epochs for ColoredObject and 50 for CIFARMNIST.
Knowl 7 — SparseIRM retains test accuracy as MLP width increases
empirical result
On the 20,000-example MLP benchmarks, increasing hidden width from 160 to 640 substantially reduces IRM and MRM test accuracy under IRMv1, while SparseIRM's accuracy rises. On ColoredMNIST, IRM changes from 62.42% to 36.80%, MRM from 65.00% to 53.58%, and SparseIRM from 65.91% to 68.54%. On FullColoredMNIST, the corresponding figures are 54.19% to 41.18% for IRM, 55.47% to 43.42% for MRM, and 59.34% to 63.26% for SparseIRM. These comparisons support the paper's claim that sparsity during training reduces the width-related overfitting seen in dense IRM and post-training sparsification.

The 50,000-example results show the same favorable scaling. On ColoredMNIST, SparseIRM under IRMv1 increases from 63.49% at hidden dimension 64 to 73.61% at dimension 640, while MRM changes from 62.02% to 69.32%. On FullColoredMNIST, SparseIRM increases from 65.06% to 69.47%, whereas MRM decreases from 62.50% to 55.91%. All values are Top-1 test accuracy.
Knowl 8 — SparseIRM improves ResNet-18 accuracy on ColoredObject and CIFARMNIST
empirical result
In the paper's ResNet-18 experiments, SparseIRM outperforms the listed IRM, MRM, and BayesianIRM baselines on both highly overparameterized benchmarks. On ColoredObject, IRMv1 SparseIRM reaches $87.4\pm0.6$ % Top-1 test accuracy, compared with IRMv1 IRM at $72.5\pm2.3$ %, MRM at $58.4\pm0.9$ %, and BayesianIRM at $78.1\pm0.6$ %; the Oracle is $87.9\pm0.3$ %. REx SparseIRM reaches $80.3\pm1.1$ %, versus REx IRM at $73.8\pm1.3$ % and MRM at $55.7\pm2.9$ %.

On CIFARMNIST, IRMv1 SparseIRM obtains $63.9\pm0.4$ %, compared with IRMv1 IRM at $51.3\pm3.0$ %, MRM at $56.7\pm2.3$ %, and BayesianIRM at $59.3\pm0.8$ %; the Oracle is $83.7\pm1.5$ %. REx SparseIRM obtains $62.7\pm0.6$ %, versus REx IRM at $50.1\pm2.2$ % and MRM at $52.6\pm1.5$ %. The results on page 8 show that the strongest reported gain is on ColoredObject, where IRMv1 SparseIRM approaches Oracle accuracy.
Knowl 9 — Representation probes show reduced encoding of spurious color
empirical result
The paper probes whether representations learned by an MLP on ColoredMNIST retain information about color, the spurious feature. A separate two-layer perceptron is trained to predict color from the learned representations using a dataset with 50% color-label correlation. The probe's training and test accuracies are $89.1\pm0.3$ % and $89.2\pm1.9$ % for IRM, and $82.8\pm1.1$ % and $83.4\pm0.4$ % for MRM. For SparseIRM they are $50.3\pm0.8$ % and $50.0\pm0.7$ %, approximately chance on this two-color task. This indicates that the SparseIRM representations retain little predictive color information in this experiment.

A second visualization compares representations of a test image before and after its color is flipped. On page 9, the absolute representation differences are visibly smaller for SparseIRM than for IRM and MRM, providing a qualitative complement to the color-prediction probe.
Knowl 10 — Weight-level sparsity is not yet an efficient training implementation
limitation
SparseIRM uses weight-level sparsity, which the paper notes is difficult to exploit efficiently in general neural-network training platforms such as TensorFlow and PyTorch. Consequently, the method does not automatically accelerate training despite using a sparse model. The authors state that training speed is not the main concern for the IRM datasets considered, but identify efficient sparse-training implementations as future work.

Coverage note — The paper's proof derivations, secondary training-accuracy tables, and future-work suggestions are omitted because they do not add load-bearing results beyond the stated guarantee, benchmark outcomes, and limitation.

References

1.Ahmed, F., Bengio, Y., van Seijen, H., and Courville, A. Systematic generalisation with group invariant predictions. In International Conference on Learning Representations, 2020.
2.Ahuja, K., Shanmugam, K., Varshney, K., and Dhurandhar, A. Invariant risk minimization games. In International Conference on Machine Learning, pp. 145–155. PMLR, 2020a.
3.Ahuja, K., Wang, J., Dhurandhar, A., Shanmugam, K., and Varshney, K. R. Empirical or invariant risk minimization? a sample complexity perspective. arXiv preprint arXiv:2010.16412, 2020b.
4.Ahuja, K., Caballero, E., Zhang, D., Bengio, Y., Mitliagkas, I., and Rish, I. Invariance principle meets information bottleneck for out-of-distribution generalization. arXiv preprint arXiv:2106.06607, 2021.
5.Arjovsky, M., Bottou, L., Gulrajani, I., and Lopez-Paz, D. Invariant risk minimization. arXiv preprint arXiv:1907.02893, 2019.
6.Bahdanau, D., Cho, K., and Bengio, Y. Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv:1409.0473, 2014.
7.Bai, H., Sun, R., Hong, L., Zhou, F., Ye, N., Ye, H.-J., Chan, S.-H. G., and Li, Z. Decaug: Out-of-distribution generalization via decomposed feature representation and semantic augmentation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 6705–6713, 2021a.
8.Bai, H., Zhou, F., Hong, L., Ye, N., Chan, S.-H. G., and Li, Z. Nas-ood: Neural architecture search for out-of-distribution generalization. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 8320–8329, 2021b.
9.Beery, S., Van Horn, G., and Perona, P. Recognition in terra incognita. In Proceedings of the European conference on computer vision (ECCV), pp. 456–473, 2018.
10.Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877–1901, 2020.
11.Chang, S., Zhang, Y., Yu, M., and Jaakkola, T. Invariant rationalization. In International Conference on Machine Learning, pp. 1448–1458. PMLR, 2020.
12.Chen, C., Shen, L., Huang, H., and Liu, W. Quantized adam with error feedback. ACM Transactions on Intelligent Systems and Technology (TIST), 12(5):1–26, 2021a.
13.Chen, K., Hong, L., Xu, H., Li, Z., and Yeung, D.-Y. Multisiam: Self-supervised multi-instance siamese representation learning for autonomous driving. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 7546–7554, 2021b.
14.Chen, Y., Rosenfeld, E., Sellke, M., Ma, T., and Risteski, A. Iterative feature matching: Toward provable domain generalization with logarithmic environments. arXiv preprint arXiv:2106.09913, 2021c.
15.Creager, E., Jacobsen, J.-H., and Zemel, R. Environment inference for invariant learning. In International Conference on Machine Learning, pp. 2189–2200. PMLR, 2021.
16.DeGrave, A. J., Janizek, J. D., and Lee, S.-I. Ai for radiographic covid-19 detection selects shortcuts over signal. Nature Machine Intelligence, pp. 1–10, 2021.
17.Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
18.Diao, S., Bai, J., Song, Y., Zhang, T., and Wang, Y. Zen: Pre-training chinese text encoder enhanced by n-gram representations. arXiv preprint arXiv:1911.00720, 2019.
19.Diao, S., Xu, R., Su, H., Jiang, Y., Song, Y., and Zhang, T. Taming pre-trained language models with n-gram representations for low-resource domain adaptation. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp. 3336–3349, 2021.
20.Gao, J., Zhou, Y., Yu, P. L., Joty, S., and Gu, J. Unison: Unpaired cross-lingual image captioning. 2022.
21.Geirhos, R., Jacobsen, J.-H., Michaelis, C., Zemel, R., Brendel, W., Bethge, M., and Wichmann, F. A. Shortcut learning in deep neural networks. Nature Machine Intelligence, 2(11):665–673, 2020.
22.Grill, J.-B., Strub, F., Altche, F., Tallec, C., Richemond, ´ P. H., Buchatskaya, E., Doersch, C., Pires, B. A., Guo, Z. D., Azar, M. G., and others. Bootstrap your own latent: A new approach to self-supervised learning. arXiv:2006.07733, 2020.
23.Gu, J., Cai, J., Joty, S. R., Niu, L., and Wang, G. Look, imagine and match: Improving textual-visual cross-modal retrieval with generative models. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 7181–7189, 2018.
24.Gu, Y., Zhang, W., Fang, C., Lee, J. D., and Zhang, T. How to characterize the landscape of overparameterized convolutional neural networks. Advances in Neural Information Processing Systems, 2020.
25.Gulrajani, I. and Lopez-Paz, D. In search of lost domain generalization. arXiv preprint arXiv:2007.01434, 2020.
26.Han, S., Mao, H., and Dally, W. J. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. International Conference on Learning Representations, 2016.
27.He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770–778, 2016.
28.He, K., Fan, H., Wu, Y., Xie, S., and Girshick, R. Momentum contrast for unsupervised visual representation learning. In CVPR, 2020.
29.Hsu, D., Kakade, S. M., and Zhang, T. Random design analysis of ridge regression. In Conference on learning theory, pp. 9–1. JMLR Workshop and Conference Proceedings, 2012.
30.Huang, M., Huang, Z., Li, C., Chen, X., Xu, H., Li, Z., and Liang, X. Arch-graph: Acyclic architecture relation predictor for task-transferable neural architecture search. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11881–11891, 2022.
31.Jacot, A., Gabriel, F., and Hongler, C. Neural tangent kernel: Convergence and generalization in neural networks. arXiv preprint arXiv:1806.07572, 2018.
32.Jacot, A., Gabriel, F., and Hongler, C. Neural tangent kernel: convergence and generalization in neural networks. In Proceedings of the 53rd Annual ACM SIGACT Symposium on Theory of Computing, pp. 6–6, 2021.
33.Jin, W., Barzilay, R., and Jaakkola, T. Domain extrapolation via regret minimization. arXiv preprint arXiv:2006.03908, 2020.
34.Kamath, P., Tangella, A., Sutherland, D., and Srebro, N. Does invariant risk minimization capture invariance? In International Conference on Artificial Intelligence and Statistics, pp. 4069–4077. PMLR, 2021.
35.Kawaguchi, K., Huang, J., and Kaelbling, L. P. Effect of depth and width on local minima in deep learning. Neural computation, 31(7):1462–1498, 2019.
36.Krizhevsky, A., Sutskever, I., and Hinton, G. E. Imagenet classification with deep convolutional neural networks. In Advances in neural information processing systems, pp. 1097–1105, 2012.
37.Krueger, D., Caballero, E., Jacobsen, J.-H., Zhang, A., Binas, J., Zhang, D., Le Priol, R., and Courville, A. Out-of-distribution generalization via risk extrapolation (rex). In International Conference on Machine Learning, pp. 5815–5826. PMLR, 2021.
38.Kusupati, A., Ramanujan, V., Somani, R., Wortsman, M., Jain, P., Kakade, S., and Farhadi, A. Soft threshold weight reparameterization for learnable sparsity. In Proceedings of the International Conference on Machine Learning, July 2020.
39.Lin, Y., Lian, Q., and Zhang, T. An empirical study of invariant risk minimization on deep models. ICML 2021 Workshop on Uncertainty and Robustness in Deep Learning, 2021.
40.Lin, Y., Dong, H., Wang, H., and Zhang, T. Bayesian invariant risk minimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16021–16030, 2022a.
41.Lin, Y., Zhu, S., and Cui, P. Zin: When and how to learn invariance by environment inference? arXiv preprint arXiv:2203.05818, 2022b.
42.Liu, J., Hu, Z., Cui, P., Li, B., and Shen, Z. Heterogeneous risk minimization. arXiv preprint arXiv:2105.03818, 2021a.
43.Liu, J., Hu, Z., Cui, P., Li, B., and Shen, Z. Kernelized heterogeneous risk minimization. arXiv preprint arXiv:2110.12425, 2021b.
44.Liu, Y., Ott, M., Goyal, N., Du, J., Joshi, M., Chen, D., Levy, O., Lewis, M., Zettlemoyer, L., and Stoyanov, V. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692, 2019.
45.Liu, Z., Han, J., Chen, K., Hong, L., Xu, H., Xu, C., and Li, Z. Task-customized self-supervised pre-training with scalable dynamic routing. In AAAI, 2022.
46.Luo, P., Wang, X., Shao, W., and Peng, Z. Towards understanding regularization in batch normalization. arXiv preprint arXiv:1809.00846, 2018.
47.Luong, T., Pham, H., and Manning, C. D. Effective approaches to attention-based neural machine translation. In Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing, pp. 1412–1421, 2015.
48.Lym, S., Choukse, E., Zangeneh, S., Wen, W., Sanghavi, S., and Erez, M. Prunetrain: fast neural network training by dynamic sparse model reconfiguration. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1–13, 2019.
49.Mei, S., Montanari, A., and Nguyen, P.-M. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33):E7665–E7671, 2018.
50.Peters, J., Buhlmann, P., and Meinshausen, N. Causal inference by using invariant prediction: identification and confidence intervals. Journal of the Royal Statistical Society. Series B (Statistical Methodology), pp. 947–1012, 2016.
51.Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019.
52.Rosenfeld, E., Ravikumar, P., and Risteski, A. The risks of invariant risk minimization. arXiv preprint arXiv:2010.05761, 2020.
53.Sagawa, S., Raghunathan, A., Koh, P. W., and Liang, P. An investigation of why overparameterization exacerbates spurious correlations. In International Conference on Machine Learning, pp. 8346–8356. PMLR, 2020.
54.Shah, H., Tamuly, K., Raghunathan, A., Jain, P., and Netrapalli, P. The pitfalls of simplicity bias in neural networks. arXiv preprint arXiv:2006.07710, 2020.
55.Shao, W., Meng, T., Li, J., Zhang, R., Li, Y., Wang, X., and Luo, P. Ssn: Learning sparse switchable normalization via sparsestmax. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2019.
56.Simonyan, K. and Zisserman, A. Very deep convolutional networks for large-scale image recognition. International Conference on Learning Representations, 2015.
57.Sun, Y., Wang, X., and Tang, X. Deep learning face representation from predicting 10,000 classes. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1891–1898, 2014.
58.Wang, H., Wu, Z., and He, J. Training fair deep neural networks by balancing influence. arXiv preprint arXiv:2201.05759, 2022.
59.Xie, C., Chen, F., Liu, Y., and Li, Z. Risk variance penalization: From distributional robustness to causality. arXiv e-prints, pp. arXiv–2006, 2020.
60.Xu, R., Cui, P., Kuang, K., Li, B., Zhou, L., Shen, Z., and Cui, W. Algorithmic decision making with conditional fairness. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 2125–2135, 2020.
61.Xu, R., Cui, P., Shen, Z., Zhang, X., and Zhang, T. Why stable learning works? a theory of covariate shift generalization. arXiv preprint arXiv:2111.02355, 2021.
62.Xu, R., Zhang, X., Cui, P., Li, B., Shen, Z., and Xu, J. Regulatory instruments for fair personalized pricing. In Proceedings of the ACM Web Conference 2022, pp. 4–15, 2022.
63.Xu, Y. and Jaakkola, T. Learning representations that support robust transfer of predictors. arXiv preprint arXiv:2110.09940, 2021.
64.Ye, N., Li, K., Bai, H., Yu, R., Hong, L., Zhou, F., Li, Z., and Zhu, J. Ood-bench: Quantifying and understanding two dimensions of out-of-distribution generalization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 7947–7958, 2022.
65.Yuan, G., Shen, L., and Zheng, W.-S. A block decomposition algorithm for sparse optimization. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 275–285, 2020a.
66.Yuan, X., Savarese, P. H. P., and Maire, M. Growing efficient deep networks by structured continuous sparsification. In International Conference on Learning Representations, 2020b.
67.Zhang, C., Bengio, S., Hardt, M., Recht, B., and Vinyals, O. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107–115, 2021a.
68.Zhang, D., Ahuja, K., Xu, Y., Wang, Y., and Courville, A. Can subnetwork structure be the key to out-of-distribution generalization? arXiv preprint arXiv:2106.02890, 2021b.
69.Zhang, G., Zhao, H., Yu, Y., and Poupart, P. Quantifying and improving transferability in domain generalization. arXiv preprint arXiv:2106.03632, 2021c.
70.Zhang, X., Xu, Z., Xu, R., Liu, J., Cui, P., Wan, W., Sun, C., and Li, C. Towards domain generalization in object detection. arXiv preprint arXiv:2203.14387, 2022a.
71.Zhang, X., Zhou, L., Xu, R., Cui, P., Shen, Z., and Liu, H. Nico++: Towards better benchmarking for domain generalization. arXiv preprint arXiv:2204.08040, 2022b.
72.Zhou, W., Zeng, Y., Diao, S., and Zhang, X. Vlue: A multi-task benchmark for evaluating vision-language models, 2022a. URL https://arxiv.org/abs/2205.15237.
73.Zhou, X., Zhang, W., Chen, Z., Diao, S., and Zhang, T. Efficient neural network training via forward and backward propagation sparsification. Advances in Neural Information Processing Systems, 2021a.
74.Zhou, X., Zhang, W., Xu, H., and Zhang, T. Effective sparsification of neural networks with global sparsity constraint. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 3599–3608, 2021b.
75.Zhou, X., Lin, Y., Pi, R., Zhang, W., Xu, R., Peng, C., and Zhang, T. Model agnostic sample reweighting for out-of-distribution learning. In International Conference on Machine Learning. PMLR, 2022b.
76.Zou, F., Shen, L., Jie, Z., Zhang, W., and Liu, W. A sufficient condition for convergences of adam and rmsprop. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11127–11135, 2019.

Citation

MLA

Zhou, X., et al. “Sparse Invariant Risk Minimization”. International Conference on Machine Learning, vol. 162, 2022, pp. 27222–44, https://proceedings.mlr.press/v162/zhou22e.html.

APA

Zhou, X., Lin, Y., Zhang, W., & Zhang, T. (2022). Sparse Invariant Risk Minimization. International Conference on Machine Learning, 162, 27222–27244. https://proceedings.mlr.press/v162/zhou22e.html

Chicago

Zhou, X., Y. Lin, W. Zhang, and T. Zhang. 2022. “Sparse Invariant Risk Minimization”. International Conference on Machine Learning 162: 27222–44. https://proceedings.mlr.press/v162/zhou22e.html.

Harvard

Zhou, X. et al. (2022) “Sparse Invariant Risk Minimization”, International Conference on Machine Learning. PMLR, pp. 27222–27244. Available at: https://proceedings.mlr.press/v162/zhou22e.html.

Vancouver

1. Zhou X, Lin Y, Zhang W, Zhang T (2022) Sparse Invariant Risk Minimization. In: International Conference on Machine Learning. PMLR, pp 27222–27244

BibTeX

@InProceedings{pmlr-v162-zhou22e,
  title = 	 {Sparse Invariant Risk Minimization},
  author =       {Zhou, Xiao and Lin, Yong and Zhang, Weizhong and Zhang, Tong},
  booktitle = 	 {Proceedings of the 39th International Conference on Machine Learning},
  pages = 	 {27222--27244},
  year = 	 {2022},
  editor = 	 {Chaudhuri, Kamalika and Jegelka, Stefanie and Song, Le and Szepesvari, Csaba and Niu, Gang and Sabato, Sivan},
  volume = 	 {162},
  series = 	 {Proceedings of Machine Learning Research},
  month = 	 {17--23 Jul},
  publisher =    {PMLR},
  pdf = 	 {https://proceedings.mlr.press/v162/zhou22e/zhou22e.pdf},
  url = 	 {https://proceedings.mlr.press/v162/zhou22e.html},
  abstract = 	 {Invariant Risk Minimization (IRM) is an emerging invariant feature extracting technique to help generalization with distributional shift. However, we find that there exists a basic and intractable contradiction between the model trainability and generalization ability in IRM. On one hand, recent studies on deep learning theory indicate the importance of large-sized or even overparameterized neural networks to make the model easy to train. On the other hand, unlike empirical risk minimization that can be benefited from overparameterization, our empirical and theoretical analyses show that the generalization ability of IRM is much easier to be demolished by overfitting caused by overparameterization. In this paper, we propose a simple yet effective paradigm named Sparse Invariant Risk Minimization (SparseIRM) to address this contradiction. Our key idea is to employ a global sparsity constraint as a defense to prevent spurious features from leaking in during the whole IRM process. Compared with sparisfy-after-training prototype by prior work which can discard invariant features, the global sparsity constraint limits the budget for feature selection and enforces SparseIRM to select the invariant features. We illustrate the benefit of SparseIRM through a theoretical analysis on a simple linear case. Empirically we demonstrate the power of SparseIRM through various datasets and models and surpass state-of-the-art methods with a gap up to 29%.}
}

Metadata:DOI registry

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: https://creativecommons.org/licenses/by/4.0/

Sparse Invariant Risk Minimization

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — SparseIRM imposes a global sparsity budget throughout IRM training

Knowl 2 — SparseIRM optimizes a stochastic mask with projected Gumbel-Softmax updates

Knowl 3 — Overparameterized linear IRM can prefer all features or only spurious and random features

Knowl 4 — A feature budget yields invariant-feature selection in the analyzed linear setting

Knowl 5 — Sparsifying during training avoids a failure mode of prune-after-training IRM

Knowl 6 — Benchmark experiments use spurious correlations that reverse at test time

Knowl 7 — SparseIRM retains test accuracy as MLP width increases

Knowl 8 — SparseIRM improves ResNet-18 accuracy on ColoredObject and CIFARMNIST

Knowl 9 — Representation probes show reduced encoding of spurious color

Knowl 10 — Weight-level sparsity is not yet an efficient training implementation

References

Citation

BibTeX

Access the Paper