Should Bias be Eliminated? A General Framework to Use Bias for OOD Generalization

## Abstract

Abstract Most existing methods for adapting models to out-of-distribution (OOD) domains rely on invariant representation learning to eliminate the influence of biased features. However, should bias always be eliminated —and if not, when should it be retained, and how can it be leveraged ? To address these questions, we first present a theoretical analysis that explores the conditions under which biased features can be identified and effectively utilized. Building on this theoretical foundation, we introduce a novel framework that strategically leverages bias to complement invariant representations during inference. The framework comprises two key components that leverage bias in both direct and indirect ways: (1) using invariance as guidance to extract predictive ingredients from bias, and (2) exploiting identified bias to estimate the environmental condition and then use it to explore appropriate bias-aware predictors to alleviate environment gaps. We validate our approach through experiments on both synthetic datasets and standard domain generalization benchmarks. Results consistently demonstrate that our method outperforms existing approaches, underscoring its robustness and adaptability.

## 1 Introduction

A widely used example of data bias is the “cow vs. camel classification” problem. As illustrated in Figure [1](https://arxiv.org/html/2507.17001v1#S1.F1)(a), models trained on imbalanced datasets often rely on the biased background features for prediction. Specifically, cows are typically photographed in grassy fields, while camels appear in desert settings. Consequently, the model may mistakenly learn to associate the background context, rather than the animal itself, with the class label. From a causal perspective, this issue arises due to spurious correlations, where the background acts as a confounding variable influencing both the images and the associated label, as shown in Figure [1](https://arxiv.org/html/2507.17001v1#S1.F1)(b).

Figure: Figure 1: Illustration of data bias in the “Cow vs. Camel Classification” problem. (a) An example of the Cow vs. Camel classification task, where both cows and camels are observed in distinct backgrounds, such as grasslands and deserts. (b) A causal graph illustrating spurious correlations introduced by confounders. (c) An intuitive example demonstrating how humans classify an image with ambiguous content as either a cow or a camel.
Refer to caption: x1.png

Most existing methods identify data bias as the key challenge in predicting unseen out-of-distribution (OOD) images and thus learn invariant representations to eliminate its influence.
For instance, Invariant Risk Minimization (IRM) irm aims to learn a representation that supports a single classifier across all domains. Domain Adversarial Neural Networks (DANN) dann employ adversarial learning to encourage domain-invariant representations. Similarly, Self-Training lee2013pseudo improves generalization by generating pseudo-labels for target domain data based on learned invariant predictions. Some de-bias methods qi2020two leverage causal inference techniques, such as the back-door and front-door rules, to discover true causal relationships from spurious correlations. More discussions on related work can be found in Appendix A.

The core assumption of these methods is that biased features are uninformative, and only features invariant across different domains should be transferred.
In this paper, we challenge this assumption and raise a critical question: *Should bias always be eliminated?*
We provide an intuitive motivation for this question in part (c) of Figure [1](https://arxiv.org/html/2507.17001v1#S1.F1).
As humans, when classifying an image with ambiguous content as either a cow or a camel, we often rely on contextual background cues to make an educated guess. This intuitively suggests that, rather than being entirely detrimental, in some scenarios, the bias derived from background information can sometimes serve as a useful reference in prediction. Mathematically, it indicates that the domain-specific (biased) features and the prediction target are conditionally dependent given the invariant content, which conflicts with the assumption of uninformative bias in invariant representation learning.

Although bias can contain useful information, using it for prediction is challenging and potentially harmful. Recent studies chang2019domain; wei2018general; kong2023partial; erm; eastwood2024spuriosity have explored the use of biased features, with the most relevant being Stable Feature Boosting (SFB) eastwood2024spuriosity, which adapts conditionally independent biased features using pseudo-labels from stable ones at test time. However, these methods lack theoretical analysis of the conditions under which data bias can be reliably identified and effectively leveraged. Overlooking this can be risky, as bias often contains both “good” components (e.g., context that aids prediction) and “bad” components (e.g., spurious noise that harms prediction), making it essential to carefully disentangle and isolate the useful information.

To address this challenge, we present a theoretical analysis for the identification and utilization conditions of the data bias.
Our results show that bias can indeed be leveraged, provided there is sufficient distributional variation across domains/labels and the relationship between the biased features and labels remains unblocked. Notably, even in more general scenarios where domain shifts affect the distributions of both the input data and the label, our analysis demonstrates that biased features can still offer valuable information for improving model performance in the inference stage.

Building on our theoretical analysis, we propose a principled generative framework to identify data bias and effectively leverage it for prediction. Specifically, we employ a Variational Autoencoder (VAE) to learn latent representations and disentangle content features from biased ones in an unsupervised manner.
We then leverage the bias for prediction through both direct and indirect strategies. For the direct approach, we extract the beneficial components of the bias by aligning them with invariant predictions. It is achieved by taking the invariant predictions as pseudo-labels to guide the refinement of biased features that filter out harmful bias components, thereby retaining only the useful context for prediction.
For the indirect approach, we utilize the biased feature to estimate the underlying environment context and adapt the prediction accordingly. In particular, we decompose the predictor into two components: an invariant predictor that takes content features as input, and a set of domain experts conditioned on biased features. Assuming that the target domain can be represented as a linear combination of source domains, we obtain a bias-aware predictor by the weighted mixture of experts using the inferred bias, enabling better alignment with the target distribution.
The main contributions of this paper are summarized as follows:

- •
We revisit the role of bias in OOD generalization. Unlike prior work that aims to eliminate bias, we highlight scenarios where bias can be a useful signal for improving prediction.
- •
We provide a theoretical analysis to explore conditions under which data bias can be reliably identified and effectively leveraged.
- •
We introduce a principled framework to leverage bias in both direct and indirect ways and validate our methods through experiments on both synthetic and real-world datasets.

## 2 Theoretical analysis

### 2.1 Problem setup and data generation Process

Figure: Figure 2: A graphical representation of the data generation process. The content variables ${\mathbf{c}}$ are invariant to environment changes, while the bias variables ${\mathbf{b}}$ vary across environments. Observed data ${\mathbf{x}}$ is generated by $g({\mathbf{c}},{\mathbf{b}})$, with ${\mathbf{c}}$ and ${\mathbf{b}}$ forming the latent variable ${\mathbf{z}}$. The environment $e$ affects ${\mathbf{b}}$ but not ${\mathbf{c}}$, and the target $y$ may also depend on $e$, reflecting complex data-environment interactions.
Refer to caption: extracted/6644939/figures/relationship.jpg

Notation and problem setup. We formulate the task as a multi-source domain generalization problem, aiming to learn a predictor that remains robust in an unseen target domain using only data from multiple source domains. Let ${\mathbf{x}}_{k}$ denote a high-dimension image observation ${\mathbf{x}}_{k}:=[x_{1},\cdots,x_{n_{x}}]\in{\mathcal{X}}\subset\mathbb{R}^{n
_{x}}$, and $y$ denotes the corresponding label. With access to $M$ source domains $\{\mathcal{S}_{1},\mathcal{S}_{2},...,\mathcal{S}_{M}\}$, with each domain represented as $({\mathbf{x}}^{\mathcal{S}_{i}},{\mathbf{y}}^{\mathcal{S}_{i}})=({\mathbf{x}}^
{\mathcal{S}_{i}}_{k},y^{\mathcal{S}i}_{k})^{m_{i}}_{k=1}$, our goal is to learn a predictor $f({\mathbf{x}})$ that can generalize well to the new domain $\mathcal{T}$ where $\mathcal{T}\nsubseteq\{\mathcal{S}_{1},\mathcal{S}_{2},...,\mathcal{S}_{M}\}$.

Data generation Process. To illustrate the origin and impact of bias, we model the data generation process using a latent variable model. Let ${\mathbf{c}}:=[c_{1},\cdots,c_{n_{c}}]\in\mathcal{C}\subset\mathbb{R}^{n_{c}}$ denotes content variables and ${\mathbf{b}}:=[b_{1},\cdots,b_{n_{b}}]\in\mathcal{B}\subset\mathbb{R}^{n_{b}}$ denotes the bias.
As shown in Figure [2](https://arxiv.org/html/2507.17001v1#S2.F2), the generation process of the observation ${\mathbf{x}}$ is defined as ${\mathbf{x}}:=g({\mathbf{c}},{\mathbf{b}})$, where ${\mathbf{z}}:=[{\mathbf{c}},{\mathbf{b}}]$ with $n_{z}=n_{c}+n_{b}$ dimension denotes the latent variables and $g:[{\mathbf{c}},{\mathbf{b}}]\mapsto{\mathbf{x}}$ denotes the generating function. To distinguish ${\mathbf{c}}$ and ${\mathbf{b}}$, we introduce an environment variable $e$, where the content ${\mathbf{c}}$ remains invariant despite changes of the environment, whereas the bias ${\mathbf{b}}$ varies as the environment changes. Taking the “cow vs. camel classification” problem for illustration, ${\mathbf{c}}$ denotes the characteristics of the object, like the color/shape of cow; while ${\mathbf{b}}$ represents the background information.
Note that we consider a more general case where $y$ may also be influenced by the environment, allowing for a more complex relationship modeling between the data and the environment. This assumption is reasonable, as the probability of encountering a cow versus a camel is strongly influenced by the environment (e.g., grasslands vs. deserts) in real-world scenarios.
Here we highlight two causal paths between bias and label target, $y\rightarrow{\mathbf{b}}$ and $y\leftarrow e\rightarrow{\mathbf{b}}$, which motivate the direct and indirect ways to leverage the bias for prediction, respectively.

### 2.2 Theoretical conditions for identifiability and inference

Bias identification.
We demonstrate that, with the advancement of representation learning, content and bias variables can be effectively disentangled with an identifiability guarantee. Since disentangling content and bias is fundamental for leveraging biased information, we briefly summarize the mild conditions under which such disentanglement can be achieved.
It is important to note that this identifiability result has already been established in previous work and is not claimed as a contribution of this paper. However, despite the identification results, existing methods rely on invariant prediction (including invariant content variable and invariant relation between bias variable and label). Different from them, building on this identifiability framework, we go further by exploring how to effectively utilize biased data to enhance OOD performance.

###### Lemma 2.1 .

(Block-wise Identification of ${\mathbf{c}}$ and ${\mathbf{b}}$ kong2023partial.) Assuming that the data generation process follows Figure [2](https://arxiv.org/html/2507.17001v1#S2.F2), and the following assumptions hold true:

- •
A1 (Smooth and Positive Density): The probability density function of latent variables is smooth and positive, i.e. $P_{\mathbf{z}|\mathbf{e}}>0$.
- •
A2 (Conditional Independence): Conditioned on $\mathbf{e}$, each $z_{i}$ is independent of any other $z_{j}$ for $i.j\in\{1,\cdots,n\},i\neq j$, i.e., $\log P_{\mathbf{z}|\mathbf{e},\mathbf{y}}=\sum_{i=1}^{n}q(z_{i},\mathbf{e},
\mathbf{y})$, where $q_{i}(z_{i},\mathbf{e},\mathbf{y})$ is the log density of the conditional distribution, i.e., $q_{i}:\log P_{z_{i}|\mathbf{e},\mathbf{y}}$.
- •
A3 (Linear Independence): For any $\mathbf{b}\in\mathcal{B}\subseteq\mathbb{R}^{n_{b}}$, there exists $n_{b}+1$ values of $\mathbf{e}$, i.e., $e_{j}$ with $j=0,1,.\cdots,n_{b}$ such that these $n_{b}$ vectors $\bm{{\mathbf{w}}}({\mathbf{b}},{\mathbf{e}}_{j})-\bm{{\mathbf{w}}}({\mathbf{b}
},{\mathbf{e}}_{0})$ with $j=1,\cdots,n_{b}$ are linearly independent, where vector $\bm{{\mathbf{w}}}({\mathbf{b}},{\mathbf{e}}_{j})$ is defined as follows:
$\bm{{\mathbf{w}}}({\mathbf{b}},{\mathbf{e}})=\left(\frac{\partial q_{1}({
\mathbf{b}}_{1},{\mathbf{e}})}{\partial{\mathbf{b}}_{1}},\cdots\frac{\partial q
_{n_{b}}({\mathbf{b}}_{n_{b}},{\mathbf{e}})}{\partial{\mathbf{b}}_{n_{b}}}
\right),$
(1)
- •
A4 (Domain Variability:) There exist two values of $\mathbf{e}$, i.e., $\mathbf{e}_{i}$ and $\mathbf{e}_{j}$, s.t., for any set $A_{\mathbf{z}}\subseteq\mathcal{Z}$ with non-zero probability
measure and cannot be expressed as $\Omega_{\mathbf{c}}\times\mathcal{B}$ for any $\Omega_{\mathbf{c}}\subset\mathcal{C}$, we have
$\int_{\mathbf{z}\in A_{\mathbf{z}}}p(\mathbf{z}|\mathbf{e}_{i})d\mathbf{z}\neq
\int_{\mathbf{z}\in A_{\mathbf{z}}}p(\mathbf{z}|\mathbf{e}_{j})d\mathbf{z}.$
(2)

Then the learned $\hat{\mathbf{c}}$ and $\hat{\mathbf{b}}$ are block-wise identifiable.

The detailed proof can be found in Appendix B.1 . Such block-wise identification ensures that the learned content and bias representations, $\hat{\mathbf{c}}$ and $\hat{\mathbf{b}}$, are functions of the ground truth content ${\mathbf{c}}$ and bias ${\mathbf{b}}$, respectively.

Use bias for inference.
After identifying bias, we further provide theoretical insights into the conditions under which it can be leveraged to enhance predictive performance. Specifically, we show that the identified bias can contribute to prediction as long as there exists an active (i.e., unblocked) causal path from the bias to the prediction target, meaning the good information carried by the bias is not fully mediated or blocked by the content.

Assumption A5 (Unblocked Influence): There exists a causal path from the bias variable ${\mathbf{b}}$ to the target variable $y$ that is not blocked by the content variable ${\mathbf{c}}$. Formally, ${\mathbf{b}}\not\!\perp\!\!\!\perp y\mid{\mathbf{c}}$.

###### Lemma 2.2 .

The bias ${\mathbf{b}}$ can be effectively utilized for prediction if Assumptions A1 through A5 hold.

As shown in Lemma [2.1](https://arxiv.org/html/2507.17001v1#S2.Thmtheorem1), under Assumptions A1–A4, both content and bias information ${\mathbf{c}},{\mathbf{b}}$ can be well identified. It indicates that we can extract ${\mathbf{c}},{\mathbf{b}}$ containing all the information as the ground-truth variables.
Assumption A5 states that

$$ $P(y\mid{\mathbf{c}},{\mathbf{b}})\;\neq\;P(y\mid{\mathbf{c}})\quad\text{on a set of nonzero measure.}$ $$

In Bayesian decision theory (for any proper loss, e.g. log-loss or squared error), the Bayes-optimal predictor using $({\mathbf{c}},{\mathbf{b}})$ is

$$ $f^{*}({\mathbf{c}},{\mathbf{b}})=\arg\min_{f}\;\mathbb{E}\bigl{[}\ell\bigl{(}y ,f({\mathbf{c}},{\mathbf{b}})\bigr{)}\bigr{]},$ $$

which depends on the full conditional $P(y\mid{\mathbf{c}},{\mathbf{b}})$. The best predictor using only ${\mathbf{c}}$ is

$$ $\arg\min_{f}\;\mathbb{E}\bigl{[}\ell(y,f({\mathbf{c}}))\bigr{]},$ $$

which depends only on $P(y\mid{\mathbf{c}})$. Since Assumption A5 ensures that the conditional distributions $p(y\mid{\mathbf{c}})$ and $p(y\mid{\mathbf{c}},{\mathbf{b}})$ differ, the joint predictor $f^{*}({\mathbf{c}},{\mathbf{b}})$ achieves a strictly lower expected loss compared to any predictor that relies solely on ${\mathbf{c}}$. This confirms that the bias variable ${\mathbf{b}}$ can indeed be exploited to improve predictive accuracy. A detailed proof, demonstrating that incorporating ${\mathbf{b}}$ leads to a better predictor under log-loss or squared error, is provided in Appendix B.2.

###### Lemma 2.1 .

###### Lemma 2.2 .

### 2.3 Understand bias’ contribution in prediction

Furthermore, we explore how bias contributes to the predictions. To illustrate the core insights of our approach, we begin with a binary classification setting. The extension to the multi-class scenario is in Appendix C. The core intuition is that $\mathbf{c}$ and $\mathbf{b}$ are conditionally independent given $y$, which indicates that we can decompose $p(y\mid\mathbf{c},\mathbf{b})$ into three parts, including a bias-aware predictor, an invariant predictor, and a label prior.

###### Theorem 2.3 .

[Decomposition of Marginal Probability]
Given the data generation process illustrated in [Figure 2](https://arxiv.org/html/2507.17001v1#S2.F2), we have:

$$ $\displaystyle\ p(y=1\mid\mathbf{c},\mathbf{b})=\sigma\biggl{(}\underbrace{ \operatorname{logit}\bigl{(}p(y=1\mid\mathbf{b})\bigr{)}}_{\text{Bias-aware Predictor}}+\underbrace{\operatorname{logit}\bigl{(}p(y=1\mid\mathbf{c})\bigr{ )}}_{\text{Invariant Predictor}}-\underbrace{\operatorname{logit}\bigl{(}p(y=1 )\bigr{)}}_{\text{Label Prior}}\biggr{)},$ $$

where $\sigma$ is the sigmoid function. The last two terms combined can be viewed as an invariant part of the model across all domains.

###### Theorem 2.3 .

##### Proof sketch of Theorem 2.3

We begin by reformulating $p(y\mid\mathbf{c},\mathbf{b})$ to make its structure more explicit. Specifically, we have:

$$ $\displaystyle\operatorname{logit}\bigl{(}p(y=1\mid\mathbf{c},\mathbf{b})\bigr{ )}=\frac{p(\mathbf{c},\mathbf{b}\mid y=1)\,p(y=1)}{p(\mathbf{c},\mathbf{b}\mid y =0)\,p(y=0)}\quad\quad\text{consider }\mathbf{c}\,\perp\!\!\!\perp\,(\mathbf{b })\,\mid\,y:$ $\displaystyle\quad=\frac{p(\mathbf{b}\mid y=1)\,p(\mathbf{c}\mid y=1)\,p(y=1)} {p(\mathbf{b}\mid y=0)\,p(\mathbf{c}\mid y=0)\,p(y=0)}=\frac{p\bigl{(}y=1\mid \mathbf{b}\bigr{)}\,\tfrac{p\bigl{(}y=1\mid\mathbf{c}\bigr{)}}{p\bigl{(}y=1 \bigr{)}}}{p\bigl{(}y=0\mid\mathbf{b}\bigr{)}\,\tfrac{p\bigl{(}y=0\mid\mathbf{ c}\bigr{)}}{p\bigl{(}y=0\bigr{)}}}.$ $$

A more detailed proof is provided in Appendix D.1.

##### Discussion on invariant prediction.

We observe that $p(\mathbf{c}\mid y)$ remains invariant across domains, since $y$ d-separate ${\mathbf{c}}$ and the environment variable $e$. Formally, $\forall\,e_{i},e_{j}\in E_{all}$, since $\mathbf{c}$ only depends on $y$, it follows that

$$ $\displaystyle p(\mathbf{c}\mid y=1,\,e=e_{i})\;=\;p(\mathbf{c}\mid y=1,\,e=e_{ j})\Rightarrow\frac{p(y=1,\,e=e_{i}\mid\mathbf{c})}{p(y=1,e=e_{i})}\;=\;\frac{ p(y=1,\,e=e_{j}\mid\mathbf{c})}{p(y=1,e=e_{j})}$ $$

which implies $\tfrac{p(y\mid\mathbf{c})}{p(y)}$ is an invariant quantity across all domains.

### 2.4 Leverage bias for better prediction

As shown in Theorem [2.3](https://arxiv.org/html/2507.17001v1#S2.Thmtheorem3), we have established that the bias affects prediction solely through the bias-dependent component $p(y=1\mid\mathbf{b})$.
Consequently, ensuring the reliability of this bias-aware prediction becomes crucial for effectively leveraging bias. However, relying on this term without additional assumptions poses challenges, as the bias-aware prediction component can vary significantly across environments. In the absence of proper constraints, incorporating it naively would revert the model to standard empirical risk minimization (ERM), failing to account for bias robustness. Furthermore, as shown in the causal graph in Figure [2](https://arxiv.org/html/2507.17001v1#S2.F2), the term $p(y=1\mid\mathbf{b})$ can be decomposed into two components: a direct path $\mathbf{b}\rightarrow y$ and an indirect path $\mathbf{b}\leftarrow e\rightarrow y$. Motivated by this structure, we propose to leverage bias for prediction in both direct and indirect ways. The direct approach uses invariant features to guide the extraction of the predictive (i.e., positively correlated) component of the bias. In contrast, the indirect approach utilizes the bias to infer the environment $e$, which in turn informs the selection of an appropriate predictor tailored to that environment.

#### 2.4.1 Using bias directly

The core idea of directly leveraging bias for the prediction is to extract its beneficial components, guided by the information encoded in invariant features.

###### Assumption 2.4 (Sufficient Information of ${\mathbf{c}}$ ) .

We assume ${\mathbf{c}}$ is informative of $y$. Formally, this means ${\mathbf{c}}\not\!\perp\!\!\!\perp y$.

###### Theorem 2.5 (Correction with Pseudo Labels) .

Consider the predictor $f_{c}:{\mathbf{c}}\rightarrow y$. Define pseudo-label prediction as $\hat{y}=f_{c}({\mathbf{c}})$. Let
$h_{0}\;=\;p(\hat{y}=0\mid y=0),\quad h_{1}\;=\;p(\hat{y}=1\mid y=1).$
Then it holds that

$$ $h_{0}+h_{1}\;>\;1,$ (3) $$

$$ $p\bigl{(}y=1\mid{\mathbf{b}}\bigr{)}\;=\;\frac{\Pr\bigl{(}\hat{y}=1\mid{ \mathbf{b}}\bigr{)}\;+\;h_{0}\;-\;1}{\,h_{0}+h_{1}\;-\;1\,}.$ (4) $$

###### Assumption 2.4 (Sufficient Information of 𝐜 𝐜 {\mathbf{c}} bold_c ) .

###### Theorem 2.5 (Correction with Pseudo Labels) .

##### Justification about Theorem 2.5 .

[Equation 3](https://arxiv.org/html/2507.17001v1#S2.E3) depends and only depends on the [2.4](https://arxiv.org/html/2507.17001v1#S2.Thmtheorem4). The detailed proof can be found in Appendix D.2. Intuitively, if $c$ helps distinguish between $y=0$ and $y=1$, then a predictor based on $c$ should be corrected more often than chance, implying $h_{0}+h_{1}>1$.

##### Proof of Equation 4 .

The key observation is that because $\hat{y}$ only depends on $y$ (through ${\mathbf{c}}$), we can relate the conditional distribution of the pseudo-label $\hat{y}$ given ${\mathbf{b}}$ to the distribution of the true label $y$. Formally,

$$ $\displaystyle p(\hat{y}=1|{\mathbf{b}})=p(\hat{y}=1|y=1,{\mathbf{b}})p(y=1|{ \mathbf{b}})+p(\hat{y}=1|y=0,{\mathbf{b}})p(y=0|{\mathbf{b}})$ $\displaystyle\textbf{Consider }\hat{y}\textbf{ only depends on }y,\hat{y}\perp \!\!\!\perp{\mathbf{b}}|y$ $\displaystyle=p(\hat{y}=1|y=1)p(y=1|{\mathbf{b}})+p(\hat{y}=1|y=0)p(y=0|{ \mathbf{b}})$ $\displaystyle=h_{1}p(y=1|{\mathbf{b}})+(1-h_{0})(1-p(y=1|{\mathbf{b}}))$ $$

Solving for $p(y=1\mid{\mathbf{b}})$ gives [Equation 4](https://arxiv.org/html/2507.17001v1#S2.E4).

#### 2.4.2 Using bias indirectly

An alternative approach is to improve prediction by more accurately estimating the underlying environmental context. Let’s start from the objective, which can be formulated with the following optimization task:

$$ $q^{*}(y\!\mid\!{\mathbf{b}})=\underset{q(\cdot\mid{\mathbf{b}})}{\arg\min}{ \textstyle\mathbb{E}_{e\sim p(e\mid{\mathbf{b}})}\bigl{[}D_{\mathrm{KL}}(p(y\! \mid\!{\mathbf{b}},e)\,\|\,q(y\!\mid\!{\mathbf{b}}))\bigr{]}}.$ (5) $$

This suggests identifying a predictor that minimizes the expected Kullback-Leibler (KL) divergence from $p(y\mid{\mathbf{b}},e)$ across all environments.
The Bayes-optimal solution ohn2023optimal is

$$ $q^{*}(y\mid{\mathbf{b}})=\sum_{e}p(e\mid{\mathbf{b}})\,p(y\mid{\mathbf{b}},e)$ $$

This implies that we can leverage the identified bias to infer the environmental context and then develop a suitable bias-aware predictor $p(y\mid{\mathbf{b}})$ with a set of domain experts $p(y\mid{\mathbf{b}},e)$ to bridge the gaps between environments. The detailed derivation can be found in Appendix D.3.

Figure: Figure 3: Overall framework of the BAG method. The framework consists of three main modules: representation learning, predictor training, and test-time adaptation. In the representation learning stage, we employ a VAE model to disentangle content and bias variables. For prediction, we construct an invariant predictor, a label prior, and a bias-aware predictor that reweights multiple domain experts using a domain estimator. These three predictor components work together for the final prediction. During test-time adaptation, the bias-aware predictor is further refined by invariant guidance, using pseudo-labels generated by the invariant predictor, enhancing its adaptation to new domains.
Refer to caption: x2.png

## 3 Method

In this section, we build on our theoretical findings in  [Section 2](https://arxiv.org/html/2507.17001v1#S2) to propose the Bias-Aware Generalization (BAG, in short) model which uses data bias for generalization. The whole pipeline can be found in Appendix E.1 . The multi-class version can be found in Appendix E.2 .

### 3.1 Overall framework

##### Representation learning and disentanglement.

Following [Section 2.2](https://arxiv.org/html/2507.17001v1#S2.SS2), we first take the observed data ${\mathbf{x}}$ from the source domains as the inputs of the image encoder, which maps the image into the latent space, such as ${\mathbf{z}}=\text{Encoder}({\mathbf{x}})$. Then, we take ${\mathbf{z}}$ into a decode, $\text{Decoder}(\cdot)$, to generate the reconstruction as $\hat{{\mathbf{x}}}$.

We apply the VAE loss to train both encoder and decoder, where the objectives can be written as

$$ $\mathcal{L}_{\text{vae}}\;=\;\mathbb{E}_{{\mathbf{x}}}\bigl{[}\|{\mathbf{x}}- \hat{{\mathbf{x}}}\|^{2}\;+\;\beta\mathrm{KL}\bigl{(}q_{\phi}(\mathbf{z}\mid{ \mathbf{x}})\,\|\,p(\mathbf{z})\bigr{)}\bigr{]},$ (6) $$

Where $\beta$ is a hyper-parameter to balance the reconstruction loss and the KL divergence, which is set to be 1 in our approach. The prior distribution
$p({\mathbf{z}})$ is chosen as an isotropic Gaussian. Minimizing equation [6](https://arxiv.org/html/2507.17001v1#S3.E6) encourages $(\mathbf{c},\mathbf{b})$ to be identifiable in a block-wise sense and captures both invariant and domain-specific factors for further modeling.

Following the method jiang2022invariant, we leverage a necessary but not sufficient condition for $\mathbf{c}\perp\!\!\!\perp\mathbf{b}\mid y$, that is $\mathbb{E}\bigl{[}\mathbf{c}\cdot(\mathbf{b}-\mathbb{E}[\mathbf{b}\mid y])
\bigr{]}=0$. Specifically, we add a regular term to constrain the condition independence, which can be written as:

$$ $\mathcal{L}_{ind}=\frac{1}{n}\sum_{i=1}^{n}\mathbf{c}_{i}\cdot\bigr{(}\mathbf{ b}_{i}-\frac{1}{\bigl{|}\{j:y_{j}=y_{i}\}\bigr{|}}\sum_{j:y_{j}=y_{i}}\mathbf{ b}_{j}\bigr{)}$ $$

##### Predictors

Based on the learned content and bias variables, ${\mathbf{c}}$ and ${\mathbf{b}}$, we can make predictions from the latent space to the label target. Guided By Theorem [2.3](https://arxiv.org/html/2507.17001v1#S2.Thmtheorem3), we implement the label predictor into three components:

$$ $f({\mathbf{z}})\;=\;f_{b}(\mathbf{b})\;+\;f_{c}(\mathbf{c})\;-\;Pr,$ $$

where $f_{b}$ is the bias-aware predictor (i.e. $p(y\mid{\mathbf{b}})$ ), whose input is the bias ${\mathbf{b}}$, $f_{c}$ is the invariant predictor to leverage the content variable ${\mathbf{c}}$. $Pr\in\mathbb{R}^{2}$ is a learnable bias term representing the logits prior $\operatorname{logit}(p(y))$. In particular, $f_{c}({\mathbf{c}})$ and $Pr$ together form an “invariant” module in all domains and can be used directly in the target domain without further modification. Since $\mathbf{b}$ depends on the domain $e$, using the $f_{s}$ only trained in the source domain naively at test time may introduce unwanted biases if the target domain differs significantly from the source domains.

### 3.2 Bias-Aware predictors

To exploit the bias $\mathbf{b}$ in the indirect manner, we introduce a set of $M$ learnable domain embeddings $\{e_{1},e_{2},\ldots,e_{M}\}$, which are corresponding to
$M$ source domains $\{\mathcal{S}_{1},\mathcal{S}_{2},...,\mathcal{S}_{M}\}$.
Each embedding $e_{k}$ captures a distinct “mode” or style that may emerge in the source training data. Intuitively, this design follows the mixture-of-experts principle, allowing each expert to specialize in a subset of the training distribution.

In terms of model design, each $e_{k}$ is a learnable vector (or a small neural component) trained to represent a particular domain bias. We implement a domain classifier $p(e=e_{i}\mid{\mathbf{b}})$ by a softmax layer over linear outputs, to produce the scalar weights
These weights indicate how well each expert $e_{i}$ aligns with the domain-specific bias $\mathbf{b}$.
For each expert, we parametrize $p\bigl{(}y\mid\mathbf{b},e=e_{i}\bigr{)}$ as an MLP network, whose input is the combination of ${\mathbf{b}}$ and $e_{i}$. It models the expert specializing in domain $e_{i}$, making specific predictions based on the latent representation ${\mathbf{b}}$.
Formally, we calculate the bias-aware predictor as a re-weighting process of different domain-specific predictors as

$$ $f_{b}(\mathbf{b})\ =\sum_{i=0}^{M}p\bigl{(}y\mid\mathbf{b},e=e_{i}\bigr{)}p(e= e_{i}\mid{\mathbf{b}}).$ $$

In training, we jointly learn all bias-aware predictions, the domain classifier, the invariant predictor, and the logits prior with a classification loss $\mathcal{L}_{\text{cls}}$.

### 3.3 Test time adaptation with invariance guidance

During test time, when accessing test data, we can further refine the model by exploiting the invariant component $\mathbf{c}$ as guided by [Theorem 2.5](https://arxiv.org/html/2507.17001v1#S2.Thmtheorem5).
Concretely, when we receive test-domain inputs $\{{\mathbf{x}}^{\mathcal{T}}\}$, we first extract their latent representations $[\mathbf{c}^{\mathcal{T}},\mathbf{b}^{\mathcal{T}}]$ using the pre-trained image encoder. We then leverage the learned invariant predictor $f_{c}(\mathbf{c}^{\mathcal{T}})$ to generate pseudo-labels $\hat{y}^{\mathcal{T}}$ on unlabeled test data.

In source domains, we can also have pseudo-labels, so we can collect $h_{0}$ and $h_{1}$ in source domains as shown in [Theorem 2.5](https://arxiv.org/html/2507.17001v1#S2.Thmtheorem5). It’s part of the probability confusion matrix between pseudo-labels and ground truth. Intuitively, these parameters adjust how the domain-invariant information can map the latent style factor $\mathbf{b}$ to valid probability estimates in the new domain.

According to [Equation 4](https://arxiv.org/html/2507.17001v1#S2.E4), recovering the corrected prediction $p(y=1\mid\mathbf{b})$ require knowing $p(\hat{y}=1\mid\mathbf{b})$ and the calibration terms $h_{0},h_{1}$. To estimate $p(\hat{y}=1\mid\mathbf{b})$, we update the bias-aware predictor $f_{b}$ with the test data $\mathbf{b}^{\mathcal{T}}$ and the generated pseudo-labels $\hat{y}^{\mathcal{T}}$.
This process can be viewed as a lightweight test-time adaptation step:

$$ $\min_{f_{b}}\;\mathbb{E}_{\mathbf{b}^{\mathcal{T}}}\Bigl{[}\mathcal{L}_{ada}= \ell\Bigl{(}\sigma\bigl{(}f_{b}(\mathbf{b}^{\mathcal{T}})\bigr{)},\;\hat{y} \Bigr{)}\Bigr{]},$ (7) $$

where the adaption loss $\mathcal{L}_{ada}$ is characterized as the cross-entropy loss with pseudo labels.
We record this post-trained predictor as $\tilde{f}_{b}(\mathbf{b}^{\mathcal{T}})$. Let $\phi(\cdot)=\operatorname{logit}((\frac{\sigma(\cdot)+h_{0}-1}{h_{0}+h_{1}-1}))$ to represent the test-time correction. We denote the corrected bias-aware predictor as $q(y|b^{\mathcal{T}})$. The final prediction is

$$ $f([\mathbf{c}^{\mathcal{T}},\mathbf{b}^{\mathcal{T}}])=f(\mathbf{c}^{\mathcal{ T}})+\phi(\tilde{f}_{b}(\mathbf{b}^{\mathcal{T}}))-Pr.$ $$

Hence, the losses of the framework in both the training and test-time adaptation stages are as follows:

$$ $\displaystyle\text{Training Stage: }\mathcal{L}_{all}=\mathcal{L}_{cls}+ \lambda_{0}\mathcal{L}_{\text{vae}}+\lambda_{1}\mathcal{L}_{ind}$ $\displaystyle\text{Test-time Adaptation Stage: }\mathcal{L}_{ada}$ $$

where $\lambda_{0}$, $\lambda_{1}$ denote the hyper-parameters.

## 4 Experiments

In this section, we test our BAG method on both synthetic data and real-world datasets that require adaptation in the target domain. To make a fair comparison, we follow the setting in jiang2022invariant.

##### Baseline methods

We compare our method with ERM and IRM irm, a classic approach that leverages invariant information for making predictions. Additionally, we compare with ACTIR jiang2022invariant, which disentangles the invariant feature and bias by an independent regularization. And with SFB eastwood2024spuriosity, which aims to exploit spurious features but in a weaker setting. We briefly summarized the difference between SFB and our method in the last part of the related work. SFB can serve as a special case of our method when 1) the environment doesn’t affect the label, and 2) don’t bias aware for prediction. For Office-Home dataset, since SFB doesn’t have results on it, we add MMD and GMDG as more baselines to show our performance.

### 4.1 Simulation experiments

##### Data simulation

To examine whether BAG can effectively identify and leverage domain shift information for better generalization, we devised a series of simulation experiments. Our data generation process strictly adheres to the causal structure depicted in [Figure 2](https://arxiv.org/html/2507.17001v1#S2.F2). First, we sample an environmental variable $e$ following a categorical distribution, $e\sim\mathrm{Categorical}(\{\pi_{1},\dots,\pi_{M}\})$, which subsequently induces variations across different domains.

Then, a binary label $y$ is generated in an environment-dependent manner:
$y=\mathbf{1}\{\,w^{\mathsf{T}}E+b_{0}>0\}.$
Next, we assign stable content $c$ based on the value of $y$: if $y=0$, then $c=c_{0}+\epsilon$; if $y=1$, then $c=c_{1}+\epsilon$. Here, $\epsilon\sim\mathcal{N}(0,\sigma_{c}^{2}I)$. Simultaneously, the bias variable $b$ is determined as:
$b=E+\mathbf{C}_{e,y},$
ensuring that the background shift is influenced jointly by the environment and the label. This setup emulates the common real-world scenario where false associations often arise. Finally, stable content and bias jointly generate the observed data:
$x=M\begin{bmatrix}c\\
b\end{bmatrix}+\eta,$
where $M$ is an invertible matrix, and $\eta\sim\mathcal{N}(0,\sigma_{x}^{2}I)$ represents measurement noise.

##### Results and Discussions

This study evaluates several methods on a synthetic dataset of 5000 samples. To ensure result reliability, each method was repeated five times, and the average performance was recorded. As shown in the right side of [Table 1](https://arxiv.org/html/2507.17001v1#S4.T1), both ERM and IRM performed well on the training set but achieved test accuracies just above 50%. We attribute this to the significant distributional disparity of environmental information $e$ across domains in the dataset, which limits the effectiveness of the class label $y$ in supporting accurate predictions. In other words, an excessive focus on invariant information may hinder model optimization, leading to marginal performance gains. In contrast, our method leverages domain-specific information (${\mathbf{b}}$) more effectively, significantly enhancing prediction performance by reducing noise compared to the class label $y$.

**Table 1: The comparison of the PACS dataset and the Synthetic dataset. All BAG results on PACS are obtained by averaging over 3 seeds. Baseline results are taken from eastwood2024spuriosity**
| Algorithm | PACS | Synthetic |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| P | A | C | S | Avg | Acc |  |
| ERM | 93.0$\pm$0.7 | 79.3$\pm$0.5 | 74.3$\pm$0.7 | 65.4$\pm$1.5 | 78.0 | 54.62$\pm$3.70 |
| ERM + PL | 93.7$\pm$0.4 | 79.6$\pm$1.5 | 74.1$\pm$1.2 | 63.1$\pm$3.1 | 77.6 | – |
| IRM (irm,) | 93.3$\pm$0.3 | 78.7$\pm$0.7 | 75.4$\pm$1.5 | 65.6$\pm$2.5 | 78.3 | 54.73$\pm$2.72 |
| IRM + PL | 94.1$\pm$0.7 | 78.9$\pm$2.9 | 75.1$\pm$4.6 | 62.9$\pm$4.9 | 77.8 | – |
| ACTIR (jiang2022invariant,) | 94.8$\pm$0.1 | 82.5$\pm$0.4 | 76.6$\pm$0.6 | 62.1$\pm$1.3 | 79.0 | 70.84$\pm$6.37 |
| SFB eastwood2024spuriosity | 95.8$\pm$0.6 | 80.4$\pm$1.3 | 76.6$\pm$0.6 | 71.8$\pm$2.0 | 81.2 | – |
| BAG | 96.6$\pm$0.8 | 86.0$\pm$0.4 | 77.9$\pm$1.0 | 73.0$\pm$0.4 | 83.4 | 97.48$\pm$1.82 |

### 4.2 Real-world experiments

##### Experimental Settings

We test our approach in commonly used PACS li2017deeper and Office-Homeoffice datasets. In the PACS dataset, we let one domain become the target and other domains as source domains. We use ResNet18 for PACS as the backbone and ResNet50 for Office-Home with MLP-based VAEs as the backbone and classifiers. We show our results over 3 random seeds. Our framework employs a single-layer and a two-layer linear layer as the encoder and decoder, respectively, and utilizes a single-layer linear layer for each specific classifier. We define three learnable parameters for the environment embedding. The experiment details can be found in Appendix F.

Table: Table 2: Results on Office-Home with ResNet-50. BAG results are averaged over 3 seeds.

##### Results and Discussions

For the PACS data, in the left side of [Table 1](https://arxiv.org/html/2507.17001v1#S4.T1), we can find that our model is far superior to other models, with an average improvement of $2.2\%$ compared with best baseline, SFB method. Our model shows consistent improvement across all target domains. Among all target domains, we achieved better improvement on the Art Painting test, which may be because of the related bias in the domains.
According to the experimental results of the Office-Home dataset on the [Table 2](https://arxiv.org/html/2507.17001v1#S4.T2), our BAG model outperforms all other baselines on three transfer tasks. It should be noted that we achieve much better performance in Clipart, which is the harder task; it shows BAG can achieve better generalization even when the domain gap is larger than the baselines. Considering that BAG involves two steps of optimization, we demonstrate the effects of each individual optimization strategy, namely BAG-BA, and BAG-TTA. Additionally, we explore the learning of latent variables by designing a model variant, BAG-VAE, which is a version of the model without the VAE component. We show that all the models performance better compared with baselines both in PACS and Synthetic data, but worse than BAG. The details of ablation study can be found in Appendix F.

## 5 Conclusion

In this paper, we challenge the conventional methods for OOD generation that treat bias solely as an obstacle. Instead, we theoretically demonstrate that bias can be strategically leveraged to enhance prediction, particularly when it retains useful dependencies on the target domain labels. Our framework demonstrates that the bias can be utilized in both direct and indirect ways. Through the empirical validation on synthetic and real-world benchmarks, we establish the effectiveness of this approach.
The results consistently show that leveraging bias in a structured manner leads to more robust and adaptable models, offering a new perspective on OOD generation.
Limitation: This work primarily aims to question the conventional treatment of bias and provide a theoretical foundation for understanding its effects. The experiments presented in this paper are intended to validate our proposed method, which don’t include large-scale experiments on extensive datasets. We acknowledge this limitation and leave systematic benchmarking on large datasets as an important direction for future work.

## References

- (1)
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz.
Invariant risk minimization.
arXiv preprint arXiv:1907.02893, 2019.
- (2)
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario March, and Victor Lempitsky.
Domain-adversarial training of neural networks.
Journal of machine learning research, 17(59):1–35, 2016.
- (3)
Dong-Hyun Lee et al.
Pseudo-label: The simple and efficient semi-supervised learning method for deep neural networks.
In Workshop on challenges in representation learning, ICML, volume 3, page 896. Atlanta, 2013.
- (4)
Jiaxin Qi, Yulei Niu, Jianqiang Huang, and Hanwang Zhang.
Two causal principles for improving visual dialog.
In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10860–10869, 2020.
- (5)
Woong-Gi Chang, Tackgeun You, Seonguk Seo, Suha Kwak, and Bohyung Han.
Domain-specific batch normalization for unsupervised domain adaptation.
In Proceedings of the IEEE/CVF conference on Computer Vision and Pattern Recognition, pages 7354–7362, 2019.
- (6)
Pengfei Wei, Yiping Ke, and Chi Keong Goh.
A general domain specific feature transfer framework for hybrid domain adaptation.
IEEE Transactions on Knowledge and Data Engineering, 31(8):1440–1451, 2018.
- (7)
Lingjing Kong, Shaoan Xie, Weiran Yao, Yujia Zheng, Guangyi Chen, Petar Stojanov, Victor Akinwande, and Kun Zhang.
Partial identifiability for domain adaptation.
arXiv preprint arXiv:2306.06510, 2023.
- (8)
Vladimir Vapnik.
Principles of risk minimization for learning theory.
Advances in neural information processing systems, 4, 1991.
- (9)
Cian Eastwood, Shashank Singh, Andrei L Nicolicioiu, Marin Vlastelica Pogančić, Julius von Kügelgen, and Bernhard Schölkopf.
Spuriosity didn’t kill the classifier: Using invariant predictions to harness spurious features.
Advances in Neural Information Processing Systems, 36, 2024.
- (10)
Ilsang Ohn and Lizhen Lin.
Optimal bayesian estimation of gaussian mixtures with growing number of components.
Bernoulli, 29(2):1195–1218, 2023.
- (11)
Yibo Jiang and Victor Veitch.
Invariant and transportable representations for anti-causal domain shifts.
Advances in Neural Information Processing Systems, 35:20782–20794, 2022.
- (12)
Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales.
Deeper, broader and artier domain generalization.
In Proceedings of the IEEE international conference on computer vision, pages 5542–5550, 2017.
- (13)
Hemanth Venkateswara, Jose Eusebio, Shayok Chakraborty, and Sethuraman Panchanathan.
Deep hashing network for unsupervised domain adaptation, 2017.
- (14)
Haoliang Li, Sinno Jialin Pan, Shiqi Wang, and Alex C Kot.
Domain generalization with adversarial feature learning.
In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 5400–5409, 2018.
- (15)
Zhaorui Tan, Xi Yang, and Kaizhu Huang.
Rethinking multi-domain generalization with a general learning objective.
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 23512–23522, 2024.
- (16)
Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar.
Invariant risk minimization games.
In International Conference on Machine Learning, pages 145–155. PMLR, 2020.
- (17)
Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski.
The risks of invariant risk minimization.
arXiv preprint arXiv:2010.05761, 2020.
- (18)
Kartik Ahuja, Ethan Caballero, Dinghuai Zhang, Jean-Christophe Gagnon-Audet, Yoshua Bengio, Ioannis Mitliagkas, and Irina Rish.
Invariance principle meets information bottleneck for out-of-distribution generalization.
Advances in Neural Information Processing Systems, 34:3438–3450, 2021.
- (19)
Pritish Kamath, Akilesh Tangella, Danica Sutherland, and Nathan Srebro.
Does invariant risk minimization capture invariance?
In International Conference on Artificial Intelligence and Statistics, pages 4069–4077. PMLR, 2021.
- (20)
David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville.
Out-of-distribution generalization via risk extrapolation (rex).
In International conference on machine learning, pages 5815–5826. PMLR, 2021.
- (21)
Ya Li, Xinmei Tian, Mingming Gong, Yajing Liu, Tongliang Liu, Kun Zhang, and Dacheng Tao.
Deep domain generalization via conditional invariant adversarial networks.
In Proceedings of the European conference on computer vision (ECCV), pages 624–639, 2018.
- (22)
Rui Shao, Xiangyuan Lan, Jiawei Li, and Pong C Yuen.
Multi-adversarial discriminative deep domain generalization for face presentation attack detection.
In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10023–10031, 2019.
- (23)
Shanshan Zhao, Mingming Gong, Tongliang Liu, Huan Fu, and Dacheng Tao.
Domain generalization via entropy regularization.
Advances in neural information processing systems, 33:16096–16107, 2020.
- (24)
Zhun Deng, Frances Ding, Cynthia Dwork, Rachel Hong, Giovanni Parmigiani, Prasad Patil, and Pragya Sur.
Representation via representations: Domain generalization via adversarially learned invariant representations.
arXiv preprint arXiv:2006.11478, 2020.
- (25)
Toshihiko Matsuura and Tatsuya Harada.
Domain generalization using a mixture of multiple latent domains.
In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 11749–11756, 2020.
- (26)
Maximilian Ilse, Jakub M Tomczak, Christos Louizos, and Max Welling.
Diva: Domain invariant variational autoencoders.
In Medical Imaging with Deep Learning, pages 322–348. PMLR, 2020.
- (27)
Guoqing Wang, Hu Han, Shiguang Shan, and Xilin Chen.
Cross-domain face presentation attack detection via multi-domain disentangled representation learning.
In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 6678–6687, 2020.
- (28)
Tan Wang, Jianqiang Huang, Hanwang Zhang, and Qianru Sun.
Visual commonsense r-cnn.
In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10760–10770, 2020.
- (29)
Riccardo Volpi and Vittorio Murino.
Addressing model vulnerability to distributional shifts over image transformation sets.
In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 7980–7989, 2019.
- (30)
Yichun Shi, Xiang Yu, Kihyuk Sohn, Manmohan Chandraker, and Anil K Jain.
Towards universal representation learning for deep face recognition.
In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 6817–6826, 2020.
- (31)
Jiaxing Huang, Dayan Guan, Aoran Xiao, and Shijian Lu.
Fsdr: Frequency space domain randomization for domain generalization.
In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 6891–6902, 2021.
- (32)
Pan Li, Da Li, Wei Li, Shaogang Gong, Yanwei Fu, and Timothy M Hospedales.
A simple feature augmentation for domain generalization.
In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 8886–8895, 2021.
- (33)
Massimiliano Mancini, Zeynep Akata, Elisa Ricci, and Barbara Caputo.
Towards recognizing unseen categories in unseen domains.
In European Conference on Computer Vision, pages 466–483. Springer, 2020.
- (34)
Yang Shu, Zhangjie Cao, Chenyu Wang, Jianmin Wang, and Mingsheng Long.
Open domain generalization with domain-augmented meta-learning.
In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 9624–9633, 2021.
- (35)
Kaiyang Zhou, Yongxin Yang, Yu Qiao, and Tao Xiang.
Mixstyle neural networks for domain generalization and adaptation.
International Journal of Computer Vision, 132(3):822–836, 2024.
- (36)
Seonguk Seo, Yumin Suh, Dongwan Kim, Geeho Kim, Jongwoo Han, and Bohyung Han.
Learning to optimize domain specific normalization for domain generalization.
In Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part XXII 16, pages 68–83. Springer, 2020.
- (37)
Mattia Segu, Alessio Tonioni, and Federico Tombari.
Batch normalization embeddings for deep domain generalization.
Pattern Recognition, 135:109115, 2023.
- (38)
Artem Rozantsev, Mathieu Salzmann, and Pascal Fua.
Beyond sharing weights for deep domain adaptation.
IEEE transactions on pattern analysis and machine intelligence, 41(4):801–814, 2018.
- (39)
Baochen Sun and Kate Saenko.
Deep coral: Correlation alignment for deep domain adaptation.
In Computer Vision–ECCV 2016 Workshops: Amsterdam, The Netherlands, October 8-10 and 15-16, 2016, Proceedings, Part III 14, pages 443–450. Springer, 2016.
- (40)
Guoliang Kang, Lu Jiang, Yi Yang, and Alexander G Hauptmann.
Contrastive adaptation network for unsupervised domain adaptation.
In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 4893–4902, 2019.
- (41)
Xiaofeng Liu, Yuzhuo Han, Song Bai, Yi Ge, Tianxing Wang, Xu Han, Site Li, Jane You, and Jun Lu.
Importance-aware semantic segmentation in self-driving with discrete wasserstein training.
In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 11629–11636, 2020.
- (42)
Jing Jiang and ChengXiang Zhai.
Instance weighting for domain adaptation in nlp.
ACL, 2007.
- (43)
Xuemiao Xu, Hai He, Huaidong Zhang, Yangyang Xu, and Shengfeng He.
Unsupervised domain adaptation via importance sampling.
IEEE Transactions on Circuits and Systems for Video Technology, 30(12):4688–4699, 2019.
- (44)
Tongtong Fang, Nan Lu, Gang Niu, and Masashi Sugiyama.
Rethinking importance weighting for deep learning under distribution shift.
Advances in neural information processing systems, 33:11996–12007, 2020.
- (45)
Eric Tzeng, Judy Hoffman, Kate Saenko, and Trevor Darrell.
Adversarial discriminative domain adaptation.
In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 7167–7176, 2017.
- (46)
Jian Shen, Yanru Qu, Weinan Zhang, and Yong Yu.
Wasserstein distance guided representation learning for domain adaptation.
In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
- (47)
Yang Zou, Zhiding Yu, BVK Kumar, and Jinsong Wang.
Unsupervised domain adaptation for semantic segmentation via class-balanced self-training.
In Proceedings of the European conference on computer vision (ECCV), pages 289–305, 2018.
- (48)
Kuniaki Saito, Donghyun Kim, Stan Sclaroff, and Kate Saenko.
Universal domain adaptation through self supervision.
Advances in neural information processing systems, 33:16282–16292, 2020.
- (49)
Guangyi Chen, Yuhao Lu, Jiwen Lu, and Jie Zhou.
Deep credible metric learning for unsupervised domain adaptation person re-identification.
In Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part VIII 16, pages 643–659. Springer, 2020.
- (50)
Ananya Kumar, Tengyu Ma, and Percy Liang.
Understanding self-training for gradual domain adaptation.
In International conference on machine learning, pages 5468–5479. PMLR, 2020.
- (51)
Hong Liu, Jianmin Wang, and Mingsheng Long.
Cycle self-training for domain adaptation.
Advances in Neural Information Processing Systems, 34:22968–22981, 2021.
- (52)
Yanghao Li, Naiyan Wang, Jianping Shi, Xiaodi Hou, and Jiaying Liu.
Adaptive batch normalization for practical domain adaptation.
Pattern Recognition, 80:109–117, 2018.
- (53)
Xiaofeng Liu, Fangxu Xing, Chao Yang, Georges El Fakhri, and Jonghye Woo.
Adapting off-the-shelf source segmenter for target medical image segmentation.
In Medical Image Computing and Computer Assisted Intervention–MICCAI 2021: 24th International Conference, Strasbourg, France, September 27–October 1, 2021, Proceedings, Part II 24, pages 549–559. Springer, 2021.
- (54)
Lingjing Kong, Shaoan Xie, Weiran Yao, Yujia Zheng, Guangyi Chen, Petar Stojanov, Victor Akinwande, and Kun Zhang.
Partial disentanglement for domain adaptation.
In International conference on machine learning, pages 11455–11472. PMLR, 2022.
- (55)
Zijian Li, Ruichu Cai, Guangyi Chen, Boyang Sun, Zhifeng Hao, and Kun Zhang.
Subspace identification for multi-source domain adaptation.
Advances in Neural Information Processing Systems, 36, 2024.
- (56)
Yu Sun, Xiaolong Wang, Zhuang Liu, John Miller, Alexei Efros, and Moritz Hardt.
Test-time training with self-supervision for generalization under distribution shifts.
In International Conference on Machine Learning, pages 9229–9248. PMLR, 2020.
- (57)
Yuejiang Liu, Parth Kothari, Bastien van Delft, Baptiste Bellot-Gurlet, Taylor Mordan, and Alexandre Alahi.
Ttt++: When does self-supervised test-time training fail or thrive?
Advances in Neural Information Processing Systems, 34, 2021.
- (58)
Muhammad Jehanzeb Mirza, Pol Jané Soneira, Wei Lin, Mateusz Kozinski, Horst Possegger, and Horst Bischof.
Actmad: Activation matching to align distributions for test-time-training.
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 24152–24161, 2023.
- (59)
Yushu Li, Xun Xu, Yongyi Su, and Kui Jia.
On the robustness of open-world test-time training: Self-training with dynamic prototype expansion.
In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 11836–11846, 2023.
- (60)
Dian Chen, Dequan Wang, Trevor Darrell, and Sayna Ebrahimi.
Contrastive test-time adaptation.
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 295–305, 2022.
- (61)
Devavrat Tomar, Guillaume Vray, Behzad Bozorgtabar, and Jean-Philippe Thiran.
Tesla: Test-time self-learning with automatic adversarial augmentation.
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 20341–20350, 2023.
- (62)
Yongyi Su, Xun Xu, and Kui Jia.
Revisiting realistic test-time training: Sequential inference and adaptation by anchored clustering.
Advances in Neural Information Processing Systems, 35:17543–17555, 2022.
- (63)
Sachin Goyal, Mingjie Sun, Aditi Raghunathan, and J Zico Kolter.
Test time adaptation via conjugate pseudo-labels.
Advances in Neural Information Processing Systems, 35:6204–6218, 2022.
- (64)
Dequan Wang, Evan Shelhamer, Shaoteng Liu, Bruno Olshausen, and Trevor Darrell.
Tent: Fully test-time adaptation by entropy minimization.
arXiv preprint arXiv:2006.10726, 2020.
- (65)
Paul Pu Liang, Terrance Liu, Liu Ziyin, Nicholas B Allen, Randy P Auerbach, David Brent, Ruslan Salakhutdinov, and Louis-Philippe Morency.
Think locally, act globally: Federated learning with local and global representations.
arXiv preprint arXiv:2001.01523, 2020.

## Appendix A Related Work

Domain generalization. Domain Generalization (DG) aims to develop models that generalize to unseen OOD target domains using only observed source domains for training. A common approach is learning invariant representations across all source domains, as seen in Invariant Risk Minimization (IRM) [1] and its extensions [16, 17, 18, 19, 20]. Other methods leverage adversarial learning [21, 22, 23, 24, 25] to extract domain-invariant features or use disentangled representation learning [26, 27, 11] to separate domain-invariant from domain-specific information. Causal inference techniques [28, 4], such as back-door and front-door rules, help mitigate bias. Furthermore techniques such as data augmentation [29, 30, 31, 32], feature augmentation [33, 34, 35], and network regularization [36, 37] can reduce overfitting and improve generalization.

Unsupervised domain adaptation Unlike DG, unsupervised domain adaptation utilizes unlabeled target data to guide model adaptation. Due to the observation of target distribution, many methods focus on distribution alignment by minimizing domain divergence [38, 39, 40, 41], using importance reweighting [42, 43, 44], and learning with adversarial discriminators, such as DANN [2], ADDA [45], and WDDA [46]. Beyond distribution alignment, some methods focus on invariant information-based adaptation, including self-training [3, 47, 48, 49, 50, 51], which uses pseudo-labels from a source-trained model to guide target domain learning, and batch normalization-based regularization [52, 53], which adapts only BN layers while keeping other parameters invariant. Recently, disentangled representation learning [54, 55] has shown strong potential for adaptation by separating domain-invariant content from domain-specific bias.

Test-Time domain adaptation In real-world scenarios, the target domain distribution is often unknown in advance, making it challenging to train an adaptive model beforehand, which motivates test time adaptation. The typical solution is Test-Time Training (TTT) [56, 57, 58, 59], which adapts models during inference by optimizing the same self-supervised learning objective used during pre-training. Besides, even without access to the target distribution, adaptation remains feasible with invariant information learned from the source domain, such as through self-training[60, 61, 62, 63] or network regularization [64, 65].
The most closely related work to ours is Stable Feature Boosting (SFB), which leverages unstable features alongside invariant representations to generate pseudo-labels when only the image is affected by environmental factors. Our approach differs in three key ways. First, we delve deeper into latent variable modeling to provide a more comprehensive explanation of bias. Second, our model accommodates more complex bias scenarios, allowing the environment variable to influence both the image and the label. Third, we integrate invariant correction and expert reweighting to effectively utilize biased data.

## Appendix B Proof of identification and better prediction

### B.1 Proof of the identification of latent variables

###### Lemma B.1 .

(Block-wise Identification of ${\mathbf{c}}$ and ${\mathbf{b}}$ [7].) Assuming that the data generation process follows Figure 2, and the following assumptions hold true:

- •
A1 (Smooth and Positive Density): The probability density function of latent variables is smooth and positive, i.e. $P_{\mathbf{z}|\mathbf{e}}>0$.
- •
A2 (Conditional Independence): Conditioned on $\mathbf{e}$, each $z_{i}$ is independent of any other $z_{j}$ for $i.j\in\{1,\cdots,n\},i\neq j$, i.e., $\log P_{\mathbf{z}|\mathbf{e},\mathbf{y}}=\sum_{i=1}^{n}q(z_{i},\mathbf{e},
\mathbf{y})$, where $q_{i}(z_{i},\mathbf{e},\mathbf{y})$ is the log density of the conditional distribution, i.e., $q_{i}:\log P_{z_{i}|\mathbf{e},\mathbf{y}}$.
- •
A3 (Linear Independence): For any $\mathbf{b}\in\mathcal{B}\subseteq\mathbb{R}^{n_{b}}$, there exists $n_{b}+1$ values of $\mathbf{e}$, i.e., $e_{j}$ with $j=0,1,.\cdots,n_{b}$ such that these $n_{b}$ vectors $\bm{{\mathbf{w}}}({\mathbf{b}},{\mathbf{e}}_{j})-\bm{{\mathbf{w}}}({\mathbf{b}
},{\mathbf{e}}_{0})$ with $j=1,\cdots,n_{b}$ are linearly independent, where vector $\bm{{\mathbf{w}}}({\mathbf{b}},{\mathbf{e}}_{j})$ is defined as follows:
$\bm{{\mathbf{w}}}({\mathbf{b}},{\mathbf{e}})=\left(\frac{\partial q_{1}({
\mathbf{b}}_{1},{\mathbf{e}})}{\partial{\mathbf{b}}_{1}},\cdots\frac{\partial q
_{n_{b}}({\mathbf{b}}_{n_{b}},{\mathbf{e}})}{\partial{\mathbf{b}}_{n_{b}}}
\right),$
(8)
- •
A4 (Domain Variability:) There exist two values of $\mathbf{e}$, i.e., $\mathbf{e}_{i}$ and $\mathbf{e}_{j}$, s.t., for any set $A_{\mathbf{z}}\subseteq\mathcal{Z}$ with non-zero probability
measure and cannot be expressed as $\Omega_{\mathbf{c}}\times\mathcal{B}$ for any $\Omega_{\mathbf{c}}\subset\mathcal{C}$, we have
$\int_{\mathbf{z}\in A_{\mathbf{z}}}p(\mathbf{z}|\mathbf{e}_{i})d\mathbf{z}\neq
\int_{\mathbf{z}\in A_{\mathbf{z}}}p(\mathbf{z}|\mathbf{e}_{j})d\mathbf{z}.$
(9)

Then the learned $\hat{\mathbf{c}}$ and $\hat{\mathbf{b}}$ are block-wise identifiable.

###### Proof.

We begin with the matched marginal distribution $p_{\mathbf{x}|\mathbf{e},\mathbf{y}}$ to bridge the relation between $\mathbf{z}$ and $\hat{\mathbf{z}}$. Suppose that $\hat{g}:\mathbf{z}\rightarrow\mathcal{X}$ is an invertible estimated generating function, we have Equation ([10](https://arxiv.org/html/2507.17001v1#A2.E10)).

$$ $\forall\mathbf{e}\in\mathcal{(}E),\quad p_{\hat{\mathbf{x}}|\mathbf{e},\mathbf {y}}=p_{\mathbf{x}|\mathbf{e},\mathbf{y}}\Longleftrightarrow p_{\hat{g}(\hat{ \mathbf{z}})|\mathbf{e},\mathbf{y}}=p_{g(\mathbf{z})|\mathbf{e},\mathbf{y}}.$ (10) $$

Sequentially, by using the change of variables formula, we can further obtain Equation ([11](https://arxiv.org/html/2507.17001v1#A2.E11))

$$ $p_{\hat{g}(\hat{\mathbf{z}}|\mathbf{e},\mathbf{y})}=p_{g(\mathbf{z}|\mathbf{e} ,\mathbf{y})}\Longleftrightarrow p_{g^{-1}\circ g(\hat{\mathbf{z}})|\mathbf{e} }|\mathbf{J}_{g^{-1}}|=p_{\mathbf{z}|\mathbf{e},\mathbf{y}}|\mathbf{J}_{g^{-1} }|\Longleftrightarrow p_{h(\hat{\mathbf{z}})|\mathbf{e},\mathbf{y}}=p_{\mathbf {z}|\mathbf{e},\mathbf{y}},$ (11) $$

where $h:=g^{-1}\circ g$ is the transformation between the ground-true and the estimated latent variables, respectively. $\mathbf{J}_{g^{-1}}$ denotes the absolute value of Jacobian matrix determinant of $g^{-1}$. Since we assume that $g$ and $\hat{g}$ are invertible, $|\mathbf{J}_{g^{-1}}|\neq 0$ and $h$ is also invertible.

According to A2 (conditional independent assumption), we can have Equation ([12](https://arxiv.org/html/2507.17001v1#A2.E12)).

$$ $p_{\mathbf{z}|\mathbf{e},\mathbf{y}}(\mathbf{z}|\mathbf{e},\mathbf{y})=\prod_{ i=1}^{n}p_{z_{i}|\mathbf{e},\mathbf{y}}(z_{i}|\mathbf{e},\mathbf{y});\quad p_{ \hat{\mathbf{z}}|\mathbf{e},\mathbf{y}}(\hat{\mathbf{z}}|\mathbf{e},\mathbf{y} )=\prod_{i=1}^{n}p_{\hat{z}_{i}|\mathbf{e},\mathbf{y}}(\hat{z}_{i}|\mathbf{e}, \mathbf{y}).$ (12) $$

For convenience, we take logarithm on both sides of Equation ([12](https://arxiv.org/html/2507.17001v1#A2.E12)) and further let $q_{i}:=\log p_{z_{i}|\mathbf{e},\mathbf{y}},\hat{q}_{i}:=\log p_{\hat{z}_{i}|
\mathbf{e},\mathbf{y}}$. Hence we have:

$$ $\log p_{\mathbf{z}|\mathbf{e},\mathbf{y}}(\mathbf{z}|\mathbf{e},\mathbf{y})= \sum_{i=1}^{n}q_{i}(z_{i},\mathbf{e},\mathbf{y});\quad\log p_{\hat{\mathbf{z}, \mathbf{y}}|\mathbf{e}}=\sum_{i=1}^{n}\hat{q}_{i}(\hat{z}_{i},\mathbf{e}, \mathbf{y}).$ (13) $$

By combining Equation ([13](https://arxiv.org/html/2507.17001v1#A2.E13)) and Equation ([11](https://arxiv.org/html/2507.17001v1#A2.E11)), we have:

$$ $p_{\mathbf{z}|\mathbf{e},\mathbf{y}}=p_{h(\hat{\mathbf{z}}|\mathbf{e},\mathbf{ y})}\Longleftrightarrow p_{\hat{\mathbf{z}}|\mathbf{e},\mathbf{y}}=p_{\mathbf{ z}|\mathbf{e},\mathbf{y}}|\mathbf{J}_{h^{-1}}|\Longleftrightarrow\sum_{i=1}^{n }q_{i}({z_{i},\mathbf{e}},\mathbf{y})+\log|\mathbf{J}_{h^{-1}}|=\sum_{i=1}^{n} \hat{q}_{i}(\hat{z}_{i},\mathbf{e},\mathbf{y}),$ (14) $$

where $\mathbf{J}_{h^{-1}}$ are the Jacobian matrix of $h^{-1}$.

Sequentially, we take the first-order derivative with $\hat{c}_{j}$ on Equation ([14](https://arxiv.org/html/2507.17001v1#A2.E14)), where $j\in\{1,\cdots,n_{c}\}$, and have

$$ $\sum_{i=1}^{n}\frac{\partial q_{i}(z_{i},\mathbf{u},\mathbf{y})}{\partial z_{i }}\cdot\frac{\partial z_{i}}{\partial\hat{c}_{j}}+\frac{\partial\log|\mathbf{J }_{h^{-1}}|}{\partial\hat{c}_{j}}=\frac{\partial q_{j}(\hat{c}_{j},\mathbf{u}, \mathbf{y})}{\partial\hat{c}_{j}}.$ (15) $$

Suppose $\mathbf{e}=e_{0},e_{1},\cdots,e_{n_{b}}$, we subtract the Equation ([15](https://arxiv.org/html/2507.17001v1#A2.E15)) corresponding to $e_{k}$ with that corresponds to $e_{0}$, and we have:

$$ $\sum_{i=1}^{n}\left(\frac{\partial q_{i}(z_{i},e_{k},\mathbf{y})}{\partial z_{ i}}-\frac{\partial q_{i}(z_{i},e_{0},\mathbf{y})}{\partial z_{i}}\right)\cdot \frac{\partial z_{i}}{\partial\hat{c}_{j}}=\frac{\partial\hat{q}_{j}(\hat{c}_{ j},e_{k},\mathbf{y})}{\partial\hat{c}_{j}}-\frac{\partial\hat{q}_{j}(\hat{c}_{ j},e_{0},\mathbf{y})}{\partial\hat{c}_{j}}.$ (16) $$

Since the distribution of estimated $\hat{\mathbf{c}}_{j}$ does not change across different domains, $\frac{\partial\hat{q}_{j}(\hat{c}_{j},e_{k},\mathbf{y})}{\partial\hat{c}_{j}}-
\frac{\partial\hat{q}_{j}(\hat{c}_{j},e_{0},\mathbf{y})}{\partial\hat{c}_{j}}=0$. Since $\frac{\partial q_{i}(c_{i},e_{k},\mathbf{y})}{\partial z_{i}}$ does not change across different domains, $\frac{\partial q_{i}(c_{i},e_{k},\mathbf{y})}{\partial c_{i}}=\frac{\partial q
_{i}(c_{i},e_{0},\mathbf{y})}{\partial c_{i}}$ for $i\in\{1,\cdots,n_{c}\}$. So we have

$$ $\sum_{i=1}^{n_{b}}\left(\frac{\partial q_{i}(b_{i},e_{k},\mathbf{y})}{\partial b _{i}}-\frac{\partial q_{i}(b_{i},e_{0},\mathbf{y})}{\partial b_{i}}\right) \cdot\frac{\partial b_{i}}{\partial\hat{c}_{j}}=0.$ (17) $$

Based on the linear independence assumption (A3), the linear system is a full-rank system. Therefore, the only solution is $\frac{\partial b_{i}}{\partial\hat{c}_{j}}=0$ for $i\in\{1,\cdots,n_{b}\}$ and $j\in\{1,\cdots,n_{c}\}$.

Since $h(\cdot)$ is smooth over $\mathbf{z}$, its Jacobian can be formalized as follows

$$ $\begin{gathered}\mathbf{J}_{h}=\begin{bmatrix}\begin{array}[]{c|c}\textbf{A}:= \frac{\partial\mathbf{b}}{\partial\hat{\mathbf{b}}}&\textbf{B}:=\frac{\partial \mathbf{b}}{\partial\hat{\mathbf{c}}}\\ \hline\cr\textbf{C}:=\frac{\partial\mathbf{c}}{\partial\hat{\mathbf{b}}}& \textbf{D}:=\frac{\partial\mathbf{c}}{\partial\hat{\mathbf{c}}}.\end{array} \end{bmatrix}\end{gathered}$ (18) $$

Note that $\frac{\partial b_{i}}{\partial\hat{c}_{j}}=0$ for $i\in\{1,\cdots,n_{b}\}$ and $j\in\{1,\cdots,n_{c}\}$ means that $\textbf{B}=0$. Since $h(\cdot)$ is invertible, $\mathbf{J}_{h}$ is a full-rank matrix. Therefore, for each $b_{i},i\in\{1,\cdots,n_{b}\}$, there exists a $h_{i}$ such that $b_{i}=h_{i}(\hat{\mathbf{b}})$.

Besides, based on A4, one can show that all entries in the submatrix C zero according to part of the proof of Theorem 4.2 in [7](Steps 1, 2, and 3). Therefore, $\mathbf{b}$ and $\mathbf{c}$ are block-wise identifiable.
∎

###### Lemma B.1 .

###### Proof.

### B.2 Better predictor with adding bias

###### Proof.

Under Assumptions A1–A5, we compare the Bayes-optimal risks for predictors based on $c$ versus $(c,b)$.

###### Proof.

##### Log-loss.

The Bayes-optimal log-loss risks satisfy

$$ $R_{\log}\bigl{(}q^{*}(\cdot\mid c)\bigr{)}=H(Y\mid C),\qquad R_{\log}\bigl{(}q ^{*}(\cdot\mid c,b)\bigr{)}=H(Y\mid C,B),$ $$

where

$$ $H(Y\mid C)=-\mathbb{E}_{C}\bigl{[}\mathbb{E}_{Y\mid C}[\log p(Y\mid C)]\bigr{] },\quad H(Y\mid C,B)=-\mathbb{E}_{C,B}\bigl{[}\mathbb{E}_{Y\mid C,B}[\log p(Y \mid C,B)]\bigr{]}.$ $$

By A5 there is a positive-measure set of $(c,b)$ on which $p(y\mid c,b)\neq p(y\mid c)$. Hence

$$ $I(Y;B\mid C)=H(Y\mid C)-H(Y\mid C,B)=\mathbb{E}_{C,B}\!\bigl{[}D_{\mathrm{KL}} (p(y\mid c,b)\,\|\,p(y\mid c))\bigr{]}>0,$ $$

so $H(Y\mid C,B)<H(Y\mid C)$ and therefore

$$ $R_{\log}\bigl{(}q^{*}(\cdot\mid c,b)\bigr{)}<R_{\log}\bigl{(}q^{*}(\cdot\mid c )\bigr{)}.$ $$

##### Squared-error loss.

The Bayes-optimal squared-error risks satisfy

$$ $R_{\mathrm{sq}}\bigl{(}f^{*}(c)\bigr{)}=\mathbb{E}\bigl{[}\mathrm{Var}(Y\mid C )\bigr{]},\qquad R_{\mathrm{sq}}\bigl{(}f^{*}(c,b)\bigr{)}=\mathbb{E}\bigl{[} \mathrm{Var}(Y\mid C,B)\bigr{]}.$ $$

By the law of total variance,

$$ $\mathrm{Var}(Y\mid C)=\mathbb{E}_{B}\bigl{[}\mathrm{Var}(Y\mid C,B)\bigr{]}\;+ \;\mathrm{Var}_{B}\!\bigl{(}\mathbb{E}[Y\mid C,B]\bigm{|}C\bigr{)}.$ $$

Since A5 ensures that $B$ influences the conditional mean of $Y$ given $C$, the second term is strictly positive on a set of positive probability. Thus

$$ $\mathbb{E}\bigl{[}\mathrm{Var}(Y\mid C,B)\bigr{]}<\mathbb{E}\bigl{[}\mathrm{ Var}(Y\mid C)\bigr{]},$ $$

and consequently

$$ $R_{\mathrm{sq}}\bigl{(}f^{*}(c,b)\bigr{)}<R_{\mathrm{sq}}\bigl{(}f^{*}(c)\bigr {)}.$ $$

This completes the proof that incorporating $b$ strictly reduces both log-loss and squared-error Bayes risks.
∎

## Appendix C Multi-class Method

Suppose we have $K\geq 2$ classes. We "one-hot encode" these classes, so that $y$ takes values in the set
$\mathcal{Y}=\{(1,0,...,0),(0,1,0,...,0),...,(0,...,0,1)\}\subseteq\{0,1\}^{K}.$
In the training stage,
let $\hat{y}=f_{\mathbf{c}}(\mathbf{c})$ $\epsilon\in[0,1]^{\mathcal{Y}\times\mathcal{Y}}$
with

$$ $\epsilon_{y_{i},y_{j}}=p[\hat{y}=y_{i}\mid y=y_{j}]$ $$

denote the class-conditional confusion matrix of the pseudo-labels. Then, we have

$$ $\displaystyle\frac{p(y=y_{i}\mid\mathbf{c},\mathbf{b})}{p(y\neq y_{i}\mid \mathbf{c},\mathbf{b})}$ (19) $\displaystyle=\frac{p(\mathbf{c},\mathbf{b},y=y_{i})}{\sum_{y_{j}\neq y_{i}}p( \mathbf{c},\mathbf{b},y=y_{j})}$ $\displaystyle=\frac{p(\mathbf{c},\mathbf{b}\mid y=y_{i})p(y=y_{i})}{\sum_{y_{j }\neq y_{i}}p(\mathbf{c},\mathbf{b}\mid y=y_{j})p(y=y_{j})}$ $\displaystyle\quad\text{consider }\mathbf{c}\perp\!\!\!\perp\mathbf{b}\mid y:$ $\displaystyle=\frac{p(\mathbf{c}\mid y=y_{i})p(\mathbf{b}\mid y=y_{i})p(y=y_{i })}{\sum_{y_{j}\neq y_{i}}p(\mathbf{c}\mid y=y_{j})p(\mathbf{b}\mid y=y_{j})p( y=y_{j})}$ $\displaystyle=\frac{p(y=y_{i}\mid\mathbf{b})p(\mathbf{c}\mid y=y_{i})}{\sum_{y _{j}\neq y_{i}}p(y=y_{j}\mid\mathbf{b})p(\mathbf{c}\mid y=y_{j})}$ $\displaystyle=\frac{p(y=y_{i}\mid\mathbf{b})\frac{p(y=y_{i}\mid\mathbf{c})}{p( y=y_{i})}}{\sum_{y_{j}\neq y_{i}}p(y=y_{j}\mid\mathbf{b})\frac{p(y=y_{j}\mid \mathbf{c})}{p(y=y_{j})}}$ $$

for $P\in R^{\mathcal{Y}}$ defined by

$$ $P_{y_{i}}=\frac{p(y=y_{i}\mid\mathbf{b})p(y=y_{i}\mid\mathbf{c})}{p(y=y_{i})} \quad\text{ for each }y_{i}\in\mathcal{Y}.$ $$

By using the logit method, we obtain:

$$ $\displaystyle\operatorname{logit}(p(y=y_{i}\mid\mathbf{c},\mathbf{b}))$ $\displaystyle=\log\left(\frac{P_{y_{i}}}{\sum_{y_{j}\neq y_{i}}P_{y_{j}}}\right)$ (20) $\displaystyle=\operatorname{logit}\left(\frac{P_{y_{i}}}{\|P\|_{1}}\right)$ $$

Applying the sigmoid function to each side, we have

$$ $p(y=y_{i}\mid\mathbf{c},\mathbf{b})=\frac{P}{\|P\|_{1}}.$ $$

Like binary classification, we only need to estimate $p(y=y_{i}\mid\mathbf{b},e)$ since we assume that $p(y=y_{i}\mid\mathbf{c})$ and $p(y=y_{i})$ can be directly used.

$$ $\displaystyle E[\hat{y}\mid\mathbf{b}]$ (21) $\displaystyle=\sum_{y_{i}\in\mathcal{Y}}E[\hat{y}\mid y=y_{i},\mathbf{b}]p[y=y _{i}\mid\mathbf{b}]$ $\displaystyle\hat{y}\textbf{ only depends on }y$ $\displaystyle=\sum_{y_{i}\in\mathcal{Y}}E[\hat{y}\mid y=y_{i}]p[y=y_{i}\mid \mathbf{b}]$ $\displaystyle=\epsilon E[y\mid\mathbf{b}]$ $$

When $\epsilon$ is non-singular, this has the unique solution $E[y\mid\mathbf{b}]=\epsilon^{-1}E[\hat{y}\mid\mathbf{b}]$,
giving a multiclass equivalent. However, it is numerically more stable to estimate $E[y\mid\mathbf{b}]$ by the least-squares solution

$$ $\operatorname{argmin}_{p\in\Delta^{\mathcal{Y}}}\left\|\epsilon p-E[\hat{y} \mid\mathbf{b}]\right\|_{2}.$ $$

## Appendix D Proof of section 2

### D.1 Decomposition and invariant

Assume by causal representation learning, we have obtained $\mathbf{c}$ (stable latent variables) and $\mathbf{b}$ (unstable latent variables). Our goal is to estimate $p(y\mid\mathbf{c},\mathbf{b})$.

$$ $\displaystyle\frac{p(y=1\mid\mathbf{c},\mathbf{b})}{p(y=0\mid\mathbf{c}, \mathbf{b})}$ (22) $\displaystyle=\frac{p(\mathbf{c},\mathbf{b},y=1)}{p(\mathbf{c},\mathbf{b},y=0)}$ $\displaystyle=\frac{p(\mathbf{b}\mid y=1)p(\mathbf{c}\mid y=1)p(y=1)}{p( \mathbf{b}\mid y=0)p(\mathbf{c}\mid y=0)p(y=0)}$ $\displaystyle=\frac{p(y=1\mid\mathbf{b})p(\mathbf{c}\mid y=1)}{p(y=0\mid \mathbf{b})p(\mathbf{c}\mid y=0)}$ $\displaystyle=\frac{p(y=1\mid\mathbf{b})\frac{p(y=1\mid\mathbf{c})}{p(y=1)}}{p (y=0\mid\mathbf{b})\frac{p(y=0\mid\mathbf{c})}{p(y=0)}}$ $$

Assume we have $e_{all}=\{e_{0},e_{1},...,e_{n}\}$, representing different domains. We assume the source domains are $e_{s}=\{e_{0},e_{1},e_{s}\}$, and target domains are $e_{t}=\{e_{s+1},...,e_{n}\}$.

$$ $\displaystyle\forall e_{i},e_{j}\in e_{all}$ (23) $\displaystyle\quad\mathbf{c}\textbf{ only depends on }y$ $\displaystyle{p(\mathbf{c}\mid y=1,e=e_{i})}=p(\mathbf{c}\mid y=1,e=e_{j})$ Use Bayes rule $\displaystyle\frac{p(y=1,e=e_{i}\mid\mathbf{c})}{p(y=1,e=e_{i})}=\frac{p(y=1,e =e_{j}\mid\mathbf{c})}{p(y=1,e=e_{j})}$ $$

$$ $p(y=1\mid\mathbf{c})=\sum_{e\in e_{s}}{\frac{p(y=1,e=e\mid\mathbf{c})}{p(y=1,e =e)}p(e\mid e_{s})}$ $$

This result allows us to estimate $p(y=1\mid\mathbf{c})$ using data from source domains and apply it in the target domain.

### D.2 Proof of theorem 2.5

In this section, we would like to prove $h_{0}+h_{1}=1$ if and only if $\mathbf{c}\perp\!\!\!\perp y$.

Forward Implication: Suppose $h_{0}+h_{1}=1$.

$$ $h_{0}+h_{1}=1\quad\text{if and only if}\quad\mathbf{c}\perp y.$ $$

Note that the entire result also holds, with almost identical proof, in the multi-environment setting of Sections 3 and 5, conditioned on a particular environment $E$.

Proof. We first prove the forward implication.

Suppose $h_{0}+h_{1}=1$.

If $\Pr[y=1]\in\{0,1\}$, then $\mathbf{c}$ and $y$ are trivially independent, because the value of $y$ is constant. In this case, the result holds trivially. So, we assume $\Pr[y=1]\in(0,1)$.

We now compute the expected value of the pseudo-label $\hat{y}$ using the Law of Total Expectation:

$$ $E[\hat{y}]=\Pr[\hat{y}=1\mid y=1]\cdot\Pr[y=1]+\Pr[\hat{y}=0\mid y=0]\cdot\Pr[ y=0].$ $$

That is,

$$ $E[\hat{y}]=h_{1}\cdot p+(1-h_{0})\cdot(1-p),$ $$

where $p=\Pr[y=1]$.

Next, we simplify this expression:

$$ $E[\hat{y}]=h_{1}\cdot p+(1-h_{0})\cdot(1-p)=(h_{0}+h_{1}-1)\cdot p+1-h_{0}.$ $$

Using the assumption $h_{0}+h_{1}=1$, we get:

$$ $E[\hat{y}]=1-h_{0}.$ $$

Thus, $E[\hat{y}]=1-h_{0}$.

Now, since the expected value of $\hat{y}$ is independent of the actual value of $y$, this suggests that $\hat{y}$ provides no information about $y$, implying that $\hat{y}$ and $y$ are independent.

Since $\hat{y}$ is a deterministic function of $\mathbf{c}$, this implies that $\mathbf{c}$ and $y$ must also be independent. Therefore, we have shown that if $h_{0}+h_{1}=1$, then $\mathbf{c}\perp y$.

To prove the reverse implication:

Suppose $\mathbf{c}\perp y$.

If $\mathbf{c}$ and $y$ are independent, then the pseudo-label $\hat{y}$ is also independent of $y$. Hence, we have:

$$ $\Pr[\hat{y}=1\mid y=1]=\Pr[\hat{y}=1\mid y=0].$ $$

This implies that:

$$ $h_{1}=1-h_{0}.$ $$

Therefore, we have:

$$ $h_{0}+h_{1}=h_{0}+(1-h_{0})=1.$ $$

Thus, we have shown that if $\mathbf{c}\perp y$, then $h_{0}+h_{1}=1$.

Putting the forward and reverse implications together, we conclude that:

$$ $h_{0}+h_{1}=1\quad\Longleftrightarrow\quad\mathbf{c}\perp y.$ $$

This completes the proof.

### D.3 The derivation of the KL optimization

We want to solve the following problem:

$$ $q^{*}(y\mid{\mathbf{b}})\;=\;\underset{q(\cdot\mid{\mathbf{b}})}{\arg\min}\; \mathbb{E}_{e\sim p(e\mid{\mathbf{b}})}\Bigl{[}D_{\mathrm{KL}}\bigl{(}p(y\mid{ \mathbf{b}},e)\,\|\,q(y\mid{\mathbf{b}})\bigr{)}\Bigr{]}.$ $$

Let

$$ $\mathcal{L}[q(\cdot\mid{\mathbf{b}})]\;=\;\mathbb{E}_{e\sim p(e\mid{\mathbf{b} })}\Bigl{[}D_{\mathrm{KL}}\bigl{(}p(y\mid{\mathbf{b}},E)\,\|\,q(y\mid{\mathbf{ b}})\bigr{)}\Bigr{]}.$ $$

By the definition of KL divergence, for discrete $y$:

$$ $D_{\mathrm{KL}}\bigl{(}P(y)\,\|\,Q(y)\bigr{)}=\sum_{y}P(y)\,\log\frac{P(y)}{Q( y)}.$ $$

Hence, for each $e$,

$$ $D_{\mathrm{KL}}\bigl{(}p(y\mid{\mathbf{b}},e)\,\|\,q(y\mid{\mathbf{b}})\bigr{) }=\sum_{y}p(y\mid{\mathbf{b}},e)\,\log\frac{p(y\mid{\mathbf{b}},e)}{q(y\mid{ \mathbf{b}})}.$ $$

Therefore,

$$ $\mathcal{L}[q]=\sum_{e}p(e\mid{\mathbf{b}})\,\sum_{y}p(y\mid{\mathbf{b}},e)\, \log\frac{p(y\mid{\mathbf{b}},e)}{q(y\mid{\mathbf{b}})}.$ $$

We want to choose $q(y\mid{\mathbf{b}})$ to minimize $\mathcal{L}[q]$, which needs $q(y\mid{\mathbf{b}})$ to satisfy
$\sum_{y}q(y\mid{\mathbf{b}})\;=\;1,\text{for each fixed }{\mathbf{b}}.$
To incorporate this constraint, we introduce a Lagrange multiplier $\lambda$ and consider the Lagrangian:

$$ $\mathcal{J}[q,\lambda]=\sum_{e}p(e\mid{\mathbf{b}})\sum_{y}p(y\mid{\mathbf{b}} ,e)\,\log\!\biggl{[}\frac{p(y\mid{\mathbf{b}},e)}{q(y\mid{\mathbf{b}})}\biggr{ ]}\;+\;\lambda\Bigl{(}\sum_{y}q(y\mid{\mathbf{b}})-1\Bigr{)}.$ $$

To find the minimizing $q$, we set

$$ $\frac{\partial\mathcal{J}}{\partial q(y\mid{\mathbf{b}})}\;=\;0.$ $$

First, note

$$ $\log\frac{p(y\mid{\mathbf{b}},e)}{q(y\mid{\mathbf{b}})}\;=\;\log p(y\mid{ \mathbf{b}},e)-\log q(y\mid{\mathbf{b}}).$ $$

Only the term $-\log q(y\mid{\mathbf{b}})$ depends on $q$. For a specific $y$,

$$ $\frac{\partial}{\partial q(y\mid{\mathbf{b}})}\bigl{[}-\log q(y^{\prime}\mid{ \mathbf{b}})\bigr{]}=\begin{cases}-\frac{1}{q(y\mid{\mathbf{b}})},&\text{if }y ^{\prime}=y,\\ 0,&\text{if }y^{\prime}\neq y.\end{cases}$ $$

Hence, for each $e$:

$$ $\frac{\partial}{\partial q(y\mid{\mathbf{b}})}\biggl{[}\sum_{y^{\prime}}p(y^{ \prime}\mid{\mathbf{b}},e)\,\log\frac{p(y^{\prime}\mid{\mathbf{b}},e)}{q(y^{ \prime}\mid{\mathbf{b}})}\biggr{]}=-\frac{p(y\mid{\mathbf{b}},e)}{q(y\mid{ \mathbf{b}})}.$ $$

Summing over $e$,

$$ $\frac{\partial}{\partial q(y\mid{\mathbf{b}})}\sum_{e}p(e\mid{\mathbf{b}})\sum _{y^{\prime}}p(y^{\prime}\mid{\mathbf{b}},e)\,\log\frac{p(y^{\prime}\mid{ \mathbf{b}},e)}{q(y^{\prime}\mid{\mathbf{b}})}=-\frac{1}{q(y\mid{\mathbf{b}})} \sum_{e}p(e\mid{\mathbf{b}})\,p(y\mid{\mathbf{b}},e).$ $$

Also,

$$ $\frac{\partial}{\partial q(y\mid{\mathbf{b}})}\Bigl{[}\lambda\bigl{(}\sum_{y^{ \prime}}q(y^{\prime}\mid{\mathbf{b}})-1\bigr{)}\Bigr{]}=\lambda.$ $$

Hence,

$$ $\frac{\partial\mathcal{J}}{\partial q(y\mid{\mathbf{b}})}=-\frac{1}{q(y\mid{ \mathbf{b}})}\sum_{e}p(e\mid{\mathbf{b}})\,p(y\mid{\mathbf{b}},e)+\lambda.$ $$

Setting this to zero:

$$ $-\frac{1}{q(y\mid{\mathbf{b}})}\sum_{e}p(e\mid{\mathbf{b}})\,p(y\mid{\mathbf{b }},e)+\lambda=0,$ $$

which gives

$$ $\frac{1}{q(y\mid{\mathbf{b}})}\,\sum_{e}p(e\mid{\mathbf{b}})\,p(y\mid{\mathbf{ b}},e)=\lambda.$ $$

Thus

$$ $q(y\mid{\mathbf{b}})=\frac{1}{\lambda}\sum_{e}p(e\mid{\mathbf{b}})\,p(y\mid{ \mathbf{b}},e).$ $$

We require $\sum_{y}q(y\mid{\mathbf{b}})=1$. Then:

$$ $1=\sum_{y}q(y\mid{\mathbf{b}})=\sum_{y}\frac{1}{\lambda}\sum_{e}p(e\mid{ \mathbf{b}})\,p(y\mid{\mathbf{b}},e)=\frac{1}{\lambda}\sum_{e}p(e\mid{\mathbf{ b}})\sum_{y}p(y\mid{\mathbf{b}},e).$ $$

Since $\sum_{y}p(y\mid{\mathbf{b}},e)=1$ for each $e$, we get

$$ $\sum_{e}p(e\mid{\mathbf{b}})\,\underbrace{\sum_{y}p(y\mid{\mathbf{b}},e)}_{=1} =\sum_{e}p(e\mid{\mathbf{b}})=1.$ $$

Hence

$$ $1=\frac{1}{\lambda}\cdot 1\quad\Longrightarrow\quad\lambda=1.$ $$

Therefore,

$$ $q(y\mid{\mathbf{b}})=\sum_{e}p(e\mid{\mathbf{b}})\,p(y\mid{\mathbf{b}},e).$ $$

## Appendix E Algorithms

### E.1 Binary Bias-Aware Generalization

Figure: Algorithm 1 Bias-Aware Generalization

### E.2 Multi-Class Bias-Aware Generalization

Figure: Algorithm 2 Multi-Class Bias-Aware Generalization

## Appendix F Experiment details and ablation Study

### F.1 Experiment details

##### PACS

We evaluate our BAG model on the PACS dataset, which comprises four distinct domains (Photo, Art painting, Cartoon, Sketch) and seven object categories. Following the standard leave-one-domain-out protocol, we train on three domains and test on the held-out domain. Our backbone is a ResNet-18 pretrained on ImageNet, used to extract visual features. On top of this, we attach a variational autoencoder that separates causal factors from spurious factors, and we introduce learnable domain embeddings to encourage independence across environments. All key hyperparameters—such as the dimensions of the latent representations, the relative weighting between reconstruction and regularization losses, and the learning rates—are chosen via Bayesian optimization, with test-domain accuracy serving as the selection criterion. We train with the Adam optimizer for 70 epochs and select the final model based on its performance under test-time augmentation on the held-out domain. We used the code based on [11].

##### Office-Home

We evaluate our BAG model on the Office-Home dataset [13], which comprises four domains (Art, Clipart, Product, Real-World) and 65 object categories. Following the leave-one-domain-out protocol, we train on three domains and test on the held-out one. Hyperparameters for each held-out domain are optimized via Bayesian optimization; the selected values for learning rate, regularization weight, number of epochs for the $z_{s}$ classifier, and VAE weight are given in Table [F.1](https://arxiv.org/html/2507.17001v1#A6.T1). The full code are provided in the supplementary material.

**Table F.1: Optimal hyperparameters on Office-Home**
| Domain | Learning Rate | $\lambda_{1}$ | # epochs (test time) | $\lambda_{0}$ |
| --- | --- | --- | --- | --- |
| Art (Ar) | $5.7\times 10^{-5}$ | 0.030 | 3 | $1.2\times 10^{-4}$ |
| Clipart (Cl) | $6.6\times 10^{-5}$ | 0.012 | 3 | $2.1\times 10^{-5}$ |
| Product (Pr) | $5.1\times 10^{-5}$ | 0.030 | 4 | $2.2\times 10^{-4}$ |
| Real-World (Rw) | $5.0\times 10^{-5}$ | 0.046 | 6 | $2.5\times 10^{-5}$ |

### F.2 Ablation Study

[Table F.2](https://arxiv.org/html/2507.17001v1#A6.T2) presents the performance of various BAG variants on the PACS dataset and in simulation experiments. All variants outperform ERM, demonstrating the effectiveness of our two bias-utilization methods. Notably, BAG-RE and BAG-TTA each exhibit their own strengths and weaknesses, which may be due to data-specific effects. BAG-VAE performs somewhat worse than BAG, underscoring the importance of accurately identifying and disentangling $\mathbf{b}$ and $\mathbf{c}$.

**Table F.2: Accuracy of different variants including BAG, BAG-VAE, BAG-RE, and BAG-TTA on four PACS domains and Synthetic data.**
| Algorithm | PACS | Synthetic |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| P | A | C | S | Avg | Acc |  |
| BAG | 96.6 | 86.0 | 77.9 | 73.0 | 83.4 | 97.7 |
| BAG-VAE | 95.8 | 84.5 | 78.9 | 66.8 | 81.5 | 94.5 |
| BAG-RE | 96.1 | 82.8 | 78.4 | 69.7 | 81.8 | 77.0 |
| BAG-TTA | 96.3 | 80.0 | 78.8 | 73.0 | 82.0 | 85.5 |

## Appendix G Statement

This work contributes to the advancement of machine learning by challenging conventional out-of-distribution (OOD) approaches and introducing a framework that strategically leverages bias rather than eliminating it. On the positive side, our methods can improve model robustness when deployed in naturally biased environments—e.g., medical imaging systems that must generalize across hospitals with different demographic distributions. By revealing how bias can be harnessed as a constructive signal, we hope to inspire new lines of research into principled bias-aware generalization across diverse domains. Because this is foundational research evaluated on controlled benchmarks and synthetic datasets, we do not identify any immediate negative societal impacts.

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 6
presentation: 4
questions: A discussion of the theoretical assumptions (A1-A4) would be useful to understand where this is applicable, and where it is not (aka, which datasets/domains does this make sense to apply to, and which do not make sense?)
soundness: 3
strengths: - The paper has a fairly decent novelty contribution, challenging entrenched assumptions in DG research about bias should be eliminated or not and reframes bias as a potentially beneficial signal. Furthermore, the formalization of when bias helps prediction (via “unblocked influence”) and the proofs of identifiability and performance are useful for technical depth and novelty. The use of causal graphs and conditional independencies is well-motivated.
- BAG unifies multiple strands: causal representation learning, mixture-of-experts modeling, and adaptive label calibration — into a coherent probabilistic framework.
- Experiments are well-structured and make sense for the pipeline.
- Paper is well written and structured.
summary: The paper challenges the conventional view in out-of-distribution (OOD) generalization that bias must always be removed. Instead, it argues that bias can sometimes help generalization and proposes a framework (BAG) that explicitly models and leverages bias.
The authors provide a theoretical analysis identifying conditions under which bias can aid prediction,  bias-aware generative model that disentangles invariant content and bias features, and a mixture-of-experts mechanism for bias-aware prediction, guided by an adaptive label prior to handle both covariate and label shift.
weaknesses: - Only tests the problem on very small datasets, making it difficult to understand or interpret whether this framework is generalisable.
- The theoretical identifiability conditions (A1–A4) require smooth, positive densities and independent latent dimensions given (e, y). These are rarely satisfied in high-dimensional deep representations. A discussion of approximate or empirical identifiability would help.
- Although disentanglement is central, the paper lacks visualizations or examples showing what the learned bias and content dimensions represent in image space. It is a bit difficult to interpret some of the results due to this oversight.
- The gains seem to be mainly on the Synthetic case, where the non-synthetic examples have minimal gain that is often within the confidence threshold.

---

confidence: 5
contribution: 1
initial_rating: 0
presentation: 1
questions: see above
soundness: 3
strengths: (1) The fundamental idea is reasonable. Covariate bias is indeed influential in classification and OOD-related tasks. I believe the authors’ perspective is justified. However, I believe this can be both helpful and risky. it may improve performance on certain tasks but could also lead to biased errors in other scenarios. 


(2) The authors conduct experiments on both synthetic and real-world datasets, demonstrating the effectiveness of their proposed method compared to existing approaches.
summary: This paper challenges the typical view that bias should be minimized for robust out-of-distribution (OOD) generalization. The authors propose Bias-Aware Generalization (BAG), a framework that leverages both stable and bias features to improve prediction performance. Through experiments on synthetic and real-world datasets, BAG outperforms traditional methods like ERM and IRM. The approach combines stable feature predictors with bias terms and uses a mixture of experts model, enhanced by soft labeling for test-time adaptation, showing that bias, when carefully managed, can enhance OOD generalization.
weaknesses: I have reviewed this paper for NeurIPS 2025, in which 5 reviewers unanimously decided to reject the paper.

In my previous review, I outlined several issues, but I found that the authors did not make substantial revisions. Even minor issues, such as typographical errors in symbols and writing, were not corrected.

Given these considerations, I have decided to reject the paper. In today’s world, where the volume of peer review work is high, I strongly recommend that the authors take each reviewer’s feedback seriously and make substantial improvements addressing the weaknesses of the work before resubmitting. The fllowing reviews also included some previous review comments that I agree with, for reference.

(1) The decomposition of content and bias representations is not empirically validated. It remains a verbal argument. How can one verify that the learned representations actually correspond to the intended content and bias factors? This is a key issue.

(2) A central part of the method’s separation of bias and content relies on the independence loss. However, this only enforces a necessary, not sufficient, condition for conditional independence.

(3) The core of the paper lies in validating the benefit of explicitly modeling and utilizing bias, which seems to be evaluable without using test-time adaptation setting. I do not see the necessity of including test-time adaptation component. If I understand correctly, the proposed method is the only one that uses test time adaptation, which gives it an unfair advantage in the empirical evaluation.

(4) I recommend that the authors use the same set of baselines across the benchmarks on which they conduct experiments. Although the authors mention that SFB did not report performance on the Office-Home dataset, it still seems feasible to compare the performance of baselines like GMDG on the PACS dataset.

(5) Large part of the claimed contributions of the paper have been derived in Stable Feature Boosting (SFB) already.

(6) Similarly, a mixture of experts method that is not too dissimilar to what the authors propose has been proposed in Prashant et al. ‘Scalable out-of-distribution robustness in the presence of unobserved confounders’ (2025), but is not referenced and at least conceptually compared to.

(7) In line 1200 of the pseudocode, is there a typo? Should it be f_c(\mathbf C) instead of f_c \mathbf C?

(8) I suggest the authors visualize e b y c x in the clearest way possible. I cannot distinguish the difference between e and b. You assume b ⊥ c | y (line 207), and also assume y is not independent of e given c (lines 245-246), but I fail to understand the rationale behind these assumptions.

(9) The writing is unclear, and there are some necessary explanations for certain concepts that would be meaningful. For example, in Lemma 2.1, what does positive support refer to, and why is this emphasized? Isn't positive support self-evident? What does $\mathcal{B}$ mean in the context of domain variability? This is not explained. Additionally, what does linear independence mean, and why is it reasonable?

(10) it is unclear what "RE" in BAG-RE refers to (w/o Regularization of IND?), and what BAG-VAE stands for. Isn't your BAG method based on VAE already? what is BAG without VAE refers to? None of these are clearly explained, and the writing is very unclear.

---

confidence: 4
contribution: 2
initial_rating: 4
presentation: 2
questions: 1. Could the authors provide detailed abations to show the independent effect of the content predictor, bias predictor, environment routing module, and adaptive label prior? Which component contributes most to the final improvement?

2. Is the environment predictor $p(e|b)$ truly necessary? Would directly predicting $p(y|b)$ yield similar performance? How sensitive is the model to the number of experts $M$ in the routing mechanism?

3. How does the proposed predictor refinement differ from SFB’s bias correction step, both theoretically and empirically? Could the authors clarify the key advantage that BAG offers over SFB beyond architectural integration?

4. Are the experiments implemented within the DomainBed codebase for fair comparison? Do the authors plan to include larger or spurious-correlation datasets (e.g., Waterbirds, CelebA, CMNIST) to validate the claim that bias utilization helps in realistic settings?
soundness: 2
strengths: * **Interesting and original perspective** — The paper takes a refreshing stance by arguing that bias or spurious features are not always detrimental. This idea of *leveraging* bias for better OOD generalization is both conceptually interesting and relevant, especially given how dominant the “bias elimination” mindset has been in this field.

* **Solid theoretical framing** — The authors provide a clear theoretical justification for when and why bias can be beneficial, using identifiability conditions and causal reasoning. This helps ground the idea beyond intuition and makes the argument more rigorous than many prior discussions on bias utilization.

* **Coherent and well-structured framework** — The proposed BAG model integrates disentanglement, environment routing, and adaptive label priors into a unified end-to-end architecture. The framework is logically consistent and easy to follow, with each component corresponding naturally to a part of the theoretical motivation.
summary: This paper challenges the common assumption in OOD generalization that spurious or bias features should always be removed. The authors provide theoretical conditions under which bias can be identifiable and beneficial for prediction, particularly when there exists an unblocked causal path between bias and label. Building on this, they propose Bias-Aware Generalization (BAG) — an end-to-end framework that disentangles invariant and bias representations via a VAE with an independence regularizer, routes predictions through inferred environments, and corrects bias using an adaptive, learnable label prior. Experiments on synthetic and standard DG benchmarks (PACS, OfficeHome) show modest gains over ERM.
weaknesses: * **Lack of ablations and component analysis** — The paper does not provide sufficient ablation studies to isolate where the performance gain comes from. For instance, it’s unclear how much each part—the content predictor, bias predictor, environment routing module, or adaptive prior—actually contributes. The absence of results like *C-only vs. B-only*, or *with vs. without routing and prior correction*, makes it hard to judge whether the added modules are necessary or if the model mainly reproduces SFB’s effect in a different architecture.

* **Limited experimental scope and dataset diversity** — The experiments are conducted only on synthetic data, PACS, and OfficeHome. These are relatively small and saturated benchmarks that no longer reflect the challenges of modern OOD evaluation. The paper should consider following the **DomainBed** codebase and benchmark protocol (VLCS, PACS, OfficeHome, TerraIncognita, DomainNet) to ensure fair hyperparameter tuning and reproducible comparisons. In addition, it would be helpful to include datasets that explicitly test spurious correlations—such as **CMNIST**, **Waterbirds**, **CelebA**, **MetaShift**, or **BAR** (as used in *Project-Probe-Aggregate: Efficient Fine-Tuning for Group Robustness*)—to more convincingly demonstrate the benefit of leveraging bias. Currently, the performance improvements appear substantial on synthetic data but modest on real-world settings.

* **Overlap with prior work and unclear novelty** — The overall framework, especially the test-time refinement step and bias correction mechanism, is conceptually very close to Stable Feature Boosting (SFB). While the paper adds disentanglement and environment routing, it is not yet clear whether these bring substantial new insights or measurable benefits. Without stronger theoretical or empirical differentiation, the work feels more like an incremental extension rather than a clear step forward.

---

confidence: 2
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. How sensitive is performance to the number of domain experts and embedding dimensionality? Any heuristics or validation procedure beyond source-domain accuracy?
2. If p(e∣b) is wrong (e.g., target domain outside the convex hull of sources), how does BAG degrade? Any fallback to invariant-only predictions?
soundness: 3
strengths: 1. The paper formalizes when bias can help rather than hurt, via identifiability and “unblocked influence” conditions, then operationalizes it in BAG.
2. The adaptive label prior and decomposition (Eq. 5) address a key weakness of prior SFB-like methods that assume invariant
summary: The paper asks whether “bias” should always be removed for out-of-distribution (OOD) generalization and proposes BAG (Bias-Aware Generalization): (1) learn disentangled content/bias features via a VAE; (2) build a bias-aware predictor using an environment-routing mixture of experts with a learned domain estimator; and (3) perform bias correction with an adaptive label prior, yielding an “invariant predictor” term that remains valid under label shift. Experiments on synthetic data, PACS, and Office-Home show gains over ERM/IRM/ACTIR and recent bias-utilization SFB.
weaknesses: 1. Conditions A1–A4 (smooth densities, conditional independence factorization, domain variability/linearly independent shifts) may be hard to meet in complex vision tasks; practical diagnostics are not provided.
2. While block-wise identifiability is cited from prior work, the paper’s empirical section offers limited stress-tests for partial violations (e.g., non-identifiable generators, entangled factors)

---

confidence: 3
contribution: 3
initial_rating: 4
presentation: 3
questions: No supplementary material that provides code, which is mentioned in L1255 that “The full code are provided in the supplementary material”.
soundness: 3
strengths: 1.	The paper provides a strong theoretical analysis of the role of bias in ODD generalization, which inspires the design of the proposed framework BAG.
2.	Experiments on two benchmarks prove the effectiveness of the proposed method.
3.	The main ideas, motivations, and methodological details are generally well-presented.
summary: This paper questions the conventional treatment of bias in out-of-distribution (OOD) generalization, proposing instead that bias can be constructively leveraged for improved performance. The authors provide a theoretical framework for identifying when and how biased features contribute positively and introduce a generative model-based approach to disentangle invariant content from bias. They further propose a bias-aware prediction scheme that incorporates environment routing and adaptive bias correction with learnable label priors. Extensive experiments on synthetic data and established domain generalization benchmarks demonstrate that their approach (BAG) achieves robust improvements over both invariance-only and prior bias-utilization methods.
weaknesses: 1.	Theoretical validation is insufficient. Beyond performance results, the theoretical conclusions in the paper lack quantitative experimental support. I recommend adding empirical analyses that directly substantiate the theoretical claims, which would also help reinforce the motivation of the proposed method.
2.	Missing comparison with SOTA DG methods, like [1,2,3]. In addition, I am curious whether the proposed methods can be effectively applied on top of SOTA DG methods.
3.	The empirical results are only demonstrated on relatively modest-scale benchmarks(PACS and OfficeHome). Experiments on larger datasets (like DomainNet or TerraIncognita) should also be provided.

[1] Zou, Yingtian, et al. "Towards robust out-of-distribution generalization bounds via sharpness." arXiv preprint arXiv:2403.06392 (2024).

[2] Zhang, Qingyang, et al. "The best of both worlds: On the dilemma of out-of-distribution detection." Advances in Neural Information Processing Systems 37 (2024): 69716-69746.

[3] Chen, Jinggang, et al. "Gaia: Delving into gradient-based attribution abnormality for out-of-distribution detection." Advances in Neural Information Processing Systems 36 (2023): 79946-79958.


Decision: Reject