Bayesian Invariant Risk Minimization

Generalization under distributional shift is an open challenge for machine learning. Invariant Risk Minimization (IRM) is a promising framework to tackle this issue by extracting invariant features. However, despite the potential and popularity of IRM, recent works have reported negative results of it on deep models. We argue that the failure can be primarily attributed to deep models' tendency to overfit the data. Specifically, our theoretical analysis shows that IRM degenerates to empirical risk minimization (ERM) when overfitting occurs. Our empirical evidence also provides supports: IRM methods that work well in typical settings significantly deteriorate even if we slightly enlarge the model size or lessen the training data. To alleviate this issue, we propose Bayesian Invariant Risk Minimization (BIRM) by introducing Bayesian inference into the IRM. The key motivation is to estimate the penalty of IRM based on the posterior distribution of classifiers (as opposed to a single classifier), which is much less prone to overfitting. Extensive experimental results on four datasets demonstrate that BIRM consistently outperforms the existing IRM baselines significantly.

Automated summary

Generalization under distributional shift is a challenge in machine learning. Invariant Risk Minimization (IRM) is a promising framework to address this issue, but recent studies have shown its poor performance on deep models, which can be attributed to overfitting. We propose Bayesian Invariant Risk Minimization (BIRM) to mitigate this problem by introducing Bayesian inference into IRM, and experimental results demonstrate its superiority over existing IRM methods.