Heterogeneous Risk Minimization

Following (Arjovsky et al., 2019; Chang et al., 2020), we consider a dataset $D=\left\{D^{e}\right\}_{e\in\mathrm{supp}(\mathcal{E}_{tr})}$, which is a mixture of data $D^{e}=\left\{(x_{i}^{e},y_{i}^{e})\right\}_{i=1}^{n_{e}}$ collected from multiple training environments $e\in\mathrm{supp}(\mathcal{E}_{tr})$, $x_{i}^{e}\in\mathcal{X}$ and $y_{i}^{e}\in\mathcal{Y}$ are the $i$-th data and label from environment $e$ respectively and $n_{e}$ is number of samples in environment $e$. Environment labels are unavailable as in most real applications. $\mathcal{E}_{tr}$ is a random variable on indices of training environments and $P^{e}$ is the distribution of data and label in environment $e$.

The goal of this work is to find a predictor $f(\cdot):\mathcal{X}\rightarrow\mathcal{Y}$ with good out-of-distribution generalization performance, which can be formalized as:

where $\mathcal{L}(f|e)=\mathbb{E}[l(f(X),Y)|e]=\mathbb{E}^{e}[l(f(X^{e}),Y^{e})]$ is the risk of predictor $f$ on environment $e$, and $l(\cdot,\cdot):\mathcal{Y}\times\mathcal{Y}\rightarrow\mathbb{R}^{+}$ is the loss function. $\mathcal{E}$ is the random variable on indices of all possible environments such that $\mathrm{supp}(\mathcal{E})\supset\mathrm{supp}(\mathcal{E}_{tr})$. Usually, for all $e\in\mathrm{supp}(\mathcal{E})\setminus\mathrm{supp}(\mathcal{E}_{tr})$, the data and label distribution $P^{e}(X,Y)$ can be quite different from that of training environments $\mathcal{E}_{tr}$. Therefore, the problem in Equation 1 is referred to as Out-of-Distribution (OOD) Generalization problem (Arjovsky et al., 2019).

Without any prior knowledge or structural assumptions, it is impossible to figure out the OOD generalization problem, since one cannot characterize the unseen latent environments in $\mathrm{supp}(\mathcal{E})$. A commonly used assumption in invariant learning literature (Rojas-Carulla et al., 2015; Gong et al., 2016; Arjovsky et al., 2019; Kuang et al., 2020; Chang et al., 2020) is as follow:

This assumption indicates invariance and sufficiency for predicting the target $Y$ using $\Phi^{*}$, which is known as invariant covariates or representations with stable relationships with $Y$ across different environments $e\in\mathcal{E}$. In order to acquire the invariant predictor $\Phi^{*}(X)$, a branch of work to find maximal invariant predictor (Chang et al., 2020; Koyama & Yamaguchi, 2020) has been proposed, where the invariance set and the corresponding maximal invariant predictor are defined as:

Here we prove that the MIP $S$ can can guarantee OOD optimality, as indicated in Theorem 2.1. The formal statement of Theorem 2.1 as well as its proof can be found in appendix.

Recently, some works suppose the availability of data from multiple environments with environment labels, wherein they can find MIP (Chang et al., 2020; Koyama & Yamaguchi, 2020). However, they rely on the underlying assumption that the invariance set $\mathcal{I}_{\mathcal{E}_{tr}}$ of $\mathcal{E}_{tr}$ is exactly the invariance set $\mathcal{I}_{\mathcal{E}}$ of all possible unseen environments $\mathcal{E}$, which cannot be guaranteed as shown in Theorem 2.2.

As shown in Theorem 2.2 that $\mathcal{I}_{\mathcal{E}}\subseteq\mathcal{I}_{\mathcal{E}_{tr}}$, the learned predictor is only invariant to such limited environments $\mathcal{E}_{tr}$ but is not guaranteed to be invariant with respect to all possible environments $\mathcal{E}$.

Here we give a toy example in Table 1 to illustrate this. We consider a binary classification between cats and dogs, where each photo contains 3 features, animal feature $X_{1}\in\left\{\text{cat},\text{dog}\right\}$, a background feature $X_{2}\in\left\{\text{on grass},\text{in water}\right\}$ and the photographer’s signature feature $X_{3}\in\left\{\text{Irma},\text{Eric}\right\}$. Assume all possible testing environments $\mathrm{supp}(\mathcal{E})=\left\{e_{1},e_{2},e_{3},e_{4},e_{5},e_{6}\right\}$ and the train environment $\mathrm{supp}(\mathcal{E}_{tr})=\left\{e_{5},e_{6}\right\}$, then $\mathcal{I}_{\mathcal{E}}=\left\{\Phi|\Phi=\Phi(X_{1})\right\}$ while $\mathcal{I}_{\mathcal{E}_{tr}}=\left\{\Phi|\Phi=\Phi(X_{1},X_{2})\right\}$. The reason is that $e_{5},e_{6}$ only tell us $X_{3}$ cannot be included in the invariance set but cannot exclude $X_{2}$. But if $e_{5}$ and $e_{6}$ can be further divided into $e_{1},e_{2}$ and $e_{3},e_{4}$ respectively, the invariance set $\mathcal{I}_{\mathcal{E}_{tr}}$ becomes $\mathcal{I}_{\mathcal{E}_{tr}}=\mathcal{I}_{\mathcal{E}}=\{\Phi(X_{1})\}$.

This example shows that the manually labeled environments may not be sufficient to achieve MIP, not to mention the cases where environment labels are not available. This limitation necessitates the study on how to exploit the latent intrinsic heterogeneity in training data (like $e_{5}$ and $e_{6}$ in the above example) to form more refined environments for OOD generalization. The environments need to be subtly uncovered, in the sense of OOD generalization problem, as indicated by Theorem B.4, not all environments are helpful to tighten the invariance set.

Besides Assumption 2.1, we make another assumption on the existence of heterogeneity in training data as:

The heterogeneity among provided environments can be evaluated by the compactness of the corresponding invariance set as $|\mathcal{I}_{\mathcal{E}}|$. Specifically, smaller $|\mathcal{I}_{\mathcal{E}}|$ leads to higher heterogeneity, since more variant features can be excluded. Based on the assumption, we come up with the problem of heterogeneity exploitation for OOD generalization.

Theorem B.4 together with Assumption 2.2 indicate that, to better constrain $\mathcal{I}_{\mathcal{E}_{tr}}$, the effective way is to generate environments with varying $P(Y|\Psi^{*}(X))$ that can exclude variant features from $\mathcal{I}_{\mathcal{E}_{tr}}$. Under this problem setting, we encounter the circular dependency: first we need variant $\Psi^{*}$ to generate heterogeneous environments $\mathcal{E}_{tr}$; then we need $\mathcal{E}_{tr}$ to learned invariant $\Phi^{*}$ as well as variant $\Psi^{*}$. Furthermore, there exists positive feedback between these two steps. When acquiring $\mathcal{E}_{tr}$ with tighter $\mathcal{I}_{\mathcal{E}_{tr}}$, more invariant predictor $\Phi(X)$ (i.e. a better approximation of MIP) can be found, which will further bring a clearer picture of variant parts, and therefore promote the generation of $\mathcal{E}_{tr}$. With this notion, we propose our framework for Heterogeneous Risk Minimization (HRM) which leverages the mutual promotion between the two steps and conduct joint optimization.

Here we introduce our invariant prediction module $\mathcal{M}_{p}$, which takes multiple environments training data $D=\left\{D^{e}\right\}_{e\in supp(\mathcal{E}_{tr})}$ as input, and outputs the corresponding invariant predictor $f$ and the indices of invariant features $M$ given current environments $\mathcal{E}_{tr}$. We combine feature selection with invariant learning under heterogeneous environments, which can select the features with stable/invariant correlations with the label across $\mathcal{E}_{tr}$. Specifically, the former module can select most informative features with respect to the loss function and latter module ensures the selected features are invariant. Their combination ensures $\mathcal{M}_{p}$ to select the most informative invariant features.

For invariant learning, we follow the variance penalty regularizer proposed in (Koyama & Yamaguchi, 2020) and simplify it in feature selection scenarios. The objective function of $\mathcal{M}_{p}$ with $M\in\{0,1\}^{d}$ is:

However, as the optimization of hard feature selection with binary mask $M$ suffers from high variance, we use the soft feature selection with gates taking continuous value in $[0,1]$. Specifically, following (Yamada et al., 2020), we approximate each element of $M=[m_{1},\dots,m_{d}]^{T}$ to clipped Gaussian random variable parameterized by $\mu=[\mu_{1},\dots,\mu_{d}]^{T}$ as

where $\epsilon$ is drawn from $\mathcal{N}(0,\sigma^{2})$. With this approximation, the objective function with soft feature selection can be written as:

where $M$ is a random vector with $d$ independent variables $m_{i}$ for $i\in[d]$. Under the approximation in Equation 6, $\|M\|_{0}$ is simply $\sum_{i\in[d]}P(m_{i}>0)$ and can be calculated as $\|M\|_{0}=\sum_{i\in[d]}\mathrm{CDF}(\mu_{i}/\sigma)$, where $\mathrm{CDF}$ is the standard Gaussian CDF. We formulate our objective as risk minimization problem:

where

Further, as for linear models, we simply approximate the regularizer $\mathrm{trace}(\mathrm{Var}_{\mathcal{E}_{tr}}(\nabla_{\theta}\mathcal{L}^{e}))$ by $\|\mathrm{Var}_{\mathcal{E}_{tr}}(\nabla_{\theta}\mathcal{L}^{e})\odot M\|^{2}$. Then we obtain $\Phi(X)$ and $\Psi(X)$ when we obtain $\mu$ as well as $M$. Further in Section 4, we theoretically prove that the prediction module $\mathcal{M}_{p}$ is able to learn the MIP with respect to given environments $\mathcal{E}_{tr}$.

Notation. $\Psi$ means the learned variant part $\Psi(X)$. $\Delta_{K}$ means $K$-dimension simplex. $f_{\theta}(\cdot)$ means the function $f$ parameterized by $\theta$.

The heterogeneity identification module $\mathcal{M}_{c}$ takes a single dataset as input, and outputs a multi-environment dataset partition for invariant prediction. We implement it with a clustering algorithm. As indicated in Theorem B.4, the more diverse $P(Y|\Psi)$ for our generated environments, the better the invariance set $\mathcal{I}$ is. Therefore, we cluster the data points according to the relationship between $\Psi$ and $Y$, for which we use $P(Y|\Psi)$ as the cluster centre. Note that $\Psi$ is initialized as $\Psi(X)=X$ in our joint optimization.

Specifically, we assume the $j$-th cluster centre $P_{\Theta_{j}}(Y|\Psi)$ parameterized by $\Theta_{j}$ to be a Gaussian around $f_{\Theta_{j}}(\Psi)$ as $\mathcal{N}(f_{\Theta_{j}}(\Psi),\sigma^{2})$:

For the given $N=\sum_{e\in\mathcal{E}_{tr}}n_{e}$ empirical data samples $\mathcal{D}=\{\psi_{i}(x_{i}),y_{i}\}_{i=1}^{N}$, the empirical distribution is modeled as $\hat{P}_{N}=\frac{1}{N}\sum_{i=1}^{N}\delta_{i}(\Psi,Y)$ where

The target of our heterogeneous clustering is to find a distribution in $\mathcal{Q}=\{Q|Q=\sum_{j\in[K]}q_{j}h_{j}(\Psi,Y),\mathbb{q}\in\Delta_{K}\}$ to fit the empirical distribution best. Therefore, the objective function of our heterogeneous clustering is:

The above objective can be further simplified to:

As for optimization, we use EM algorithm to optimize the centre parameter $\Theta$ and the mixture weight $\mathbb{q}$. After optimizing equation 13, for building $\mathcal{E}_{tr}$, we assign each data point to environment $e_{j}\in\mathcal{E}_{tr}$ with probability:

In this way, $\mathcal{E}_{tr}$ is generated by $\mathcal{M}_{c}$.

We design two mechanisms to simulate the varying correlations among covariates across environments, named by selection bias and anti-causal effect.

Selection Bias In this setting, the correlations between variant covariates and the target are perturbed through selection bias mechanism. According to Assumption 2.1, we assume $X=[\Phi^{*},\Psi^{*}]^{T}\in\mathbb{R}^{d}$ and $Y=f(\Phi^{*})+\epsilon$ and that $P(Y|\Phi^{*})$ remains invariant across environments while $P(Y|\Psi^{*})$ changes arbitrarily. For simplicity, we select data points according to a certain variable set $V_{b}\subset\Psi^{*}$:

where $|r|>1$, $V_{b}\in\mathbb{R}^{n_{b}}$ and $\hat{P}(x)$ denotes the probability of point $x$ to be selected. Intuitively, $r$ eventually controls the strengths and direction of the spurious correlation between $V_{b}$ and $Y$(i.e. if $r>0$, a data point whose $V_{b}$ is close to its $y$ is more probably to be selected.). The larger value of $|r|$ means the stronger spurious correlation between $V_{b}$ and $Y$, and $r\geq 0$ means positive correlation and vice versa. Therefore, here we use $r$ to define different environments.

In training, we generate $sum=2000$ data points, where $\kappa=95\%$ points from environment $e_{1}$ with a predefined $r$ and $1-\kappa=5\%$ points from $e_{2}$ with $r=-1.1$. In testing, we generate data points for 10 environments with $r\in[-3,-2.7,-2.3,\dots,2.3,2.7,3.0]$. $\beta$ is set to 1.0. We compare our HRM with ERM, DRO, EIIL and IRM for Linear Regression. We conduct extensive experiments with different settings on $r$, $n_{b}$ and $d$. In each setting, we carry out the procedure 10 times and report the average results. The results are shown in Table 4.

From the results, we have the following observations and analysis: ERM suffers from the distributional shifts in testing and yields poor performance in most of the settings. DRO surprisingly has the worst performance, which we think is due to the over-pessimism problem (Frogner et al., 2019). EIIL has the similar performance with ERM, which indicates that its inferred environments cannot reveal the spurious correlations between $Y$ and $V_{b}$. IRM performs much better than the above two baselines, however, as IRM depends on the available environment labels to work, it uses much more information than the other three methods. Compared to the three baselines, our HRM achieves nearly perfect performance with respect to average performance and stability, especially the variance of losses across environments is close to 0, which reflects the effectiveness of our heterogeneous clustering as well as the invariant learning algorithm. Furthermore, our HRM does not need environment labels, which verifies that our clustering algorithm can mine the latent heterogeneity inside the data and further shows our superiority to IRM.

Besides, we visualize the differences between environments using Task2Vec (Achille et al., 2019) in Figure 2, where larger value means the two environments are more heterogeneous. The pooled training data are mixture of environments with $r=1.9$ and $r=-1.1$, the difference between whom is shown in yellow box. And the red boxes show differences between learned environments by HRM^s and HRM. The big promotion between $\mathcal{E}_{init}$ and $\mathcal{E}_{learn}$ verifies our HRM can exploit heterogeneity inside data as well as the existence of the positive feedback. Due to space limitation, results of varying $sum,\kappa,n_{b}$ as well as experimental details are left to appendix.

Anti-causal Effect Inspired by (Arjovsky et al., 2019), we induce the spurious correlation by using anti-causal relationship from the target $Y$ to the variant covariates $\Psi^{*}$. In this experiment, we assume $X=[\Phi^{*},\Psi^{*}]^{T}\in\mathbb{R}^{d}$ and firstly sample $\Phi^{*}$ from mixture Gaussian distribution characterized as $\sum_{i=1}^{k}z_{k}\mathcal{N}(\mu_{i},I)$ and the target $Y=\theta_{\phi}^{T}\Phi^{*}+\beta\Phi_{1}\Phi_{2}\Phi_{3}+\mathcal{N}(0,0.3)$. Then the spurious correlations between $\Psi^{*}$ and $Y$ are generated by anti-causal effect as

where $\sigma(\mu_{i})$ means the Gaussian noise added to $\Psi^{*}$ depends on which component the invariant covariates $\Phi^{*}$ belong to. Intuitively, in different Gaussian components, the corresponding correlations between $\Psi^{*}$ and $Y$ are varying due to the different value of $\sigma(\mu_{i})$. The larger the $\sigma(\mu_{i})$ is, the weaker correlation between $\Psi^{*}$ and $Y$. We use the mixture weight $Z=[z_{1},\dots,z_{k}]^{T}$ to define different environments, where different mixture weights represent different overall strength of the effect $Y$ on $\Psi^{*}$.

In this experiment, we set $\beta=0.1$ and build 10 environments with varying $\sigma$ and the dimension of $\Phi^{*},\Psi^{*}$, the first three for training and the last seven for testing. We run experiments for 10 times and the averaged results are shown in Table 3. EIIL achieves the best training performance with respect to prediction errors on training environments $e_{1}$, $e_{2}$, $e_{3}$, while its performances in testing are poor. ERM suffers from distributional shifts in testing. DRO seeks for over-considered robustness and performs much worse. IRM performs much better as it learns invariant representations with help of environment labels. HRM achieves nearly uniformly good performance in training environments as well as the testing ones, which validates the effectiveness of our method and proves its excellent generalization ability.

We test our method on three real-world tasks, including car insurance prediction, people income prediction and house price prediction.