SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We note that when evaluation on validation and test data for Bayes-SHD and AUC metrics, we hard threshold  $P_{\alpha_t}(G|Z)$ . We find that through training this the final  $\alpha_t$  is typically small enough in magnitude such that  $P_{\alpha_t}(G|Z)$  does not yield a full threshold of  $Z$ . To this end, we select large  $\alpha_t$  when computing the KL metric to mimic hard threshold behaviour as experienced during training. We use  $\alpha_t = 1 \times 10^8$ . In DynDiBS the parameter  $\gamma$  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};\psi)$  and  $h_\phi$  with MLP architectures.  $P_F(s_i|s_{i-1};\psi)$  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};\psi)$ . 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};\psi)$  MLP architecture. We use a uniform backward policy for  $P_B(s_{i-1}|s_i;\psi)$ . To parameterize  $h_\phi$ , we use a 3 layer MLP with hidden layer dimensions of
Table 7: Hyper-parameters for DynGFN. We define learning rate as  $\epsilon$ ,  $m_{train}$  as number of training samples, and  $m_{eval}$  the number of evaluation samples.
Linear System
Model	ε	λ0	T	mtrain	meval
l-DynGFN	0.0001	100	0.01	1024	5000
h-DynGFN	0.00001	100	0.005	256	3000
Non-linear System
Model	ε	λ0	T	mtrain	meval
l-DynGFN	0.0001	150	0.01	1024	5000
h-DynGFN	0.00001	150	0.005	256	3000
Cellular System
Model	ε	λ0	T	mtrain	meval
l-DynGFN	0.00005	45	0.01	1024	1000
h-DynGFN	0.0001	10	0.1	1024	1000
$\{64, 64, 64\}$  and exponential linear unit activations (ELU). We consider two parametrizations for  $f_{\theta}$ : single linear parameters, i.e.  $dx = \theta x$ , and a single hidden layer neural network  $dx = f_{\theta}(x)$ . We use these parametrizations to model linear and non-linear node-wise parent-child structural models, where  $x \in \mathbb{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,  $\lambda_0$  (sparsity encouragement for identified graphs), a temperature parameter  $T$  that scales the magnitude of the reward likelihood (e.g.  $\frac{1}{T^2}\mathrm{MSE}(dx,\widehat{dx})$ ), learning rate  $\epsilon$ , 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  $\ell$ -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  $\pi(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};\psi)$  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};\psi)$  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 \leq c \leq 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.
Figure 5: Validation curves throughout training of  $l$ -DynGFN on a linear system. Mean squared error (left) and Bayes-SHD (right).
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.  $\ell$ -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  $\ell$ -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  $\ell$ -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  $\sim 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  $\rho > 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  $\frac{dx(t)}{dt} = f(x(t))$ . We define the instantaneous graph as:
$$
g _ {i, j} := \cup_ {x} \mathbf {1} \left(\frac {d f (x) _ {j}}{d x _ {i}}\right), \quad i = 1, \dots , d, j = 1, \dots , d. \tag {6}
$$
Here,  $\mathbf{1}(z) = 1$  if  $z \neq 0$ , otherwise  $\mathbf{1}(z) = 0$ . We can define the temporal summary of the system drift  $f$  as:
$$
F (T, x) = x _ {0} + \int_ {0} ^ {T} f (x (t)) d t. \tag {7}
$$
Then, we can define a long-horizon graph over long time-scales as:
$$
G _ {i, j} := \cup_ {x} \mathbf {1} \left(\frac {\partial F (T , x) _ {j}}{\partial x _ {i}}\right), \quad i = 1, \dots , d, j = 1, \dots , d. \tag {8}
$$
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 DynGNF is currently limited to smaller systems, we foresee approaches that would enable some degree of scaling DynGNF to larger systems. One approach is to leverage biological information of known gene-gene connections as a more informative prior for DynGNF. Currently, DynGNF learns a forward stochastic policy for  $Q(G|D)$  starting from an initialized state  $s_0$  of all zeros. Instead, we can define  $\tilde{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 DynGNF would need to search over, thus improving the potentially scalability of DynGNF. 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, DynGNF 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.