SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Adrián Javaloy<sup>1,*</sup> Pablo Sánchez-Martín<sup>1,2</sup> Isabel Valera<sup>1,3</sup>
$^{1}$ Department of Computer Science of Saarland University, Saarbrücken, Germany
$^{2}$ Max Planck Institute for Intelligent Systems, Tübingen, Germany
<sup>3</sup>Max Planck Institute for Software Systems, Saarbrücken, Germany
In this work, we deepen on the use of normalizing flows for causal inference. Specifically, we first leverage recent results on non-linear ICA to show that causal models are identifiable from observational data given a causal ordering, and thus can be recovered using autoregressive normalizing flows (NFs). Second, we analyse different design and learning choices for causal normalizing flows to capture the underlying causal data-generating process. Third, we describe how to implement the do-operator in causal NFs, and thus, how to answer interventional and counterfactual questions. Finally, in our experiments, we validate our design and training choices through a comprehensive ablation study; compare causal NFs to other approaches for approximating causal models; and empirically demonstrate that causal NFs can be used to address real-world problems—where mixed discrete-continuous data and partial knowledge on the causal graph is the norm. The code for this work can be found at https://github.com/psanch21/causal-flows.
Deep learning is increasingly used for causal reasoning, that is, for finding the underlying causal relationships among the observed variables (causal discovery), and answering what-if questions (causal inference) from available data [28]. Our focus in this paper is to solve causal inference problems using only observational data and (potentially partial) knowledge on the causal graph of the underlying structural causal model (SCM). This is exemplified in Fig. 1, where our proposed framework is able to estimate the (unobserved) causal effect of externally intervening on the sensitive attribute (red and yellow distributions), using solely observed data (blue distribution) and partial information about the causal relationship between features.
In this context, previous works have mostly relied on different deep neural networks (DNNs)—e.g., nor
malizing flows (NFs) [23, 26, 27], generative adversarial networks (GANs) [20, 39], variational autoencoders (VAEs) [15, 40], Gaussian processes (GPs) [15], or denoising diffusion probabilistic models (DDPMs) [2]—to iteratively estimate the conditional distribution of each observed variable given its causal parents, thus using an independent DNN per observed variable. Hence, to predict the effect of an intervention in the causal data-generating process, these approaches fix the value of the intervened variables when computing the new value for their children. However, they may also suffer
German Credit - Checking account
Figure 1: Observational and interventional distributions of the categorical variable checking account of the German Credit dataset [8], and their estimated values according to a causal normalizing flow.  $\mathbf{x}_S$  is a binary variable representing the users' sex.
*Correspondence to ajavaloy@cs.uni-saarland.de.
from error propagation—which worsens with long causal paths—and a high number of parameters, which is addressed in practice with ad-hoc parameter amortization techniques [26, 27]. Moreover, several approaches also rely on implicit distributions [2, 20, 27, 32], and thus do not allow evaluating the learnt distribution.
In contrast, and similar to [16, 32, 34, 41], we here aim at learning the full causal-generating process using a single DNN and, in particular, using a causal normalizing flow. To this end, we first theoretically demonstrate that causal NFs are a natural choice to approximate a broad class of causal data-generating processes (§3). Then, we design causal NFs that inherently satisfy the necessary conditions to capture the underlying causal dependencies (§4), and introduce an implementation of the do-operator that allows us to efficiently solve causal inference tasks (§5). Importantly, our causal NF framework allows us to deal with mixed continuous-discrete data and partial knowledge on the causal graph, which is key for real-world applications. Finally, we empirically validate our findings and show that causal NFs outperform competing methods also using a single DNN to approximate the causal data-generating process (§6).
Related work To the best of our knowledge, the closest works to ours are [16, 34, 41], as they all capture the whole causal data-generating process using a single DNN. Our approach generalizes the result from Khemakhem et al. [16], which also relies on autoregressive normalizing flows (ANF), but only considers affine ANFs and data with additive noise. In contrast, our work provides a tighter connection between ANFs and SCMs (affine or not), more general identifiability results, and sound ways to both embed causal knowledge in the ANF, and to apply the do-operator. Another relevant line of works connect SCMs with GNNs [32, 41], and despite making little assumptions on the underlying SCM, they lack identifiability guarantees, and interventions on the GNN are performed by severing the graph, which we show in App. C may not work in general. Nevertheless, it is worth noting that the way we use  $A$  in the network design (\$4) is inspired by these works.
2 Preliminaries and background
2.1 Structural causal models, interventions, and counterfactuals
A structural causal model (SCM) [28] is a tuple  $\mathcal{M} = (\tilde{\mathbf{f}}, P_{\mathbf{u}})$  describing a data-generating process that transforms a set of  $d$  exogenous random variables,  $\mathbf{u} \sim P_{\mathbf{u}}$ , into a set of  $d$  (observed) endogenous random variables,  $\mathbf{x}$ , according to  $\tilde{\mathbf{f}}$ . Specifically, the endogenous variables are computed as follows:
$$
\mathbf {u} := \left(\mathrm {u} _ {1}, \mathrm {u} _ {2}, \dots , \mathrm {u} _ {d}\right) \sim P _ {\mathbf {u}}, \quad \mathrm {x} _ {i} = \tilde {f} _ {i} \left(\mathrm {x} _ {\text {p a} _ {i}}, \mathrm {u} _ {i}\right), \quad \text {f o r} i = 1, 2, \dots , d. \tag {1}
$$
In other words, each  $i$ -th component of  $\tilde{\mathbf{f}}$  maps the  $i$ -th exogenous variable  $\mathbf{u}_i$  to the  $i$ -th endogenous variable  $\mathbf{x}_i$ , given the subset of the endogenous variables that directly cause  $\mathbf{x}_i$ ,  $\mathbf{x}_{\mathrm{pa}_i}$  (causal parents).
An SCM also induces a causal graph, a powerful tool to reason about the causal dependencies of the system. Namely, the causal graph of an SCM  $\mathcal{M} = (\tilde{\mathbf{f}}, P_{\mathbf{u}})$  is the directed graph that describes the functional dependencies of the causal mechanism. We can define the adjacency matrix of the causal graph as  $\mathbf{A} := \nabla_{\mathbf{x}} \tilde{\mathbf{f}}(\mathbf{x}, \mathbf{u}) \neq \mathbf{0}$ , where  $\mathbf{0}$  is the constant zero function, and the comparisons are made elementwise. Furthermore, the direct causes of the  $i$ -th variable (pa $_i$  in Eq. 1) are the parent nodes of the  $i$ -th node in  $\mathbf{A}$ , and the ancestors of this node (which we denoted by  $\mathrm{an}_i$ ) are its (in)direct causes. See Fig. 2 for an example of a causal chain.
In the case that  $A$  is acyclic, we can pick a causal ordering describing
which variables do not cause others, and which ones may cause them. Namely, a permutation  $\pi$  is said to be a causal ordering of an SCM  $\mathcal{M}$  if, for every  $\mathbf{x}_i$  that directly causes  $\mathbf{x}_j$ , we have  $\pi(i) < \pi(j)$ . Note that this definition equals that of a topological ordering and, without loss of generality, we will assume throughout this work that the variables are sorted according to a causal ordering.
Importantly, besides describing the (observational) data-generating process, SCMs enable causal inference by allowing us to answer what-if questions regarding: i) how the distribution over the observed variables would be if we force a fixed value on one of them (interventional queries); and ii) what would have happened to a specific observation, if one of its dimensions would have taken a different value (counterfactual queries).
Figure 2: Causal graph, and its causal ordering  $\pi$  and adjacency matrix  $A$ .
$$
\pi = (1 \quad 2 \quad 3)
$$
$$
\boldsymbol {A} = \left( \begin{array}{c c c} 0 & 0 & 0 \\ 1 & 0 & 0 \\ 0 & 1 & 0 \end{array} \right)
$$
Structural equivalence To reason about causal dependencies, we introduce the notion of structural equivalence. We say that two matrices  $S$  and  $R$  are structurally equivalent, denoted  $S \equiv R$ , if both matrices have zeroes exactly in the same positions. Similarly, we say that  $S$  is structurally sparser than  $R$ , denoted as  $S \preceq R$ , if whenever an element of  $R$  is zero, the same element of  $S$  is zero.
2.2 Autoregressive normalizing flows
Normalizing flows (NFs) [25] are a model family that express the probability density of a set of observations using the change-of-variables rule. Given an observed random vector  $\mathbf{x}$  of size  $d$ , a normalizing flow is a neural network with parameters  $\theta$  that takes  $\mathbf{x}$  as input, and outputs
$$
T _ {\boldsymbol {\theta}} (\mathbf {x}) =: \mathbf {u} \sim P _ {\mathbf {u}} \quad \text {w i t h} \log \text {- d e n s i t y} \quad \log p (\mathbf {x}) = \log p \left(T _ {\boldsymbol {\theta}} (\mathbf {x})\right) + \log | \det  (\nabla_ {\mathbf {x}} T _ {\boldsymbol {\theta}} (\mathbf {x})) |, \tag {2}
$$
where  $P_{\mathbf{u}}$  is a base distribution that is easy to evaluate and sample from. Since Eq. 2 provides the log-likelihood expression, it naturally leads to the use of maximum likelihood estimation (MLE) [1] for learning the network parameters  $\theta$ . While many approaches have been proposed in the literature [25], here we focus on autoregressive normalizing flows (ANFs) [18, 24]. Specifically, in ANFs the  $i$ -th output of each layer  $l$  of the network, denoted by  $\mathbf{z}_i^l$ , is computed as
$$
\mathbf {z} _ {i} ^ {l} := \tau_ {i} ^ {l} \left(\mathbf {z} _ {i} ^ {l - 1}; \mathbf {h} _ {i} ^ {l}\right), \quad \text {w h e r e} \quad \mathbf {h} _ {i} ^ {l} := c _ {i} ^ {l} \left(\mathbf {z} _ {1: i - 1} ^ {l - 1}\right), \tag {3}
$$
and where  $\tau_{i}$  and  $c_{i}$  are termed the transformer and the conditioner, respectively. The transformer is a strictly monotonic function of  $\mathbf{z}_i^{l - 1}$ , while the conditioner can be arbitrarily complex, yet it only takes the variables preceding  $\mathbf{z}_i$  as input. As a result, ANFs have triangular Jacobian matrices,  $\nabla_{\mathbf{x}}T_{\theta}(\mathbf{x})$ .
3 Causal normalizing flows
Problem statement Assume that we have a sequence of i.i.d. observations  $X = \{x_{1}, x_{2}, \ldots, x_{N}\}$  generated according to an unknown SCM  $\mathcal{M}$ , from which we have partial knowledge of its causal structure. Specifically, we know at least its causal ordering  $\pi$ , and at most the whole causal graph  $\mathbf{A}$ . Our objective in this work is to design and learn an ANF  $T_{\theta}$ , with parameters  $\theta$ , that captures  $\mathcal{M}$  by maximizing the observational likelihood (MLE), i.e.,
$$
\underset {\boldsymbol {\theta}} {\operatorname {m a x i m i z e}} \frac {1}{N} \sum_ {n = 1} ^ {N} \left[ \log p \left(T _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right)\right) + \log | \det  \left(\nabla_ {\mathbf {x}} T _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right)\right) | \right], \tag {4}
$$
and that can successfully answer interventional and counterfactual queries during deployment, thus enabling causal inference. We refer to such a model as a causal normalizing flow.
Assumptions We restrict the class of SCMs considered by making the following fairly common assumptions: i) diffeomorphic data-generating process, i.e.,  $\tilde{\mathbf{f}}$  is invertible, and both  $\tilde{\mathbf{f}}$  and its inverse are differentiable; ii) no feedback loops, i.e., the induced causal graph is acyclic; and iii) causal sufficiency, i.e., the exogenous variables are mutually independent,  $p(\mathbf{u}) = \prod_{i}p(\mathbf{u}_{i})$ .
SCMs as TMI maps To achieve our objective, and bridge the gap between SCMs and normalizing flows, we resort to triangular monotonic increasing (TMI) maps, which are autoregressive functions whose  $i$ -th component is strictly monotonic increasing with respect to its  $i$ -th input. TMI maps hold a number of useful properties, such as being closed under composition and inversions. Conveniently, a layer of an ANF (Eq. 3) is a parametric TMI map that can approximate any other TMI map arbitrarily well, which makes ANFs also TMI maps approximators;² a fact that has been exploited in the past to prove that ANFs are universal density approximators [25].
We now show that any SCM can be rewritten as a tuple  $(\mathbf{f}, P_{\mathbf{u}}) \in \mathcal{F} \times \mathcal{P}_{\mathbf{u}}$ , where  $\mathcal{F}$  is the set of all TMI maps, and  $P_{\mathbf{u}}$  is the set of all fully-factorized distributions,  $p(\mathbf{u}) = \prod_{i} p(\mathbf{u}_i)$ . First, given an acyclic SCM  $\mathcal{M} = (\tilde{\mathbf{f}}, P_{\mathbf{u}})$  with  $\tilde{\mathbf{f}}: \mathbb{X} \times \mathbb{U} \to \mathbb{X}$  as in Eq. 1, we can always unroll  $\tilde{\mathbf{f}}$  by recursively replacing each  $\mathbf{x}_i$  in the causal equation by its function  $\tilde{f}_i$  (see Fig. 4b for an example), obtaining an equivalent non-recursive function  $\hat{\mathbf{f}}: \mathbb{U} \to \mathbb{X}$ . This function  $\hat{\mathbf{f}}$  writes each  $\mathbf{x}_i$  as a function of its exogenous ancestors  $\mathbf{u}_{\mathrm{an}_i}$  and, since  $\mathcal{M}$  is acyclic,  $\hat{\mathbf{f}}$  is a triangular map. For simplicity, assume that  $P_{\mathbf{u}}$  is a standard uniform distribution. Then, following the causal ordering, we can apply a Darmois
2While it is a common to shuffle the inputs for each layer, we keep the same order across the network.
construction [7, 13] and replace each function  $\hat{f}_i$  by the conditional quantile function of the variable  $\mathbf{x}_i$  given  $\mathbf{x}_{\mathrm{pa}_i}$  (which depends on  $\mathbf{u}_{\mathrm{an}_i}$ ) eventually arriving to a TMI map  $\mathbf{f}$ . This procedure follows the proof for non-identifiability in ICA [13], but restricted to one ordering. The case for a general  $P_{\mathbf{u}}$  follows a similar construction, but using a Knothe-Rosenblatt (KR) transport [19, 31] instead.
Isolating the exogenous variables Now that we have SCMs and causal NFs under the same family class-i.e., the family  $\mathcal{F}\times \mathcal{P}_{\mathbf{u}}$  of TMI maps with fully-factorized distributions—we leverage existing results on identifiability to show that we can find a causal NF  $T_{\theta}$  such that the  $i$ -th component of  $T_{\theta}(\mathbf{x})$  is a function of the true exogenous variable  $\mathbf{u}_i$  that generated the observed data. More precisely, note that, since we can rewrite  $\mathcal{M}$  as an element of the family  $\mathcal{F}\times \mathcal{P}_{\mathbf{u}}$ , identifying the true exogenous variables of an SCM  $\mathcal{M}$  is equivalent to solving a non-linear ICA problem with TMI generators, for which Xi and Bloem-Reddy [38] proved the following (re-stated to match our setting):
Theorem 1 (Identifiability). If two elements of the family  $\mathcal{F} \times \mathcal{P}_{\mathbf{u}}$  (as defined above) produce the same observational distribution, then the two data-generating processes differ by an invertible, component-wise transformation of the variables  $\mathbf{u}$ .
Thm. 1 implies that, if we can find a causal NF  $(T_{\theta}, P_{\theta}) \in \mathcal{F} \times \mathcal{P}_{\mathbf{u}}$  that matches the observational distribution generated by  $\mathcal{M} = (\mathbf{f}, P_{\mathcal{M}}) \in \mathcal{F} \times \mathcal{P}_{\mathbf{u}}$ , then we know that the exogenous variables of the flow differ from the real ones by a function of each component independently, i.e.,  $T_{\theta}(\mathbf{f}(\mathbf{u})) \sim P_{\theta}$  with  $\mathbf{u} \sim P_{\mathcal{M}}$  and  $T_{\theta}(\mathbf{f}(\mathbf{u})) = h(\mathbf{u}) = (h_1(\mathfrak{u}_1), h_2(\mathfrak{u}_2), \ldots, h_d(\mathfrak{u}_d))$ , where each  $h_i$  is an invertible function. Fig. 3 graphically illustrates Thm. 1. Furthermore, Thm. 1 also implies that the functional dependencies of the causal NF must agree with that of the SCM, i.e., that  $T_{\theta}$  needs to be causally consistent with  $\mathcal{M}$ . We formally present this result in the following corollary (proof can be found in App. A), where  $\pmb{I}$  denotes the identity matrix:
Figure 3: Thm. 1 as a commutative diagram
Corollary 2 (Causal consistency). If a causal NF  $T_{\theta}$  isolates the exogenous variables of an SCM  $\mathcal{M}$ , then  $\nabla_{\mathbf{x}}T_{\theta}(\mathbf{x}) \equiv \mathbf{I} - \mathbf{A}$  and  $\nabla_{\mathbf{u}}T_{\theta}^{-1}(\mathbf{u}) \equiv \mathbf{I} + \sum_{n=1}^{\mathrm{diam}(A)}\mathbf{A}^{n}$ , where  $A$  is the causal adjacency matrix of  $\mathcal{M}$ . In other words,  $T_{\theta}$  is causally consistent with the true data-generating process,  $\mathcal{M}$ .
A sketch of the proof goes as follows: since Fig. 3 is a commutative diagram, we can write the result of  $T_{\theta}$  and  $T_{\theta}^{-1}$  in terms of the true  $f$ ,  $h$ , and their inverses. Then, we can use the chain rule to compute their Jacobian matrices, and since  $h$  has a diagonal Jacobian matrix, it preserves the structure of the Jacobian matrices of  $f$  and its inverse. To sum up, we have shown that causal NFs are a natural choice to estimate an unknown SCM by showing that: i) both SCMs and causal NFs fall within the same family  $\mathcal{F} \times \mathcal{P}_{\mathbf{u}}$ ; ii) any two elements of this family with identical observational distributions are causally consistent; and iii) they differ by an invertible component-wise transformation.
3.1 Causal NFs for real-world problems
To bring theory closer to practice, we need to extend causal NFs to handle mixed discrete-continuous data and partial knowledge on the causal graph, which are common properties of real-world problems. Due to space limitations, we provide here a brief explanation, and formalize these ideas in App. A.2.
Discrete data To extend our results to also account for discrete data, we take advantage of the general model considered by Xi and Bloem-Reddy [38] that includes observational noise (independent of the exogenous variables), and consider a continuous version of the observed discrete variables by adding to them independent noise  $\varepsilon \in [0,1]$  (e.g., from a standard uniform), such that the real distribution is still recoverable. Intuitively, our approach assumes that discrete variables correspond to the integer part of (noisy) continuous variables generated according to an SCM fulfilling our assumptions, such that both our theoretical and practical insights still apply.
Partial knowledge While we rarely know the entire causal graph  $A$ , we often have a good grasp on causal relationships between a subset of observed variables—e.g., sex and age are not causally related—while missing the rest. When only partial knowledge on the graph is available—i.e., we only know the causal relationship between a subset of the observed variables, we can instead work with a modified acyclic graph  $\tilde{A}$  obtained by finding the strongly connected components as in [35], where subsets of variables with unknown causal relationships are treated as a block (see §7 for an example).
$$
\mathbf {x} = G \mathbf {x} + I \mathbf {u}
$$
$$
\mathbf {x} = G _ {3} \left(G _ {2} \left(G _ {1} \mathbf {u}\right)\right)
$$
$$
\mathbf {x} = (G ^ {2} + G + I) \mathbf {u}
$$
$$
\mathbf {u} = (\boldsymbol {I} - \boldsymbol {G}) \mathbf {x}
$$
(a) Recursive.
(b) Unrolled.
(c) Compacted.
Figure 4: Example of the linear SCM  $\{\mathrm{x}_1\coloneqq \mathrm{u}_1;\mathrm{x}_2\coloneqq 2\mathrm{x}_1 + \mathrm{u}_2;\mathrm{x}_3\coloneqq 3\mathrm{x}_2 + \mathrm{u}_3\}$  written (a) in its usual recursive formulation; (b) without recursions, with each step made explicit; (c) without recursions, as a single function; and (d) writing  $\mathbf{u}$  as a function of  $\mathbf{x}$ . The red dashed arrows show the influence of  $\mathbf{u}_1$  on  $\mathbf{x}_3$  for all equations from  $\mathbf{u}$  to  $\mathbf{x}$ , with the compacted version exhibiting shortcuts (see §4). Note that in the linear case we have  $A\coloneqq G\neq 0$ , and that  $G_{1},G_{2},G_{3}\preceq G + I$  are any three matrices such that their product equals  $G^{2} + G + I$ .
(d) Inverted.
This allows us to reuse our theoretical results for known parts of the graph, thus generalizing the block identifiability results from von Kugelgen et al. [36].
4 Effective design of causal normalizing flows
We showed in §3 that causal NFs are a natural choice to learn the underlying SCM generating the data. Importantly, Thm. 1 assumes that we can find a causal NF whose observational distribution perfectly matches the true data distribution (according to the underlying SCM). In practice, however, reaching the optimal parameters may be tricky as: i) we only have access to a finite amount of training data; and ii) the optimization process for causal NFs (like for any neural network) may converge to a local optima. In this section, we analyse different design choices for causal NFs to guide the optimization towards solutions that do not only provide an accurate fit of the observational distribution, but allow us to also accurately answer to interventional and counterfactual queries.
Let us start with an illustrative example. Suppose that we are given the linear SCM in Fig. 4a, and we want to write the SCM equations as a TMI map to approximate them with a causal NF. As discussed in §3, we can unroll the causal equations (Fig. 4b)—resulting in a composition of functions structurally as sparse as  $I + A$ . These functions can be compacted into a single transformation (Fig. 4c), such that each  $x_{i}$  depends on its ancestors,  $\mathbf{u}_{\mathrm{an}_i}$ . However, note that in this step shortcuts appear, making direct and indirect causal paths in this representation indistinguishable—in our example, the indirect causal path from  $\mathbf{u}_1$  to  $\mathbf{x}_3$  present in Fig. 4a and Fig. 4b does not go anymore through the path that generates  $\mathbf{x}_2$ , but instead via a shortcut that directly connects  $\mathbf{u}_1$  to  $\mathbf{x}_3$ . Alternatively, we can invert the equations to write  $\mathbf{u}$  as a function of  $\mathbf{x}$  (Fig. 4d), which is structurally equivalent to  $I - A$ .
We remark that the above steps can be applied to any considered acyclic SCM (refer to App. B for a more detailed discussion). In particular, we can unroll the equations in a finite number of steps, and we can similarly reason about the causal dependencies through the Jacobian matrices of the generators,  $\nabla_{\mathbf{x}}T_{\theta}(\mathbf{x})$  and  $\nabla_{\mathbf{u}}T_{\theta}^{-1}(\mathbf{u})$ . Moreover, note that the diffeomorphic assumption implies that we can invert the causal equations. Next, inspired by the different representations of an SCM (exemplified in Fig. 4), we consider the following design choices for causal NFs:
Generative model The first architecture imitates the unrolled equations (Fig. 4b), i.e., the causal NF is defined as a function from  $\mathbf{u}$  to  $\mathbf{x}$ . Importantly, when full knowledge of the causal graph  $\mathbf{A}$  is assumed, we also replicate the structural sparsity per layer by adequately masking the flow with  $\mathbf{I} + \mathbf{A}$ . In this way, the information from  $\mathbf{u}$  to  $\mathbf{x}$  is restricted to flow as if we were unrolling the causal model [34] and, as a result, the output of the  $l - 1$ -th layer of the causal NF is given by:
$$
\mathbf {z} _ {i} ^ {l - 1} = \tau_ {i} \left(\mathbf {z} _ {i} ^ {l}; \mathbf {h} _ {i} ^ {l - 1}\right), \quad \text {w h e r e} \quad \mathbf {h} _ {i} ^ {l - 1} = c _ {i} \left(\mathbf {z} _ {\mathrm {p a} _ {i}} ^ {l}\right). \tag {5}
$$