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

We sincerely thank the reviewer for their helpful comments and would like to provide answers to all the concerns raised.

---

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 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, d =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.

Answer: As we have demonstrated in Figure 1, a solution of the optimization problem over a fixed (random) graph ordering is far from the solution of the optimization problem minimized over all DAGs. Therefore, we claim that the results of our paper are already meaningful.

We restricted our setup to linear SEMs just for the sake of simplicity. Extensions to nonlinear SEMs are possible and are the focus of our ongoing research.

---

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).

Answer: We would like to apologize for the missing proof of Theorem 2. The proof is a direct consequence of the application of the algorithm Universal Stochastic Gradient Method (Alg.3 of [SGD]).

In particular, Theorem 4.2 of [SGD] guarantees that applied to convex L-Hölder continuous functions in the domain of radius $R$, $||x-y|| \leq R$, in $k$ steps decreases functional loss at rate $\mathcal{O} (\frac{LR}{k^2} + \frac {\sigma R}{\sqrt k} )$. In our case, once our Algorithm 1 fixes the ordering $\pi_k$ in line 4, DAGs allowed by ordering $\pi_k$ form a linear subspace of $R^{d \times d}$ and in that subspace loss (11) is smooth and convex, and $k$ steps Universal Stochastic Gradient Method converge to a minimum of the ordering at rate given by Theorem 4.2 of [SGD]. This minimum of the ordering is a local minimum of the loss (11).

[SGD] Universal Gradient Methods for Stochastic Convex Optimization, Anton Rodomanov, Ali Kavis, Yongtao Wu, Kimon Antonakopoulos, Volkan Cevher, link:https://openreview.net/pdf?id=Wnhp34K5jR

---

...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.

Answer: Thank you. We will include these papers in our paper.

Reference [3] investigates the NOTEARS framework for continuous optimization in Bayesian network structure learning, providing new theoretical insights into the KKT conditions for such formulations. They propose combining NOTEARS with a local search method informed by KKT conditions (KKTS), leading to improvements in SHD.

While [3] focuses on enhancing existing algorithms with a post-processing step, our work takes an orthogonal approach, reformulating DAG learning as a stochastic optimization problem. We develop an entirely new framework that seamlessly incorporates stochastic gradient methods, allowing efficient processing of large graphs (up to 10,000 nodes), exceeding the computational limits of NOTEARS and [3].

Reference [4] presents a bi-level algorithm that iteratively refines topological orders through node swaps. By combining off-the-shelf linear constraint solvers with a novel method for identifying node pairs for swaps, it finds local minima or KKT points under relaxed conditions, achieving solutions with lower scores than state-of-the-art methods. This approach can also serve as a post-processing step to refine the outputs of other algorithms.

Comparing our work to [4], we again adopt an orthogonal approach by reformulating a minimization over DAGs as a stochastic optimization task. While the topological swaps method emphasizes refining solutions and improving local optimality, our work addresses the challenges of large-scale graph learning, providing a scalable framework.

We would like to address the raised concerns.

In the case of a linear model, the optimal solution of this step is typically close to the identity matrix I, which is significantly distant from satisfying the directed acyclic graph (DAG) constraints.

We believe there is a misunderstanding. The minimization problem (11) is defined over directed acyclic graphs, and if W represents a valid DAG, all elements on its diagonal are necessarily zero. So, during the whole algorithm, the diagonal of $W_k$ is enforced to be zero. We added these details in the updated paper. Therefore, the solution of (11) is not close to the identity matrix.

---

Theorem 2 & local optimas

You are correct; the solution within a fixed ordering is not necessarily a local solution of (11). We sincerely apologize for this oversight. To address this, we are correcting the claim in Theorem 2 and relegating it to Appendix B.

We are grateful to the reviewer for the thoughtful review. We are happy to respond to each of the concerns raised.

---

W1: 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.

Answer: We would like to apologize for that. The proof of Theorem 2 is a direct application of the optimization algorithm Universal Stochastic Gradient Method (Alg.3 of [SGD]). In particular, Theorem 4.2 of [SGD] guarantees that applied to convex L-Hölder continuous functions in the domain of radius $R$, $||x-y|| \leq R$, in $k$ steps decreases functional loss at rate $\mathcal{O} (\frac{LR}{k^2} + \frac {\sigma R}{\sqrt k} )$. In our case, once our Algorithm 1 fixes the ordering $\pi_k$ in line 4, DAGs allowed by ordering $\pi_k$ form a linear subspace of $R^{d \times d}$ and in that subspace loss (11) is smooth and convex, and $k$ steps Universal Stochastic Gradient Method converge to a minimum of the ordering at rate given by Theorem 4.2 of [SGD]. This minimum of the ordering is a local minimum of the loss (11).

[SGD] Universal Gradient Methods for Stochastic Convex Optimization, Anton Rodomanov, Ali Kavis, Yongtao Wu, Kimon Antonakopoulos, Volkan Cevher, link:https://openreview.net/pdf?id=Wnhp34K5jR

---

W2: 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.

Answer: We would like to point out that the space of all DAGs is a union of numerous linear subspaces, which is extremely nonconvex. Finding the projection (i.e., the closest DAG to a given matrix) is a nontrivial problem that is likely NP-hard. Currently, we do not know how to quantify how close the output of our Alg.2 is to a stationary point of (any) optimization problem.

It is important to note that our framework is not tied to this specific projection method. Any alternative projection method can be seamlessly integrated into our approach.

---

W3: 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.

Answer: We would like to point out that the statistical properties (SHD, TPR, FPR, etc.) depend on the solution of the optimization problem (11) rather than the optimization algorithms for solving it. In particular, all algorithms minimizing the same optimization problem will output an identical point.

We would like to clarify that the primary contribution of the paper was to establish a flexible and scaleable computational framework for minimizing objective function (11). Adjustments of the objective function to optimize statistical metrics (SHD, TPR, FPR, etc.) were beyond the project's scope and are a direction of our ongoing research.

---

W4: A few aspects of the real data example need clarification. First, the reason for having training and testing sets of size n =857 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.

Answer: We choose only the first subset ($n=853$) to train the models to be consistent with papers [1,2,3,4], which consider the full dataset to be a challenging benchmark (see lines 499-503 in our paper). As per the reviewer's request, we added the results for the whole dataset in Appendix B.

W4: ...Finally, I have seen DAGMA run on this dataset before, so I am confused about why it fails now.

Answer: We have been puzzled by that as well.

To clarify the situation, the original DAGMA paper does not present a comparison on real datasets. We are aware of only one paper presenting DAGMA evaluation on real datasets, namely [4]. We used the implementation of DAGMA published on their GitHub without any modifications, but we were unable to make it work for neither sample size $n=853$ nor $n=7466$.

---

Regarding questions, typos and presentation suggestions:

Answer: We are very grateful to you for helping us improve our paper. Your feedback has been invaluable in polishing our work.

[1] Xun Zheng, Bryon Aragam, Pradeep Ravikumar, and Eric Xing. DAGs with no tears: Continuous optimization for structure learning. Advances in Neural Information Processing Systems, 31, 2018.

[2] Ignavier Ng, AmirEmad Ghassami, and Kun Zhang. On the role of sparsity and DAG constraints for learning linear DAGs. Advances in Neural Information Processing Systems, 33:17943–17954,2020.

[3] Yinghua Gao, Li Shen, and Shu-Tao Xia. DAG-GAN: Causal structure learning with generative adversarial nets. In ICASSP 2021-2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 3320–3324. IEEE, 2021.

[4] Misiakos, Panagiotis, Chris Wendler, and Markus Püschel. "Learning DAGs from data with few root causes." Advances in Neural Information Processing Systems 36 (2024).

I thank the authors for their response. However, my main concern regarding the statistical performance of the new approach remains unresolved. In particular, (11) is a nonconvex optimization problem, so the algorithms applied to solve it need not output identical solutions as they do not guarantee a globally optimal solution.

I'm still positive about the scalability of the method and believe the paper is worthy of publication in the future after some additional work.

We appreciate the reviewer’s constructive feedback and are grateful for the opportunity to address all the raised concerns.

Q1: 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

Answer: For simplicity, we have restricted our results to linear SEMs in this paper. However, extending the proposed method to nonlinear SEMs is indeed possible and forms the focus of our ongoing research efforts.

---

Q2: 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?

Answer: We would like to clarify that Theorem 2 quantifies the rate of decrease using an off-the-shelf optimization algorithm Universal Stochastic Gradient Method (Alg. 3 of [SGD]), without any modifications or exploitation of the structure of the considered optimization problem.

[SGD] Universal Gradient Methods for Stochastic Convex Optimization, Anton Rodomanov, Ali Kavis, Yongtao Wu, Kimon Antonakopoulos, Volkan Cevher, link:https://openreview.net/pdf?id=Wnhp34K5jR

---

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 $O(d^2)$. Given that the numerical experiments involve cases where $d$ can reach $10^4$, the complexity could escalate to approximately $10^8$. This computational burden seems significant and may pose challenges for scalability in practice.

Q3: Could the authors elaborate further on the scalability of the proposed method, particularly given the projection method’s complexity?

Answer: We would like to clarify that the projection method was not a bottleneck of the algorithm. Observe that the model W itself is a matrix with $d^2$ elements; therefore, for large $d$, $d \approx 10^4$, all direct manipulations of W become computationally challenging. In our simulations, the limiting factor for scalability was not the computation complexity, but the memory requirements to store matrix W itself (see lines 318-319)

---

Q4: 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.

Answer: We would like to thank the reviewer for the suggested reference; we are currently adding the method to our experimental comparison. We want to discuss the method in detail.

Reference [A] investigates the NOTEARS framework for continuous optimization in Bayesian network structure learning, providing new theoretical insights into the KKT conditions for such formulations. They propose combining NOTEARS with a local search method informed by KKT conditions (KKTS), leading to improvements in SHD.

While [A] focuses on enhancing existing algorithms with a post-processing step, our work takes an orthogonal approach, reformulating DAG learning as a stochastic optimization problem. We develop an entirely new framework that seamlessly incorporates stochastic gradient methods, allowing efficient processing of large graphs (up to 10,000 nodes), exceeding the computational limits of NOTEARS and [A].

We would like to thank the reviewer for the constructive review. We would like to answer all raised concerns.

The presentation is poor, many concepts are used without introduction, e.g. transitive closure, full DAG, R, in convergence statement.

Answer: We would like to thank the reviewer for pointing those out, we are correcting them immediately.

---

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.

Answer: We believe there may be a misunderstanding. The minimization problem (11) is defined over directed acyclic graphs (DAGs). The identity matrix does not qualify as a DAG. Furthermore, if W represents a valid DAG, all elements on its diagonal are necessarily zero (this is enforced during the whole optimization process) and therefore, the solution of (11) is not close to the identity matrix.

---

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:

Answer: We thank the reviewer for pointing out these relevant algorithms, which have helped us better situate our work within the broader literature.

Reference [1] introduces a procedure for estimating the topological ordering of the underlying graph under the assumption of equal error variances. This approach naturally extends to high-dimensional settings where $p>n$, providing an alternative to greedy search methods with competitive accuracy and computational efficiency. However, unlike our method, it relies on controlling the maximum in-degree of the graph, which can become computationally intensive for graphs with high in-degree.

Our approach, on the other hand, reformulates the DAG learning problem as a stochastic optimization task, leveraging gradient-based methods to scale efficiently to larger graph sizes. While variance-ordering methods usually depends on various error variance assumptions, our framework avoids any such requirements. Additionally, variance-ordering methods can be enhanced with greedy searches to obtain improved performance. Yet, our experimental results show that our method $\psi$DAG achieves state-of-the-art scalability and accuracy without the need for such hybrid extensions. This positions our method as a robust alternative to variance-ordering and greedy approaches, particularly for large-scale DAG learning tasks.

Reference [2] presents a theoretical analysis of Gaussian DAG models, deriving minimax optimal bounds for structure recovery under the equal varience assumption. Their work emphasizes sample efficiency and offers theoretical guarantees for structure learning in Gaussian settings. While this contribution is insightful, it primarily focuses on theoretical aspects and does not address the computational challenges posed by large-scale DAGs. In contrast, our method introduces a scalable stochastic optimization framework that can efficiently handle graphs with up to 10,000 nodes, extending its applicability to a wider range of practical problems.