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

We sincerely thank the Reviewer for the thorough and thoughtful evaluation of our work. We appreciate your recognition of our paper’s clarity, comprehensive experiments, and the novelty and significance of our approach to scalable DAG learning. Your insights are invaluable, and we believe that addressing them will significantly strengthen our work.

W1 [Significance] The role and justification of vertex ordering require further clarification. Lines 172--193 suggest that using the true vertex ordering for the true DAG (solution) can improve the objective value, but in real-world scenarios, the true ordering is generally unavailable during optimization.

Our method does not assume the true vertex ordering; it explores candidate orderings to guide optimization.

The motivation for using vertex orderings stems from the fact that the space of possible orderings (d!) is exponentially smaller than the space of all DAGs ($2^{d^2–d}$). This makes projection into the ordering space significantly more tractable than directly optimizing over the full DAG space. We would like to clarify that our method does not assume access to the true vertex ordering, which, as the Reviewer correctly points out, is not available in practice. Instead, our algorithm iteratively explores candidate orderings, and for each one, solves a convex optimization problem to find the optimal DAG consistent with that ordering. This procedure provides meaningful upper bounds and helps guide the search toward better solutions within a highly non-convex landscape.

We will update the manuscript to clarify that the true ordering is not used or assumed during optimization. We hope this addresses your concerns, and we would be happy to elaborate further if needed.

W2 [Significance] Extending $\Psi$DAG to nonlinear SEMs would further increase the method’s utility, as many recent approaches support both linear and nonlinear formulations.

Extending $\Psi$DAG to nonlinear SEMs is indeed a valuable direction and is entirely feasible within our framework. However, in this work, we deliberately focused on the linear setting to highlight the core methodological contributions and scalability advantages of our approach. We have also acknowledged this limitation in the paper.

Q1. Does a vertex ordering always exist for $W_{k}^{(1/3)}$?

Yes, a vertex ordering always exists due to the construction of the projection step. While the ordering may not be unique, the projection is explicitly designed to yield an acyclic structure, which guarantees at least one valid topological ordering.

Q2. As I understand it, $W_{k}^{(1/3)}$, $W_{k}^{(2/3)}$ and $W_{k+1}$ become dense matrices. Is this correct? If so, how is sparsity induced in $W_{K}$? I could not find a sparsity-inducing regularization term similar to the one used in NOTEARS.

Sparsity is not enforced, but can be easily added.

Both $W_{k}^{(2/3)}$ and $W_{k+1}$ are DAGs by construction. While our current implementation does not include a sparsity-inducing regularizer, such as an $l_1$-penalty, it can easily be added to encourage sparsity if desired. Our framework is fully compatible with sparsity regularization.

Q3. Can the projection method described in Algorithm 3 also be applied to deterministic settings such as NOTEARS and DAGMA? I am particularly interested in a comparison of the computational time between regularization-based approaches (e.g., NOTEARS, DAGMA) and the projection-based method proposed in Algorithm 3 under a deterministic setting.

Our projection can be applied to deterministic settings, but $\Psi$-DAG remains more efficient due to stochastic updates.

Yes, the projection method described in Algorithm 3 can, in principle, be applied to deterministic settings such as NOTEARS by removing the acyclicity constraint h(A) and instead performing a projection onto the DAG space after each update. However, one of the key advantages of our approach lies in its computational efficiency: unlike NOTEARS and DAGMA, which rely on computing full gradients over the entire dataset at each step, our method uses Stochastic Gradient Descent (SGD)-type updates. This significantly reduces the computational burden, especially in high-dimensional or large-sample regimes. A direct empirical evaluation under a fully deterministic setting falls beyond the scope of the current work.

Q4. What is the role of weights L in Algorithm 3?

The weight matrix $L$ encodes the relative importance of edges for the projection.

Great question. The weights L originate from an optimization perspective, particularly inspired by coordinate descent methods, where some directions may have a larger impact on the loss than others. For instance, consider the optimization problem $f(x,y)=100(x−1)^2+0.01(y+1)^2$. In this case, an error of $\varepsilon=0.1$ in the x-coordinate ($x=x^*+ \varepsilon$) results in a significantly higher increase in loss than the same error in the y-coordinate. This reflects the fact that the x-direction is more important to the objective function.

Analogously, in our framework, the weight matrix L captures the relative importance of each edge during the projection step. Additional discussion on the benefits and role of L can be found in Appendix E.

We hope that we answered all of the questions posed in the review. We are happy to answer any additional questions in the discussion period.

We sincerely thank the Reviewer for their comments and for recognizing the novelty, significance, and robust performance of our method. We are particularly grateful for the acknowledgment of its stability and strong gains over prior continuous optimization approaches. We will incorporate high-dimensional PC and LiNGAM baselines in the camera-ready version as suggested. While our approach is extendable to nonlinear settings, we focused on the linear case to clearly highlight the core contributions, a point we will make more explicit in the revised manuscript.

1. Is there a correctness guarantee that can be stated formally?

Exact DAG recovery is NP-hard, but we provide local convergence guarantees.

We agree that theoretical guarantees are important. Recovering the true underlying DAG is known to be NP-hard in general. As such, no algorithm, including ours, can guarantee exact recovery of the true DAG in polynomial time without performing an exponential number of steps. Moreover, the DAG learning problem is inherently non-convex and may contain exponentially many local minima, making it theoretically hard to guarantee convergence to a global minimum in reasonable time.

From a theoretical perspective, most existing continuous optimization-based approaches (NOTEARS, GOLEM, NoCurl) can only guarantee convergence to stationary points due to this non-convexity. In contrast, our method benefits from a structured reformulation, which allows us to guarantee convergence to a local minimum, as stated in Theorem 8.

2. Have the authors tried running their approach in the nonlinear setting? If so, how do these results compare?

We focused on establishing a novel and scalable framework for differentiable DAG learning, leveraging a stochastic formulation to improve computational efficiency and scalability. To clearly isolate and evaluate the core contributions, our experiments are currently limited to linear SEMs. However, extending the proposed approach to nonlinear SEMs, such as those considered in Zheng et al. (2020) [1], is a natural direction we left for future work.

3. Did the authors consider any additional experiments to specifically measure the stability of the optimization to other approaches?

We evaluated stability across multiple runs and under varying noise; our method is consistently robust.

Great question. All results presented in the main paper are averaged over multiple runs with different random initializations. We report the mean values and indicate the standard error using shaded regions. We note that some baseline methods failed to converge under certain random seeds, as described in lines 273–279, highlighting their sensitivity to initialization.

In addition, we explicitly evaluated robustness under varying noise ratios, shown in Figure 4. While other continuous optimization baselines (GOLEM and NOCURL) exhibit increasing SHD (lower is better) as the noise ratio increases, our method maintains a consistently low SHD across all noise levels. This demonstrates the robustness and stability of our approach under more challenging conditions. Notably, DAGMA failed to converge for noise ratios above 15, further emphasizing the advantage of our method in terms of both convergence and stability.

[1] Zheng, X., Dan, C., Aragam, B., Ravikumar, P., and Xing, E. Learning sparse nonparametric dags. In International Conference on Artificial Intelligence and Statistics, pp. 3414–3425. Pmlr, 2020.487

Dear Reviewer M4a3,

As the discussion period is approaching its end, we wanted to kindly follow up to ask whether our responses have addressed all of your concerns. We greatly appreciate your engagement and would be happy to clarify any remaining questions if needed.

We thank the Reviewer for recognizing the strengths of our work, including its scalability, accuracy, and mathematical soundness. Your insights are invaluable, and we believe that addressing them will significantly strengthen our work.

(W1+Q3): Applicability beyond linear SEMs

Yes, $\Psi$DAG supports conditional SEMs with deterministic or fixed inputs. Below we show empirical results where $\Psi$DAG outperforms baselines in such settings.

SHD ($\downarrow$)

F1 Score ($\uparrow$)

The proposed method is fully compatible with conditional SEMs where some variables are deterministic or fixed across samples. From an optimization perspective, the algorithm operates on observed data, regardless of whether the input values are drawn stochastically or held constant. However, it is important that the overall model includes a stochastic component in the generation of each node (i.e., additive noise) to maintain identifiability. If some structural weights were strictly deterministic and the noise component entirely absent, identifiability would be compromised, not due to our method specifically, but as a general limitation in the causal discovery literature. In this regard, our method adopts the same assumptions as widely used methods (e.g., NOTEARS, GOLEM, NOCURL), and remains applicable to conditional SEM settings where inputs are partly deterministic.

(W2+Q4): Limited sample performance

$\Psi$-DAG performs well even in low-sample regimes, despite leveraging stochastic optimization.

From an optimization perspective, our method handles sum-type objectives such as the mean squared error $F(w) = \frac{1}{n}\sum_{i=1}^{n} f(w, x_i)$ by treating them as stochastic expectation problems of the form $E_\xi[f(w,x_{\xi})]$ where $\xi \in \lbrace 1,\ldots,n \rbrace$. This allows the use of stochastic optimization techniques even when the number of samples is limited.

Although $\Psi$-DAG leverages stochastic gradient updates, which are typically associated with large datasets, we observe that it remains effective in low-sample regimes. In particular, our experiments on a real-world dataset (lines 283-300), where the sample size is small, show that $\Psi$-DAG outperforms baselines, indicating strong robustness to limited data.

W3: Non-i.i.d

$\Psi$-DAG is compatible with non-i.i.d. noise and shows strong empirical robustness in such settings. We include below preliminary results under non-i.i.d. noise (up to $d=100$), showing $\Psi$-DAG remains competitive.

Our primary focus in this work was on addressing the DAG constraint, which we consider to be the most computationally challenging component of the problem. The noise distribution primarily influences the form of the loss function, and our framework is agnostic to whether the noise is i.i.d. or non-i.i.d. across nodes.

We conducted additional experiments using synthetic data with heterogeneous noise. We first sample a random DAG from the Erdős–Rényi model and randomly select two variables whose structural mechanisms vary across subgroups. For each node, data is generated following a topological order: if a node has parents, its value is computed as a weighted sum of its parents passed through a randomly selected nonlinear function ($\sin$, $\tanh$, square, or linear), plus noise. The noise is sampled from either a Gaussian or uniform distribution with varying scales, introducing non-i.id noise across nodes and subgroups.

We report SHD (lower is better) and F1 score (higher is better), below, and $\Psi$-DAG consistently outperforms all baselines. We will update the paper to include these new results and clarify the generality of our method to non-i.i.d. noise settings.

SHD ($\downarrow$)

F1 Score ($\uparrow$)

Q1: Linear SEM applicability

Linear SEMs with additive noise are broadly applicable in real-world domains and remain essential for identifiability.

• In genomics, linear SEMs are used to infer gene regulatory networks from gene expression data, where each gene’s expression is influenced by others plus latent biological noise [1]

• In biology, SEMs have been used by Sachs et al. [2] to model causal protein-signaling networks from single-cell data, capturing interactions between molecules under measurement noise.

• In economics and social science, SEMs model interactions among variables like education, income, and policy indicators [3,4]

These models rely on additive noise to capture unobserved variability and ensure identifiability, as discussed in works such as [5,6].

Identifiability requirement: The presence of stochasticity (i.e., a noise term) at every node is a key assumption for ensuring identifiability in many causal discovery methods, including NOTEARS and GOLEM.

Q2: Nosie distribution

Our method is distribution-agnostic and works with both continuous and discrete, i.i.d. and non-i.i.d. noise.

$\Psi$-DAG does not require any explicit assumptions about the form of the noise distribution. In particular, it does not assume that the noise is i.i.d. across nodes, nor does it require the noise distribution to be known. The formulation is flexible and can accommodate both i.i.d. and non-i.i.d. noise structures. As shown above, it works under non-i.i.d. settings as well, although for benchmarking purposes we focused on standard assumptions.

Furthermore, $\Psi$-DAG does not impose restrictions on the type of noise; it can handle continuous and discrete distributions. While our experiments focused on standard continuous noise settings (e.g., Gaussian, Gumbel), the framework itself is general and applicable to settings with discrete noise as well.

Q5: Limitations

$\Psi$-DAG does not guarantee global optimality, but can serve as a fast and effective subroutine in exact solvers.

The primary goal of $\Psi$-DAG is to efficiently find high-quality feasible solutions in a short amount of time, leveraging stochastic optimization and projection techniques. However, the method does not guarantee globally optimal solutions due to the combinatorial nature of the DAG learning problem.

For settings where exact recovery of the true DAG is essential, $\Psi$-DAG can be effectively combined with a branch-and-bound strategy over the vertex ordering space. While the space of DAGs is of size $2^{d^2-d}$, the space of vertex orderings is significantly smaller at $d!$. In such frameworks, $\Psi$-DAG can be used to provide both upper and lower bounds for each ordering, which are crucial for efficient pruning during the search. This makes it particularly suitable as a bounding subroutine within exact combinatorial solvers.

We hope that we answered all of the questions posed in the review. We are happy to answer any additional questions in the discussion period.

[1] Zhang, B., Gaiteri, C., Bodea, L.-G., Wang, Z., McElwee, J., Podtelezhnikov, A., Zhang, C., Xie, T.,Tran, L., Dobrin, R., et al. Integrated systems approach identifies genetic nodes and networks in late-onset Alzheimer’s disease.

[2] Sachs, K., Perez, O., Pe’er, D., Lauffenburger, D., and Nolan, G. Causal protein-signaling networks derived from multiparameter single-cell data.

[3] Bollen, K. A. Structural Equations with Latent Variables. Wiley.

[4] Morgan, S. and Winship, C. Counterfactuals and causal inference. Cambridge University Press, 2015.

[5] Peters et al., “Causal discovery with continuous additive noise models,”

[6] Kalisch & Bühlmann, “Estimating High-Dimensional Directed Acyclic Graphs with the PC-Algorithm,”

We sincerely thank the reviewer for their valuable feedback. Your insights help us refine our work and improve its clarity and impact. Below, we address the key points raised:

(W1 + Q3): Comparison with NOCURL

Our method avoids acyclicity constraints and uses a fundamentally different projection and optimization strategy for better scalability.

NOCURL first solves the unconstrained optimization problem $A^{\text{pre}} = \arg\min F(A, X) + \lambda h(A)$, with the NOTEARS acyclicity constraint $h(A) = \mathrm{tr}[\exp(A \circ A)] - d$ (costly matrix exponential), then applies Hodge decomposition to project to a DAG. It requires computing gradients of $h(A)$ twice with different $\lambda$s, which increases the cost (Section 4: Algorithm: DAG with No Curl [6]).

Our method avoids this entirely:

• We never compute $h(A)$ or its gradients at any stage.

• We perform a lightweight projection step (Algorithm 3) to directly obtain a topological ordering.

• We solve a stochastic convex optimization (Problem 10) using this fixed ordering, which avoids expensive matrix operations and allows scalable SGD-based updates.

A detailed comparison with NOCURL is provided in Appendix A (lines 543–550) due to space limitations in the main text. We will move it into the main body in the camera-ready version for better visibility.

(W2 + Q4): Normalization of Synthetic Data

We show that our method works under realistic, non-equal noise variance conditions.

We thank the reviewer for raising this important point regarding standardization, which is sometimes referred to as a type of normalization of the simulated data. As highlighted by [1], non-standardized benchmarks can artificially inflate performance due to variance patterns (varsortability) that some algorithms exploit. Building on this, [2] demonstrated that continuous structure learning methods often fail under standardization because it inflates the effective noise ratio, breaking assumptions of equal noise variance. In our experiments (see lines 229–233), we deliberately used a non-equal noise ratio variance setting to provide a more realistic and challenging evaluation that avoids reliance on marginal variance cues. This setting aligns with the more comprehensive evaluation approach suggested by recent works and reflects real-world noise heterogeneity.

Our model demonstrated consistent performance across a wide range of noise ratios, as shown in Figure 4, indicating that it is robust to the distributional changes introduced by standardization. This suggests that our method would remain effective even when applied to standardized data, as it does not rely on marginal variance patterns or uniform noise levels. We apologize if this was not sufficiently clear in the current manuscript and will include a more detailed explanation to explicitly clarify our handling of standardization and noise variance.

W3: I might have missed it, but it is not clear to me how the "stochastic reformulation" build in Section 3 differs from the usual MSE problem.

The objective is similar, but our optimization is stochastic, enabling SGD-based scalability.

The stochastic reformulation introduced in Section 3 shares a similar objective with the standard MSE formulation, as both aim to minimize reconstruction error. However, the key difference lies in how the optimization is conducted. The stochastic formulation enables the use of Stochastic Gradient Descent (SGD)-type updates, which are computationally efficient and scalable to large graphs. In contrast, most benchmark methods, such as NOTEARS, GOLEM, and NOCURL, rely on full-batch gradient steps, which are more computationally demanding and often suffer from poorer scalability.

From a theoretical standpoint, this distinction aligns with the difference between Stochastic Approximation (SA) and Sample Average Approximation (SAA). SA methods, such as the one used in our formulation, optimize toward the true, unbiased solution by treating the data as stochastic samples. In contrast, SAA methods optimize over fixed empirical losses, which can lead to biased estimates and, consequently, the inability to consistently recover the true underlying DAG structure [3, 4]. SA-based strategies are widely adopted in other machine learning domains (e.g., deep learning), where they have been shown to outperform deterministic alternatives [5]. In our implementation (lines 199–200), we employed an advanced variant of Stochastic Gradient Descent, the Universal Stochastic Gradient Method proposed by Rodomanov et al. (2024), which adapts learning rates automatically.

Q1: Why are results of NOTEARS on the synthetic dataset missing in the main body? When reported in the Appendix (Figure 7), NOTEARS seems to perform better than $\Psi$-DAG.

NOTEARS is significantly slower and less scalable. We focused on DAGMA since it improves on NOTEARS in both accuracy and scalability.

We would like to clarify the purpose of Figure 7 in the appendix. This figure illustrates that, under the equal-variance setting, the conclusions of [1] and [2] hold, specifically, that continuous optimization methods such as NOTEARS and GOLEM achieve good structural recovery. However, as shown in Figure 7b, the runtime of NOTEARS and GOLEM increases substantially with the number of nodes, limiting their scalability. While NOTEARS achieves SHD comparable to our method in this simplified setting, our approach does so with significantly lower computational cost, as shown in Figure 7b. Given NOTEARS' scalability limitations and the fact that DAGMA improves upon NOTEARS both in accuracy and efficiency (see lines 213–214), we focused on DAGMA in the main experiments. Nonetheless, we emphasize that DAGMA becomes slower than our method as the number of nodes increases, as demonstrated in Figures 5 and 9. We will include NOTEARS in the relevant figures, despite its limited scalability, and will update the appendix to clarify the intent of Figure 7 and to highlight these distinctions more explicitly.

Q2: How is the acyclic optimization for the second step of your algorithm performed? Which acyclicity constraint are you employing?

Acyclicity is guaranteed via a projection step that yields a topological order, no gradients or acyclicity constraints needed.

We would like to clarify the acyclic optimization procedure in our method. Optimization is performed only in Steps 1 and 3 of our algorithm. In Step 2, as shown in Algorithm 3, we apply a projection operator that transforms the intermediate matrix into the acyclic graph and returns a corresponding topological ordering. This projection step enforces acyclicity by construction and does not involve gradient-based optimization. Once the topological order is obtained, Step 3 solves an optimization problem with a fixed vertex ordering, as detailed in Section 4.1 (Optimization for Fixed Vertex Ordering). Because the ordering is fixed, the optimization is constrained to the space of acyclic graphs, ensuring that every iterate during the optimization remains a DAG. More specifically, the optimization problem is reformulated as Problem (10), a stochastic convex optimization problem that can be solved using any stochastic gradient method. In our implementation (lines 199–200), we use the Universal Stochastic Gradient Method [Rodomanov et al., 2024] to ensure stable and efficient convergence.

We hope that we answered all of the questions posed in the review. We are happy to answer any additional questions in the discussion period.

[1] Reisach, Alexander, Christof Seiler, and Sebastian Weichwald. "Beware of the simulated dag! causal discovery benchmarks may be easy to game." Advances in Neural Information Processing Systems 34 (2021): 27772-27784.

[2] Ng, Ignavier, Biwei Huang, and Kun Zhang. "Structure learning with continuous optimization: A sober look and beyond." Causal Learning and Reasoning. PMLR, 2024.

[3] Nemirovski, A., Juditsky, A., Lan, G. and Shapiro, A., 2009. Robust stochastic approximation approach to stochastic programming. SIAM Journal on optimization, 19(4), pp.1574-1609.

[4] Lan, G., 2012. An optimal method for stochastic composite optimization. Mathematical Programming, 133(1), pp.365-397.

[5] Bottou, L., Curtis, F.E. and Nocedal, J., 2018. Optimization methods for large-scale machine learning. SIAM review, 60(2), pp.223-311.

[6] Yu, Yue, et al. "DAGs with no curl: An efficient DAG structure learning approach." International conference on machine learning. Pmlr, 2021.

I thank the authors for their rebuttal.

I am satisfied with the highlighted differences between their method and NOCURL, which justifies the novelty of $\psi$-DAG. Entirely avoiding the use of acyclicity constraints is valuable, and it significantly improves scalability of the method. To better frame their work in the context of unconstrained acyclic optimization, I would suggest the authors to also comment on papers sharing the same motivation, which are currently missing from the related works, e.g., [1–3, 5]. Furthermore, I agree that the trade-off between performance and scalability is significant when dealing with NOTEARS, and I encourage your decision of highlighting this matter in the main body.

I also thank the authors for their SA vs SAA discussion. If I had to be overly naïve, for most practitioners this would reduce to the choice of an optimization scheme to tackle the MSE problem. While lacking your theoretical emphasis, other works already explored the use of mini-batch optimization in causal structure learning [1, 3–6]. Which mini-batch size are you employing? Would your methodology still hold if solving MSE in a full-batch scheme? What would be the downsides?

Finally, I still have concerns on the empirical evaluation of the method. While I tend to agree that Ng et al. (2024) "suggests that our method would remain effective even when applied to standardized data", it would help to empirically validate this claim even in a small setting.

Overall, I thank the authors for better framing their work, of which I can better appreciate the novelty once clarified that the method is entirely unconstrained. In this case, it would help to better position the paper in the context of other unconstrained approaches. Together with my latest remaining concerns on the empirical evaluation, I am willing to revise my score accordingly to borderline accept.

1. Annadani, Yashas, et al. "Bayesdag: Gradient-based posterior inference for causal discovery." Advances in Neural Information Processing Systems 36 (2023): 1738-1763.
2. Charpentier, Bertrand, Simon Kibler, and Stephan Günnemann. "Differentiable DAG Sampling." International Conference on Learning Representations.
3. Duong, Bao, et al. "Reinforcement Learning for Causal Discovery without Acyclicity Constraints." Transactions on Machine Learning Research.
4. Lopez, Romain, et al. "Large-scale differentiable causal discovery of factor graphs." Advances in Neural Information Processing Systems 35 (2022): 19290-19303.
5. Massidda, Riccardo, et al. "Constraint-Free Structure Learning with Smooth Acyclic Orientations." The Twelfth International Conference on Learning Representations.
6. Nazaret, Achille, et al. "Stable differentiable causal discovery." Proceedings of the 41st International Conference on Machine Learning. 2024.