Out-of-distribution Generalization for Total Variation based Invariant Risk Minimization
Authors: Yuanchao Wang
Zhao-Rong Lai
Tianqi Zhong
Venue: ICLR 2025 Poster (2025)
Keywords: Out-of-distribution generalization
total variation
invariant risk minimization
primal-dual optimization
Primary area: other topics in machine learning (i.e., none of the above)
TL;DR: We propose an out-of-distribution generalization methodology for the total variation based invariant risk minimization.

Abstract: 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 hyperparameter 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.

Published as a conference paper at ICLR 2025
OUT-OF-DISTRIBUTION GENERALIZATION FOR TOTAL VARIATION BASED INVARIANT RISK MINIMIZATION
Yuanchao Wang Zhao-Rong Lai' Tianqi Zhong²
"Guangdong Key Laboratory of Data Security and Privacy Preserving,
College of Cyber Security, Jinan University
Prrat School of Engineering, Duke University
SSino-French Engineer School, Beihang University
yuanchac. wang@duke. edu, baizhr@jnu.co edu. cn, zhongtiang anqqi @buaa edu. cn
ABSTRACT
Invariant risk minimization 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, extend IRM-TV to Lagrangian multiplier model named OOD-TV-IRM. We find that the autonomous TV penalty hyperpa- rameter is exactly the Lagrangian multiplier. Thus OOD-TV-IRA 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 ob- jective is to reach 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 re- sults show that OOD-TV-IRM outperforms IRM-TV in most situations.
INTRODUCTION
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. However. the distributions between training and data or even within the training or test data are often different in practical situations, which affects the model's performance in new environments. Besides, ERM tends to exploit correlations in the training data, even if these correlations do not hold in different environments. This may cause the model to learn some spurious features that are irrelevant to the target, so as to perform poorly in new environments. Because ERM only focuses on the average performance on given data distribution, the model may overfit to specific training data and thus perform unstably when facing distribution changes or outliers. These problems lead to the poor generalization ability of ERM across different environments (Recht al., 2019; Arjovsky et al., 2019; Lin et al., 2022).
The key solving the above problems is distinguish between invariant features and spurious features that cause distribution shifts, so that a trained model can be generalized to an unseen do- main. This leads to the concept of out-of-distribution (OOD) generalization (Ben-Tal et al., 2009; Ye et al., 2022; Huang et al., 2023). One important methodology is the Invariant Risk Minimization (IRM, Arjovsky et al. 2019) criterion. It aims extract invariant features across different envi- ronments to improve the generalization and robustness of the model (Yang et al., 2023; Xin et al., 2023). Specifically, it introduces a gradient norm penalty that measures the optimality of the virtual classifier in each environment. Then it minimizes the risk in all potential environments to ensure that the model can still work effectively when facing distribution drift. summary, IRM copes with distribution changes by introducing environmental invariance constraints in order to make the trained model more robust in wider range of application scenarios (Ahuja et al., 2020). There are various variants and improvements of IRM, such as Heterogeneous Risk Minimization (HRM. Liu et al. 2021), Risk Extrapolation (REx, Krueger et al. 2021), SparselRM (Zhou et al., 2022), jointly
"Corresponding author.

---

Published as a conference paper at ICLR 2025
learning with auxiliary information (ZIN, Lin et al. 2022), and invariant feature learning through independent variables (TIVA, Tan et al. 2023a). Diversifying spurious features (Lin et al., 2024) can also be effective approach to improve IRM. These approaches not only provide new IRM set- tings or new application scenarios, but also improve the robustness and extensibility of IRM. They contribute the theoretical understanding of OOD generalization.
A recent work reveals that the mathematical essence of IRM is total variation (TV) model (Lai & Wang, 2024). TV measures the locally varying nature of function, which is widely applied to different areas of mathematics and engineering, such as signal processing and image restoration. Its core idea to eliminate noise by minimizing the TV of function while preserving sharp dis- continuities in the function (e.g., edges in an image (Dey et al., 2006)). The interpretation of IRM as a TV model not only provides a unified mathematical framework, but also reveals why the IRM approach can work effectively across different environments. The original IRM actually contains a TV term based on the lg norm (TV-(2). Compared with TV-22, TV-L, further has the coarea formula that provides geometric nature of preserving sharp discontinuities. This property is well-suited for the invariant feature extraction of IRM. Hence an IRM-TV-(2 framework proposed and it shows better performance than the IRM-TV-lz framework. When considering the learning risk part of the full-variance model, IRM-TV-é1 TV- performs better in OOD generalization. This finding helps to address distributional drift and improve model robustness in machine learning-
However, IRM-TV-[1 may not achieve complete OOD generalization, due to insufficient diversity in the training environments and the inflexible TV penalty. Lai & Wang (2024) investigaae some additional requirements for IRM-TV-(, to achieve OOD generalization. A key point is to let the penalty parameter vary according to the invariant feature extractor; however, the authors do not specify a tractable implementation. It leaves unsolved how to construct and optimize this framework for broader practical applications, resulting in an urgent need for effective solutions.
To fill this gap, we have found that the TV penalty hyperparameter can be used as a Lagrangian multiplier. We propose to extend the IRM-TV framework Lagrangian multiplier framework, named OOD-TV-IRM. is essentially a primal-dual optimization framework. The primal optimiza- tion reduces the entire invariant risk, in order to learn invariant features, while the dual optimization adversarial interference with spurious features.
strengthens the TV penalty, in order to provide The objective is to reach semi-Nash equilibrium where the balance between the training loss and OOD generalization is maintained. We also develop convergent primal-dual algorithm that fa- cilitates an adversarial learning scheme. This work not only extends the theoretical foundation of previous work, but also provides concrete implementation scheme to improve OOI) generalization of machine learning methods across unknown environments.
PRELIMINARIES AND RELATED WORKS
2.1 INVARIANT RISK MINIMIZATION
In machine learning tasks, often work with data set consisting of multiple samples, where each sample includes input and output variables. These samples are typically drawn from various environments, but in many scenarios, the specific environment labels are not explicitly available. As a result, the challenge is to train a model that can generalize well across these different environments, even when the environment information is unknown.
IRM (Arjovsky et al., 2019) addresses this challenge by structuring the model into two components: feature extractor and classifier. The feature extractor is responsible for identifying invariant fea- tures from the input data, while the classifier makes predictions based on these features. IRM aims minimize the invariant risk across different training environments, in order find a representation that performs consistently well across all environments.
For given training data set of n samples D == {(x,, w) E X'X > the empirical risk in a given environment is computed using a loss function L, which measures the discrepancy between the predicted value and the true output. This risk is expressed as:
1 Σ L(w ((xx), Vi. R(w o D, = n
(1)
i=
2

---

Published as a conference paper at ICLR 2025
which represents the mean loss samples in environment IRM simultaneously learns an invariant feature extractor & and a classifier w across multiple training environments, so that the risk minimized in all environments:
min
Σ CE
R(w o 2, t. argmin R(w' o e),,, Ve E Eir,
where the classifier w should be the optimal solution for minimizing the risk in each environment e. However, it is a challenging bilevel optimization problem. Hence Arjovsky et al. (2019) further propose a practical variant of IRM as follows:
min Σ (R(1o ( e CE
In this version, the classifier w set to a scalar value of and a gradient norm penalty term introduced to convert the objective into single-level optimization problem.
ZIN (Lin et al., 2022) further extends the IRM framework by introducing auxiliary information E 2 to simultaneously learn environment partitioning and invariant feature representation. It assumes that the space of training environments & the convex hull E linearly independent fundamental environments. By learning mapping p Z Ag that provides weights for different environments, ZIN optimizes the following objective:
min 02, +
R(w & E (()) 1(E) + À Amax Ve 2113
(4)
By this means, ZIN can effectively learn invariant features even without explicit environment parti- tion information.
2.2 TOTAL VARIATION
TV a mathematical operator widely recognized for its ability to measure the global variability of function. It has been extensively applied in various fields such as optimal control, data transmission, and signal processing. The central concept behind TV is to quantify the overall amount of change in function across its domain. It is able preserve sharp discontinuities (such edges in images) while removing noise, making it a useful tool in image restoration and signal denoising. This prop- erty is further elucidated through the coarea formula (Chen al., 2006). For function f (?)). where Ω is an open subset of Rd
/ V = / n
ds dr,
(5)
(x)
where fr)) represents the level set of the value y. This formulation shows that TV integrates over all the contours of the function, reinforcing its capability in capturing piecewise-constant fea- tures. Mumford & Shah (1985) and Rudin et (1992) propose the following TV-&2 and TV-6, models, respectively.
(/IIP+AJUU-T dy
+
inf SELP(D)
inf SELA
(6)
where f (9)) is the ground-truth signal and A is a hyperparameter that affects accuracy. These models aim to preserve target features in the approximation f and leave noise in the residual (f- In general, TV-,, can better preserve useful features and sharp discontinuities in the approximation f.
2.3 INVARIANT RISK MINIMIZATION BASED ON TOTAL VARIATION
can be seen that the penalties of (3) and (4) are similar to a TV- term, which has been verified by Lai & Wang (2024) under some mild conditions, such measurability of the involved functions, non-correlation between the feature extractor and the environment variable, representability of w by etc. In this sense, (3) and (4) are actually IRM-TV-&2 and Minimax-TV-A models. Considering
3

---

Published as a conference paper at ICLR 2025
the superiority of TV-2 over TV- in sharp feature preservation, Lai & Wang (2024) further propose the following IRM-TV-I1 (7) and Minimax-TV-4, (8) models:
min
{Eu [R o D) + X(E I V R(w o })
(7)
min E Fr [R( 6)] + ( R(w +))+)
(8)
where Ewep denotes the mathematical expectation with respect to (w.r..)) whose measure is induced by p. and ( denotes the uniform probability measure for the environment e. The maxi- mization w.r.. p identifies the worst-case inferred enviromment causing the greatest variation in the risk function. Hence, Minimax-TWe TV-4 can be seen as the TV-& version of ZIN. The corresponding coarea formula for IRM-TV-L1 Minimax-TV-ei is
Ja V w R(w T) du
f.
ds dy,
((w():((wo+)))))
where (w Q: R(w o b) >} denotes the level set, and 8 is the Hausdorff measure in the corre- sponding dimensions. This formula ensures that IRM-TV-E) Minimax-TV-e, preserves essential structural features while maintaining robustness across different environments. In the rest of this paper, we abbreviate IRM-TV-& as IRM-TV if not specified otherwise.
METHODOLOGY
3.1
AUTONOMOUS TOTAL VARIATION PENALTY AND LAGRANGIAN MULTIPLIER
OOD generalization represents the generalization ability of trained model to an unseen domain. From the perspective IRM-TV, it can be represented by (Ben-Tal et al., 2009; Arjovsky et al., 2019; Lin et al., 2022; Lai & Wang, 2024)
min max li(w Φ).
(10)
In general, IRM-TV cannot achieve OOD generalization without additional conditions, due to insuf- ficient diversity in the training environments and the inflexible TV penalty. However, Lai & Wang (2024) indicate that the penalty hyperparameter Xạ of IRM-TV should be variable according to the invariant feature extractor d:
min {E [R(w Φ) + Ao (Ewl V.. R
(11)
min
{ --
E
[R(w 4)] + X max(Eue-pll +)1)
(12)
The main reason that a variable Xo can fill the gap between the objective value of (11) or (12) and the objective value of (10). Lai & Wang (2024) further indicate that this variable Xg exists in general situations. However, they do not specifyy any tractable and concrete implementation of Ag.
To fill this gap, we find that Ay actually serves Lagrangian multiplier in (11) and (12). Hence we use it as an autonomous parameter, which is also taken into the training scheme. Specifically, the parameters for the invariant feature extractor, also denoted by Φ are directly taken as inputs for X. Next, À can be parameterized by another set of parameters V. By this means, can be represented function of both V and Φb λ(V, D). Then À not only depends t. but also adjusts its strength through V. Now it can be deployed in a primal-dual optimization model, as illustrated in the next section.
3.2 PRIMAL-DUAL OPTIMIZATION AND SEMI-NASH EQUILIBRIUM
Replacing X by A(V, +) in (11) and (12), we obtain the corresponding Lagrangian functions:
( ④) E,, RR w 4) + A( 4) (V +)]))
(13)
h(, V, ④) E 1 [R( e T) + X(V, + (F ( R(w +)]
(14)
Applying the OOD generalization criterion (10) in the literature of IRM, we adopt a minimax scheme to optimize t and V/
4

---

Published as a conference paper at ICLR 2025
Outer Minimization (Primal): the primal variable + is trained to minimize the entire invariant risk, in order to capture invariant features across different environments.
Inner Maximization (Dual): the dual variable vi (or (p. V is trained to maximize the au- tonomous TV penalty ( ②) (E,.I/V R(wo))) (or A(V, V +) p [V R(wo@))))
which confronts the most adverse scenarios with spurious features.
They lead to the following proposed OOD-TV-IRM (15) and OOD-TV-Minimax (16) models:
min max SVV V, +) := min
Bu
[R( }) + max [A 4 (E | e R(w [
(15)
W
min max hpp V, +b == min E 4m
R[ D] max n.
[A(V, () R(w .)
A.
(16)
They are essentially primal-dual optimization models. To further understand their functions, we in- troduce the concept of semi-Nash equilibrium in the context of this paper, without loss of generality.
Definition 1. A semi-Nash equilibrium (V'*, Φ*) of Lagrangian function g( V 2) satisfies the following conditions:
1. g (V"", x0)) g (V, Φ*) for any Ti in the parameter space.
2. g M Φ6) g( L mas ( (Ø*,, Φ*) & g maz ()) +) for any / in the parameter space, where V man ( ( E arg max [A(V, ④b) (E.r [I R(w o [)]] In other words, V Wan (+)) is a solution to the dual optimization with fixed.
A semi-Nash equilibrium (p", Mi" Φ²) of a Lagrangian function h(p v b) follows the same defi- nition with the unified dual variable (p, V).
Remark: in the context of game theory, OOD-TV-IRM (15) (or OOD-TV-Minimax Eq. 16) can be seen two-player zero-sum game, where VVV +) and -9( -V(V, +) are the gains for players v and 4. respectively. Item 1 in Definition 1 indicates that 4 is the best strategy for dual optimization given 4 (which independent of Vf). On the other hand, Item indicates that 6 the conditional best strategy for primal optimization only when Vi mez ( is optimized over each क. In this sense, (V"", (b)) is only a semi-Nash equilíbrium instead of a full Nash equilibrium. It accords with our task where & is the feature extractor and primary variable that should confront the most adverse scenar- ios, while V is the TV penalty parameter and secondary variable that provides OOD generalization depending on Φ.
Theorem 2. Assume that R(w o 4) is continuous w.r.t. (w, 1 and differentiable w.r.t w, A(V, 2) is continuous w.r.t. (V, 2). and the feasible set of parameters (p, V 2) is bounded and closed. Then any solution OOD-TV-IRM (15) or OOD-TV-Minimx (16) semi-Nash equilibrium.
Proof. The proof is given in Appendix A.1.
Most of the widely-used loss functions and deep neural networks satisfy the conditions of Theorem 2, such as the squared loss, cross-entropy loss, multilayer perceptron, convolutional neural networks, etc. In practical computation, a bounded and closed feasible parameter set is also satisfied due to the maximum float point value. Hence Theorem is applicable to general situations. It indicates that OOD-TV-IRM and OOD-TV-Mininax balance between the training loss and OOD generalization at semi-Nash equilibrium.
3.3 PRIMAL-DUAL ALGORITHM AND ADVERSARIAL LEARNING
We develop a primal-dual algorithm to solve OOD-TV-IRM (15) or OOD-TV-Minimax (16). To do this, we further assume that R(w o 4) and X(V, 4) are differentiable w.r.t. their corresponding arguments w, & and V. As indicated in (Lai & Wang, 2024), V R., &) in (15) or (16) is non- differentiable w.r.t. Φ. Hence we adopt the subgradient descent method update these parameters, as illustrated in Appendix A.2. Then the primal and dual updates for (15) or (16) are:
((++1) (A) mi (M) ( ((()),
(17)
(((k-1) = (k) + 112
(k) V Vog( (V(k) (((+1));
(0(++1) ((+)_ (1()) )ph() ((()) so()))
(o(tt)), v((+1) = (,(t), ()())) + 1122
(1)
V (Cnne)((p()) (T(A)) (((+1)),
(18)
5

---

Published as a conference paper at ICLR 2025
(R) > 0 are the learning rates for the primal and dual updates, respectively. The com- where (8) _122
putation of the above (sub)gradients is illustrated in Appendix A.2. In general machine learning scenarios, the smoothness and convexity of the Lagrangian functions g and h are usually unknown, thus 7/1 and 1/2 should be set according to different specific tasks. In addition, we provide a convergent scheme as follows.
Theorem 3. The same assumptions in Theorem 2 are used. By setting
( x(k))= + 0; ( x(k)) O. is Vvg(T((M) @l(+t))))+))))))))))))))))) + 0;
(k)
kP (00000 (1
71
0,
KP Vagl [ gllttaaee
(k)
1/2
0,
if Vvg(vv( 1((+1)) =
(19)
ke 11 (1l Pllcmie
is Ooh(pth) w(kk x(k)) + 0:
(k)
or
Th
0,
is aph(plh) wilk) Φ(())) = 0.
if ((o.))((ll)) w((), (((+1))) 0;
(k)
KP V u hi 0ib,, w
1/2
(20)
0,
if (o, wilk), (((+1)))+ 0.
with an arbitrary p. > 1, the primal-duc dual algorithm (17) or (18) achieves convergent sequences {(D()) (1())) or ((plk), T((), ((()))= respectively. Moreover, its computational complexiry
O(((p (1(pp-10 1)c) to achieve convergence tolerance of O.
Proof. The proof is given in Appendix A.2.
In fact, (17) (18) facilitates adversarial learning to dynamically adjust model parameters (vp +) or (p, Mr 2) to balance training loss and OOD generalization. This approach extends the theoretical work of IRM-TV by providing practical, dynamic implementation of X Φ) that adapts to the most adverse environments.
models as well the primal- al- gorithm can also be applied to the TV-E2 versions, which are omitted here since they are simpler differentiable cases.
Remark: the above
and
3.4 IMPLEMENTATION WITH NEURAL NETWORKS
First, we note that all the theoretical results in this section hold with any functions satisfying the corresponding conditions (e.g., continuity, differentiability). not limited neural networks. To pro- vide specific instance, we implement A( ②) with neural networks that take ! and Φ as network trainable parameters and network inputs, respectively. Besides A(V, 4). & and p can also be devel- oped by neural networks. Specifically, $ X + H. is the feature extractor that maps sample from the sample space to the feature space H. It can be implemented by different types of neural networks depending the practical situation, as shown in Table 1. The data sets correspond to different synthetic or real-world tasks, which will be illustrated Section For example, & can be multilayer perceptron in Simulation, larly, p Z - A is the environment inference operator that maps the auxiliary variables the probability space A F-ddimensional real space (i.e., E environments). Moreover, the architectures of p are developed by multilayer perceptrons, which compatible with the corresponding 4 shown in Table 1.
convolutional neural network in Landcover. Simi-
Z O(E) and ((E) 1}. where ER denotes the
{v E E
Without ambiguity, we can also use + to denote the parameters of the feature extractor. Then A b) can be implemented as a neural network that takes & input and v/ as network param- eters. We provide tractable architectures of À in Table 1, which should also be compatible with the corresponding 4. With the primal-dusl dual algorithm Section 3.3, & and v can be learned in an adversarial way to improve OOD generalization. which puts Theorem 2 into practice.
4 EXPERIMENTAL RESULTS
We conduct experiments on 7 simulation and worldwordd data sets to evaluate the effectiveness of the proposed OOD-TV-IRM and OOD-TV-Minimax in different OOD generalization tasks. These

---

Published as a conference paper at ICLR 2025
Table 1: Network architectures of the invariant feature extractor 4 the environment inference oper- ator and the TV penalty strength À w.r.t. their corresponding inputs.
DATA SET
SIMULATION
1S, 1)
LINEARC 16]-RELUUD -ONNNELLLLLLAAAAAAAAA LINEER 16 1) SIGMOIDU
L1NEAR(512, -LIN-RLL 32,
LINEAR(7, 10) SIGMOIDO)
CONV 32, 5) -0--RLLL---COOC
SEELUOO-AXPOOLBBRR 2tw.../.....-CCiiooooo. C)-YCoVVD 61,
LINEAR(2, 16)-RRL.U0-LNNAARI S0-TTAAM
-RELUO -MAXPOOLIDC 2) CCNVVOD(64,,
-RELU)) AVVELo--VVE
ADULT
LINEAR(59, >RELU) >LINEAR(16, I)
LINEARI6, 16) :RELED .INEAR :SOFTMAX))
HOUSE PRICE
LINEAR(15,33) -RELU() LIN1AR(32,1)
LINEAR(1,64) RELUO LINEAR(64 SOPIMANÓ
CONV2D(3,32) 3,3333 MAXPno12n(2.2)
COLORED MNIST
-CoNN22x332,64.3 MxxPo0L2D(2.2)
LINEAR(3.16) RELUO (INEAR(16.1) SOFTMAX)
CoNv20(64.128.3 MAXPO0L22(222)
3 LINEAR(128.10)
RESNET-180
LINEAR(3.16) RELUO LINEAR(16. SOFTMAX(
DATA SET
+ À
SIMULATION
1) RELUU LINEAR(1, 1) SOFTPLUS()
CELERA
ANEEAR8825,, R - LINEAR( 16, 1) SOFTPLUSO)
LANDCOVER
LINEAR 32) RELU() - LINEAR(32, 1) - SOFTPLUS()
ADULT
L977AR(977, -1 16)-RELUU)) - LINEAR(16, 1) SOFTPLUS()
HOUSE PRICE
LINEAR(545, 32) - RELUO + LINEAR(32, 16) - SOFTPLUS() - 1I6EAR(16,1)
COLORED MNIST
LINEAR(242122.32) -> RELU() LINEAR(32. - SOFTPLUSO)
NICO
LINEAR(I 180113,32) RREUU0 LINEAR(32,1) SOFTPLUSO)
tasks largely follow the settings in the literature of IRM (Lin al., 2022; Lai & Wang, 2024). We also include the original IRM (Arjovsky et al., 2019). ZIN (Lin et al., 2022), IRM-TV-(1, Minimax- TV-6, (Lai & Wang. 2024), EIIL (Creager et al., 2021). LIF (Nam et al., 2020), and TIVA (Tan et al., 2023b) frameworks in comparisons to verify the effectiveness of the proposed primal-dual optimization and adversarial learning scheme. Note that IRM and ZIN are equivalent to IRM-TV- l2 and Minimax-TVaa
respectively. Each experiment is repeated 10 times to record the mean and standard deviation (STD) of the results for each compared method. The code is available at https: //gi thub. com/ a /
4.1 SIMULATION DATA
The simulation data consists of temporal heterogeneity observations with distributional shift w.r.t time, which is used in (Lin et al., 2022; Tan et al., 2023a) to evaluate OOD generalization. Details of generating this data set are provided in Appendix A.3. Among all the parameter settings, the most challenging one is (p, P .P.) (0.999, 0.9, 0.8), thus we use this setting in the evaluation.
Table (left) shows the mean and worst accuracies of different methods over the four test environ- ments simulation data. The proposed OOD-TV-based methods outperform their counterparts in all the mean and worst accuracy cases. Moreover, TV-IRM-- achieves the highest accuracies among all the compared methods in all the cases, while OOD- MOD-ii-MMnimax-fy achieves the highest accuracies among all the compared methods without environment partition in all the cases. Hence the OOD-TV-based methods achieve the best performance regardless of whether the environment partition is available or not.
Table 2: Accuracies of different methods on simulation data (left) and CelebA (right).
Mean STD
STD
+
EBIL
0.8161 = 0.0361
0.7159
0.0540
EIIL
0.7323
0.0281
0.72 + 0.0291
0.6965
0.60192
0.5983
TIVA
0.8274 = 0.0154
0.8139 = 0.0179
TIVA
0.7693 = 0.0199
0.7253
0.0209
ZIN
0.7338
0.06673
0.7155
0.0544
OOD Minimax,
0.8916 + 0.0202
0.8858
0.0223
OOD OOD-TV Minimaxx--
0.7608 0.0187
0.6895
0.0243
Minimax-T TV-7,
0.7636 = 0.0742
0.6031
0.0666
Minimax-TV-yy
0.6376 0.0190
0.5341
0.000
DVD-TV-Minimaa
0.8379 + 0.0418
0.7805
0.0612
DOD-TV-Minll
Mi-nima---n
0.7044 -: 0.0361
0.6180
0.0580
IRM
0,9670 1: 0.01993
(L3553
0.0235
IRM
0.7631 0.0281
0.7410
0.025
DOD-TV-IRM-6,
0.8904 + 0150
0.8831
:: 0.0195
DOD-TV-IRM-G
0.7786 = 0.0157
0.7549 = 0.0227
TV-2,
0,8713 1 0.12001
0.8618
0.0200
IV.
0.7151 1 0.0200
0.7461
0.017
-VD--Y-IRM-6
0.8944 1 0.0150
0.8877 1 0.0168
OOD-TV-IRM-i,
0.7945 = 0.0162
0.7560 0.0170
7

---

Published as a conference paper at ICLR 2025
4.2 CELEBA
The CelebA data set (Liu et al., 2015) contains face images of celebrities. The task is to identify smiling faces, which are deliberately correlated with the gender variable. The 312-dimensional deep features of the face images are extracted using pre-trained ResNet18 88 model (He et al., 2016), while the invariant features are learned using subsequent multilayer perceptrons (MIP). EIII,,, LfF, TIVA, ZIN,
Minimax-TV-é, and OOD-TV-Minima--4; take seven additional de- scriptive variables for environment inference, including Young, Blond Hair, Eyeglasses, High Cheek- bones, Big Nose, Bags Under Eyes, and Chubby. As for IRM, OOD-TV-IRM-C2, IRM-TV-(, and OOD-TV-IRM-4.. they use the gender variable as the environment indicator.
Table 2 (right) shows the accuracies of different methods on CelebA. The proposed OOD- DOD-sVd methods outperform their counterparts in all the 8 cases with both mean and worst accuracies. More- over, OOD-TV-IRM--, achieves the highest accuracies among all the competitors. Hence OOD-TV- IRM improves OOD generalization compared with IRM-TV.
4.3 LANDCOVER
The Landcover data set consists of time series data and the corresponding land cover types derived from satellite images (Gislason et al., 2006; Russwurm et al., 2020; Xie et al., 2021). The input data has dimensions of 46 x 8 and is used to identify one of six land cover types. The invariant feature extractor & is implemented as ID-CNN to process the time series input, following the approaches of Xie et al. (2021) and Lin et al. (2022). In scenarios where ground-truth environment partitions are unavailable, latitude and longitude are used as auxiliary information for environment inference. non-African data, and then tested on both non-African (from regions
All methods are trained not overlapping with the training data) and African regions. This is a complex and challenging experiment among the 7 experiments of this paper.
Table (left) shows the accuracies of different methods on Landcover. The proposed OOD-TV- based methods outperform their counterparts in all the 8 cases with both mean and worst accuracies. Hence the OOD-TV-based methods show competitive performance in this challenging task.
Table 3: Accuracies of different methods on Landcover (left) and adult income prediction (right).
Method
Mean STD
Worst STD
Method
Mean STD
Worst + STD
EML
0.6603 = 0.0112
0.6614 = 0.0156
EII
0.7734 = 0:0120
0.7396 0.0132
LIF
0.5087 0.0567
0.5017
0.0580
LIF
0.7785 0.0080
0.7545 0.0084
TIVA
0.5388 1 0.0154
0.3351
0.0175
TIVA
0.8226 1: 0.0071
0.8085
0.0088
7IN
0.6511 + 0.0254
0.6468
0.0225
ZIN
0.8201 = 0.0024
0.8758
0.0001
MOD-TV-Minimax-fy
0.6530 + 0.0178
0.6490
+ 0.0124
DOD-IV-Mininax-ff
0.8345 = 0.0055
8105 = 0.0063
Minimax-VV. TV-6,
0.6573 = 0.12119
0.6551
0.0176
Minimax-TV-
0.8101 = 0.0082
0,7958
0.0103
MOD-TV-Minimaree
0.6604 + 0.051
0.6566
+ 0.0150
0.8330 = 0.0058
8093 0.0075
TRM
0.6472 + 0.01
0.6422 0.0154
TRM
0.8257 0.0068
0.8052
+ 0.0112
OOD IVM-
0.6705 0.0132
0.6672
0.0174
OOD TV-RM--,
0.8369 0,0109
8103
0.0091
IRM-TV-A
0.6422 - 0.0196
0.6300 - 0.0208
TRM-TV-
0.8208 + 0.0074
0.8064
0.0073
OOD-TV-IRM-G
0.6672
0.6636 =: 0.0197
OOD-TV-IRM-G
0.8435 0.0089
0.8197 0.0091
4.4 ADULT INCOME PREDICTION
This task uses the Adult data set' to predict whether an individual's 's income exceeds $50K per year based on census data. The data set split into four subgroups representing different environments based on race E [Black, NlncBlack] and sex E [Male. Female). Two-thirds of the data from the Black Male and Non-Black Female subgroups are randomly selected for training, and the compared methods are verified across all four subgroups using the remaining data. Six integer variables Capital-Gain, Capital-Loss, and Hours- Week - are fed
Age, FNLWGT, Education- into EIIL, LIF, TIVA. ZIN, OOD-TV-Mimima-- Minimax-TV-L, and OOD-TV-Minimx--4/ for environment inferennee Ground- truth environment indicators are provided for IRM, OOD-TV-IRM- fz. IRM-TV- and
Categorical variables, excluding race and sex, are encoded using one-hot encoding, followed by principal component analysis (PCA), retaining over 99% of the cumulative explained variance. The transformed features are then combined with the six integer
8

---

Published as a conference paper at ICLR 2025
variables, resulting in 19-dimensional representations, which normalized to have and unit variance for invariant feature learning.
Table 3 (right) shows the accuracies of different methods this income prediction task. The pro- posed DOD-TV-based methods outperform their counterparts in all the cases with both mean and worst accuracies. Hence they are more effective and robust than the original IRM-TV methods in this prediction task regardless of whether the environment partition is available or not.
4.5 HOUSE PRICE PREDICTION
The above four experiments are all classification tasks. We also perform regression task with the House Prices data set?. This task uses 15 variables, such as the number of bathrooms, locations, and other similar features, to predict the house price. The training and test sets consist of samples with built years in periods [1900, 1950] and (1950, 2000), respectively. The house prices within the same built year are normalized. The built year variable is fed into EIIL, TIVA, ZIN, OOD-TV-Miimax- l2. Minimax-TV--e and TED-TV-Minimax-4 for environment inference. LfF is not suitable for this regression task due to the unstable training process and poor performance. The training set partitioned into 5 subsets with 10-year range in each subset. Samples in the same subset can be seen as having the same environment.
Table 4 (left) provides the mean squared errors (MSE) of different methods in this regression task. The OOD-TV-based methods achieves lower MSEs than their counterparts in the mean result of all the environments and the worst environment. Hence they are more effective and robust than the original IRM-TV methods the regression task.
Table Mean squared errors of different methods on house price prediction (left) and accuracies of different methods Colored MNIST (right).
Method
Mean
STD
Worst STD
BIIL
0.3644
0.0512
0.4953 1 0.0587
TIVA
0.3869
0.0413
0.5519 0.0508
Method
Mean STD
ZIN
0.3203
0.0378
0.4465
0.0552
EIL
0.8236 = 0.0265
DOD-TV-Minimax-f;
0.2743
0438
0.3983
050666
LIF
0,8964 1 0.0250
Minmax- TV-1,
0,33339
0.0308
0.4741 + 0.0320
ZIN
0.9302 + 0.0133
NOD-TV-Minina--
0.2621 = 0.0510
0.3706 = 0. 0.0549
-in-in--Mn
0.9671 = 0.0060
IRM
0.4589 = 0.0762
0.6378 = 0.0873
Minimax-
0.9322 = 0.0362
OVD-TV-IM--2
0.3549 + 0.0519
0.4772 + 0.0591
OVD-TV-Mininx--4
0.9514 + 0.0242
IRM-TV-4
0.4009 + 0.0355
0.5642 + 0.0615
OOD-TV-IRM-I
0.3383 = 0.0303
0.4763 # 0.0428
4.6 COLORED MNIST
We use the Colored MNIST data set evaluate the performance of our approach in multi-group classification of more general scenario. It consists of hand-written digits of 0 to 9 with different background colors of red, green, and blue. To make it more challenging, remove the background color information of samples, so that only the mean values of the red, green, and blue channels of sample can be used three auxiliary variables for environment inference. EIIL, LiF, OOD-TV- Minimax, and Minimax-TV are suitable for this setting, while TIVA is unavailable due to unknown auxiliary feature extractions. Their classification accuracies are shown Table (right). The two OOD-TV-Minima methods achieve the highest accuracies among all the compared methods in this multi-group classification task. Thus OOD-TV-Minimax is effective in improving OOD generaliza- tion for such complex scenario.
4.7 NICO
The NICO data set (He et al., 2021) a widely-used benchmark in Non-Independen and Identically Distributed (Non-L.L.D) image classification with contexts. This data set presents challenges due to both correlation shift and diversity shift. There are two superclasses in NICO: Vehicle and Animal,

---

Published as a conference paper at ICLR 2025
which consist of 9 classes and 10 classes, respectively. Each class has several contexts, while dif- ferent classes may have different context sets. We randomly split each context into 80% training samples and 20% test samples. Then we remove the context information of samples and unite all the contexts to form a whole class. Therefore, only the mean values of the red, green, and blue channels of a sample can be used as three auxiliary variables for environment inference. EIll.,, OOD-TV- Minimax, and Minimax-Ti can be used in this setting, while LIF has unstable training process and poor performance, and TIVA is unavailable due unknown auxiliary feature extractions. As shown in Table 5, the two POD-TV-Minima methods outperform other competitors in terms of classifi- cation accuracy. Thus DOD-TV-Minimal is effective improving OOD generalization with both correlation shift and diversity shift.
Table 5: Accuracies of different methods on NICO.
Method
Vehicle (Mean + STD)
Animal (Mean + STD)
BIIL
0.6902 + 0.0357
0.8664 0.0110
ZIN
0.7176 + 0.0374
0.8541 = 0.0250
DOD-TV-Minimax-12
0.7872 + 0.0275
0.8826 0.0249
Minimax-TV-L.
0.7104 + 0.0236
0.8289 = 0.0311
BOD-TV-Minimay.
0.7289 + 0.0206
0.9112 + 0.0154
4.8 ADVERSARIAL LEARNING PROCESS
To visualize the adversarial learning process, we plot the objective values and parameter changes for OOD-TV-IRM and DOD-TV-Minimax on simulation data in Figure AI. Specifically, the overall objectives are the g( VV, 4) and h(p. V, Φ) defined in (13) and (14), respectively. The penalty terms correspond to the second terms (13) and (14), respectively. The parameter changes correspond to [[+((++] (1())12 and ( (1()) respectively. We follow the annealing strategy of (Lin et al., 2022) in the early epochs, thus the adversarial learning starts from the 2001st epoch. The Adam scheme (Kingma Ba, 2015) is adopted as the optimizer. Figure A1 indicates that the adversarial learning process becomes stable with 400 epochs. Both of the overall objective and the penalty term converge to form a stable gap, which corresponds to the fidelity term be optimized.
Besides, both (k+1) (<1) and (k+1) ()()) I|a converge to small values, which indicates that the parameters (k)) and ()(1) are rarely updated.
We also demonstrate the adversarial learning process for OOD-TV-Minimax (since OOD-TV-IRM not applicable) on Colored MNIST in Figure A2. The Adam optimizer is used for the primal parameter 4. while the dual update (18, 20) is used for & to ensure convergence. Results show that the adversarial learning process effectively converges after 600 epochs for OOD-TV-Mimmmx- l and after 300 epochs for OOD-TV-Minimax-f, All these results indicate that the adversarial learning process is feasible.
CONCLUSION AND DISCUSSION
We extend the invariant risk minimization based on total variation model (IRM-TV) Lagrangian multiplier model OOD-TV-IRM, in order to improve out-of-distribution (OOD) generalization. OOD-TV-IRM essentially primal-dual optimization model. The primal optimization minimizes the entire invariant risk to extract invariant features, while the dual optimization strengthens the au- tonomous TV penalty to provide an adversarial interference. The objective of OOD-TV-IRM is to find semi-Nash equilibrium that balances the training loss and OOD generalization. We further develop a convergent primal-dual algorithm to facilitate adversarial learning for invariant features and adverse environments. Experimental results show that the proposed OOD-TV-IRM framework improves effectiveness and robustness in most experimental tasks, achieving competitive mean ac- curacies, worst accuracies, and mean squared errors across different test environments.
One potential future direction is to maximize the benefits of OOD-TV-IRM by improving the di- versity and representation of training environments. Another potential approach is to improve the effectiveness of the primal-duse optimization scheme. A third potential approach is to develop new penalties with different variations beyond the widely-used TV penalty.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work is supported in part by the National Natural Science Foundation of China under grant 62176103, and in part by the Science and Technology Planning Project Guangzhou under grant 2024A0419896.
REFERENCES
Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. Invariant risk min- imization games. In Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 145-155.PMLR, 13-18 Jul 2020.
Martin Arjovsky, Léon Bottou, Ishaan Guirajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint wwXXiv:907.028997 2019.
Aharon Ben-Tal, Laurent El Ghaoui, and Arkadi Nemirovski. Robust optimization, volume 28.
Princeton university press, 2009.
Terrence Chen, Wotao Yin, Xiang Zhou, and Dorin Comaniciu. Total variation models for variable lighting face recognition. JEEE transactions on pattern analysis and machine intelligence, 28: 1519-1524, 2006.
Elliot Creager, Joern-Henrik Jacobsen, and Richard Zemel. Environment inference for invariant learning. Marina Meila and Tong Zhang (eds.). Proceedings of the 38th International Confer- Machine Learning, volume 139, pp. 2189-2280000 2200, 18 24 Jul 2021.
Nicolas Dey, Laure Blanc-Feraud, Christophe Zimmer, Pascal Roux, Zvi Kam, Jean-Christophe Olivo-Marin, and Josiane Zerubia. Richardson-lucy algorithm with total variation regularization for 3d confocal microscope deconvolution. Microscopy research and technique, 69(4):260-266, 2006.
Pall Oskar Gislason, Jon Atli Benediktsson, and Johannes R. Sveinsson. Random forests for land cover classification. Pattern Recognition Letters, 27(4):294-300, 2006.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEFE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016.
Yue Hc, Zheyan Shen, and Peng Cui. Towards non-iid image classification: A dataset and baselines. Pattern Recognition, 110: 107383, 2021.
Zhuo Huang, Xiaobo Xia, Li Shen, Bo Han, Mingming Gong, Chen Gong, and Tongliang Liu. Harnessing distr--iistiitutionn examples via augmenting content and style. In The Eleventh Inter- national Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023, 2023.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.
David Krueger, Ethan Caballero, Joern-Henri Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. diitti--iittibuttion generalization via risk extrapo- lation (rex). In International conference machine learning, pp. 5815-5826. PMLR, 2021.
Zhao-Rong Lai and Weiwen Wang. Invariant risk minimization a total variation model. Pro- ceedings of the 41st International Conference on Machine Learning, volume 235, pp. 25913- 25935, PMLR, 21-27 Jul 2024.
Yong Lin, Shengyu Zhu, Lu Tan, and Peng Cui. Zin: When and how leam invariance without environment partition? Advances in Neural Information Processing Systems, 35:24529-22 224522, 2022.
11

---

Published as a conference paper at ICLR 2025
Yong Lin, Lu Tan, Yifan Hao, Honam Wong, Hanze Dong, Weizhong Zhang, Yujiu Yang, and Tong Zhang. Spurious feature diversification improves dustr-udist generalization. In The Twelfth International Conference Learning Representations, 2024.
Jiashuo Liu, Zheyuan Hu, Peng Cui, Bo Li, and Zheyan Shen. Heterogeneous risk minimization. In Proceedings of the 38th International Conference on Machine Learning, volume 139, pp. 6804- 6814, 18-24 24 Jul 2021.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), December 2015.
David Mumford and Jayant Shah. Boundary detection by minimizing functionals. In IEEE Confer- ence on computer vision and pattern recognition, volume 17, pp. 137-154. San Francisco, 1985.
Junhyun Nam, Hyuntak Cha, Sungsoo Ahn, Jacho Lee, and Jinwoo Shin. Learning from failure: De- biasing classifier from biased classifier. In Advances in Neural Information Processing Systems, volume 33, pp. 20673-20684, 2020.
Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet classifiers generalize to imagenet? In International conference on machine learning, pp. 1389-5400. PMIR, 2019.
Leonid Rudin, Stanley Osher, and Emad Fatemi. Nonlinear total variation based noise removal algorithms. Physica D: Nonlinear Phenomena, 60:259-268, 11 1992.
Marc Russwurm, Sherrie Wang, Marco Korner, and David Lobell. Meta-learning for few-shot land cover classification. In Proceedings the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops, June 2020.
Xiaoyu Tan, Lin Yong, Shengyu Zhu, Chao Qu, Xihe Qiu, Xu Yinghui, Peng Cui, and Yuan Qi. Provably invariant learning without domain information. In International Conference on Machine Learning. pp. 33563-33580. PMLR, 2023a.
Xiaoyu Tan, Lin Yong, Shengyu Zhu, Chao Qu, Xihe Qiu, Xu Yinghui, Peng Cui, and Yuan Qi. Provably invariant learning without domain information. In Proceedings of the 40th International Conference on Machine Learning, pp. 33563-33580, 2023b.
Sang Michael Xie, Ananya Kumar, Robbie Jones, Fereshte Khani, Tengyu Ma, and Percy Liang. In- n-out: Pre-training and self self-training using auxiliary information for aut-of-distribution robustness. In International Conference on Learning Representations, 2021.
Shiji Xin, Yifci Wang. Jingtong Su, and Yisen Wang. On the connection between invariant learning and adversarial training for aatrof-disttibution generalization. Proceedings of the AAAI Confer- ence on Artificial Intelligence, 37(9):10519 19-11111191177.1---111111111 Jun. 2023.
Mengyue Yang. Zhen Fang, Yonggang Zhang, Yali Du. Furui Liu, Wang, and Jun Wang. Invariant learning via probability of sufficient and necessary causes. In Advances in Neural Information Processing Systems, volume 36, pp. 79832-798577 2023.
Ton, Jianhong
Nanyang Ye, Kaican Li, Haoyue Bai, Runpeng Yu, Lanqing Hong, Fengwei Zhou, Zhenguo Li, and Jun Zhu. Ood-bench: Quantifying and understanding two dimensions distriiiitiinttoonn generalization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2022. Oral.
Xiao Zhou, Yong Lin, Weizhong Zhang, and Tong Zhang. Sparse invariant risk minimization. In International Conference on Machine Learning, pp. 27222-27244. PMLR, 2022.
12

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 PROOF OF THEOREM 2
Proof. Part (a).
First, verify that there exists at least one solution to (15). Since A(V: +) is continuous w.r.f (V, b) and (V, +) belongs to bounded and closed set, a solution to max A(V, Φ) can be obtained for any given by the Weierstrass extreme value theorem, denoted as max (#). Let A(Φ)
X L maz ($), b). We then turn to prove that A($) is a continuous function w.r.t. ②. For 21 # 12.
we have
A maz (22), 11 < X mar (21),
(21)
From the continuity of X, for any given there exists some ö1 >> 0, such that for any Φb B( 51), we have
X maz 22), b) < X VV maz ( +1);
(22)
where B 61) denotes an open ball centered at o with radius 61. Letting +2 E 1( in (22)
and combining it with (21), have
A( man (22), 12) - c < A( maz (91). P1).
(23)
Following similar deduction, for the above > 0, there exists some 82 > O, such that for c; 2 82), we have
A(T mar 11), 11) - < X( maz (12), T2))
(24)
Combining (23) and (24), for any given c 0, as long as || 42 ll minfō1, 62], we have
AA man (1) L1) - A( mar (12), D2)l <
(25)
This proves that A()) is continuous w.r.t. +. As for (Ew IV a R(w w o I))? the TV operator and the expectation are taken on the variable w, which does not affect the continuity w.r.t. 4. Define
G( = g( max (t). ①) = [R( D) +
o +)1))
(26)
It can be that G( is the sum, multiplication, and composition of continuous functions w.r.. l, thus it is also continuous w.r.t. +. Using the Weierstrass extreme value theorem again, there exists solution mine G(D), denoted by $h Let v = Va (Tb), then (V/", (T) is solution to (15).
Part (b).
Next, we turn to investigate h(p, W. Φ) in (16). From Appendix A.3 of IRM-TV (Lai & Wang, 2024),
E
Ew+pp R(w
V R(wod,
(27)
where e; denotes the i-th training environment the literature IRM, and p; denotes the environ- ment inference operator for the i-th training environment. (27) indicates that Eut-pl R(w }
is a linear function w.r.t. p E A F. where A is the E-dimensional simplex defined in Section 3.4. Thus Emee-pl [[[ .. R.. (w o [ well as (E.re-p/VV R(w {]] is continuous w.r.t. p. On the other
hand, A(V, +) is continuous w.r.t. V. Define
E(p, 2) == (Ewepll V th R(w +)
(28)
To examine the continuity of A(4, 4 E)p b) w.r.t the entire dual variable (p, V) check its values at two different points (pi, 11 and (P2, U2):
AA () )) A(22, ()E((2 Ø)
< <|> ()E(p1 1)
3) + [A(W1, D)E(p, b) A( D) Φ)
=|X(V,, Φ) (((p,, )) E(p,, 3) (((2, T) 1 ②) A(T2, 1)
(29)
13

---

Published as a conference paper at ICLR 2025
When (p2, U2) + (p1, 11), the right hand side of (29) tends to zero, thus the left hand side of (29) also tends to zero. This verifies that A(V, p @) is continuous w.r.t. (p, V) for any given ④.
Using the Weierstrass extreme value theorem, there exists at least max,, v [A(M, 2) (p, D)] for any given 2 denoted by (p, V) (Φ). Define
solution to
H( = h((,, V man (2) (D), 4) = - " [R( D)1 + max[A(V; 2 E D)}, p.
(30)
which continuous w.r.t. 4 Following similar deductions to Part (a), there exists solution to ming H denoted by D* Let (p", ((* (p. Mm (*), then (p", w,, 2T* is solution to (16).
Part (c).
Lastly, we turn to verify that any solutions (*) to (15) and (p", vV, (-) to (16) are semi- Nash equilibria. We first fix & since & v maz (o*) and (p", ve) (p. L maz ($)). we have "(",", > g( 2*) or h(p", " -*) > h(p, V, 2*) for any v or (p. v in the parameter space. Hence Item 1 in Definition is satisfied.
On the other hand, from (26) and (30), we have G( & G and H()*) & H in the parameter space. Hence Item in Definition is also satisfied. In summary, Theorem 2 is proved.
A.2 PROOF OF THEOREM 3
Before starting the proof, we compute some gradients or subgradients place of the conventional gradients at non-differentiable points for and h. This allows us to maintain the optimization flow with the autograd module of mainstream learning architectures (like Pytorch"). First, the subgradient of V R(w + w.r.t. is:
Jg VV.. V. R(wo))) V V R(mo0)]
if V V. o 4) + 0,
d V R ( 2)|=
(31)
if V. 1) 0,
where Ja [-] is the Jacobian matrix w.r.t. 4, and * is the matrix multiplication. By the chain rule of derivative, the (sub)gradients of g w.r.t. & and V
(v v 2) =E,, R(w +) 2X( 4) E. V,, R(w o +] Ewldg V VwN(w o +)
VoA(W, +) (E.,. V R( +)1))
Vug(v, 4) = A 2) (E [ R +))
(32)
Similarly, the (sub)gradients of / w.r.t. V. and $ are:
vph(o, V, = 3) Ru V R(wob)]] IV Vyh(p, v, 2) = Vv, 1) (Ew+oll V. R(w o D)1)
Oph(p, V +)+)== =E w
Tr [ R)
E
+ 2 D) Bu R( D)
+ VoA(W, 4) (Ewe-p R(w 1)1).
(33)
Proof of Theorem 3. We first investigate the sequence (*))). Combining (17) and (19), we have
1
| ( 1|2 k
(34)
Summing up both sides of (34) from ie = 2 to yields
(K)
dc
< 00.
(35)
kp
CP
çp-1
p-1
2=
hhupss//pyorcch.orr/
14

---

Published as a conference paper at ICLR 2025
It means that the sum of infinite series Σ k= [ <b(h)l2 12 is convergent. Then for any given
0, there exists some N such that 2 [b(k+)) o0(1)12 1|2 Hence for any l >> N and any
m E N_,,
k+m-1
| - 12
Σ
p0(+1) - -
[ - 1 <e.
(36)
J=N
It indicates that { is Cauchy sequence. From the completeness of the parameter space for o ((k) converges to some point d.. Following similar deductions, the sequences { (v)(k) associated with g. { associated with h, and {(o(k), (r(k))} associated with h also converge.
Next, we turn to analyze the computational complexity of this primal-dual algorithm. Without consider one iteration having complexity of O(1), since the complex-
loss of generality, ity within one iteration depends on specific forms of or and relatively constant across different (sub)gradient-type algorithms. Suppose the algorithm needs N iterations in order to achieve a vergence tolerance of c. From the above deduction, we have
|| [
22
Σ
[ 20(1 1 1 dc= N- 5P
1 p p-1
(N - 11p- <e.
(37)
The last inequality in (37) indicates that the convergence tolerance can be achieved long
N > ((p - +1
(38)
By taking the smallest integer for N in (38), we obtain the computational complexity O(((p
F-)).
A.3 SIMULATION DATA GENERATION
The binary outcome of interest (t) with time t [0. 1] is caused and influenced by the invariant features (t) E and the spurious features Xt) E R:
[N(1, 1), w.p. 0.5; xt( (1-1,,1) w.p. 0.5.
sign(( (x))),, sign( (x)))),
W.p. Poi
Y(t)
w.p. 1 1- Pa-
Xx(t) (N(Y(), 1), (t),),
w.p. (();
w.p. 1 (t)
(t) and (t) are further extended to 5 and 10 dimensional sequences by adding standard Gaus- sian noise, respectively. We control the correlation between (t) and (t) by parameter Pa, keeping constant, but change the correlation between (t) and (t) as a function of over the interval [0, 1] by parameter (t). Hence, the training data is generated with parameter settings (PT.P" Pc). where P, and p denote the (t) setting for [0, 0.5) and E [0.5, 1]. respectively. As for the test data, we p. {0.999, 0.8, 0.2, 0.001 } and keep the same pa- Time can be ex- Minimax-TV-61,
ploited as the auxiliary variable EIIL, LIF, TIVA, ZIN, DOD-TV-Minimax-f. and OOD-OVOOVVVA DDDDTT-MMiiaaa--- for environment inference.
A.4 MORE RESULTS OF ADVERSARIAL LEARNING PROCESS
15

---

Published as a conference paper at ICLR 2025
Objective Training OOD-TV-IRM
Parameter Change Adversarial Training OOD-TVIRM
Penalty
|0/t+1)
6 lz
Whole Objective
|u
Value
Value
0.10
0.05
0.00
2100
2150
2200
2300
2350
2100
2350 2400
Epach
Epoch
Objective Value of Adversarial Training for 3.0-TVVMnimmx.
Parameter Change Adversarial Training 000-Tv-Minimas
1.0
Penalty Term
| 1/
Whole Objective
1 - v
0.8
0.04
Value
0.6
Vol bie
0.03
0.4
0.02
w
0.2
0.0
2000
2050
2100
21500
22.00 Epoch
2250
2300
2350
2400
2000
2050
2100
2150
2200
2250
2300
2350 2400
Epoch
Figure Al: Objective value and parameter change of adversarial learning for OOD-TV-IRM and OOD-TV-Minima on simulation data.
16

---

Published as a conference paper at ICLR 2025
Objective Value of Advercarial Training for OQD-TV/ nimax
Parameter Change of Adversarial Training OOD-TV---VVOOO 00D-TV-Minimas
|0/+1) 1/
2.5
2.0
2.0
Value
Penalty Temm
1.5
Whole Objective
1.0
1.0-
0.5
0.5
0.0
0.0
200
400
600
800
1000
0
200
400
600
800
1000
Epoch
Epech
Objective Value of Adversarial Training for DOD-TV-Minimxx aly
Para meter Change Adversarial Training for GOD-TV-Minimax
| e
Whole Objective
0.175
|yit- - 4
0.150
0.125
Value
Valtee
0.100
0.075
20
0.050
10
0.025
0.000
2250
2300
2350
2400
2450
2500
2550
2250
2300
2350
2400
2450
2500
2550
Epoch
Epech
Figure A2: Objective value and parameter change of adversarial learning for OOD-TV-Minimax on the Colored MNIST data set. The top two figures correspond OOD-TV-Minimax--6 while the bottom two figures correspond to
17


Reviews:
[{"review_id": "9TCKRWQPez", "reviewer": "Reviewer_kusC", "summary": "This paper proposes to learn the penalty term $\\lambda$ in IRM-TV using a neural network to improve the OOD generalization performance of IRM-TV. As $\\lambda$ is associated with the feature extractor $\\Phi$, which also requires learning, $\\lambda$ and $\\Phi$ are learned alternatively via minimax adversarial learning. Experiments on a simulated dataset and real-world datasets show that learning $\\lambda$ can improve IRM-TV's performance when distribution shifts.", "questions": "1. Is there any repetition of experiments presented in section 4.2 to 4.4?\n2. How do $g(\\Psi, \\Phi)$, $h(\\Psi, \\Phi, \\rho)$, $||\\Phi^{(k+1)}-\\Phi^{(k)}||_2$ and $||\\Psi^{(k+1)}-\\Psi^{(k)}||_2$ change in real-world datasets as training goes by? Are $g(\\Psi, \\Phi)$ and $h(\\Psi, \\Phi, \\rho)$ able to converge and how long does it take to converge?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Developing methods to learn $\\lambda$ is important to improve IRM, which is a commonly used algorithm in OOD generalization. \n2. The proposed minimax adversarial learning is novel and looks like a promising method to learn $\\lambda$ well. \n2. The writing is clear to follow.", "weaknesses": "My concern mainly falls in the experimental results:\n1. The experiment seems only to be performed once as no standard deviation of the performance is reported. \n2. Adversarial learning could be hard to converge. Analysis like Figure 1 on real-world datasets could be given to facilitate the understanding of how capable OOD-TV is when facing larger $\\Phi$ and more complex datasets.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DLy4GUGlPI", "reviewer": "Reviewer_wMCo", "summary": "Previously proposed Invariant Risk Minimization aims to extract invariant features across different environments to improve the generalization and robustness of the model.  A recent work reveals that the mathematical essence of IRM is a total variation (TV) model. TV measures the locally varying nature of a function, which is widely applied to different areas of mathematics and engineering, such as signal processing and image restoration. In this paper, the authors proposed to generalized versions of IRM with the TV framework and demonstrates the improvements of IRM.", "questions": "See weakness.", "limitations": "", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The generalized framework of IRM provides a novel perspective for IRM and provide new possibilities for improvements.", "weaknesses": "There are several weakness for this paper.\n\n1.The performance improvements are quite marginal. As shown in Table 3, the TV framework IRM can only outperform baselines by a slight margin. Besides, as shown in Figure 1, the curves are quite close with each other. It is suggested to add experiments to demonstrate what the previous IRM cannot do while the novel proposed IRM is good at.\n\n2.It is not clear why this framework demonstrates improvements. As min-max optimizes on the worst-case scenarios, is it possible this could bring potential overfitting issue. Additionally, there are already lots of paper incorporating min-max optimization procedure into IRM. By googling, it can be found that these following papers all use minimax methods and invariant learning:\n[1]Heterogeneous risk minimization\n[2]Invariant risk minimization games\n\n3.No comprehensive experiments on OOD generalization are conducted. This paper only evaluate on the spurious correlation dominated datasets. It can also evaluate the methods on diversity shift dominated datasets with widely-used OOD generalization evaluation benchmark, such as ood-bench. It is known that the IRM cannot deal very well on diversity shift dominated datasets. This can potentially strengthen this paper to point  out clearly the advantage of the proposed method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NZ5JOPHgde", "reviewer": "Reviewer_U2jR", "summary": "This paper studies invariant risk minimization by introducing an additional penalty named TV penalty. The proposed TV penalty aims to quantify the global variation of a function and further enhance the performance of IRM. To realize such a penalty, the authors propose to leverage additional learnable parameters and train them in an adversarial manner. As a result, the learning performance and be further improved thanks to the variance penalty and enforces learning of invariant features. Through extensive quantitative results, the effectiveness of the proposed OOD-TV-IRM is carefully justified.", "questions": "Please see the weaknesses part.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- By incorporating an additional penalty with learnable parameters, the proposed methodology is novel and interesting. The introduction of TV penalty is also novel in the field of invariant learning.\n- The proposed method is theoretically justified, which ensure the training stability when adversarial process is incorporated.\n- The writing is clear and sound, it is not hard to understand this paper.\n- Experimental results shows satisfactory improvement compared to baseline methods.", "weaknesses": "- Motivation needs to be further justified. It is interesting to introduce such a novel penalty to invariant learning. However, it is unclear that why such a penalty is beneficial compared to other techniques, such as ZIN. If such a penalty applies stronger regularization to extract invariant features compared to other methods, the intuitive reason should be further explained. \n- Introducing additional parameters just for a penalty is not computationally friendly. If the application is for large-scale machine learning problems, such a method would be unfavorable. Moreover, it requires extra effort on deciding the hypothesis class of the parameters.\n- Limited number of baseline methods. There are only a few baselines are chosen for comparison. It would be better if more recent and state-of-the-art invariant learning methods are considered.\n- Computational efficiency of the proposed method is not discussed.\n- Missing some related references:  \nHuang et al., Harnessing Out-Of-Distribution Examples via Augmenting Content and Style, in ICLR 2023.  \nYang et al., Invariant learning via probability of sufficient and necessary causes, in NeurIPS 2023.    \nXin et al., On the connection between invariant learning and adversarial training for out-of-distribution generalization, in AAAI 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RZJKxSZZWk", "reviewer": "Reviewer_Hkpi", "summary": "This manuscript is motivated by the practical challenge of Out-of-Distribution (OOD) generalization within the context of the IRM-TV framework. Specifically, prior work on IRM-TV aimed to interpret Invariant Risk Minimization (IRM) through the lens of total variation models. Despite its theoretical contributions, the previous research did not offer any concrete algorithms or practical solutions to tackle the OOD generalization problem (according to the authors' assertion)—an area where IRM is often regarded as a classic approach.\n\nIn response to this gap, the current study introduces OOD-TV-IRM, an innovative methodology that builds upon the insights drawn from IRM-TV. The authors highlight a critical observation: that the penalty parameter, which plays a vital role in regularization, should be tailored to vary across different extractors. This consideration is particularly important for effectively addressing the nuances associated with OOD tasks.\n\nTo implement this idea, the manuscript employs a neural network to dynamically fit the penalty parameter, thereby ensuring adaptability in response to the model's requirements during training. The proposed model engages in an adversarial training process, balancing empirical risk with the total variation penalty. This strategic approach not only enhances the model's robustness but also facilitates its generalization capabilities.\n\nTo validate the effectiveness of their proposed method, the authors conducted experiments using several toy datasets. These datasets serve as controlled environments that allow for a clear examination of the model's performance and its ability to generalize to unseen distributions. The results presented in the manuscript demonstrate the potential efficacy of OOD-TV-IRM in navigating the complexities of OOD scenarios, thus contributing valuable insights to the field.", "questions": "See above.", "limitations": "", "rating": 3, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The manuscript is commendably well written, providing a clear and logical progression of ideas. The motivation behind this study is articulated effectively, emphasizing the goal of addressing the shortcomings present in prior research. This clarity helps to engage the reader right from the beginning, setting the stage for a meaningful exploration of the topic at hand. Additionally, I would like to highlight Section 2, which is particularly well-executed. This section offers valuable insights and serves as an excellent guide for understanding the core concepts of the Invariant Risk Minimization through Total Variation (IRM-TV) framework. The organization and clarity in this section significantly enhance the reader's ability to grasp the foundational ideas that underpin the subsequent analysis.\n\n2. The theoretical analysis presented in this manuscript is notably comprehensive and rigorous. The authors begin with a clear definition of Invariant Risk Minimization (IRM) and Total Variation (TV), systematically leading to the derivation of the IRM-TV formulation. This step-by-step approach allows readers to follow the logical flow of the authors' reasoning easily. Furthermore, the manuscript provides compelling evidence to address the limitations of IRM-TV in effectively managing OOD generalization tasks. By articulating these limitations clearly, the authors lay a solid foundation for the necessity of their proposed method. Subsequently, the manuscript delves into a detailed explanation of the proposed approach, presenting it in a theoretical context. This thorough exploration not only clarifies the mechanics of the new method but also reinforces its relevance in tackling the identified shortcomings of previous works.", "weaknesses": "1. The significance of this work seems to be limited.\n2. The contributions of this work are not enough.\n3. The necessary summary on the theoretical analysis is missing.\n4. The experimental results are not enough.\n\nConcretely,\n1. The overall significance of this work to the broader research community appears insufficient. While the manuscript aims to implement an insight from the IRM-TV framework, it does not convincingly demonstrate the explicit advantages of applying this framework specifically to OOD generalization tasks. Assuming that IRM-TV is a significant advancement—although it is yet to gain citations despite being accepted to ICML-24—the manuscript fails to articulate why its implementation in OOD generalization is strategically important. Clarity on this point would greatly enhance the work's relevance to the community.\n\n2. When evaluating the contributions of this manuscript, it seems that they are primarily derived from insights proposed in IRM-TV, along with the design of an adversarial training paradigm. While the thorough theoretical analysis is commendable, it does not present novel findings that advance the field. The approach of using a neural network to fit hyperparameters is a well-established concept, seen in various contexts such as Neural Architecture Search. Without distinctive contributions that push the boundaries of current knowledge, the manuscript struggles to establish a strong impact.\n\n3. The manuscript presents an extensive theoretical analysis, but it suffers from a lack of necessary summarization. The inclusion of many intermediate derivative processes in the main text can overwhelm readers and challenge their ability to follow the core arguments. While the authors do provide remarks at the end of Section 3.2, there is no clear conclusion to encapsulate the preceding derivations effectively. A concise summary that highlights the key findings would greatly enhance the readability and coherence of the theoretical analysis.\n\n4. The experimental results presented in the manuscript are insufficient to draw robust conclusions. This inadequacy arises from two dimensions: the choice of datasets and the range of methods for comparison. The authors primarily use simulated toy data, CelebA, and Landcover for evaluation, yet they neglect to include mainstream datasets commonly used in the OOD community, such as NICO and Colored MNIST. Given that IRM is known to perform well in scenarios involving correlation shifts, it is imperative to include datasets that exemplify these characteristics—for instance, NICO and Colored MNIST should be at least considered. Furthermore, expanding the range of compared methods would enhance the significance of the experimental findings. Including additional baseline approaches would provide a more comprehensive understanding of the proposed method’s performance relative to existing techniques.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)