Causality Pursuit from Heterogeneous Environments via Neural Adversarial Invariance Learning

Let us revisit the thought experiment from the perspective of a hyper-intelligent alien, Alice. Alice knows nothing about cows and camels except for 1000 images with annotated labels highly associated with the background, for example, $r=90\%$ cows/camels on grass/sand. It’s impossible for her to know that the background cannot determine the object given this limited information. In other words, both $X_{1}$ and $X_{2}$ can be regarded as causality out of pragmatic considerations. However, if she receives another set of 1000 labeled images, where $r=70\%$ cows/camels on grass/sand, she might begin to question the causality role of the background: the emerging evidence of the varying associations between $X_{2}$ and $Y$ falsify the hypothesis that $X_{2}$ is causality if she believes that causality persists across diverse environments.

When there is no supervision of the cause-effect relationship, the observation from heterogeneous sources is essential. We consider the following multi-environment regression problem that mimics human causality learning. Let $\mathcal{E}$ be the set of sources/environments. For each environment $e\in\mathcal{E}$, we observe $n$ i.i.d. data $(X_{1}^{(e)},Y_{1}^{(e)}),\ldots,(X_{n}^{(e)},Y_{n}^{(e)})\sim\mu^{(e)}$, where $\mu^{(e)}$, the joint distribution of $(X^{(e)},Y^{(e)})$, satisfies

Here $S^{\star}$, the unknown true important variable set, and $m^{\star}:\mathbb{R}^{|S^{\star}|}\to\mathbb{R}$, the target regression function, are both invariant across different environments; but the joint distributions $\mu^{(e)}$ can vary. We aim to learn the set of quasi-causal variables $S^{\star}$ and estimate the invariant regression function $m^{\star}$ using data $\{\{(X_{i}^{(e)},Y_{i}^{(e)})\}_{i=1}^{n}\}_{e\in\mathcal{E}}$ from $|\mathcal{E}|$ heterogeneous environments. The same $n$ in the problem formulation is just for expository simplicity, the extension to varying $n^{(e)}$ is straightforward. We refer to the above problem as nonparametric invariance pursuit or nonparametric causality pursuit exchangeably, as based on the data alone, without prior knowledge, we can not differentiate these concepts.

Here, we temporarily refrain from causal discussions. Under particular scenarios, such a problem can be instantiated to causal discovery in the Structural Causal Model (SCM) framework (Peters et al.,, 2016) and transfer learning with a more realistic assumption (Rojas-Carulla et al.,, 2018); see the details in Section A.1. We offer in Section 3 a rigorous and comprehensive interpretation of what $S^{\star}$ is in the SCM with interventions on $X$. It is also notable to mention that model (1.1) only requires invariance in the first moment instead of full distributional invariance, i.e., $\varepsilon^{(e)}\sim F_{\varepsilon}$ and independent of $X_{S^{\star}}^{(e)}$, as typically required for causal discovery (Peters et al.,, 2016). It is more realistic and allows for between-environment heteroscedastic errors.

It is important to note that the standard nonparametric regression generally diverges from our target $m^{\star}$, i.e., $\mathbb{E}[Y^{(e)}|X^{(e)}=x]\neq m^{\star}(x_{S^{\star}})$. This discrepancy arises because $\mathbb{E}[\varepsilon^{(e)}|X^{(e)}]\neq 0$. Such a “curse of endogeneity” problem is the main challenge we need to address. Including even one of endogenous spurious variables, for example, $X_{2}$ background color in the above thought experiment, in the regression function will create an inconsistent estimation of $m^{\star}$. Thus, it is essential to design an algorithm to eliminate all endogenous spurious variables.

This paper proposes a unified estimation framework – the Focused Adversarial Invariance Regularized (FAIR) estimator. It regularizes the user-specified risk loss $\ell(y,v)$ by a novel regularizer. Specifically, the FAIR estimator is the solution of the following minimax optimization program

Here $\ell(\cdot,\cdot)$ is a loss whose population solution leads to the conditional expectation, $\gamma>0$ is the regularization hyper-parameter to be determined, $(\mathcal{G},\mathcal{F})$ are the function classes to be specified by the user satisfying $\mathcal{G}\subseteq\mathcal{F}$. The first part is the risk minimization, and the second component is the test of exogeneity of the variables $S_{g}=\mathrm{supp}(g)$ used by the regression function $g$, where $\mathcal{F}_{S_{g}}=\{f\in\mathcal{F}:f(x)=h(x_{S_{g}})\text{ for some }h:\mathbb{R}^{|S_{g}|}\to\mathbb{R}\}$ is the testing function class for the prediction functions in $\mathcal{G}$ that only “focuses” on the variables $S_{g}$ that $g$ used. Two useful classes of functions are linear and square-integrable classes for $(\mathcal{G},\mathcal{F})$, which correspond respectively to linear models and nonparametric regression models; see Section 4.1 for additional details. Note that the second component is nonnegative after maximization by comparing with $f^{(e)}=0$ so that the penalty is nonnegative. For the empirical counterpart, we solve a similar minimax optimization program that substitutes $\mathbb{E}_{\mu^{(e)}}[\cdot]$ with the corresponding sample means.

To see why such a FAIR penalty works, let us consider the nonparametric regression setting in which $\mathcal{F}=\{f:\mathbb{E}_{\mu^{(e)}}[f^{2}(X_{S_{g}})]<\infty\}$. By conditioning on $X_{S_{g}}$, for $f^{(e)}\in\mathcal{F}_{S_{g}}$, we have

Then, the supremum in (1.2) can be explicitly found and the objective now becomes

Therefore, $g(X)=m^{\star}(X_{S^{\star}})$ is a minimax solution.

To motivate (1.2), let us first consider the additional constraint $\mathbb{E}_{\mu^{(e)}}[f^{(e)}(X_{S_{g}})^{2}]=1$ so that the first part of the second component in (1.2) is basically the maximal correlation between the residual $\{Y-g(X_{S_{g}})\}$ and testing functions $f^{(e)}(X_{S_{g}})$. Hence, the criterion (1.2) is to find a set of variables $X_{S_{g}}$ as exogenous (weakly correlated) with the residuals as possible for all testing functions in $\mathcal{F}_{S_{g}}$. By the Lagrange multiplier method, the constrained maximization problem can be written as

Choosing the multiplier $\lambda=1/2$ (justified in the above paragraph) gives rise to the object function (1.2).

FAIR penalty screens out all endogenous spurious variables when $\gamma$ is sufficiently large. This is easily seen when the penalty in (1.2) is not zero, such a $g$ is dominated by $g=m^{\star}$ when $\gamma$ is sufficiently large. After endogenous spurious variables, we can apply the commonly-used statistical variable selection methods (Hastie et al.,, 2009; Wainwright,, 2019; Fan et al.,, 2020) to further eliminate exogenous spurious or weak causal variables. In addition, we will show that under the SCM with arbitrary and nondegenerate interventions on $X$, our proposed FAIR estimator can unveil $S^{\star}$ being precisely expressed by the graph structure of the SCM, which can be interpreted as the “pragmatic” direct cause of the response $Y$ in general and will coincide with the direct causes if all the root children are intervened. The obtained result is clearly distinguished from what least squares, or even its worst-case variants like distribution robust optimization (Duchi & Namkoong,, 2021), Maximin (Meinshausen & Bühlmann,, 2015), can obtain. Our method indeed learns certain data-driven causality, while others cannot go beyond learning associations.

We propose a unified, algorithmic, and sample-efficient methodological framework that can discover the invariant regression function, i.e. to solve a generalized version of the problem in Section 1.1. The method is simple, universal, fully algorithmic, and sample-efficient: It is just one optimization objective (1.2) complemented by one extra hyper-parameter $\gamma$; it accommodates many losses and can be seamlessly integrated by various machine learning algorithms; it does not require any prior structural knowledge, and it is almost as statistically efficient as standard regression under various cases.

As a special instance in our framework, the FAIR neural network (FAIR-NN) estimator is proposed for which $\mathcal{G}$ and $\mathcal{F}$ are neural networks to unveil $m^{\star}$ in (1.1). It is the first theoretically guaranteed estimator that can efficiently recover $m^{\star}$ under a single general and minimal identification condition associated with the heterogeneity of the environments. Its sample efficiency can be understood in several notable aspects: it requires the minimal identification condition, leading to fewer required environments; it exhibits the same $L_{2}$ error rate as if directly regressing $Y$ on known $X_{S^{\star}}$, regardless of the complexity of spurious associations; and it adapts to the unknown low-dimension structure of the invariant association $m^{\star}$ in a same manner as Kohler & Langer, (2021). In summary, the FAIR-NN estimator circumvents the “curse of dimensionality” and “curse of endogeneity” simultaneously in a fully algorithmic manner, which does not rely on the prior knowledge of $m^{\star}$ structure or cause-effect relationships among variables.

While the complicated combinatorial constraint and minimax optimization are introduced in (1.2), we show that a variant of gradient descent – gradient descent ascent with Gumbel approximation to handle the combinatorial-nature “focused” constraint $f\in\mathcal{F}_{S_{g}}$ – continues to apply to our specifically designed algorithm and neural network estimators with no curse-of-dimension in implementation. Numerical results in Section 5 support this.

Though our framework is designed for algorithmic learning, it is versatile in that the user can also incorporate their strong prior structural knowledge such as linearity or additivity of $m^{\star}$ into the FAIR estimation. This can be realized by restricting the function class $\mathcal{G}$ within this known structure and designating $\mathcal{F}$ as a more expansive class. We demonstrate that harnessing such strong structural knowledge can relax the condition for identification. It is worth pointing out that identification is viable even when $|\mathcal{E}|=1$ corresponding to observational data; see examples in Section B.6. At the methodology level, our method bridges the invariance principle (Peters et al.,, 2016) and asymmetry principle (Janzing et al.,, 2016) for observational data into a unified framework.

Starting from the pioneering work of Peters et al., (2016), there is considerable literature proposing methods to estimate $m^{\star}$ in (1.1), predominantly when $m^{\star}$ is linear. These methods broadly fall into two categories: hypothesis test-based methods and optimization-based methods. For the hypothesis test-based methods (Peters et al.,, 2016; Heinze-Deml et al.,, 2018; Pfister et al.,, 2019), the Type-I error is controlled for an estimator $\widehat{S}$ with $\mathbb{P}(\widehat{S}\subseteq S^{\star})\geq 1-\alpha$. Nonetheless, these procedures may result in missing important variables or conservative solutions like $\widehat{S}=\emptyset$ due to the inherent worst-case construction in the algorithm. Additionally, the introduction of hypothesis tests also hinders its seamless integration by machine learning algorithms, limiting their scalability. On the other hand, some optimization-based methods (Ghassami et al.,, 2017; Rothenhäusler et al.,, 2019, 2021) focus on linear $m^{\star}$ and tackle the problem under additional structures such as linear SCMs with additive interventions (Rothenhäusler et al.,, 2019). This limitation curtails its applicability to a broader nonparametric setting. Some optimization-based methods (Pfister et al.,, 2021; Yin et al.,, 2021) designed for linear models are heuristic and lack finite sample guarantees. In summary, there is still a crucial gap towards efficiently estimating $m^{\star}$ without additional assumptions on the underlying model. Although Fan et al., (2023) recently bridged this gap for linear $m^{\star}$ through an optimization-based method, it is still unclear under the general nonparametric setting. This paper is the first to attain sample-efficient estimation for the general model with non-asymptotic guarantees in terms of both $|\mathcal{E}|$ and $n$.

Arjovsky et al., (2019) considers a general task, which aims to search for a data representation such that the optimal solution given that representation is optimal across diverse environments. They propose an optimization-based approach called invariant risk minimization (IRM), with many subsequent variants proposed later. However, their method comes with no statistical guarantees and requires at least $d$ environments even for the linear model, and the improvement over standard empirical risk minimization is not clear (Rosenfeld et al.,, 2021; Kamath et al.,, 2021). Our paper is the first to offer a comprehensive theoretical analysis of general invariance learning when the representation class is $\{(x_{1},\ldots,x_{d})\to(a_{1}x_{1},\ldots,a_{d}x_{d}):a_{1},\ldots,a_{d}\in\{0,1\}\}$ and to show that sample efficient estimation is in general viable even when $|\mathcal{E}|=2$. The main reason why this is attainable is due to the exact invariance pursued by our FAIR penalty and its “focused” nature, see the discussion in Section A.2.

Under the SCM framework, there is considerable literature on causal discovery using observational data (Spirtes et al.,, 2000; Richardson,, 1996; Chickering,, 2002; Hyttinen et al.,, 2013, 2014), but they cannot go beyond Markov equivalent class (Geiger & Pearl,, 1990) and thus fail to establish the exact cause-effect direction in general. Such a problem can be resolved by imposing additional assumptions under the circumstances that the algorithm can only passively observe data rather than performing intervention actively. These methods can be divided into two categories – one based on the invariance principle and the other based on the asymmetry principle. The invariance-based approaches (Peters et al.,, 2016) use samples from multiple experiments where some unknown intervention may apply to the variables other than $Y$. It leverages the idea that the cause-effect mechanism will remain constant while the reverse effect-cause association may vary. On the other hand, the asymmetry-based approaches (Shimizu et al.,, 2006; Hoyer et al.,, 2008; Zhang & Hyvärinen,, 2009; Janzing et al.,, 2012; Peters et al.,, 2014) only observe one sample of observational data and use the idea that the cause-effect mechanism admits a simple prior known structure, whereas its inverse does not, example includes the additive noise structure (Hoyer et al.,, 2008). These two principles for causal discovery seem to have been orthogonal before. Our estimation framework is the first to offer a unified methodological perspective on these two principles with theoretical guarantees. It demonstrates the ability to simultaneously leverage both principles for identification and estimation.

Adversarial estimation is introduced in Goodfellow et al., (2014) for generative modeling. Its application in the statistics spans distribution estimation (Liang,, 2021), instrumental variable regression (Dikkala et al.,, 2020), estimating the (implicit) influence function (Chernozhukov et al.,, 2020; Hirshberg & Wager,, 2021), and so on. The idea of minimizing the worst-case reward among diverse environments can also be considered as “an algorithmic remedy” for out-of-distribution generalization. There are different considerations of the “reward” such as risk (Sagawa et al.,, 2020), excess risk (Agarwal & Zhang,, 2022), and the negative of the explained variance (Meinshausen & Bühlmann,, 2015). However, these methods are conceptually similar to running least squares in regression and thus cannot go beyond just learning associations. We adopted the adversarial estimation in our estimation from two novel aspects. Firstly, it allows us to use a simple objective function that homogenizes different tasks and prediction models for estimation. Moreover, such a minimax optimization objective and the Gumbel approximation in the implementation jointly relax the combinatorial nature in (1.3) and make a variant of gradient descents continue to work numerically.

This paper is structured as follows. We first provide the proposed method with non-asymptotic theoretical analysis, and causal interpretations for our canonical nonparametric causality (invariance) pursuit problem in Sections 2–3, respectively. Such a special instance of our framework also helps to illustrate the main idea and philosophy of our general invariance pursuit problem and FAIR estimation framework, which will be formally presented in Section 4. In the main text, we provide a sketch of the abstract unified result, from which all non-asymptotic results are derived as corollaries, along with its other applications in Section 4.3 and defer the detailed statements to the Appendix. We provide a computationally efficient implementation using variants of gradient descent and Gumbel approximation, followed by its application to the simulation and real data analysis in Section 5. All the proofs are collected in the supplemental material.

We use upper case $(X,Y,Z)$ to represent random variables/vectors and denote their instances as $(x,y,z)$. Define $[n]=\{1,\ldots,n\}$. For a vector $x=(x_{1},\ldots,x_{d})^{\top}\in\mathbb{R}^{d}$, we let $\|x\|_{2}=(\sum_{j=1}^{d}x_{j}^{2})^{1/2}$. For given index set $S=\{j_{1},\ldots,j_{|S|}\}\subseteq[d]$ with $j_{1}<\cdots<j_{|S|}$, we denote $[x]_{S}=(x_{j_{1}},\ldots,x_{j_{|S|}})^{\top}\in\mathbb{R}^{|S|}$ and abbreviate it as $x_{S}$ if there is no ambiguity. We let $a\lor b=\max\{a,b\}$ and $a\land b=\min\{a,b\}$. We use $a(n)\lesssim b(n)$, $b(n)\gtrsim a(n)$, or $a(n)=O(b(n))$ if there exists some constant $C>0$ such that $a(n)\leq Cb(n)$ for any $n\geq 3$. Denote $a(n)\asymp b(n)$ if $a(n)\lesssim b(n)$ and $a(n)\gtrsim b(n)$. In the theorem statement and proof, we will use $C$ to represent the universal constants that may vary from line to line and will use $\widetilde{C},\widetilde{C}_{1},\ldots$ to represent the constant that may depend on the other constants defined in the paper.

In the context of the multi-environment setup, consider the following notations. For each $e\in\mathcal{E}$, let $\Theta^{(e)}=L_{2}(\mu^{(e)}_{x}):=\{f:\int f^{2}(x)\mu_{x}^{(e)}(dx)<\infty\}$, and denote $\|f\|_{2,e}=\{\int f^{2}(x)\mu^{(e)}_{x}(dx)\}^{1/2}$. Given $n$ observations $\{(X_{i}^{(e)},Y_{i}^{(e)})\}_{i=1}^{n}\subseteq\mathbb{R}^{d}\times\mathbb{R}$ drawn i.i.d. from $\mu^{(e)}$, we define $\mathbb{E}[f(X^{(e)},Y^{(e)})]=\int f(x,y)\mu^{(e)}(dx,dy)$ and $\widehat{\mathbb{E}}[f(X^{(e)},Y^{(e)})]=\frac{1}{n}\sum_{i=1}^{n}f(X^{(e)}_{i},Y^{(e)}_{i})$ for any $f\in\Theta^{(e)}$. We assume $\mathbb{E}[|Y^{(e)}|^{2}]<\infty$. Let $\bar{\mu}=\frac{1}{|\mathcal{E}|}\sum_{e\in\mathcal{E}}\mu^{(e)}$, and $\Theta=L_{2}(\bar{\mu}_{x})$ equipped with the norm $\|\cdot\|_{2}=\{\int f^{2}(x)\bar{\mu}_{x}(dx)\}^{1/2}$. It is easy to verify that $\Theta=\bigcap_{e\in\mathcal{E}}\Theta^{(e)}$.

Let $S\subseteq[d]$ be any index set. Given a function class $\mathcal{H}\subseteq\{h:\mathbb{R}^{d}\to\mathbb{R}\}$, we define $\mathcal{H}_{S}$ be the class of functions in $\mathcal{H}$ that only depend on variables $x_{S}$, i.e., $\mathcal{H}_{S}=\{h\in\mathcal{H},h(x)\equiv u(x_{S})\text{ for some }u:\mathbb{R}^{|S|}\to\mathbb{R}~{}~{}\mu^{(e)}\text{-}a.s.\forall e\in\mathcal{E}\}$. We sometimes also write $h(x_{S})$ instead of $h(x)$ for $h\in\mathcal{H}_{S}$ since $h$ only depends on $x_{S}$. For any $h\in\mathcal{H}$, we use $S_{h}\subseteq[d]$ to represent the index set of the variables $h$ depends on. We let $\{\mathcal{H}\}^{k}=\{(h_{1},\ldots,h_{k}):h_{i}\in\mathcal{H}~{}\forall i\in[k]\}$. For any $(X,Y)$’s joint distribution $\nu$, we use $\nu_{x}$ to denote the marginal distribution of $X$, and $\nu_{x,S}$ to denote the marginal distribution of $X_{S}$.

We introduce some notations. Recall that $\mu^{(e)}$ is the joint distribution of $(X,Y)$ in environment $e$. Let $m^{(e,S)}(x):=\mathbb{E}[Y^{(e)}|X_{S}^{(e)}=x_{S}]$ be the conditional expectation of $Y$ given $X_{S}$ in environment $e$. Recall that $\nu_{x,S}$ is the marginal distribution of $X_{S}$ for $(X,Y)\sim\nu$. It is easy to see that $\mu^{(e)}_{x,S}$ is absolutely continuous with respect to $\bar{\mu}_{x,S}=[\frac{1}{|\mathcal{E}|}\sum_{e\in\mathcal{E}}\mu^{(e)}]_{x,S}$ for any $S\subseteq[d]$ hence $\rho^{(e)}_{S}$, the Radon–Nikodym derivative of $\mu_{x,S}^{(e)}$ with respect to $\bar{\mu}_{x,S}$, is well defined. We define $\bar{m}^{(S)}(x)=\sum_{e\in\mathcal{E}}\rho^{(e)}_{S}(x_{S})m^{(e,S)}(x)$, which can be interpreted as the population-level least squares that regress $Y$ on $X_{S}$ using all the data in $\mathcal{E}$.

1 (a)–(b) is just a restatement of (1.1) together with i.i.d. data within each environment; data across different environments may be dependent. (c)–(d) are standard in nonparametric regression. (e) rules out some degenerate cases, for example, $m^{\star}(x_{1})=x_{1}^{2}$ with $S^{\star}=\{1\}$ and $X_{2}=X_{1}^{4}$, or $m^{\star}(x_{1},x_{2})=f(x_{1})$ with $S^{\star}=\{1,2\}$, and is imposed for technical convenience. The target (invariant) regression function in nonparametric causality pursuit is $m^{\star}$.

Given all the data $\{\{(X_{i}^{(e)},Y_{i}^{(e)})\}_{i=1}^{n}\}_{e\in\mathcal{E}}$ from heterogeneous environments, we consider using the following FAIR-NN least squares estimator to learn $m^{\star}$ in (1.1). Specifically, the FAIR-NN least squares estimator is the solution to the subsequent minimax optimization objective