DAG L EARNING ON THE P ERMUTAHEDRON
A BSTRACT
We introduce Daguerro , a strategy for learning directed acyclic graphs (DAGs).
In contrast to previous methods, our problem formulation (i) guarantees to learn a
DAG, (ii) does not propagate errors over multiple stages, and (iii) can be trained end-to-end without pre-processing steps.
Our algorithm leverages advances in differentiable sparse structured inference for learning a total ordering of the variables in the simplex of permutation vectors (the permutahedron), allowing for a substantial reduction in memory and time complexities w.r.t.
existing permutation-based continuous optimization methods.
I NTRODUCTION
Learning a directed acyclic graph (DAG) from observational data is very challenging yet useful for many applications, e.g., in biology (Sachs et al., 2005), genetics (Zhang et al., 2013), or finance (Sanford & Moosa, 2012).
The challenge comes from the combinatorial nature of the solution space, whose size grows super-exponentially with the number of variables d .
While exact scoredbased algorithms exist for small d (Singh & Moore, 2005; Xiang & Kim, 2013; Cussens, 2011), approximate methods (Scanagatta et al., 2015; Aragam & Zhou, 2015; Ramsey et al., 2017) rely on global or local search heuristics in order to scale to problems with thousands of nodes.
For instance, a large body of recent works (Zheng et al., 2018; Yu et al., 2019; Zheng et al., 2020;
Ng et al., 2020; Brouillard et al., 2020;
He et al., 2021) formulate DAG learning as a continuous optimization problem, where the acyclicity constraint is expressed as a smooth function and relaxed to allow efficient optimization with multi-purpose solvers.
As such, the absence of cycles is no longer guaranteed and solutions often require post-processing.
Assuming identifiability of the underlying graph structure, another prominent line of works (Friedman & Koller, 2003; Gao et al., 2020; Reisach et al., 2021; Cundy et al., 2021; Charpentier et al., 2022) learns DAGs by (i) finding an ordering of the variables, and (ii) selecting the best scoring graph among (or marginalizing over) the structures that are consistent with the found ordering.
Their benefit is to work on the space of orderings which is smaller and more regular than the space of DAGs (Friedman & Koller, 2003) while guaranteeing acyclicity.
Existing works, however, are based either on two-steps procedures (Gao et al., 2020; Reisach et al., 2021) that do not guarantee optimality of the solution, or on end-to-end ones that involve higher computational costs (Friedman & Koller, 2003; Cundy et al., 2021; Charpentier et al., 2022).
In particular, the recent work of Cundy et al.
(2021) models a total ordering distribution on the Birkhoff polytope (the convex hull of permutation matrices) to obtain a differentiable operator, which has O ( d 3 ) time and O ( d 2 ) memory complexities.
Charpentier et al.
(2022) proposes using another operator with improved time complexity ( O ( d 2 ) ) by constraining the permutation matrix to be row-stochastic (and not bistochastic).
In this work, we propose Daguerro , an end-to-end score-based strategy that belongs to the permutation-based family of works.
Contrary to Cundy et al.
(2021); Charpentier et al. , we build our method on the SparseMAP operator (Niculae et al., 2018) for learning the total ordering of [Figure]
the variables. This choice of differentiable operator allows us to work in the simplex of permutation vectors (a.k.a.
Permutahedron) with significant gains in time ( O ( d log d ) ) and space ( O ( d ) )
complexities (Blondel et al., 2020).
Preliminary results on synthetic data show that Daguerro generally recovers the topological sort of the nodes but tends to return overly dense graphs.
M ETHOD
Let X ∈ R n × d be a data matrix consisting of n i.i.d.
observations for d nodes, where we denote the j -th column of X by x j ∈ R n , and G ∈ D a DAG, where D denotes the discrete space of DAGs
G =
( V , E ) on d nodes.
We represent G by its adjacency matrix A ∈ R d × d , where an element
A ij ̸ =
iff a directed edge exists from node i to node j .
Learning DAGs via ranking and masking We formulate our score-based problem as follows
[Equation]
where ℓ : R n × R n → R is a point-wise loss function, { f ϕ } dj
=1 is a set of functions parameterized by ϕ ∈ R d ϕ that computes x j given its parents pa ( j )
= { i ∈ [ d ] | A ( θ ) ij ̸ = 0 } , and θ ∈ R
d is an utility vector whose role will be clarified later. To enforce the constraint A ( θ ) ∈ D , we learn a total ordering of the variables and mask out all edges that are not consistent with it: if x j ≺ x i no edge can be drawn from i to j .
As this procedure returns only transitive closures, we further encourage the removal of unnecessary edges via l 1 regularization (as later defined in eq.
We decompose A ( θ ) as follows:
[Equation]
where S ∈ R d × d is a full and unconstrained weighted adjacency matrix, σ ( θ ) ∈
Σ d is a parametrized permutation with Σ d being the set of all d -permutations, and M ∈ { 0 , 1 } d × d is an upper strictly triangular masking matrix such that M ij = 1 if i < j and 0 otherwise.
Further, ◦ denotes the Hadamard product, and M σ ( θ ) indicates that the rows and columns of M have been permuted according to σ ( θ ) .
Soft ranking on the permutahedron.
In contrast to Cundy et al.
(2021) that operates in the Birkhoff polytope, to learn the optimal permutation we turn to an efficient formulation of ranking as constrained optimization on the permutahedron, following Blondel et al.
The ranking of a vector θ ∈ R d , i.e., the permutation σ ∈ Σ d that arranges the elements of θ in decreasing order, can be written as (Blondel et al., 2020, Lemma 1)
[Equation]
where ρ =
( d, d − 1 , . . .
The permutahedron of a vector w ∈ R d is the convex hull of all possible permutations of w :
[Equation]
Blondel et al.
(2020) show that projection onto P ( ρ ) yields a soft ranking operator:
[Equation]
This mapping is differentiable w.r.t.
θ and, in the absence of ties, as the regularization parameter τ goes to zero, σ τ ( θ ) → σ ( θ ) (Blondel et al., 2020, Proposition 2).
An algorithm based on isotonic regression permits O ( d log d ) computation of σ τ ( θ ) .
However, soft rankings are not suitable for our goal, since we cannot replace M σ ( θ ) by M σ τ ( θ ) .
In fact, the latter is only defined for the vertices of the permutahedron, i.e., discrete permutations ρ ∈ Σ d .
The same problem arises when working with soft permutation matrices, as done in Cundy et al.
(2021) and Charpentier et al. When leveraging Mena et al. (2018)’s work, the authors resort to the straight-through gradient estimator (Bengio et al., 2013), which boils down to (i) finding a hard permutation matrix in the forward pass by solving an optimal matching program via the Hungarian algorithm (Kuhn, 2010), and (ii) computing the gradients of its continuous approximation obtained through the Sinkhorn operator (Sinkhorn, 1964) for the backward pass.
The straight-through gradient estimator is also deployed when leveraging the SoftSort operator (Prillo & Eisenschlos, 2020), as done in (Charpentier et al., 2022) where the one-hot encoded argsort is used for the forward pass.
This strategy however induces additional bias in the training procedure.
SparseMAP distributions over permutations.
To handle the issue that our loss can only be evaluated meaningfully at discrete permutations, we note that any µ ∈ P ( ρ ) can be written as a convex combination of vertices, [Equation]
[Equation]
where α ( · ) ≥ 0 , and (cid:88)
[Equation]
SparseMAP (Niculae et al., 2018) is a strategy for differentiable structured inference that solves problems of the form of eq.
(5), yielding not only the optimal solution σ τ ( θ ) = µ ∗ , but also a sparse decomposition α ∗ ( θ ) , such that µ ∗ =
E σ ∼ α ∗ [ ρ σ ] .
This is achieved via an active set algorithm, which maintains a sparse decomposition, and iteratively adds or removes permutations based on linear oracles of the form (3), implemented by a sorting algorithm.
This iterative algorithm is more costly than the direct isotonic regression one, but provides the decomposition α ∗ , which allows us to compute exact marginalizations of the form E σ ∼ α ∗ ( θ )
[ f ( σ )] by explicitly summing over the sparse set of permutations found by SparseMAP, as proposed for structured latent variable learning by
Correia et al.
learning on the permutahedron.
Our final bi-level optimization problem is:
[Equation]
The inner problem eq.
(8) is optimized via SparseMAP, which returns a sparse categorical distribution α ∗ and its gradients w.r.t.
θ .
We can solve the outer optimization via proximal adaptive gradient methods (Parikh & Boyd, 2014).
A schematic of the overall approach is presented in Figure 1. [Figure]
Computational details.
Each SparseMAP iteration involves a length- d argsort and a Cholesky update of a s -by- s matrix, where s is the size of the active set, bounded by the iteration number.
With a constant number of iterations (as in our implementation), this leads to an overall time complexity of O ( sd log d + s 2 ) and space complexity O ( sd ) .
If run until exact convergence, Carath ´ eodory’s convex hull theorem guarantees s ≤ d + 1 leading to cubic worst-case complexity; however, in practice, we almost always reach convergence within 100 iterations.
In addition, we warm-start the sorting algorithm with the last selected permutation.
This is better both in theory and in practice than the O ( d 3 ) complexity of maximization over the Birkhoff polytope.
A tighter analysis or specialized algorithm based on isotonic regression is a promising direction for future work.
E XPERIMENTS
We report a set of preliminary experiments to validate the proposed method for learning DAGs from observational data.
We consider synthetic settings with simulated linear additive noise models.
Following common practices from the literature (Zheng et al., 2018; Gao et al., 2020), we generate a random graph with 2 ∗ d expected edges according to the Erd¨os-R´enyi, Scale-Free or BiPartite models and assign to each edge a weight uniformly drawn in [ − 2 , − 0 .
5] ∪ [0 . 5 , 2] .
We then sample data according to the linear model XA + ε , where ε ∈ R d is an exogenous random variable whose elements are independently distributed as Normal(0 , 0 . 5) ( Gaussian ), Normal(0 , σ j ) with σ j ∈ [0 , 0 . 5] ( Gaussian-Heterogeneous ), Exponential(0 . 5) ( Exponential ) or Gumbel(0 , 0 . 5) ( Gumbel ).
As baselines, we consider state-of-the-art approximate score-based methods: NoTears (Zheng et al., 2018), the first continuous optimization method, which optimizes the Frobenius reconstruction loss and where the DAG constrain is enforced via the Augmented Lagrangian approach; Golem (Ng et al., 2020), another continuous optimization method which optimizes the data likelihood (under
Gaussian equal variance error assumptions) regularized by NoTears ’s DAG penalty; NPVAR (Gao et al., 2020), an iterative algorithm that learns topological layers and then prunes edges based on node residual variance (with Generalized Additive Models regressor backend to estimate conditional variance)
; sortnregress (Reisach et al., 2021), another two-steps strategy that orders nodes by increasing variance and selects the parents of a node among all its predecessors using the Least Angle Regressor (Efron et al., 2004).
Before evaluation, we post-process the graphs found by NoTears and Golem by iteratively removing edges ordered by increasing weight until obtaining a DAG, as they often contain cycles.
For our method, we instantiate f ϕ j
( X , A ( θ ) j ) =
XA ( θ ) j and optimize the data likelihood (under Gaussian equal variance error assumptions).
We set the hyper-parameters of all methods to their default values, apart from sortnregress that uses the Bayesian Information Criterion for model selection.
In particular, for Daguerro we set λ = 0 .
1 , τ = 1 and remove all edges with weights ≤ 0 .
3 as commonly done in the literature.
We additionally apply a l 2 regularization to all learned parameters (and set its hyper-parameter to 0 .
01 ) as we find it stabilizes training.
In Figure 2, we compare the methods by two metrics: the Structural Hamming Distance ( SHD ) between true and estimated graphs, as standard in the literature; and the Topological Ordering Pearson
Correlation ( TOPC ) that we introduce to measure how well a method is able to estimate the true topological layers of the nodes.
We believe it is essential to study both metrics because the SHD alone does not distinguish between wrong edges that are consistent with the variable ordering from wrong edges that are not.
For instance, the transitive closure of a sparse graph has high SHD even though it preserves the true ordering of the nodes.
TOPC computes the Pearson’s correlation coefficient between the true topological layers and the learned ones.
Formally, given a DAG G =
( V , E ) and its adjacency matrix A we define the topological layers ( or layer decomposition)
L ( G ) =
( L 1 , . . .
, L R ) following Gao et al. (2020), where L A = { x j | pa ( j ) ∈ ( ∪ r − 1 r ′ =1 L r ′ ) ∧ x j ∈ ( V \∪ r − 1 r ′
=
1 L r ′ ) } is the set of nodes that are sources of the
DAG G [ V \
∪ rr ′ =
1 L r ′ ] and 1 ≤ R ≤ d is the depth of the graph.
Note that this layer decomposition is unique, as any node x j belongs to a single layer.
We can then assign to each node x j the index of the layer it belongs to, obtaining the assignment vectors r A .
TOPC between estimated ˆ A and true
A adjacency matrices is computed as
[Equation] with Cov( · ) denoting the covariance and Var( · ) denoting the variance. In terms of TOPC , we observe that Daguerro provides better orderings than those of NoTears and NPVAR ; it sightly improves over sortnregress and it is on-par with Golem on the equal-variance noise models, while clearly improving on Golem in the Gaussian-Heterogeneous setting.
In terms of SHD , we find that Daguerro never provides the best performance, although it still improves over NoTears and NPVAR .
Given these results, we believe that a better strategy for pruning edges is key to improve our method.
Additional improvement may come from tuning our hyper-parameters.
C ONCLUSION AND FUTURE WORK
In this work, we presented Daguerro , an efficient score-based method for end-to-end learning of directed acyclic graphs and structural equation models.
Daguerro reliably finds correct total ordering, yet tends to yield too dense graphs.
We anticipate that incorporating a l 0 regularization term (Louizos et al., 2018) could help mitigate this latter issue.
Future directions include adapting our strategy to non-linear SEMs (less amenable to continuous optimization approaches) and experimenting with interventional data.
While our current implementation uses SparseMAP, other recent strategies for structured discrete distributions such as SST (Paulus et al., 2020) or I-MLE (Niepert et al., 2021) could offer compelling alternatives to optimize the structure parameters.
A CKNOWLEDGEMENTS
We thank the anonymous reviewers for their constructive feedback. VZ is supported by the
French Project APRIORI ANR-18-CE23-0015. JK acknowledges support by the Engineering and Physical Sciences Research Council with grant number EP/S021566/1. VN is partially supported by the Hybrid Intelligence Centre, a 10-year programme funded by the Dutch Ministry of Education, Culture and Science through the Netherlands Organisation for Scientific Research ( https://hybrid-intelligence-centre.nl ). A D ETAILED RESULTS
[Figure] [Figure] [Figure] [Figure]