SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lazar Atanackovic $^{1,2*}$
Alexander Tong $^{3,4*}$
Bo Wang $^{1,2,5}$
Leo J. Lee<sup>1,2</sup>
Joshua Bengio $^{3,4,6}$  Jason Hartford $^{3,4,7}$
<sup>1</sup>University of Toronto, <sup>2</sup>Vector Institute
<sup>3</sup>Mila - Quebec AI Institute, <sup>4</sup>Université de Montréal
$^{5}$ University Health Network,  $^{6}$ CIFAR Fellow,  $^{7}$ Valence Labs
One of the grand challenges of cell biology is inferring the gene regulatory network (GRN) which describes interactions between genes and their products that control gene expression and cellular function. We can treat this as a causal discovery problem but with two non-standard challenges: (1) regulatory networks are inherently cyclic so we should not model a GRN as a directed acyclic graph (DAG), and (2) observations have significant measurement noise, so for typical sample sizes there will always be a large equivalence class of graphs that are likely given the data, and we want methods that capture this uncertainty. Existing methods either focus on challenge (1), identifying cyclic structure from dynamics, or on challenge (2) learning complex Bayesian posteriors over DAGs, but not both. In this paper we leverage the fact that it is possible to estimate the "velocity" of gene expression with RNA velocity techniques to develop an approach that addresses both challenges. Because we have access to velocity information, we can treat the Bayesian structure learning problem as a problem of sparse identification of a dynamical system, capturing cyclic feedback loops through time. Since our objective is to model uncertainty over discrete structures, we leverage Generative Flow Networks (GFlowNets) to estimate the posterior distribution over the combinatorial space of possible sparse dependencies. Our results indicate that our method learns posteriors that better encapsulate the distributions of cyclic structures compared to counterpart state-of-the-art Bayesian structure learning approaches.
Inferring gene regulatory networks (GRNs) is a long standing problem in cell biology [25, 44]. If we knew the GRN, it would dramatically simplify the design of biological experiments and the development of drugs because it would serve as a map of which genes to perturb to manipulate protein and gene expression. GRNs concisely represent the complex system of directed interactions between genes and their regulatory products that govern cellular function through control of RNA (gene) expression and protein concentration. We can treat GRN inference as a causal discovery problem by treating the regulatory structure between genes (variables) as causal dependencies (edges) that we infer / rule out by using gene expression data. Structure learning methods aim to automate this task by inferring a set of directed acyclic graphs (DAGs) that are consistent with the conditional independencies that we can measure among the variables [14, 41, 42]. While there may be multiple
Correspondence to: (1.atanackovic@mail.utoronto.ca)
Figure 1: Architecture for Bayesian structure learning of dynamical systems. DynGFN consists of three main components: A GFlowNet modeling a posterior distribution over graphs  $Q_{\psi}(G|\mathcal{D})$ , a HyperNetwork modeling a posterior over parameters given a graph  $Q_{\phi}(\theta |G,D)$ , and the structural equation model scoring  $G$  and  $\theta$  according to how well they fit the data. Although the figure shows the case where  $Q_{\psi}(G|\mathcal{D})$  is modelled with a GFlowNet, this can be any arbitrary graph sampler that can sample discrete structures  $G \sim Q_{\psi}(G|\mathcal{D})$ .
DAGs in this set—the "Markov equivalence class"—when we are able to perturb the variables with enough experimental interventions, it is possible to uniquely identify a causal graph [18].
However, structure learning for inferring GRNs comes with two non-standard challenges: (1) gene regulation contains inherent cyclic feedback mechanisms, hence we should not model a GRN as a DAG, and (2) observations are limited and have significant measurement noise, hence there exists a large equivalence class of graphs that are likely given datasets with typical sample sizes. Existing methods either focus on (1) - identifying graphs with cyclic structure by leveraging dynamics [15, 13] or assuming the system is in equilibrium [36], or (2) - learning complex Bayesian posteriors over explanatory DAGs [12], but not both. In this work, we address both challenges concurrently in a fully differentiable end-to-end pipeline (see Figure 1).
To accomplish this, we treat structure learning as a problem of sparse identification of a dynamical system. From a dynamical systems perspective, one can model both causal structure between variables as well as their time-dependent system response with the drift function [37, 43]. We leverage the fact that we can estimate the rate of change of a gene's expression (velocity) with RNA velocity methods [9]. This data takes the form of dynamic tuple pairs  $(x, dx)$ , which we can use to pose the dynamical system learning problem as a regression task (see Figure 1). This significantly simplifies the learning objective as we can model system dynamics while also learning structure without the need for numerically intensive differential equation solvers. We view this as a step towards Bayesian structure learning from continuous dynamics – we term this Bayesian dynamic structure learning.
Our approach estimates the posterior over the sparse dependencies and parameters of the dynamical system. This is important in scientific applications because it is usually prohibitively expensive to acquire a enough data to uniquely identify the true graph underlying a data generating process. Capturing the complex distribution over candidate structure is critical for downstream scientific applications and is an essential step in active causal discovery [39, 52, 19]. This is especially important in settings where experiments are expensive, e.g. conducting genetic perturbations for inference of GRNs. Bayesian structure learning is a class of methods that try to model this distribution over structure from observed data. These methods model posteriors over admissible structures  $P(G|D)$  that explain the observations [32, 11, 4, 12, 31], but focus on modelling distributions over DAGs.
Our approach leverages Generative Flow Networks (GFlowNets) to model complex distributions over cyclic structures. GFlowNets [7, 8] parameterize the distribution over any discrete object (e.g. graphs) through a sequential policy, and as a result avoid needing to make restrictive parametric assumptions on the distribution. This makes them a useful tool in structure learning, particularly in cases where  $P(G|D)$  is discrete and complex [12]. In this work, we use GFlowNets to learn posteriors over the sparse structure in a dynamical system, and separately learn the posteriors over the parameters of the drift function via a HyperNetwork [16] that conditions on inferred structures. Our main contributions are summarized as follows:
- We develop a novel framework for Bayesian structure learning under the lens of dynamical system identification for modelling complex posteriors over cyclic graphs. We consider flexible parameterizations for the structural model such that we can capture both linear and non-linear dynamic relationships.
- We design a novel GFlowNet architecture, Dynamic GFlowNet (DynGFN), tailored for modelling posteriors over cyclic structures. We propose a per-node factorization within DynGFN that enables efficient search over the discrete space of cyclic graphs.
- We empirically evaluated DynGFN on synthetic dynamic data designed to induce highly multimodal posteriors over graphs.
- We showcase the use of DynGFN on a real biological system using single-cell RNA-velocity data for learning posteriors of GRNs.
There are many works on the problem of identifying causal structure  $G$  from either observational [e.g. 49, 54, 35] or interventional [e.g. 26, 29, 38] data, but the majority of existing methods return only the most likely DAG under the observed data. By returning only the most likely graph, these methods are often overconfident in their predictions. Bayesian approaches attempt to explicitly model a posterior distribution over DAGs given the data and model specification.
Bayesian Structure Learning: Recently, there has been significant interest in fully differentiable Bayesian methods for structure learning in the static case. DiBS [32], BCD-Nets [11], VCN [4], and DAG-GFlowNet [12] all attempt to learn a distribution over structural models from a fully observed system. The key difference is in how these methods parameterize the graph. DiBS is a particle variational inference method that uses two matrices  $U$  and  $V$  where  $G = \text{sigmoid}(U^T V)$  where the sigmoid is applied elementwise which is similar to graph autoencoders. BCD-Nets and DP-DAG use the Gumbel-Sinkhorn distribution to parameterize a permutation and direct parameterization of a lower triangular matrix. VCN uses an autoregressive LSTM to generate the graph as this gets rid of the standard uni-modal constraint of Gaussian distributed parameters. DAG-GFN has shown success for modelling  $P(G|D)$  [12]. However, it remains restrictive to assume the underlying structure of the observed system is a DAG as natural dynamical systems typically contain regulating feedback mechanisms. This can be particularly challenging for GFlowNets since including cycles in the underlying structure exponentially increases the discrete search space. We show that under certain assumptions we can in part alleviate this shortcoming for learning Bayesian posteriors over cyclic structures for dynamical systems. In small graphs, these methods can model the uncertainty over possible models (including over Markov equivalence classes).
Dynamic and Cyclic Structure Learning: There has been comparatively little work towards Bayesian structure learning from dynamics. Recent works in this direction based on NeuralODEs [10] propose a single explanatory structure [50, 6, 1, 2]. CD-NOD leverages heterogeneous non-stationary data for causal discovery when the underlying generative process changes over time [53, 21]. A similar approach uses non-stationary time-series data for causal discovery and forecasting [20]. DYNOTEARS is a score-based approach that uses time-series to learn structure [40]. However, these methods do not attempt to explicitly model a distribution over the explanatory structure. Other methods aim to learn cyclic dependencies in the underlying graph [24, 36, 28, 3]. For instance, [24] propose an iterative method that leverages interventional data to learn directed cyclic graphs. It is suggested that CD-NOD is also extendable to learn cyclic structure [21]. But these methods do not model a posterior over structure. In general, there remains a gap for the problem of Bayesian structure learning over cyclic graphs.
We include further discussion on related work for GRN inference from single-cell transcriptomic data and cell dynamics in Appendix C.1.
3.1 Bayesian Dynamic Structure Learning
Problem Setup: We consider a finite dataset,  $\mathcal{D}$ , of dynamic pairs  $(x,dx)\in \mathbb{R}^d\times \mathbb{R}^d$  where  $x$  represents the state of the system sampled from an underlying time-invariant stochastic dynamical system governed by a latent drift  $\frac{dx}{dt} = f(x,\epsilon)$  where  $\epsilon$  is a noise term that parameterizes the SDE;  $x$  and  $\epsilon$  are mutually independent. The latent drift has some fixed sparsity pattern i.e.  $\frac{\partial f_i}{\partial x_j}\neq 0$  for a small set of variables, which can be parameterized by a graph  $G$  such that  $g_{ij} = \mathbf{1}[\frac{\partial f_i}{\partial x_j}\neq 0]$ , where  $g_{ij}\in G,i = 1,\ldots ,d,j = 1,\ldots ,d$ . The variables  $x_{j}$  for which  $\frac{\partial f_i}{\partial x_j}\neq 0$  can be interpreted as the causal parents of  $x_{i}$ , denoted  $\mathrm{Pa}(x_i)$ . This lets us define an equivalent dynamic structural model [37, 43] of the form,
$$
\frac {d x _ {i} (t)}{d t} = f _ {i} \left(\operatorname {P a} \left(x _ {i}\right), \epsilon_ {i}\right), \tag {1}
$$
for  $i = 1, \dots, d$ . For the graph  $G$  to be identifiable, we assume that all relevant variables are observed, such that causal sufficiency is satisfied.
Our goal is to model our posterior over explanatory graphs  $Q(G|\mathcal{D})$  given the data. We aim to jointly learn distribution over parameters  $\theta$  that parameterize the latent drift  $f(x)$ ; these parameters will typically depend on the sparsity pattern such that  $p(\theta |G)\neq p(\theta)$ . We can factorize this generative model as follows,
$$
p (G, \theta , \mathcal {D}) = p (\mathcal {D} | G, \theta) p (\theta | G) p (G) \tag {2}
$$
This factorization forms the basis of our inference procedure. We learn a parameterized function  $f_{\theta}(x):\mathbb{R}^{d}\to \mathbb{R}^{d}$  that approximates the structural model defined in (1). To model this joint distribution, we need a way of representing  $P(G)$ , a distribution over the combinatorial space of possible sparsity patterns, and  $P(\theta |G)$ , the posterior over the parameters of  $f_{\theta}$ . We use GFlowNets [7] to represent  $P(G)$ , and a HyperNetwork to parameterize  $P(\theta |G)$ .
3.2 Generative Flow Networks
GFlowNets are an approach for learning generative models over spaces of discrete objects [7, 8]. GFlowNets learn a stochastic policy  $P_F(\tau)$  to sequentially sample an object  $\mathbf{x}$  from a discrete space  $\mathcal{X}$ . Here  $\tau = (s_0, s_1, \dots, s_n)$  represents a full Markovian trajectory over plausible discrete states, where  $s_n$  is the terminating state (i.e. end of a trajectory) [34]. The GFlowNet is trained such that at convergence, sequential samples from the stochastic policy over a trajectory,  $\mathbf{x} \sim P_F(\tau)$ , i.e.  $\mathbf{x} = s_n$ , are equal in distribution to samples from the normalized reward distribution  $P(\mathbf{x}) = \frac{R(\mathbf{x})}{\sum_{\mathbf{x}' \in \mathcal{X}} R(\mathbf{x}')}$ . The GFlowNet policies are typically trained by optimizing either the Trajectory Balance (TB) loss [34], Subtrajectory Balance (Sub-TB) loss [33], or the Detailed Balance (DB) loss [12]. In this work, we exploit the DB loss to learn a stochastic policy for directed graph structure.
Detailed Balance Loss: The DB loss [12] leverages the fact that the reward function can be evaluated for any partially constructed graph (i.e. any prefix of  $\tau$ ), and hence we get intermediate reward signals for training the GFlowNet policy. The DB loss is defined as:
$$
\mathcal {L} _ {\mathrm {D B}} \left(s _ {i}, s _ {i - 1}\right) = \left(\log \frac {R \left(s _ {i}\right) P _ {B} \left(s _ {i - 1} \mid s _ {i} ; \psi\right) P _ {F} \left(s _ {n} \mid s _ {i - 1} ; \psi\right)}{R \left(s _ {i - 1}\right) P _ {F} \left(s _ {i} \mid s _ {i - 1} ; \psi\right) P _ {F} \left(s _ {n} \mid s _ {i} ; \psi\right)}\right) ^ {2}, \tag {3}
$$
where  $P_F(s_i | s_{i-1}; \psi)$  and  $P_B(s_{i-1} | s_i; \psi)$  represent the forward transition probability and backward transition probability, and a trainable normalizing constant, respectively. Under this formulation, during GFlowNet training the reward is evaluated at every state. For this reason, the DB formulation is in general advantageous for the structure learning problem where any sampled graph can be viewed as a complete state, hence more robustly inform gradients when training the stochastic policy than counterpart losses. Previous work has shown GFlowNets are useful in settings with multi-modal posteriors. This is of particular interest to us where many admissible structures can explain the observed data equally well. We model  $Q_\psi(G)$  using  $P_F(s_i | s_{i-1}; \psi)$  and learn the parameters  $\psi$ .
Algorithm 1 Batch update training of DynGFN
1: Input: Data batch  $(x_{b},dx_{b})$  , initial NN weights  $\psi ,\phi ,L^0$  sparsity prior  $\lambda_0$  , and learning rate  $\epsilon$    
2:  $s_0\gets \mathbf{0}_{B\times d\times d}$  Training is paralleled over  $B$  graph trajectories   
3:  $a\sim P_F(s_1|s_0;\psi)$  Sample initial actions vector   
4: while a not  do   
5: Compute  $P_F(s_i|s_{i - 1};\psi),P_B(s_{i - 1}|s_i;\psi)$    
6:  $\theta \leftarrow h_{\phi}(s_i)$    
7:  $\widehat{dx_b}\gets f_\theta (x,s_i)$    
8:  $R_{i}(s_{i})\gets e^{-\| dx_{b} - \widehat{dx}_{b}\|_{2}^{2} - \lambda_{0}\| s_{i}\|_{0}}$    
9:  $\psi \gets \psi -\epsilon \nabla_{\psi}\mathcal{L}_{DB}(s_i,s_{i - 1})$ $\triangleright \mathcal{L}_{DB}(s_i,s_{i - 1})$  computed as in Equation 3   
10:  $a\sim P_F(s_i|s_{i - 1};\psi),s_i\rightarrow s_{i + 1}$  Take action step to go to next state   
11:  $\phi \gets \phi +\epsilon \nabla_{\phi}\log R$    
return Updated GFN weights  $\psi$  and updated HyperNetwork weights  $\phi$
4 DynGFN for Bayesian Dynamic Structure Learning
We present a general framework for Bayesian dynamic structure learning and propose a GFlowNet architecture, DynGFN, tailored for modelling a posteriori over discrete cyclic graphical structures. We summarize our framework in Figure 1 and Algorithm 1. DynGFN consists of 3 key modules:
1. A graph sampler that samples graphical structures that encode the structural dependencies among the observed variables. This is parameterized with a GFlowNet that iteratively adds edges to a graph.
2. A model that approximates the structural equations defined in (1) to model the functional relationships between the observed variables, indexed by parameters  $\theta$ . This is a class of functions that respect the conditional independencies implied by the graph sampled in step 1. We enforce this by masking inputs according to the graph.
3. Because the functional relationships between variables may be different depending on which graph is sampled, we use a HyperNetwork architecture that outputs the parameters  $\theta$  of the structural equations as a function of the graph. We also show that under linear assumptions of the structural modules, we can solve for optimal  $\theta$  analytically (i.e. we do not need the HyperNetwork).
For training, we assume  $L^0$  sparsity of graphs  $G$  to constrain the large discrete search space over possible structures. We use a reward  $R$  for a graph  $G$  and  $L^0$  penalty of the form:  $R(G) = e^{-\|dx - \widehat{dx}\|^2 + \lambda_0 \|G\|_0}$ . We motivate this set-up so we can estimate  $\widehat{dx}$  close to  $dx$  in an end-to-end learning pipeline. Since estimates for  $\widehat{dx}$  are dependent on  $G$  and  $\theta$ , this reward informs gradients to learn a policy that can approximate  $Q(G)$  given dynamic data.
The main advantage of DynGFN comes when modelling complex posteriors with many modes. Prior work has shown GFlowNets are able to efficiently model distributions where we can share information between different modes [34]. The challenge we tackle is how to do this with a changing objective function, as the GFlowNet objective is a function of the current parameter HyperNetwork and the structural equations. We use multilayer perceptrons (MLPs) to parameterize the stochastic GFlowNet policy, HyperNetwork architecture, and the dynamic structural model<sup>1</sup>.
4.1 Graph Sampler
DynGFN models a posterior distribution over graphs  $Q(G|\mathcal{D})$  given a finite set of observations. To learn  $Q(G|\mathcal{D})$ , DynGFN needs to explore over a large discrete state space. Since we aim to learn a bipartite graph between  $x$  and  $dx$ , DynGFN needs to search over  $2^{d^2}$  possible structures, where  $d$  denotes the dimensionality of the system and  $2^{d^2}$  the number of possible edges in  $G$ . For even moderate  $d$ , this discrete space is very large (e.g. for  $d = 20$  we have  $2^{400}$  possible graphs).
<sup>1</sup When we assume linear dynamic structural relationships, we can solve for the parameters analytically, thus do not need MLPs for the HyperNetwork and dynamic structural model. This is further discussed in section 4.2
However, under the assumption of causal sufficiency, we can significantly reduce this search space, by taking advantage of the fact that  $Q(G|\mathcal{D})$  factorizes as follows,
$$
Q (G | D) = \prod_ {i \in [ 1, \dots , d ]} Q _ {i} (G [, i ] | D) \tag {4}
$$
By using this model, we reduce the search space from  $2^{d^2} \to d2^d$ . For  $d = 20$  this reduces the search space from  $2^{400}$  to  $\approx 2^{24.3}$ . While still intractable to search over, it is still a vast improvement over the unfactorized case. We call this model a per-node posterior, and we use a per-node GFlowNet going forward. We discuss details regarding encouraging forward policy exploration during training in Appendix B.6.
4.2 HyperNetwork and Structural Model
We aim to jointly learn the structural encoding  $G$  and parameters  $\theta$  that together model the structural relationships  $dx = f_{\theta}(x, G)$  of the dynamical system variables. To accomplish this, we propose learning an individual set of parameters  $\theta$  for each graph  $G$ , independent of the input data  $x$ . This approach encapsulates  $P(\theta | G)$  in (2). We use a HyperNetwork architecture that takes  $G$  as input and outputs the structural equation model parameters  $\theta$ , i.e.  $\theta = h_{\phi}(G)$  hence  $P(\theta | G) = \delta(\theta | G)$  allowing us to learn a separate  $\theta$  for each  $G$ . This HyperNetwork model does not capture uncertainty in the parameters, however the formulation may be extended to the Bayesian setting by placing a prior on the HyperNetwork parameters  $\phi$ . Although  $h_{\phi}$  allows for expressive parameterizations for  $\theta$ , it may not be easy to learn<sup>2</sup>. HyperNetworks have shown success in learning parameters for more complex models (e.g. LSTMs and CNNs) [16], hence motivates their fit for our application.
Linear Assumption on Dynamic Structural Model: In some cases it may suffice to assume a linear differential form  $\frac{dx}{dt} = Ax$  to approximate dynamics. In this setting, given a sampled graph  $G \sim Q(G)$  and  $n$  i.i.d. observations of  $(x, dx)$  we can solve for  $\theta = A$  analytically. To induce dependence on the graph structure, we use the sampled  $G$  as a mask on  $x$  and construct  $\tilde{x}_i = G_i^T \odot x$ . Then we can solve for  $\theta$  on a per-node basis as
$$
\theta_ {i} = \left(\tilde {x} _ {i} ^ {T} \tilde {x} _ {i} + \lambda I\right) ^ {- 1} \tilde {x} _ {i} ^ {T} d x _ {i}, \tag {5}
$$
where  $i = 1, \dots, d$ ,  $\lambda > 0$  is the precision of an independent Gaussian prior over the parameters, and  $I$  is the identity matrix. We use  $\lambda = 0.01$  throughout this work.
5 A Useful Model of Indeterminacy
In order to evaluate the ability of DynGFN to model complex posteriors over graphs, we need a structure learning problem with a large equivalence class of admissible graphs. We present a simple way to augment a set of identifiable dynamics under some model to create a combinatorial number of equally likely dynamics under the same model. More specifically, this creates a ground truth posterior  $Q^{*}(G|D) \propto \sum T(G^{*})$  where  $T(\cdot): \mathcal{G} \to \mathcal{G}$  is an analytically computable transformation over graphs and  $G^{*}$  is the identified graph under the original dynamics. We use this system to test how well we can learn a posterior over structures that matches what we see in single-cell data.
Specifically, given a dataset of  $(x,dx)\in \mathbb{R}^d\times \mathbb{R}^d$  pairs, we create a new dataset with  $d + 1$  variables where the 'new' variable  $v^{\prime}$  is perfectly correlated with an existing variable  $v$ . In causal terms, this new variable inherits the same parents as  $v$ , that is  $\mathrm{Pa}(v^{\prime})\coloneqq \mathrm{Pa}(v)$  and the same structural equations as  $v$ , that is  $dv^{\prime} = dv$ . This is depicted in Figure 2. This creates a number of new possible explanatory graphs, which we generalize with the following proposition.