text:
\section{Introduction}

Drawing conclusions by analyzing data from a single source carries serious challenges. For example, if a gene expression-phenotype correlation is found to be statistically significant in one lab, does it imply that we will be able to replicate it in other settings? It is possible that some shared latent characteristic between the samples explains such correlation, rather than a true biological connection between the gene and the disease. Confounding is a pervasive problem, where unobserved variables that change across environments invalidate any result made on a single dataset~\citep{zou2014epigenome}. Leveraging multiple data sources---also called environments---can help to identify the relationship between variables that are common across diverse settings and hence are more likely to be unconfounded true relations~\citep{ben2010theory}. However, how best to integrate multiple environments and when would it work are very much open areas of research. This question has recently generated wide interest in the machine learning community due to its connections to causal inference and robustness to domain shift~\citep{ben2010theory,imbens2015causal,kouw2016feature}. Environments can correspond to different batches, labs where the experiment is done, perturbations, etc. One approach is to first fit a regression model on each environment separately to obtain the environment $e$ specific parameters $\beta_e$ between the covariates and the response. Then some meta-analysis can be performed to aggregate the different $\beta_e$, or at the level of p-values if those are available. In the presence of unknown confounders, however, such two-stage meta-analysis can produce erroneous results. Other models introduce the environment as part of the model, and fit all the data where the heterogeneity across environments is captured by methods such as mixed-effects models \citep{pinheiro2000linear}, time varying coefficients \citep{hastie1993varying,fan1999statistical}, etc. However, these models assume a structured transformation across environments, and cannot capture a potential radical change in the covariate structure from one environment to another due to an experimental manipulation. Another approach is to focus on the invariant element of the model that we want to identify, and look for such invariant signal across environments \citep{peters2016causal}. We follow this line of work, where techniques differ from each other through their scope of application and their goal: some estimate the invariant vector of coefficients, while others provide a vector of coefficients that solves a minimax risk problem over the different environments. We will further discuss the relevant literature at the end of Section~\ref{section:section-2}.
\paragraph{Our Contributions.} We present a novel method for identifying invariant factors across multiple environments that we call \emph{Kullback-Leibler regression} (KL regression) focusing on identifying the covariates that have a causal effect on the response. Unlike other approaches, KL regression simultaneously allows for the existence of \emph{latent confounders}, is valid in settings where the different environments arise from distribution shifts in the observed covariates \emph{and} shifts in the response, and \emph{provably recovers the invariant causal vector} of coefficients. We provide asymptotic convergence results and demonstrate its improved performance compared to standard approaches in extensive experiments. Additionally, we introduce a \emph{Lasso KL regression} estimator, where a sparsity inducing penalty is added to the problem defining our estimator. In practice, this variant under a sparsity assumption on $\beta^*$ shows remarkable stability and robustness to model misspecification. We begin in Section~\ref{section:section-2} by presenting the problem of linear regression estimation with confounding due to latent variables based on linear Structural Equation Models (SEM) and describing the relevant literature. In Section~\ref{section:section-3} we describe KL regression in its general form, then present the required parametrization steps tailored to our SEM model that lead to our KL regression estimator (both sparse and non-sparse versions). We thoroughly validate the theoretical results empirically in Section~\ref{section:section-4} and compare our method to other commonly used techniques. KL regression achieves the top score for a DREAM5 challenge.

\section{Confounding with Structural Equation Models}\label{section:section-2}

\subsection{Generating Multiple Environments Via Linear Structural Equation Models}

Latent variables that introduce a confounding effect on the observed data is a common framework to represent spurious effects. In a linear setting, we may assume that the unconfounded vector of coefficients---the target of interest---which relates the response to the observed covariates represents an invariant structure that has a causal interpretation \citep{rojas2018invariant}. However, estimators like least squares regression may be biased due to the presence of latent confounding effects. We consider the case where the data follows a linear Structural Equation Model (SEM) \citep{bollen1989structural,kaplan2008structural}, which defines a distribution $\pi$ over the observed covariates $\boldsymbol{X}\in \mathbb{R}^D$, the hidden covariates $\boldsymbol{H}\in \mathbb{R}^Q$ and the observed response $Y\in \mathbb{R}$ as a solution to the system of equations:
\begin{align}
    \begin{cases}{}\label{equation-linear-sem-individual}
    H_q = \epsilon_{H,q} \quad \text{for $1\leq q \leq Q$}
    \\ X_d  = \sum_{1\leq d' \leq D} [B_{XX}]_{dd'}X_{d'}\! +\! \sum_{1\leq q\leq Q}[B_{XH}]_{dq}H_q
    \\ \hspace{2cm} + \; \epsilon_{X,d} \qquad \text{for $1\leq d \leq D$}
    \\ Y = \sum_{1\leq d \leq D} \beta^*_{d}X_d + \sum_{1\leq q \leq Q} \eta_{0,q} H_q + \epsilon_Y
    \end{cases}
\end{align}{}
where the coefficients are given by the \emph{connectivity} matrices $B_{XX}\in\mathbb{R}^{D\times D}$, $B_{XH}\in\mathbb{R}^{D\times Q}$ and the vectors $\beta^*\in \mathbb{R}^D$, $\eta_0 \in \mathbb{R}^Q$. We define $\epsilon := (\epsilon_X,\epsilon_Y,\epsilon_H) \in \mathbb{R}^{D+Q+1}$ as a random variable with finite second moments, where its three components are pairwise independent random vectors (one could equivalently represent the latent variable effect implicitly with a SEM defined only for $(\boldsymbol{X}, Y)$ where $\epsilon_X$ and $\epsilon_Y$ are correlated). We can write this model in a matrix form:
\begin{equation}\label{equation-linear-sem}
    \begin{pmatrix}
    \boldsymbol{X} \\ Y \\ \boldsymbol{H}
    \end{pmatrix} \!=\! \mathbf{B} \!\begin{pmatrix}
    \boldsymbol{X} \\ Y \\ \boldsymbol{H}
    \end{pmatrix} \!+ \!\begin{pmatrix}
    \epsilon_X \\ \epsilon_Y \\ \epsilon_H
    \end{pmatrix}\!; \; \mathbf{B}\! = \!\begin{pmatrix}
    B_{XX} & 0 & B_{XH}
    \\ \beta^{*T} & 0 & \eta_0^T
    \\ 0 & 0 & 0
    \end{pmatrix}
\end{equation}
where $(\Sigma_{\epsilon_X},\Sigma_{\epsilon_Y}, \Sigma_{\epsilon_H})$ are the covariance matrices of $\epsilon_X,\epsilon_Y,\epsilon_H$ respectively. We assume that $(I -\mathbf{B})$ is invertible, and thus the solution $\pi$ to the system is characterized by $\mathbf{B}$ and the distribution of $\epsilon$ through $\pi := \mathcal{D}\big((I -\mathbf{B})^{-1}\epsilon\big)$ where we denote by $\mathcal{D}(U)$ the distribution of the random variable $U$. The SEM model presented above can be reformulated with a probabilistic graphical model, where the variables (both observed or latent) correspond to the nodes and the edges are determined by the connectivity matrix $\mathbf{B}$. Our objective is to identify the parameter $\beta^*$ that captures the causal dependence between the observed covariates and the response, and especially to identify the non-zero entries: this is the problem of causal discovery. In the context of graphical models we want to identify the parent nodes of the response node, especially so when $\beta^*$ is assumed to be sparse and it is of importance to identify the few causal factors. Based only on $\boldsymbol{X},Y$ samples from $\pi$, the least squares regression vector of $Y$ on the covariates $\boldsymbol{X}$ is equal to $\beta^*$ plus an additive confounding term (we provide an explicit formulation in Prop.~\ref{proposition:covariance-parameters}). Our method recovers $\beta^*$ by leveraging information from multiple environments, and under sparsity assumptions we propose a regularized version to specifically identify relevant covariates. Environments in practice correspond to different data sources: if the collected data are patient medical records, different hospitals may correspond to different environments. When analyzing cellular genetic pathways, measurements from different cell cultures where different genes have been knocked out correspond to data from different environments. The crucial \emph{invariance} assumption is that the mechanism generating the response variable $Y$, encoded in the parameters of the structural equation of $Y$, is constant across environments. In particular, those parameters relating the response to the observed covariates, encoded by $\beta^*$, are invariant. We observe data from $E\geq 1$ different environments, which are characterized by their underlying respective distributions over $\mathbb{R}^D \times \mathbb{R}$. We use the subscript $(\boldsymbol{X}_e, Y_e)$ to refer to the observed variables in environment $e \in [E]:=\{1,\dots,E\}$ (and no longer to the covariate index as in \ref{equation-linear-sem-individual}). Our main assumption defines how models relate to each other -- i.e. which elements in equation~\ref{equation-linear-sem} may change--, which we call a \emph{model of environments}. This represents our beliefs on how interventions may perturb the initial observational data generating distribution. In our work, we only allow for $(\Sigma_{e,\epsilon_X}, \Sigma_{e,\epsilon_Y})$ to depend on the environment, so that each $e$-th environment distribution $\pi_e$ is the solution of the linear SEM:
\begin{equation}\label{eq:model-of-environments}
    \begin{pmatrix}
    \boldsymbol{X}_e \\ Y_e \\ \boldsymbol{H}
    \end{pmatrix} = \mathbf{B} \begin{pmatrix}
    \boldsymbol{X}_e \\ Y_e \\ \boldsymbol{H}
    \end{pmatrix} + \begin{pmatrix}
    \epsilon_{e,X} \\ \epsilon_{e,Y} \\ \epsilon_H
    \end{pmatrix}
\end{equation}
This model of environments can be understood as transformations from a base SEM via ``shift interventions'' on variables $\boldsymbol{X},Y$ \citep{rothenhausler2015backshift, rothenhausler2019causal}, where an environment-specific additive noise term is added to $\epsilon$ in the SEM: the distribution of the $e$-th environment is given by $\pi_e := \mathcal{D}\big((I -\mathbf{B})^{-1}(\epsilon + v_e)\big)$ where $v_e = (v_{e,X}, v_{e,Y}, 0)$ is a centered random variable in $\mathbb{R}^{D+Q+1}$ where only the first $D + 1$ coordinates are non-zero, independent of $\epsilon$. These interventions are also called ``parametric'' interventions (on $\boldsymbol{X}$ and $Y$) \citep{eberhardt2007interventions} or ``soft'' interventions \citep{eaton2007exact}. Constraining $v_{e,Y} = 0$ would correspond to shift interventions only on $\boldsymbol{X}$ (cf. next paragraph and ``Shift $\boldsymbol{X}$'', ``Shift $Y$'' in Table~\ref{table-methods} for further details). These contrast with ``stuctural'' or ``perfect'' interventions that also allow for modifications in the matrix $\mathbf{B}$, which occur for example with do-interventions \citep{judea2000causality}. A graphical representation of this model of environments is presented in Figure~\ref{fig:DAG}: different environments are obtained via changes in the distributions of $\epsilon_X, \epsilon_Y$ via shift interventions, represented via red circles. The remaining structure of the model is unchanged across environments. Defining a model of environments carries a trade-off: the smaller the set of elements that change across environments is, the easier it is to find invariant quantities across environments, but conversely that corresponds to stronger assumptions that will restrict the scope of applications. From now on we no longer explicitly mention the latent variables and hence $\pi_e$ refers to the marginal distribution over $\boldsymbol{X}_e, Y_e$.
\begin{figure}
    \centering
    \includegraphics[width=.4\linewidth]{figure_dag.png}
    \caption{Graphical representation of the process generating multiple environments.}
    \label{fig:DAG}
\end{figure}
\begin{table}
  \caption{Multi-environment analysis methods and their scope of application.}
  \label{table-methods}
  \centering
  \begin{tabular}{llllll}
    Method     & Shift $\boldsymbol{X}$ & Shift $Y$ & Confounding & Recovers $\beta^*$\\
    \midrule
    KL Regression  & Yes & Yes & Yes & Yes \\
    Invariant Causal Prediction  & Yes & No &  No & Yes \\
    Causal Dantzig  & Yes & Yes & Yes & Yes\\
    Invariant Risk Minimization  & Yes & Yes & No & No \\
    Anchor Regression  & Yes & Yes & Yes & No\\
    Causal Transfer Learning & Yes & No & No & Yes \\
    \bottomrule
  \end{tabular}
\end{table}
\paragraph{Related Literature}
The idea that true causal mechanisms are \emph{invariant} across different experimental settings \citep{peters2016causal, buhlmann2018invariance} has generated a lot of interest recently in the causal inference community. Based on samples collected from different ``environments'', it is possible to identify covariates whose effect on the response variable is universal, which are defined as causal. Solving the causal discovery problem problem with purely observational data is not possible, and a popular line of work in causal inference which uses Bayesian networks to represent causal models \citep{pearl1995causal,spirtes2000causation} relies on experimental or interventional data recover the graph structure. Pearl's do-calculus \citep{judea2000causality} identifies causal effects based on observational data but assuming perfect knowledge of the graphical structure of the Bayesian network. Other methods assume different data sources are generated through a change in the structure of the Bayesian network, and under some assumptions the causal graph can be recovered \citep{eberhardt2010combining,hyttinen2012learning}. This requires detailed knowledge on the perturbations that originate the distinct data sources. On the other hand, under a more limited knowledge of the mechanism that generates different data sources, the different distributions can be modeled with ``uncertain'' interventions \citep{korb2004varieties,eaton2007exact,eberhardt2010combining,tian2013causal}, for which we do not have a clear knowledge, and is the setting we assume in this work. Also, instead of estimating the complete graph structure \citep{spirtes2000causation,chickering2002optimal,hauser2012characterization,rothenhausler2015backshift}, we rather focus on identifying the sub-graph that connects the response variable to the other covariates. This leads us to formulating a linear SEM as in equation~\ref{equation-linear-sem}, and we now present some methods that solve this problem (we provide in Table~\ref{table-methods} a quick summary of how these relate to each other). The main modeling assumption for any procedure leveraging multiple environments is the definition of how such environments relate one to another, i.e. what the \emph{model of environments} is. Usually, procedures that focus on recovering $\beta^*$ (using our linear SEM framework) usually do not allow for any modification of the structural equation of $Y$: this is the basis of Invariant Causal Prediction (ICP) \citep{peters2016causal} and further methods based on the same principle \citep{heinze2018invariant,pfister2019invariant}. This body of work is based on the assumption that the conditional distribution of the response on a subset of covariates that are termed ``causal'' is constant across environments: i.e. if $S^*\subset[D]$ are the set of indices of such causal variables, denoting $\boldsymbol{U}_{A}:= (U_a)_{a\in A}$ for a vector $\boldsymbol{U}$, and $A$ subset of indices, then
\begin{equation}\label{invariant-assumption}
    Y_e|\boldsymbol{X}_{e,S^*} \stackrel{d}{=}Y_f|\boldsymbol{X}_{f,S^*} \quad \forall e, f\in [E]
\end{equation}{}
This assumption, translated into our SEM model, implies that the row corresponding to $Y$ is invariant across environments (i.e. no ``shift $Y$'' interventions in Table~\ref{table-methods}), hence our assumptions are weaker in that aspect. This weakened assumption for the model of environments is in line with what methods such as Invariant Risk Minimization (IRM) \citep[Definition 6, p.10]{arjovsky2019invariant} assume and argue is needed in real-world applications: changes in the noise level of the response across environments. The structure relating the covariates to the response---within the connectivity matrix $\mathbf{B}$---is invariant in any case (indeed, $\beta^*$ is defined because of such invariance). However, our method fundamentally differs from those other methods that allow shifts in $Y$, as their goal is not to provide an estimate of $\beta^*$ (i.e. ``Recovers $\beta^*$'' in Table~\ref{table-methods}), but instead solve distributional robustness problems:
\begin{align*}
    \hat{\beta} := \min_{b \in \mathbb{R}^D} \max_{e\in [E]} \mathbb{E}[(Y_e - \boldsymbol{X}_e b)^2]
\end{align*}{}
These estimators minimize the worst-case prediction risk, and are related to adversarial training \citep{sinha2017certifying,gao2017wasserstein}. Whenever the set of environments contains all possible modifications of the structural equations except that of the target, the output of such optimization problem is $\beta^*$, as described in Causal Transfer Learning (CTL) \citep{rojas2018invariant}. Besides this setting, theoretical guarantees of those estimators will usually focus on hedging against a particular subset of interventions instead of estimating $\beta^*$, as in Anchor Regression \citep{rothenhausler2018anchor}. Finally, going back to those methods that focus on recovering $\beta^*$, most of them do not allow for confounding (i.e. ``Confounding'' in Table~\ref{table-methods}). The invariance assumption~(\ref{invariant-assumption}) fails whenever a latent variable $\boldsymbol{H}$ simultaneously affects $Y$ or $\boldsymbol{X}$, or whenever its distribution changes across environments \citep[Example 1, p.6]{christiansen2020switching}. Several of the previously mentioned techniques fail to recover $\beta^*$ under confounding, and recently proposed methods such as \cite{christiansen2020switching} restrict the latent variables to be discrete (and do not allow for changes in the noise level $\Sigma_{\epsilon_Y}$ of the response). KL regression shares many characteristics with Causal Dantzig \citep{rothenhausler2019causal}, which has an unregularized and a sparsity-inducing regularization versions. These two methods are related in a way that is analogous to the relationship between the Lasso linear regression \citep{tibshirani1996regression} and the Dantzig Selector \citep{candes2007dantzig}. Causal Dantzig without regularization leads to an estimator that is similar to ours in its scope of applications. However, as we show in the experiments, KL regression outperforms unregularized Causal Dantzig in terms of stability. Using a sparsity-inducing regularization in Causal Dantzig greatly improves its behavior and matches KL regression performance. However, this regularized Causal Dantzig \emph{needs} sparsity in $\beta^*$, otherwise it does not properly recover it. KL regression does not require sparsity to perform well, although the regularized KL regression version substantially outperforms any other method in sparse settings.

\subsection{Invariant Elements for our Model of Environments}

Assume that variables $\epsilon$ defined in the SEM (cf. equation~\ref{equation-linear-sem}) have finite second moments. Denote by $\Sigma_e^{joint}$ the covariance matrix of the observed variables in environment $e$, which we assume has full rank.
\begin{align*}
    \Sigma_e^{joint} := \begin{pmatrix}
    \Sigma_{e,X} & \Sigma_{e,XY}
    \\ \Sigma_{e,XY}^T & \Sigma_{e,Y}
    \end{pmatrix} := \text{Cov}\begin{pmatrix}
    \boldsymbol{X}_e
    \\Y_e
    \end{pmatrix}
\end{align*}
Our model of environments assumption, that relies on a linear SEM, leads to a specific structure of these covariances across environments. Our method recovers $\beta^*$ only via the analysis of $\big(\Sigma_e^{joint}\big)_{e \in [E]}$, i.e. the second moments, regardless of the specific distribution of $\epsilon$. As a first step, we observe that in any given environment we can decouple the confounding bias from $\beta^*$ in the regression vector estimated by least squares.
\begin{proposition}\label{proposition:covariance-parameters}
Under our model of environments (cf. equation~\ref{eq:model-of-environments}) we have:
\begin{align*}
    & \Sigma_{e,X} =  C\Big(\Sigma_{e,\epsilon_{X}} + B_{XH}\Sigma_{\epsilon_H}B_{XH}^T\Big)C^T
    \\ & \Sigma_{e,XY} = \Sigma_{e,X}\beta^* + \eta^*
\end{align*}{}
where $C := (I -B_{XX})^{-1}$ and $\eta^* := CB_{XH}\Sigma_{\epsilon_{H}}\eta_0$.
\end{proposition}
We prove this result in Appendix~\ref{proof:proposition:covariance-parameters}. In particular, $\Sigma_{e,X}$ depends on the environment $e$ through $\Sigma_{e, \epsilon_X}$, and $\beta^*, \eta^*$ do not depend on $e$. Given data from the $e$-th environment, the least squares estimator $\hat{\beta}_e$ estimates the population regression coefficient given by $\beta_e := \Sigma_{e,X}^{-1}\Sigma_{e,XY} = \beta^* + \Sigma_{e,X}^{-1}\eta^*$. The confounding due to the additional term $\Sigma_{e,X}^{-1}\eta^*$ implies that, based on purely observational data (i.e. data from one single environment), least squares is biased for $\beta^*$.

\section{Kullback-Leibler Regression}\label{section:section-3}

\subsection{General Framework}
We reformulate the joint covariance of environment $e$ through the regression coefficient $\beta_e$ and the residual variance given by \mbox{$\sigma_e^2 := \Sigma_{e,Y} - \beta_e^T\Sigma_{e,X}\beta_e \in \mathbb{R}$}:
\begin{align*}
    \Sigma_e^{joint} =
    \begin{pmatrix}
    \Sigma_{e,X} & \Sigma_{e,X}\beta_e
    \\ \beta_e^T\Sigma_{e,X} & \sigma_e^2 + \beta_e^T\Sigma_{e,X}\beta_e
    \end{pmatrix}
\end{align*}
Let $\Pi$ be a mapping that we call the \emph{regression function} associated to the $e$-th environment. Its input consists on the pair $(\Sigma_{e,X}, \sigma_e^2)$ --that we call \emph{regressors}, and the \emph{regression coefficient} $\theta$:
\begin{equation}\label{distribution-map}
    \Pi:\begin{cases}{}\big(\mathbb{R}^{D\times D}\times\mathbb{R}_+\big)\times\mathbb{R}^D \rightarrow \mathbb{R}^{(D+1)\times (D+1)}
    \\(\Sigma_{e,X}, \sigma_e^2), \theta \mapsto
    \begin{pmatrix}
    \Sigma_{e,X} & \Sigma_{e,X}\theta
    \\ \theta^T\Sigma_{e,X} & \sigma_e^2 + \theta^T\Sigma_{e,X}\theta
    \end{pmatrix}
    \end{cases}
\end{equation}{}
The output is a positive semi-definite matrix such that $\Sigma_e^{joint} = \Pi\big((\Sigma_{e,X}, \sigma_e^2),\beta_e\big)$. For notation simplicity, let $\Pi_e(\theta):=\Pi\big((\Sigma_{e,X}, \sigma_e^2),\theta\big)$. Given a distance $\delta(\cdot, \cdot)$ over the set of positive semi-definite matrices, we can construct a loss function over $\theta$:
\begin{align}\label{loss-eq-simple}
    & \mathcal{L}(\theta) := \sum_{e\in [E]} \delta\big(\Pi_e(\theta), \Sigma_e^{joint}\big)
\end{align}{}
Although the distributions $\pi_e$ are not Gaussian, we can follow a quasi-maximum likelihood approach and define $\delta$ as the Kullback-Leibler divergence between two multivariate Gaussian variables with covariances $\Pi_e(\theta)$ and $\Sigma_e^{joint}$. This leads to the following environment-wise Mahalanobis distance in $\theta$ parametrized by the regressors $\Sigma_{e,X},\sigma_e^2$.
\begin{align*}
    \delta\big(\Pi_e(\theta), \Sigma_e^{joint}\big) := KL\Big(\mathcal{N}\big(\Pi_e(\theta)\big)\Big\Vert \mathcal{N}\big(\Sigma_e^{joint}\big)\Big)
\end{align*}
We provide in Appendix~\ref{appendix:regression-analogy} additional insights on this choice of distance and why our method can be understood as performing linear regression in a space of distributions. We then construct an estimator $\hat{\theta}$ by minimizing the corresponding empirical risk by plugging-in the empirical counterparts of $\Sigma_{e,X}, \sigma_e^2,\beta_e$ (denoted by a hat):
\begin{align}\label{loss-eq-simple-empirical}
    & \hat{\mathcal{L}}(\theta) := \sum_{e\in [E]} \delta\Big(\Pi\big((\hat{\Sigma}_{e,X},\hat{\sigma}_e^2), \theta\big), \hat{\Sigma}_e^{joint}\Big)
    \\ & \hat{\theta} := \argmin_{\theta\in \mathbb{R}^D} \hat{\mathcal{L}}(\theta)
\end{align}{}
This particular optimization problem in $\theta$ is convex, and has a closed-form solution. This procedure looks for a parameter $\theta$ that makes the environment-wise distributions as compatible as possible with a \emph{model of distributions}. Here, by minimizing $\mathcal{L}(\theta)$, we implicitly assume that there is a coefficient $\theta^*$ that defines for each environment $e$ a distribution $\pi_e^*$ whose covariance matrix is given by $\Pi_e(\theta^*)$. The observed environment distributions $\pi_e$ are then just ``unstructured'' perturbations of such $\pi_e^*$, that is there is no relationship between the regressors (i.e. $(\Sigma_{e,X}, \sigma_e^2)$) and the perturbation of the coefficient $\beta_e$ with respect to $\theta^*$.  We here understand the role of the structural causal model: individually, each of the environment distributions do not allow us to recover the invariant vector. It is only through the joint analysis of the environments that we identify it, and thus one needs to account for this causal structure in the parametrization in the loss. The problem formulation above is not adapted to our setting where the confounding bias in $\beta_e$ is environment-dependent, related to the regressors.

\subsection{Kullback-Leibler Regression for our Model of Environments}

\begin{figure*}[t]
\centering
\includegraphics[width=0.9\linewidth]{figure_a.pdf}
\caption{\textbf{Baseline comparison} Plots of the (log-scaled) MSE of the estimator of the regression parameter (with respect to the true $\beta^*$) as a function of the number of samples. The $L_1$ regularized KL regression consistently achieved the lowest MSE.}
\label{fig:figure_a}
\end{figure*}

CV the choice of lambda for KL regression to get also a decrease with number of samples.

Whenever our environment distributions $\pi_e$ come from a latent variable model as in the previous section (cf. equation~\ref{eq:model-of-environments}), minimizing the parameter $\theta$ in the mapping $\Pi_e(\theta)$ in \eqref{distribution-map} that defines our loss function does not yield the desired result, basically because there is no $\theta^*$ that is equal to all $\beta_e$ simultaneously. However, as we saw in Proposition~\ref{proposition:covariance-parameters}, the expression of the confounded linear regression coefficient $\beta_e$ in environment $e$ can be decomposed into two terms $\beta_e = \beta^* + \Sigma_{e,X}^{-1}\eta^*$, where $\eta^*$ no longer depends on $e$. As such, we can modify the loss in \eqref{loss-eq-simple} by expressing $\theta = \beta + \Sigma_{e,X}^{-1}\eta$, and then minimize the reparametrized loss over $(\beta,\eta)$, which is still jointly convex in both arguments:
\begin{align}\label{loss-DR}
    & \mathcal{L}(\beta, \eta) \!:=\! \sum_{e\in [E]}\! \delta\Big(\!\Pi_e(\beta + \Sigma_{e,X}^{-1}\eta))\Vert \Sigma_e^{joint}\!\Big)
    \\ & (\tilde{\beta}^{KL}, \tilde{\eta}) := \argmin_{\beta, \eta \in \mathbb{R}^D} \mathcal{L}(\beta, \eta)
\end{align}{}
Importantly, the reparametrization only requires knowledge of the covariance of $\boldsymbol{X}_e$, which can be estimated. This is the key insight of our method: because of the invariant structure of the model of environments and the judicious choice of distance in the space of covariances, we can construct a tractable loss that is minimized at the invariant vector. We denote by $\tilde{\beta}^{KL}$ the solution to the above optimization problem, and look for conditions that guarantee uniqueness of the solution, thus identifying $\beta^*$.
\begin{proposition}\label{proposition:solution-DR}
Provided that the following matrix is invertible
\begin{align}
    S_{\beta}:= \Big(\sum_{e\in [E]}\frac{\Sigma_{e,X}^{-1}}{\sigma_e^2}\Big)\Big(\sum_{e\in [E]}\frac{\Sigma_{e,X}}{\sigma_e^2}\Big) - \Big(\sum_{e\in [E]}\frac{1}{\sigma_e^2}\Big)^2I_D
\end{align}
the unique solution to the KL regression problem~\ref{loss-DR} has a closed form given by
\begin{align*}
    \tilde{\beta}^{KL} = S_{\beta}^{-1}\Bigg(&\Big(\sum_{e\in [E]}\frac{\Sigma_{e,X}^{-1}}{\sigma_e^2}\Big)\Big(\sum_{e\in [E]}\frac{\Sigma_{e,X}\beta_e}{\sigma_e^2}\Big)  -  \Big(\sum_{e\in [E]}\frac{1}{\sigma_e^2}\Big)\Big(\sum_{e\in [E]}\frac{\beta_e}{\sigma_e^2}\Big)\Bigg)
\end{align*}{}
Assuming that our environment distributions $\pi_e$ follow the SEM model above, then we have $\tilde{\beta}^{KL} = \beta^*$ (at the population level, i.e. assuming we can perfectly estimate the joint covariance under $\pi_e$).
\end{proposition}{}
\begin{proof}\label{proof:proposition:distribution-DR}
We prove here Proposition~\ref{proposition:distribution-DR}. Again, we drop the subscript $X$ from the covariance matrix $\Sigma_{e,X}$ for notation simplicity. Let the following matrices in $\mathbb{R}^{dE\times d}$:
\begin{align*}
    \mathbb{U} =  \begin{pmatrix}{}
    \frac{\hat{\Sigma}_1}{\sigma_1}
    \\\vdots
    \\ \frac{\hat{\Sigma}_E}{\sigma_E}
    \end{pmatrix}; \quad
    \mathbb{V} =  \begin{pmatrix}{}
    \frac{\hat{\Sigma}^{-1}_1}{\sigma_1}
    \\\vdots
    \\ \frac{\hat{\Sigma}^{-1}_E}{\sigma_E}
    \end{pmatrix} ; \quad
    \mathbb{I} =  \begin{pmatrix}{}
    \frac{I_D}{\sigma_1}
    \\\vdots
    \\ \frac{I_D}{\sigma_E}
    \end{pmatrix}
\end{align*}{}
We have that
\begin{align*}
    \hat{S}_{\beta} = & \mathbb{I}^T\mathbb{V}\mathbb{U}^T\mathbb{I} - \mathbb{I}^T\mathbb{I}\mathbb{I}^T\mathbb{I}
    \\ \hat{\beta} = & (\mathbb{I}^T\mathbb{V}\mathbb{U}^T\mathbb{I} - \mathbb{I}^T\mathbb{I}\mathbb{I}^T\mathbb{I})^{-1}(\mathbb{I}^T\mathbb{V}\mathbb{U}^T - \mathbb{I}^T\mathbb{I}\mathbb{I}^T)\hat{\mathbb{B}}
    \\ = & \mathbb{W}\hat{\mathbb{B}}
\end{align*}{}
where
\begin{align*}
    \mathbb{W} = & (\mathbb{I}^T\mathbb{V}\mathbb{U}^T\mathbb{I} - \mathbb{I}^T\mathbb{I}\mathbb{I}^T\mathbb{I})^{-1}(\mathbb{I}^T\mathbb{V}\mathbb{U}^T - \mathbb{I}^T\mathbb{I}\mathbb{I}^T) \in \mathbb{R}^{d\times dE}
    \\ \hat{\mathbb{B}} = & \begin{pmatrix}{}
    \frac{\hat{\beta}_1}{\sigma_1}
    \\ \vdots
    \\ \frac{\hat{\beta}_E}{\sigma_E}
    \end{pmatrix} \in \mathbb{R}^{dE}; \qquad
    \mathbb{B} =  \begin{pmatrix}{}
    \frac{\beta_1}{\sigma_1}
    \\ \vdots
    \\ \frac{\beta_E}{\sigma_E}
    \end{pmatrix} \in \mathbb{R}^{dE}
\end{align*}
so that each stacked vector in $\hat{\mathbb{B}}$ is independent of the others, and for every $e\in [E]$, we have that $\hat{\beta}_e = \beta_e + \frac{\sigma_e}{\sqrt{n_e}}\hat{\Sigma}^{-1/2}_{e,X}z_e$ where $(z_e)_{1\leq e \leq E}$ are independent standard Gaussians. We therefore get that
\begin{align*}
    \hat{\mathbb{B}} = & \begin{pmatrix}{}
    \frac{\beta_1}{\sigma_1}
    \\ \vdots
    \\ \frac{\beta_E}{\sigma_E}
    \end{pmatrix} + \begin{pmatrix}
    \frac{1}{\sqrt{n_1}}\hat{\Sigma}_1^{-1/2}z_1
    \\ \vdots
    \\ \frac{1}{\sqrt{n_E}}\hat{\Sigma}_E^{-1/2}z_E
    \end{pmatrix} = \mathbb{B} + \mathbb{Z}
\end{align*}{}
Therefore $\hat{\beta}^{KL} = \mathbb{W}\hat{\mathbb{B}} = \mathbb{W}\mathbb{B} + \mathbb{W}\mathbb{Z}$. But we have that $\mathbb{W}\mathbb{B} = \mathbb{E}[\hat{\beta}^{KL}|(X_e)_e]$, and $\mathbb{W}\mathbb{Z}$ is a multivariate Gaussian, whose covariance is given by :
\begin{align*}
    V = & \mathbb{W}\begin{pmatrix}{}
    \frac{1}{n_1}\hat{\Sigma}^{-1}_1 & 0&0 &\cdots & 0
    \\0 & \frac{1}{n_2}\hat{\Sigma}^{-1}_2& 0 &\cdots & 0
    \\ 0 & 0 & \ddots& \cdots & 0
    \\ \vdots & \vdots&&\vdots&
    \\ 0 & \cdots & 0& 0 &\frac{1}{n_E}\hat{\Sigma}_E^{-1}
    \end{pmatrix}\mathbb{W}^T
    \\ = & \sum_{e=1}^E \frac{1}{n_e}A_e
\end{align*}{}
for some symmetric matrix $A_e$. Now, if all the environments have the same sample size, $n_e = n$, then
\begin{align*}
    V & = \frac{1}{n}\hat{S}_{\beta}^{-1}\mathbb{I}^TA\mathbb{I}\hat{S}_{\beta}^{-1,T}
    \\ A & = (\mathbb{V}\mathbb{U}^T - \mathbb{I}\mathbb{I}^T)\begin{pmatrix}{}
    \hat{\Sigma}^{-1}_1 & 0&0 &\cdots & 0
    \\0 & \hat{\Sigma}^{-1}_2& 0 &\cdots & 0
    \\ 0 & 0 & \ddots& \cdots & 0
    \\ \vdots & \vdots&&\vdots&
    \\ 0 & \cdots & 0& 0 &\hat{\Sigma}_E^{-1}
    \end{pmatrix}(\mathbb{U}\mathbb{V}^T - \mathbb{I}\mathbb{I}^T)
    \\ & = \Big[\frac{\hat{\Sigma}_i^{-1}\hat{\Sigma}_j - I}{\sigma_i \sigma_j} \Big]_{i,j}\Big[ \hat{\Sigma}_j^{-1}1_{j=k}\Big]_{j,k}\Big[\frac{\hat{\Sigma}_k\hat{\Sigma}_l^{-1} - I}{\sigma_k \sigma_l} \Big]_{k,l}
    \\ & = \Big[ \sum_{j=1}^E \frac{\hat{\Sigma}_i^{-1}\hat{\Sigma}_j\hat{\Sigma}_l^{-1} - \hat{\Sigma}_i^{-1} - \hat{\Sigma}_l^{-1} + \hat{\Sigma}_j^{-1}}{\sigma_i\sigma_j^2\sigma_l}\Big]_{i,l}
\end{align*}
\begin{align*}
    \mathbb{I}^TA\mathbb{I} & = \sum_{i,j,l = 1}^E \frac{\hat{\Sigma}_i^{-1}\hat{\Sigma}_j\hat{\Sigma}_l^{-1} - \hat{\Sigma}_i^{-1} - \hat{\Sigma}_l^{-1} + \hat{\Sigma}_j^{-1}}{\sigma_i^2\sigma_j^2\sigma_l^2}
    \\  &=\Big(\sum_{e\in [E]}\frac{\hat{\Sigma}_e^{-1}}{\sigma_e^2}\Big)\Big(\sum_{f\in [E]}\frac{\hat{\Sigma}_f}{\sigma_f^2}\Big)\Big(\sum_{e\in [E]}\frac{\hat{\Sigma}_e^{-1}}{\sigma_e^2}\Big) - \Big(\sum_{e\in [E]}\frac{1}{\sigma_e^2}\Big)^2\Big(\sum_{e\in [E]}\frac{\hat{\Sigma}_e^{-1}}{\sigma_e^2}\Big)
    \\  &=\Big(\sum_{e\in [E]}\frac{\hat{\Sigma}_e^{-1}}{\sigma_e^2}\Big)\hat{S}_{\beta}^{T}
\end{align*}{}
Therefore we conclude
\begin{equation*}
V = \frac{1}{n} \hat{S}_{\beta}^{-1}\Big(\sum_{e\in [E]}\frac{\hat{\Sigma}_e^{-1}}{\sigma_e^2}\Big)
\end{equation*}
\end{proof}{}

The conditions for this proposition are mild: instead of assuming specific types of interventions for each individual covariate to guarantee identifiability of $\beta^*$ \citep{hyttinen2012learning, peters2016causal}, we just require diverse enough environments in terms of the covariates covariance. Identifiability of $\beta^*$ holds as long as $S_{\beta}$ is invertible, regardless of the actual distribution of $\epsilon$. Under which conditions is the matrix $S_{\beta}$ invertible? In particular, if only one environment is available then $S_{\beta}$ is not invertible: we can not make any statement based on purely observational data. However, KL regression can output $\tilde{\beta}^{KL}$ based on only two different environments under some conditions, where all covariates are simultaneously perturbed in one of the two environments. We now reformulate a result from \cite{bhagwat_subramanian_1978} under our framework, which we prove in Appendix~\ref{proof:proposition:invertibility_S} for completeness.
\begin{proposition}[\cite{bhagwat_subramanian_1978}]\label{proposition:invertibility_S}
A sufficient condition for $S_{\beta}$ to be invertible is that there exists $e^*\in [E]$ such that the following matrix is invertible
\begin{align*}
\Sigma_{e^*,X}^{-1}  - \Big(\sum_{e\in [E]}\frac{1}{\sigma_e^2}\Big)^{-1}\Big(\sum_{e\in [E]}\frac{\Sigma_{e,X}^{-1}}{\sigma_e^2}\Big)
\end{align*}
\end{proposition}{}
Assume for simplicity that the residual variances are equal to 1 (corresponding to independent measurement errors), if $\Sigma_{2,X}^{-1}\Sigma_{1,X} - I_D$ is invertible then $S_{\beta}$ is invertible: consider the following example where such scenario may arise. A scientist can experimentally down-regulate simultaneously all (measured) expressions of genes of interest, without interfering with unobserved latent genes. Covariances of gene expressions will be on different scales (eg. $\Sigma_{1,X} \succ \Sigma_{2,X}$), and the latent confounding effect will remain unaltered (assuming a linear gene regulatory network model). For these two environments (observational and interventional), the condition above can be easily satisfied. More generally, simulations show that whenever the environments are more heterogeneous, KL regression performs better.
In practice, all the quantities that define the loss $\mathcal{L}$ are estimated from the empirical covariances $\hat{\Sigma}_e^{joint}$: define $\hat{\beta}^{KL}$ as the plug-in estimator of $\tilde{\beta}^{KL}$ where we replace the environment population covariances (and derived quantities) by their empirical counterparts, where analogously $\hat{S}_{\beta}$ is the empirical version of $S_{\beta}$. For each environment $e \in [E]$ we observe \mbox{$n_e\geq 1$} i.i.d. samples. We can provide a distributional result: assuming that we have access to the population residual variances $\sigma_e^2$, conditioned on the sample covariates (i.e. as in fixed design linear regression), the KL regression estimator $\hat{\beta}^{KL}$ satisfies the following result:
\begin{proposition}\label{proposition:distribution-DR}
We have that, for $\epsilon$ Gaussian,
\begin{align*}
    \hat{\beta}^{KL}\mid(X_e)_e \sim \mathcal{N}\big(\mathbb{E}[\hat{\beta}^{KL}|(X_e)_e], V \big)
\end{align*}{}
where $V = \sum_{e=1}^E \frac{1}{n_e}A_e$, for some symmetric matrices $A_e$ that depend only on the empirical covariance matrices. In the case where $n_e = n$ for all $e$, we have
\begin{align*}
V = \frac{1}{n} \hat{S}_{\beta}^{-1}\Big(\sum_{e\in [E]}\frac{\hat{\Sigma}_{e,X}^{-1}}{\sigma_e^2}\Big)
\end{align*}
\end{proposition}{}
We prove this proposition in Appendix~\ref{proof:proposition:distribution-DR}. If all environments have the same number of samples $n$, the plug-in estimator $\hat{\beta}^{KL}$ converges to $\tilde{\beta}^{KL}$ at rate $1/ \sqrt{n}$. Also, whenever $S_{\beta}$ is ill-conditioned, the variance of the estimator $\hat{\beta}^{KL}$ will be large. In general, we do not have access to the population residual variances $\sigma_e^2$, and the distribution is not Gaussian. We leave this extension to future work.

\subsection{Kullback-Leibler Regression with Sparsity Inducing Penalty}

\begin{figure*}[t]
\centering
\includegraphics[width=0.8\linewidth]{figure_c.pdf}
\caption{\textbf{Diverse environments improve KL regression performance} Plots of the (log-scaled) MSE as a function of the environment-wise random component in $\Sigma_{e,\epsilon_X}$. }
\label{fig:environment-diversity}
\end{figure*}

Often times it is reasonable to assume that the invariant vector $\beta^*$ is sparse. Modifying KL regression to incorporate this prior knowledge can substantially improve the accuracy and stability of the method. We add an $L^1$-norm sparsity-inducing penalty to the optimization problem (\ref{loss-DR}), which is convex in $(\beta,\eta)$, to define the following \emph{Lasso KL regression}:
\begin{align}\label{loss-DR-sparse}
    & \mathcal{L}_{\lambda}(\beta, \eta) \!:=\!  \mathcal{L}(\beta, \eta) + \lambda\Vert\beta\Vert_1
    \\ & (\tilde{\beta}^{KL}(\lambda), \tilde{\eta}(\lambda)) := \argmin_{\beta, \eta \in \mathbb{R}^D} \mathcal{L}_{\lambda}(\beta, \eta)
\end{align}{}
Solving this problem in practice is easily done with off-shelf optimization packages such as CVXPY \citep{diamond2016cvxpy} in Python. Surprisingly, Lasso KL regression outperforms all the competing methods in sparse settings, is stable in small sample settings and robust to model misspecification. We leave for future work a deeper analysis of finite sample guarantees of the Lasso KL regression.

\subsection{Robustness of Kullback-Leibler Regression to Model Misspecification}

If we assume that the object of interest in the linear SEM (\ref{equation-linear-sem}), $\beta^*$, remains invariant across environments then we can bound its estimation error under misspecification. Assume we allow small perturbations in the remaining coefficients of our SEM model~(\ref{eq:model-of-environments}) for some of the environments: we define \mbox{$\eta_e := (I-B_{XX,e})^{-1}B_{XH,e}\Sigma_{\epsilon_X ,e}\eta_{0,e}$}. Define \mbox{$\delta_e := \eta_e - \eta^*$}, we have the following bound:
\begin{proposition}\label{proposition:robustness} We have:
\begin{align*}
& \Vert S_{\beta} (\tilde{\beta}^{KL} - \beta^*)\Vert^2 \leq C \sup_{e\in [E]}\Vert \delta_e\Vert^2
\end{align*}
where \mbox{$C:= \big(\sum_{e\in [e]}\frac{1}{\sigma_e^2}\big)\Big(\Vert\sum_{e\in [E]}\frac{\Sigma_e^{-1}}{\sigma_e^2}\Vert + \sum_{e\in [e]}\frac{\Vert\Sigma_e^{-1}\Vert}{\sigma_e^2}\Big)$}
\end{proposition}
In practice, we can bound the misspecification error in $\tilde{\beta}^{KL}$ with a bound on the (unknown) term $\sup_{e\in [E]}\Vert \delta_e\Vert^2$, as the remaining terms in $C$ can be estimated from the data. In particular, the effect of the model misspecification is entirely captured by $\delta_e$, even though the covariate covariance is also modified by such perturbations.

\begin{figure}
     \centering
     \begin{subfigure}
\centering
\includegraphics[width=.8\linewidth]{figure_f.pdf}
\end{subfigure}
     \vfill
     \begin{subfigure}
\centering
\includegraphics[width=.8\linewidth]{figure_b.pdf}
\end{subfigure}
     \vfill
     \begin{subfigure}
\centering
\includegraphics[width=.8\linewidth]{figure_j.pdf}
\end{subfigure}
\label{fig:model-misspecification}
\caption{\textbf{Effect of model misspecification.} Plots of the (log-scaled) MSE of the estimator of the regression parameter (with respect to the true $\beta^*$) as a function of the noise level added to $B_{XX}$ (top), $B_{XH}$ (center) and $\Sigma_{\epsilon_H}$ (bottom). Surprisingly, KL regression performs remarkably well (in particular the $L^1$-penalized version), showing a certain level of robustness to model misspecification. The x-axis scaling is not linear.}
\end{figure}
Given that in practice this bound is hard to compute, we empirically evaluate the robustness of KL regression though simulations. We perturb the connectivity matrix $B_{XX}$ by adding a random noise to it, which varies in strength: denoting $s$ the scaling of the random component, we have that $B_{e,XX} = B_{0,XX} + sZ_e$, where $B_{0,XX}$ is shared across environments, and $Z_e$ is a matrix with standard Gaussian entries. We run analogous simulations perturbing the connectivity matrix $B_{XH}$, and perturbing the latent variable covariance matrix $\Sigma_{\epsilon_H}$. We report the results of the simulation in Figure~\ref{fig:model-misspecification}, where we vary $s \in [0,5]$ (where $s=0$ corresponds to our baseline configuration described in section~\ref{appendix:simulation-environments}) and select a few values that are representative of the general behavior of the estimators. As expected, KL regression is robust to small perturbations in the connectivity matrix and outperforms most other methods. As the strength of perturbation in $B_{e,XX}$ increases KL regression performance degrades substantially. Surprisingly, Lasso KL regression seems to be very robust to perturbations in $B_{XH}$, $\Sigma_{\epsilon_H}$. A noticeable effect is that, under a slight perturbation in $B_{XX}$, Anchor Regression improves substantially and performs surprisingly well, although it is designed only to protect against a certain shift-interventions in some particular directions. We hypothesize that the changes due to model misspecification lead to a more conservative hedging in multiple directions, bringing Anchor regression closer to Causal Transfer Learning \citep{rojas2018invariant}, which in theory recovers $\beta^*$.

\section{Experiments}\label{section:section-4}

\subsection{Validation on Synthetic Data}

We extensively evaluate KL regression on synthetic data. We compare the following methods: \textbf{1) KL regression}, both Lasso KL regression and unregularized KL regression. \textbf{2) Invariant Causal Prediction (ICP)} \citep{peters2016causal} The implementation we used in the R package (\texttt{InvariantCausalPrediction}) detects that the assumptions for using ICP are not satisfied, and suggests running the \texttt{hiddenICP} version. \textbf{3) Anchor Regression} \citep{rothenhausler2018anchor} \textbf{4) Causal Dantzig} \citep{rothenhausler2019causal}  We run experiments with both the regularized and unregularized versions. \textbf{5) ComBat} \citep{johnson2007adjusting} \textbf{6) Averaging the regression coefficients}. Appendix~\ref{appendix:simulation-environments} provides a detailed description of the data generating process of our synthetic data simulations, and precise description of all competing methods and the following simulations.
\paragraph{Simulations with Varying Number of Environments.} Figure~\ref{fig:figure_a} plots the mean squared error (MSE) for the different estimators as a function of the number of samples for 2 and 4 environments. KL regression recovers $\beta^*$ with relatively few samples in several environments. Compared to alternatives in Table~\ref{table-methods}, KL regression is the only method that provably recovers $\beta^*$ at the population level. Compared to unregularized Causal Dantzig and hidden ICP (which behave similarly in 2 environments), KL regression is substantially more efficient as the number of environments increases.
\paragraph{Increasing Diversity Across Environments.} As suggested in Proposition~\ref{proposition:invertibility_S}, having diverse covariance matrices $\Sigma_{e,X}$ implies that $S_{\beta}$ is likely to be invertible and well-conditioned, and thus the asymptotic variance of our KL regression estimator will be lower.  Figure~\ref{fig:environment-diversity} shows how MSE changes as the environment-wise covariances $\Sigma_{e,\epsilon_{X}}$ differ more one of another (cf. Appendix~\ref{appendix:diversity} for further details). All methods improve with the diversity of the environments, but KL regression particularly benefits from diversity.

\paragraph{Additional Simulations.} We further explore the limitations and the robustness of our method in Appendix~\ref{appendix:simulations}, and summarize here our findings:
\begin{itemize}[noitemsep,nolistsep,leftmargin=*]
    \item We analyze the performance under various degrees of confounding  in Appendix~\ref{appendix:varying-confounding}. KL regression is robust as confounding gets stronger, while other methods degrade rapidly for higher confounding.
    \item We report in Appendix~\ref{appendix:sparsity} simulations where the sparsity level changes: regularized Causal Dantzig, and regularized KL regression which perfectly recover sparse vectors, degrade in settings where $\beta^*$ is not sparse. Unregularized KL regression works well in both sparse and non-sparse settings.
    \item Artificially splitting the data from one environment to generate new environments does not help KL regression, it actually worsens it as it does not increase environment diversity (Appendix~\ref{appendix:split}).
    \item KL regression does not require data to originate from a Gaussian SEM. We validate the previous results on data sampled from Student's t-distribution with different degrees of freedom and report the results in Appendix~\ref{appendix:student}.
\end{itemize}

\subsection{Recovering Gene Regulatory Networks: DREAM Challenge Datasets}
Validating KL regression on real-world datasets is challenging as there is often no ground truth causal mechanism. Gene Regulatory Networks (GNR) can provide us with some partial ground truth, indicating only whether two genes, a transcription factor (TF) and a regulated gene, are related or not (i.e. whether the corresponding $\beta^*_i$ is non-null). We therefore focus on Lasso KL regression and regularized Causal Dantzig for this experiment, as they automatically select features. This type of datasets also fits well with our model of environments (cf. equation~\ref{eq:model-of-environments}): deletions on upstream genes provide the sort of perturbation that shifts the global covariate distribution of gene expressions, without altering the true regulation mechanism of a specific target gene. We run our method on some of the DREAM4 and DREAM5 datasets released as part of an well-known online challenge. Importantly, we can also benchmark KL Regression against the leaderboard of the challenge via the AUPR metric (area under precision-recall curve) used in the challenge. We were limited in our experiments to datasets for which the environment-wise covariance matrix $\Sigma_{e,X}$ is invertible, and we expect to relax this assumption in future work. We generated different environments depending on whether genes had been knocked-down. This required assumption on $\Sigma_{e,X}$ does not hold whenever the number of gene expression measurements is smaller than the number of TFs. However, single cell gene sequencing has the potential to make KL regression widely applicable in future GNR inference tasks. We provide in Appendix~\ref{appendix:DREAM} additional information for these real-data experiments.

In the well-known DREAM5 challenge, our method \emph{ranked at the top of the leaderboard} by a substantial margin for the analyzed dataset. For DREAM4, KL regression still ranked in the top 40\% of submissions: however the semi-synthetic datasets from DREAM4 contained one single knock-down experiments for every single gene, making our model of environments less realistic (compared to DREAM5 where a single knock-down experiment had multiple replicates).

\begin{table}
  \caption{AUPR Performance on DREAM Challenge.}
  \label{table:DREAM}
  \centering
  \begin{tabular}{lllllll}
    \toprule
    Method    & DREAM 4 & DREAM 5    \\
    \midrule
    Lasso KL Regr. & 0.26 & \textbf{0.45} \\
    C. Dantzig (reg.) &  0.15   & 0.27  \\
    Top Submission & \textbf{0.37} & 0.31 \\
    \bottomrule
  \end{tabular}
\end{table}

\section{Discussion}
We provide a detailed mathematical foundation for KL regression, as well as extensive validation on synthetic and real data. KL has strong mathematical guarantees, and works better across many settings compared to existing methods. We derived our theoretical guarantee for KL regression in a standard SEM setting commonly used in causal inference, without assuming that data is sampled from any particular parametric distribution (i.e. we do not assume data is Gaussian). Empirically, KL regression performs well even with large model mismatch; we are very encouraged that it \emph{achieved the top score on a Dream5 dataset}, which has complex nonlinear interactions. KL regression introduces a new regression-flavored perspective to the invariance-based causal inference framework, opening the door to potentially leverage other regression-related concepts such as bootstrapping, model selection techniques (AIC, BIC criterion) and various regularization methods (beyond the current Lasso KL regression $L^1$-penalized version). Interesting future research directions will also consist in developing a non-linear version of KL regression for generalized linear models.

\subsubsection*{Acknowledgments}

The authors thank Dominik Rothenhäusler and Bryan He for their helpful comments, and Kailas Vodrahalli for his help with his implementation of invariant risk minimization. J.R.G. is supported by a Stanford Graduate Fellowship. J.Z. is supported by a Chan–Zuckerberg Biohub Investigator grant and National Science Foundation (NSF) Grant CRII 1657155.

\section{Introduction}

Drawing conclusions by analyzing data from a single source carries serious challenges. For example, if a gene expression-phenotype correlation is found to be statistically significant in one lab, does it imply that we will be able to replicate it in other settings? It is possible that some shared latent characteristic between the samples explains such correlation, rather than a true biological connection between the gene and the disease. Confounding is a pervasive problem, where unobserved variables that change across environments invalidate any result made on a single dataset~\citep{zou2014epigenome}. Leveraging multiple data sources---also called environments---can help to identify the relationship between variables that are common across diverse settings and hence are more likely to be unconfounded true relations~\citep{ben2010theory}. However, how best to integrate multiple environments and when would it work are very much open areas of research. This question has recently generated wide interest in the machine learning community due to its connections to causal inference and robustness to domain shift~\citep{ben2010theory,imbens2015causal,kouw2016feature}. Environments can correspond to different batches, labs where the experiment is done, perturbations, etc. One approach is to first fit a regression model on each environment separately to obtain the environment $e$ specific parameters $\beta_e$ between the covariates and the response. Then some meta-analysis can be performed to aggregate the different $\beta_e$, or at the level of p-values if those are available. In the presence of unknown confounders, however, such two-stage meta-analysis can produce erroneous results. Other models introduce the environment as part of the model, and fit all the data where the heterogeneity across environments is captured by methods such as mixed-effects models \citep{pinheiro2000linear}, time varying coefficients \citep{hastie1993varying,fan1999statistical}, etc. However, these models assume a structured transformation across environments, and cannot capture a potential radical change in the covariate structure from one environment to another due to an experimental manipulation. Another approach is to focus on the invariant element of the model that we want to identify, and look for such invariant signal across environments \citep{peters2016causal}. We follow this line of work, where techniques differ from each other through their scope of application and their goal: some estimate the invariant vector of coefficients, while others provide a vector of coefficients that solves a minimax risk problem over the different environments. We will further discuss the relevant literature at the end of Section~\ref{section:section-2}.
\paragraph{Our Contributions.} We present a novel method for identifying invariant factors across multiple environments that we call \emph{Kullback-Leibler regression} (KL regression) focusing on identifying the covariates that have a causal effect on the response. Unlike other approaches, KL regression simultaneously allows for the existence of \emph{latent confounders}, is valid in settings where the different environments arise from distribution shifts in the observed covariates \emph{and} shifts in the response, and \emph{provably recovers the invariant causal vector} of coefficients. We provide asymptotic convergence results and demonstrate its improved performance compared to standard approaches in extensive experiments. Additionally, we introduce a \emph{Lasso KL regression} estimator, where a sparsity inducing penalty is added to the problem defining our estimator. In practice, this variant under a sparsity assumption on $\beta^*$ shows remarkable stability and robustness to model misspecification. We begin in Section~\ref{section:section-2} by presenting the problem of linear regression estimation with confounding due to latent variables based on linear Structural Equation Models (SEM) and describing the relevant literature. In Section~\ref{section:section-3} we describe KL regression in its general form, then present the required parametrization steps tailored to our SEM model that lead to our KL regression estimator (both sparse and non-sparse versions). We thoroughly validate the theoretical results empirically in Section~\ref{section:section-4} and compare our method to other commonly used techniques. KL regression achieves the top score for a DREAM5 challenge.

\section{Confounding with Structural Equation Models}\label{section:section-2}

\subsection{Generating Multiple Environments Via Linear Structural Equation Models}

Latent variables that introduce a confounding effect on the observed data is a common framework to represent spurious effects. In a linear setting, we may assume that the unconfounded vector of coefficients---the target of interest---which relates the response to the observed covariates represents an invariant structure that has a causal interpretation \citep{rojas2018invariant}. However, estimators like least squares regression may be biased due to the presence of latent confounding effects. We consider the case where the data follows a linear Structural Equation Model (SEM) \citep{bollen1989structural,kaplan2008structural}, which defines a distribution $\pi$ over the observed covariates $\boldsymbol{X}\in \mathbb{R}^D$, the hidden covariates $\boldsymbol{H}\in \mathbb{R}^Q$ and the observed response $Y\in \mathbb{R}$ as a solution to the system of equations:
\begin{align}
    \begin{cases}{}\label{equation-linear-sem-individual}
    H_q = \epsilon_{H,q} \quad \text{for $1\leq q \leq Q$}
    \\ X_d  = \sum_{1\leq d' \leq D} [B_{XX}]_{dd'}X_{d'}\! +\! \sum_{1\leq q\leq Q}[B_{XH}]_{dq}H_q
    \\ \hspace{2cm} + \; \epsilon_{X,d} \qquad \text{for $1\leq d \leq D$}
    \\ Y = \sum_{1\leq d \leq D} \beta^*_{d}X_d + \sum_{1\leq q \leq Q} \eta_{0,q} H_q + \epsilon_Y
    \end{cases}
\end{align}{}
where the coefficients are given by the \emph{connectivity} matrices $B_{XX}\in\mathbb{R}^{D\times D}$, $B_{XH}\in\mathbb{R}^{D\times Q}$ and the vectors $\beta^*\in \mathbb{R}^D$, $\eta_0 \in \mathbb{R}^Q$. We define $\epsilon := (\epsilon_X,\epsilon_Y,\epsilon_H) \in \mathbb{R}^{D+Q+1}$ as a random variable with finite second moments, where its three components are pairwise independent random vectors (one could equivalently represent the latent variable effect implicitly with a SEM defined only for $(\boldsymbol{X}, Y)$ where $\epsilon_X$ and $\epsilon_Y$ are correlated). We can write this model in a matrix form:
\begin{equation}\label{equation-linear-sem}
    \begin{pmatrix}
    \boldsymbol{X} \\ Y \\ \boldsymbol{H}
    \end{pmatrix} \!=\! \mathbf{B} \!\begin{pmatrix}
    \boldsymbol{X} \\ Y \\ \boldsymbol{H}
    \end{pmatrix} \!+ \!\begin{pmatrix}
    \epsilon_X \\ \epsilon_Y \\ \epsilon_H
    \end{pmatrix}\!; \; \mathbf{B}\! = \!\begin{pmatrix}
    B_{XX} & 0 & B_{XH}
    \\ \beta^{*T} & 0 & \eta_0^T
    \\ 0 & 0 & 0
    \end{pmatrix}
\end{equation}
where $(\Sigma_{\epsilon_X},\Sigma_{\epsilon_Y}, \Sigma_{\epsilon_H})$ are the covariance matrices of $\epsilon_X,\epsilon_Y,\epsilon_H$ respectively. We assume that $(I -\mathbf{B})$ is invertible, and thus the solution $\pi$ to the system is characterized by $\mathbf{B}$ and the distribution of $\epsilon$ through $\pi := \mathcal{D}\big((I -\mathbf{B})^{-1}\epsilon\big)$ where we denote by $\mathcal{D}(U)$ the distribution of the random variable $U$. The SEM model presented above can be reformulated with a probabilistic graphical model, where the variables (both observed or latent) correspond to the nodes and the edges are determined by the connectivity matrix $\mathbf{B}$. Our objective is to identify the parameter $\beta^*$ that captures the causal dependence between the observed covariates and the response, and especially to identify the non-zero entries: this is the problem of causal discovery. In the context of graphical models we want to identify the parent nodes of the response node, especially so when $\beta^*$ is assumed to be sparse and it is of importance to identify the few causal factors. Based only on $\boldsymbol{X},Y$ samples from $\pi$, the least squares regression vector of $Y$ on the covariates $\boldsymbol{X}$ is equal to $\beta^*$ plus an additive confounding term (we provide an explicit formulation in Prop.~\ref{proposition:covariance-parameters}). Our method recovers $\beta^*$ by leveraging information from multiple environments, and under sparsity assumptions we propose a regularized version to specifically identify relevant covariates. Environments in practice correspond to different data sources: if the collected data are patient medical records, different hospitals may correspond to different environments. When analyzing cellular genetic pathways, measurements from different cell cultures where different genes have been knocked out correspond to data from different environments. The crucial \emph{invariance} assumption is that the mechanism generating the response variable $Y$, encoded in the parameters of the structural equation of $Y$, is constant across environments. In particular, those parameters relating the response to the observed covariates, encoded by $\beta^*$, are invariant. We observe data from $E\geq 1$ different environments, which are characterized by their underlying respective distributions over $\mathbb{R}^D \times \mathbb{R}$. We use the subscript $(\boldsymbol{X}_e, Y_e)$ to refer to the observed variables in environment $e \in [E]:=\{1,\dots,E\}$ (and no longer to the covariate index as in \ref{equation-linear-sem-individual}). Our main assumption defines how models relate to each other -- i.e. which elements in equation~\ref{equation-linear-sem} may change--, which we call a \emph{model of environments}. This represents our beliefs on how interventions may perturb the initial observational data generating distribution. In our work, we only allow for $(\Sigma_{e,\epsilon_X}, \Sigma_{e,\epsilon_Y})$ to depend on the environment, so that each $e$-th environment distribution $\pi_e$ is the solution of the linear SEM:
\begin{equation}\label{eq:model-of-environments}
    \begin{pmatrix}
    \boldsymbol{X}_e \\ Y_e \\ \boldsymbol{H}
    \end{pmatrix} = \mathbf{B} \begin{pmatrix}
    \boldsymbol{X}_e \\ Y_e \\ \boldsymbol{H}
    \end{pmatrix} + \begin{pmatrix}
    \epsilon_{e,X} \\ \epsilon_{e,Y} \\ \epsilon_H
    \end{pmatrix}
\end{equation}
This model of environments can be understood as transformations from a base SEM via ``shift interventions'' on variables $\boldsymbol{X},Y$ \citep{rothenhausler2015backshift, rothenhausler2019causal}, where an environment-specific additive noise term is added to $\epsilon$ in the SEM: the distribution of the $e$-th environment is given by $\pi_e := \mathcal{D}\big((I -\mathbf{B})^{-1}(\epsilon + v_e)\big)$ where $v_e = (v_{e,X}, v_{e,Y}, 0)$ is a centered random variable in $\mathbb{R}^{D+Q+1}$ where only the first $D + 1$ coordinates are non-zero, independent of $\epsilon$. These interventions are also called ``parametric'' interventions (on $\boldsymbol{X}$ and $Y$) \citep{eberhardt2007interventions} or ``soft'' interventions \citep{eaton2007exact}. Constraining $v_{e,Y} = 0$ would correspond to shift interventions only on $\boldsymbol{X}$ (cf. next paragraph and ``Shift $\boldsymbol{X}$'', ``Shift $Y$'' in Table~\ref{table-methods} for further details). These contrast with ``stuctural'' or ``perfect'' interventions that also allow for modifications in the matrix $\mathbf{B}$, which occur for example with do-interventions \citep{judea2000causality}. A graphical representation of this model of environments is presented in Figure~\ref{fig:DAG}: different environments are obtained via changes in the distributions of $\epsilon_X, \epsilon_Y$ via shift interventions, represented via red circles. The remaining structure of the model is unchanged across environments. Defining a model of environments carries a trade-off: the smaller the set of elements that change across environments is, the easier it is to find invariant quantities across environments, but conversely that corresponds to stronger assumptions that will restrict the scope of applications. From now on we no longer explicitly mention the latent variables and hence $\pi_e$ refers to the marginal distribution over $\boldsymbol{X}_e, Y_e$.
\begin{figure}
    \centering
    \includegraphics[width=.4\linewidth]{figure_dag.png}
    \caption{Graphical representation of the process generating multiple environments.}
    \label{fig:DAG}
\end{figure}
\begin{table}
  \caption{Multi-environment analysis methods and their scope of application.}
  \label{table-methods}
  \centering
  \begin{tabular}{llllll}
    Method     & Shift $\boldsymbol{X}$ & Shift $Y$ & Confounding & Recovers $\beta^*$\\
    \midrule
    KL Regression  & Yes & Yes & Yes & Yes \\
    Invariant Causal Prediction  & Yes & No &  No & Yes \\
    Causal Dantzig  & Yes & Yes & Yes & Yes\\
    Invariant Risk Minimization  & Yes & Yes & No & No \\
    Anchor Regression  & Yes & Yes & Yes & No\\
    Causal Transfer Learning & Yes & No & No & Yes \\
    \bottomrule
  \end{tabular}
\end{table}
\paragraph{Related Literature}
The idea that true causal mechanisms are \emph{invariant} across different experimental settings \citep{peters2016causal, buhlmann2018invariance} has generated a lot of interest recently in the causal inference community. Based on samples collected from different ``environments'', it is possible to identify covariates whose effect on the response variable is universal, which are defined as causal. Solving the causal discovery problem problem with purely observational data is not possible, and a popular line of work in causal inference which uses Bayesian networks to represent causal models \citep{pearl1995causal,spirtes2000causation} relies on experimental or interventional data recover the graph structure. Pearl's do-calculus \citep{judea2000causality} identifies causal effects based on observational data but assuming perfect knowledge of the graphical structure of the Bayesian network. Other methods assume different data sources are generated through a change in the structure of the Bayesian network, and under some assumptions the causal graph can be recovered \citep{eberhardt2010combining,hyttinen2012learning}. This requires detailed knowledge on the perturbations that originate the distinct data sources. On the other hand, under a more limited knowledge of the mechanism that generates different data sources, the different distributions can be modeled with ``uncertain'' interventions \citep{korb2004varieties,eaton2007exact,eberhardt2010combining,tian2013causal}, for which we do not have a clear knowledge, and is the setting we assume in this work. Also, instead of estimating the complete graph structure \citep{spirtes2000causation,chickering2002optimal,hauser2012characterization,rothenhausler2015backshift}, we rather focus on identifying the sub-graph that connects the response variable to the other covariates. This leads us to formulating a linear SEM as in equation~\ref{equation-linear-sem}, and we now present some methods that solve this problem (we provide in Table~\ref{table-methods} a quick summary of how these relate to each other). The main modeling assumption for any procedure leveraging multiple environments is the definition of how such environments relate one to another, i.e. what the \emph{model of environments} is. Usually, procedures that focus on recovering $\beta^*$ (using our linear SEM framework) usually do not allow for any modification of the structural equation of $Y$: this is the basis of Invariant Causal Prediction (ICP) \citep{peters2016causal} and further methods based on the same principle \citep{heinze2018invariant,pfister2019invariant}. This body of work is based on the assumption that the conditional distribution of the response on a subset of covariates that are termed ``causal'' is constant across environments: i.e. if $S^*\subset[D]$ are the set of indices of such causal variables, denoting $\boldsymbol{U}_{A}:= (U_a)_{a\in A}$ for a vector $\boldsymbol{U}$, and $A$ subset of indices, then
\begin{equation}\label{invariant-assumption}
    Y_e|\boldsymbol{X}_{e,S^*} \stackrel{d}{=}Y_f|\boldsymbol{X}_{f,S^*} \quad \forall e, f\in [E]
\end{equation}{}
This assumption, translated into our SEM model, implies that the row corresponding to $Y$ is invariant across environments (i.e. no ``shift $Y$'' interventions in Table~\ref{table-methods}), hence our assumptions are weaker in that aspect. This weakened assumption for the model of environments is in line with what methods such as Invariant Risk Minimization (IRM) \citep[Definition 6, p.10]{arjovsky2019invariant} assume and argue is needed in real-world applications: changes in the noise level of the response across environments. The structure relating the covariates to the response---within the connectivity matrix $\mathbf{B}$---is invariant in any case (indeed, $\beta^*$ is defined because of such invariance). However, our method fundamentally differs from those other methods that allow shifts in $Y$, as their goal is not to provide an estimate of $\beta^*$ (i.e. ``Recovers $\beta^*$'' in Table~\ref{table-methods}), but instead solve distributional robustness problems:
\begin{align*}
    \hat{\beta} := \min_{b \in \mathbb{R}^D} \max_{e\in [E]} \mathbb{E}[(Y_e - \boldsymbol{X}_e b)^2]
\end{align*}{}
These estimators minimize the worst-case prediction risk, and are related to adversarial training \citep{sinha2017certifying,gao2017wasserstein}. Whenever the set of environments contains all possible modifications of the structural equations except that of the target, the output of such optimization problem is $\beta^*$, as described in Causal Transfer Learning (CTL) \citep{rojas2018invariant}. Besides this setting, theoretical guarantees of those estimators will usually focus on hedging against a particular subset of interventions instead of estimating $\beta^*$, as in Anchor Regression \citep{rothenhausler2018anchor}. Finally, going back to those methods that focus on recovering $\beta^*$, most of them do not allow for confounding (i.e. ``Confounding'' in Table~\ref{table-methods}). The invariance assumption~(\ref{invariant-assumption}) fails whenever a latent variable $\boldsymbol{H}$ simultaneously affects $Y$ or $\boldsymbol{X}$, or whenever its distribution changes across environments \citep[Example 1, p.6]{christiansen2020switching}. Several of the previously mentioned techniques fail to recover $\beta^*$ under confounding, and recently proposed methods such as \cite{christiansen2020switching} restrict the latent variables to be discrete (and do not allow for changes in the noise level $\Sigma_{\epsilon_Y}$ of the response). KL regression shares many characteristics with Causal Dantzig \citep{rothenhausler2019causal}, which has an unregularized and a sparsity-inducing regularization versions. These two methods are related in a way that is analogous to the relationship between the Lasso linear regression \citep{tibshirani1996regression} and the Dantzig Selector \citep{candes2007dantzig}. Causal Dantzig without regularization leads to an estimator that is similar to ours in its scope of applications. However, as we show in the experiments, KL regression outperforms unregularized Causal Dantzig in terms of stability. Using a sparsity-inducing regularization in Causal Dantzig greatly improves its behavior and matches KL regression performance. However, this regularized Causal Dantzig \emph{needs} sparsity in $\beta^*$, otherwise it does not properly recover it. KL regression does not require sparsity to perform well, although the regularized KL regression version substantially outperforms any other method in sparse settings.

\subsection{Invariant Elements for our Model of Environments}

Assume that variables $\epsilon$ defined in the SEM (cf. equation~\ref{equation-linear-sem}) have finite second moments. Denote by $\Sigma_e^{joint}$ the covariance matrix of the observed variables in environment $e$, which we assume has full rank.
\begin{align*}
    \Sigma_e^{joint} := \begin{pmatrix}
    \Sigma_{e,X} & \Sigma_{e,XY}
    \\ \Sigma_{e,XY}^T & \Sigma_{e,Y}
    \end{pmatrix} := \text{Cov}\begin{pmatrix}
    \boldsymbol{X}_e
    \\Y_e
    \end{pmatrix}
\end{align*}
Our model of environments assumption, that relies on a linear SEM, leads to a specific structure of these covariances across environments. Our method recovers $\beta^*$ only via the analysis of $\big(\Sigma_e^{joint}\big)_{e \in [E]}$, i.e. the second moments, regardless of the specific distribution of $\epsilon$. As a first step, we observe that in any given environment we can decouple the confounding bias from $\beta^*$ in the regression vector estimated by least squares.
\begin{proposition}\label{proposition:covariance-parameters}
Under our model of environments (cf. equation~\ref{eq:model-of-environments}) we have:
\begin{align*}
    & \Sigma_{e,X} =  C\Big(\Sigma_{e,\epsilon_{X}} + B_{XH}\Sigma_{\epsilon_H}B_{XH}^T\Big)C^T
    \\ & \Sigma_{e,XY} = \Sigma_{e,X}\beta^* + \eta^*
\end{align*}{}
where $C := (I -B_{XX})^{-1}$ and $\eta^* := CB_{XH}\Sigma_{\epsilon_{H}}\eta_0$.
\end{proposition}
We prove this result in Appendix~\ref{proof:proposition:covariance-parameters}. In particular, $\Sigma_{e,X}$ depends on the environment $e$ through $\Sigma_{e, \epsilon_X}$, and $\beta^*, \eta^*$ do not depend on $e$. Given data from the $e$-th environment, the least squares estimator $\hat{\beta}_e$ estimates the population regression coefficient given by $\beta_e := \Sigma_{e,X}^{-1}\Sigma_{e,XY} = \beta^* + \Sigma_{e,X}^{-1}\eta^*$. The confounding due to the additional term $\Sigma_{e,X}^{-1}\eta^*$ implies that, based on purely observational data (i.e. data from one single environment), least squares is biased for $\beta^*$.

\section{Kullback-Leibler Regression}\label{section:section-3}

\subsection{General Framework}
We reformulate the joint covariance of environment $e$ through the regression coefficient $\beta_e$ and the residual variance given by \mbox{$\sigma_e^2 := \Sigma_{e,Y} - \beta_e^T\Sigma_{e,X}\beta_e \in \mathbb{R}$}:
\begin{align*}
    \Sigma_e^{joint} =
    \begin{pmatrix}
    \Sigma_{e,X} & \Sigma_{e,X}\beta_e
    \\ \beta_e^T\Sigma_{e,X} & \sigma_e^2 + \beta_e^T\Sigma_{e,X}\beta_e
    \end{pmatrix}
\end{align*}
Let $\Pi$ be a mapping that we call the \emph{regression function} associated to the $e$-th environment. Its input consists on the pair $(\Sigma_{e,X}, \sigma_e^2)$ --that we call \emph{regressors}, and the \emph{regression coefficient} $\theta$:
\begin{equation}\label{distribution-map}
    \Pi:\begin{cases}{}\big(\mathbb{R}^{D\times D}\times\mathbb{R}_+\big)\times\mathbb{R}^D \rightarrow \mathbb{R}^{(D+1)\times (D+1)}
    \\(\Sigma_{e,X}, \sigma_e^2), \theta \mapsto
    \begin{pmatrix}
    \Sigma_{e,X} & \Sigma_{e,X}\theta
    \\ \theta^T\Sigma_{e,X} & \sigma_e^2 + \theta^T\Sigma_{e,X}\theta
    \end{pmatrix}
    \end{cases}
\end{equation}{}
The output is a positive semi-definite matrix such that $\Sigma_e^{joint} = \Pi\big((\Sigma_{e,X}, \sigma_e^2),\beta_e\big)$. For notation simplicity, let $\Pi_e(\theta):=\Pi\big((\Sigma_{e,X}, \sigma_e^2),\theta\big)$. Given a distance $\delta(\cdot, \cdot)$ over the set of positive semi-definite matrices, we can construct a loss function over $\theta$:
\begin{align}\label{loss-eq-simple}
    & \mathcal{L}(\theta) := \sum_{e\in [E]} \delta\big(\Pi_e(\theta), \Sigma_e^{joint}\big)
\end{align}{}
Although the distributions $\pi_e$ are not Gaussian, we can follow a quasi-maximum likelihood approach and define $\delta$ as the Kullback-Leibler divergence between two multivariate Gaussian variables with covariances $\Pi_e(\theta)$ and $\Sigma_e^{joint}$. This leads to the following environment-wise Mahalanobis distance in $\theta$ parametrized by the regressors $\Sigma_{e,X},\sigma_e^2$.
\begin{align*}
    \delta\big(\Pi_e(\theta), \Sigma_e^{joint}\big) := KL\Big(\mathcal{N}\big(\Pi_e(\theta)\big)\Big\Vert \mathcal{N}\big(\Sigma_e^{joint}\big)\Big)
\end{align*}
We provide in Appendix~\ref{appendix:regression-analogy} additional insights on this choice of distance and why our method can be understood as performing linear regression in a space of distributions. We then construct an estimator $\hat{\theta}$ by minimizing the corresponding empirical risk by plugging-in the empirical counterparts of $\Sigma_{e,X}, \sigma_e^2,\beta_e$ (denoted by a hat):
\begin{align}\label{loss-eq-simple-empirical}
    & \hat{\mathcal{L}}(\theta) := \sum_{e\in [E]} \delta\Big(\Pi\big((\hat{\Sigma}_{e,X},\hat{\sigma}_e^2), \theta\big), \hat{\Sigma}_e^{joint}\Big)
    \\ & \hat{\theta} := \argmin_{\theta\in \mathbb{R}^D} \hat{\mathcal{L}}(\theta)
\end{align}{}
This particular optimization problem in $\theta$ is convex, and has a closed-form solution. This procedure looks for a parameter $\theta$ that makes the environment-wise distributions as compatible as possible with a \emph{model of distributions}. Here, by minimizing $\mathcal{L}(\theta)$, we implicitly assume that there is a coefficient $\theta^*$ that defines for each environment $e$ a distribution $\pi_e^*$ whose covariance matrix is given by $\Pi_e(\theta^*)$. The observed environment distributions $\pi_e$ are then just ``unstructured'' perturbations of such $\pi_e^*$, that is there is no relationship between the regressors (i.e. $(\Sigma_{e,X}, \sigma_e^2)$) and the perturbation of the coefficient $\beta_e$ with respect to $\theta^*$.  We here understand the role of the structural causal model: individually, each of the environment distributions do not allow us to recover the invariant vector. It is only through the joint analysis of the environments that we identify it, and thus one needs to account for this causal structure in the parametrization in the loss. The problem formulation above is not adapted to our setting where the confounding bias in $\beta_e$ is environment-dependent, related to the regressors.

\subsection{Kullback-Leibler Regression for our Model of Environments}

\begin{figure*}[t]
\centering
\includegraphics[width=0.9\linewidth]{figure_a.pdf}
\caption{\textbf{Baseline comparison} Plots of the (log-scaled) MSE of the estimator of the regression parameter (with respect to the true $\beta^*$) as a function of the number of samples. The $L_1$ regularized KL regression consistently achieved the lowest MSE.}
\label{fig:figure_a}
\end{figure*}

CV the choice of lambda for KL regression to get also a decrease with number of samples.

Whenever our environment distributions $\pi_e$ come from a latent variable model as in the previous section (cf. equation~\ref{eq:model-of-environments}), minimizing the parameter $\theta$ in the mapping $\Pi_e(\theta)$ in \eqref{distribution-map} that defines our loss function does not yield the desired result, basically because there is no $\theta^*$ that is equal to all $\beta_e$ simultaneously. However, as we saw in Proposition~\ref{proposition:covariance-parameters}, the expression of the confounded linear regression coefficient $\beta_e$ in environment $e$ can be decomposed into two terms $\beta_e = \beta^* + \Sigma_{e,X}^{-1}\eta^*$, where $\eta^*$ no longer depends on $e$. As such, we can modify the loss in \eqref{loss-eq-simple} by expressing $\theta = \beta + \Sigma_{e,X}^{-1}\eta$, and then minimize the reparametrized loss over $(\beta,\eta)$, which is still jointly convex in both arguments:
\begin{align}\label{loss-DR}
    & \mathcal{L}(\beta, \eta) \!:=\! \sum_{e\in [E]}\! \delta\Big(\!\Pi_e(\beta + \Sigma_{e,X}^{-1}\eta))\Vert \Sigma_e^{joint}\!\Big)
    \\ & (\tilde{\beta}^{KL}, \tilde{\eta}) := \argmin_{\beta, \eta \in \mathbb{R}^D} \mathcal{L}(\beta, \eta)
\end{align}{}
Importantly, the reparametrization only requires knowledge of the covariance of $\boldsymbol{X}_e$, which can be estimated. This is the key insight of our method: because of the invariant structure of the model of environments and the judicious choice of distance in the space of covariances, we can construct a tractable loss that is minimized at the invariant vector. We denote by $\tilde{\beta}^{KL}$ the solution to the above optimization problem, and look for conditions that guarantee uniqueness of the solution, thus identifying $\beta^*$.
\begin{proposition}\label{proposition:solution-DR}
Provided that the following matrix is invertible
\begin{align}
    S_{\beta}:= \Big(\sum_{e\in [E]}\frac{\Sigma_{e,X}^{-1}}{\sigma_e^2}\Big)\Big(\sum_{e\in [E]}\frac{\Sigma_{e,X}}{\sigma_e^2}\Big) - \Big(\sum_{e\in [E]}\frac{1}{\sigma_e^2}\Big)^2I_D
\end{align}
the unique solution to the KL regression problem~\ref{loss-DR} has a closed form given by
\begin{align*}
    \tilde{\beta}^{KL} = S_{\beta}^{-1}\Bigg(&\Big(\sum_{e\in [E]}\frac{\Sigma_{e,X}^{-1}}{\sigma_e^2}\Big)\Big(\sum_{e\in [E]}\frac{\Sigma_{e,X}\beta_e}{\sigma_e^2}\Big)  -  \Big(\sum_{e\in [E]}\frac{1}{\sigma_e^2}\Big)\Big(\sum_{e\in [E]}\frac{\beta_e}{\sigma_e^2}\Big)\Bigg)
\end{align*}{}
Assuming that our environment distributions $\pi_e$ follow the SEM model above, then we have $\tilde{\beta}^{KL} = \beta^*$ (at the population level, i.e. assuming we can perfectly estimate the joint covariance under $\pi_e$).
\end{proposition}{}
\begin{proof}\label{proof:proposition:distribution-DR}
We prove here Proposition~\ref{proposition:distribution-DR}. Again, we drop the subscript $X$ from the covariance matrix $\Sigma_{e,X}$ for notation simplicity. Let the following matrices in $\mathbb{R}^{dE\times d}$:
\begin{align*}
    \mathbb{U} =  \begin{pmatrix}{}
    \frac{\hat{\Sigma}_1}{\sigma_1}
    \\\vdots
    \\ \frac{\hat{\Sigma}_E}{\sigma_E}
    \end{pmatrix}; \quad
    \mathbb{V} =  \begin{pmatrix}{}
    \frac{\hat{\Sigma}^{-1}_1}{\sigma_1}
    \\\vdots
    \\ \frac{\hat{\Sigma}^{-1}_E}{\sigma_E}
    \end{pmatrix} ; \quad
    \mathbb{I} =  \begin{pmatrix}{}
    \frac{I_D}{\sigma_1}
    \\\vdots
    \\ \frac{I_D}{\sigma_E}
    \end{pmatrix}
\end{align*}{}
We have that
\begin{align*}
    \hat{S}_{\beta} = & \mathbb{I}^T\mathbb{V}\mathbb{U}^T\mathbb{I} - \mathbb{I}^T\mathbb{I}\mathbb{I}^T\mathbb{I}
    \\ \hat{\beta} = & (\mathbb{I}^T\mathbb{V}\mathbb{U}^T\mathbb{I} - \mathbb{I}^T\mathbb{I}\mathbb{I}^T\mathbb{I})^{-1}(\mathbb{I}^T\mathbb{V}\mathbb{U}^T - \mathbb{I}^T\mathbb{I}\mathbb{I}^T)\hat{\mathbb{B}}
    \\ = & \mathbb{W}\hat{\mathbb{B}}
\end{align*}{}
where
\begin{align*}
    \mathbb{W} = & (\mathbb{I}^T\mathbb{V}\mathbb{U}^T\mathbb{I} - \mathbb{I}^T\mathbb{I}\mathbb{I}^T\mathbb{I})^{-1}(\mathbb{I}^T\mathbb{V}\mathbb{U}^T - \mathbb{I}^T\mathbb{I}\mathbb{I}^T) \in \mathbb{R}^{d\times dE}
    \\ \hat{\mathbb{B}} = & \begin{pmatrix}{}
    \frac{\hat{\beta}_1}{\sigma_1}
    \\ \vdots
    \\ \frac{\hat{\beta}_E}{\sigma_E}
    \end{pmatrix} \in \mathbb{R}^{dE}; \qquad
    \mathbb{B} =  \begin{pmatrix}{}
    \frac{\beta_1}{\sigma_1}
    \\ \vdots
    \\ \frac{\beta_E}{\sigma_E}
    \end{pmatrix} \in \mathbb{R}^{dE}
\end{align*}
so that each stacked vector in $\hat{\mathbb{B}}$ is independent of the others, and for every $e\in [E]$, we have that $\hat{\beta}_e = \beta_e + \frac{\sigma_e}{\sqrt{n_e}}\hat{\Sigma}^{-1/2}_{e,X}z_e$ where $(z_e)_{1\leq e \leq E}$ are independent standard Gaussians. We therefore get that
\begin{align*}
    \hat{\mathbb{B}} = & \begin{pmatrix}{}
    \frac{\beta_1}{\sigma_1}
    \\ \vdots
    \\ \frac{\beta_E}{\sigma_E}
    \end{pmatrix} + \begin{pmatrix}
    \frac{1}{\sqrt{n_1}}\hat{\Sigma}_1^{-1/2}z_1
    \\ \vdots
    \\ \frac{1}{\sqrt{n_E}}\hat{\Sigma}_E^{-1/2}z_E
    \end{pmatrix} = \mathbb{B} + \mathbb{Z}
\end{align*}{}
Therefore $\hat{\beta}^{KL} = \mathbb{W}\hat{\mathbb{B}} = \mathbb{W}\mathbb{B} + \mathbb{W}\mathbb{Z}$. But we have that $\mathbb{W}\mathbb{B} = \mathbb{E}[\hat{\beta}^{KL}|(X_e)_e]$, and $\mathbb{W}\mathbb{Z}$ is a multivariate Gaussian, whose covariance is given by :
\begin{align*}
    V = & \mathbb{W}\begin{pmatrix}{}
    \frac{1}{n_1}\hat{\Sigma}^{-1}_1 & 0&0 &\cdots & 0
    \\0 & \frac{1}{n_2}\hat{\Sigma}^{-1}_2& 0 &\cdots & 0
    \\ 0 & 0 & \ddots& \cdots & 0
    \\ \vdots & \vdots&&\vdots&
    \\ 0 & \cdots & 0& 0 &\frac{1}{n_E}\hat{\Sigma}_E^{-1}
    \end{pmatrix}\mathbb{W}^T
    \\ = & \sum_{e=1}^E \frac{1}{n_e}A_e
\end{align*}{}
for some symmetric matrix $A_e$. Now, if all the environments have the same sample size, $n_e = n$, then
\begin{align*}
    V & = \frac{1}{n}\hat{S}_{\beta}^{-1}\mathbb{I}^TA\mathbb{I}\hat{S}_{\beta}^{-1,T}
    \\ A & = (\mathbb{V}\mathbb{U}^T - \mathbb{I}\mathbb{I}^T)\begin{pmatrix}{}
    \hat{\Sigma}^{-1}_1 & 0&0 &\cdots & 0
    \\0 & \hat{\Sigma}^{-1}_2& 0 &\cdots & 0
    \\ 0 & 0 & \ddots& \cdots & 0
    \\ \vdots & \vdots&&\vdots&
    \\ 0 & \cdots & 0& 0 &\hat{\Sigma}_E^{-1}
    \end{pmatrix}(\mathbb{U}\mathbb{V}^T - \mathbb{I}\mathbb{I}^T)
    \\ & = \Big[\frac{\hat{\Sigma}_i^{-1}\hat{\Sigma}_j - I}{\sigma_i \sigma_j} \Big]_{i,j}\Big[ \hat{\Sigma}_j^{-1}1_{j=k}\Big]_{j,k}\Big[\frac{\hat{\Sigma}_k\hat{\Sigma}_l^{-1} - I}{\sigma_k \sigma_l} \Big]_{k,l}
    \\ & = \Big[ \sum_{j=1}^E \frac{\hat{\Sigma}_i^{-1}\hat{\Sigma}_j\hat{\Sigma}_l^{-1} - \hat{\Sigma}_i^{-1} - \hat{\Sigma}_l^{-1} + \hat{\Sigma}_j^{-1}}{\sigma_i\sigma_j^2\sigma_l}\Big]_{i,l}
\end{align*}
\begin{align*}
    \mathbb{I}^TA\mathbb{I} & = \sum_{i,j,l = 1}^E \frac{\hat{\Sigma}_i^{-1}\hat{\Sigma}_j\hat{\Sigma}_l^{-1} - \hat{\Sigma}_i^{-1} - \hat{\Sigma}_l^{-1} + \hat{\Sigma}_j^{-1}}{\sigma_i^2\sigma_j^2\sigma_l^2}
    \\  &=\Big(\sum_{e\in [E]}\frac{\hat{\Sigma}_e^{-1}}{\sigma_e^2}\Big)\Big(\sum_{f\in [E]}\frac{\hat{\Sigma}_f}{\sigma_f^2}\Big)\Big(\sum_{e\in [E]}\frac{\hat{\Sigma}_e^{-1}}{\sigma_e^2}\Big) - \Big(\sum_{e\in [E]}\frac{1}{\sigma_e^2}\Big)^2\Big(\sum_{e\in [E]}\frac{\hat{\Sigma}_e^{-1}}{\sigma_e^2}\Big)
    \\  &=\Big(\sum_{e\in [E]}\frac{\hat{\Sigma}_e^{-1}}{\sigma_e^2}\Big)\hat{S}_{\beta}^{T}
\end{align*}{}
Therefore we conclude
\begin{equation*}
V = \frac{1}{n} \hat{S}_{\beta}^{-1}\Big(\sum_{e\in [E]}\frac{\hat{\Sigma}_e^{-1}}{\sigma_e^2}\Big)
\end{equation*}
\end{proof}{}

The conditions for this proposition are mild: instead of assuming specific types of interventions for each individual covariate to guarantee identifiability of $\beta^*$ \citep{hyttinen2012learning, peters2016causal}, we just require diverse enough environments in terms of the covariates covariance. Identifiability of $\beta^*$ holds as long as $S_{\beta}$ is invertible, regardless of the actual distribution of $\epsilon$. Under which conditions is the matrix $S_{\beta}$ invertible? In particular, if only one environment is available then $S_{\beta}$ is not invertible: we can not make any statement based on purely observational data. However, KL regression can output $\tilde{\beta}^{KL}$ based on only two different environments under some conditions, where all covariates are simultaneously perturbed in one of the two environments. We now reformulate a result from \cite{bhagwat_subramanian_1978} under our framework, which we prove in Appendix~\ref{proof:proposition:invertibility_S} for completeness.
\begin{proposition}[\cite{bhagwat_subramanian_1978}]\label{proposition:invertibility_S}
A sufficient condition for $S_{\beta}$ to be invertible is that there exists $e^*\in [E]$ such that the following matrix is invertible
\begin{align*}
\Sigma_{e^*,X}^{-1}  - \Big(\sum_{e\in [E]}\frac{1}{\sigma_e^2}\Big)^{-1}\Big(\sum_{e\in [E]}\frac{\Sigma_{e,X}^{-1}}{\sigma_e^2}\Big)
\end{align*}
\end{proposition}{}
Assume for simplicity that the residual variances are equal to 1 (corresponding to independent measurement errors), if $\Sigma_{2,X}^{-1}\Sigma_{1,X} - I_D$ is invertible then $S_{\beta}$ is invertible: consider the following example where such scenario may arise. A scientist can experimentally down-regulate simultaneously all (measured) expressions of genes of interest, without interfering with unobserved latent genes. Covariances of gene expressions will be on different scales (eg. $\Sigma_{1,X} \succ \Sigma_{2,X}$), and the latent confounding effect will remain unaltered (assuming a linear gene regulatory network model). For these two environments (observational and interventional), the condition above can be easily satisfied. More generally, simulations show that whenever the environments are more heterogeneous, KL regression performs better.
In practice, all the quantities that define the loss $\mathcal{L}$ are estimated from the empirical covariances $\hat{\Sigma}_e^{joint}$: define $\hat{\beta}^{KL}$ as the plug-in estimator of $\tilde{\beta}^{KL}$ where we replace the environment population covariances (and derived quantities) by their empirical counterparts, where analogously $\hat{S}_{\beta}$ is the empirical version of $S_{\beta}$. For each environment $e \in [E]$ we observe \mbox{$n_e\geq 1$} i.i.d. samples. We can provide a distributional result: assuming that we have access to the population residual variances $\sigma_e^2$, conditioned on the sample covariates (i.e. as in fixed design linear regression), the KL regression estimator $\hat{\beta}^{KL}$ satisfies the following result:
\begin{proposition}\label{proposition:distribution-DR}
We have that, for $\epsilon$ Gaussian,
\begin{align*}
    \hat{\beta}^{KL}\mid(X_e)_e \sim \mathcal{N}\big(\mathbb{E}[\hat{\beta}^{KL}|(X_e)_e], V \big)
\end{align*}{}
where $V = \sum_{e=1}^E \frac{1}{n_e}A_e$, for some symmetric matrices $A_e$ that depend only on the empirical covariance matrices. In the case where $n_e = n$ for all $e$, we have
\begin{align*}
V = \frac{1}{n} \hat{S}_{\beta}^{-1}\Big(\sum_{e\in [E]}\frac{\hat{\Sigma}_{e,X}^{-1}}{\sigma_e^2}\Big)
\end{align*}
\end{proposition}{}
We prove this proposition in Appendix~\ref{proof:proposition:distribution-DR}. If all environments have the same number of samples $n$, the plug-in estimator $\hat{\beta}^{KL}$ converges to $\tilde{\beta}^{KL}$ at rate $1/ \sqrt{n}$. Also, whenever $S_{\beta}$ is ill-conditioned, the variance of the estimator $\hat{\beta}^{KL}$ will be large. In general, we do not have access to the population residual variances $\sigma_e^2$, and the distribution is not Gaussian. We leave this extension to future work.

\subsection{Kullback-Leibler Regression with Sparsity Inducing Penalty}

\begin{figure*}[t]
\centering
\includegraphics[width=0.8\linewidth]{figure_c.pdf}
\caption{\textbf{Diverse environments improve KL regression performance} Plots of the (log-scaled) MSE as a function of the environment-wise random component in $\Sigma_{e,\epsilon_X}$. }
\label{fig:environment-diversity}
\end{figure*}

Often times it is reasonable to assume that the invariant vector $\beta^*$ is sparse. Modifying KL regression to incorporate this prior knowledge can substantially improve the accuracy and stability of the method. We add an $L^1$-norm sparsity-inducing penalty to the optimization problem (\ref{loss-DR}), which is convex in $(\beta,\eta)$, to define the following \emph{Lasso KL regression}:
\begin{align}\label{loss-DR-sparse}
    & \mathcal{L}_{\lambda}(\beta, \eta) \!:=\!  \mathcal{L}(\beta, \eta) + \lambda\Vert\beta\Vert_1
    \\ & (\tilde{\beta}^{KL}(\lambda), \tilde{\eta}(\lambda)) := \argmin_{\beta, \eta \in \mathbb{R}^D} \mathcal{L}_{\lambda}(\beta, \eta)
\end{align}{}
Solving this problem in practice is easily done with off-shelf optimization packages such as CVXPY \citep{diamond2016cvxpy} in Python. Surprisingly, Lasso KL regression outperforms all the competing methods in sparse settings, is stable in small sample settings and robust to model misspecification. We leave for future work a deeper analysis of finite sample guarantees of the Lasso KL regression.

\subsection{Robustness of Kullback-Leibler Regression to Model Misspecification}

If we assume that the object of interest in the linear SEM (\ref{equation-linear-sem}), $\beta^*$, remains invariant across environments then we can bound its estimation error under misspecification. Assume we allow small perturbations in the remaining coefficients of our SEM model~(\ref{eq:model-of-environments}) for some of the environments: we define \mbox{$\eta_e := (I-B_{XX,e})^{-1}B_{XH,e}\Sigma_{\epsilon_X ,e}\eta_{0,e}$}. Define \mbox{$\delta_e := \eta_e - \eta^*$}, we have the following bound:
\begin{proposition}\label{proposition:robustness} We have:
\begin{align*}
& \Vert S_{\beta} (\tilde{\beta}^{KL} - \beta^*)\Vert^2 \leq C \sup_{e\in [E]}\Vert \delta_e\Vert^2
\end{align*}
where \mbox{$C:= \big(\sum_{e\in [e]}\frac{1}{\sigma_e^2}\big)\Big(\Vert\sum_{e\in [E]}\frac{\Sigma_e^{-1}}{\sigma_e^2}\Vert + \sum_{e\in [e]}\frac{\Vert\Sigma_e^{-1}\Vert}{\sigma_e^2}\Big)$}
\end{proposition}
In practice, we can bound the misspecification error in $\tilde{\beta}^{KL}$ with a bound on the (unknown) term $\sup_{e\in [E]}\Vert \delta_e\Vert^2$, as the remaining terms in $C$ can be estimated from the data. In particular, the effect of the model misspecification is entirely captured by $\delta_e$, even though the covariate covariance is also modified by such perturbations.

\begin{figure}
     \centering
     \begin{subfigure}
\centering
\includegraphics[width=.8\linewidth]{figure_f.pdf}
\end{subfigure}
     \vfill
     \begin{subfigure}
\centering
\includegraphics[width=.8\linewidth]{figure_b.pdf}
\end{subfigure}
     \vfill
     \begin{subfigure}
\centering
\includegraphics[width=.8\linewidth]{figure_j.pdf}
\end{subfigure}
\label{fig:model-misspecification}
\caption{\textbf{Effect of model misspecification.} Plots of the (log-scaled) MSE of the estimator of the regression parameter (with respect to the true $\beta^*$) as a function of the noise level added to $B_{XX}$ (top), $B_{XH}$ (center) and $\Sigma_{\epsilon_H}$ (bottom). Surprisingly, KL regression performs remarkably well (in particular the $L^1$-penalized version), showing a certain level of robustness to model misspecification. The x-axis scaling is not linear.}
\end{figure}
Given that in practice this bound is hard to compute, we empirically evaluate the robustness of KL regression though simulations. We perturb the connectivity matrix $B_{XX}$ by adding a random noise to it, which varies in strength: denoting $s$ the scaling of the random component, we have that $B_{e,XX} = B_{0,XX} + sZ_e$, where $B_{0,XX}$ is shared across environments, and $Z_e$ is a matrix with standard Gaussian entries. We run analogous simulations perturbing the connectivity matrix $B_{XH}$, and perturbing the latent variable covariance matrix $\Sigma_{\epsilon_H}$. We report the results of the simulation in Figure~\ref{fig:model-misspecification}, where we vary $s \in [0,5]$ (where $s=0$ corresponds to our baseline configuration described in section~\ref{appendix:simulation-environments}) and select a few values that are representative of the general behavior of the estimators. As expected, KL regression is robust to small perturbations in the connectivity matrix and outperforms most other methods. As the strength of perturbation in $B_{e,XX}$ increases KL regression performance degrades substantially. Surprisingly, Lasso KL regression seems to be very robust to perturbations in $B_{XH}$, $\Sigma_{\epsilon_H}$. A noticeable effect is that, under a slight perturbation in $B_{XX}$, Anchor Regression improves substantially and performs surprisingly well, although it is designed only to protect against a certain shift-interventions in some particular directions. We hypothesize that the changes due to model misspecification lead to a more conservative hedging in multiple directions, bringing Anchor regression closer to Causal Transfer Learning \citep{rojas2018invariant}, which in theory recovers $\beta^*$.

\section{Experiments}\label{section:section-4}

\subsection{Validation on Synthetic Data}

We extensively evaluate KL regression on synthetic data. We compare the following methods: \textbf{1) KL regression}, both Lasso KL regression and unregularized KL regression. \textbf{2) Invariant Causal Prediction (ICP)} \citep{peters2016causal} The implementation we used in the R package (\texttt{InvariantCausalPrediction}) detects that the assumptions for using ICP are not satisfied, and suggests running the \texttt{hiddenICP} version. \textbf{3) Anchor Regression} \citep{rothenhausler2018anchor} \textbf{4) Causal Dantzig} \citep{rothenhausler2019causal}  We run experiments with both the regularized and unregularized versions. \textbf{5) ComBat} \citep{johnson2007adjusting} \textbf{6) Averaging the regression coefficients}. Appendix~\ref{appendix:simulation-environments} provides a detailed description of the data generating process of our synthetic data simulations, and precise description of all competing methods and the following simulations.
\paragraph{Simulations with Varying Number of Environments.} Figure~\ref{fig:figure_a} plots the mean squared error (MSE) for the different estimators as a function of the number of samples for 2 and 4 environments. KL regression recovers $\beta^*$ with relatively few samples in several environments. Compared to alternatives in Table~\ref{table-methods}, KL regression is the only method that provably recovers $\beta^*$ at the population level. Compared to unregularized Causal Dantzig and hidden ICP (which behave similarly in 2 environments), KL regression is substantially more efficient as the number of environments increases.
\paragraph{Increasing Diversity Across Environments.} As suggested in Proposition~\ref{proposition:invertibility_S}, having diverse covariance matrices $\Sigma_{e,X}$ implies that $S_{\beta}$ is likely to be invertible and well-conditioned, and thus the asymptotic variance of our KL regression estimator will be lower.  Figure~\ref{fig:environment-diversity} shows how MSE changes as the environment-wise covariances $\Sigma_{e,\epsilon_{X}}$ differ more one of another (cf. Appendix~\ref{appendix:diversity} for further details). All methods improve with the diversity of the environments, but KL regression particularly benefits from diversity.

\paragraph{Additional Simulations.} We further explore the limitations and the robustness of our method in Appendix~\ref{appendix:simulations}, and summarize here our findings:
\begin{itemize}[noitemsep,nolistsep,leftmargin=*]
    \item We analyze the performance under various degrees of confounding  in Appendix~\ref{appendix:varying-confounding}. KL regression is robust as confounding gets stronger, while other methods degrade rapidly for higher confounding.
    \item We report in Appendix~\ref{appendix:sparsity} simulations where the sparsity level changes: regularized Causal Dantzig, and regularized KL regression which perfectly recover sparse vectors, degrade in settings where $\beta^*$ is not sparse. Unregularized KL regression works well in both sparse and non-sparse settings.
    \item Artificially splitting the data from one environment to generate new environments does not help KL regression, it actually worsens it as it does not increase environment diversity (Appendix~\ref{appendix:split}).
    \item KL regression does not require data to originate from a Gaussian SEM. We validate the previous results on data sampled from Student's t-distribution with different degrees of freedom and report the results in Appendix~\ref{appendix:student}.
\end{itemize}

\subsection{Recovering Gene Regulatory Networks: DREAM Challenge Datasets}
Validating KL regression on real-world datasets is challenging as there is often no ground truth causal mechanism. Gene Regulatory Networks (GNR) can provide us with some partial ground truth, indicating only whether two genes, a transcription factor (TF) and a regulated gene, are related or not (i.e. whether the corresponding $\beta^*_i$ is non-null). We therefore focus on Lasso KL regression and regularized Causal Dantzig for this experiment, as they automatically select features. This type of datasets also fits well with our model of environments (cf. equation~\ref{eq:model-of-environments}): deletions on upstream genes provide the sort of perturbation that shifts the global covariate distribution of gene expressions, without altering the true regulation mechanism of a specific target gene. We run our method on some of the DREAM4 and DREAM5 datasets released as part of an well-known online challenge. Importantly, we can also benchmark KL Regression against the leaderboard of the challenge via the AUPR metric (area under precision-recall curve) used in the challenge. We were limited in our experiments to datasets for which the environment-wise covariance matrix $\Sigma_{e,X}$ is invertible, and we expect to relax this assumption in future work. We generated different environments depending on whether genes had been knocked-down. This required assumption on $\Sigma_{e,X}$ does not hold whenever the number of gene expression measurements is smaller than the number of TFs. However, single cell gene sequencing has the potential to make KL regression widely applicable in future GNR inference tasks. We provide in Appendix~\ref{appendix:DREAM} additional information for these real-data experiments.

In the well-known DREAM5 challenge, our method \emph{ranked at the top of the leaderboard} by a substantial margin for the analyzed dataset. For DREAM4, KL regression still ranked in the top 40\% of submissions: however the semi-synthetic datasets from DREAM4 contained one single knock-down experiments for every single gene, making our model of environments less realistic (compared to DREAM5 where a single knock-down experiment had multiple replicates).

\begin{table}
  \caption{AUPR Performance on DREAM Challenge.}
  \label{table:DREAM}
  \centering
  \begin{tabular}{lllllll}
    \toprule
    Method    & DREAM 4 & DREAM 5    \\
    \midrule
    Lasso KL Regr. & 0.26 & \textbf{0.45} \\
    C. Dantzig (reg.) &  0.15   & 0.27  \\
    Top Submission & \textbf{0.37} & 0.31 \\
    \bottomrule
  \end{tabular}
\end{table}

\section{Discussion}
We provide a detailed mathematical foundation for KL regression, as well as extensive validation on synthetic and real data. KL has strong mathematical guarantees, and works better across many settings compared to existing methods. We derived our theoretical guarantee for KL regression in a standard SEM setting commonly used in causal inference, without assuming that data is sampled from any particular parametric distribution (i.e. we do not assume data is Gaussian). Empirically, KL regression performs well even with large model mismatch; we are very encouraged that it \emph{achieved the top score on a Dream5 dataset}, which has complex nonlinear interactions. KL regression introduces a new regression-flavored perspective to the invariance-based causal inference framework, opening the door to potentially leverage other regression-related concepts such as bootstrapping, model selection techniques (AIC, BIC criterion) and various regularization methods (beyond the current Lasso KL regression $L^1$-penalized version). Interesting future research directions will also consist in developing a non-linear version of KL regression for generalized linear models.

\subsubsection*{Acknowledgments}

The authors thank Dominik Rothenhäusler and Bryan He for their helpful comments, and Kailas Vodrahalli for his help with his implementation of invariant risk minimization. J.R.G. is supported by a Stanford Graduate Fellowship. J.Z. is supported by a Chan–Zuckerberg Biohub Investigator grant and National Science Foundation (NSF) Grant CRII 1657155.

redpajama_set_name: RedPajamaArXiv