FCAUSE: FLOW-BASED CAUSAL DISCOVERY
Year: 2021

Abstract:
Current causal discovery methods either fail to scale, model only limited forms of functional relationships, or cannot handle missing values. This limits their reliability and applicability. We propose FCause a new flow-based causal discovery method that addresses these drawbacks. Our method is scalable to both high-dimensional as well as large volume of data, is able to model complex nonlinear relationships between variables, and can perform causal discovery under partially observed data. Furthermore, our formulation generalizes existing continuous optimization-based causal discovery methods, providing a unified view of such models. We perform an extensive empirical evaluation, and show that FCause achieves state of the art results in several causal discovery benchmarks under different conditions reflecting real-world application needs.

heading: 1 INTRODUCTION
text: Understanding causal relationships between variables is crucial in many applications, including biology (Koller & Friedman, 2009; Sachs et al., 2005), economics (Varian, 2016; Cunningham, 2020), and healthcare (Tu et al., 2019). In addition, such information can also be used to advance other machine learning domains, such as fairness (Kusner et al., 2017; Chiappa, 2019), privacy (Tople et al., 2020; Chandrasekaran et al., 2021) and robustness (Arjovsky et al., 2019; Zhang et al., 2020; Zheng et al., 2021). In practice, however, we often do not have a priori knowledge of causal relationships. We can gain this knowledge experimentally, e.g., through randomized control trials (Hariton & Locascio, 2018), but experiments are sometimes cost-prohibitive or unethical. In such cases, causal discovery, the task of learning the causal relationships from existing data is essential. While causal discovery methods have been studied for decades, classical approaches tend to have scaling issues and rely on domain-specific assumptions.
Recent advances in causal discovery use continuous optimization methods to overcome the scalability issues in classical approaches (Zheng et al., 2018; 2020). They frame the combinatorial optimization problem of learning a directed acyclic graph (DAG) as a smooth optimization problem under constraints, allowing the use of efficient algorithms for continuous optimization. In addition to scalability advantages, this represents a first step in bridging the gap between deep learning and causal discovery. However, continuous optimization-based causal discovery remains experimental. These methods cannot handle datasets with missing values and suffer from simple distributional assumptions, making performance sensitive to changes in the dataset’s scale (Reisach et al., 2021; Kaiser & Sipos, 2021). These limitations directly impact these methods’ reliability and applicability.
In contrast, when the causal graph—the structure but not functional forms of causal relationships—is already known, causal-aware (Bayesian) deep learning methods have many application advantages. One recent framework, Carefl (Khemakhem et al., 2021), fits an autoregressive flow model using the variables’ causal ordering for the autoregressive transformations. This yields a model that can be used to answer complex causal queries, such as counterfactuals. Its performance for end-toend causal discovery, however, is limited by its reliance on classical causal discovery methods for inferring causal structures, the cost of fitting several flow-based models, and an inability to handle datasets with missing values.
In this work, we show how continuous optimization methods can be combined with flow-based models to address these deficiencies. We propose FCause, an efficient flow-based causal discovery method able to capture complex nonlinear relations between variables, robust to data scaling, and able to handle datasets with missing values (missing at random). FCause uses powerful flow-based
models to capture complex relationships between variables (Khemakhem et al., 2021), performs efficient causal discovery using ideas from continuous optimization-based causal discovery methods (Zheng et al., 2018), and handles partially-observed datasets using amortized approximate inference to estimate distributions over missing values (Kingma & Welling, 2013).
In section 6, we perform extensive empirical evaluations of FCause, and observe that it achieves state-of-the-art results on multiple causal discovery benchmarks, and that its performance is robust with respect to changes in the scaling of the data (e.g., standardized data or not). In addition, FCauseremains competitive when used on datasets with 30% of missing values, even comparing against baselines used on the corresponding fully-observed dataset.
Finally, we present a unified formulation of several continuous optimization-based causal discovery algorithms based on flow-based models (Khemakhem et al., 2021). This unified perspective facilitates the development of general new techniques to improve these methods, and allows a simple comparison between methods, shedding light on the benefits and drawbacks of each one.

heading: 2 RELATED WORK
text: Approaches for causal discovery from observational data can be classified into three broad groups: constraint-based, score-based and functional causal models (Glymour et al., 2019). Constraint-based methods employ conditional independence tests to discover the underlying causal structure (Spirtes & Glymour, 1991; Spirtes et al., 2000). Score-based methods find a causal structure by optimizing a score function (Chickering & Meek, 2015; Chickering, 2020). Functional causal models represent each variable as a function of its direct causes and some noise term (Shimizu et al., 2006; Hoyer et al., 2008). Classical causal discovery methods in these groups often struggle to scale to high dimensions, as the combinatorial space of possible causal orderings among variables grows superexponentially (Peters et al., 2017). Thus, scalable causal discovery algorithms tend to reduce the search space through reliance on domain-specific assumptions.
Recently, Zheng et al. (2018) introduced Notears, a new family of score-based method based on continuous optimization. Notears uses a novel algebraic characterization of directed acyclic graphs (DAG), which allows an equality-constrained optimization problem to jointly learn the model parameters and adjacency relationships between variables. Notears inspired the development of other methods, Notears-MLP and Notears-Sob (Zheng et al., 2020), Grandag (Lachapelle et al., 2019), and DAG-GNN (Yu et al., 2019), which extend the original formulation to model nonlinear relationships between variables. The methods’ main benefits are its scalability and simplicity, a consequence of the fact that standard numerical solvers can be used to solve the resulting optimization problem. However, they have often been observed to be sensitive to different data scalings (Kaiser & Sipos, 2021), and cannot handle missing values.
Additionally, normalizing flows (Rezende & Mohamed, 2015) have been used to build causal aware models (Cai et al., 2018; Khemakhem et al., 2021). These are based on the fact that the variables’ causal ordering can be used for the transformations used by autoregressive flows (Kingma et al., 2016; Huang et al., 2018).1 These methods are able to model complex nonlinear relationships between variables. However, they rely on prior traditional methods or local search algorithms for causal discovery. Specifically, Self (Cai et al., 2018) uses a hill climbing procedure, and Carefl (Khemakhem et al., 2021) proposes to use a constraint-based method (e.g. PC) to find the graphs skeleton with as many oriented edges as possible and to fit several flow models to determine the orientation of the remaining edges. Apart from aforementioned methods, causal aware deep learning models have, in general, shown better properties regarding generalization and robustness (Arjovsky et al., 2019; Zhang et al., 2020; Kyono et al., 2020; Tople et al., 2020).

heading: 3 PRELIMINARIES
text: In this section, we describe the representations of causal relationships (SEM and DAGs) and explain how causal discovery can be formulated as an optimization task. We bring these components together when we introduce FCause in section 4.
1While the formal connection to a restricted type of autoregressive flows was proposed by (Khemakhem et al., 2021), some previous methods (Mooij et al., 2011; Cai et al., 2018) use closely related ideas.
Causal graphical models (Peters et al., 2017) are widely used to represent causal relationships between variables. Each variable is assigned a node in a graph (typically a DAG), and causal relationships are represented as directed edges between nodes. These models are encoded using the graph’s adjacency matrix A ∈ {0, 1}d×d, where d is the number of variables, or by explicitly stating the set of parents for each variables. If each of the edges is additionally assigned a weight, a weighted adjacency matrix is used W ∈ Rd×d instead. Here, zero entries indicate the absence of an edge, and non-zero entries indicate the presence of an edge with the corresponding weight.
Structural Equation Models (SEM) are commonly used to describe a causal system. They characterize the value of a variable as a function of its causal parents and some external noise. Let x = [x1, . . . , xd] represent d causally related variables, A ∈ {0, 1}d×d the binary adjacency matrix between variables, and z = [z1, . . . , zd] pairwise independent random noise. SEMs model each variable as xi = fi(xpa(i), zi), where fi a scalar function and xpa(i) = {xj : Aj,i = 1} represents the set of parents of xi according to A.
Additive Noise Models (ANM) represent a specific type of SEM for which the noise is additive:
xi = fi(xpa(i)) + zi, i = 1, . . . , d. (1)
Equivalently, given a fixed adjacency matrix A, these d equations can be expressed jointly as
x = fA(x) + z, (2)
where fA(x) outputs a d-dimensional vector whose i-th component is given by fi(xpa(i)). This is one of the most common functional forms in causal discovery. A crucial question regarding ANMs involves the identifiability of the underlying causal graph under observational data. Can the true adjacency matrix be recovered in the limit of infinite data? Hoyer et al. (2008) showed that it is possible except when the underlying system is a combination of linear functions with Gaussian noise. In this work, we assume that the underlying functions are nonlinear without any specific assumption on their functional form. In this setting, the causal graph can be identified.
A new algebraic characterization of DAGs was recently proposed by Zheng et al. (2018). They showed that, given a weighted adjacency matrix W ∈ Rd×d, the quantity
h(W ) = tr ( eW W ) − d (3)
is non-negative, and zero if and only if W represents a DAG. Since its introduction, this algebraic characterization of DAGs has been widely used to frame causal discovery problems with ANMs as continuous optimization tasks (Zheng et al., 2018; 2020; Lachapelle et al., 2019). All these methods propose to train a model’s parameters θ by maximizing a score subject to the constraint h(W (θ)) = 0, whereW (θ) represents the weighted adjacency matrix as a function of the parameters θ. Due to the non-convexity of the set {θ : h(W (θ)) = 0}, these methods use h(W (θ)) as a penalty term added to the loss, whose weight is increased as optimization proceeds.

heading: 4 FCAUSE: FLOW-BASED CAUSAL DISCOVERY
text: We present FCause, a novel flow-based causal discovery method for nonlinear additive noise models. Our method learns a binary adjacency matrix jointly with the model parameters, is able to capture complex nonlinear relationships between variables, is stable to data re-scalings, and can handle datasets with missing values. We present the method for fully-observed datasets in section 4.1, and its extension to datasets with missing values in section 4.2. For the latter we assume that values are missing completely at random or missing at random, meaning that the missingness pattern either has no cause or has fully observed causes, which is a common setting (Rubin, 1976; Stekhoven & Bühlmann, 2012; Ma et al., 2018; Mattei & Frellsen, 2019; Strobl et al., 2018).

heading: 4.1 FCAUSE WITH FULLY-OBSERVED DATA
text: FCause takes a Bayesian approach to causal discovery (Heckerman et al., 1999). We model the adjacency matrix jointly with the observations as
pθ(x 1, . . . , xN , A) = p(A) ∏ n pθ(x n|A). (4)
We propose to fit the model parameters θ, and to use the posterior distribution over binary adjacency matrices pθ(A|x1, . . . , xN ) to characterize the causal structure learned by the model. There are two challenges to applying this approach: (i) The posterior distribution over A is intractable, and (ii) maximum likelihood cannot be used to fit the model parameters, due the the presence of the latent variable A. We overcome these using variational inference (Jordan et al., 1999; Blei et al., 2017; Zhang et al., 2018). We define a variational distribution qφ(A) to approximate the intractable posterior pθ(A|x1, . . . , xN ), and use it to build the evidence lower bound (ELBO), a lower bound on the marginal likelihood that can be used as a surrogate objective. This yields
ELBO(θ, φ) = Eqφ(A)
[ log p(A)
∏ n pθ(x n|A)
] +H(qφ) ≤ log pθ(x1, . . . , xN ), (5)
whereH(qφ) represents the entropy of the distribution qφ. (Details for the derivation in appendix B.)
Modeling and optimization details Following Carefl, we set pθ(xn|A) to an autoregressive flow with base distribution pz and transformation z = gA(x) = x − fA(x) (cf. equations 1 and 2), and p(A) to some prior over adjacency matrices that places zero mass on matrices that do not represent DAGs (needed to guarantee invertibility of the transformation, cf. appendix E). We get FCause’s final objective by replacing pθ(xn|A) in eq. (5) using the change of variable formula from random variables (Kingma et al., 2016). This yields
ELBO(θ, φ) = Eqφ(A)
[ log p(A)
∏ n pz (x n − fA(xn; θ)) |det JA(xn; θ)|
] +H(qφ), (6)
where JA(x; θ) is the Jacobian of the transformation used, gA(x; θ) = x − fA(x; θ). We set the base noise distribution pz to be a factorized Gaussian with mean zero and learnable variances, the variational distribution qφ(A) to be the product of independent Bernoulli distributions, one for each entry in the matrix A,2 and use the Gumbel-softmax to get stochastic estimates of the ELBO’s gradient (Maddison et al., 2016; Jang et al., 2016).3
Interestingly, the objective from eq. (6) may be simplified by noting that the Jacobian-determinant term is always one for matrices A that represent a DAG (see e.g. (Mooij et al., 2011). We formalize this in lemma 1 and include a proof in appendix D). Thus, since the prior over A only allows DAGs, this term does not need to be computed.
Lemma 1. Let A represent a binary adjacency matrix, fA : Rd → Rd a function whose i-th output only depends on the parents of xi, and JA(x) the Jacobian of gA(x) = x− fA(x). If the adjacency matrix A represents a DAG, then |det JA(x)| = 1.
The choice for fA : Rd → Rd must satisfy the adjacency relations specified by A. That is, if Aj,i = 0, then the function fi(x) (i.e., the i-th component of the output of fA(x)) must satisfy ∂fi(x)/∂xj = 0. Inspired by Graph Neural Networks (Hamilton, 2020), we propose a flexible parameterization that satisfies this by setting
fi(x) = hi  d∑ j=1 Aj,i gj(xj)  , (7) where gi and hi (i = 1, . . . , d) are multi-layer perceptrons. A naive implementation of this requires training 2d neural networks. To avoid this, we propose to parameterize these functions as hi(·) = h(ui, ·) and gi(·) = g(ui, ·), where ui is an d-dimensional trainable embedding. This simple idea reduces the number of MLPs needed from 2d to just 2.
Finally, as mentioned above, the prior p(A) must assign zero probability to matrices that do not represent a DAG. While such a prior can be obtained in practice, it leads to a poorly conditioned
2We use the parameterization proposed by Lippe et al. (2021), which uses two parameters per each edge, one to model the existence and the other for the orientation.
3The use of the Gumbel-softmax to learn binary a adjacency matrix was also previously used by Ng et al. (2019a) and Brouillard et al. (2020).
optimization problem. Therefore, following Zheng et al. (2018), we use a “soft” alternative instead and gradually anneal it towards the “hard” DAG-forcing prior during training. This soft prior is
p(A) ∝ exp ( −λs‖vec(A)‖1 − ρ h(A)2 − αh(A) ) , (8)
where h(·) is the non-DAG penalty from eq. (3) (h(A) > 0 for non DAGs and exactly zero for DAGs). The first term favours sparsity in the adjacency matrix A, and the latter two favour DAGs. We anneal this soft prior by increasing the weights ρ and α as optimization proceeds. We include details on the full optimization procedure used in appendix A.
Note that the use of this soft prior means that only in the latter stages of training does the soft prior approach the hard DAG-enforcing constraint. While this means that in the initial training stages the assumptions of lemma 1 do not hold, we find that in latter stages all sampled matrices A are DAGs, satisfying lemma 1’s assumptions and providing strong empirical performance.

heading: 4.2 FCAUSE WITH PARTIALLY-OBSERVED DATA
text: We introduce FCause for partially-observed datasets. Missing values are naturally present in many domains (e.g. online education (Wang et al., 2021)). Thus, developing methods to handle them effectively is crucial. Traditional deletion-based methods fail in many applications since, after deletion, very few data points are left. In contrast, FCause handles missing values efficiently, by using approximate inference to estimate distribution over missing values.
We use xno to denote the observed components of sample x n, xnu to denote the unobserved components, and pθ(xno , x n u) to denote the probability of the sample obtained by combining the values in xno and x n u. As in the formulation without missing values, exact maximum likelihood is intractable. Thus, we propose to introduce a distribution to approximate the true posterior over the latent variables and to optimize the ELBO instead. However, in this case the latent variables include not only the adjacency matrix but also all unobserved values in the dataset. We define the variational approximation as
qφ,ψ(A, x 1 u, . . . , x N u |x1o, . . . , xNo ) = qφ(A) ∏ n qψn(x n u|xno ), (9)
which leads to the objective
ELBO(θ, φ, ψ) = Eqφ,ψ
[ log p(A)
∏ n pθ(x n o , x n u|A)
] +H(qφ) +
∑ n H(qψn). (10)
This is the objective maximized by FCause when the data has missing values. As explained before, we compute p(xno , x n u|A) via the change of variable formula for random variables, and set qφ(A) to be the product of independent Bernoulli distributions. Finally, we set qψn(x n u|xno ) to be a factorized Gaussian, and use reparameterization to get unbiased gradients with respect to ψn.
For efficiency, we use an amortization network (Kingma & Welling, 2013) to avoid training N independent Gaussian approximations. To do so, we train a single neural network with parameters ψ, which receives as input the zero imputed sample xi concatenated with its “missingness” mask— which indicates indices of components missing for that sample—and outputs the mean and scale of q(xiu|xio). This allows us to train a single set of parameters ψ, which are shared across samples in the dataset, instead of N set of independent parameters, {ψ1, . . . , ψN}, one for each sample.

heading: 5 UNIFIED FLOW-BASED FORMULATION AND ANALYSIS
text: This section introduces a simple analysis showing that, similarly to FCause, most causal discovery methods based on continuous optimization can be framed from a probabilistic perspective as fitting a flow. The benefits of this unified perspective are twofold. First, it allows a simple comparison between methods, and sheds light on the different assumptions used by each one, their benefits and drawbacks. Second, it simplifies the development of new tools to improve these methods, since any improvements to one of them can be easily mapped to the others by framing them in this unified framework (e.g. our extension to handle missing values can be easily integrated with Notears).
The connection between causal discovery methods based on continuous optimization and flow-based models uses the concept of a weighted adjacency matrixW (θ) ∈ Rd×d linked to a function f(x; θ) : Rd → Rd. Loosely speaking, these matrices can be seen as characterizing how likely is each output of f(x; θ) to depend on each component of the input x. For instance, W (θ)j,i = 0 indicates that fi(x; θ) is completely independent of xj . Such adjacency matrices can be constructed efficiently for a wide range of parameterizations for f , such as multi layer perceptrons and weighted combinations of nonlinear functions (Zheng et al., 2018; 2020). We give precise details on their definition and construction in appendix C. Lemma 2. Let f(x; θ) : Rd → Rd be a θ-parameterized function with weighted adjacency matrix W (θ) ∈ Rd×d. Given a dataset {x1, . . . , xN}, fitting a flow with the transformation z = x−f(x; θ), base distribution pz and a hard acyclicity constraint on W (θ) is equivalent to solving
max θ ∑ n log pz(x− f(x; θ)) s.t. h(W (θ)) = 0, (11)
where h(·) is the algebraic characterization of DAGs from eq. (3).
Proof. The acyclicity constraint is enforced by constraining the optimization domain to Θ = {θ : h(W (θ)) = 0}. Then, the maximum likelihood objective can be written as∑
n
log pθ(x n) = ∑ n log pz(x n − f(xn; θ)) + log ∣∣∣∣det d(xn − f(xn; θ))dxn ∣∣∣∣ (12)
= ∑ n log pz(x n − f(xn; θ)) (lemma 1), (13)
where the first equality we use the change of variable formula, valid because the transformation z = g(x; θ) = x− f(x; θ) is invertible for any θ ∈ Θ.
Lemma 2 is the main building block in the formulation of continuous optimization-based causal discovery methods from a probabilistic perspective as fitting flow models. This is simply because the objective used by each of the methods can be exactly recovered from eq. (11) with specific choices for f(x; θ) and pz . Notears (Zheng et al., 2018) uses a standard Gaussian for pz and a linear transformation for f(x, θ). (See appendix C for more details.) Notears-MLP (Zheng et al., 2020) uses a standard Gaussian for pz and d independent multi-layer perceptrons, one for each component of f(x, θ). Notears-Sob (Zheng et al., 2020) uses a standard Gaussian for pz and a weighted linear combination of nonlinear basis functions. GAE (Ng et al., 2019b) uses a standard Gaussian for pz and a GNN for f(x, θ). Grandag (Lachapelle et al., 2019) uses a factorized Gaussian with mean zero and learnable scales for pz and d independent multi layer perceptrons, one for each component of f(x, θ). Golem (Ng et al., 2020). This is a linear method whose original formulation was already in a probabilistic perspective, using the linear transformation for f(x; θ). In summary, recently proposed causal discovery methods based on continuous optimization can be formulated from a probabilistic perspective as fitting a flow with different constraints, transformations, and base distributions. From these three components, the base distribution used by the flow is sometimes overlooked but plays an important role. It has been observed, both theoretically and empirically, that methods that use a standard Gaussian (e.g. Notears, Notears-MLP, Notears-Sob) may fail to recover the true underlying causal graph in simple cases, and that their performance may change significantly when re-scaling variables in the dataset (e.g. standardizing data or not) (Loh & Bühlmann, 2014; Kaiser & Sipos, 2021; Reisach et al., 2021). These issues can be overcome by setting the flow’s base distribution to the actual noise distribution used to generate the data. While this distribution is typically unknown, methods such as FCause, Grandag and Golem take a step in this direction by optimizing some of its parameters, instead of fixing all of them to arbitrary values.
Finally, as mentioned above, this unified formulation simplifies the development of new tools to improve these methods. For instance, in this work we introduce a method to handle datasets with
missing values for FCause. Turns out that this idea is readily applicable to any of the causal discovery methods described in this section. This can be seen by formulating them as probabilistic models pθ(x), and replacing the likelihood objective used for fully observed datasets with the ELBO from eq. (10) (and possibly removing the random variable A from the objective).

heading: 6 EXPERIMENTS
text: This section presents results that empirically validate FCause for causal discovery. All simulations were performed using NVIDIA V100 GPUs.
Datasets. We consider synthetic, pseudo-real, and real data. For the synthetic data we follow the approach from Lachapelle et al. (2019) and Zheng et al. (2020). We sample a DAG following two different random graph models, Erdos-Renyi (ER) and scale-free (SF). For each graph we simulate xi = fi(xpa(i)) + zi, where fi is sampled from a Gaussian process and zi ∼ N (0, σ2i ), with σi ∼ U(0.2, 1). We consider three different dimensionalities d ∈ {16, 32, 64}, and for each d we consider two possible number of edges e ∈ {d, 4d}. Each resulting dataset is identified as ER(d, e) or SF(d, e). All datasets have n = 1000 training samples. (In appendix F we show results in the exact same setting but with higher levels of noise, achieved by sampling σi ∼ U(0.2, 2).) For the pseudo-real dataset we consider data generated with the SynTReN generator (Van den Bulcke et al., 2006), which creates synthetic transcriptional regulatory networks and produces simulated gene expression data that approximates experimental data. We use the datasets generated by Lachapelle et al. (2019) (dimension 20), and take n = 400 for training. Finally, for the real dataset we use the widely used dataset that measures the level of several proteins in human cells from Sachs et al. (2005). We use a training set with n = 800 observational samples of dimension d = 11.
Baselines. We compare against PC (Kalisch & Bühlman, 2007), (linear) Notears (Zheng et al., 2018), the nonlinear variants Notears-MLP and Notears-Sob (Zheng et al., 2020), Grandag (Lachapelle et al., 2019) (without the preliminary neighborhood search), GES (Chickering, 2002), and ICALiNGAM (Shimizu et al., 2006). GES was obtained from the GES package, and all other baselines were obtained from the gcastle package.
Causality metrics used. Following common practice (Glymour et al., 2019; Tu et al., 2019) we report adjacency and orientation metrics (recall, precision, and F1 score), and causal accuracy (Claassen & Heskes, 2012). For FCause, which returns a distribution over adjacency matrices, we report the expected values of these metrics estimated using samples A ∼ qφ(A).
In all cases we standardize data and repeat all simulations for four different random seeds. (For the synthetic and pseudo-real datasets this includes generating four datasets, one for each seed.) Table 1 presents a high level summary of the results, by ranking the methods according to different metrics. It can be observed that FCause ranks highest, followed by Notears-MLP, Notears-Sob, GES and PC (all four perform similarly), and finally followed by Grandag, ICALiNGAM and Notears.
Fig. 1 presents more detailed results. It shows the adjacency F1 score, orientation F1 score, and causal accuracy achieved by each method in each of the datasets considered. (We report full results, including standard deviations and all remaining metrics —precision and recall metrics for adjacency and orientation— in appendix G). Again, it can be observed that FCause tends to perform best across most datasets, followed by Notears-MLP, Notears-Sob, GES and PC.
The figure also shows the results achieved by FCause ran on the exact same datasets with 30% of the values missing. This is identified as FCause (PO). As expected, the method performance gets slightly worse in the presence of missing values. However, it can be observed that it is still very competitive, and that on many datasets its performance is better than that of the baselines, even when these are ran using the dataset with no missing values. This shows the effectiveness of FCause’s way of dealing with missing values.
Results using datasets generated with a higher level of noise in are shown in fig. 2 (appendix F). As expected, as a consequence of the higher level of noise, all methods perform slightly worse. However, the main conclusions remain unchanged.
Stability of FCause with respect to data scaling Since FCause learns the scales of the base distribution, we expect its performance to be robust to different re-scaling of the variables in the dataset (cf. section 5). This is in contrast to other methods, such as Notears-MLP and Notears-Sob, who fix the base distribution to a standard Gaussian, and whose performance has been observed to change when re-scaling variables (Kaiser & Sipos, 2021; Reisach et al., 2021). We consider a simple experiment to verify this empirically. We ran FCause, Notears-MLP and Notears-Sob on all datasets with and without data standardization, and compare the results obtained in both cases.
Table 2 shows the performance achieved by each method when ran with and without data standardization. It can be observed that FCause’s mean performance is essentially the same in both cases.
The same cannot be said for Notears-MLP and Notears-Sob, whose performances vary (sometimes significantly) depending on whether the data is standardized or not.
To quantify this variation, table 3 shows the relative difference between the performances achieved by each method with and without standardizing the data. For each method, metric and dataset, the relative difference between the performance obtained with and without data standardization is computed as RD(a, b) = 2|a − b|/(a + b). The table reports mean values across datasets. It can be observed that FCause achieves the lowest relative difference across all metrics by a large margin, followed by Notears-MLP, and finally by Notears-Sob.
Finally, we note that FCause can be used for missing value imputation, using the mean of qψ(xu|xo) to impute missing values. We show imputation results in appendix H, where it can be observed that FCause performs similarly to mice (Van Buuren & Groothuis-Oudshoorn, 2011) and missforest (Stekhoven & Bühlmann, 2012), and outperforms PVAE (Ma et al., 2018).

heading: 7 DISCUSSION AND FUTURE WORK
text: We proposed FCause, a scalable causal discovery method able to model complex nonlinear relationships between variables that can handle datasets with missing values. Our results indicate that FCause could be impactful in real-world applications, where the functional form for the underlying model is not known, data normalization is out of the practitioner’s control, and there may be missing values in the datasets. FCause has shown to yield good results in all these scenarios.
As future work, we identify two possible paths for further improvement. One involves the use of more powerful density approximators to model the noise variables. While we use a Gaussian with learnable scales, other options, such as one dimensional flows, could be tried. The other improvement involves extending the method to models more general than ANMs, such as the post nonlinear (Zhang & Hyvarinen, 2012) or affine (Khemakhem et al., 2021). These extensions should be possible, since all transformations involves in these kinds of models are invertible. Interestingly, once developed, these extensions could be readily applied to other causal discovery methods based on continuous optimization, thanks to the unifying framework presented in section 5.

heading: A OPTIMIZATION DETAILS
text: As mentioned in the main text, we gradually increase the values of ρ and α as optimization proceeds, so that non DAGs are heavily penalized. Inspired by Notears, we do this with a method that resembles the updates used by the augmented Lagrangian procedure for optimization (Nemirovsky, 1999). The optimization process interleaves two steps: (i) Optimize the objective for fixed values of ρ and α for a certain number of steps; and (ii) Update the values of the penalty parameters ρ and α. The whole optimization process involves running the sequence (i)-(ii) until convergence, or until the maximum allowed number of optimization steps is reached.
Step (i). Optimizing the objective for some fixed values of ρ and α using Adam (Kingma & Ba, 2014). We optimize the objective for a maximum of 500 epochs or until convergence, whichever happens first (we assume convergence if the loss does not improve during 100 epochs. If so, we move to step (ii)). We use Adam, initialized with a step-size of 0.01. During training, we reduce the step-size by a factor of 10 if the training loss does not improve for 50 steps. We do this a maximum of two times. If we reach the condition a third time, we do not decrease the step-size and assume optimization converged, and move to step (ii).
Iterating (i)-(ii). We initialize ρ = 1 and α = 0. At the beginning of step (i) we measure the dag-penalty P1 = Eqφ(A)h(A). Then, we run step (i) as explained above. At the beginning of step (ii) we measure the dag-penalty again, P2 = Eqφ(A)h(A). If P2 < 0.65P1, we leave ρ unchanged and update α← α+ρP2. Otherwise, if P2 ≥ 0.65P1, we leave α unchanged and update ρ← 10 ρ. We repeat the sequence (i)-(ii) for a maximum of 25 steps or until convergence (measured as α or ρ reaching some max value), whichever happens first. The approach is summarized in algorithms 1 and 2.
Other details. We use λs = 5 (the regularization factor favouring sparsity).
Algorithm 1 Outer training loop (iterating (i) and (ii)) Require: Sparsity factor λs, temperature for Gumbel-softmax τ , dataset D = {xn}Nn=1, maximum
values ρmax, αmax Initialize (ρ, α) = (1, 0) Initialize each component of qφ(A) to a Bernoulli with p = 0.1 Initialize Niters = 0 while ρ < ρmax & α < αmax & Niters < 25 do
Estimate DAG-penalty P1 = 1M ∑M m=1 h(Am), with Ai ∼ qφ(A) Run inner optimization loop for α and ρ . See algorithm 2 Estimate DAG-penalty P2 = 1M ∑M m=1 h(Am), with Ai ∼ qφ(A) if P2 < 0.65× P − 1 then α = α+ ρP2 elseN is odd ρ = 10× ρ
end ifNiters = Niters + 1 end while

heading: B ELBO DERIVATION
text: The goal of maximum likelihood involves maximizing the likelihood of the observed variables. For FCause (with fully observed datasets) this corresponds to the log-marginal likelihood
log pθ(x 1, . . . , xN ) = log ∑ A p(A) ∏ n pθ(x n|A). (14)
Computing the marginalization from the equation above is intractable, even for moderately low dimensions, since the number of terms in the sum grows exponentially with the size of A (which grows quadratically with the problem’s dimensionality).
Algorithm 2 Inner training loop (step (i)) Require: Sparsity factor λs, temperature for Gumbel-softmax τ , dataset D = {xn}Nn=1
Initialize Nepochs = 0 Initialize step-size η = 0.01 while Not converged & Nepochs < 500 do
for XB ∈ D do . Batch of size 100 Sample A ∼ qφ(A) with gumbel-softmax trick with temperature τ Compute ELBO from eq. (6) using sample A and fA(·; θ) for each x ∈ XB Update φ and θ using∇φ,θELBO with Adam update and step-size η if Condition to reduce step-size then
η = η/10 end if
end for Nepochs = Nepochs + 1
end while
Variational inference proposes to use a distribution qφ(A) to build the ELBO, a lower bound of the objective from eq. (14), as follows:
log pθ(x 1, . . . , xN ) = log ∑ A p(A) ∏ n pθ(x n|A) (15)
= log ∑ A qφ(A) p(A) ∏ n pθ(x n|A) qφ(A)
(16)
= logEqφ(A) [ p(A) ∏ n pθ(x
n|A) qφ(A)
] (17)
≥ Eqφ(A) [ log p(A) ∏ n pθ(x
n|A) qφ(A)
] (Jensen’s inequality) (18)
= Eqφ(A)
[ log p(A)
∏ n pθ(x n|A)
] +H(qφ) (19)
= ELBO(φ, θ), (20)
where we used that H(qφ) = −Eqφ(A) log qφ(A) is the entropy of the distribution qφ. Interestingly, the distribution qφ that maximizes the ELBO is exactly the one that minimizes the KL-divergence between the approximation and the true posterior, KL(qφ(A)‖pθ(A|x1 . . . , xN )) (see, e.g., (Blei et al., 2017)). This is why qφ can be used as a posterior approximation.

heading: C DETAILS ON THE UNIFIED FLOW-BASED FORMULATION
text: Details on weighted adjacency matrices These are used by Notears, Notears-MLP, Notears-Sob and Grandag, among other methods. We will present the formulation from (Zheng et al., 2020) since it is the more general one. Simply put, given a function f(x; θ) : Rd → Rd, Zheng et al. (2020) propose to compute each entry of the weighted adjacency W (θ) as
W (θ)j,i = ‖∂jfi‖L2 , (21)
where ∂jfi denotes the partial derivative of fi(x; θ) with respect to xj . Clearly, if the output of fi(x; θ) does not depend on the value of xj , then we get W (θ)j,i = 0.
While this is a useful general notion to build these weighted adjacency matrices, it is not easy to apply in practice. Zheng et al. (2020) propose some specific alternatives for a few cases, including the case where each fi is given by a multi layer perceptron (MLP) , or as a weighted combination of nonlinear basis functions. We include the details for MLPs here. In this case, each fi is given by an independent MLP. Let Q(1)i represent the weights of the first layer of fi. If the j-th column of Q
(1) i is all zeros, then the output of f1 cannot depend on xj . Using this insight Zheng et al. (2020)
propose to build the weighted adjacency matrix as
W (θ)j,i = ‖jthcolumn(Q(1)i )‖L2 . (22)
Details for the Notears case We show full details for the flow formulation of (linear) Notears. The details for the other methods follow this closely. Notears (linear) finds the optimal weighted adjacency W by solving
min W∈Rd×d N∑ n=1 ‖xn −W xn‖2 s.t. h(W ) = 0, (23)
where h(W ) = 0 exactly characterizes DAGs (eq. (3)). As proposition 1 states, the exact same optimization problem is obtained when fitting a flow under acyclicity constraints. Proposition 1. Fitting a flow using a standard Gaussian as base distribution pz , the transformation z = x − Wx, and a hard constraint enforcing acyclicity on W results in exactly the same optimization problem as the one in eq. (23) used by Notears (Zheng et al., 2018).
Proof. We enforce the acyclicity constraint on W by restricting the optimization domain to Ω = {W : h(W ) = 0}. Then, the maximum likelihood objective can be written as∑
n log p(xn) = ∑ n log pz(x n −Wxn) + log |det(I −W )| (change of variable) (24)
= ∑ n log pz(x n −Wxn) (lemma 1) (25)
= − ∑ n ‖xn −Wxn‖2 (pz is std Gaussian), (26)
where eq. (24) is valid because the transformation z = g(x) = x−Wx is invertible for anyW ∈ Ω. We see that we recover exactly the Notears objective from eq. (23).

heading: D PROOF OF LEMMA 1
text: We split the proof in several simple steps.
1. gA(x) = x − fA(x) has Jacobian-determinant det(I − JA(x)), where JA(x) is the Jacobian of fA(x). 2. JA(x) has non-zero entries exactly in the positions where A> is non-zero. Therefore, it retains the DAG structure.
3. Matrices with a DAG structure are nilpotent (i.e., all eigenvalues are zero). Thus, JA(x) can be factorized as JA(x) = QUQ∗, where Q is unitary and U is strictly upper triangular (Schur factorization).
4. Finally, det(I − fA(x)) = det(I −QUQ∗) = det(I − U) = 1.

heading: E DAGS AND INVERTIBILITY
text: As mentioned in section 4, if A is not a DAG the transformation z = gA(x) = x − fA(x) may not be invertible. A simple example suffices to show this. Consider the two-variable scenario with
A = [ 0 1 1 0 ] and fA(x) = [f1(x2), f2(x1)], where f1(x2) = x2 and f2(x1) = x1. Then, computing z = x− fA(x) we get
z1 = x1 − x2 z2 = x2 − x1.
Clearly, the function z = gA(x) = x− fA(x) is not invertible, since for any pair (x1, x2) satisfying x1 = x2 we get z1 = z2 = 0.

heading: F RESULTS WITH HIGHER LEVEL OF NOISE
text: This section shows results in the exact same setting as the one described in section 6, with the difference being that synthetic datasets are generated with a higher level of noise (the real and pseudo-real datasets are exactly the same as before). This is achieved by sampling the noise scale σi ∼ U(0.2, 2) (instead of σi ∼ U(0.2, 1)). Results are shown in fig. 2. As expected, as a consequence of the higher level of noise, all methods perform slightly worse. However, the main conclusion remains unchanged, FCause outperforms other methods and is still competitive even when ran on the datasets with 30% of missing values.

heading: G FULL RESULTS WITH STANDARD DEVIATIONS
text: Results for all metrics, including standard deviations, are shown in fig. 3. It can be observed that some methods achieve higher precision values than FCause for some datasets. Looking at the results in detail, this is not surprising. The methods that achieve better precision values always have extremely low recall. Meaning that they detect very few edges, which tend to be correct. This leads to quite high precision values, but very low recall, F1 and causal accuracy. That is, their overall performance is quite poor, since they fail to detect most edges. (e.g. Suppose that a problem has 100 edges. A method that detects just one edge, and that edge is correct, gets a precision of one, but a recall and F1 of almost zero). On the other hand, while FCause tends to achieve slightly lower precision values, its overall performance is significantly better.

heading: H MISSING VALUES IMPUTATION
text: Datasets. We use the same datasets described in section 6, all standardized. We remove 30% of the values in the training set randomly. The testing sets use are composed of n = 500 samples for the
synthetic datasets, n = 100 samples for the pseudo-real data obtained with the SynTReN generator, and n = 53 samples for the protein cells dataset.
Baselines. We compare FCause against mean imputing, mice (Van Buuren & Groothuis-Oudshoorn, 2011), missforest (Stekhoven & Bühlmann, 2012) and PVAE (Ma et al., 2018). We also include as baseline FCause with an adjacency matrix fixed to a strictly upper triangular. We call this method UTFlow.
Metrics. We report the normalized Root Mean Squared Error (NRMSE) computed on the test set. All results are computed as the mean across four different random seeds.
Results are shown in fig. 4. For clarity, all values reported are the mean across four runs with different random seeds. There is no single method that dominates across all datasets. Overall, it can be observed that Missforest achieves the best performance, followed by FCause and UTFlow, and finally Mice. PVAE, while better than mean imputing, performs worse than other methods in most datasets. Perhaps surprisingly, UTFlow and FCause achieve very similar performances.

Review summaries:
This paper proposes a general flow-based approach to learn DAGs from data which provides a unified view of existing continuous optimization methods for structure learning. As a side benefit, the authors demonstrate that the proposed method could naturally be modified to handle missing data. The authors provide empirical studies to show that their proposed method outperforms the other baselines.

This paper proposes to combine the continuous optimization-based causal discovery approach from notears with flow-based function learning. An extension is given to data missing (completely) at random. Also, a unifying framework is presented facilitating the comparison and exchange of ideas between different continuous optimization-based causal discovery methods, viewing all in terms of flows.