Neural structure learning with stochastic differential equations

"If one uses existing methods for uncovering SDEs ... is there a way to estimate the graph structure from the learned function (e.g. by computing the partial derivatives between the function and the input)?"

Attempting to extract a graph from the learned SDE in this way, without the explicit concept of a SDE graph, is challenging for a few reasons.

Firstly, to use partial derivatives, one would have to evaluate the partial derivative of the drift and diffusion networks at every input point in the input domain, which is not practical. Secondly, the learned drift and diffusion functions may have different graphs, and it is unclear how we should combine these. Thirdly, there are no theoretical results to justify this approach (prior to our paper's theory). For these reasons, prior latent SDE work (e.g. Li et al.) does not admit an easy way to extract structure.

In order to construct an empirical baseline for comparison, one way in which the graph can be extracted from the latent SDE is if the each output dimension of the drift and diffusion functions is implemented by a different neural network. Then we can simply extract the weights of the first layer of the networks, following the approach of (Bellot et al. 2022) for ODEs. As shown in the table, this approach (which we call LSDE) performed reasonably compared to some of the other baselines, but far worse than SCOTCH. We have added this to the appendix and provide more details there.

"Do the authors have an idea of how the performance scales to higher dimensional observations?"

In our experiments, we do test on the 100-dimensional DREAM3 dataset. This is a very challenging dataset due to its high dimensionality with small number of observations (below 1000 observations). However, our model significantly outperforms the baselines, and to the best of our knowledge, SCOTCH achieves the state-of-the-art performance. In addition, as far as we are aware, this is the highest dimensionality dataset tackled among the related works, both regarding structure learning and SDE learning. Hasan et al. does consider modelling MNIST observed data, but the latent dimensionality (upon which the SDE operates) is very small, consisting of only digit positions.

If we have addressed your concerns to your satisfaction, we would be grateful if you would consider increasing your score. We are of course happy to clarify further or answer any additional questions you may have.

We thank the reviewer for their hard work and helpful feedback.

"What is the main difference with other latent SDE models that have been proposed? ..."

The main conceptual difference in comparison to prior work on latent SDEs is that we tackle the problem of structure learning, or in other words, the qualitative dependencies between variables in the SDE. In this respect, our theoretical results are crucial for showing when this can be achieved (regardless of the specifics of the SDE learning method). This complements existing results on structural identifiability in the discrete-time structure learning literature, as well as the ODE structural identifiability of Bellot et al. (2022).

In more detail, Theorem 4.1 is a new result that shows (under assumptions 1, 2) the identifiability of the graphical structure when observing data generated directly from the corresponding SDEs. Then, Theorem 4.2 leverages the proof of Hasan et al. (as stated in our proof), and shows the identifiability of the graph (and SDE) even when the SDE is over latent variables, if the time range is bounded with infinite data or the observation interval is fixed. On the other hand, Hasan et al. focus on the identifiability of the model with fixed time intervals and cannot show the identifiability of the graph since they lack Theorem 4.1.

Finally, Theorem 4.3 shows the consistency of our particular method, SCOTCH, in recovering the correct graph in the limit of infinite data, justifying its use for structure learning. This consistency in continuous time has not been shown in Hasan et al.

We should also emphasize that although Hasan et al. presented its framework under SDEs, its nature is a discrete model that requires observation at fixed time intervals. This is because they assume an Euler discretization of the SDE so that the corresponding model is a discrete state-space model. On the other hand, our model is much more general and inherently continuous in time, not restricted to fixed-interval Euler discretization.

"Aside from the graph interpretation in the drift/diffusion, the work is similar to a lot of existing works, so I wonder a bit how it is practically different from some of the others"

In addition to the conceptual and theoretical results above, we now aim to clarify the practical/methodological distinctions. In short, SCOTCH share some similarities to e.g. Li et al. in terms of the SDE formulation, but with some important differences:

1. Bayesian Structure Learning: In SCOTCH, we employ Bayesian structure learning and explicitly model a variational posterior distribution over graphs. Practically, this means that we sample from this distribution during training, effectively "masking" the learned neural network drift and diffusion functions in different ways; this provides a signal to learn the structure. On the other hand, Li et al. treated the model parameters as a point estimate and do not have the concept of the graph. Even if one tries to extract the "graph" from the weights, it is a point-estimate graph rather than a distribution over graphs.
2. Objective Function and Network Design: The incorporation of a learnable graph distribution leads to a distinct (though related) ELBO (Eq. 15) compared to Li et al. Correspondingly the network designs are also different; for example, the prior drift and diffusion functions follow the design of Eq. 8, while the posterior drift function $h_{\psi}$ takes as input the graph as explained in Section 3.1.
3. Latent Dimensionality: A more straightforward difference is that SCOTCH always has the same latent and observed dimensionality, unlike other works. This ensures that the graphical structure of the latent system can be identified from observed data (under assumptions).

Please acknowledge author replies, and in particular, indicate whether your concerns have been addressed or require further discussion.

We thank the reviewer for their valuable feedback, and recognition of our contributions.

"The paper would benefit from a more in-depth analysis and discussion regarding the sensitivity of the proposed method to the SDE solvers."

The reviewer is correct that there are a number of factors related to the accuracy of the SDE solution from the SDE solver. Firstly, there is the matter of discretization step size of the solver; a smaller step size generally leads to more accurate SDE solution, but at the cost of time and space complexity. The computational cost (with default Euler discretization) will scale inversely with the step size. We conducted an additional experiment for the Ecoli1 dataset (in DREAM3) with a smaller step size and added it to the appendix of the revised paper and below.

We can see that, at least for the Ecoli1 dataset, a finer discretization does not necesarily improves the performance, while increasing computational cost. Therefore, for Dream3 dataset, we use $\Delta t=0.05$ as default.

Secondly, we chose to use a pathwise gradient estimator rather than the adjoint method, as we found this was more efficient time-wise and we did not run into memory limitations. Although theoretically, they should give the same performance, in practice, the pathwise gradient estimator may have an advantage that computing its gradient does not require solving another SDE, which is subject to the accuracy of the SDE solver. Finally, it is also possible to use higher-order numerical solvers such as the Milstein method, in place of Euler-Maruyama; however we have not thoroughly explored this in the current work. We have added discussion on SDE solvers to the experimental section and appendix.

"...a thorough discussion of the reasons behind the choice of priors for structure sampling... whether achieving sparsity in G leads to improved estimation performance, or if it primarily results in more interpretable results."

The reason we need the sparsity prior is because in theory, a fully connected graph will always have the greatest capacity to fit data. The purpose of incorporating a sparsity term is to avoid false positive edges, i.e. spurious edges that may slightly improve the ELBO but are not present in the underlying system. In real-world systems (e.g. DREAM3), we found this helped to more accurately predict the ground-truth graph.

L1 sparsity priors (Lasso) are the most common type of prior used in structure learning, so we adopt this for our method. That said, there are other possibilities. For example, another prior sometimes used in DAG structure learning (Eggeling et al. 2019) assigns probability $p(G) \propto \prod_{d=1}^{D} \frac{1}{{D-1 \choose |pa(d)|}}$, where $pa(d)$ is the set of parents of the $d^{\textnormal{th}}$ variable. In theory we can use any graph prior that assigns positive probability to every graph. The L1 prior turns out to be particularly convenient (in combination with our variational graph posterior) as the KL-divergence in Eq. (14) is tractable.

(Eggeling et al. 2019) Ralf Eggeling, Jussi Viinikka, Aleksis Vuoksenmaa, Mikko Koivisto. On Structure Priors for Learning Bayesian Networks. AISTATS 2019.

We thank the reviewer for their positive assessment of our paper.

"I am more curious about the quality of produced graphs from SCOTCH compared to other models, and how we understand their structures. I wonder if there is any way to visualize the graphs."

In our experiments, we do quantitatively compare the quality of the graphs to the ground truth from SCOTCH vs other models using the AUROC metric, where a higher value corresponds to a greater degree of conformity to the true graph. For visualizing the graph, one can simply sample the adjacency matrix from SCOTCH and use python packages like networkx to automatically plot the corresponding graphs.

We hope that we have understood your concerns regarding this. Please let us know if not, or if you have any other feedback that could improve the paper - we would be very grateful.

We thank the reviewer for their encouraging assessment of the strengths of our paper and for their comments, which we address below.

"I think the assumptions (1, 2) should be more clearly stated"

Agreed; we have updated the manuscript to place the assumptions in the main text.

"homogenous drift and diffusion processes assumption may not be able to model some types of dynamics"

This is definitely an interesting point, which we now discuss in the conclusion of the revised manuscript. In a nutshell, with homogeneous drift and diffusion functions, SCOTCH cannot model non-stationary dynamics. For example, systems with obvious changepoints (e.g. climate science models) will not have time-homogenous drift/diffusion functions. Unfortunately, our theoretical results no longer hold if time-dependent graphs are allowed.

"...the explanation of the improved performance of RHINO on the Netsim data"

With regards to the performance of SCOTCH vs Rhino, we copy (part of) Table 3 of Gong et al. (2022) above, showing that Rhino can achieve almost perfect AUROC on the Netsim dataset when allowed to use instantaneous effects, even when restricted to Gaussian noise (Rhino+g), it can achieve 0.974. Both Rhino+NoInst (no instantaneous effects), and the Granger causality method cMLP (which cannot model instantaneous effects) can only achieve 0.93 AUROC, indicating the importance of instataneous effects on this dataset. As explained in Section 3.2, SCOTCH also does not incorporate instantaneous effects, so the performance is worse compared with Rhino. We also highlight the corresponding sections in the revised paper.

As for SCOTCH vs Rhino+NoInst with both complete and missing data, we can see that the performance is broadly similar and superior to other baselines. Although the performance of Rhino is slightly better than SCOTCH with missing data, we are not entirely sure why SCOTCH does not outperform Rhino+NoInst. We suspect it may be due to pecularities of the NetSim dataset, or the fact that Rhino+NoInst uses a more complex noise distribution transformed by a normalizing flow, whereas SCOTCH uses a Gaussian noise modulated by a diffusion matrix.

"What is the effect of the sparsity prior on the structure? Can a different type of sparsity be used?"

The reason we need the sparsity prior is because in theory, a fully connected graph will always have the greatest capacity to fit data. The purpose of incorporating a sparsity term is to avoid false positive edges, i.e. spurious edges that may slightly improve likelihood but are not present in the underlying system. In real-world systems (e.g. DREAM3), we found this helped to more accurately predict the ground-truth graph.

L1-based sparsity priors (Lasso, group Lasso, adaptive group Lasso) are the most common type of prior used in structure learning, so we adopt this for our method. That said, there are other possibilities. For example, another prior sometimes used in DAG structure learning (Eggeling et al. 2019) assigns probability $p(G) \propto \prod_{d=1}^{D} \frac{1}{{D-1 \choose |pa(d)|}}$, where $pa(d)$ is the set of parents of the $d^{\textnormal{th}}$ variable. In theory we can use any graph prior that assigns positive probability to every graph. The L1 prior turns out to be particularly convenient (in combination with our variational graph posterior) as the KL-divergence in Eq. (14) is tractable.

(Eggeling et al. 2019) Ralf Eggeling, Jussi Viinikka, Aleksis Vuoksenmaa, Mikko Koivisto. On Structure Priors for Learning Bayesian Networks. AISTATS 2019.

As the discussion period comes to a close, we would like to once again thank all reviewers for their feedback, which has been invaluable in enabling us to improve the paper.

Our latest manuscript revision is to conform to the paper length constraints, and includes an additional edit pass on the Appendix as mentioned by reviewer gu2c.