Title: DynGFN: Towards Bayesian Inference of Gene Regulatory Networks with GFlowNets
Venue: NeurIPS 2023
Accepted: yes
Keywords: Bayesian Structure Learning

Generative Flow Networks

Single-cell

Dynamical Systems
Abstract: 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 directed acyclic graphs, but not both. In this paper we leverage the fact that it is possible to estimate the ``velocity'' of the expression of a gene 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. 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.

DynGFN: Towards Bayesian Inference of Gene Regulatory Networks with GFlowNets

Abstract

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.

1 Introduction

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

(

,

,

) =

(

,

)

(

)

(

)

̂

G

∼

Q ψ

(

G

)

1

0

1

…

0

1

0

…

⋮

⋮

/uni22F1

1

0

0

…

⋮

G

θ

θ

Structural Model

{

G

θ

x

}

⋮

⋮

x i

d x i

G

p

G

θ

D

p

D

|

G

θ

p

θ

|

G

p

G

Qψ

(

G

) ∝

p

(

G

) →

G

p

(

θ

|

G

) =

δ θ

h ϕ

(

G

) →

θ

p

(

D

|

G

,

θ

) ∝

e

-∥

dx

-

f θ

(

x

,

G

)∥

f θ

(

x

,

G

) →

d x

x

d x

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.

2 Related Work

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 = 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 Preliminaries

3.1 Bayesian Dynamic Structure Learning

Problem Setup: We consider a finite dataset, D , of dynamic pairs ( x, dx ) ∈ R d × R d where x respresents the state of the system sampled from an underlying time-invariant stochastic dynamical system governed by a latent drift dx dt = f ( x, ϵ ) where ϵ is a noise term that parameterizes the SDE; x and ϵ are mutually independent. The latent drift has some fixed sparsity pattern i.e. ∂f i ∂x j = 0 for a small set of variables, which can be parameterized by a graph G such that g ij = 1 [ ∂f i ∂x j = 0] , where g ij ∈ G,i = 1 , . . . , d, j = 1 , . . . , d . The variables x j for which ∂f i ∂x j = 0 can be interpreted as the causal parents of x i , denoted Pa ( x i ) . This lets us define an equivalent dynamic structural model [37, 43] of the form,

̸

for i = 1 , . . . , 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 |D ) given the data. We aim to jointly learn distribution over parameters θ that parameterize the latent drift f ( x ) ; these parameters will typically depend on the sparsity pattern such that p ( θ | G ) = p ( θ ) . We can factorize this generative model as follows,

This factorization forms the basis of our inference procedure. We learn a parameterized function f θ ( x ) : R d → 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 ( θ | G ) , the posterior over the parameters of f θ . We use GFlowNets [7] to represent P ( G ) , and a HyperNetwork to parameterize P ( θ | 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 ( τ ) to sequentially sample an object x from a discrete space X . Here τ = ( s 0 , s 1 , . . . , 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, x ∼ P F ( τ ) , i.e. x = s n , are equal in distribution to samples from the normalized reward distribution P ( x ) = R ( x ) ∑ x ′ ∈X R ( 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 τ ), and hence we get intermediate reward signals for training the GFlowNet policy. The DB loss is defined as:

where P F ( s i | s i -1 ; ψ ) and P B ( s i -1 | s i ; ψ ) 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 ψ ( G ) using P F ( s i | s i -1 ; ψ ) and learn the parameters ψ .

̸

̸

Algorithm 1 Batch update training of DynGFN

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 posterior over discrete cyclic graphical structures. We summarize our framework in Figure 1 and Algorithm 1. DynGFN consists of 3 key modules:

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.

A model that approximates the structural equations defined in (1) to model the functional relationships between the observed variables, indexed by parameters θ . 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.

Because the functional relationships between variables may be different depending on which graph is sampled, we use a HyperNetwork architecture that outputs the parameters θ 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 θ 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 -̂ dx ∥ 2 2 + λ 0 ∥ G ∥ 0 . We motivate this set-up so we can estimate ̂ dx close to dx in an end-to-end learning pipeline. Since estimates for ̂ dx are dependent on G and θ , 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 1 .

4.1 Graph Sampler

DynGFN models a posterior distribution over graphs Q ( G |D ) given a finite set of observations. To learn Q ( G |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).

However, under the assumption of causal sufficiency, we can significantly reduce this search space, by taking advantage of the fact that Q ( G |D ) factorizes as follows,

By using this model, we reduce the search space from 2 d 2 → d 2 d . For d = 20 this reduces the search space from 2 400 to ≈ 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 θ that together model the structural relationships dx = f θ ( x, G ) of the dynamical system variables. To accomplish this, we propose learning an individual set of parameters θ for each graph G , independent of the input data x . This approach encapsulates P ( θ | G ) in (2). We use a HyperNetwork architecture that takes G as input and outputs the structural equation model parameters θ , i.e. θ = h ϕ ( G ) hence P ( θ | G ) = δ ( θ | G ) -allowing us to learn a separate θ 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 ϕ . Although h ϕ allows for expressive parameterizations for θ , it may not be easy to learn 2 . 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 dx dt = A x to approximate dynamics. In this setting, given a sampled graph G ∼ Q ( G ) and n i.i.d. observations of ( x, dx ) we can solve for θ = A analytically. To induce dependence on the graph structure, we use the sampled G as a mask on x and construct ˜ x i = G T i ⊙ x . Then we can solve for θ on a per-node basis as

where i = 1 , . . . d , λ > 0 is the precision of an independent Gaussian prior over the parameters, and I is the identity matrix. We use λ = 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 ) ∝ ∑ T ( G ∗ ) where T ( · ) : G → 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 ) ∈ R d × R d pairs, we create a new dataset with d +1 variables where the 'new' variable v ′ is perfectly correlated with an existing variable v . In causal terms, this new variable inherits the same parents as v , that is Pa ( v ′ ) := Pa ( v ) and the same structural equations as v , that is dv ′ = dv . This is depicted in Figure 2. This creates a number of new possible explanatory graphs, which we generalize with the following proposition.

Proposition 1. Given any d dimensional ODE system with G ∗ identifiable under f ∈ F , the D = d + a dimensional system dx dt = A x , denote the vector of multiplicities m ∈ N d with m i as the number of repetitions of each variable. Then this construction creates an admissible family of graphs G ′ where |G ′ | = ∏ i ∈ d (2 m i -1) G i 1 . Furthermore, under an L 0 penalty on G , this reduces to ∏ i ( m i ) G i 1 .

See Appendix A for full proof. The intuition behind this proposition can be seen from the case of adding a single copied variable. This corresponds to A = [ δ v I d ] where δ v is a vector with a 1 on

Static Graph

Dynamic Graph

𝑮

∗

𝑣

!

𝑣

"

𝑣

#

𝒙

𝒅𝒙

𝑣

!

𝑣

"

𝑣

#

𝑣

#

$

𝑮

"

𝒙

𝒅𝒙

𝒙

𝒅𝒙

𝒙

𝒅𝒙

node v and zeros elsewhere, and I d is the d -dimensional identity matrix. Let v have c children, such that v ∈ Pa ( c ) in the identifiable system, then any of those c child nodes could depend either on v or on the new node v ′ or both. This creates 3 c possible explanatory graphs. If we restrict ourselves to the set of graphs with minimal L 0 norm, then we eliminate the possibility of a child node depending on both v and v ′ , this gives 2 c possible graphs, choosing either v or v ′ as a parent.

6 Experimental Results

In this section we evaluate the performance of DynGFN against counterpart Bayesian structure learning methods (see Appendix B.2 for details). Since our primary objective is to learn Bayesian posteriors over discrete structure G , we compare to Bayesian methods that can also accomplish this task, i.e. versions of BCD-Nets [11] and DiBS [32]. We show in certain cases, DynGFN is able to better capture the true posterior when there are a large number of modes. We evaluate methods according to four metrics: Bayes-SHD, area under the receiver operator characteristic curve (AUC), Kullback-Leibler (KL) divergence between learned posteriors Q ( G ) and the distribution over true graphs P ( G ∗ ) , and the negative log-likelihood (NLL) P ( D | G,θ ) (in our setting this reduces to the mean squared error between ̂ dx and dx , given θ and sampled G ′ s ). Since the analytic linear solver requires data at run-time to compute optimal parameters for the structural model, we include the NLL metric only for models using the HyperNetwork solver. Bayes-SHD measures the average distance to the closest structure in the admissible set of graphs according to the structural hamming distance, which in this case is simply the hamming distance of the adjacency matrix representation to the closest admissible graph. We assume P ( G ∗ ) is uniform over G ∗ and include further details about evaluating the quality of learned posteriors in Appendix B.8.

6.1 A Toy Example with Synthetic Data

To validate the ability of DynGFN to learn cyclic dependencies we consider identifiable acyclic and cyclic 3 variable toy systems and provide a comparison with a DAG structure learning method (NOTEARS [54]). We show results for this toy example in Figure 3. NOTEARS does not model cyclic dependencies and therefore struggles to yield accurate predictions of A in the cyclic setting. We can also verify this result by considering a conditional independence test over cyclic dependencies. It is easy to see that the conditional independence test fails in the cyclic setting: in the acyclic case, we can identify the v-structure by observing that x 1 ⊥ x 3 and x 1 ̸⊥ x 3 | x 2 , which implies that x 2 is a collider (i.e. x 1 and x 3 are marginally independent and conditionally dependent); while in the cyclic example, we introduce time dependencies such that there are cycles in the summary graph that render these variables marginally dependent. We show that DynGFN is able to identify the true cyclic dependencies in this toy example. We note that in this toy example DynGFN exhibits a degree

Acyclic

Cyclic

NOTEARS

-

0.21

0.57

Acyclic

1.0

0.58

Cyclic

AUC-ROC

𝑥

!

𝑥

"

𝑥

#

𝑥

!

𝑥

"

𝑥

#

of convergence sensitivity on a per-run basis. In the following sections we provide comprehensive results over 5 random seeds and consider larger systems.

6.2 Experiments with Synthetic Data

We generated synthetic data from two systems using our indeterminacy model presented in section 5: (1) a linear dynamical system dx = A x , and (2) a non-linear dynamical system dx = sigmoid ( A x ) . We consider ℓ -DynGFN and h -DynGFN, i.e. DynGFN with the linear analytic parameter solver as shown in (5), and DynGFN with the HyperNetwork parameter solver h ϕ . Likewise, we compare ℓ -DynGFN and h -DynGFN to counterpart Bayesian baselines which we call ℓ -DynBCD, ℓ -DynDiBS, h -DynBCD, and h -DynDiBS. To constrain the discrete search procedure, we assume a sparse prior on the structure (i.e. the graphs G ), using the L 0 prior. Due to challenging iterative optimization dynamics present when using θ = h ϕ ( G ) for DynGFN, to train initialize the forward policy P F ( s i | s i -1 ; ψ ) using the ψ learned in ℓ -DynGFN to provide a more admissible starting point for learning h ϕ (we

discuss further details in Appendix B.7). We do not need to do this for h -DynBCD and h -DynDiBS as we are able to train both models end-to-end without iterative optimization. In Table 1 we show results of our synthetic experiments for learning posteriors over multi-modal distributions of cyclic graphs. We observe the DynGFN is most competitive on both synthetic systems for modelling the true posterior over structure. Details about DynGFN, baselines, and accompanying hyper-parameters can be found in Appendix B.

6.3 Ablations Over Sparsity and Linearity of Dynamic Systems

We conduct two ablations: (1) ablation over sparsity of the dynamic system structure, and (2) ablation over ∆ t , the time difference between data points of dynamic simulation. For a sparsity level of 0 . 9 , the ground truth graphs have 50 edges out of d 2 possible edges. In these experiments, P ( G ∗ ) for d = 20 and sparsity 0 . 9 has 1024 modes. We conduct the ablations over 5 random seeds for each set of experiments.

Sparsity: DynGFN uses the L 0 prior on G throughout training. Under this setting, system sparsity carries significant weight on the ability to learn posteriors over the structured dynamics of a system. We show this trend in Table 2. We note that computing the KL-divergence for DynGFN, specifically computing the probability of generating a true G , becomes computationally intractable as G is less sparse 3 . For systems of 0 . 9 and 0 . 95 sparsity, we observe a decreasing trend in KL and Bayes-SHD, and an increasing trend in AUC. This result is expected as DynGFN can better traverse sparse graphs as the combinatorial space over possible trajectories is smaller relative to denser systems.

Linearity: Training DynGFN via the linear solver for the structural model parameters is an easier objective due to simplified

training dynamics. Because of this, we explore the performance of ℓ -DynGFN assuming f θ for modelling equation (1) to be linear in the non-linear system. We do this by conducting an ablation over ∆ t and find that the performance of ℓ -DynGFN on the non-linear system improves as ∆ t → 0 . We show a portion of this trend in Table 3.

6.4 Experiments on Single-Cell RNA-velocity Data

To show how DynGFN can be applied to single cell data we use a cell cycle dataset of human Fibroblasts [46]. As a motivating example we show the correlation structure of single-cell RNA-seq data from human Fibroblast cells [46] Figure 4. We show both the raw correlation and the correlation over cell cycle time, which is significantly higher. With such a pure cell population whose primary axis of variation is state in the cell cycle by aggregating over cell cycle time we expect observation noise to be averaged out, leading to a 'truer' view of the correlation between latent variables. Further details for this experimental set-up are provided in Appendix C.1.3. Since there are many genes which are affected by the cell cycle phase, there are many correlated variables that are downstream of the true cell cycle regulators. This provides a natural way of using cell cycle data to evaluate a model's ability to capture the Bayesian posterior. In Table 4 we show results for learning posteriors over an undetermined GRN using RNA velocity data. We find that ℓ -DynGFN and h -DynGFN yield low KL

(a)

(b)

(c)

Cdc25A

Cdk1

Cdc25C

Cell Cycle Module

Mcm2

Mcm5

𝜌 > 0.75

True edges

Potential edges

(d)

HST1H3G -

HNRNPA3 -

GOLGA8A -

CYP51A1

CARMIL2-

SLC40A1-

ADRAID

SNHG10 -

ADGRL2

ZCCHC8 -

SLC8A1-

COOBA

ANXA3-

SYDE2

KCNK2-

HMMR

P4HA2

USP46 -

PMP22

FOSB

GDF5-

LSMS-

CCNF-

RRM2-

ORC6-

PLAT

PLAT

3

DXDO

PRSS35

SAMD

ID

PHLDA1

135

DON

PRDM1

TIVH

CD2

- 0.50

- 0.25

- 0.75

0.75

- 0.00

0.50

1.00

0.25

1.00

HIST1H3G

GOLGABA

HNRNPA3

CYP51A1

CARMIL2

ADRA1D

SLC40A1

ADGRL2

SNHG10

ZCCHC8

SLCBA1

COQBA

P4HA2

ANXA3

SYDE2

USP46

KCNK2

HMMR

PMP22

LSM5

CCNF

RRM2

FOSB

GDF5

ORC6

PLAT

nns

NABP

DIOO

MT1X

ID

I3S

DON

- 0.00

- 0.75

0.50

= 0.25

0.75

0.25

0.50

1.00

1.00

3.5

2.5

0.5

4.0

3.0

2.0

1.5

1.0

0.0

1e6

0.0

0.2

0.4

0.6

Time Pearson-p

Pearson-p

0.8

1.0

and moderate Bayes-SHD. While ℓ -DynBCD performs well in terms of identify a small distribution of true G 's, it falls short in modelling the true posterior (this can be seen from low Bayes-SHD, high KL).

7 Conclusion

We presented DynGFN, a method for Bayesian dynamic structure learning. In low dimensions we found that DynGFN is able to better model the distribution over possible explanatory structures than counterpart Bayesian structure learning baseline methods. As a proof of concept, we presented an example of learning the distribution over likely explanatory graphs for linear and non-linear synthetic systems where complex uncertainty over explanatory structure is prevalent. We demonstrate the use of DynGFN for learning gene regulatory structure from single-cell transcriptomic data where there are many possible graphs, showing DynGFN can better model the uncertainty over possible explanations of this data rather than capturing a single explanation.

Limitations and Future Work: We have demonstrated a degree of efficacy when using DynGFN for Bayesian structure learning with dynamic observational data. A key limitation of DynGFN is scaling to larger systems. To effectively model P ( G,θ, D ) , DynGFN needs to search over an environment state space of possible graphs. This state space grows exponentially with the number of possible edges, i.e. 2 d 2 or d 2 d for per-node-GFN where d is the number of variables in the system. Therefore, DynGFN is currently limited to smaller systems. Nevertheless, there are many applications where Bayesian structure learning, even over 5-20 dimensional examples that we explore here, could be extraordinarily useful. We include further discussion of scaling DynGFN in Appendix C.3 with some ideas on how to approach this challenge. We found that training DynGFN requires some selection of hyper-parameters and in particular parameters that shape the reward function. Selecting hyper-parameters for the baseline methods prove more difficult for this task.

8 Acknowledgments and Disclosure of Funding

This research was enabled in part by the computational researches provided by Mila ( mila.quebec ) and Compute Canada ( ccdb.computecanada.ca ). In addition, resources used in preparing this research were in part provided by the Province of Ontario and companies sponsoring the Vector Institute ( vectorinstitute.ai/partners/ ). All authors are funded by their primary academic institution. We also acknowledged funding from the Natural Sciences and Engineering Research Council of Canada, Recursion Pharmaceuticals, CIFAR, Samsung, and IBM. We are grateful to Cristian Dragos Manta for catching some typos in the manuscript.

References

Aliee, H., Theis, F. J., and Kilbertus, N. Beyond predictions in neural odes: Identification and interventions. arXiv preprint 2106.12430 , 2021.

Aliee, H., Richter, T., Solonin, M., Ibarra, I., Theis, F., and Kilbertus, N. Sparsity in continuousdepth neural networks. Advances in Neural Information Processing Systems (NeurIPS) , 2022.

Am´ endola, C., Dettling, P., Drton, M., Onori, F., and Wu, J. Structure learning for cyclic linear causal models. Uncertainty in Artificial Intelligence (UAI) , 2020.

Annadani, Y., Rothfuss, J., Lacoste, A., Scherrer, N., Goyal, A., Bengio, Y., and Bauer, S. Variational causal networks: Approximate bayesian inference over causal structures. arXiv preprint , 2021.

Beaumont, P., Horsburgh, B., Pilgerstorfer, P., Droth, A., Oentaryo, R., Ler, S., Nguyen, H., Ferreira, G. A., Patel, Z., and Leong, W. CausalNex, 2021. URL https://github.com/ quantumblacklabs/causalnex .

Bellot, A. and Branson, K. Neural Graphical Modelling in Continuous Time: Consistency Guarantees and Algorithms. International Conference on Learning Representations (ICLR) , 2022.

Bengio, E., Jain, M., Korablyov, M., Precup, D., and Bengio, Y. Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation. Advances in Neural Information Processing Systems (NeurIPS) , 2021.

Bengio, Y., Deleu, T., Hu, E. J., Lahlou, S., Tiwari, M., and Bengio, E. GFlowNet Foundations. arXiv preprint 2111.09266 , 2022.

Bergen, V., Lange, M., Peidli, S., Wolf, F. A., and Theis, F. J. Generalizing RNA velocity to transient cell states through dynamical modeling. BioRxiv preprint 820936 , 2019.

Chen, R. T. Q., Rubanova, Y., Bettencourt, J., and Duvenaud, D. Neural Ordinary Differential Equations. Advances in Neural Information Processing Systems (NeurIPS) , 2018.

Cundy, C., Grover, A., and Ermon, S. BCD Nets: Scalable Variational Approaches for Bayesian Causal Discovery. Advances in Neural Information Processing Systems (NeurIPS) , 2021.

Deleu, T., G´ ois, A., Emezue, C., Rankawat, M., Lacoste-Julien, S., Bauer, S., and Bengio, Y. Bayesian Structure Learning with Generative Flow Networks. Uncertainty in Artificial Intelligence (UAI) , 2022.

Friedman, N., Murphy, K., and Russell, S. Learning the structure of dynamic probabilistic networks. Uncertainty in Artificial Intelligence (UAI) , 1998.

Glymour, C., Zhang, K., and Spirtes, P. Review of Causal Discovery Methods Based on Graphical Models. Frontiers in Genetics , 2019.

Granger, C. W. Investigating causal relations by econometric models and cross-spectral methods. Econometrica: journal of the Econometric Society , 1969.

Ha, D., Dai, A. M., and Le, Q. V. Hypernetworks. International Conference on Learning Representations (ICLR) , 2017.

Hashimoto, T. B., Gifford, D. K., and Jaakkola, T. S. Learning Population-Level Diffusions with Generative Recurrent Networks. International Conference on Machine Learning (ICML) , 2016.

Hauser, A. and B¨ uhlmann, P. Characterization and greedy learning of interventional markov equivalence classes of directed acyclic graphs. Journal of Machine Learning Research (JMLR) , 13(1), 2012.

He, Y.-B. and Geng, Z. Active learning of causal networks with intervention experiments and optimal designs. Journal of Machine Learning Research (JMLR) , 9, 2008.

Huang, B., Zhang, K., Gong, M., and Glymour, C. Causal discovery and forecasting in nonstationary environments with state-space models. International Conference on Machine Learning (ICML) , 2019.

Huang, B., Zhang, K., Zhang, J., Ramsey, J., Sanchez-Romero, R., Glymour, C., and Sch¨ olkopf, B. Causal discovery from heterogeneous/nonstationary data. The Journal of Machine Learning Research (JMLR) , 21, 2020.

Huguet, G., Magruder, D. S., Tong, A., Fasina, O., Kuchroo, M., Wolf, G., and Krishnaswamy, S. Manifold interpolating optimal-transport flows for trajectory inference. Advances in Neural Information Processing Systems (NeurIPS) , 2022.

Huguet, G., Tong, A., Zapatero, M. R., Wolf, G., and Krishnaswamy, S. Geodesic Sinkhorn: Optimal transport for high-dimensional datasets. arXiv preprint 2211.00805 , 2022.

Itani, S., Ohannessian, M., Sachs, K., Nolan, G. P., and Dahleh, M. A. Structure learning in causal cyclic networks. Proceedings of Workshop on Causality: Objectives and Assessment at NIPS 2008 , 2010.

Karlebach, G. and Shamir, R. Modelling and analysis of gene regulatory networks. Nature reviews Molecular cell biology , 9, 2008.

Ke, N. R., Bilaniuk, O., Goyal, A., Bauer, S., Larochelle, H., Pal, C., and Bengio, Y. Learning neural causal models from unknown interventions. arXiv preprint 1910.01075 , 2019.

La Manno, G., Soldatov, R., Zeisel, A., Braun, E., Hochgerner, H., Petukhov, V., Lidschreiber, K., Kastriti, M. E., L¨ onnerberg, P., Furlan, A., Fan, J., Borm, L. E., Liu, Z., van Bruggen, D., Guo, J., He, X., Barker, R., Sundstr¨ om, E., Castelo-Branco, G., Cramer, P., Adameyko, I., Linnarsson, S., and Kharchenko, P. V. RNA velocity of single cells. Nature , 560, 2018.

Lacerda, G., Spirtes, P. L., Ramsey, J., and Hoyer, P. O. Discovering cyclic causal models by independent components analysis. arXiv preprint 1206.3273 , 2012.

Lachapelle, S., Brouillard, P., Deleu, T., and Lacoste-Julien, S. Gradient-Based Neural DAG Learning. International Conference on Learning Representations (ICLR) , 2020.

Liu, Q. and Wang, D. Stein variational gradient descent: A general purpose bayesian inference algorithm. Advances in Neural Information Processing Systems (NeurIPS) , 2016.

Lopez, R., H¨ utter, J.-C., Pritchard, J. K., and Regev, A. Large-scale differentiable causal discovery of factor graphs. Advances in Neural Information Processing Systems (NeurIPS) , 2022.

Lorch, L., Rothfuss, J., Sch¨ olkopf, B., and Krause, A. DiBS: Differentiable Bayesian Structure Learning. Advances in Neural Information Processing Systems (NeurIPS) , 2021.

Madan, K., Rector-Brooks, J., Korablyov, M., Bengio, E., Jain, M., Nica, A., Bosc, T., Bengio, Y., and Malkin, N. Learning gflownets from partial episodes for improved convergence and stability. arXiv preprint 2209.12782 , 2022.

Malkin, N., Jain, M., Bengio, E., Sun, C., and Bengio, Y. Trajectory Balance: Improved Credit Assignment in GFlowNets. Advances in Neural Information Processing Systems (NeurIPS) , 2022.

Monti, R. P., Zhang, K., and Hyv¨ arinen, A. Causal discovery with general non-linear relationships using non-linear ica. Uncertainty in Artificial Intelligence (UAI) , 2020.

Mooij, J. M., Janzing, D., Heskes, T., and Sch¨ olkopf, B. On causal discovery with cyclic additive noise models. Advances in Neural Information Processing Systems (NeurIPS) , 2011.

Mooij, J. M., Janzing, D., and Sch¨ olkopf, B. From Ordinary Differential Equations to Structural Causal Models: The deterministic case. Uncertainty in Artificial Intelligence (UAI) , 2013.

Mooij, J. M., Magliacane, S., and Claassen, T. Joint causal inference from multiple contexts. The Journal of Machine Learning Research (JMLR) , 21, 2020.

Murphy, K. P. Active learning of causal bayes net structure. Technical report, UC Berkeley, 2001.

Pamfil, R., Sriwattanaworachai, N., Desai, S., Pilgerstorfer, P., Georgatzis, K., Beaumont, P., and Aragam, B. Dynotears: Structure learning from time-series data. International Conference on Artificial Intelligence and Statistics (AISTATS) , 2020.

Pearl, J. Causality . Caimbridge University Press, second edition, 2009.

Peters, J., Janzing, D., and Sch¨ olkopf, B. Elements of causal inference: foundations and learning algorithms . The MIT Press, 2017.

Peters, J., Bauer, S., and Pfister, N. Causal models for dynamical systems. Probabilistic and Causal Inference: The Works of Judea Pearl , 2022.

Pratapa, A., Jalihal, A. P., Law, J. N., Bharadwaj, A., and Murali, T. Benchmarking algorithms for gene regulatory network inference from single-cell transcriptomic data. Nature methods , 17, 2020.

Qiu, X., Zhang, Y., Martin-Rufino, J. D., Weng, C., Hosseinzadeh, S., Yang, D., Pogson, A. N., Hein, M. Y., Hoi (Joseph) Min, K., Wang, L., Grody, E. I., Shurtleff, M. J., Yuan, R., Xu, S., Ma, Y., Replogle, J. M., Lander, E. S., Darmanis, S., Bahar, I., Sankaran, V. G., Xing, J., and Weissman, J. S. Mapping transcriptomic vector fields of single cells. Cell , 185, 2022.

Riba, A., Oravecz, A., Durik, M., Jim´ enez, S., Alunni, V., Cerciat, M., Jung, M., Keime, C., Keyes, W. M., and Molina, N. Cell cycle gene regulation dynamics revealed by RNA velocity and deep-learning. Nature Communications , 13, 2022.

Saelens, W., Cannoodt, R., Todorov, H., and Saeys, Y. A comparison of single-cell trajectory inference methods. Nature Biotechnology , 37, 2019.

Schiebinger, G., Shu, J., Tabaka, M., Cleary, B., Subramanian, V., Solomon, A., Gould, J., Liu, S., Lin, S., Berube, P., Lee, L., Chen, J., Brumbaugh, J., Rigollet, P., Hochedlinger, K., Jaenisch, R., Regev, A., and Lander, E. S. Optimal-Transport Analysis of Single-Cell Gene Expression Identifies Developmental Trajectories in Reprogramming. Cell , 176, 2019.

Spirtes, P. and Glymour, C. An algorithm for fast recovery of sparse causal graphs. Social Science Computer Review , 9, 1991.

Tank, A., Covert, I., Foti, N., Shojaie, A., and Fox, E. Neural Granger Causality. IEEE Transactions on Pattern Analysis and Machine Intelligence , 2021.

Tong, A., Huang, J., Wolf, G., van Dijk, D., and Krishnaswamy, S. TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics. International Conference on Machine Learning (ICML) , 2020.

Tong, S. and Koller, D. Active learning for structure in bayesian networks. International Joint Conference on Artificial Intelligence (IJCAI) , 17, 2001.

Zhang, K., Huang, B., Zhang, J., Glymour, C., and Sch¨ olkopf, B. Causal discovery from nonstationary/heterogeneous data: Skeleton estimation and orientation determination. International Joint Conference on Artificial Intelligence (IJCAI) , 2017.

Zheng, X., Aragam, B., Ravikumar, P., and Xing, E. P. DAGs with NO TEARS: Continuous Optimization for Structure Learning. Advances in Neural Information Processing Systems (NeurIPS) , 2018.

Supplementary Material

A Proof of Proposition 1

Proposition 1 calculates the number of admissible structure graphs for a linear ODE system with correlated variables. We will first show the general case this is ∏ i ∈ d (2 m i -1) G i 1 , then analyze the case of an L 0 penalty on the edges of G , which reduces the size of the set of admissible graphs to ∏ i ( m i ) G i 1 .

Proof. Consider an identifiable linear system dx dt = Ax where we directly observe ( x, dx dt ) with G ∗ identifiable. Then the system with m = 1 d has exactly one admissible graph by definition. For each node, we analyze its set of child nodes in G , i.e. c ( u ) = { v ∈ V s.t. u → v ∈ G} . For an identifiable system, each child v must have an incoming edge from its parent.

Next, we consider the process of adding a correlated variable, i.e. consider the situation of w.l.o.g. consider m = ( s, 1 , 1 , . . . , 1) for some s > 1 . Then for each child of c j ( v 1 ) , there are now s possible parents. This has multiplied the number of possible graphs by 2 s -1 . Since each element of m is independent, this leads to the first statement, i.e. |G ′ | = ∏ i ∈ d (2 m i -1) G i 1 .

Under an L 0 penalty, then we constrain the possible graphs to s different graphs, where each child node picks exactly one of the s possible parents. This leads to the second statement, |G ′ | = ∏ i ( m i ) G i 1

B Experimental Details

B.1 Single Cell Dataset Preprocessing

We start with the processed data from [46]. We first filter it applying steps from the ScVelo tutorial. We then sub-select the genes of interest and use the 'Ms' and 'velocity' layers, which we normalize to mean zero standard deviation one for the states and scale the dx with the same parameters.

B.2 Baselines for Bayesian Dynamic Structure Learning

Existing Bayesian structure learning methods are typically constrained to learning DAGs. Temporal information about the the dynamic relationships amongst variables in a system can help alleviate this constraint. DiBS and BCD-Nets are two state-of-the-art Bayesian structure learning approaches for static systems. We apply versions of DiBS and BCD-Nets such that they are applicable in our Bayesian dynamic structure learning framework for learning cyclic graph structure from dynamic data. We use the approach taken in DiBS and parameterize the distribution over graphs as P α t ( G | Z ) = ∏ i ∏ j P α t ( G ij | Z ij ) , where Z = U T V , U, V ∈ R k × d . Here P α t ( G ij = 1 | Z ij ) = σ ( α t Z ij ) , σ ( x ) = 1 / (1 + e -α t x ) , and α t = αc ( t ) ( t denotes the training iteration. We use c ( t ) = √ t ). As t → ∞ , P α t ( G | Z ) → δ ( G | Z ) . In DiBS, Stein variational gradient decent (SVGD) [30] is used to iteratively transport particles Z to learn the target distribution. Following from the above parameterization, we implement a version of BCD-Nets by treating U ∼ N ( µ u , σ 2 u ) , V ∼ N ( µ v , σ 2 v ) , and learning µ u , µ v , σ u , and σ v . Since our framework uses dynamic data, we incorporate DiBS and BCD-Nets within the framework (labeled DynDiBS and DynBCD, respectively) to leverage dynamic information for Bayesian structure learning of cyclic graphs.

B.3 Hyper-parameters for Baselines

For both DynBCD and DynDiBS we use k = d across datasets. Since DynDiBS is an ensemble based method, we use 1024 samples for the linear and non-linear synthetic systems and 1000 samples for the cellular system (both training and evaluation). Since DynBCD is a variational approach and doesn't require parallelized model ensembles, we use a large quantity of samples for training and evaluation. In the case of the cellular system, since there is a significantly smaller quantity of admissible graphs, we use less samples for DynBCD. We use graph sparsity regularization denoted by λ 0 and a temperature parameter T that scales the magnitude of the likelihood (e.g. 1 T 2 MSE ( dx, ̂ dx ) ).

In Table 5 and Table 6 we outline the hyper-parameters we found to yield the most competitive results. We use grid search to tune DynBCD and DynDiBS. All baselines are trained for 1000 epochs.

We note that when evaluation on validation and test data for Bayes-SHD and AUC metrics, we hard threshold P α t ( G | Z ) . We find that through training this the final α t is typically small enough in magnitude such that P α t ( G | Z ) does not yield a full threshold of Z . To this end, we select large α t when computing the KL metric to mimic hard threshold behaviour as experienced during training. We use α t = 1 × 10 8 . In DynDiBS the parameter γ helps control separation of particles Z during training. In general, we found DynBCD and DynDiBS baselines are challenging to train and to find hyper-parameter settings with good performance. In part, we believe this is due to the numerous hyper-parameters required to tune as well as the general difficulty of the objective.

B.4 Neural Network Architectures and Hyper-parameters

We parameterize P F ( s i | s i -1 ; ψ ) and h ϕ with MLP architectures. P F ( s i | s i -1 ; ψ ) takes the current state as input and first computes common representations using a 3 layer MLP. Then a 2 layer MLP with a softmax output activation takes the representations as input and outputs a distribution over possible actions. The latter MLP is used to parameterize one head for each distribution P F ( s i | s i -1 ; ψ ) . We use a hidden unit dimension of 128 and leaky rectified linear unit (Leaky ReLU) activation functions for the P F ( s i | s i -1 ; ψ ) MLP architecture. We use a uniform backward policy for P B ( s i -1 | s i ; ψ ) . To parameterize h ϕ , we use a 3 layer MLP with hidden layer dimensions of

Linear System

{ 64 , 64 , 64 } and exponential linear unit activations (ELU). We consider two parametrizations for f θ : single linear parameters, i.e dx = θx , and a single hidden layer neural network dx = f θ ( x ) . We use these parametrizations to model linear and non-linear node-wise parent-child structural models, where x ∈ R d are the node-wise input observations.

B.5 Hyper-parameters for DynGFN

DynGFN requires setting a variety of hyper-parameters that lead to different trade offs in model performance. In particular, λ 0 (sparsity encouragement for identified graphs), a temperature parameter T that scales the magnitude of the reward likelihood (e.g. 1 T 2 MSE ( dx, ̂ dx ) ), learning rate ϵ , softmax tempering c (we always use a cosine schedule for c, with a discrete period of 5 epochs), and number of training epochs. In our experiments we select hyper-parameter values that we find lead to competitive performance (this pertains to ℓ -DynGFN and h -DynGFN models) by observing performance over a few hyper-parameter values. We outline the selected hyper-parameters for each respective model in Table 7. For GFlowNet sampling, due to computational limits, we use less training samples m train than evaluation samples m eval for DynGFN. We train DynGFN for 1000 epochs.

B.6 GFlowNet Exploration vs. Exploitation

The general procedure for training GFlowNets is inspired from reinforcement learning where the primary objective is to learn a stochastic policy π ( a | s ) to sample actions from an action space given a current state. In our setting, the action space represents possible locations where an additional edge can be placed to an existing graph and each state is represented by a current graph. Since under this training procedure we are sampling from the GFlowNet policy P F ( s i | s i -1 ; ψ ) at every iteration then attributing a reward associated to the sampled state/graph, the policy is susceptible to exploitation: if P F ( s i | s i -1 ; ψ ) samples a graph(s) with a high reward, it becomes easy for the policy to focus on sampling said graphs since they yield high reward. To alleviate this we encourage exploration using softmax tempering on our stochastic policy, by multiplying the logits of our forward policy by 1 /c before applying the softmax function. A larger c flattens the stochastic policy such that exploration within the action space is encouraged. However, setting the parameters c is challenging and there exists a trade-off between exploring and exploiting the stochastic policy during optimization. We address this by using a cosine schedule for c such that 1 ≤ c ≤ 1000 . We treat the period of the cosine schedule as a hyper-parameter.

B.7 Discussion of Training Dynamics

GFlowNets are a relatively recent class of models that can be challenging to optimize. We discuss some of the challenges with training them especially in the context of a learned energy function.

0

200

400

600

800

1000

0.01

0.02

0.03

0.04

0.05

0.06

0.07

0.08

MSE

0

200

400

600

800

1000

25

50

75

100

125

150

175

Bayes-SHD

We observed that in settings where the energy reward is fixed and we could proportionally penalize missing edges as well as the addition of incorrect edges (e.g. ℓ -DynGFN), we were able to better learn posteriors over admissible graphs over models that require sparse priors and/or trainable energies. This suggests that DynGFN may be limited by an inadequate energy reward. However, we found training DynGFN with a trainable energy function challenging since the GFlowNet stochastic policy depends on the rewards, and vice versa. Further investigation and experimentation into this alternating optimization procedure is required. To in part alleviate the constraints of this challenge, in this work we presented results for h -DynGFN which used initialized GFlowNet parameters from a trained ℓ -DynGFN model on the same task. We found that this initialization technique enabled more effective iterative optimization of the forward policy and energy reward. In Figure 5 we show validation curves throughout training of ℓ -DynGFN on a linear system with d = 20 and a sparsity value of 0 . 9 .

B.8 Evaluating Quality of Learned Posteriors

Using our indeterminacy model defined in section 5, we can determine P ( G ∗ ) for a given set of correlated variables, where G ∗ denotes the set of true equally admissible structures. Here we assume P ( G ∗ ) is a uniform distribution over G ∗ and determine the KL between Q ( G ) and P ( G ∗ ) . We compute the KL considering only the probabilities of a trained models to generate all structures in G ∗ , i.e. Q ( G ∗ ) . This is due to the computational constraints for calculating the KL for DynGFN since even for sparse graphs of moderate size this is a combinatorial computation. Nonetheless, this approach allows us to directly compare the learned posteriors Q ( G ) to a ground truth discrete distribution over structure G to evaluate the effectiveness of Bayesian structure learning approaches.

B.9 Implementation Details

Our model is implemented in Pytorch and Pytorch Lightning and is available at https://github. com/lazaratan/dyn-gfn . Models were trained on a heterogeneous mix of HPC clusters for a total of ~ 1,000 GPU hours primarily on NVIDIA RTX8000 GPUs.

C Additional Details

C.1 Single-cell Biology and Gene Regulatory Network Inference

C.1.1 Gene Regulatory Networks and Cell Dynamics

One dynamical system of interest is that of cells. Cellular response to environmental stimuli or genetic perturbations can be modelled as a complex time-varying dynamical system [17, 1]. In general, dynamical system models are a useful tool for downstream scientific reasoning. In this work we are primarily interested in identifying the underlying cell dynamics from data. A reasonable model for cell dynamics is as a stochastic dynamical system with many, possibly unobserved, components. There are many data collection models for gaining insight into this system from single-cell RNA-sequencing data. We will primarily focus on RNA velocity type methods, where both x and an estimate of dx are available in each cell, but note that there are other assumptions to infer dynamics and regulation such as pseudotime-based methods [47, 1], and optimal transport methods [17, 48, 51, 22, 23]. After learning a possible explanatory regulation, this is used in downstream tasks, but the resulting conclusions drawn from these models are necessarily conditional on the inferred regulation. Motivated by gene regulatory networks, we explicitly model uncertainty over graphs which allows us to propagate the resulting uncertainty to downstream conclusions.

C.1.2 Learning Gene Regulatory Networks From Single-cell Data

Single-cell transcriptomics has an interesting property in that from a single measurement we can estimate both the current state x and the current velocity dx . Because mechanistically RNA undergoes a splicing process, we can measure the quantities of both the unspliced (early) and spliced (late) RNA in the cell. From these two quantities we can estimate the current RNA content for each gene and the current transcription rate. There exist many models for denoising and interpreting this data [27, 9, 45]. Furthermore, there exist more elaborate measurement techniques to extract more accurate velocity estimates [45]. The fact that we have an estimate of the current velocity is exceptionally useful for continuous time structure discovery because it allows us to avoid explicitly unrolling the dynamical system.

Learning the underlying causal structure from data is one of the open problems in biology. There are many works that attempt to learn the effect of a change in a gene, or the addition of a drug. These works often build models that directly predict the outcome of an intervention. This may be useful for certain applications, but often does not generalize well out of distribution. We would like to learn a model of the underlying instantaneous dynamics that give rise to effects at longer time scales. This approach has a number of advantages. (1) it is closer to the mechanistic model; it may be easier to learn a model of the instantaneous dynamics rather than the dynamics over long time scales (details in Appendix C.2). (2) One model can be trained and applied to data from many sources including RNA-velocity, Pseudotime, Single-cell time series, and steady state perturbational data. (3) The instantaneous graph may be significantly sparser (and therefore easier to learn) than the summary graph or the equilibrium graph.

C.1.3 Further Details on Experiments with Single-cell RNA-velocity Data

The process of Eukaryotic cell division can be divided into four well regulated stages based on the phenotype, Gap 1 ( G 1 ), Synthesis ( S ), Gap 2 ( G 2 ), and Mitosis ( M ). This process is a good starting point for GRN discovery as it is (1) relatively well understood, (2) deterministic, and (3) well studied with plentiful data. While there is an underlying control loop controlling the progression of the cell cycle, there are many other genes that also change during this cycle. To rediscover the true control process from data we must disentangle the true causal genes from the downstream correlated genes. This may become very difficult when we only observe dynamics at longer time scales. We hide a cell cycle regulator among two downstream genes that are highly correlated (Spearman ρ > 0 . 75 ) and test whether we can model the Bayesian posterior - namely that we are uncertain about which of the three genes (Cdc25A, Mcm2, or Mcm5) is the true causal parent of Cdk1.

C.2 Instantaneous Graph and Long-horizon Graph

The graph recovered depends on the time scale considered. We make a distinction between the conditional structure of the graph based on the time scale. Consider a system governed by dx ( t ) dt = f ( x ( t )) . We define the instantaneous graph as:

̸

Here, 1 ( z ) = 1 if z = 0 , otherwise 1 ( z ) = 0 . We can define the temporal summary of the system drift f as:

Then, we can define a long-horizon graph over long time-scales as:

If we are integrating the drift f over long time-scales, the long-horizon graph may be less sparse than the instantaneous graph. In cellular systems, this equates to observing cell dynamics over long time-scales, in turn observing increase quantities of correlations between variables of the system. Thus, trying to delineate the instantaneous dynamics from long time-scales may be difficult depending on the underlying system dynamics.

C.3 Further Discussion on Future Work

Although DynGFN is currently limited to smaller systems, we foresee approaches that would enable some degree of scaling DynGFN to larger systems. One approach is to leverage biological information of known gene-gene connections as a more informative prior for DynGFN. Currently, DynGFN learns a forward stochastic policy for Q ( G | D ) starting from an initialized state s 0 of all zeros. Instead, we can define ˜ s 0 using a prior of high confidence biological connections and sequentially add edges starting from this new initial state. This would reduce the number of possible structures DynGFN would need to search over, thus improving the potentially scalability of DynGFN. Another approach is to learn structure between sets of genes (variables) rather than single genes. Since GRNs are generally very sparse, it makes sense to group genes in sets. Consequently, we can then learn structure between these grouped genes, rather than just individual genes. In turn, DynGFN can explore/learn the structure over a smaller space while effectively capturing structure between a significantly larger set of genes. To group genes, we would use prior biological information, either form existing literature or expert domain knowledge.

In this work we exploit the use of a minimal prior, i.e. L 0 sparsity prior, for learning Bayesian dynamic structure between variables. In general, the aforementioned approaches for scaling DynGFN to larger systems involve the use of more informative priors on G . Although we mention two ways we foresee approaching this in the biological context of GRN inference, the general approach of using more informative priors can help scale DynGFN to larger systems across applications.

C.4 Broader Impacts

While it is important to acknowledge the potential risks of drawing incorrect scientific conclusions due to incorrect assumptions, our work embraces a Bayesian perspective for structure learning. A key component of our work is to account for uncertainty within our method, aiming to minimize the chances of incorrect conclusions. It is important to note that the accuracy of conclusions relies on applying the method in settings that align with the underlying assumptions, such as causal sufficiency and the use of dynamic observational data. By adhering to these guidelines, our approach holds promise for producing robust and reliable scientific outcomes.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this work, the authors propose an extension of GFlowNets to enable posterior sampling of cyclic graph structures and apply their approach to infer gene regulatory networks (GRN). While existing methods are able to infer cyclic graph structures or sample from Bayesian posteriors over DAGS, they cannot do both simultaneously. The authors formulate the problem as a dynamical systems identification problem which can be factored into the priors and model likelihood. The priors are learned using either linear or non-linear Hyper Networks while the model likelihood is estimated using the GFlowNets framework with detailed balance loss. The authors show their model is robust to several parameters including edge sparsity and time intervals between data points. Furthermore, they introduce several optimization ideas enabling a tractable search over admissible graph structures. Finally, they benchmark their approach against several baseline methods and apply their approach to an RNA velocity dataset to show their method is able to recover known and putative gene-gene interactions.\n\nSTRENGTHS:\nThe paper addresses an interesting/important question which has been the center of much research for decades: posterior sampling over directed graphs, this time with cycles and not DAGs. The authors appear knowledgable in this field and they take advantage of GFlowNets, building a method to assess posterior structures with cycles as they unfold them across time.\n\nWEAKNESSES:\nOur biggest grievance is that the paper is just not reader friendly for anyone who is not an expert in that specific line of work. \nAs noted the work seems technically very solid, one that offers specific new solution for the stated problem. Still, we as researchers not in this specific field, remain unsure after reading it of the overall innovation/significance of creating posterior sampling of structures over GFlowNets. We add more details below.\n\n(1)\nThe termination criteria of Algorithm 1 is not clear. Under what conditions does the transition probability become null? Although one of the main innovations of the paper is sampling from the posterior of GFlowNet structures, the authors don’t reference any standard Bayesian analysis. How are posterior samples summarized (for example, how is a point estimate of p obtained from posterior samples in Figure 3)? How and when does Algorithm 1 converge to the posterior? Is the approach sensitive to variable initializations? \n\n(2)\nAlthough the authors compare their approach to sensible baselines using sensible metrics, it is not clear the approach has a clear advantage over the original GFlowNet in this applied setting. For example, how does the MAP parameter estimate (or equivalent)  of this method compare to the original GFlowNet under the AUC metric (these are presumably comparable). Does GFlowNet produce different insight on gene-gene interactions in the RNA velocity dataset?\n\n(3)\nThe authors assume causal sufficiency as a criteria for model identifiability. Is this a reasonable assumption given the sparsity of single cell data and limitations in the scalability of the model? In other words, how can we be reasonably certain that all casual gene measurements are observed given the model effectively only functions on a gene space of size 20 or lower?\n\n(4)\nSection 5 on graph augmentation is very abstruse. Efforts could be made to improve readability/accessibility for an audience that doesn’t work directly in this field. The related Figure 2 is also very cryptic. \n\n(5)\nIt is not clear or obvious how the factorization into the “per-node” posterior in section 4.1 leads to the reduction in search space from 2^d^2 to d2^d. A brief explanation would be useful.\n\n(6)\nLine 154: “Previous work has shown GFlowNets are useful in settings with multi-modal posteriors.” This needs a citation.\n\n(7)\nThere is inconsistency in the notation in Figure 1 which leads to some confusion. Specifically, it appears the Q variable is overloaded. How is the graph prior (represented as an adjacency matrix) sampled? Is there a hyperprior that is used or are graph structures uniform? \n\n(8)\nThe patterns that the authors suggest that are evident in the Figure 3 heatmaps are not convincing. Perhaps it would be better to sort both heatmaps the same way instead of independent clustering.\n\nQUESTIONS:\nsee above.\n\nLIMITATIONS:\nsee above.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors tackle the well-known problem of Gene Regulatory Network (GRN) inference from time series data. Starting from an estimate of RNA velocity, GRN inference is framed as a causal discovery task, with two specificities: the inferred graphs may be cyclic, and a probabilistic distribution over candidate graphs is sought rather than a point estimate, so as to account for the high amounts of noise inherently present in biological data. The method is articulated around two assumptions:\n\n- a dynamic structural model $\\frac{dx_i(t)}{dt} = f_i(\\text{Pa}(x_i), \\epsilon_i)$ for each gene. $\\text{Pa}(x_i)$ is directly obtained from the regulatory graph $G$, and the structural causal model (SCM) $f_i$ is parametrized by $\\theta$\n- a factorisation of the joint distribution of the graph $G$, the structural causal model parameters $\\theta$, and the observed data $\\mathcal{D}$ as $p(G,\\theta,\\mathcal{D}) = p(\\mathcal{D}|G,\\theta)p(\\theta|G)p(G)$. The graph sampler $p(G)$ is modeled as a GFlowNets. Next, since the parameters of the SCM depend on the graph structure, they are computed by a HyperNetwork, i.e. $p(\\theta|G)$ is substituted to a neural network which takes as input a graph $G$ and outputs the structural equation model parameters $\\theta$: $p(\\theta|G) = \\delta(\\theta|G)$.\n \nNext, the authors evaluate the proposed method, DynGFN, on several synthetic examples as well as a real-world one involving single-cell RNA velocity data. The authors introduce other baselines by using different approaches than a GFlowNets for the graph sampler $p(G)$. This being said, DynGFN offers substantial improvements when it comes to jointly recovering the ground truth structure while characterizing the uncertainty around it.\n\nSTRENGTHS:\nThe big picture of the method is well explained. The work described here is likely to have a significant impact for the GRN inference community, as I believe that not many probabilistic methods are available in this field. I also believe that DynGFN could be used as a starting point for the more general problem of physical/biological dynamical system discovery, where one is not only concerned with the graph adjacency matrix but with providing a mechanistic description of the structural causal model $f_i$ [1,2]\n\n[1] Discovering governing equations from data by sparse identification of nonlinear dynamical systems, PNAS, 2016\n[2] Identification of dynamic mass-action biochemical reaction networks using sparse Bayesian methods, PLoS Comp. Biol., 2022\n\nWEAKNESSES:\n- The paper builds on GFlowNets for the graph sampler $P(G)$ and Hypernetworks to generate the structural causal model parameters $\\theta$ given a sampled graph $G$. I found it slightly difficult to understand everything $in$ $detail$ given that these tools are each introduced only in a small paragraph.\n\n- Even though the method focuses on Bayesian Inference, it would have been interesting to assess DynGFN's performances compared to SOTA methods for GRN inference from time series data such as DynGENIE3 [1], BINGO [2] or others. Lastly, these methods use time-resolved gene expression data, from which one can get an estimate of the RNA velocity using for instance finite differences, which may be a quite rough estimate for sparsely sampled measurements. Do I understand right that the RNA velocity data employed in this submission somehow constitutes a ``more principled'' way to estimate the RNA velocity?\n\n[1] dynGENIE3: dynamical GENIE3 for the inference of gene networks from time series expression data, Scientific Reports, 2018\n[2] Gene regulatory network inference from sparsely sampled noisy data, Nature Communications, 2020\n\nQUESTIONS:\n- Any insights on why $h$-DynGFN yields significantly better results than $\\ell$-DynGFN on the linear system presented in table 1?\n\n-  In the single cell experiment, where does the ground truth network (used to compute AUC / Bayes-SHD) come from? Could you explain a little more what does \"correlation over cell cycle time\" means in Section 6.3?\n\nI have read the rebuttal done by the authors.\n\nLIMITATIONS:\nThe authors mention as main limitation the difficulties to scale to large gene regulatory networks due to the combinatorial explosion. Causal sufficiency is also assumed here, which means that all relevant variables are observed. This assumption hardly ever holds in real-world biological problems. This issue is also mentioned by the authors in the Appendix.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAuthors introduced a principled methodology called DynGFN, which effectively identified cyclic structures and concurrently modeled the Bayesian posteriors over directed acyclic graphs (DAGs). Leveraging RNA velocity, the authors formulated a dynamic system that unveiled the underlying gene regulatory networks. DynGFN was meticulously designed with three modules, and its superior performance was demonstrated through synthetic experiments and real-world analysis.\n\nSTRENGTHS:\nThe real biological systems can be rarely formulated as DAGs and there are always feedback loops to make the system work. DynGFN was appropriately motivated with a real biological thinking.\n\nWEAKNESSES:\nThe overall manuscript was hard to follow as there was too much content included. There were still a lot of technical details that needed to be clarified. It would improve substantially if the manuscript can be carefully restructured.\n\nQUESTIONS:\nMajor concerns\n1. The synthetic experiments and real analysis did not explicitly present any results related to feedback loop identification, which serves as one of the motivations behind DynGFN. Could you elaborate on this point?\n2. How many epochs did it take to stabilize the training of l-DynGFN?\n3. Given the inference of RNA velocity involves multiple free parameters, how do these parameters impact the inference results obtained from DynGFN?\n4. In the synthetic experiments simulating dynamic systems, how were the ground truth regulatory networks constructed based on the simulated dynamic systems?\n5. Appendix B.3 mentions the construction of the validation and test sets to fine-tune the hyperparameters. Could you explain how these sets were constructed for the experiment?\n6. Regarding Table 4, could you clarify what the ground truth graph was used for calculating Bayes-SHD and AUC? Was it formulated based on external knowledge?\n7. In Figure 3, it is important to note that gene regulation is not solely determined by the direct interaction between corresponding proteins. In fact, it has been reported that \"CDK1 targets MCM2-7 complex...\" [1].\n\n[1] Enserink, Jorrit M., and Richard D. Kolodner. \"An overview of Cdk1-controlled targets and processes.\" Cell division 5.1 (2010): 1-41.\n\n\n\nMinor concerns\n1. Figure 1 lacks definitions for numerous parameters, which may cause confusion in understanding the illustration.\n\nLIMITATIONS:\nThe authors mentioned two limitations: 1) scalability issues concerning larger systems and 2) hyperparameter tuning. To address these challenges, the authors suggested employing more informative priors or grouping genes together as nodes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose to learn dynamical systems with cyclic dependencies. They factorise the generative model using a variant of the GFlowNet model, a HyperNetwork and MLPs.\n\nSTRENGTHS:\nPositives include:\n1. extending the formulation from DAGs to cyclic graphs\n2. introducing a per-node posterior formulation  to improve the computational complexity\n3. using GFlowNets to address the multimodality issues\n\nWEAKNESSES:\nNegatives include:\n1. due to the high density the paper is hard to understand: a lot is explained and detailed in the Appendix\n2. it is hard to gauge the quality of the results: even for the artificial case it is hard to assess intuitively if the approach solves the problem in a satisfactory way. Perhaps the authors should use simple systems with few nodes, with and without cycles, and show how well the true dependencies are recovered. \n3. it is unclear how important it is to being able to model cyclic structures: where are the cyclic dependencies in the examples shown in Fig. 2 and Fig. 3? The authors should show a clear example where failing to model the cyclic dependencies results in a poor fit.  \n4. it is unclear how to translate the declared advantage of being able to model multi-modal posteriors  in practical terms: how would we make use of this representation capacity when explaining the cause dependencies in a gene regulatory network?\n\nQUESTIONS:\nTo illustrate the advantages of the proposed approach:\n1.   use simple systems with few nodes, with and without cycles, and show how well the true dependencies are recovered. \n2. show a clear example where failing to model the cyclic dependencies results in a poor fit.  \n3. show how to use the multimodal representation capacity when explaining the cause dependencies in a gene regulatory network.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-24"}