D
IFFERENTIABLE DAG S AMPLING
Anonymous authors
Paper under double-blind review
A BSTRACT
We propose a new differentiable probabilistic model over DAGs (DP-DAG).
DPDAG allows fast and differentiable DAG sampling suited to continuous optimization.
To this end, DP-DAG samples a DAG by successively (1) sampling a linear ordering of the node and (2) sampling edges consistent with the sampled linear ordering.
We further propose VI-DP-DAG, a new method for DAG learning from observational data which combines DP-DAG with variational inference.
Hence,
VI-DP-DAG approximates the posterior probability over DAG edges given the observed data.
VI-DP-DAG is guaranteed to output a valid DAG at any time during training and does not require any complex augmented Lagrangian optimization scheme in contrast to existing differentiable DAG learning approaches.
In our extensive experiments, we compare VI-DP-DAG to other differentiable DAG learning baselines on synthetic and real datasets.
VI-DP-DAG signiﬁcantly improves DAG structure and causal mechanism learning while training faster than competitors.
I NTRODUCTION
Directed Acyclic Graphs (DAGs) are important mathematical objects in many machine learning tasks.
For example, a direct application of DAGs is to represent causal relationships in a system of variables.
In this case, variables are represented as nodes and causal relationships are represented as directed edges.
Hence, DAG learning has found many applications for causal discovery in biology, economics or planning (Pearl, 1988; Ramsey et al., 2017; Sachs et al., 2005; Zhang et al., 2013).
However,
DAG learning is a challenging problem for two reasons.
First, while DAG learning with data from randomized and controlled experiments is the gold-standard for causal discovery, experimental data might be hard or unethical to obtain in practice.
Hence, a more common but also challenging setting is
DAG learning from observational data which is possible under proper conditions (Pearl, 2009; Spirtes et al., 2000).
Second, the number of possible DAGs scales super-exponentially with the number of variables which makes DAG learning an NP hard problem (Chickering et al., 2012; Robinson, 1973).
A ﬁrst traditional family of models for DAG learning are discrete score-based approaches.
These approaches aim at solving the following discrete optimization problem:
[Equation]
where X denotes the observed data and the discrete DAGs space is composed of DAGs with unweighted (present or absent) edges.
Examples of score functions are Bayesian Information Criteria (BIC) (Chickering & Heckerman, 1997) or Minimum Description Length (MDL) (Bouckaert, 1993).
Discrete approaches have two main limitations: (1) the optimization search space of discrete DAGs is large and constrained which often makes the problem intractable without further assumptions, and (2) the learning procedure is not differentiable and thus not amenable to gradient-based optimization, as done by most deep learning approaches.
To mitigate these issues, a second more recent family of models for DAG learning proposes to leverage continuous optimization by using an augmented
Lagrangian formulation (Lachapelle et al., 2020; Ng et al., 2019; Wehenkel & Louppe, 2021; Yu et al., 2019; Zheng et al., 2018).
These approaches aim at solving the following optimization problem:
[Equation]
where h ( G ) is a smooth function over weighted graphs G which is equal to zero when the graph G satisﬁes the DAG constraints.
Standard examples of score functions are (negative) cross-entropy (Yu et al., 2019) or Mean Squared Error (MSE) (Ng et al., 2019) when reconstructing the data.
During
[Figure]
optimization, G is a weighted graph where the discreteness and acyclicity constraints are relaxed.
Hence, these continuous optimization approaches have two main limitations: (1) the augmented
Lagrangian optimization is computationally expensive as it requires multiple complex dual ascent iterations, and (2) the discrete and acyclicity constraints are relaxed during optimization which does not guarantee valid discrete DAGs without non-differentiable pre- and post-processing as proposed by Causal Additive Model (CAM) (Bühlmann et al., 2014).
For a more comprehensive description, we recall the augmented Lagrangian optimization method in detail in App. A.
In this paper, we focus on differentiable DAG learning methods and make the following contributions:
• We propose a new probabilistic model over DAGs (DP-DAG) which is capable of fast and differentiable sampling.
DP-DAG can be implemented in few lines of code using GumbelSinkhorn, Gumbel-Top-k and Gumbel-Softmax distributions to parametrize differentiable sampling over permutations and edges (see Fig. 1).
• We propose a new method for DAG learning from observational data (VI-DP-DAG) which instantiates a general probabilistic formulation for DAG learning with DP-DAG and variational inference.
VI-DP-DAG guarantees valid DAG outputs at any time during training.
• We show in our experiments on established synthetic and real datasets that DP-DAG outperforms other differentiable DAG learning baselines for DAG structure and causal mechanisms learning while training one order of magnitude faster.
R ELATED W ORK
We differentiate between three types of DAG learning approaches: the discrete optimization approaches, the continuous optimization approaches and the sampling-based approaches.
We refer to the survey (Vowels et al., 2021) for a more detailed overview of DAG learning approaches.
Discrete optimization.
First, to make the search space more tractable, discrete optimization approaches modify the original problem with additional assumptions on DAG treewidth (Nie et al., 2014; Scanagatta et al., 2016), ancestral constraints (Chen et al., 2016) or on the number of parents of each variable (Viinikka et al., 2020).
Other methods are based on greedy search (Chickering, 2002) or discrete optimization of the topological order (Park & Klabjan, 2017; Scanagatta et al., 2015; Teyssier & Koller, 2005).
Another type of discrete optimization approaches are constraint-based methods.
These methods explore the discrete DAG space by performing independence tests between observed variables (Bühlmann et al., 2014; Spirtes et al., 2001).
Continuous optimization.
Second, continuous approaches usually relax the discreteness and acyclicity constraints by using an augmented Lagrangian formulation of the optimization problem (Lachapelle et al., 2020;
Ng et al., 2019; Wehenkel & Louppe, 2021;
Yu et al., 2019; Zheng et al., 2018).
Some approaches deﬁne the DAG structure from neural network weights (Lachapelle et al., 2020; Zheng et al., 2018) while other approaches directly learn the DAG adjacency matrix (Ng et al., 2019; Wehenkel & Louppe, 2021; Yu et al., 2019).
In contrast to these methods, VI-DP-DAG learns a probabilistic model over the DAG structure.
Further, these approaches penalize DAG constraints violation in the augmented Lagrangian formulation but do not guarantee that they are fulﬁlled during training.
Recently, Yu et al.
(2021) propose to complement the augmented Lagrangian optimization with a second step projecting the learned graph on admissible solutions.
Hence, contrary to VI-DPDAG, most of these approaches use non-differentiable processing steps – e.g. removing cycles and spurious edges – to output valid and high-quality DAGs.
Common examples of processing steps are
Preliminary Neighbors Selection (PNS) and CAM pruning (Bühlmann et al., 2014).
Sampling-based optimization.
Third, other works use DAG sampling to estimate the posterior distribution over DAGs with MCMC (Kuipers et al., 2020; Niinimäki et al., 2011; 2016;
Talvitie et al., 2020; Viinikka et al., 2020).
While previous works improve the quality and speed of MCMC computations by sampling (partial) orders or making assumptions on the number of parents per node, they are still computationally extremely expensive (Kuipers et al., 2020; Niinimäki et al., 2011; 2016; Viinikka et al., 2020).
E.g., Viinikka et al.
(2020) recommend to run MCMC methods during 12 hours to sample from the posterior distribution over DAGs with 100 nodes.
In contrast,
VI-DP-DAG approximates the posterior distribution over DAG edges with variational inference and can sample very fast.
E.g., our VI-DP-DAG trains in around 190 seconds and samples in less than 1 second for a DAG with 100 nodes.
Further, while the construction of the MCMC chains are generally non-differentiable, our DP-DAG is capable of fully differentiable DAG learning and can leverage gradient-based optimization.
Other works propose optimization of discrete problems using differentiable probabilistic distribution over various discrete objects like subsets or spanning trees but not on DAG structures (Grathwohl et al., 2021; Karalias & Loukas, 2020; Paulus et al., 2020).
Further, recent works combine differentiable edge sampling with Gumbel trick and Lagrangian optimization but do not deﬁne valid distributions over the full DAG structure (Brouillard et al., 2020; Ng et al., 2019).
In contrast, DP-DAG does not require complex Lagrangian optimization and guarantees valid
DAGs solutions at any time during training.
Finally, Grosse et al.
(2021) explores an orthogonal direction where the search space in sequential decision making problems is represented with a DAG.
P ROBABILISTIC M ODEL OVER DAG S
A Directed Acyclic Graph (DAG) is a graph G =
( V, E ) with n nodes x 1 , ..., x n and m directed edges which does not exhibit directed cycles.
A DAG always admits a valid permutation (or linear ordering) π : [[1 , n ]] →
[[1 , n ]] of the nodes such that a node cannot have a direct edge toward a node with lower rank i.e., π ( i ) <
π ( j ) implies no directed edge from node x π ( j ) to node x π ( i ) .
Valid permutations are often not unique.
Interestingly, this property has an equivalent matrix formulation:
Theorem 1.
Lets A ∈ { 0 , 1 } n be the (a priori non-triangular) adjacency matrix associated with an arbitrary node labelling of the DAG G .
The adjacency matrix A always admits a permutation matrix Π ∈ { 0 , 1 } n × n and an upper triangular matrix U ∈ { 0 , 1 } n × n such that A = Π T UΠ .
The permutation matrix Π directly corresponds to a valid component-wise permutation π .
Hence,
Th.
1 simply states that the matrix U is the new adjacency matrix where the new node labels are a valid permutation of the original node labels i.e. A ij =
U π ( i ) π
( j ) such that U π ( i ) π
( j )
=
if π ( i ) < π ( j ) .
The decomposition of the DAG adjacency matrix A in Th.
1 is generally not unique as a DAG G generally admits multiple valid linear permutations Π .
Hence, we deﬁne the following probabilistic model over DAGs (DP-DAG) based on the adjacency matrix decomposition in Th.
[Equation]
where P ( G ) is the set of valid permutation matrices for the DAG G , U n is the space of binary upper-triangular matrices of size n × n , P ( Π ) is the distribution over permutations and P ( U ) is the distribution over edges consistent with the sampled permutation of the nodes.
Note that the number of valid permutations |P
( G ) | can be exponentially large in the number of nodes which makes the exact computation of the probability of a given DAG adjacency matrix P ( A ) intractable for large graphs.
However, DAG sampling does not require any enumeration of the valid linear permutations.
Indeed, we propose a new differentiable DAG sampling method (i.e. A ∼ P φ,ψ ( A ) ) based on differentiable edge sampling (i.e. U ∼ P φ ( U ) ) and differentiable permutation sampling (i.e. Π ∼ P ψ ( Π ) ).
The variables ψ and φ denote the parameters of the edge and permutation distributions.
Differentiable edge sampling.
The Bernoulli distribution is a well-suited distribution to model randomness over a discrete binary variable like an edge U ij ∈ { 0 , 1 } ∼ Ber( p ) where p ∈ [0 , 1] is the probability for the edge to exist.
Unfortunately, standard random sampling operations from the
Bernoulli distribution are not differentiable.
In contrast, the Gumbel-Softmax distribution allows for differentiable sampling and approximates the Bernoulli distribution (Jang et al., 2017).
The GumbelSoftmax distribution is deﬁned on continuous variables, i.e. ˆ U ij ∈ [0 , 1] ∼ Gumbel-Softmax τ ( φ ) with φ ∈ [0 , 1] , where the temperature parameter τ allows to interpolate between a one-hot-encoded categorical distribution ( τ → 0 ) and continuous categorical densities ( τ → + ∞ ).
For differentiable sampling, we can use the straight-through estimator (Bengio et al., 2013): we use the discrete variable U ij = arg max[1 − ˆ U ij , ˆ U ij ] in the forward pass, and the continuous approximation ˆ U ij in the backward pass.
Thus, sampling all the upper triangular indices of U ∈ { 0 , 1 } n × n has complexity
O ( n 2 ) .
We recall the deﬁnition of the Gumbel-Softmax distribution in detail in App.
B.1.
Differentiable permutation sampling.
Similarly to an edge, a permutation Π is discrete, making differentiable sampling challenging.
We describe two alternative methods which allow for differentiable permutation sampling.
First, the Gumbel-Sinkhorn (Mena et al., 2018) is deﬁned on a continuous relaxation of the permutation matrix, i.e. ˆ Π ∈ [0 , 1] n × n ∼ Gumbel-Sinkhorn τ ( ψ ) with ψ ∈ [0 , 1] n × n , where the temperature parameter τ also allows to interpolate between discrete and continuous distributions similarly to the Gumbel-Softmax distribution.
For differentiable sampling, we can use the straight-through estimator (Bengio et al., 2013): we use the discrete permutation Π =
Hungarian ( ˆ Π ) by applying the Hungarian algorithm (Munkres, 1957) to compute a discrete permutation in the forward pass, and the continuous approximation ˆ Π in the backward pass.
Sampling a permutation matrix Π ∈ { 0 , 1 } n × n is dominated by the Hungarian algorithm and has a complexity of O ( n 3 ) .
We recall the deﬁnition of the Gumbel-Sinkhorn distribution in detail in App.
B.2.
A second method orthogonal to the Gumbel-Sinkhorn method is to use the combination of the
Gumbel-Top-k trick (Kool et al., 2019) and the SoftSort operator (Prillo & Eisenschlos, 2020) which also deﬁnes a distribution on a continuous relaxation of the permutation matrix.
For k = n , the
Gumbel-Top-n distribution states the sorted perturbed log-probabilities, i.e. π = Sort ( ψ + G ) where parameters ψ are log-probabilities and G ∈ R n are i.i.d.
Gumbel noise, deﬁnes a distribution over component-wise permutation (a.k.a.
linear ordering without replacement).
Instead of the Sort operator, we apply the SoftSort operator to the perturbed log-probabilities which outputs a continuous relaxation of the permutation matrix , i.e. ˆ Π = SoftSort ( ψ + G ) ∈ R n × n .
For differentiable sampling, we use the straight-through estimator (Bengio et al., 2013): we use the discrete permutation Π = arg max ˆ Π by applying the (one-hot) argmax operator
row-wise (Prillo & Eisenschlos, 2020) in the forward pass, and the continuous approximation ˆ Π in the backward pass.
Sampling a permutation matrix Π ∈ { 0 , 1 } n × n is dominated by the SoftSort operation and has a complexity of O ( n 2 ) .
The permutation sampling complexity with Gumbel-Top-k combined with SoftSort is thus lower than the permutation sampling complexity with Gumbel-Sinkhorn.
We recall the deﬁnition of the
Gumbel-Top-k distribution and SoftSort operator in detail in App.
B.3 and App. C.
Differentiable DAG sampling.
Given the aforementioned methods for differentiable edge and permutation sampling, we propose a new simple and valid sampling procedure for DAG sampling which consists in three steps: (1) Sample a permutation Π from a probabilistic model over permutations
P ψ ( Π ) i.e. Π ∼ P ψ ( Π ) .
(2) Sample an upper triangular matrix U by sampling the upper triangular elements from a probabilistic model over edges P φ ( U ij ) i.e. U ij ∼ P φ ( U ij ) .
(3) Compute the ﬁnal adjacency matrix A from the permutation matrix Π and the upper triangular matrix U i.e. A = Π T UΠ .
This procedure is capable of sampling any possible DAGs of n nodes due to Th.
In practice, we propose to parametrize the distribution P ψ ( Π ) using the Gumbel-Sinkhorn or the Gumbel-Top-k trick which deﬁne valid distributions over permutations, and parametrize the distributions P φ ( U ij ) using the Gumbel-Softmax trick which deﬁnes a valid distribution over edges.
Given these parametrizations, the sampling procedure allows fast and differentiable sampling and can be implemented in a few lines of code (see Fig. 1).
The total DAG sampling complexity is dominated by the permutation sampling step which has a complexity of O ( n 3 ) using Gumbel-Sinkhorn and O ( n 2 ) using Gumbel-Top-k combined with SoftSort.
Finally, the DAG sampling procedure of DP-DAG guarantees a valid DAG output without additional pre- or post-processing at any time during training.
V ARIATIONAL DAG L EARNING FROM O BSERVATIONAL D ATA
Structural Equation Model (SEM).
We assume that the data follow a SEM i.e. x i =
f ∗ i ( x pa ∗ ( i ) , (cid:15) i ) where (cid:15) i are independent noises and f ∗ i is the ground-truth causal mechanism which deterministically transforms the parent causes x pa ∗ ( i ) and noise (cid:15)
i in the consequence value x i .
We assume that all variables are observed.
We aim at learning both the true parents pa ∗ ( i ) for each node i describing the direct cause-consequence relationship, and the true causal mechanisms f ∗ i .
Alternatively, a SEM can equivalently be written in a matrix form.
Indeed, one can use the DAG adjacency matrix A ∗ as a mask before applying the causal mechanism f ∗ i i.e. X i
= f ∗ i
( A ∗ i ◦ X , (cid:15) i ) .
Similar masking formulations of a SEM have been used in previous works (Brouillard et al., 2020;
Ng et al., 2019).
Probabilistic DAG learning loss.
We propose a new general formulation for DAG learning based on differentiable DAG sampling which intuitively aims at maximizing the expected score
E G [ score ( X , G )] under the the probability distribution over DAGs P φ,ψ ( G ) i
[Equation] [Equation]
This formulation allows a rich probabilistic solution that assigns a conﬁdence score to all possible
DAGs.
In contrast, the discrete and continuous DAG learning formulations in Eq.
1 and Eq.
2 only compute a single point estimate DAG solution and do not model any uncertainty on the ﬁnal DAG output.
A speciﬁc instance of the Eq.
4 is the optimization of the following ELBO loss:
[Equation]
where λ ∈ R + is a regularization factor and θ, φ, ψ are the parameters of the model to optimize.
Indeed, similarly to Eq. 4, both terms (i) and (ii) can be formulated as an expectation over the DAG probabilistic model P φ,ψ ( A ) .
Importantly, the optimum of the variational inference problem in Eq.is reached when the learned probability distribution over DAGs is equal to the posterior probability distribution over DAGs i.e. P ( A | φ, ψ ) ≈ P ( A |D ) where D denotes the dataset of observations.
Variational inference with DP-DAG.
We propose VI-DP-DAG, a new method combining DP-DAG and variational inference to learn the matrix form of a SEM from observational data.
At training time,
VI-DP-DAG consists of three steps: (1) It differentiably samples a valid DAG adjacency matrix A from a probabilistic model over DAGs P φ,ψ ( A ) .
In practice, we parametrize P φ,ψ ( A ) with DP-DAG.
(2) It applies the n transformations f i,θ on the observations X masked with the sampled DAG
A i.e. ˆ
X i =
f i,θ ( A i ◦ X , (cid:15) i ) .
In practice, we parametrize f
i,θ with neural networks.
(3) It jointly updates all parameters θ, φ, ψ by maximizing at each iteration an approximation of the ELBO loss in Eq.
In practice, we approximate the term (i) by sampling a single DAG matrix A at each iteration and assume a Gaussian distribution with unit variance around ˆ X (i.e. (i) =
||
X − ˆ
X
|| 2 ).
We compute the term (ii) by setting a small prior P prior ( U ij ) on the edge probability (i.e. (ii) =
(cid:80) ij KL ( P φ ( U ij ||
P
prior ( U ij )) )
thus acting as a natural sparsity regularization.
We set no prior on the permutation probability for two reasons:
In theory, a permutation prior is likely biased toward graphs that are compatible with a larger number of orderings (Kuipers & Moffa, 2017; Viinikka et al., 2020).
In practice, the closed-form computation of the permutation probability P θ ( Π ) is generally intractable (Mena et al., 2018).
Thus,
VI-DP-DAG approximates the true posterior distribution over DAG edges probability only, which achieves excellent predictive performance in practice.
Beyond variational inference, VI-DP-DAG is theoretically motivated from a second perspective.
Assuming that the data comes from a SEM with additive noise, Ng et al.
(2019) showed that minimizing the term (i) enforces the sampled matrix
A to represent a super-graph of the ground-truth DAG adjacency matrix A ∗ (i.e. all the edges in
A ∗ are also in A ) thus suggesting to add a sparsity regularization term to remove spurious edges.
Interestingly, the term (ii) – which arises naturally from the variational inference framework – indeed acts a a sparsity regularizer and pushes the probability P φ ( U ij ) to the ﬁxed low prior value P prior ( U ij ) .
At inference time, we can use VI-DP-DAG in two ways.
On one hand, we can extract the cause/consequence relationships from, e.g., one sampled DAG adjacency matrix A .
On the other hand, we can estimate the missing values x i from its learned parents and causal mechanisms i.e. x i
≈
ˆ x i =
f i,θ ( x pa ( i )) .
Intuitively, a small prediction error
|| x i − ˆ x
i || 2 indicates that the learned parents x pa ( i ) forecast well the child node x i i.e. the parent nodes x pa ( i ) are Granger-causes of the child node x i .
E XPERIMENTS
In this section, we compare VI-DP-DAG to existing methods for differentiable DAG learning on extensive experiments including undirected and directed link prediction, consequence prediction from learned causal mechanisms, and training time.
This set-up aims at highlighting the speed and quality of VI-DP-DAG for DAG structure learning and causal mechanisms learning.
Further, we evaluate the speed and the optimization quality of the probabilistic model DP-DAG alone.
5. S ET - UP
In our experiments, VI-DP-DAG parametrizes the permutation probability P ψ ( Π ) with GumbelSinkhorn or Gumbel-Top-k trick, the edge probability P φ ( U ) with Gumbel-Softmax distribution and the causal mechanisms f
i,θ with a 3 layers Multi-Layer Perceptron (
MLP
We use early stopping and perform a grid search over the permutation probability parametrization (i e. Gumbel-Sinkhorn or
Gumbel-Top-k), the ﬁxed prior probability P prior ( U ij ) ∈
[1 e − 2 , 1 e − 1 ] and the regularization factor λ ∈ [0 , 1 e − 1 ] .
Finally, all temperature parameters are ﬁxed to τ = 1 in all experiments
Baselines.
We focus our evaluation on differentiable methods for DAG learning and compared to
DAG-GNN (Yu et al., 2019), DAG-NF (Wehenkel & Louppe, 2021), GraN-DAG (Lachapelle et al., 2020) and Masked-DAG (Ng et al., 2019)
by using their ofﬁcial implementation when available.
Unless otherwise stated, we did not use non-differentiable pre- or post-processing steps for any of the models.
For the sake of completeness, we also compare DP-DAG with models using nondifferentiable pre- and post-processing like Preliminary Neighbor Selection (PNS) and CAM pruning (Bühlmann et al., 2014) in App.
F.3.
Similarly to VI-DP-DAG, we perform a grid-search for the hyper-parameters of all models (including learning rate, number of hidden units per linear layer, or regularization factors).
For all metrics, we report the mean and the standard error of the mean of the results over 10 runs.
We give further details on the models in App.
E.
Datasets.
We use a total of 11 synthetic and real datasets.
For the synthetic datasets, we closely follow (Lachapelle et al., 2020; Ng et al., 2019; Peters et al., 2014).
The graph structures A ∗ were generated with the Erdös-Rényi (ER) or Scale-Free (SF) network models, and the causal mechanisms f ∗ i were generated from a (non-linear) Gaussian Process with RBF kernel of bandwidth one and with independent zero-mean Gaussian noises (cid:15)
i .
We considered different sizes of graphs including number of nodes n in { 10 , 50 , 100 } and number of edges m in { 10 , 40 , 50 , 100 , 200 , 400 } .
We denote the synthetic datasets by SF/ER- n - m .
We used 10 sampled datasets per setting.
For the real datasets, we closely follow (Koller & Friedman, 2009; Lachapelle et al., 2020; Ng et al., 2019;
Yu et al., 2021; Zheng et al., 2018).
We use the Sachs dataset which measures the expression level of different proteins and phospholipids in human cells (Sachs et al., 2005).
We also use the pseudo-real SynTReN dataset sampled from a generator that was designed to create synthetic transcriptional regulatory networks and produces simulated gene expression data that approximates experimental data (Van den
Bulcke et al., 2006).
For all these datasets, we can access the ground-truth DAG adjacency matrices
A ∗ used to generate the data or estimated by experts.
We split all datasets in training/validation/test sets with 80% / 10% / 10% .
We give further details on the datasets in Sec.
D in the appendix.
5. R ESULTS
DAG sampling.
We show results for the DAG sampling time in
Fig. 2 and Fig.
3 in the appendix.
While both Gumbel Sinkhorn and Gumbel-Top-k parametrizations can sample a DAG with 200 nodes in less than 1 second, Gumbel-Top-k is signiﬁcantly faster than Gumbel-Sinkhorn as expected from the complexity analysis.
Further, we show the results of differentiable sampling optimization using DP-DAG for DAG learning when the ground-truth DAG is observed in Tab. 10 in the appendix.
On this toy task, DP-DAG recovers almost perfectly the groundtruth DAG structure.
These observations validate that DP-DAG is reliable method for fast and differentiable DAG sampling.
Further details about these two experiments are in App. F.1.
DAG structure.
We evaluate the learning of the DAG structure by comparing the ground-truth adjacency matrix A ∗ ij to the learned score S ij .
For baselines, we use entries from the weighted adjacency matrix S ij =
A ij as scores for directed edges, and S ij =
A ij + A ji as scores for undirected edges.
For DP-DAG, we use the edge probability S ij =
P φ ( U π ( i ) π
( j ) ) as scores for directed edges and S ij =
P φ ( U π ( i ) π
( j ) )
+ P φ ( U π ( j ) π ( i ) ) as scores for undirected edges.
In this case, the permutation π (or equivalently Π ) is deterministically computed by removing Gumbel noise in the forward sampling step.
The directed scores (Dir-) indicate if an edge with a speciﬁc direction exists between two nodes.
The undirected scores (Un-) indicates if an edge without a speciﬁc direction exists between two nodes.
We compare the directed and undirected scores to the ground-truth binary adjacency matrix by using the area under the curve of precision-recall (AUC-PR) and the area under the receiver operating characteristic curve (AUC-ROC) .
These metrics have the important advantage to be independent of a threshold choice (Vowels et al., 2021).
Thus, AUC-PR and AUC-ROC are better indicators of the true performance of the models contrary to other structure metrics like Structural Hamming Distance which needs to arbitrary select a threshold.
We show the results for differentiable DAG structure learning on synthetic and real datasets in Tab.and in Tab. 7 in the appendix.
In total, VI-DP-DAG achieves 27/44 best scores and second-best scores otherwise.
On synthetic datasets, only GraN-DAG competes with VI-DP-DAG on small datasets with 10 nodes.
DP-DAG outperforms all baselines on all larger synthetic datasets with 50 or 100 nodes for both ER and SF graph types.
E.g. on the ER-100-400 dataset, VI-DP-DAG gives an absolute improvement of +17% for Un-AUC-PR, +25% for Un-AUC-ROC, +13% for Dir-AUC-PR and +9% for Dir-AUC-ROC.
We observe that DAG-GNN and DAG-NF perform comparably to random guessing corresponding to an AUC-ROC of 50% on those datasets.
We hypothesize that their poor performances might be due to their incorrect modeling assumptions or their incapacity to model this type of causal relationships as mentioned in Lachapelle et al.
Yu et al.
On real datasets, no baseline consistently outperforms VI-DP-DAG which obtains 4/8 best scores and second-best scores otherwise.
We observe on both synthetic and real datasets that VI-DP-DAG performs better with undirected edges scores than with directed edges scores indicating that VI-DP-DAG might invert the edge directions.
This can be intuitively explained by the presence of node pairs where x i is an equally good predictor of x j as x j is a good predictor of x i .
Finally, we show further evidence that
VI-DP-DAG learns reliable DAG structures under perturbations in App. F.5.
Indeed it conﬁdently assigns lower conﬁdence scores to perturbed versions of the ground-truth DAG.
We show the results for DAG structure learning with non-differentiable pre- and post-processing in
Tab. 11 and Tab. 12 in the appendix.
In these experiments, we apply the default pre-processing PNS and post-processing CAM pruning steps to CAM (Bühlmann et al., 2014), GraN-DAG (Lachapelle et al., 2020),
Masked-DAG (Ng et al., 2019) as suggested by the original papers.
We similarly apply CAM pruning to the ﬁnal results of VI-DP-DAG.
In this case, we observe that the greedy CAM is the best performing model for most synthetic datasets while VI-DP-DAG maintains competitive performance.
Interestingly, VI-DP-DAG with non-differentiable pre- and post-processing is the best performing model on real datasets (i.e. 7/8 best scores).
Further, we notice that the pre- and post-processing steps raise signiﬁcantly the scores of GraN-DAG and Masked-DAG which becomes more competitive with VI-DP-DAG, while VI-DP-DAG beneﬁts from it to a lower extent (e.g. see
ER-100-400 in Tab. 1 and Tab. 11).
This suggests that methods based on augmented Lagrangian optimization are more dependent on additional non-differentiable processing steps than VI-DP-DAG.
Lachapelle et al.
(2020) make similar observations on the dependence on PNS and CAM pruning steps.
Causal mechanisms.
We evaluate the causal mechanisms learning by computing the Mean Squared
Error (MSE) between the ground-truth node values x i and the estimated node values ˆ x i
= f
i,θ ( x pa ( i )) on a new test set.
It can be computed in a vectorized way i.e. || X − ˆ
X
|| 22 where ˆ X i
=
f i,θ ( A i ◦ X ) .
This task is particularly important when predicting missing values from the values of known parent causes.
For this experiment, we compare to other baselines which jointly learn the adjacency matrix A and the causal mechanisms f
i,θ explicitly.
Hence, we compare to Masked-DAG and GraN-DAG.
At testing time, we remove the remaining cycles from the adjacency matrix A learned with augmented
Lagrangian optimization to guarantee a valid DAG structure as suggested by Lachapelle et al.
Ng et al.
Further, we also compare to GT-DAG which uses the ground-truth DAG adjacency matrix
A ∗ at training and testing time.
GT-DAG can be considered as a lower bound for the MSE scores.
It indicates if other methods manage to jointly learn well both the DAG structure A and the causal mechanisms f i,θ .
For DP-DAG, the adjacency matrix A is computed by deterministically computing π (or equivalently Π ) without Gumbel noise in the forward sampling step, and second keeping the edges such that P φ ( U π ( i ) π ( j ) ) > . 5 .
These edges intuitively correspond to parent relationships which are more likely to be present than absent according to the model.
We show the results for causal mechanisms learning in Tab. 2.
Overall, DP-DAG is the most competitive model and achieves 9/11 best scores.
Similarly to DAG structure learning, GraN-DAG is competitive for small graphs with 10 nodes but VI-DP-DAG outperforms all baselines for most larger graphs with 50 and 100 nodes.
Interestingly, DP-DAG brings a particularly signiﬁcant improvement on real datasets compared to GraN-DAG and Masked-DAG with scores very close to GT-DAG.
This suggests that the potentially noisy ground-truth DAGs estimated by experts on these real datasets have an equivalent predictive power as the DAG learned by VI-DP-DAG.
This aligns well with the intuition that VI-DP-DAG learns Granger causal relationships.
Finally, we evaluate in App.
{+F.
+}
the [Table] impact of the threshold choice t on the edge probabilities P φ ( U π ( i ) π ( j ) ) > t .
We observe that the
MSE score of DP-DAG achieves stable high performance regardless of the threshold value.
Training time.
We evaluate the training time of all models on a single GPU (NVIDIA GTX 1080 Ti, 11 GB memory).
Training time is particularly relevant for DAG structure learning since predicting the adjacency matrix A from a new dataset represents a full training.
Thus, real-world applications of
DAG learning generally aim at fast training time.
We show the results in Tab. 3.
The training of VI-DP-DAG is one order of magnitude faster than all models on all datasets.
Hence, assuming that training time is typically linearly correlated with compute cost, energy consumption and CO 2 emission, the fast training DP-DAG signiﬁcantly improves the applicability of DAG learning methods.
In particular, VI-DP-DAG is from × 5 to × 18 faster to train than GraN-DAG which is the second-best performing model for DAG structure learning and causal mechanisms learning.
The training speed difference can be explained by the different optimization schemes.
Indeed, Lagrangian optimization typically requires solving T ≈ 10 successive optimization problems (Zheng et al., 2018) while VI-DP-DAG only requires solving a single optimization problem.
Finally, by aggregating training and sampling time, VI-DP-DAG requires around 191 seconds to sample a DAG with 100 nodes from the learned approximate posterior distribution.
In contrast, other sampling methods based on MCMC sampling require running the MCMC methods for 12 hours to sample a DAG with 100 nodes (Viinikka et al., 2020).
[Table] [Table]
C ONCLUSION
We propose a Differentiable Probabilistic model over DAGs (DP-DAG) which allows fast and differentiable DAG sampling and can be implemented in few lines of code.
To this end, DP-DAG uses differentiable permutation and edges sampling based on Gumbel-Sinkhorn, Gumbel-Top-K and
Gumbel-Softmax tricks.
We propose VI-DP-DAG, a new method combining variational inference and DP-DAG for DAG learning from observational data with continuous optimization.
VI-DP-DAG guarantees valid DAG prediction at any time during training.
In our extensive experiments, VI-DPDAG performs favorably to other differentiable DAG learning baselines for DAG structure and causal mechanisms learning on synthetic and real datasets while training one order of magnitude faster
[Table]
E THICS S TATEMENT
The Assessment List for Trustworthy AI (ALTAI) (Com., 2020) includes reliability, transparency and
Environmental well-being.
Accurate and fast causal discovery from observational data is important to give meaningful and interpretable predictions at a low energy cost.
Therefore, VI-DP-DAG brings a signiﬁcant improvement regarding these values by performing high-quality DAG structure and causal mechanisms predictions at a low training cost.
While DP-DAG achieves high-performance for DAG structure learning and causal mechanisms learning, there is always a non-negligible risk that DP-DAG does not fully capture the real-world complexity.
Thus, we raise awareness about the risk of excessive trust in causal predictions from
Machine learning models.
This is crucial when causal predictions are applied in domains subject to privacy or fairness issues (e.g. ﬁnance, medicine, policy decision making, etc).
In particular,
VI-DP-DAG does not provide any guarantee on the causal interpretation of the edges in the ﬁnal predicted DAG.
Therefore, we encourage practitioners to proactively confront the model predictions for DAG structure and causal mechanisms to desired behaviors in real-world use cases.
R EPRODUCIBILITY S TATEMENT
We provide all datasets at the following anonymous link 1 .
In App.
D in the appendix, we give a detailed description for each dataset used in this paper.
This description includes the task description, the dataset size, the number of nodes and edges and the train/validation/test splits used in the experiments.
We further explain the hardware used for the experiments.
We provide the model code at the following anonymous link 2 .
In App.
E in the appendix, we give a detailed description of the architecture and grid search performed for each model used in this paper.
We describe the metrics used in the experiments in App. 5.
A A UGMENTED L AGRANGIAN O PTIMIZATION
In this section we recall the augmented Lagrangian optimization used by (Lachapelle et al., 2020;
Ng et al., 2019; Wehenkel & Louppe, 2021;
Yu et al., 2019; Zheng et al., 2018).
To this end, we denote the learned adjacency matrix A and the the other model parameters by θ .
Hence, we can rewrite the augmented Lagrangian optimization problem as follows:
[Equation]
where h ( A ) =
Tr ( e A ) − n or h ( A ) =
Tr ( I + α A ◦ A ) n − n with α ∈ R + .
The function h is a smooth function over weighted adjacency matrices A ∈ R n × n which is equal to zero when the corresponding graph G satisﬁes the DAG constraints.
The augmented Lagrangian of the problem 6 is deﬁned as:
[Equation]
where λ is the Lagrangian multiplier and c is the penalty parameter.
During training, a practical optimization scheme is to update the parameters λ and c with the iteration:
[Equation]
where η > 1 and γ > 1 are hyper-parameters.
The augmented Lagrangian optimization procedure is expensive since each Lagrangian iteration requires to solve a minimization problem which can be done with stochastic optimization methods.
In practice, previous works set that η = 10 and γ =and use a stopping criterion h ( A k ) < (cid:15) ∈ { 1 e − 6 , 1 e − 8 , 1 e − 10 } to decide for the end of optimization where (cid:15) is a constraint tolerance (Lachapelle et al., 2020; Ng et al., 2019;
Yu et al., 2019; Zheng et al., 2018).
This generally leads to up to 10 Lagrangian iterations (Zheng et al., 2018) and does not guarantee acyclicity of the ﬁnal adjacency matrix (Lachapelle et al., 2020;
Ng et al., 2019; Zheng et al., 2018).
Hence, the augmented Lagrangian requires a ﬁnal non-differentiable thresholding step to remove remaining cycles.
In contrast, DP-DAG can always output valid sampled DAG solutions.
B G UMBEL -S OFTMAX , G UMBEL -S INKHORN , G UMBEL -T OP -K DISTRIBUTIONS B. G UMBEL -S OFTMAX DISTRIBUTION (J ANG ET AL ., 2017)
We assume that a variable U on the ( k − 1) -dimensional simplex follows a Gumbel-Softmax distribution with class probabilities φ 1 , ..., φ k and temperature parameter τ i.e.:
[Equation]
The Gumbel-Softmax distribution becomes identical to the categorical distribution when the temperature τ approaches 0 .
The density of the Gumbel-Softmax can be expressed as follows:
[Equation]
where Γ( k ) denotes the Gamma function.
We can obtain a sample U from the Gumbel-Softmax distribution by computing:
[Equation]
whre G i are i.i.d.
samples from a Gumbel (0) distribution.
For differentiable sampling, we can use the straight-through estimator (Bengio et al., 2013).
Hence, the discrete variable one_hot (arg max( U )) is used for the forward pass, and the continuous variable U
B. G UMBEL -S INKHORN DISTRIBUTION (M ENA ET AL ., 2018)
We assume that a variable Π on the space of permutation matrix follows the Gumbel-Sinkhorn distribution with parameters ψ ∈ R n × n and temperature parameter τ i.e.:
[Equation]
The Gumbel-Sinkhorn distribution becomes identical to the Gumbel-Matching distribution when the temperature τ approaches 0 .
Neither the Gumbel-Matching or the Gumbel-Sinkhorn distributions have tractable densities.
We can obtain a sample Π from the Gumbel-Sinkhorn distribution by computing:
[Equation]
where S is the Sinkhorn operator which iteratively normalizes rows and columns of a matrix and
G ∈ R n × n is a matrix of standard i.i.d.
Gumbel noise.
For differentiable sampling, we can use the straight-through estimator (Bengio et al., 2013).
Hence, the discrete variable Π =
Hungarian
( ˆ Π ) after applying the Hungarian algorithm (Munkres, 1957) to get a discrete permutation is used for the forward pass, and the continuous variable Π is used for the backward pass.
Finally, we can deterministically compute a permutation by removing all Gumbel noise in the forward pass.
B. G UMBEL -T OP -K DISTRIBUTION (K OOL ET AL ., 2019)
We assume that the variable π on the space of ordered sample of size k without replacement follows the Gumbel-Top-k distribution with parameters ψ ∈ R n i.e.:
[Equation]
where τ is an optional temperature parameter usually set to 1 .
The density of the Gumbel-Top-k can be expressed as follows:
[Equation]
where N j = { 1 , ..., n } (cid:31) { π 1 , ..., π j − 1 } .
This corresponds to subsequently sampling k times without replacement from the categorical distribution over the nodes parametrized with ψ with renormalization after each step.
We can compute a sample π from the Gumbel-Top-k distribution by computing:
[Equation]
where arg top-k is the operator selecting the indices with the top-k largest values and G ∈ R n is a vector of i.i.d.
Gumbel noise.
Importantly, remark that the Gumbel-Top-n distribution (i.e. k = n ) states the sorted perturbed log-probabilities π = Sort ( ψ + G ) where G ∈ R n ∼ Gumbel (0) deﬁnes a distribution over component-wise permutation (or linear ordering without replacement).
Finally, we can deterministically compute a permutation by removing all Gumbel noise in the forward pass.
C S OFT S ORT O PERATOR (P RILLO & E ISENSCHLOS , 2020)
The SoftSort operator with temperature parameter τ is deﬁned as:
[Equation]
where d is any differentiable almost everywhere, semi-metric function (e.g. d = | · | ).
It deﬁnes a family as simple continuous relaxation for the permutation matrix corresponding to the argsort operator Π arg sort ( s ) .
Indeed, we have:
[Equation]
In particular, this limits holds almost surely if the entries of s are drawn from a distribution that is absolutely continuous w.r.t.
the Lebesgue measure on R .
Further, note that SoftSort dτ ( s ) is a unimodal row stochastic matrix and the permutation arg max( SoftSort dτ ( s )) is equal to arg sort ( s ) .
D D ATASETS
All the datasets are split in train/validation/test splits with 80% / 10% / 10% .
We use 10 different seeds for the dataset splits.
We make all datasets available at the following anonymous link 3 .
Synthetic We follow the data generation used in related works (Lachapelle et al., 2020; Ng et al., 2019; Peters et al., 2014).
The graph structures A ∗ are generated with the Erdös-Rényi (ER) or ScaleFree (SF), and the causal mechanisms f ∗ i are generated from a (non-linear) Gaussian Process with
RBF kernel of bandwidth one and with independent zero-mean Gaussian noises (cid:15) i .
We considered different sizes of graphs including number of nodes n in { 10 , 50 , 100 } and number of edges m in { 10 , 40 , 50 , 100 , 200 , 400 } .
In our experiments, we use 10 sampled datasets per setting.
Each dataset contains 1000 samples.
Sachs (Sachs et al., 2005):
The Sachs dataset has been widely used in the DAG learning literature (Koller & Friedman, 2009; Lachapelle et al., 2020; Ng et al., 2019;
Yu et al., 2021; Zheng et al., 2018).
It contains 853 observational data points.
Each data point measures the expression level of different proteins and phospholipids in human cells.
The ground-truth DAG has 11 nodes andedges.
SynTReN (Van den Bulcke et al., 2006):
The SynTReN dataset has already been used in (Lachapelle et al., 2020).
Ten datasets are generated using the Syntren generator 4 with software default parameters except for th complex 2-regulator interactions probability which was set to 1 .
Each dataset contains 500 samples.
The ground-truth DAG has always 20 nodes.
E M ODELS
In this section, we provide more details about the training of all the models.
We train all models on a single GPU (NVIDIA GTX 1080 Ti, 11 GB memory).
All numbers are averaged over 10 seeds for the model initialization.
We make available the implementation of VI-DP-DAG at the following anonymous link 5 .
VI-DP-DAG (ours): VI-DP-DAG parametrizes the permutation probability P ψ ( Π ) with GumbelSinkhorn or Gumbel-Top-K tricks, the edge probability P φ ( U ) with Gumbel-Softmax distribution and the causal mechanisms
f i,ψ with neural networks with 3 linear layers.
We performed a grid search over the learning rate lr ∈ [1 e − 4 , 1 e − 2 ] , the number of hidden units per linear layer h ∈ { 8 , 16 , 64 } , the permutation probability parametrization (i.e. Gumbel-Sinkhorn or Gumbel-Top-K), the ﬁxed prior probability P prior ( U ij ) ∈
[1 e − 2 , 1 e − 1 ] and the regularization factor λ ∈ [0 , 1 e − 1 ] .
The temperature parameter for Gumbel-Softmax distribution was ﬁxed to τ = 1 in all experiments.
VI-DP-DAG is trained using the Adam optimizer (Kingma & Ba, 2015).
We perform early stopping by checking loss improvement on the validation set every two epochs and a patience p = 10 .
DAG-GNN (Yu et al., 2019): DAG-GNN model proposes a Graph Neural Network (GNN) autoencoder using variational inference.
The latent variable is not the DAG adjacency matrix A but a noise matrix Z .
DAG-GNN uses the augmented Lagrangian approach for optimization.
We follow the recommendation for augmented Lagrangian optimization in (Lachapelle et al., 2020; Ng et al., 2019;
Yu et al., 2019; Zheng et al., 2018) and set η = 10 and γ = 14 .
We used a neural network with 3 linear layers for the encoder and decoder architectures.
Further, we performed a grid-search on the learning rate lr ∈ [1 e − 4 , 1 e − 2 ] and the number of hidden units per linear layer h ∈ { 8 , 16 , 64 } .
DAG-GNN is trained using the Adam optimizer (Kingma & Ba, 2015).
This model does not directly apply for consequence prediction from learned causal mechanisms as it does not predict the consequence values from the values of the learned parents.
We use the implementation provided in https://github.com/fishmoon1234/DAG-GNN .
DAG-NF (Wehenkel & Louppe, 2021): DAG-NF proposes to combine Normalizing Flows with an augmented Lagrangian optimization to learn DAG structures from observational data.
We follow the recommendation for augmented Lagrangian optimization in (Lachapelle et al.,
Ng et al.
, https com
/b442986b
-q9j87qze-kydrb7wn-
; Yu et al., 2019; Zheng et al., 2018) and set η = 10 and γ = 14 .
Further, we performed a grid-search on the monotonic or afﬁne normalizer architectures, the conditioner and normalizer learning rate lr ∈ [1 e − 4 , 1 e − 2 ] , the number of hidden units per linear layer in the conditioner and normalizer architectures h ∈ { 50 , 100 , 150 , 200 } and different l1 regularization λ l 1 ∈ [0 , 60] .
We used a neural network with 3 linear layers for the conditioner and normalizer architectures and the embedding size was set to set to 30 as suggested by the authors Wehenkel & Louppe (2021).
DAG-NF is trained using the Adam optimizer (Kingma & Ba, 2015).
This model does not directly apply for consequence prediction from learned causal mechanisms as it does not predict the consequence values from the values of the learned parents.
We use the ofﬁcial implementation provided in https://github.com/AWehenkel/Graphical-Normalizing-Flows .
Gran-DAG (Lachapelle et al., 2020):
Gran-DAG proposes to compute the adjacency matrix from the learned weight of the non-linear causal relationships.
Gran-DAG uses the augmented Lagrangian approach for optimization.
We follow the recommendation for augmented Lagrangian optimization in (Lachapelle et al., 2020; Ng et al., 2019; Yu et al., 2019; Zheng et al., 2018) and set η = 10 and γ = 14 .
We used a neural network with 3 linear layers.
Further, we performed a grid-search on the learning rate lr ∈ [1 e − 4 , 1 e − 2 ] and the number of hidden units per linear layer h ∈ { 8 , 16 , 64 } .
Gran-DAG is trained using the RmsProp optimizer (Tieleman & Hinton, 2012).
This model can be used for consequence prediction from learned causal mechanisms.
Indeed, it uses an architecture similar to DPDAG which aims at predicting the values of consequence values from values of the learned parents.
We use the ofﬁcial implementation provided in https://github.com/kurowasan/GraN-DAG .
Masked-NN (Ng et al., 2019):
Masked-NN proposes to learn the binary adjacency matrix A by combining the Gumbel-Softmax distribution with the augmented Lagrangian approach for optimization.
We follow the recommendation for augmented Lagrangian optimization in (Lachapelle et al., 2020; Ng et al., 2019; Yu et al., 2019; Zheng et al., 2018) and set η = 10 and γ = 14 .
We used a neural network with 3 linear layers.
Further, we performed a grid-search on the learning rate lr ∈ [1 e − 4 , 1 e − 2 ] and the number of hidden units per linear layer h ∈ { 8 , 16 , 64 } .
Masked-NN is trained using the Adam optimizer (Kingma & Ba, 2015).
This model can be used for consequence prediction from learned causal mechanisms.
Indeed, it uses an architecture similar to DP-DAG which aims at predicting the values of consequence values from values of the learned parents.
GT-DAG:
GT-DAG is a new baseline with the exact same architecture as VI-DP-DAG except that it uses the ground-truth DAG
A ∗ instead of DAGs sampled from DP-DAG during learning.
Hence, this baseline can be considered as the best-case scenario when learning the causal mechanisms f
i,ψ from observational data.
Further, we performed a grid-search on the learning rate lr ∈ [1 e − 4 , 1 e − 2 ] and the number of hidden units per linear layer h ∈ { 8 , 16 , 64 } .
GT-DAG is trained using the Adam optimizer (Kingma & Ba, 2015).
We perform early stopping by checking loss improvement on the validation set every two epochs and a patience p = 10 .
CAM (Bühlmann et al., 2014): CAM proposes to learns DAGs structure by (1) using Preliminary
Neighbor Selection (PNS), (2) greedily searching for edges with the largest Likelihodd gain, and ﬁnally (3) pruning spurious edges with signiﬁcance testing of covariates (Bühlmann et al., 2014).
The preprocessing PNS step and the postprocessing steps are not differentiable but are decisive for high performance.
This motivates other approaches like Lachapelle et al.
Ng et al.
(2019) to use the same processing steps.
F A DDITIONAL RESULTS F. DAG SAMPLING
In this section, we compare the DAG sampling performances of DP-DAG using Gumbel-Sinkhorn or
Gumbel-Top-k parametrization for permutation sampling.
To this end, we sample 30 DAGs withto 200 nodes and compute the mean and the variance of the sampling time.
We show results for the sampling time in Fig.
The Gumbel-Top-k parametrization is signiﬁcantly faster than the GumbelSinkhorn parametrization.
This is expected from the complexity analysis of the two permutation parametrization.
Gumbel-Sinkhorn has a complexity of O ( n 3 ) while Gumbel-Top-k has a complexity of O ( n 2 ) .
Further, we evaluate the capacity of the two differentiable sampling methods for learning DAGs when the ground-truth DAG is observed.
To this end, the direct MSE error between a new sampled DAG
A and the ground-truth DAG
A ∗ is used as the loss at each iteration.
This task aims at evaluating whether the DAG sampling optimization scheme is a reliable optimization scheme on a simple toy task.
This is a priori not guaranteed since the DAG sampling space is very large.
However, we observe in Tab. 10 that both Gumbel-Sinkhorn and Gumbel-Top-k parametrizations recover almost perfectly the ground-truth DAG
A ∗ for all graph sizes and graph types.
These results are particularly signiﬁcant since each dataset setting is averaged over 10 sampled graphs and 4 learning rates
lr ∈ { 10 e − 1 , 10 e − 2 , 10 e − 3 , 10 e − 4 } .
These experiments validate that DP-DAG is reliable method for fast and differentiable DAG sampling.
[Figure]
[Table] [Table]
truth DAG
A ∗ is used as loss.
Each results is averaged over 10 sampled graphs and 4 learning rates.
Both Gumbel-Sinkhorn and Gumbel-Top-k sampling methods lead to high performance for all graph sizes and graph types.
F. DAG STRUCTURE LEARNING WITHOUT NON - DIFFERENTIABLE PROCESSING
We show the additional results for differentiable DAG structure learning on synthetic and real datasets in Tab. 7.
While GraN-DAG competes with VI-DP-DAG on small synthetic datasets with aroundnodes, DP-DAG outperforms all baselines on all larger synthetic datasets with 50 or 100 nodes for both ER and SF graph types.
On real datasets, no baseline consistently outperforms VI-DP-DAG which obtains 4/8 best scores and second-best scores on the remaining evaluations.
Furthermore, we show results for DAG structure learning for VI-DP-DAG, GraN-DAG and MaskedDAG using the mean and the standard deviation of the SHD scores between the ground-truth dag adjacency matrix A ∗ and the thresholded predicted adjacency matrix P φ ( U π ( i ) π ( j ) )
> t or A ij > t for different threshold choices in Fig. 4.
Thresholds t are ordered from sparser graphs (i
[Table]
e. larger thresholds) to denser graphs (i
[Table]
e. smaller thresholds).
We observe that models might be sensitive to the threshold choice thus motivating the use of threshold agnostic metrics such that AUC-PR and
AUC-ROC.
However, we observe that VI-DP-DAG generally achieves the best performances for different graphs types and graph sizes given the best possible threshold selection.
[Table]
F. DAG STRUCTURE LEARNING WITH NON - DIFFERENTIABLE PROCESSING
We show the additional results for DAG structure learning with non-differentiable pre- and postprocessing in Tab. 11 and Tab. 12.
In these experiments, we apply the default pre-processing PNS and post-processing CAM pruning steps to CAM, GraN-DAG, Masked-DAG as suggested by the original papers.
We similarly apply CAM pruning to the ﬁnal results of VI-DP-DAG.
While the greedy
CAM is the best performing model for most synthetic datasets, VI-DP-DAG maintains competitive performance on all datasets and is the best performing model on real datasets (i.e. 7/8 best scores).
The pre- and post-processing steps are decisive for the augmented Lagrangian optimization baselines
GraN-DAG and Masked-DAG
to achieve performance competitive with VI-DP-DAG.
Lachapelle et al.
(2020) makes similar observations on the dependence on PNS and CAM pruning steps.
We show the comparison of GraN-DAG and VI-DP-DAG with and without additional processing steps
[Figure]
[Figure]
[Table]
[Figure] [Figure] [Table] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
in Tab. 13 and Tab. 14.
GraN-DAG is indeed more dependent on additional processing steps than
VI-DP-DAG to achieve high performance.
F. N ON - DIFFERENTIABLE PROCESSING TIME
We show additional results of the processing time (incl.
the mean and standard error of the mean over 5 runs) for the non-differentiable PNS pre-processing, the CAM algorithm and the DAG pruning post-processing (Bühlmann et al., 2014) for ER-10-40, ER-50-200, ER-100-400 with with different number of nodes in Tab. 15.
Each of this processing step becomes signiﬁcantly slower tha VI-DPDAG for larger number of nodes.
In particular on ER-100-400, the PNS pre-processing is around ×slower than VI-DP-DAG, the CAM algorithm is at least more than × 900 slower than VI-DP-DAG since it did not ﬁnish in less than 2 days, and the DAG pruning post-processing is around × 23 slower than VI-DP-DAG.
F. DAG STRUCTURE LEARNING UNDER PERTURBATIONS .
We evaluate if the models conﬁdently assign lower conﬁdence scores to perturbed versions of the ground-truth adjacency matrix A ∗ .
To this end, we create a perturbed matrix ˜
A ∗ by randomly moving different number of edges from A ∗ , and compute its relative conﬁdence score.
The relative conﬁdence score assigned by a model is the sum of the edges scores of the perturbed matrix ˜
A ∗ normalized by the score of the clean matrix A ∗ i.e. (cid:80) ij S ij 1 ˜ Aij
[Equation]
ij S ij 1 ˜ A ∗ ij where S ij =
P φ ( U π ( i ) π
( j ) ) for VI-DP-DAG and
S ij =
A ij for other models using Lagrangian optimization.
We report the mean and standard deviation of the relative conﬁdence score over 10 sampled perturbed graphs in Fig.
As desired, VI-DP-DAG conﬁdently assigns a higher score to the ground-truth matrix and lower scores to graphs with larger perturbations.
This indicates that the DAG sampling optimization of DP-DAG converges to a reliable probabilistic model over DAGs.
In contrast, GraNDAG and Masked-DAG assign signiﬁcantly more noisy scores to perturbed graphs (see e.g. Fig. 5b).
Further, other baselines do not conﬁdently assign strictly decreasing scores to graphs with larger perturbations.
F. C AUSAL MECHANISMS LEARNING
We show causal mechanisms learning results when varying the threshold t used to compute the adjacency matrix A from the edge probabilities
A i j = 1 P φ ( U ij | Π ) >t of VI-DP-DAG in Fig.
The
MSE scores of VI-DP-DAG are almost insensitive to the threshold value except for probability [Figure] [Figure] [Table]
[Table]
[Table]
[Table] [Table]