Differentiable DAG Sampling

1 INTRODUCTION . 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 first traditional family of models for DAG learning are discrete score-based approaches . These approaches aim at solving the following discrete optimization problem : max G score ( X , G ) s.t . G ∈ discrete DAGs ( 1 ) 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 : max G score ( X , G ) s.t . h ( G ) = 0 ( 2 ) where h ( G ) is a smooth function over weighted graphs G which is equal to zero when the graph G satisfies 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 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 . 2 RELATED WORK . 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 define 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 fulfilled 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 define 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 . 3 PROBABILISTIC MODEL OVER DAGS . A Directed Acyclic Graph ( DAG ) is a graph G = ( V , E ) with n nodes x1 , ... , xn 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 can not 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 = ΠTUΠ . 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 . Aij = Uπ ( i ) π ( j ) such that Uπ ( i ) π ( j ) = 0 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 define the following probabilistic model over DAGs ( DP-DAG ) based on the adjacency matrix decomposition in Th . 1 : P ( A ) = ∑ Π∈P ( G ) , U∈Un P ( U ) P ( Π ) s.t . A = ΠTUΠ ( 3 ) where P ( G ) is the set of valid permutation matrices for the DAG G , Un 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 Uij ∈ { 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 defined 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 Uij = 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 ( n2 ) . We recall the definition 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 defined 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 ( n3 ) . We recall the definition 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 defines 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 ∈ Rn are i.i.d . Gumbel noise , defines 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 ) ∈ Rn×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 ( n2 ) . 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 definition 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φ ( Uij ) i.e . Uij ∼ Pφ ( Uij ) . ( 3 ) Compute the final adjacency matrix A from the permutation matrix Π and the upper triangular matrix U i.e . A = ΠTUΠ . This procedure is capable of sampling any possible DAGs of n nodes due to Th . 1 . In practice , we propose to parametrize the distribution Pψ ( Π ) using the Gumbel-Sinkhorn or the Gumbel-Top-k trick which define valid distributions over permutations , and parametrize the distributions Pφ ( Uij ) using the Gumbel-Softmax trick which defines 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 ( n3 ) using Gumbel-Sinkhorn and O ( n2 ) 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 .

Summary: This paper presents a differentiable probabilistic model (DP-DAG) over DAGs that retains accuracy over comparable differentiable DAG models. In combination with variational inference for training, VI-DP-DAG demonstrates large gains in computational efficiency at the cost of an intractable exact scoring function (ie computing $P_{\phi,\psi}(A)$). The computational gains come from the factorization of the DAG distribution into the product of orderings and edges, which results in intractable scoring as all valid permutations must be marginalized over in order to score a graph. However, this is not an issue: the pathwise derivative is used during training, requiring only a differentiable sampling procedure, and evaluation does not require scoring.