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: Using the recently developed differentiable sampling techniques for discrete variables (e.g. Gumbel-Softmax) and variable orderings (e.g. Gumbel-Sinkhorn and Gumbel-Top-k), this paper proposes a differentiable DAG sampling strategy, and applies it to the problem of DAG learning.