$\psi$DAG : Projected Stochastic Approximation Iteration for DAG Structure Learning

## Abstract

Abstract Learning the structure of Directed Acyclic Graphs (DAGs) presents a significant challenge due to the vast combinatorial search space of possible graphs, which scales exponentially with the number of nodes. Recent advancements have redefined this problem as a continuous optimization task by incorporating differentiable acyclicity constraints. These methods commonly rely on algebraic characterizations of DAGs, such as matrix exponentials, to enable the use of gradient-based optimization techniques. Despite these innovations, existing methods often face optimization difficulties due to the highly non-convex nature of DAG constraints and the per-iteration computational complexity. In this work, we present a novel framework for learning DAGs, employing a Stochastic Approximation approach integrated with Stochastic Gradient Descent (SGD)-based optimization techniques. Our framework introduces new projection methods tailored to efficiently enforce DAG constraints, ensuring that the algorithm converges to a feasible local minimum. With its low iteration complexity, the proposed method is well-suited for handling large-scale problems with improved computational efficiency. We demonstrate the effectiveness and scalability of our framework through comprehensive experimental evaluations, which confirm its superior performance across various settings.

## 1 Introduction

Learning graphical structures from data using Directed Acyclic Graphs (DAGs) is a fundamental challenge in machine learning (Koller & Friedman, ; Peters et al., ; Arjovsky et al., ; Sauer & Geiger, ). This task has a wide range of practical applications across fields such as economics, genome research (Zhang et al., ; Stephens & Balding, ), social sciences (Morgan & Winship, ), biology (Sachs et al., ), causal inference (Pearl, ; Spirtes et al., ).
Learning the graphical structure is essential because the resulting models can often be given causal interpretations or transformed into representations with causal significance, such as Markov equivalence classes.
When graphical models cannot be interpreted causally (Pearl, ; Spirtes et al., ), they can still offer a compact and flexible representation for decomposing the joint distribution.

Structure learning methods are typically categorized into two approaches: score-based algorithms searching for a DAG minimizing a particular loss function and constraint-based algorithms relying on conditional independence tests. Constraint-based methods, such as the PC algorithm (Spirtes & Glymour, ) and FCI (Spirtes et al., ; Colombo et al., ), use conditional independence tests to recover the Markov equivalence class under the assumption of faithfulness. Other approaches, like those described in Margaritis & Thrun () and Tsamardinos et al. (), employ local Markov boundary search. On the other hand, score-based methods frame the problem as an optimization of a specific scoring function, with typical choices including BGe (Kuipers et al., ), BIC (Chickering & Heckerman, ), BDe(u) (Heckerman et al., ), and MDL (Bouckaert, ). Given the vast search space of potential graphs, many score-based methods employ local heuristics, such as Greedy Equivalence Search (GES) (Chickering, ), to efficiently navigate this complexity. Additionally, Tsamardinos et al. (), Gámez et al. () propose hybrid methods combining elements of both constraint-based and score-based learning.

Recently, Zheng et al. () introduced a smooth formulation for enforcing acyclicity, transforming the structure learning problem from its inherently discrete nature into a continuous, non-convex optimization task. This formulation allows for the use of gradient-based optimization techniques, enabling various extensions and adaptations to various domains, including nonlinear models (Yu et al., ; Ng et al., ; Kalainathan et al., ), interventional datasets (Brouillard et al., ; Faria et al., ), unobserved confounders (Bhattacharya et al., ; Bellot & Van der Schaar, ), incomplete datasets (Gao et al., ; Wang et al., ), time series analysis (Sun et al., ; Pamfil et al., ), multi-task learning (Chen et al., ), multi-domain settings (Zeng et al., ), federated learning (Ng & Zhang, ; Gao et al., ), and representation learning (Yang et al., ).
With the growing interest in continuous structure learning methods (Vowels et al., ), a variety of theoretical and empirical studies have emerged. For instance, Ng et al. () investigated the optimality conditions and convergence properties of continuously constrained approaches such as Zheng et al. (). In the bivariate case, Deng et al. () demonstrated that a suitable optimization strategy converges to the global minimum of the least squares objective. Additionally, Zhang et al. () and Bello et al. () identified potential gradient vanishing issues with existing DAG constraints (Zheng et al., ) and proposed adjustments to overcome these challenges.

#### Contributions.

In this work, we focus on the graphical models represented as Directed Acyclic Graphs (DAGs). Our main contributions can be summarized as follows:

- 1.
Problem reformulation: We introduce a new reformulation (11) of the discrete optimization problem for finding DAG as a stochastic optimization problem and we discuss its properties in detail in Section 3.1.
We demonstrate that the solution of this reformulated problem recovers the true DAG (Section 3.1).
- 2.
Novel algorithm: Leveraging insights from stochastic optimization, we present a new framework (Algorithm 1) for DAG learning (Section 4) and present a simple yet effective algorithm $\psi{}$DAG (Algorithm 3) within the framework. We proved that Algorithm $\psi{}$DAG converges to a local minimum of problem (11) at the rate
$\mathbb{E}\left[F(\mathbf{W}_{T})-F(\mathbf{W}^{\ast})\right]\leq\mathcal{O}
\left(\tfrac{\sigma_{1}R}{\sqrt{T}}+\tfrac{L_{1}R^{2}}{T}\right),$
where $T$ is a number of SGD-type steps, as detailed in (Section 4.1).
- 3.
Experimental comparison: In Section 5, we demonstrate that the method $\psi{}$DAG scales very well with graph size, handling up to $10000$ nodes. At that scale, the primary limitation is not computation complexity but the memory required to store the DAG itself.
As a baseline, we compare $\psi{}$DAG with established DAG learning methods, including NOTEARS (Zheng et al., ), GOLEM (Ng et al., ) and DAGMA (Bello et al., ).
We show a significant improvement in scalability, as baseline methods struggle with larger graphs. Specifically, NOTEARS (Zheng et al., ), GOLEM (Ng et al., ) and DAGMA (Bello et al., ) require more than $100$ hours for graphs with over $3000$ nodes, exceeding the allotted time.

## 2 Background

### 2.1 Graph Notation

Before discussing the connection to the most relevant literature, we formalize the graph notation associated with DAGs.

Let $\mathcal{G}\stackrel{{\scriptstyle def}}{{=}}(V,E,w)$ represent a weighted directed graph, where $V$ denotes the set of vertices with cardinality $d\stackrel{{\scriptstyle def}}{{=}}\left|V\right|$, $E\in 2^{V\times V}$ is the set of edges, and $w:V\times V\to\mathbb{R}\setminus\{0\}$ assigns weights to the edges. The adjacency matrix $\mathbf{A}(\mathcal{G}):\mathbb{R}^{d\times d}$ is defined such that $[\mathbf{A}(\mathcal{G})]_{ij}=1$ if $(i,j)\in E$ and 0 0 otherwise. Similarly, the weighted adjacency matrix $\mathbf{W}(\mathcal{G})$ is defined by $[\mathbf{W}(\mathcal{G})]_{ij}=w(i,j)$ if $(i,j)\in E$ and 0 0 otherwise.

When the weight function $w$ is binary, we simplify the notation to $\mathcal{G}\stackrel{{\scriptstyle def}}{{=}}(V,E)$. Similarly, when the graph $\mathcal{G}$ is clear from context, we shorthand the notation to $\mathbf{A}\stackrel{{\scriptstyle def}}{{=}}\mathbf{A}(\mathcal{G})$ and $\mathbf{W}\stackrel{{\scriptstyle def}}{{=}}\mathbf{W}(\mathcal{G})$.

We denote the space of DAGs as $\mathbb{D}$. Since we will be utilizing topological sorting of DAGs(^1^11Topologial sorting of a graph $\mathcal{G}\stackrel{{\scriptstyle def}}{{=}}(V,E,w)$ refers to vertex ordering $V_{1},V_{2},\dots,V_{d}$ such that $E$ contains no edges of the form $V_{i}\to V_{j},$ where $i\leq j$. Importantly, every DAG has at least one topological sorting.), we also denote the space of vertex permutations $\Pi$.

### 2.2 Linear DAG and SEM

A Directed Acyclic Graph (DAG) model, defined on a set of $n$ random vectors $X\in\mathbb{R}^{n\times d}$, where $X=(X_{1},\dots,X_{n})$ and $X_{i}\in\mathbb{R}^{d}$, consists of two components:

- 1.
A DAG $\mathcal{G}=(V,E)$, which encodes a set of conditional independence relationships among the variables.
- 2.
The joint distribution $P(X)$ with density $p(x)$, which is Markov with respect to the DAG $\mathcal{G}$ and factors as $p(x)=\prod_{i=1}^{d}p(x_{i}\mid x_{\text{PA}_{\mathcal{G}}(i)}),$ where $\text{PA}_{\mathcal{G}}(i)=\{j\in V:X_{j}\to X_{i}\in E\}$ represents the set of parents of $X_{i}$ in $\mathcal{G}$.

This work focuses on the linear DAG model, which can be equivalently represented by a set of linear Structural Equation Models (SEMs). In matrix notation, the linear DAG model can be expressed as

$$ $X=X\mathbf{W}+N,$ (1) $$

where $\mathbf{W}=[\mathbf{W}_{1}|\cdots|\mathbf{W}_{d}]$ is a weighted adjacency matrix, and $N=(N_{1},\dots,N_{n})$ is a matrix where each $N_{i}\in\mathbb{R}^{d}$ represents a noise vector with independent components. The structure of graph $\mathcal{G}$ is determined by the non-zero coefficients in $\mathbf{W}$; specifically $X_{j}\to X_{i}\in E$ if and only if the corresponding coefficient in $\mathbf{W}_{i}$ for $X_{j}$ is non-zero. The classical objective function is based on the least squares loss applied to the linear DAG model,

$$ $l(\mathbf{W};X)\stackrel{{\scriptstyle def}}{{=}}\frac{1}{2n}\|X-X\mathbf{W}\| _{F}^{2}.$ (2) $$

### 2.3 Most Related Literature

A body of research most connected to our work is based on the non-convex continuous framework of Zheng et al. (), Ng et al. (), and Bello et al. (), which we will discuss in detail.

Zheng et al. () addressed the constrained optimization problem

$$ $\min_{\mathbf{W}\in\mathbb{R}^{d\times d}}\ell(\mathbf{W};X)_{{\sf{NOTEARS}}{} }\stackrel{{\scriptstyle def}}{{=}}\frac{1}{2n}\|X-X\mathbf{W}\|_{F}^{2}+ \lambda\|\mathbf{W}\|_{1}\quad\text{subject to}\quad h(\mathbf{W})=0,$ (3) $$

where $\ell(\mathbf{W};X)$ represents the least squares objective and $h(\mathbf{W}):=\text{tr}(e^{\mathbf{W}\odot\mathbf{W}})-d$ enforces the DAG constraint. Additionally, an $\ell_{1}$ regularization term $\lambda\|\mathbf{W}\|_{1}$, where $\|\cdot\|_{1}$ is the element-wise $\ell_{1}$-norm and $\lambda$ is a hyperparameter incorporated into the objective function. This formulation addresses the linear case with equal noise variances, as discussed in Loh & Bühlmann () and Peters & Bühlmann (). This constrained optimization problem is solved using the augmented Lagrangian method (Bertsekas et al., ), followed by thresholding the obtained edge weights.
However, since this approach computes the acyclicity function via the matrix exponential, each iteration incurs a computational complexity of $\mathcal{O}(d^{3})$, which significantly limits the scalability of the method.

Ng et al. () introduced the GOLEM method, which enhances the scoring function by incorporating an additional log-determinant term, $\log|\det(I-\mathbf{W})|$ to align with the Gaussian log-likelihood,

$$ $\min_{\mathbf{W}\in\mathbb{R}^{d\times d}}\ell(\mathbf{W};X)_{{\sf{GOLEM}}{}} \stackrel{{\scriptstyle def}}{{=}}\frac{d}{2}\log\|X-X\mathbf{W}\|_{F}^{2}- \log|\det(I-\mathbf{W})|+\lambda_{1}\|\mathbf{W}\|_{1}+\lambda_{2}h(\mathbf{W}),$ (4) $$

where $\lambda_{1}$ and $\lambda_{2}$ serve as regularization hyperparameters within the objective function. Although the newly added log-determinant term is zero when the current model $\mathbf{W}$ is a DAG, this score function does not provide an exact characterization of acyclicity. Specifically, the condition $\log|\det(I-\mathbf{W})|=0$ does not imply that $\mathbf{W}$ represents a DAG.

Bello et al. () introduces a novel acyclicity characterization for DAGs using a log-determinant function,

$$ $\min_{\mathbf{W}\in\mathbb{R}^{d\times d}}\ell(\mathbf{W};X)_{{\sf{DAGMA}}{}} \stackrel{{\scriptstyle def}}{{=}}\frac{1}{2n}\|X-X\mathbf{W}\|_{F}^{2}+ \lambda_{1}\|\mathbf{W}\|_{1}\quad\text{subject to}\quad h^{s}_{ldet}(\mathbf{ W})=0,$ (5) $$

where
$h^{s}_{ldet}(\mathbf{W})\overset{\text{def}}{=}-\log\det(sI-\mathbf{W}\circ
\mathbf{W})+d\log s$, and it is both exact and differentiable.

In practice, the augmented Lagrangian method enforces the hard DAG constraint by increasing the penalty coefficient towards infinity, which requires careful parameter fine-tuning and can lead to numerical difficulties and ill-conditioning (Birgin et al., ; Ng et al., ). As a result, existing methods face challenges across several aspects of optimization, including the careful selection of constraints, high computational complexity, and scalability issues.

To overcome these challenges, we propose a novel framework for enforcing the acyclicity constraint, utilizing a low-cost projection method. This approach significantly reduces iteration complexity and eliminates the need for expensive hyperparameter tuning.

## 3 Stochastic Approximation for DAGs

Our framework is built on a reformulation of the objective function as a stochastic optimization problem, aiming to minimize the stochastic function $F(w)$,

$$ $\min_{w\in\mathbb{R}^{d}}\left\{F(w)\stackrel{{\scriptstyle def}}{{=}}\mathbb{ E}_{\xi}\left[f(w,\xi)\right]\right\},$ (6) $$

where $\xi\in\Xi$ is a random variable drawn from the distribution $\Xi$. This formulation is common in stochastic optimization where computing the exact expectation is infeasible, but the values of $f(w,\xi)$ and its stochastic gradients $g(w,\xi)$ can be computed. Linear and logistic regressions are classical examples of such problems.

To address this problem, two main approaches exist: Stochastic Approximation (SA) and Sample Average Approximation (SAA). The SAA approach involves sampling a fixed number $n$ of random variables or data points $\xi_{i}$ and then minimizing their average $\tilde{F}(w)$:

$$ $\min_{w\in\mathbb{R}^{d}}\left\{\tilde{F}(w)\stackrel{{\scriptstyle def}}{{=}} \tfrac{1}{n}\sum_{i=1}^{n}f(w,\xi_{i})\right\}.$ (7) $$

Now, the problem (7) becomes deterministic and can be solved using various optimization methods, such as gradient descent. However, the main drawback of this approach is that the solution of (7) $\tilde{w}^{\ast}$ is not necessarily equal to the solution of the original problem (6). Even with a perfect solution of (7), there will still be a gap $\|\tilde{w}^{\ast}-w^{\ast}\|=\delta_{x}$ and $F(\tilde{w}^{\ast})-F^{\ast}=\delta_{F}$ between approximate and true solution. These gaps are dependent on the sample size $n$.
Stochastic Approximation (SA) minimizes the true function $F(w)$ by utilizing the stochastic gradient $g(w,\xi)$. Below, we provide the formal definition of a stochastic gradient.

###### Assumption 1 .

For all $w\in\mathbb{R}^{d}$, we assume that stochastic gradients $g(w,\xi)\in\mathbb{R}^{d}$ satisfy

$$ $\mathbb{E}[g(w,\xi)\mid w]=\nabla F(w),\quad\mathbb{E}\left[\|g(w,\xi)-\nabla F (w)\|^{2}\mid w\right]\leq\sigma_{1}^{2}.$ (8) $$

We use these stochastic gradients in SGD-type methods:

$$ $w_{t+1}=w_{t}-h_{t}g(w_{t},\xi_{i}),$ (9) $$

where $h_{t}$ is a step-size schedule.
SA originated with the pioneering paper by Robbins & Monro (). For convex and $L$-smooth function $F(w)$, Polyak (); Polyak & Juditsky (); Nemirovski et al. (); Nemirovski & Yudin () developed significant improvements to SA method in the form of longer step-sizes with iterate averaging, and obtained the convergence guarantee

$$ $\mathbb{E}\left[F(w_{T})-F(x^{\ast})\right]\leq\mathcal{O}\left(\tfrac{\sigma_ {1}R}{\sqrt{T}}+\tfrac{L_{1}R^{2}}{T}\right).$ $$

Lan () developed an optimal method with a guaranteed convergence rate $\mathcal{O}\left(\tfrac{\sigma_{1}R}{\sqrt{T}}+\tfrac{L_{1}R^{2}}{T^{2}}\right)$, matching the worst-case lower bounds. The key advantage of SA is that it provides convergence guarantees for the original problem (6). Additionally, methods effective for the SA approach tend to perform well for the SAA approach as well.

###### Assumption 1 .

### 3.1 Stochastic Reformulation

Using the perspective of Stochastic Approximation, we can rewrite the linear DAG (1) as

$$ $x=X_{i}=\left[I-\mathbf{W}_{\ast}^{\top}\right]^{-1}N_{i},$ (10) $$

where $\mathbf{W}^{\ast}$ is a true DAG that corresponds to the full distribution, and our goal is to find DAG $\mathbf{W}$ that is close to $\mathbf{W}^{\ast}$. If we assume that $x=X_{i}$ is a random vector sampled from a distribution $\mathcal{D}$, we can express the objective function as an expectation,

$$ $\min_{\mathbf{W}\in\mathbb{D}}\mathbb{E}_{x\sim\mathcal{D}}\left[l(\mathbf{W}; x)\stackrel{{\scriptstyle def}}{{=}}\tfrac{1}{2}\|x-\mathbf{W}^{\top}x\|^{2}= \tfrac{1}{2}\|x^{\top}-x^{\top}\mathbf{W}\|^{2}\right].$ (11) $$

For $x$ from (10) we can calculate $\|x-\mathbf{W}^{\top}x\|=\|(I-\mathbf{W}^{\top})x\|=\|(I-\mathbf{W})\left[I-
\mathbf{W}_{\ast}^{\top}\right]^{-1}N_{i}\|$, which implies that the minimizer of (11) recovers the true DAG. Conversely, this is not the case for methods such as Zheng et al. (), Ng et al. (), and Bello et al. (), which are based on SAA approaches with losses (2), (3), (4), (5).

## 4 Scalable Framework

Instead of strictly enforcing DAG constraints throughout the entire iteration process, we propose a novel, scalable optimization framework that consists of three main steps:

- 1.
Running an optimization algorithm $\mathcal{A}_{1}$ without any DAG constraints, $\mathcal{A}_{1}:\mathbb{R}^{d\times d}\to\mathbb{R}^{d\times d}.$
- 2.
Finding a DAG that is “closest” to the current iterate using a projection $\psi:\mathbb{R}^{d\times d}\to(\mathbb{D},\Pi)$, which also returns its topological sorting.
- 3.
Running the optimization algorithm $\mathcal{A}_{2}$ while preserving the vertex order, $\mathcal{A}_{2}:(\mathbb{D};\Pi)\to\mathbb{D}.$

Figure: Algorithm 1 $\psi{}$DAG framework

### 4.1 Optimization for the fixed vertex ordering

Let us clarify how to optimize while preserving the vertex order in step 3 of the framework.
Given a DAG $\mathcal{G}$, we can construct its topological ordering, denoted as $ord(\mathcal{G})$. In this ordering, for every edge, the start vertex appears earlier in the sequence than the end vertex. In general, this ordering is not unique. In the space of DAGs with $d$ vertices $\mathbb{D}$, there are $d!$ possible topological orderings.

Once we have a topological ordering of the DAG, we can construct a larger DAG, $\hat{\mathcal{G}}$, by performing the transitive closure of $\mathcal{G}$. This new DAG $\hat{\mathcal{G}}$ contains all the edges of the original DAG, and additionally, it includes an edge between vertices $V_{i}$ and $V_{j}$ if there exists the path from $V_{i}$ to $V_{j}$ in $\mathcal{G}$. Thus, $\hat{\mathcal{G}}$ is an expanded version of $\mathcal{G}$.

Now, the question arises: is it possible to construct an even larger DAG that contains both $\mathcal{G}$ and $\hat{\mathcal{G}}$? The answer is yes! We call this graph the Full DAG, denoted by $\tilde{\mathcal{G}}$, which is constructed via full transitive closure. In $\tilde{\mathcal{G}}$, there is an edge from vertex $V_{i}$ to vertex $V_{j}$ if $i<j$ in topological ordering $ord(\mathcal{G})$. This makes $\tilde{\mathcal{G}}$ the maximal DAG that includes $\mathcal{G}$. Note that for every topological sort, there is a corresponding full DAG. So, there are a total of $d!$ different full DAGs in the space of DAGs with $d$ vertices $\mathbb{D}$.

We are now ready to discuss the optimization part. Let us formulate the following optimization problem

$$ $\min_{\mathbf{W}\in\mathbb{R}^{d\times d}}\mathbb{E}_{x\sim\mathcal{D}}\left[l (\mathbf{W}\cdot\mathbf{A};x)=\tfrac{1}{2}\|x-(\mathbf{W}\cdot\mathbf{A})^{ \top}x\|^{2}\right],$ (12) $$

where ($\cdot$) denotes elementwise matrix multiplication. In this formulation, $\mathbf{A}$ acts as a mask, specifying coordinates that do not require gradient computation. The problem (12) is a quadratic convex stochastic optimization problem, which can be efficiently solved using stochastic gradient descent (SGD)-type methods. These methods guarantee convergence to the global minimum, with a rate of $\mathcal{O}\left(\tfrac{\sigma_{1}R}{\sqrt{T}}+\tfrac{L_{1}R^{2}}{T}\right).$

Assume that $\mathcal{G}^{\ast}$ is the true DAG with a weighted adjacency matrix $\mathbf{W}^{\ast}$, which is the solution we aim to find. Next, we can have the true ordering $ord(\mathcal{G}^{\ast})$ and the true full DAG $\tilde{\mathcal{G}}^{\ast}$ with its adjacency matrix $\mathbf{A}(\tilde{\mathcal{G}}^{\ast})$. The optimization problem (11), with the solution $\mathbf{W}^{\ast}$, can be addressed by solving the optimization problem (12) with $\mathbf{A}=\mathbf{A}(\mathcal{G}^{\ast})$. This result indicates that, if we know the true topological ordering $ord(\mathcal{G}^{\ast})$, then we can recover the true DAG $\mathbf{W}^{\ast}$ with high accuracy. From a discrete optimization perspective, this approach significantly reduces the space of constraints from $2^{d^{2}-d}$ to $d!$.
To illustrate the specificity of the minimizer of the proposed problem, Figure 1 demonstrates that minimizing (11) over a fixed random vertex ordering does not approach the true solution of (11). "Correct order" curve demonstrates the convergence of (12) when the true ordering $ord(\mathcal{G}^{\ast})$ is known.

Figure: Figure 1: Minimization of (11) using SGD over a fixed topological ordering of vertices on graph type ER4 with $d=100$ vertices with Gaussian noise. Plots demonstrate that minimizing (11) over a fixed random vertex ordering does not approach the true solution of (11).
Refer to caption: x1.png

In Algorithm (1), in step 4, we find the closest ordering to the current point, and next in step 5 solve (12) for this ordering and its full DAG. The solution obtained for the given ordering is an approximate local minimum for the original problem (11), which leads us to the next theorem

###### Theorem 2 .

For an $L_{1}$-smooth function $F(\mathbf{W})=\mathbb{E}_{x\sim\mathcal{D}}\left[l(\mathbf{W};x)\right]$, Algorithm (1), with access to $\sigma_{1}$-stochastic gradients, converges to a local minimum of problem (11) at the rate

$$ $\mathbb{E}\left[F(\mathbf{W}_{T})-F(\mathbf{W}^{\ast})\right]\leq\mathcal{O} \left(\tfrac{\sigma_{1}R}{\sqrt{T}}+\tfrac{L_{1}R^{2}}{T}\right),$ $$

where $T$ is a number of SGD-type steps.

Note that for a fixed vertex ordering and fixed adjacency matrix $\mathbf{A}$, the objective (12) becomes separable, enabling parallel computation for large-scale problems. In this work, we solved the minimization problem (12) for the number of nodes up to $d=10^{4}$, at which point the limiting factor was the memory to store $\mathbf{W}\in\mathbb{R}^{d\times d}$. Through parallelization and efficient memory management, it is possible to solve even larger problems.

###### Theorem 2 .

### 4.2 Method

We now introduce the method $\psi{}$DAG, which implements the framework outlined in Algorithm 1.

For simplicity, we select algorithm $\mathcal{A}_{1}$ as $\tau_{1}$ steps of Stochastic Gradient Descent (SGD). Similarly, $\mathcal{A}_{2}$ consists of $\tau_{2}$ steps SGD, where gradients are projected onto the space spanned by DAG’s topological sorting, thus preserving the vertex order. It is important to reiterate that SGD is guaranteed to converge to the neighborhood of the solution. In the implementation, we employed an advanced version of SGD, Universal Stochastic Gradient Method from (Rodomanov et al., ).

The implementation of the projection method is simple as well. We compute a “closest” topological sorting and remove all edges not permitted by this ordering.
The topological sorting is computed by a heuristic that calculates norms of all rows and columns to find the lowest value $v_{i}$. The corresponding vertex $i$ is then assigned to the ordering based on the following rule:

- •
If $v_{i}$ was the column norm, $i$ is assigned to the beginning of the ordering.
- •
If $v_{i}$ was the row norm, $i$ is assigned to the end of the ordering.

This step reduces the number of vertices, and the remaining vertices are topologically sorted through by a recursive call. We formalize this procedure in Algorithm 2. Note that this procedure can be efficiently implemented without recursion and with the computation cost $\mathcal{O}(d^{2})$.

Figure: Algorithm 2 Projection $\psi(\mathbf{W})$ computing the “closest” vertex ordering (recursive form)

Figure: Algorithm 3 $\psi{}$DAG

## 5 Experiments

We experimentally compare our newly proposed algorithm $\psi{}$DAG(^2^22Code implementing the proposed algorithm is available at [https://github.com/kz29/-psiDAG](https://github.com/kz29/-psiDAG). We use the Universal Stochastic Gradient Method from (Rodomanov et al., ) as the inner optimizer. The code is available in the OPTAMI package [https://github.com/OPTAMI/OPTAMI](https://github.com/OPTAMI/OPTAMI) (Kamzolov et al., ).) to other score-based methods for computing linear DAGs, NOTEARS (Zheng et al., ), GOLEM(^3^33In all experiments, we consider GOLEM-EV, where the noise variances are equal.) (Ng et al., ) and DAGMA (Bello et al., ).
As it is established that DAGMA Bello et al. () is an improvement over NOTEARS Zheng et al. (), we use mostly the former one in our experiments.
As the baseline algorithms were implemented without extensive hyperparameter tuning, we avoided hyperparameter tuning as much as possible.
In particular, we apply the same threshold as the one in Zheng et al. (), Ng et al. (), Bello et al. () across all scenarios.

Figure 2 shows that $\psi{}$DAG consistently exhibits faster convergence across different noise distributions.
Appendix B extends this result across different graph sizes and graph types.

### 5.1 Synthetic Data Generation

Figure: Figure 2: Linear SEM methods of $\psi{}$DAG, GOLEM and DAGMA on graphs of type ER4 with $d=1000$ number of nodes and with different noise distributions: Gaussian (first), exponential (second), and Gumbel (third).
Refer to caption: x3.png

We generate ground truth DAGs to have $d$ nodes and an average of $k\times d$ edges, where $k\in\{2,4,6\}$ is a sparsity parameter.
The graph structure is determined by the choice of the graph models to be either Erdős-Rényi (ER) or Scale Free (SF), and together with sparsity parameter $k$, we refer to them as $\text{ER}k$ or $\text{SF}k$.
Each of the edges has assigned a random weight uniformly sampled from the interval $[-1,-0.05]\cup[0.05,1]$.

Following the linear Structural Equation Model (SEM), the observed data $X$ has form $X=N(I-W)^{-1}$, where $N\in\mathbb{R}^{n\times d}$ represents $n$ $d$-dimensional independent and identically distributed (i.i.d.) noise samples drawn from either Gaussian, exponential or Gumbel distributions.
In this study, we focus on an equal variance (EV) noise setting, with a scale factor of 1.0 applied to all variables. Unless otherwise specified, we generate the same number of samples $n\in\{5000,10000\}$ for training and validation datasets, respectively.
A more detailed description can be found in Appendix A.

### 5.2 Scalability Comparison

Figure: (a) ER2 graph type
Refer to caption: x9.png

Figure: (a) ER2
Refer to caption: x15.png

In this section, we discuss the runtime speed of the proposed $\psi{}$DAG algorithm. We run the proposed algorithm and baselines until the function value converges close to the solution,
$f(x_{k})-f(\overline{x})\leq 0.1\cdot f(\overline{x})$.

Figures 3(a) and 3(b) compare the performance of $\psi{}$DAG against GOLEM and DAGMA on smaller graphs with various structures and noise distributions. Meanwhile, Figure 4 illustrates the scalability of $\psi{}$DAG on the large graphs.

Both Figures 3 and 4 clearly demonstrate that $\psi{}$DAG significantly outperforms GOLEM and DAGMA in terms of runtime across both sparse and dense ER graphs in nearly all considered scenarios. The only exception occurs with very small graphs $d<100$ and high sparsity (ER2, SF2), where DAGMA is marginally faster than $\psi{}$DAG.
However, as graph size and density increase, $\psi{}$DAG scale efficiently across all scenarios even up until $d=10000$ nodes. In the case of sparse graphs (Figures 4(a), 4(b)), $\psi{}$DAG consistently converges within a few hours, even for $d=10000$ nodes.

In contrast, increasing graph size causes the computational cost of both GOLEM and DAGMA to skyrocket. Notably, across all tested graphs (4(a)), GOLEM exceeds allocated runtime of $36$ hours for $d\geq 3000$ nodes, while DAGMA exceeds for $d\geq 5000$ nodes.

In several experimental scenarios, we also observed that both GOLEM and DAGMA occasionally failed to meet the stopping criterion, even for smaller graphs. For very small ER6 graphs ($d=100$, Figure 4(c)), neither method consistently achieved the stopping criterion – DAGMA failed to converge once, and GOLEM failed twice out of three random seeds. With $d=1000$ nodes, DAGMA again failed to converge in one out of the three runs.
All non-converging runs were excluded from the figures.

### 5.3 Real Data

We also evaluate the proposed method against baselines NOTEARS (Zheng et al., ), GOLEM (Ng et al., ), and DAGMA (Bello et al., ) on a real-world dataset, *causal protein signaling network data*, provided by Sachs et al. () that captures the expression levels of proteins and phospholipids in human cells.
This dataset is widely used in the literature on probabilistic graphical models, with experimental annotations that are well-established in the biological research community.

The dataset comprises 7,466 samples, of which we utilize the first 853, corresponding to a network with 11 nodes representing proteins and 17 edges denoting their interactions. Despite its relatively small size, it is considered to be a challenging benchmark in recent studies (Zheng et al., ; Ng et al., ; Gao et al., ).
For all experiments, we used the first 853 samples for training and the subsequent 902 samples for testing. After the training phase, we employed the same default threshold of 0.3 as was used by the other baseline approaches NOTEARS, GOLEM, DAGMA.

As shown in LABEL:real_data, our method outperforms both baselines GOLEM (Ng et al., ) and NOTEARS (Zheng et al., ) in all metrics, the SHD (lower is better), TPR (higher is better) and FPR (smaller is better). More detailed description can be found in Appendix A.
We report the total number of edges of the output DAG. We do not report the performance of DAGMA because it fails to optimize the problem (its iterate $\mathbf{W}$ diverges from the feasible domain during the first iteration).

**Table 1: Performance of the top-performing methods on the causal protein signaling network dataset Sachs et al. (). The threshold for all methods is 0.3.**
|  | SHD$(\downarrow)$ | TPR $(\uparrow)$ | FPR $(\downarrow)$ | Total edges | Reference |
| --- | --- | --- | --- | --- | --- |
| GOLEM | 26 | 0.294 | 0.47 | 23 | Ng et al. () |
| NOTEARS | 15 | 0.294 | 0.26 | 15 | Zheng et al. () |
| $\psi{}$DAG | 14 | 0.411 | 0.18 | 14 | Algorithm 3 |

## 6 Conclusion

We introduce a novel framework for learning Directed Acyclic Graphs (DAGs) that addresses the scalability and computational challenges of existing methods. Our approach leverages Stochastic Approximation techniques in combination with Stochastic Gradient Descent (SGD)-based methods, allowing for efficient optimization even in high-dimensional settings. A key contribution of our framework is the introduction of new projection techniques that effectively enforce DAG constraints, ensuring that the learned structure adheres to the acyclicity requirement without the need for computationally expensive penalties or constraints seen in prior works.

The proposed framework is theoretically grounded and guarantees convergence to a feasible local minimum. One of its main advantages is its low iteration complexity, making it highly suitable for large-scale structure learning problems, where traditional methods often struggle with runtime and memory limitations. By significantly reducing the per-iteration cost and improving convergence behavior, our framework demonstrates superior scalability when applied to larger datasets and more complex graph structures.

We validate the effectiveness of our method through extensive experimental evaluations across a variety of settings, including both synthetic and real-world datasets. These experiments show that our framework consistently outperforms existing methods such as GOLEM (Ng et al., ), NOTEARS (Zheng et al., ), and DAGMA (Bello et al., ), both in terms of objective loss and runtime, particularly in scenarios involving large and dense graphs. Furthermore, our method exhibits robust performance across different types of graph structures, highlighting its potential applicability to various practical fields such as biology, finance, and causal inference.

## Appendix A Detailed Experiment Description

#### Computing.

Our experiments were carried out on a machine equipped with 80 CPUs and one NVIDIA Quadro RTX A6000 48GB GPU. Each experiment was allotted a maximum wall time of 36 hours as in DAGMA Bello et al. ().

#### Graph Models.

In our experimental simulations, we generate graphs using two established random graph models:

- •
Erdős-Rényi (ER) graphs: These graphs are constructed by independently adding edges between nodes with a uniform probability. We denote these graphs as ERk, where $kd$ represents the expected number of edges.
- •
Scale-Free (SF) graphs: These graphs follow the preferential attachment process as described in Barabási & Albert (). We use the notation SFk to indicate a scale-free graph with expected $kd$ edges and an attachment exponent of $\beta=1$, consistent with the preferential attachment process. Since we focus on directed graphs, this model corresponds to Price’s model, a traditional framework used to model the growth of citation networks.

It is important to note that ER graphs are inherently undirected. To transform them into Directed Acyclic Graphs (DAGs), we generate a random permutation of the vertex labels from 1 to $d$, then orient the edges according to this ordering. For SF graphs, edges are directed as new nodes are added, ensuring that the resulting graph is a DAG. After generating the ground-truth DAG, we simulate the structural equation model (SEM) for linear cases, conducting experiments accordingly.

#### Metrics.

The performance of each algorithm is assessed using the following four key metrics:

- •
Structural Hamming Distance (SHD): A widely used metric in structure learning that quantifies the number of edge modifications (additions, deletions, and reversals) required to transform the estimated graph into the true graph.
- •
True Positive Rate (TPR): This metric calculates the proportion of correctly identified edges relative to the total number of edges in the ground-truth DAG.
- •
False Positive Rate (FPR): This measures the proportion of incorrectly identified edges relative to the total number of absent edges in the ground-truth DAG.
- •
Runtime: The time taken by each algorithm to complete its execution provides a direct measure of the algorithm’s computational efficiency.
- •
Stochastic gradient computations: Number of gradient computed.

#### Linear SEM.

In the linear case, the functions are directly parameterized by the weighted adjacency matrix $W$. Specifically, the system of equations is given by $X_{i}=X\mathbf{W}_{i}+N_{i},$ where $\mathbf{W}=[\mathbf{W}_{1}|\cdots|\mathbf{W}_{d}]\in\mathbb{R}^{d\times d}$, and $N_{i}\in\mathbb{R}$ represents the noise. The matrix $\mathbf{W}$ encodes the graphical structure, meaning there is an edge $X_{j}\to X_{i}$ if and only if $W_{j,i}\neq 0$. Starting with a ground-truth DAG $B\in\{0,1\}^{d\times d}$ obtained from one of the two graph models, either ER or SF, edge weights were sampled independently from $\text{Unif}[-1,-0.05]\cup[0.05,1]$ to produce a weight matrix $\mathbf{W}\in\mathbb{R}^{d\times d}$. Using this matrix $\mathbf{W}$, the data $X=X\mathbf{W}+N$ was sampled under the following three noise models:

- •
Gaussian noise: $N_{i}\sim N(0,1)$ for all $i\in[d]$,
- •
Exponential noise: $N_{i}\sim\text{Exp}(1)$ for all $i\in[d]$,
- •
Gumbel noise: $N_{i}\sim\text{Gumbel}(0,1)$ for all $i\in[d]$.

Using these noise models, random datasets $X\in\mathbb{R}^{n\times d}$ were generated by independently sampling the rows according to one of the models described above. Unless specified otherwise, each simulation generated $n=5000$ training samples and a validation set of $10,000$ samples.

The implementation details of the baseline methods are as follows:

- •
NOTEARS Zheng et al. ()was implemented using the authors’ publicly available Python code, which can be found at [https://github.com/xunzheng/notears](https://github.com/xunzheng/notears). This method employs a least squares score function, and we used their default set of hyperparameters without modification.
- •
GOLEM Ng et al. () was implemented using the authors’ Python code, available at [https://github.com/ignavierng/golem](https://github.com/ignavierng/golem), along with their PyTorch version at [https://github.com/huawei-noah/trustworthyAI/blob/master/gcastle/castle/algorithms/gradient/notears/torch/golem_utils/golem_model.py](https://github.com/huawei-noah/trustworthyAI/blob/master/gcastle/castle/algorithms/gradient/notears/torch/golem_utils/golem_model.py). We adopted the default hyperparameter settings provided by the authors, specifically $\lambda_{1}=0.02$ and $\lambda_{2}=5$. For additional details of their method, we refer to Appendix F in Ng et al. ()
- •
DAGMA Bello et al. () was implemented using the authors’ Python code, which is available at [https://github.com/kevinsbello/dagma](https://github.com/kevinsbello/dagma). We used the default hyperparameters provided in their implementation.

#### Thresholding.

Following the approach taken in previous studies, including the baseline methods (Zheng et al., ; Ng et al., ; Bello et al., ), for all the methods, we apply a final thresholding step of 0.3 to effectively reduce the number of false discoveries.

## Appendix B Experiments

Plots show performance of the algorithms $\psi{}$DAG, GOLEM, DAGMA for combinations of number of vertices $d\in\{10,50,100,500,1000,3000,5000,10000\}$, graph types $\in\{ER,SF\}$, average density of graphs $k\in\{2,4,6\}$, and noise types to be either Gaussian, Exponential, or Gumbel.

Plots are grouped by the noise type and the number of vertices of the graph and arranged into figures.

ER graph types: Figures 5 and 6 show performance on ER2 graphs, Figures 7 and 8 show performance on ER4 graphs,
Figure 9 shows performance on ER6 graphs.

SF graph types: Figures 10 and 11 show performance on SF2 graphs, Figures 12 and 13 show performance on SF4 graphs,
Figure 14 shows performance on SF6 graphs.

We report a functional value decrease compared to i) time elapsed and ii) number of gradients computed, which also serves as a proxy of time.

Figure 6(b) shows that DAGMA requires a significantly larger amount of gradient computations compared to both $\psi{}$DAG and GOLEM.

### B.1 Small to Moderate Number of Nodes

Our experiments demonstrate that while number of nodes is small, $d<100$, GOLEM is more stable than DAGMA, and $\psi{}$DAG method is the most stable.
While DAGMA shows impressive speed for smaller node sets, the number of iterations required is still higher than both GOLEM and our method.
Across all scenarios, $\psi{}$DAG consistently demonstrates faster convergence compared to the other approaches, requiring fewer iterations to reach the desired solution.

### B.2 Large Number of Nodes

For graphs with a large number of nodes $d\in\{5000,10000\}$, we were unable to run neither of the baselines, and consequently, Figure 15 includes only one algorithm. GOLEM was not feasible due to its computation time exceeding $350$ hours. DAGMA was impossible as its runs led to kernel crashes. In all cases, we utilized a training set of 5,000 samples and a validation set of 10,000 samples.

### B.3 Denser Graphs

For a thorough comparison, in Figures 9 and 14, we compare graph structures ER6 and SF6 under the Gaussian noise type.
Plots indicate that while DAGMA exhibits a fast runtime when the number of nodes is small, $d<100$, it requires more iterations to achieve convergence. Algorithm $\psi{}$DAG consistently outperforms GOLEM and DAGMA in both training time and a number of stochastic gradient computations, and the difference is more pronounced for lager number of nodes and denser graphs.

Figure: (a) $d=10$ vertices
Refer to caption: x18.png

Figure: (a) $d=500$ vertices
Refer to caption: x36.png

Figure: (a) $d=10$ vertices
Refer to caption: x50.png

Figure: (a) $d=500$ vertices
Refer to caption: x68.png

Figure: (a) $d=10$ vertices
Refer to caption: x82.png

Figure: (a) $d=10$ vertices
Refer to caption: x94.png

Figure: (a) $d=500$ vertices
Refer to caption: x112.png

Figure: (a) $d=10$ vertices
Refer to caption: x126.png

Figure: (a) $d=500$ vertices
Refer to caption: x144.png

Figure: (a) $d=10$ vertices
Refer to caption: x156.png

Figure: (a) ER2
Refer to caption: x168.png

Figure: Figure 16: Comparison of $\psi{}$DAG, $\psi{}$DAG weighted and DAGMA for ER2 graph with $d=3000$ nodes and Gaussian noise.
Refer to caption: x176.png

## Appendix C Weighted Projection

Inspired by the importance of sampling, we considered adjustment of the projection method by weights. Specifically, we considered the elements of the $\mathbf{W}{}$ to be weighted element-wisely by the second directional derivatives of the objective function, $\mathbf{L}[i][j]\stackrel{{\scriptstyle def}}{{=}}\left(\frac{d}{d\mathbf{W}[i
][j]}\right)^{2}\mathbb{E}_{X\sim\mathcal{D}}\left[l(\mathbf{W};X)\right]$.
As we don’t have access to the whole distribution $\mathcal{D}$, we approximate it by the mean of already seen samples,

$$ $\mathbf{L}_{k}[i][j]\stackrel{{\scriptstyle def}}{{=}}\left(\frac{d}{d\mathbf{ W}[i][j]}\right)^{2}\frac{1}{k}\sum_{k=0}^{k-1}{l\left(\mathbf{W};X_{k}\right) }=\frac{1}{k}\sum_{t=0}^{k-1}\left(X_{k}[j]\right)^{2}.$ (13) $$

Weights (13) are identical for whole columns; hence, they impose storing only one vector. Updating them requires a few element-wise vector operations.

Figure: (a) $d=3000$ vertices
Refer to caption: x178.png

Figures 16 and 17 show that this weighting can lead to an improved convergence (slightly faster convergence to a slightly lower functional value) without imposing extra cost. However, this improvement over runtime is not consistent; hence, for simplicity, we deferred this to the appendix.

## Peer Reviews

confidence: 2
contribution: 2
initial_rating: 5
presentation: 1
questions: - Where is the proof for Theorem 2? What is the $R$ in the statement? How do we relate this to other optimization results on DAG learning in literature?
- What is the intuition behind Alg 2? There is not explanation why it comes from and why we want to recover the ordering in this way. Any guarantee if it consistently outputs a valid ordering?
soundness: 2
strengths: - The proposed method relies on SGD and is fast compared to other continuous optimization approaches for DAG learning.
- Simulation experiments are conducted to show the proposal gives faster convergence speed relative to the optimum.
summary: This paper studies structure learning in linear DAG models, e.g. to recover the underlying DAG from observational data. 
It follows score based learening approach and tries to optimize the square loss. It proposes to apply stochastic gradient descent to solve the optimization problem, and then introduces a procedure to estimate the topological ordering based on learned coefficient matrix.
The experiments are conducted on equal variance noise setting to compare with existing methods.
weaknesses: - The presentation is poor, many concepts are used without introduction, e.g. transitive closure, full DAG, $R$ in convergence statement.
- In the first step, SGD is applied to solve (11). But actually, the solution to (11) is simply W=I analytically. Also, fixing ordering, the third step  has analytical solution that is simply linear regression of each node onto preceding nodes. All SGD optimizations can be avoided.
- Ordering based DAG learning is a long line of research. This paper lacks citation and discussion on that. Especially, for the equal variance setup considered in this paper, there are scalable algorithms:

[1] Chen, W., Drton, M., & Wang, Y. S. (2019). On causal discovery with an equal-variance assumption. Biometrika, 106(4), 973-980.

[2] Gao, M., Tai, W. M., & Aragam, B. (2022). Optimal estimation of Gaussian DAG models. In International Conference on Artificial Intelligence and Statistics (pp. 8738-8757). PMLR.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. Can the stochastic approximation method proposed in this paper be extended to nonlinear SEMs? If so, a discussion on the potential challenges or limitations in this direction would be insightful.
  
2. It appears that Theorem 2 may have been previously developed by Nemirovski and Yudin (1983). Could the authors clarify the novelty of this result, or highlight any differences from the original work?
  
3. Could the authors elaborate further on the scalability of the proposed method, particularly given the projection method’s complexity?
  
4. Would the authors consider including the topological swap method from [A] as a baseline? This comparison would provide additional insights into the performance of the proposed method relative to related approaches.
soundness: 2
strengths: 1. This paper is well-written and easy to follow.
  
2. This paper introduces a novel stochastic approximation-based method for learning the DAG structure.
  
3. The paper provides comprehensive numerical experiments, which effectively demonstrate both the performance and efficiency of the proposed methods.
summary: This paper proposes a novel stochastic approximation framework for learning DAGs, integrating SGD with efficient projection methods to enforce acyclicity. By addressing the non-convexity and complexity challenges of prior approaches, the method achieves low iteration complexity and demonstrates superior scalability and performance in extensive experiments.
weaknesses: 1. The stochastic approximation approach presented in this paper is limited to linear structural equation models (SEMs). In contrast, other methods such as GOLEM and DAGMA offer extensions to nonlinear SEMs. Expanding the proposed method to nonlinear settings could enhance its contribution.
  
2. The projection technique introduced in Algorithm 2 appears to lack a clear theoretical guarantee. Moreover, in line 339, the complexity of the projection method is stated as $\mathcal{O}(d^2)$. Given that the numerical experiments involve cases where $d$ can reach $10^4$, the complexity could escalate to approximately 108. This computational burden seems significant and may pose challenges for scalability in practice.
  
3. The paper could benefit from additional relevant references. For example, Deng et al. [A] propose a topological swap method to accelerate the optimization of the NOTEARS objective by identifying the topological order first and then learning the edge weights of the directed graph. This process aligns with the ideas explored in this paper, and including such a baseline would strengthen the completeness and contextualization of the work.

[A] Deng, Chang, et al. "Optimizing NOTEARS objectives via topological swaps." *International Conference on Machine Learning*. PMLR, 2023.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. Page 1 Line 37: Missing an "and" at the end of the sentence.
2. Page 2 Line 86: Should it say "handling up to 10000 *nodes*"?
3. Page 3 Line 114: Is $X=(X_1,\ldots,X_n)$ a concatenation of column vectors? If so, I think the quantity on the RHS should be transposed. Same for $N$ on Line 126.
4. Page 3: I found it strange that the matrix $W$ is boldface while the matrix $X$ is not.
5. Page 4 Line 168: The $\overset{def}{=}$ looks different here compared with other equations.
6. Page 4 Line 185: The bit that says "where $\xi\in\Xi$ is a random variable drawn from the distribution $\Xi$" confused me. I think it should read "where $\xi$ is a random variable that follows the distribution $\Xi$".
7. Assumption 1: The quantity $\sigma_1^2$ is not defined.
8. Page 4 Line 211: The quantity $R$ is not explained.
9. Page 5 Line 263: $V_J$ should be $V_j$.
10. Figure 1: What is $\bar{x}$?
11. Page 7 Line 330: Can you be more precise than "closest"? Since the projection is heuristic, the word "close" might be better.
12. Page 7 Line 337: What does "topologically through" mean?
13. Figures 2, 3, and 4: The font size is too small for me to read without zooming in.
14. Page 15 Line 799: I am confused about the reference to GPUs in the computing section. Do any of the algorithms use GPUs?
15. Page 17 Line 866: If the edges in the ground truth DAG can be as small as 0.05, why apply a threshold of 0.3? That would make it impossible to ever recover the true graph.
16. Appendix A: I might have missed it, but I could not see the value of $\lambda$ used by $\psi$DAG.
soundness: 2
strengths: The scalability of $\psi$DAG is impressive, and I found the computational experiments convincing. The paper is well-written and contains a good mix of material.
summary: This paper introduces $\psi$DAG, a new method for learning a directed acyclic graph (DAG) from observational data. The new method aims to improve on existing approaches in terms of scalability. Its theoretical complexity is $O(d^2)$, improving on $O(d^3)$ for popular alternatives. This theoretical improvement is demonstrated in computational experiments.
weaknesses: 1. One of the main contributions outlined on page 2 is a proof of Algorithm 1's convergence rate (Theorem 2). I could not see a proof of Theorem 2 in the paper or appendix.
2. The projection algorithm (Algorithm 2) is a key component of the computational machinery, but it remains opaque. It would be insightful to write this projection in terms of an optimization problem and discuss whether the algorithm attains a stationary point of this problem.
3. My primary concern with this paper is the lack of synthetic experiments comparing the *statistical* properties of the competing approaches. Only the computational merits of each approach are compared. It is essential to understand how $\psi$DAG performs on the synthetic data in terms of SHD, TPR, and FPR.
4. A few aspects of the real data example need clarification. First, the reason for having training and testing sets of size $n=853$ and $n=902$ instead of the whole dataset ($n=7466$) is not explained. Second, the role of the testing set is unclear, since the metrics reported do not require a testing set. Finally, I have seen DAGMA run on this dataset before, so I am confused about why it fails now.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: No further quesions, questions are addressed above.
soundness: 2
strengths: This paper is the first I have encountered in differentiable DAG learning that integrates a stochastic algorithm, presenting a novel and compelling approach. A primary advantage of the stochastic algorithm lies in its computational efficiency, offering a marked improvement over conventional methods. The authors have also conducted an extensive set of experiments to substantiate the algorithm's efficiency, which I find to be thorough and convincing.
summary: The paper leverages stochastic optimization within a newly proposed framework for DAG learning, introducing a novel algorithm. The authors effectively demonstrate the efficiency of this algorithm through comprehensive experimentation.
weaknesses: The paper’s approach is limited by its reliance on a simple linear model with equal variance noise, resulting in an easy optimization problem (least squares with $\ell_1$ penalty). When the graph order is fixed, the problem reduces to a Lasso optimization, which is computationally trivial for large dimensions (e.g., $d =1000,n =5000$), taking only a few seconds or minutes to solve. This simplicity undermines the significance of the study. To make the work more meaningful, the authors might consider applying their stochastic algorithm to more complex structural equation models (SEM), where optimization challenges are more substantial.


Additionally, several statements in the paper lack rigor. For example, the abstract and Theorem 2 claim that the algorithm converges to a feasible local minimum, but no theoretical proof or justification is provided for Algorithm 1, which is crucial given the NP-complete nature of the problem in (11). Original guarantee of local optimal solution from [1][2] can not be directly applied here, since it is unconstrained optimization and objective function is convex. Please refer to [3][4], these two paper study the local optimality of problem (11) indetails, they should be included and discussed appropirately in the paper. Furthermore, the transitive closure of $\mathcal{G}$ in lines 261-262 lacks a formal definition, which should be added for clarity.

The experimental setup also introduces concerns. The authors sample edge weights from a non-standard range, uniformly from $[−1, −0.05] \cup [0.05, 1]$, which deviates from established practices. Applying the same threshold (0.3) as methods like NOTEARS, DAGMA, and GOLEM in this unconventional setting risks removing correctly recovered edges, potentially skewing baseline performance and inflating the loss gap. To address this, the authors should consider aligning with the edge weight sampling practices of prior studies to ensure comparability. Additionally, incorporating Structural Hamming Distance (SHD) as a recovery accuracy measure and including error bars to capture result variability would improve the robustness and interpretability of findings.

[1] Arkadi Nemirovski and David Yudin, Problem Complexity and Method Efficiency in Optimization
[2] Guanghui Lan, An optimal method for stochastic composite optimization
[3] Dennis Wei, Tian Gao, Yue Yu, DAGs with No Fears: A Closer Look at Continuous Optimization for Learning Bayesian Networks
[4] C Deng, K Bello, B Aragam, PK Ravikumar, Optimizing NOTEARS objectives via topological swaps
[4] PL Loh, P Bühlmann, High-dimensional learning of linear causal networks via inverse covariance estimation

Meta-review: The paper studies the problem of learning the structure of directed acyclic graphs from data. The main contribution is an optimization based approach for the learning problem that leverages stochastic gradient descent to obtain an efficient and scalable algorithm.

The main strength of this work is that the SGD based algorithm proposed is more efficient and scalable than existing continuous optimization methods for the problem. The reviewers noted several weaknesses and concerns, including that the proposed approach is limited to linear models, some components of the approach lack theoretical justification, the discussion on related work needs to be significantly expanded, and the experimental methodology lacks an evaluation of the statistical performance of the algorithm. Overall, the reviewers appreciated the scalability of the proposed approach but there was consensus that the paper would need additional work and improvements in order for it to be accepted.

Decision: Reject