Causally Inspired Regularization Enables Domain General Representations

A key feature of machine learning is its capacity to generalize across new domains. When these domains present different data distributions, the algorithm must leverage shared structural concepts to achieve out-of-distribution (OOD) or out-of-domain generalization. This capability is vital in numerous important real-world machine learning applications. For example, in safety-critical settings such as autonomous driving, a lack of resilience to unfamiliar distributions could lead to human casualties. Likewise, in the healthcare sector, where ethical considerations are critical, an inability to adjust to shifts in data distribution can result in unfair biases, manifesting as inconsistent performance across different demographic groups.

An influential approach to domain generalization is Invariant Causal Prediction (ICP; (Peters et al., 2016)). ICP posits that although some aspects of data distributions (like spurious or non-causal mechanisms (Pearl, 2010)) may change across domains, certain causal mechanisms remain constant. ICP suggests focusing on these invariant mechanisms for prediction. However, the estimation method for these invariant mechanisms suggested by (Peters et al., 2016) struggles with scalability in high-dimensional feature spaces. To overcome this, Arjovsky et al. (2019) introduced Invariant Risk Minimization (IRM), designed to identify these invariant mechanisms by minimizing an objective. However, requires strong assumptions for identifying the desired domain-general solutions (Ahuja et al., 2021; Rosenfeld et al., 2022); for instance, observing a number of domains proportional to the spurious features’ dimensions is necessary, posing a significant challenge in these high-dimensional settings.

Subsequent variants of IRM have been developed with improved capabilities for identifying domain-general solutions (Ahuja et al., 2020; Krueger et al., 2021; Robey et al., 2021; Wang et al., 2022; Ahuja et al., 2021). Additionally, regularizers for Distributionally Robust Optimization with subgroup shift have been proposed (GroupDRO) (Sagawa et al., 2019). However, despite their solid theoretical motivation, empirical evidence suggests that these methods may not consistently deliver domain-general solutions in practice Gulrajani and Lopez-Paz (2020); Kaur et al. (2022); Rosenfeld et al. (2022).

Kaur et al. (2022) demonstrated that regularizing directly for conditional independencies implied by the generative process can give domain-general solutions, including conditional independencies beyond those considered by IRM. However, their experimental approach involves regularization terms that require direct observation of spurious features, a condition not always feasible in real-world applications. Our proposed methodology also leverages regularizers inspired by the conditional independencies indicated by causal graphs but, crucially, it does so without necessitating prior knowledge (or proxies) of the spurious features.

The source of distribution shift can be isolated to components of the joint distribution. One special case of distribution shift is covariate shift (Shimodaira, 2000; Zadrozny, 2004; Huang et al., 2006; Gretton et al., 2009; Sugiyama et al., 2007; Bickel et al., 2009; Chen et al., 2016; Schneider et al., 2020), where only the covariate distribution $P(X)$ changes across domains. Ben-David et al. (2009) give upper-bounds on target error based on the $\mathcal{H}$-divergence between the source and target covariate distributions, which motivates domain alignment methods like the Domain Adversarial Neural Networks (Ganin et al., 2016) and others (Long et al., 2015; Blanchard et al., 2017). Others have followed up on this work with other notions of covariate distance for domain adaptation, such as mean maximum discrepancy (MMD) (Long et al., 2016), Wasserstein distance (Courty et al., 2017), etc. However, Kpotufe and Martinet (2018) show that these divergence metrics fail to capture many important properties of transferability, such as asymmetry and non-overlap** support. Furthermore, Zhao et al. (2019) shows that even with the alignment of covariates, large distances between label distributions can inhibit transfer; they propose a label conditional importance weighting adjustment to address this limitation. Other works have also proposed conditional covariate alignment (des Combes et al., 2020; Li et al., 2018c, b).

Another form of distribution shift is label shift, where only the label distribution changes across domains. Lipton et al. (2018) propose a method to address this scenario. Schrouff et al. (2022) illustrate that many real-world problems exhibit more complex ’compound’ shifts than just covariate or label shifts alone.

One can leverage domain adaptation to address distribution shifts; however, these methods are contingent on having access to unlabeled or partially labeled samples from the target domain during training. When such samples are available, more sophisticated domain adaptation strategies aim to leverage and adapt spurious feature information to enhance performance (Liu et al., 2021; Zhang et al., 2021; Kirichenko et al., 2022). However, domain generalization, as a problem, does not assume access to such samples (Muandet et al., 2013).

To address the domain generalization problem, Invariant Causal Predictors (ICP) leverage shared causal structure to learn domain-general predictors (Peters et al., 2016). Previous works, enumerated in the introduction (Section 1), have proposed various algorithms to identify domain-general predictors. Arjovsky et al. (2019)’s proposed invariance risk minimization (IRM) and its variants motivated by domain invariance:

where $R^{e}(w\circ\Phi)=\mathbb{E}\big{[}\ell(y,w\cdot\Phi(x))\big{]}$, with loss function $\ell$, feature extractor $\Phi$, and linear predictor $w$. This objective aims to learn a representation $\Phi$ such that predictor $w$ that minimizes empirical risks on average across all domains also minimizes within-domain empirical risk for all domains. However, Rosenfeld et al. (2020); Ahuja et al. (2020) showed that this objective requires unreasonable constraints on the number of observed domains at train times, e.g., observing distinct domains on the order of the rank of spurious features. Follow-up works have attempted to improve these limitations with stronger constraints on the problem – enumerated in the introduction section.

Our method falls under domain generalization; however, unlike the domain-general solutions previously discussed, our proposed solution leverages different conditions than domain invariance directly, which we show may be more suited to learning domain-general representations.

We often represent causal relationships with a causal graph. A causal graph is a directed acyclic graph (DAG), $G=(V,E)$, with nodes $V$ representing random variables and directed edges $E$ representing causal relationships, i.e., parents are causes and children are effects. A structural equation model (SEM) provides a mathematical representation of the causal relationships in its corresponding DAG. Each variable $Y\in V$ is given by $Y=f_{Y}(X)+\varepsilon_{Y}$, where $X$ denotes the parents of $Y$ in $G$, $f_{Y}$ is a deterministic function, and $\varepsilon_{Y}$ is an error capturing exogenous influences on $Y$. The main property we need here is that $f_{Y}$ is invariant to interventions to $V\backslash\{Y\}$ and is consequently invariant to changes in $P(V)$ induced by these interventions. Interventions refer to changes to $f_{Z}$, $Z\in V\backslash\{Y\}$.

In this work, we focus on domain-general predictors $d_{g}$ that are linear functions of features with domain-general mechanisms, denoted as $g_{\text{dg}}\coloneqq w\circ\Phi_{\text{dg}}$, where $w$ is a linear predictor and $\Phi_{\text{dg}}$ identifies features with domain-general mechanisms. We use domain-general rather than domain-invariant since domain-invariance is strongly tied to the property: ${Y\perp\!\!\!\perp e\,|\,Z_{\text{dg}}}$ (Arjovsky et al., 2019). As shown in the subsequent sections, this work leverages other properties of appropriate causal graphs to obtain domain-general features. This distinction is crucial given the challenges associated with learning domain-general features through domain-invariance methods (Rosenfeld et al., 2020).

Given the presence of a distribution shift, it’s essential to identify some common structure across domains that can be utilized for out-of-distribution (OOD) generalization. For example, Shimodaira (2000) assume $P(Y|X)$ is shared across all domains for the covariate shift problem. In this work, we consider a setting where each domain is composed of observed features and labels, $X\in\mathcal{X},Y\in\mathcal{Y}$, where $X$ is given by an invertible function $\Gamma$ of two latent random variables: domain-general $Z_{\text{dg}}\in\mathcal{Z}_{\text{dg}}$ and spurious $Z_{\text{spu}}\in\mathcal{Z}_{\text{spu}}$. By construction, the conditional expectation of the label $Y$ given the domain-general features $Z_{\text{dg}}$ is the same across domains, i.e.,

Conversely, this robustness to $e$ does not necessarily extend to spurious features $Z_{\text{spu}}$; in other words, $Z_{\text{spu}}$ may assume values that could lead a predictor relying on it to experience arbitrarily high error rates. Then, a sound strategy for learning a domain-general predictor – one that is robust to distribution shifts – is to identify the latent domain-general $Z_{\text{dg}}$ from the observed features $X$.

The approach we take to do this is motivated by the Reichenbach Common Cause Principle, which claims that if two events are correlated, there is either a causal connection between the correlated events that is responsible for the correlation or there is a third event, a so-called (Reichenbachian) common cause, which brings about the correlation (Hitchcock and Rédei, 2021; Rédei, 2002). This principle allows us to posit the class of generative processes or causal mechanisms that give rise to the correlated observed features and labels, where the observed features are a function of domain-general and spurious features. We represent these generative processes as causal graphs. Importantly, the map** from a node’s causal parents to itself is preserved in all distributions generated by the causal graph (Equation 1), and distributions can vary arbitrarily so long as they preserve the conditional independencies implied by the DAG (Markov Property (Pearl, 2010)).

We now enumerate DAGs that give observe features with spurious correlations with the label.

We have now clarified that hard DAGs (i.e., those not solved by ERM) share the TCRI property. The challenge is that $Z_{\text{dg}}$ and $Z_{\text{spu}}$ are not independently observed; otherwise, one could directly regularize. Existing work such as Kaur et al. (2022) empirically study semi-synthetic datasets where $Z_{\text{spu}}$ is (partially) observed and directly learn $Z_{\text{dg}}$ by regularizing that $\Phi(X)\perp\!\!\!\perp Z_{\text{spu}}\,|\,Y,e$ for feature extractor $\Phi$. To our knowledge, we are the first to leverage the TCRI property without requiring observation of $Z_{\text{spu}}$. Next, we set up our approach with some key assumptions. The first is that the observed distributions are Markov to an appropriate DAG.

These structural causal models (Equation 2-4) correspond to causal graphs Figures 2(a)-2(c), respectively.

Given Assumption 4.2, we aim to leverage TCRI property ($Z_{\text{dg}}\perp\!\!\!\perp Z_{\text{spu}}\,|\,Y\forall e\in\mathcal{E}_{tr}$) to learn the latent $Z_{\text{dg}}$ without observing $Z_{\text{spu}}$ directly. We do this by learning two feature extractors that, together, recover $Z_{\text{dg}}$ and $Z_{\text{spu}}$ and satisfy TCRI (Figure 3). We formally define these properties as follows.

In other words, a feature extractor that satisfies the total information criterion recovers the complete latent feature sets $Z_{\text{dg}},\,Z_{\text{spu}}$. This allows us to define the proposed implementation of the TCRI property non-trivially – the conditional independence of subsets of the latents may not have the same implications on domain generalization. We note that $X\perp\!\!\!\perp Y|Z_{\text{dg}},Z_{\text{spu}}$, so $X$ has no information about $Y$ that is not in $Z_{\text{dg}},Z_{\text{spu}}$.

Proposition 4.5 shows that TCRI is necessary and sufficient to identify $Z_{\text{dg}}$ from a set of training domains. We note that we can verify if $\Phi_{\text{dg}}(X)$ and $\Phi_{\text{spu}}(X)$ are correlated with $Y$ by checking if the learned predictors are equivalent to chance. Next, we describe our proposed algorithm to implement the conditions to learn such a feature map. Figure 3 illustrates the learning framework.

We begin by evaluating with simulated data, i.e., with known ground truth mechanisms; we use Equation 5 to generate our simulated data, with domain parameter $\sigma_{e_{i}}$; code is provided in the supplemental materials.

We observe 2 domains with parameters $\sigma_{e=0}=0.1$, $\sigma^{e=1}=0.2$ with $\sigma_{y}=0.25$, 5000 samples, and linear feature extractors and predictors. We use partial covariance as our conditional independence penalty $\mathcal{L}_{TCRI}$. Table  2 shows the learned value of $\Phi_{\text{dg}}$, where ‘Oracle’ indicates the true coefficients obtained by regressing $Y$ on domain-general $Z_{\text{dg}}$ directly. The ideal $\Phi_{\text{dg}}$ recovers $Z_{\text{dg}}$ and puts zero weight on $Z_{\text{spu}}$.

Now, we evaluate the efficacy of our proposed objective on non-simulated datasets.