Unsupervised Representation Learning - an Invariant Risk Minimization Perspective
Authors: Yotam Norman
Ron Meir
Venue: ICLR 2026 Poster (2026)
Keywords: Unsupervised Learning
Invariant Risk Minimization
Variational Autoencoder
Principal Components Analysis
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We propose and study Invariant Risk Minimization (IRM) in the context of unsupervised learning

Abstract: We propose a novel unsupervised framework for Invariant Risk Minimization (IRM), extending the concept of invariance to settings where labels are unavailable. Traditional IRM methods rely on labeled data to learn representations that are robust to distributional shifts across environments. In contrast, our approach redefines invariance through feature distribution alignment, enabling robust representation learning from unlabeled data. We introduce two methods within this framework: Principal Invariant Component Analysis (PICA), a linear method that extracts invariant directions under Gaussian assumptions, and Variational Invariant Autoencoder (VIAE), a deep generative model that separates environment-invariant and environment-dependent latent factors. Our approach is based on a novel ``unsupervised'' structural causal model and supports environment-conditioned sample-generation and intervention. Empirical evaluations on synthetic dataset, modified versions of MNIST, and CelebA demonstrate the effectiveness of our methods in capturing invariant structure, preserving relevant information, and generalizing across environments without access to labels.

Published as a conference paper at ICLR 2026
UNSUPERVISED REPRESENTATION LEARNING AN INVARIANT RISK MINIMIZATION PERSPECTIVE
Yotam Norman, Ron Meir
Department of Electrical & Computer Engineering Technion Israel Institute of Technology yotamnore & gma com
ABSTRACT
arX arX anXivssubmit/7320453 [cs.LG] 2 Mar 2026
We propose a novel unsupervised framework for Invariant Risk Minimization (IRM), extending the concept of invariance to settings where labels are unavail- able. Traditional IRM methods rely on labeled data to learn representations that are robust to distributional shifts across environments. In contrast, our ap- proach redefines invariance through feature distribution alignment, enabling ro- bust representation learning from unlabeled data. We introduce two methods within this framework: Principal Invariant Component Analysis (PICA), lin- ear method that extracts invariant directions under Gaussian assumptions, and Variational Invariant Autoencoder (VIAE), deep generative model that sepa- rates environment-inng inivarrant and environment-dependerr latent factors. Our ap- proach is based on novel "unsupervised" structural causal model and supports ewwironment-condiiionn
cumple-generation and intervention. Empirical evalua- tions on synthetic dataset, modified versions of MNIST, and CelebA demonstrate the effectiveness of our methods in capturing invariant structure, preserving rele- vant information, and generalizing across environments without access to labels.
INTRODUCTION
Invariant Risk Minimization (IRM), introduced by Arjovsky et al. (2019), addresses the challenge of learning robust models in the presence of distribution shifts across environments (domains). This framework forms cornerstone of the broader field of seeks enable models generalize effectively unseen environments. IRM considers scenarios with accessible training environments and inaccessible test environments, assuming that the under- lying data distribution changes across environments. Within this setting, certain latent features of the data remain stable across environments (invariant features), while others change and are referred to as environmental or spurious features. The goal IRM is to learn a representation of the data that preserves invariant features while filtering out environmenta ones. A predictor or classifier trained on this representation is designed to exhibit robustness to distribution shifts, including shifts occur- ring in unseen environments that share the same underlying generative process. This work extends the concept of invariance in IRM to unsupervised settings, removing the dependency on labels or target values. By redefining invariance in the unsupervised context, unlock new opportunities for research in IRM. Specifically, we propose methods to explore unsupervised algorithms within the IRM framework and demonstrate the feasibility of achieving invariant representations for unla- beled data. We present two novel approaches: Principal Invariant Component Analysis (PICA) and Variational Invariant Autoencoder (VIAE). PICA, grounded in Gaussian and linearity assumptions common in PCA literature, identifies linear transformation that achieves invariant projections of the data. This dimensionality reduction method effectively filters out dimensions subject distri- butional shifts while retaining invariant ones. VIAE, variational autoencoder adapted for unsuper- vised IRM, is aimed at separating the latent space into invariant and environmental parts, allowing causal interventions for both generated samples and data points. This method is empirically tested on datasets inspired by common benchmarks in IRM literature (Gulrajani & Lopez- Paz 2020). mod- ified suit the unsupervised framework. These contributions pave the way for new explorations in invariant representation learning, offering tools to handle distribution shifts in scenarios where labeled data is unavailable and/or expensive.
(OOD) learning, which

---

Published as a conference paper at ICLR 2026
1.1 DEFINITIONS & FORMALIZATION
We consider set of environments Eall, each inducing a distinct probability distribution over the data: Xe pe for each e E &all- The set &all is partitioned into two disjoint subsets: &iaain containing the environments observed during training, and &est, containing environments that are not available during training. At inference time, samples may originate from either Eraain or Etests depending the context and the application.
Random vectors are denoted using uppercase letters (e.g., Zav: Ze), where the subscripts inv and e indicate whether the variable is invariant to the environment Deterministic variables are denoted using lowercase letters. When deterministic variable written function of e, it implies direct dependence on the environment, i.e., f(e) may vary across different environments
wertrooment-deeendeeee respectively.
1.2 SUPERVISED AND UNSUPERVISED IRM
the supervised IRM setup, the data consists of triplets (X, e),,, where X denotes the input vector, Y is the label, and specifies the environment. The data assumed to be generated according
(x, Y) (x,
where is a known deterministic parameter that modulates the joint distribution over inputs and labels (X, Y). The IRM optimization objective, originally formulated by Arjovsky et al. (2019). given by:
min d: +H
Σ ow ; s.l. EEEerain
arg min (1 o)) Ve Etrain-
#
u: FL
Where (.) denotes the risk under environment e. Unlike standard empirical risk minimization (ERM), IRM introduces an additional constraint: the learned predictor w must be optimal for each environment independently. Subsequent work (e.g. Zhou et (2022 Lin et al. (2022): Ahuja al. (202T); Chen al. (2022)) has focused on proposing surrogate objectives and algorithms to approximate solutions to this challenging bi-level optimization problem.
In this work, we extend the IRM framework to the unsupervised setting. Specifically, we investigate whether it is possible to learn invariant representations from unlabeled data drawn from multiple environments
X. (*)),
where the goal is to learn feature map d(X) such that
pei (ø(X)) ped (o(X)) Vel, 62 E Ean-
In other words, the learned representation should be invariant to the environment. To this end, we consider generative model parameterized by 8, with environment--peccffe likelihoods Pj (X). Our proposed objective is to maximize the sum of log-likelihoods across all environments, subject to the constraint that the induced distribution of the learned features of (XX) (also parametrized by 0) is identical across environments. This leads to the unsupervised IRM optimization problem:
max
log P7 (X)
s.t. P)()(X)) Pj (o(X)) Vi,j Errains (2)
eEErrain
where for x. E R and = ((x) Rd. the notation (X, 7) denotes distribution over D+d such that Pr(X dr, Z E dz) = Px()6((= dxz))dxdz.
This formulation shares similarities with unsupervised representation learning approaches such as Variational Autoencoders (VAE) (Kingma & Welling 2014) and Probabilistic PCA (Vidal et al.) 2016), but crucially introduces an explicit constraint enforcing representation invariance across en- vironments. Analogous to supervised IRM, the maximum likelihood term plays the role of the em- pirical risk, while the constraint enforces an invariance property, here defined as equality of feature distributions across environments.
2

---

Published as a conference paper at ICLR 2026
Tuv
Y
Y
ZIne
Zine
X
X
Ze
Ze
Z,
Z,
Anti-Causal FIlF SCM
Causal FIlF SCM
Causal PIlF SCM
Unsupervised SCM
Figure I: IRM Generative Structural Causal Models for supervised (left 3 figures) and unsupervised
(right figure) cases
1.3 GENERATIVE CAUSAL GRAPHS
The foundational assumptions underlying this work, and much of the broader IRM literature, rooted in causality theory (Peters el al., 2017 Schölkopf et al., (02T), specifically particular fam- ily of causal generative processes. These processes are typically formalized using Structural Causal Models (SCMs), which characterize the couse-and-effees relationships among the variables in sys- tem. the IRM setting, prior works commonly assume SCMs that distinguish between "Fully Informative Invariant Features" (FIIF) and "Partially Informative Invariant Features" (PIIF) (Ahuja cct al. 2021). as well as between causal and anti-causal label structures. A central concept across these variants is the decomposition of the latent space into two components: Ze the environment- dependent features, and Zinv, the invariant features that remain stable across environments.
In this work, introduce supervised SCM. This model generalizes previous assumptions, providing a unified framework that encompasses both FIIF and PIIF structures, as well causal and anti-causal generative mecha- nisms. By doing so, lays the foundation for unsupervised invariant representation learning under broader and more flexible generative model.
SCM tailored for the unsupervised setting, which term Un-
2 RELATED WORK
The study of causality predates the formulation of Invariant Risk Minimization (IRM). In particular, Peters al. (2015) established connection between causal relationships and invariance principles, laying the foundation for the development of IRM. IRM was formally introduced by Arrjvsky et al.) (2019), along with its first approximating objective, IRMv1. Following this, substantial research has locused on designing improved objectives and algorithms. Notably, Zhou al. (2022) and Lin et al. (002) proposed methods that perform better in the over-paramest wrer-parameterized regime, while Ahuja et al.
2021 introduced bottleneck-ber basedd approach and also considered the differences between the FIIF and PiIF cases. Lin et al. (2022) also leveraged stochastic networks, similarly to our method, though their work remained within the supervised setting. In Salaudeen & Koyejo (2024), both the invariant component Ziav and the environment-depente parameterized and learned. We adopt similar modeling choice, as data-point reconstruction and generation cannot be achieved without the environmental part of the latent space. A more detailed discussion on this observation appears in Section
component Z. of the latent representation were
On the theoretical side, Rosenfeld et al. (2020) highlighted limitations of supervised IRM, demon- strating that under mild assumptions, an impractically large number of environments may be re- quired to guarantee generalization to unseen environments. Wald et al. (2022) further showed that the interpolation property, common in over-parametrized learning algorithms, precludes invariance, providing theoretical justification for the strategies proposed in Lin et al. (2022) and Zhou et al. (2022). Finally, Toyota & Fukumizu (2023) established that, under suitable assumptions, the IRM objective indeed leads to environment-robus solutions.
Outside the IRM framework, Neria & Nir (024) explore an unsupervised approach to learning representations that optimized for downstream tasks rather than for robustness to distribution shifts. Unsupervised invariant representation learning was also studied prior to the introduction of IRM. For example, Lopez al. (2018) and Moyer et al. (2018) also used the variational auto-
3

---

Published as a conference paper at ICLR 2026
encoder framework toward this goal. Other notable works, relying different sets of assumption and frameworks, include Sun et al. (2016) and Muandet et al. (2013).
PRINCIPAL INVARIANT COMPONENT ANALYSIS
The general objective for the unsupervised IRM problem, as given in equation equation 2 to maximize the likelihood under the constraint of invariant fcaturization. Before tackling the gen- cral case, we focus on a simpler yet instructive setting: the linear and Gaussian case, where Xg _' (( > Ve &. A well-known dimensionality reduction technique under sim- ilar assumptions is Principal Component Analysis (Vidal et al. 2016). PCA falls within the broader class of unsupervised learning algorithms, and in this section, we propose variant tailored to the IRM setting, which we denote as Principal Invariant Component Analysis (PICA). PICA aims to eliminate "environmental dimensions," " equivalently, find an invariant projection across envi- ronments.
Before we dive into the problem at hand, let us make an additional simplifying assumption. We assume that the data in each environment is mean-centered, i.c., Me = 0 Ve E. This can be easily achieved by centering the data using X. X. Ez- p. [X] for each environment. The resulting PICA optimization problem is
Pu Ee [(u x) ; s.t. ||0|22 11 pe p((u X) Vi.j E Erain- (3)
max
eEEerain
The objective seeks a vector such that the random variable's u X second moment (variance) is maximized across all training environments. Intuitively, this ensures that retain the direction containing the most information across environments, similar to standard PCA. The constraint has two parts, first, ||u|| eliminates scaling redundancy. The second constraint where the IRM part of the problem comes into play, this constraint limits to be an invariant direction/dimension.. Given the zero-mean assumption, the problem simplifies to
max ΣuTI, 33g
s.t.
||u122 1, u. T u E Vi,j E Etrain.
(4)
For simplicity, we focus on the two environments case |Erain = 2. In this case, the invariance
constraint reduces to u² (M]- - > 0, meaning that must lie in the null space of (2) >g) i.e, u E ker(E) )g). The objective, meanwhile, simplifies to u. (Σ) + Σz)u. Thus, the PICA solution can be found via the following two-step procedure:
1. Find ker (Σ} -3)
2. Choose u according to max T (Σ) + Σ2)
well
The second step of the exact solution appears in the appendix. Finally, if we're interested in finding dimensionality reduction scheme that keeps the d, dny most varying invariant components, we can simply repeat step two for d, times, choosing each time the next best vector in the null space the one which maximizes the objective.
Algorithm PICA (Principal Invariant Component Analysis)
1: Input: Covariance matrices ΣJ, Σ3, desired reduced invariant dimension d, S diav
2: Output: Matrix U, E Ruxd, containing the top d, invariant principal directions
3: Initialize U, +-
4: Compute the invariant subspace: U {u E R" Σ})u
5: for = 1 to d, do
6: Find ut* arg max u (E) Σ2)u
7: Update U, + concat Ur, u²)
8: Update + 24 {u*}
9: end for
10: Return U,
4

---

Published as a conference paper at ICLR 2026
X_1 Samples 2_2 Samples
K_2 Samplees
K_2 Serrptes
(90) 1401 (3i4 1621
Taxt
Tart
Output of the PICA algorithm with d, = on the synthetic dataset. The projection captures the invariant component shared across both environments.
Figure 2:
3.1 EXPERIMENTING WITH PICA
We illustrate PICA on a simple synthetic example based on the generative process
X = Me(e) Ainy + A.Ze + C,
where
(e(1) [0,0,0]", Mc(2) -0,0,,00,,,,) AND = 1,, A. A. (1,1, (1) 10, (2) 2, Znn NO, 1), -.~ o
N(000.0255)
We generate 1000 samples of X, and 1000 samples of X2. We then apply PICA with target reduced dimension d, = 1. The results are shown in Figure 2 Although the original data divided between two environments that are clearly characterized by different covariance matrices, the projection of the data exhibits a constant distribution across environments.
4 VARIATIONAL INVARIANT AUTO-ENCODER
The main algorithm introduced in this work is the Variational Invariant Autoencoder (VIAE). It follows the standard VAE (Kingma & Welling 2014) framework, but with a key modification- is specifically designed to recover the structure of the proposed "unsupervised" SCM shown in Figure right.. The block diagram in Figure reflects this structure. The algorithm's 's architecture has the following favorable properties:
1. Factorized Latent Space: VIAE explicitly factorizes the latent space into two subspaces: an invariant component Zinv and an environment-specife, component Ze.
2. Latent Space Interventions: VIAE enables interventions on the latent space by sampling
Z. from different priors.
3. Bottleneck Architecture: VIAE enforces narrow latent representation, acting as an in-
formation bottleneck. This aligns with findings in IRM literature (c.g.. (Ahuja al. (021)) suggesting that bottlenecks improve the identification of invariant predictors.
Where properties 2 and 3 are inherited from the baseline VAE framework. VIAE contains one decoder, one Invariant encoder to produce the Invariant part of the latent space, and |Errainl environ- mental encoders (one for each environment), to produce the environmental part of the latent space. Another way look at it is that the decoder's and invariant encoder's parameters shared across environments, while the environmental encoder's parameters are unique for each environment.
By utilizing the causality constraints induced by the unsupervised SCM, we define the latent space and the encoder and decoder probabilities.
For the Encoder, for a given environment e, the prior over latent variables is
P (7 Za) Po (Pmn)) Pe (Za) P(Zin) Pe (Ze);
5

---

Published as a conference paper at ICLR 2026
Invariant Enooder
Zun
Envésnmental Encoder
Riua (Z N.)
Decoder
x. P((
X
Ze
R
PeiXi,
de (dolaeei
Figure 3: VIAE architecture. A shared invariant encoder produces Zinv, while environment-specifie encoders produce Ze. The decoder reconstructs X from both components.
where we used the causal graph to determine that Zinv TL Ze, Zun IL For the invariant encoder, we use the standard prior Zinv NNO, T). As for each environmental encoder, we introduce an environment-sppeinee prior Ze ~ (N) I). He(i) 1 Me(j) Vi,j Eall, enabling controlled sampling from a chosen environment during generation. The encoder posterior probability factor- izes
Pd (Zivv, Ze(X) Ze
X)P = P(Zinv Ze, (Z.
This factorization reflects two key causal properties (Peters et al. 2017). (i) Given Ze, Zinv is conditionally independent of the environment (ii) Zinv and Ze become statistically dependent when conditioning on X, due to the collider structure Zinv + + Ze.
As consequence, the invariant encoder takes both X and Ze as input, while each environmental encoder depends only on X.
The relevant distribution for the Decoder is the conditional likelihood probability (*(X Zive, Ze). Invoking the causal structure, we note that when conditioned on Zinv and Ze, the environment e provides no additional information about X. Thus
pe (X Z.) P(X Z.) = = Z()
This is causal mechanism, meaning that it's independent of its source Z distribution, which means that any interventions (changes in the distribution) Z won't affect the decoder. Note that the decoder does not receive any explicit information of the environment, the invariant (causal) mechanism is stable across different environments, as assumed causality theory.
4.1 DATASETS
We evaluate VIAE on two synthetic datasets of our own design, inspired by benchmarks commonly used the IRM and domain adaptation literature (Gulrajani & Lopez-Paz 2020
We introduce SMNIST (Squares MNIST), variant of MNIST (Bottou et al. 1994) with artificial spurious features. includes two training environments, each using half of the MNIST training set. In 1, a 7x7 white square is added to the top-left corner: in 2, the square appears in the bottom-right. Two analogous test environments use the MNIST test set with squares in the top-right (e 3) and bottom-left (e 4) corners, respectively.
SMNIST:
SCMNIST: Inspired by the commonly known CMNIST ((rrjvvsy et al. 2019). Gulrajani & Lopez-Pat Paz 2020) and many more), define SCMNIST (Single Colored MNIST). SCMNIST con- tains two training environments and one test environment. In training, MNIST digits are encoded into either the red RGB channel for 1 to the green RGB channel for (others set to zero). The test environment (e = 3) encodes digits in the blue channel.
For experiments involving the realistic CelebA dataset, please The fariness section D

---

Published as a conference paper at ICLR 2026
Env romment
Environment
9
9
9
9
9
3
3
B
3
B
En romment
Environment
9
9
9 SMNIST
9
9
3
3
3
8
3
SCMNIST
Figure 4: Generated samples conditioned on fixed Zuv Top row: samples with different Z. drawn
from pi (Z.). Bottom row: samples with different Ze sampled from p2 (Ze). Left side: SMNIST dataset, right side: SCMNIST dataset. Invariant features (in our case, the digits) are preserved for all samples, with stable environment for each row.
4.2 SAMPLE GENERATION
To illustrate the capabilities of VIAE, we demonstrate that for single instance of Zuw, sampled from its prior, the output generated samples from the decoder have the same invariant properties regardless of Ze. On the other hand, the generated samples environment is consistent with the chosen prior of Ze. which means that the decoder able to correctly produce samples from a specific environment without it being provided any explicit information over the desired environment. This is demonstrated in Figure
4.3 ENVIRONMENT TRANSFER
Beyond sample generation, we leverage the VIAE framework address the IRM problem. Con- sidering our unsupervised framework, and specifically the VAE settings, some modifications to the IRM goals are necessary. Traditionally, IRM seeks an invariant representation d(X) such that downstream predictor w 6(X) performs well across all environments. In our case, the invariant encoder naturally plays the role of the featurizer such that 2mn = <(X), but the decoder cannot be interpreted as a classifier nor a regressor. This raises an important question what should a satis- fying restoration look like? and/or what should the IRM objective be in setting? We propose that the IRM problem can be considered as "solved" if all data-points can be "transferred" to single environment, while preserving their invariant content. Since distribution shifts are assumed to arise only across environments, if we manage to take a dataset that is divided between environ- ments, and convert it to an equivalent dataset that contains a single environment, then the problem reduces to an ordinary learning problem, and the IRM problem can be considered as solved. More formally, the goal is to transform each data-point from its original environment e, E Eall into an equivalent sample in predetermined target environment E Etrain, while maintaining the invari- ant features Zun and adapting only the environmental ones Ze- That is, we would like to generate Xe, ~ per Z = ENCi.v (Xe,, Ze.)) )). where ENCiuv is the invariant encoder function. To illustrate, consider the well-known "camels and cows" example, where cows often appear in green pastures and camels in deserts (as appears in many works, including Arjovsky et al.) (2019)). Such dataset contains spurious correlations between the classes and the image background. However, we can transform all images so that both cows and camels appear in the same environment (e.g.. desert), these spurious background cues lose their predictive power. such a scenario, we argue that the IRM objective can be considered as solved, not by removing spurious features entirely, but by aligning them across environments. An alternative, seemingly more "straightforward" idea might be to generate "purely invariant" samples. However, we argue that such representations are ill-defined in many domains. For instance, medical imaging, brightness might be a spurious factor correlated with equipment rather than pathology. What does brightness-invariant X-ray look like? Removing brightness altogether would destroy the data. Thus, transferring all data to a single environment, rather than stripping away the environmental features, serves practical alternative that gets the job done. This task resembles domain adaptation as explored in Tzeng et al.) (2017) and Zhu et al. 2017), but with a key distinction- VIAE enables, in some circumstances, perform environment
7

---

Published as a conference paper at ICLR 2026
Invariant Encodar
iinv
Envirosmental
Oe_ (B Xc)
Decoder
Xei
Eneader
Ze.
Xer
(2.
Zee
Cs
Ct
Mip-ted.
Figure 5: Environment Transfer for 6, Etrain
0
6
9
0
6
9
0
6
9
0
9
SMNIST
6
8
6
6
8
6
8
6
8
0
6
9
0
6
9
0
9
2
6
9
SCMNIST
8
8
6
8
E Epváá
og Etest
Figure 6: Environment transfer examples. e, Etrain for the left part and es Etest for the right part, demonstrated for the SMNIST (up) and SCMNIST (down) datasets.
transfer from unseen source environments e, Etest. To our knowledge, previous methods typically require both source and target environments to be seen during training.
4.3.1 ENVIRONMENT TRANSFER FOR SEEN & UNSEEN ENVIRONMENTS
Consider the case which the source environment has been seen during training, i.e., es E Etrain- Under this assumption, environment transfer can be performed using the following procedure
Sample a data-point from the source environment Xe, ~ p e, Etrain
2. Pass Xe, through the environmental encoder corresponding the source environment to obtain the source environmental features 2., ~ pe, (ZelXo.)
3. Use Xz, and 2r, as inputs to the invariant encoder to get 2ime Zin ~ P(Zinel 2c,)
4. Sample Ze, from the target environment encoder prior: Ze. ped (Ze) = -N(ueeet)
5. Feed Żinv and Ee, to the Decoder get Xz,; Xe, Dec( Gine: Zc.)
This process is illustrated in Figure 5 and demonstrated in Figure 0left).
When the source environment not part of the training set, i.c., e, E Etests we encounter fundamental limitation we do not possess an environmental encoder trained for e,, Consequently, generalization to arbitrary unseen environments not guaranteed the general case. However, for some easier scenarios, may attempt to estimate the environmental features by leveraging the existing environmental encoders. Specifically, by using the approximation
Σ
2,
Ze
P5 (...)...)
Etrain
eEEtroin
That is, we pass X., through every one of the environmental encoders and average their outputs get estimation of Ze.. As shown in Figure 6 (right part), this approach works reasonably
8

---

Published as a conference paper at ICLR 2026
Table 1: Accuracy of the four defined classifiers, on test data-points from training environments. Reported results are mean + standard deviation over 10 runs.
Classifier
SMNIST
SCMNIST
Y12L
0.845 + 0.050
0.832 = 0.072
Y021.
0.362 + 0.041
0.345 + 0.045
ê12 0.556 0.066
0.583 = 0.055
Ee2c 1.0 = 0
1.0
well for the easier case of the SMNIST dataset. However, fails for the more complex SCMNIST dataset. This disparity can be understood using the insights TrommRosenfele et al. (2020). In essence, generalization to unseen environments is possible only when the training environments sufficiently "cover" the space of all possible environments. When this coverage is lacking, such generalization is fundamentally out of reach. For example, in the SCMNIST case, for all of the training environments, the blue color channel always equals zero. This means that even in the basic linear- algebraic sense, the training environments don't span the "blue dimension".
4.4 BACK TO SUPERVISED LEARNING
Recall the original supervised IRM problem as formalized in equation [ The goal is learn an invariant features extractor Z = Φ Xx) such that the prediction y = w o of (X.) is invariant. Al- though our own approach does not adopt this supervised setup, there is natural similarity between the invariant feature extractor in supervised IRM and the encoders in our model. The question we would like to answer is whether our encoders perform well such feature extractors, and if so, does that mean our completely unsupervised algorithm solves the original supervised IRM problem? To
investigate this, consider the feature extractor Ż Zinv Ze] o(X)). Where 2e is the out- put of the environmental encoder and 2nv the output of the invariant encoder. We train four linear classifiers top of these features:
Yr2r. WT2L Zinv
(Label prediction from invariant features)
Yo2r. whI Ze
(Label prediction from environmental features)
ê1 1172
(Environment prediction from invariant features)
(Environment prediction from environmental features)
By "environment prediction" we mean that the original environment from which each data point came serves as label, and the objective is to perform classification based on it. This experiment aims examine whether the model successfully learns a separated latent space, with invariant infor- mation captured Zinv and environmental information in Z.. We evaluate the performance of the different classifiers on test data from the training environments Etrain- We train the VIAE network from scratch 10 times. After each training run, we the aforementioned linear classifiers on top of the encoders output and evaluate their accuracy on the test set. The results are reported in Ta- ble I. Let's analyze these results. The label classifier using invariant features (Yr21) achieves high accuracy (around 0.84 - 0.83), indicating that the invariant encoder successfully retains the relevant class information, in our case, the digit identity, which is invariant by construction. In contrast, the label classifier using environmental features (Ye2t) achieves much lower accuracy (0.36 0.34). While this is higher than random chance (0. 1), suggesting some residual label correlation in the en- vironmental features, is significantly worse than the invariant- classifier. The third classifier (érze) attempts to predict the environment from the invariant features. It yields near-random accu- racy (0.55- 0.58), close to the haseline of 0.5 for trivial classifier with no information. We argue that this is the most important result in this section. Much of the IRM literature focuses on filter- ing out spurious (environmental) features to achieve robust generalization even under distribution shifts where spurious correlations invert (c.g., coWS sandy beaches at test time after only seeing cows on grass and camels in deserts during training). While this result does not theoretically guar- antee the absence of spurious correlations Zmn. it provides a strong empirical indication that the invariant space is indeed, invariant. Finally, the fourth classifier (6.2a) predicts the environment from
9

---

Published as a conference paper at ICLR 2026
environmental features and achieves perfect accuracy (1.0) for both datasets. This demonstrates that the environmental latent space captures all the information needed to distinguish between environ- ments, exactly intended. These results further demonstrate that our model achieves the desired separation in latent space, and motivates the possibility of using similar methods for downstream supervised IRM objectives.
5 CONCLUSIONS AND FUTURE WORK
In this paper, we introduced an unsupervised approach to the IRM problem. We proposed two algo- rithms, each designed to solve the problem under different assumptions and requirements. Looking ahead, we hope this new unsupervised framework opens the door to novel insights and solutions for IRM, ones that are difficult or impossible to achieve within the well-studied supervised paradigm. We outline two concrete directions for future research. (i) Develop theoretically complete scheme for performing environment transfer from previously unseen environments, i.e., é, E &eett (ii) In- corporate more advanced learning architectures improve empirical performance. Regarding the first direction, we helieve that meta-learning approaches (such as MAMI. (Finn et al.) 2017) may enable effective unseen-envinonment transfer in few-shot one-shot settings. For instance, it may be possible to fine-tune environmental encoders to adapt to a new environment using only a few examples. For true zero-shot transfer, however, a different model architecture probably neces- sary. As for the second direction, our work relied on the vanilla VAE architecture as baseline for the VIAE scheme. Leveraging more modern generative models (such as GANS (Goodfellow al.] 2014) and diffusion models (Ho et al.) (0200) could allow us to extend the success achieved on SMNIST, SCMNIST and CelebA to more complex and realistic datasets.
The code we used for the experiments is available at http://ggithhb.com/oommmoo///
6 ACKNOWLEDGMENTS
This research was partially supported by grant no. 2022330 from the United States Israel Bina- tional Science Foundation (BSF), by Israel Science Foundation (ISF) grant no. 1693/22, and by the Skillman chair (RM).
REFERENCES
Kartik Ahuja, Ethan Caballero, Dinghuai Zhang, -han-Chrittophe Gagnon-Audet, Yoshua Bengio, loannis Mitliagkas, and Irina Rish. Invariance Principle Meets Information Bottleneck for Out- of-Distribution Generalization. 2021. URL
CRL https: http::://xxvvvvvvv abs /2tp/22166.. 06607
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant Risk Minimiza- tion. 2019. URL https::t:://// http:///wwwwwww.r/./////
907. 02893
Léon Bottou, Corinna Cortes, John S. Denker, Harris Drucker, Isabelle Guyon, Lawrence D. Jackel, Yann Le Cun, Urs A. Muller, Eduard Säckinger, Patrice Simard, and Vladimir Vapnik. Compar- ison of classifier methods: A case study in handwritten digit recognition. In Proceedings of the 12th IAPR International Conference on Pattern Recognition (ICPR), 1994.
Yongqiang Chen, Kaiwen Zhou, Yatao Bian, Binghui Xie, Bingzhe Wu, Yonggang Zhang, Kaili Ma, Han Yang, Peilin Zhao, Han, and James Cheng. Pareto Invariant Risk Minimization: Towards Mitigating the Optimization Dilemma in
Generalization. 2022.
URL /esttss://arxet org/ abw//206...02 07766
Elliot Creager, David Madras, Joern-Henrik Jacobsen, Marissa Weis, Kevin Swersky, Toniann Pitassi, and Richard Zemel. Flexibly fair representation learning by disentanglement. In Pro- ceedings of the 36th International Conference on Machine Learning (ICML), 2019.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In Doina Precup and Yee Whye Teh (eds.). Proceedings of the 34th International Conference Machine Learning. volume 70 of Proceedings of Machine Learning Research, pp. 1126-1135, 1135. PMLR, 06-11 Aug 2017. URL https: / httss://pccceedingg. mir. wtt.sppeess././//
finn 7a h.ml
10

---

Published as a conference paper at ICLR 2026
Ian J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde- Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. In Advances in Neural Information Processing Systems (NeurIPS), 2014.
Ishaan Guirajani and David Lapez-Paz. In Search of Lost Domain Generalization 2020. URL
attps://azzae org/ ht/p///baa 20..a//2000.. 01434
Moritz Hardt, Eric Price, and Nati Srebro. Equality of opportunity in supervised learning. In Advances in Neural Information Processing Systems (NeurIPS), volume 29, 2016.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic In Advances in Neural Information Processing Systems (NeurIPS), 2020. URL /
2006.11239
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. International Conference on Learning Representations, 2014.
Yong Lin, Hanze Dong, Hao Wang. and Tong Zhang. Bayesian invariant risk minimization. IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR). 2022, pp. 16021- 16030., 2022.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
Francesco Locatello, Gabriele Abbati, Tom Rainforth, Stefan Bauer, Bernhard Schölkopf, and Olivier Bachem. On the faimess of disentangled representations. In Advances in Neural In- formation Processing Systems (NeurIPS), 2019.
Romain Lopez, Jeffrey Regier, Michael Jordan, and Nir Yosef. Information constraints on auto- encoding variational bayes. In Advances in Neural Information Processing Systems, 2018.
Daniel Moyer, Shuyang Gao, Greg Ver Steeg Rob Brekelmans, and Aram Galstyan. Invariant rep- resentations without adversarial training. In Advances in Neural Information Processing Systems, 2018.
Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant feature representation. In Proceedings of the 30th International Conference Machine Learning (ICML), 2013.
Uzan Neria and Weinberger Nir. A representation-leammigg game for classes of prediction tasks. arXiv preprint arXiv preprint arXhv:2403.06977, 2024.
Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference using invariant predic- tion: identification and confidence intervals 2015. URL/atpss///rriiv.o/////
01 01332
Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of Causal Inference: Founda- fions and Learning Algorithms. The MIT Press, Cambridge, MA, 2017.
Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. The risks of invariant risk minimization. International Conference on Learning Representations, 2020., 2020.
Olawale Salaudeen and Sanmi Koyejo.
Causally Inspired Regularization Enables Domain
General Representations 2024. URL
SallaueenzZa/salaudeee
pdd
Bernhard Schölkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio. Towards causal representation learning. Proceedings of
the IEEE, 2021.
A.K. Subramanian. Pytorch-vae.
http:///pss//.ttt.oo..o.////atts VAE 2020.
Baochen Sun, Jiashi Feng, and Kate Saenko. Conrelation alignment for unsupervised domain adap- tation. 2016. htlh:tttps//aarrttvrrr////////1p

---

Published as a conference paper at ICLR 2026
Shoji Toyota and Kenji Fukumizu. 2023. /ta.attpes//xxxx.ooo////////////
Optimality of Invariant Risk Minimization.
Eric Tzeng, Judy Hoffman, Kate Saenko, and Trevor Darrell. Adversarial discriminative domain adaptation. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2017.
René Vidal, Yi Ma, and S.S. Sastry. Generalized Principal Component Analysis. Springer, 2016.
Yoav Wald, Gal Yona, Uni Shalit, and Yair Carmon. Malign overfitting: Interpolation can provably preclude invariance. 2022. URL
Rich Zemel, Yu Wu, Kevin Swersky, Toni Pitassi, and Cynthia Dwork. Learning fair representations. In Proceedings of the 30th International Conference Machine Learning (ICML), 2013.
Xiao Zhou, Yong Lin, Weizhong Zhang, and Tong Zhang. Sparse invariant risk minimization. 39th International Conference on Machine Learning, PMLR 162:27222-27244, 2022., 2022.
Jun- Yan Zhu, Taesung Park, Phillip Isola, and Alexei A. Efros. Unpaired image-to-image translation using cycle-consistent adversarial networks. Proceedings of the IEEE International Conference Computer Vision (ICCV), 2017.
12

---

Published as a conference paper at ICLR 2026
Appendix
A PRINCIPAL INVARIANT COMPONENT ANALYSIS
A.1 ALGORITHM'S S SOLUTION
In this section, we'll 'lll focus on solving the two-steps procedure of the PICA algorithm described in The first step done by simply solving for (Σ) Σn 0. The solution to the second step can be obtained by setting u = Uv, where v E Rddov a normalized vector, and U E [R Dxduv is a matrix whose columns span U. Substituting this term into the optimization objective yiclds
max v U (5 E
Uv
s.t. 110122 1.
This equivalent to the standard PCA objective, with U (Σ) + Σ2) U playing the role of the covariance matrix. The optimal u is given by Ur. where is the eigenvector associated with the largest eigenvalue of U (E) + For the subsequent u2 U3...... ud, dimensions, we simply set u, Uv, where v, is the i-th eigenvector of U (21 22)U, corresponding the i-th largest cigenvalue.
As final note, address the case with more than two environments, |E] 2. Although the derivations for this case were not carried out as part of this work, we believe that they can be conducted as natural extension for the environments setting. The solution in this case would most likely involve finding the direction that maximizes the quadratic form over the sum of all environments, subject to the constraint that the projection lies in the intersection of the null spaces of the pairwise difference matrices between all environments. The resulting two-step procedure is as follows:
Compute
u = n ker 5/ t.fEErain
2. Choose u according to
Σ
max u well
u
VeE
complete derivation and empirical validation of this proposed extension left for future work.
A.2 ANALYTICAL DERIVATION OF PICA
We motivate the PICA algorithm using the following assumed underlying linear Gaussian model
X. = =ee(e) + A Zine + A +
(5)
Where
X. E R2D
Zinv E Rdo,,
Zin (0,
Ze E 1dhh
Z. No0,
V(O, Gaussain noise
From the causality graph] 1 Zin 11. Ze, = = A3 Vi,j
Here, (e(e), Ame A, are deterministic with appropriate dimensions. Importantly, (e( and (e) vary across environments, hence written explicitly functions of e, while the matrices Any and
13

---

Published as a conference paper at ICLR 2026
A. remain stable across environments. From this model, the covariance in enviromment can be derived as
D A +o(()).AA + +
Consequently, we have
Σ E = (o2
- (2) + =
2A 2n A
( (2()) + 20
single environment case and we simply perform standard PCA as demonstrated in Vidal et al. (206). Thus, we henceforth assume (1) # (2). We can see that Σ Σ contains only the environmental data component. There- fore, by choosing principal components from its null space, we remove environmentt-eependen dimensions. Further evaluating the objective yields
First, observe that when (1) (2) the problem reduces
u (E} + Σ 2u (Aiav- o. u+ (12(1) + (AcA.)
2u (Ainv + o² u.
Where the second equality is due the second term in vanishing for in the null space of A. Assuming that the invariant information signal Ain Ziny dominates the noise e, we get that maxi- mizing the objective under the IRM constraint will yield a vector in the direction that maximizes the invariant variance, as we would expect from an IRM version of PCA.
A.3 PROBABILISTIC PICA ALGORITIM
Building on the derivation presented Vidal al. (2016), we proceed develop a generative probabilistic (linear) model tailored our IRM setting. By learning a generative model, we can sample new synthetic examples from the underlying data distribution.
Consider the generative model presented in 5 The environment-specific expecttttttttttttttttttt and covariance matrix are
Px. = Me(e), - Aiv A + (e), + + = Amy A (e()), + oZ1.
For simplicity, assume from now on that diav de = d. In order to reconstruct the generative model, the objective Probabilistic PICA is to extract the model's parameters (e(e), (e)), Anv, Ae, o² from the population mean and the population covariance. An immediate result is that for every environment, the environmental mean simply equals to the population mean over the environment. Consider the covariance matrix, Since both Ainv and A. have rank of d, their D - d smallest eigenvalues are equal to zero, which means that the D 2d smallest eigenvalues of the covariance matrix equal o?. In order to simplify the upcoming derivations, we consider from now on the two environments case [ Consider
E (02(1) (2(2)AA44
Without loss of generality, assume that (1 2 (2). which yields that the subtraction matrix
Σ3 Σ is positive semi-definite (PSD). The special case of (1) = (2) corresponds to the degenerate case in which there is no distribution shift. Therefore, we'll assume o² > o2)... By using the spectral decomposition of the subtraction matrix, we get that the estimation for Ag is
1
Vo ((() U.A:/2.
With U, containing the d eigen vectors corresponding the the highest eigenvalues of Σ Σ and Ag is diagonal matrix containing the eigenvalues of Σ/ Σ Consider the mean of the covariance matrices
1
Σ²
(2 + (2))
A ATT
A.A + +021
2
Σ1 +
02(1) + (2
AT
[ --2-
- 00?
2
210 02(2))
Define
El + E
o2(1 02 [1 - )の) - -021
M(O2(L). 2 A
AT
2
(102(1)-o
14

---

Published as a conference paper at ICLR 2026
The resulting estimation of Any is
Any = U
With Ug containing the cigen vectors corresponding the the highest eigenvalues of M
and Aa is diagonal matrix containing the eigenvalues of M (1) arranged in decreasing order. In order to complete the generative model, an estimation for o and (2) is needed. Without loss of generality, we can set (1) = 1, as this is simply a matter of rescaling Åg appropriately. Now, all that is left is to find population-level term for ae (2). Consider the following term:
tr 1r D2
(o2(1) - (2)
(6)
tr tr(),) Ir Ar(Ainv A inv) (( o²
(1 (2 o??D
tr(Aiuv ) + tr
In order to proceed, we must make further assumptions. Consider the case for which the envi- ronmental signal is significantly stronger than the invariant one. this could also be looked as assuming that most of the information originates from the environmental features. This assumption reflects more challenging IRM setup, where weaker invariant signal is buried much stronger environmental "noise". Consequently, assume that the environmental features are "stronger" in the following sense:
tr(A,A)) tr(Ainv Anv o² D.
Using the above assumption together with (1) = 1, we get from that
tr tr(E) Σ
tr (2))
(2)
tr tr(2
tr [-
Further evidence that this value represents form of "worst-case" scenario can be obtained by noting that both tr (Ayy and o² D are positive, meaning that
tr > tr EE
1 == 2 o2)
Consider the assumption (2) < (1). By setting o (2) to its lower bound, the resulting model corresponds to the largest possible distribution shift between environments for the given data.
B VARIATIONAL INVARIANT AUTO-ENCODER
B. ELBO DEVELOPMENT
First introduced by Kingma & Welling 2014), the Evidence Lower Bound (ELBO) in our setting is given by
(z X.e) (log po X Z) (T. (o(e))
For notational convenience in the derivations of this section, we treat the environment as given conditional variable, writing po(- e) rather than embedding it in the probability function in (.).
As reminder from the encoder derivation in section 4
(2( = 90 (iav. Ze|X, e) = (s( X, e)) (Zno Ze. X),
we get
(Z(e)e
=Egg (Zun. Bal ZelX log Ze)] KL/gs ZZ e)gs (Zinv Iee X) |l po((zle)ppoZZnn
15

---

Published as a conference paper at ICLR 2026
Consider the KL term
KL [L Ze X, e 99 (Zivv Ze. X) || po Ze e (o
( [inv))
1o 9o (iuvv Ze: X) )) (Zo)
gol (Zin Ze. 9o XX, log
džnnyZz
pel (Zel e) (Znv)
Siov
/ Fiav
Txl e) dziny
q6l (6 () log
po (o(ee)
9 (Zuvv IZe x)
46 Zine Zos (Z. e)
džædZinv
po (oin)n))
Zlov
90(Ze, (e)
4s (Z X log
(o Zey X X)dzinvdee
Zinve
/
qol go(Zin Zes X)
40 (o,
(Zinel Ze; X) log
dzinv
po (ine)
20
Zinv
KT (Z,(x,e) e) || Po (Zclee)) (x [TT(y0 Zino iZe X) || po ((in)))
The total ELBO term is
Eg pl
Ze e) pa
(Zinv Zey X) po (Znv))
It's important to recall that isn't a random variable, but an index indicating a change in environment meaning a distribution shift. Therefore, for each such distribution, a different set of parameters used, leading each probability function conditioned on to distinct set of parameters unique to that environment. order to differentiate between environment specific parameters and invariant parameters, the environmental parameters will have subscript while the causal ones will have subscript inv
Eag, (Z) (Fgenao (Zuv X2) [log Po... (X Zanv, Zc)]]
(Tg (Zel (X) po. (Ze)) E. (Ze X) [K[ (Zin Zey X)
|| pome (Zinv)
Assume that everything is Gaussian distributed as a function of the given variables, meaning that
po (Xl Ze)= NNinn (Z.,, Einv (7))
- (Zinv Z, X) = (N (X, Zo), Einv (X, Ze)), Pine (Zin)) = =0( 9d. Zg X)) (X)(()). (x)),, po. (Za) (N(u(ee) Ve
Resulting in
log Joh (Ze
(Zelx) ||
po, (2e))
Eu (2.lx)
pa_ (Ze)
(2TT) 1// 112 (2T))
1/2 (=-MM (x( (
Ego. (2.
log
d/2g -d/2 (2m)_d/2ett(-mm(()) (==- (x-m((e)
Using the common assumption that the covariance matrices are diagonal, we get
de
Ex (X) (ma(X); (a 1-1--llg-- log Ex (X):z)
Similarly, the invariant KL term be written
Egg. (Z. (L (Zinvl Ze, X) x) po... (Zavv)
diav
= =Ego. (2,)X)
Ein (X, Ze)se Minv (X, 2c) - -1 loog Ein (X, Zo)
2 1=1
16

---

Published as a conference paper at ICLR 2026
The expectation Eg,, be approximated using sampled data point, analogous the procedure in stochastic gradient descent
Einy (x, (x())
0
E(X)
Ex(X)
Minw (X, Ze (X))
(xx)
(X())
O (c(e)
mee) =
The resulting total KL term is simply
/nvl/ee (Z(gg(Ze|Xxe) po po(Zele)) + E (KL (Znvl Ze, X) || po (inv))
[K
[E()a + (m(X); u(e):)² - -log E(X)s)
C ENVIRONMENT TRANSFER ANALYSIS
C.I PRELIMINARIES
Let us return to the linear case. Assume the data generating process
Xe = (e(e) + AineZinu + Ae + E.
Where
X₂ D
Zinv E Ruy,, Zinv ~ I)
Ze E Rd, Ze ~ (o.
NN(0,, Io Gaussain noise
From the causality graph: Zinv Ze, = Vi,jEε
Using this model, we would like to find closed-form analytical expressions for the en- coders and the decoder operations under linearity constraints. The environmental en- coder, invariant encoder and decoder weights are obtained by maximizing the probabilities Pa (Z.) P(Zmyl Xe), P(Xe|Zine Ze) respectively. In the linear Gaussian case, this objec- tive reduces to minimizing the mean squared error (MSE), for which the optimal estimators are the corresponding conditional expectations-
Ep- (Ze X2), E(Zinv Ze. X.). E(XelZan Ze)-
Under the linearity assumption, the conditional expectation estimator is the Wiener estimator of the form:
E|A|B] E[A] + COV(A, B)VAR (B)B E(B)
Let us derive it for our model.
The environmental encoder is Ep- Xx]. The different components are:
B|2.] 0
E[X.] = (e(E)
COV Xe] = de(e)AT
VAR(X.) Ain A + + +
Leading to
2. = P7 de(e) (Amn ATT + =c(e)AAAA + 2(x, Ha(e))
17

---

Published as a conference paper at ICLR 2026
It's important to mention that the variance is invertible only due to the noise, the other matrices which compose are of smaller rank. we assume zero noise and/or want to avoid inverting this close-to-sine -iniguullnyy matrix, computational method should be used instead of an analytic one. The invariant encoder is F Zmv Xe, Ze]- The different components are:
E ime]
0
X
Me (e)
E
Ze
X
COV
Zinvs
[A
Ze
Xxe
Ame + Ae) + + [e(e)
(c) oe(e)1
VAR
Ze
Using the fact that the block matrices on the diagonal of the variance matrix are invertible (again, thanks to the noise), we can invert the variance matrix using block-wise inversion formula and get
Ainv A iuv
A. + Te [e(e)A
Ac(
(e))))
(Ainv iav
+ 0. Ae)
-
)-1 (Aim iny + + - -2(1- AA -e de(ee
-A (Ain + (e
+ +o
(Ainv A + +021)-1
-AAine my + +) Ar
AA (Amn A + +o2.) 0)
[e(e)]
Which yields
Zinv E Zine Xx, Z.] A (Aine + (X. me(e) AnZa).
Observe that the invariant encoder turned out to be independent of (e).
The decoder function E[Z.|X.J. The different components
E[ Me(e)
Zinv Ze
E
Zinv Ze
COV
Xex
[Auvv (e))A)
7inv
I,
0
VAR
Ze
0,
(e)))
We can easily invert the variance and get
I
I. =
0.
de)()
0.
of (c)1
and we finally get
X.
Zine; =c]= + Ainy Zin A.Zo
That is exactly the generative process (minus the noise), which is nice and intuitive result. After obtaining analytical expressions for each of the functions of the autoencoder components, the next step analyze the specific cases for transferring between seen environments and transferring from unseen environments.
C.2 ENVIRONMENT TRANSFER FOR SEEN ENVIRONMENTS
For seen environments, the environmental encoder that of the source environment C. E Etrain'
2., = Ep., X.] (e,)) (Ai + +ee + + -1 (X, - -Me(es))
The invariant encoder directly dependent e, solely through the centering of X, (the mean term):
2m = Zinvl Xe,: Ze,1 = A (Ainv + o2 (Xe, - (e.())) - Ae
Finally, the decoder environmental component from the target environment:
X Fx Zinvy Zea]] (ex) + Ainy 7un A.Ze,
18

---

Published as a conference paper at ICLR 2026
C.3 ENVIRONMENT TRANSFER FOR UNSEEN ENVIRONMENTS
For 6s E Erett. we assume that Ae; Aie, of and pr(e), Ve E Etrain are known, and (e(es), o ea E Eeet are unknown. This requires us derive the environmental encoder again Ze, Eper IZe Xe].
E El
E[Xe,] =Me(es) unkown
COVIZe, Xe,] (ea))A de(es) unkown
VAR(X,.) Ana Abav +o] +oo:: oe(e.) unkown
C.3.1 HEURISTIC MOTIVATION
Let us give a quick motivation for our heuristic method for unseen environment transfer. For unseen environment transfer, we're trying to estimate an "equivalent" target data point from a seen envi- ronment Xe,, e E Etrain data point from an unseen environment X e, &vess using the data we gathered from our training environments. We can therefore consider the following MSE optimization problem
Xz, arg min X
Σ
Epe
[1
X 12
eEEtroin
The minimum obtained by differentiating and locating a stationary point
Σ eEEtrnin
(2Xe,
2Ep (x.)
and the resulting estimation is
Ep. [Xo]Xe,]
Etrain
egEreain
This derivation can also be repeated for Ze, instead of Xo., that way, we can derive an expression for the environmental encoder.
For Ze, we get that
arg min Ze.
Σ Ep
[116 Ze, 11
Ze
eEErrnin
Ze,
Ep [Zo|
|Errain
eEEtrain
(e) (A + +0?1) (Xe, Mo(e)).
Errain
eEEtratn We use this result as motivation behind our heuristic estimation function of Ze, for unseen source environment, which is
Ze,
Σ Zo
Zo p=(()Xe.)-
Errain
eEEtruin
D APPLICATIONS TO FAIRNESS
Until now, we have examined the IRM framework primarily through the lens of robustness to distri- bution shifts across environments. However, previous works, such as Hardd et al. (2016) demonstrate that similar approaches can be effectively applied in the context of algorithmic fairness between sub- populations. In this view, the environmental features correspond to sensitive attributes such as race, gender, or age, variables we would like the model to be invariant to. This perspective especially relevant in socially sensitive applications such hiring, admissions lending, where it is essential
19

---

Published as a conference paper at ICLR 2026
to ensure that decisions made based meritocratic and relevant factors (e.g., qualifications, perience, or financial stability). rather than irrelevant or discriminatory ones. Extensive research has been devoted to learning fair representations, underscoring the significance of this problem. Notable contributions include Zemel et al. (2013), Creager et al. (2019) Locatello et al. (2019).
Within our framework, or discriminatory attributes as environmental features. In this section, we explore the potential of the VIAE algorithm for promoting fairness by separating invariant (relevant) and environment-specffic (irrelevant) components in the learned representation.
interpret these relevant factors as invariant features, and the sensitive
We use the architecture suggested in [Subramanian (2020) for the decoder and invariant encoder, while using a reduced encoder variant for each environmental encoder.
D.1 CELEBA DATASET
To demonstrate VIAE in the context of fairness, we use the CelebA dataset (Liu et al. 2013). which contains over 200, 000 face images of celebrities annotated with 40 binary attributes. We focus on the "Male" attribute to define two subpopulations: "male" and "female". In our setup, this attribute serves as the environmental variable, representing sensitive feature aim to disentangle from invariant content. This is natural choice, as gender-related biases in facial recognition and classification systems can be a source of significant ethical concerns in real-world applications. Our goal is to encompass the gender-related characteristic of the CelebA images in the environmental part of the latent space, while the invariant part contains the other, non gender specific features. Under this setting, we expect VIAE to generate gender- speccffc reconstructions conditioned on the appropriate prior, enabling independent control over invariant and environment--pecifee aspects of the generated samples. Furthermore, VIAE allows for environment transfer: converting given image from one subpopulation (gender) to the other, while preserving the invariant features such as facial structure, expression, pose and on. This setup provides concrete use case for testing the fairness- prosooiin capacity of VIAE: a model that can modify sensitive attributes without affecting adentity-releven features preclude discriminatory prediction biases and promote fairness.
D.2 SAMPLE GENERATION
Following the approach used for SCMNIST and SMNIST datasets, Figure illustrates the genera- tion behavior of VIAE on the CelebA dataset. Similarly to the SMNIST/SCMNIST cases, we draw single sample of the invariant features Zin and generate images using five different samples from each the two environmmnt-speccfficc priors. The model successfully identifies the unique character- istics of each subpopulation, corresponding to the "male" and "female" environments. It produces coherent samples for each group, with the top five images generated from the "male" environment prior and the bottom five from the "female" prior. Notably, we notice a similarity between the different samples, suggesting that they share some set of features which the model learned to be "invariant gender" in one way another.
D.3 ENVIRONMENT TRANSFER
We further evaluate VIAE's ability to perform environment transfer between seen environments. Specifically, we change the environmental component the latent space to transfer images from the "male" domain to the "female" domain, while keeping the invariant representation fixed. As shown in Figure S the algorithm successfully transforms "male" images into "female" counterparts. Im- portantly, several identity-related features-stf as facial structure, expression, and pema-remain somewhat consistent between the original and transferred images, indicating that these attributes are captured in the invariant representation. We emphasize that this experiment is not intended to com- pete with current sfate-of-thh-ar art generative models in terms visual fidelity or realism. Rather, serves qualitative validation of VIAE's capacity to disentangle and manipulate environmental features independently. These preliminary results highlight the potential of the VIAE framework for controlled generation and fairness applications, motivating future work to explore more advanced architectures and training strategies.
20

---

Published as a conference paper at ICLR 2026
Environment 1
Environment 2
Figure 7: CelebA sample generation. The top five images created using five samples from Ze ~ (1), I) (first environment, "male"), and the bottom five using five samples from Z ~ NM (2), I) (second environment, "female"), while keeping Znny fixed.
21

---

Published as a conference paper at ICLR 2026
source:
source:
source:
Original Samples
source: 1
source:
source: 1
source: 1
source:
source:
Recon- structions
source: 1
source:
source: 1
Environment Transfer
Figure 8: Environment transfer from the "male" subpopulation to the "female" subpopulation. The top two rows show original images X the next two rows show their respective reconstructions X and the last two rows show their respective transferred versions X..-
22


Reviews:
[{"review_id": "EzpbfBh59J", "reviewer": "Reviewer_vC27", "summary": "The paper proposes an invariance principle for unsupervised learning, where an optimal reconstruction is sought from a latent representation that is constrained to be identically distributed along the training environments. The principle is demonstrated on Gaussian data, then a tailored VAE architecture for the problem is presented and demonstrated on MNIST variants and CelebA.\nThe empirical results show the model learns latent representations that hold domain invariant features, and another latent representation that holds domain dependent features.", "questions": "What is the conceptual difference from methods mentioned in the \"weaknesses\" part? The presence of a label is a technical difference, but the motivation for replacing the label with a reconstruction loss seems somewhat weak.\n\nIt is mentioned that a separate encoder is trained for each environment. Is it actually a separate set of weights being trained, and if so then why? It seems more natural to train an encode that takes a one hot encoding of the environment, and will leverage the larger combined dataset to train its weights, while also possible learning similarities between the domains.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Overall I liked the topic of the paper, the discussion presented and the developed methods.\nI think there are some original contributions that are presented clearly, and the work could attract some interest from crowd interested in these topics.", "weaknesses": "There are some apparent weaknesses that I think the authors should take into account when revising the paper:\n1. The motivation for the problem is not entirely clear. \"Risk minimization\" is a term mostly used in the context of a prediction problem, and I think the unsupervised setting is inherently different, hence a more suitable name for the work or solution might be something like an Invariant autoencoder/Environment-Invariant autoencoder etc. Now given this framing, it is not entirely clear what one has to gain from invariance in the unsupervised case, and what is the spurious correlation present at training time. Are we expecting the reconstruction error to be stable or min-max optimal on new environments, or some other form of robustness? The unclarity about motivation is especially apparent in section 4.2 where the authors link the problem back to supervised learning. I was not fully able to follow the claims in this section, and I think some more conceptual clarity is required either via math (as suggested in point 3 below), or a more convincing empirical problem. \n2. The idea seems quite similar to works from the domain adaptation and generalization literature, like DICA, DANN, CORAL, and Domain Separation Networks [1, 2, 3, 4]. Some of them are more closely related than others to the work under review, but it seems important that the authors refer to these works and explain the conceptual differences. Other than the presence of a label, the invariance principle where one wishes to learn a representation $\\Phi(X)$ such that $P_{e}(\\Phi(X)) = P_{e'}(\\Phi(X))$ for each $e,e'\\in{\\mathcal{E}_{train}}$ is shared across these works.\n3. The paper could benefit from some additions like a mathematical result demonstrating that the approach achieves some well-motivated desirable property on the linear-gaussian setting. Another interesting aspect could be to discuss the causal version of the \"unsupervised\" graph in Fig. 1, i.e. $Z_e \\leftarrow X \\rightarrow Z_{inv}$ and derive the corresponding architecture.\n4. Finally, a clear drawback is that experiments are performed in rather small datasets and carefully designed problems that follow the assumptions of the method.\n\nSmall comments: It might be worthwhile enlarging Figure 1 and explaining the terms FIIF and PIIF formally. There are several places that use terms like \"recover the causal structure\" (line 248), \"causality constraints\" etc. I think it's better to keep the word invariance rather than causality, because recovering the causal structure might allude some readers to think the work tried to do causal discovery or some sort of structure learning, which is not the case. \n\nOverall, while I have several reservations about the paper, I gave an overall score of borderline accept and will reconsider it upon the authors' response. \n\n[1] Bousmalis, Konstantinos, et al. \"Domain separation networks.\" Advances in neural information processing systems 29 (2016).\n[2] Muandet, Krikamol, David Balduzzi, and Bernhard Schölkopf. \"Domain generalization via invariant feature representation.\" International conference on machine learning. PMLR, 2013.\n[3] Sun, Baochen, Jiashi Feng, and Kate Saenko. \"Correlation alignment for unsupervised domain adaptation.\" Domain adaptation in computer vision applications. Cham: Springer International Publishing, 2017. 153-171.\n[4] Sicilia, Anthony, Xingchen Zhao, and Seong Jae Hwang. \"Domain adversarial neural networks for domain generalization: When it works and how to improve.\" Machine Learning 112.7 (2023): 2685-2721.", "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": "5meKVoEK4Z", "reviewer": "Reviewer_wMwF", "summary": "This paper proposes unsupervised Invariant Risk Minimization (IRM), extending the IRM framework to settings without labels by redefining invariance through feature distribution alignment across environments. The authors introduce two methods: Principal Invariant Component Analysis (PICA), a linear Gaussian approach that identifies invariant directions via null-space projections, and Variational Invariant Autoencoder (VIAE), a deep generative model that factorizes latent representations into environment-invariant (Z_inv) and environment-specific (Z_e) components. The framework is evaluated on synthetic data, modified MNIST variants (SMNIST, SCMNIST), and preliminary experiments on CelebA.", "questions": "- Can the authors provide some grounded justification for distributional equivalence?\n- Can the authors demonstrate a practical use-case, backed with results, of the presented method?", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 1, "strengths": "- The problem formulation--extending IRM to unlabeled multi-environment data--appears to be novel.\n- The two-method approach (linear PICA + nonlinear VIAE) provides complementary perspectives with clear mathematical exposition.\n- The fairness application demonstrates a natural use case where environment-invariant features correspond to removing sensitive attributes.", "weaknesses": "- One of the core weaknesses is that the objective is conceptually ill-defined. The paper redefines invariance as matching the marginals of the representations across environments but does not justify how this invariance serves IRM's goal of robust prediction. For instance, the learnt invariant features could be useless for any feasible downstream tasks. \n- Connecting to the previous point, there is no identifiability analysis of the learned representation; the model could learn arbitrary rotations of the invariant features [see 1] and there is no theoretical or empirical justification along those lines. I am also unsure if the learned representation would be coherent in all cases, for instance when data distribution is unbalanced across environments (as an extreme case, consider all 3's in one environment, all 1's in another). \n- The paper is missing extremely relevant prior literature on learning disentangled representations [1, 2, 3, 4] and fairness via disentanglement [11, 12, 13], as well as less critical but appropriate citations works on causal representation learning, IRM and domain generalization [5, 6, 7, 8, 9, 10], among others. Similarly, links between section 4.2 and references in domain adaptation should be fleshed out.\n- Empirical justification is very limited: (1) MNIST-based simple datasets, no baselines - zero comparison to disentangled representations such as $ \\beta $-VAE, or supervised IRM methods (2) No quantitative disentanglement metrics, evaluation relies on visual inspection; (3) Section 4.3's claim that 84% linear probe accuracy \"validates\" the approach lacks context without supervised baselines showing achievable performance.\n\nOverall: The paper explores an interesting question but suffers from fundamental conceptual ambiguity (what is \"invariance\" without labels?) and insufficient empirical rigor (simplistic datasets, no baselines, no disentanglement metrics, no supervised comparisons). \n\nReferences:\n\n[1] Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations - Locatello et al. 2019  \n[2] beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework - Higgins et al., 2017  \n[3] Isolating Sources of Disentanglement in Variational Autoencoders - Chen et al, 2018  \n[4] Disentangling by Factorising - Kim et al, 2018  \n[5] Towards Causal Representation Learning - Schölkopf et al., 2021  \n[6] Weakly Supervised Disentangled Generative Causal Representation Learning - Shen et al, 2022  \n[7] On Learning Invariant Representations for Domain Adaptation - Zhao et al, 2019  \n[8] Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization - Sagawa et al, 2020  \n[9] Learning Optimal Features via Partial Invariance - Choraria et al., 2023  \n[10] Context is Environment - Gupta et al, 2023  \n[11] Learning Fair Representations - Zemel et al., 2013  \n[12] Flexibly Fair Representation Learning by Disentanglement - Creager et al., 2019   \n[13] On the Fairness of Disentangled Representations - Locatello et al., 2019", "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": "pjyuNeaFzn", "reviewer": "Reviewer_XTfo", "summary": "The paper proposes two unsupervised algorithms inspired by the invariant risk minimization (IRM) paper.  PICA (principal invariant component analysis) is in fact quite constrained and therefore less interesting than VIAE (variational invariant auto-encoder) which works by splitting the latent into an invariant component and an environment dependent component.  VIAE can also be used for environment transfer and to possibly recover IRM", "questions": "Imprecise notations: \n- line 88. $(X,e)\\sim P_X^e(X)$.  If $P^e_X$ is a distribution over $X$, then it does not generate pairs $(X,e)$\n- line 91. What's the relation between $P^e_X$ and $P_X$ ?\n- line 98. Should there be an additional sum or an empirical average on the data for each environment.\n- line 102. If the definition of $P(X,Y)$ implicitly depends on $\\Phi$, why write $P(X,\\Phi(X))$ in line 98?\n\nQuestions:\n- What is the relation of PICA and CCA (canonical correlation analysis, Hoteling 1936)?\n  Hoteling is also the inventor of PCA btw.\n- Isn't PICA very constrained as the difference of two empirical convariance matrices  $\\Sigma_1-\\Sigma_2$ can easily have a null kernel.\n- Line 286. Why not the other direction $P_e(Z_{inv},Z_e|X)=P(Z_{inv}|X) P_e(Z_e|Z_{inv},X)$ ?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The notion of unsupervised invariant algorithms is compelling\n- The VIAE approach seems promising (as illustrated by the environment transfer algorithm.).", "weaknesses": "- Some notations are confused (see below)\n- PICA is very constrained because the difference of two empirical covariance matrices $\\Sigma_1-\\Sigma_2$ is likely to have a null kernel.  VIAE is far more satisfactory.\n- Experiments in section 4.2.1 are insufficient. For instance, you could construct a CMNIST problem with labels that only depend on the shape of the digit, and a RevCMNIST problem using the same patterns but with labels that only depend on the color. The invariant features in the sense of IRM would be the shape for CMNST and the color for RevCMNIST. But since both datasets have the same input patterns, the unsupervised approach cannot make this difference.", "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)