Out-of-distribution Generalization for Total Variation based Invariant Risk Minimization

Invariant risk minimization is an important general machine learning framework that has recently been interpreted as a total variation model (IRM-TV). However, how to improve out-of-distribution (OOD) generalization in the IRM-TV setting remains unsolved. In this paper, we extend IRM-TV to a Lagrangian multiplier model named OOD-TV -IRM. We find that the autonomous TV penalty hyperpa-rameter is exactly the Lagrangian multiplier. Thus OOD-TV -IRM is essentially a primal-dual optimization model, where the primal optimization minimizes the entire invariant risk and the dual optimization strengthens the TV penalty. The objective is to reach a semi-Nash equilibrium where the balance between the training loss and OOD generalization is maintained. We also develop a convergent primal-dual algorithm that facilitates an adversarial learning scheme. Experimental results show that OOD-TV -IRM outperforms IRM-TV in most situations. Traditional risk minimization methods such as Empirical Risk Minimization (ERM) are widely used in machine learning. ERM generally assumes that both training and test data come from the same distribution. Based on this assumption, ERM learns model parameters by minimizing the average loss on the training data.