q-bio.MNDec 11, 2025
SaveTracking large chemical reaction networks and rare events by neural networks
Organizations: Institute of Data Science, University of Hong Kong, Hong Kong · Department of Systems Science, Faculty of Arts and Sciences, Beijing Normal University, Zhuhai 519087, China · School of Physics, University of Electronic Science and Technology of China, Chengdu 611731, China · School of Physical Science and Technology, Beijing University of Posts and Telecommunications, Beijing 102206, China · Institute of Theoretical Physics, Chinese Academy of Sciences, Beijing 100190, China · School of Cyber Science and Technology, University of Science and Technology of China, Hefei 230027, China · School of Fundamental Physics and Mathematical Sciences, Hangzhou Institute for Advanced Study, UCAS, Hangzhou 310024, China · Institute of Fundamental and Frontier Sciences, University of Electronic Science and Technology of China, Chengdu 611731, China · Non-classical Information Science Basic Discipline Research Center of Sichuan Province, University of Electronic Science and Technology of China, Chengdu 611731, China
Abstract
Chemical reaction networks are widely used to model stochastic dynamics in biophysical processes. Their exact probabilistic description via the chemical master equation becomes intractable for high-dimensional systems due to the exponential growth of the state space. This difficulty is further exacerbated when metastable transitions and rare events dominate the dynamics. Here, we present an advanced neural-network approach that effectively tackles large reaction networks and rare-event dynamics within a unified framework. Building on autoregressive representations of the full joint probability distribution, we achieve efficient training through second-order optimization methods, including natural gradient descent and the time-dependent variational principle. We leverage enhanced-sampling strategies that enable accurate resolution of low-probability regions governing rare transitions. We demonstrate the method on a multistable genetic toggle switch and a mitogen-activated protein kinase (MAPK) cascade consisting of 16 species and 35 reactions with dense coupling. We further quantify rare events in spatially extended Schlögl models in both one- and two-dimensional lattices, reaching system sizes with extremely large state spaces up to . These results highlight the potential of neural modeling as a flexible and scalable framework for studying high-dimensional chemical reaction networks and rare events.
Figures & tables
Figure 1: The present work efficiently tracks large reaction networks and rare events. (a) The chemical reaction system follows the CME, whose exponentially large state space makes the exact probability distribution intractable. (b) The variational autoregressive network factorizes the joint distribution into conditional probabilities. We use NADE-based architectures to improve training efficiency over sequential RNNs. (c) Our advances include both reducing computational cost to deal with large networks and quantifying rare-event statistics. (d) The VAN is trained by minimizing the KL-divergence between time steps. Moving from stochastic gradient descent (SGD), we now use natural gradient (NG), which improves convergence by rescaling updates with the inverse Fisher information matrix and the time-dependent variational principle (TDVP) with the projected temporal evolution. In the genetic toggle switch model [ 24 , 25 ] , NG and TDVP achieve significant speedup over SGD. (e) To capture rare events, we augment the training procedure with enhanced-sampling strategies, including mixture sampling with configurations drawn from a uniform distribution , diffusive sampling based on a kernel , and sampling using an exponent-based modification. In the Schlögl model, these sampling schemes markedly reduce the Hellinger distance to the baseline from expensive Gillespie simulations.
| Example | Optimize by | Species | Reactions | Depth | Width | Batch | Epochs per step | Time (hr) | ||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Toggle switch | SGD | 4 | 8 | 80 | 8001 | 0.005 | 1 | 16 | 2000 | 100 | 0.005 | 3.234 |
| NG | 4 | 8 | 80 | 8001 | 0.005 | 1 | 16 | 2000 | 5 | 0.5 | 0.806 | |
| TDVP | 4 | 8 | 80 | 8001 | 0.005 | 1 | 16 | 2000 | 1 | 1.0 | 0.185 | |
| MAPK cascade | NG | 16 | 35 | 10 | 0.01 | 1 | 8 | 2000 | 5 | 0.5 | 127.584 | |
| 1D Schlögl (8 Sites) | NG | 8 | 46 | 85 | 1 | 16 | 5000 | 5 | 0.8 | 79.433 | ||
| 2D Schlögl (2x4 Sites) | NG | 8 | 52 | 85 | 1 | 16 | 5000 | 5 | 0.8 | 102.82 |
Table 1: The models solved by NNCME-2 and the corresponding computational cost under the chosen hyperparameters. The time-step length is measured in units of the inverse reaction rates, and the physical time is . The upper count limit is . Depth and width refer to the NADE architecture, and is the learning rate for optimization algorithms. NG denotes natural gradient and TDVP denotes the time-dependent variational principle. All computations were performed on a single Tesla V100 GPU.
Figure 2: NNCME-2 achieves at least a 5-fold acceleration while maintaining high accuracy, as demonstrated in the genetic toggle switch. (a) A schematic of the chemical reactions. (b) The time evolution of the average counts for the genes and proteins, comparing the VAN (dots), Gillespie simulation (solid line), FSP results from SSIT (dashed), and ACME results (dash-dotted). (c) Comparison of the standard deviations obtained from VAN, FSP (SSIT), and ACME against Gillespie at time points . The Gillespie simulation has trajectories. (d) Joint distribution of the two proteins from VAN at time points , showing four stable states in the long-time limit. The initial distribution assigns equal probability to the bound and unbound promoter states for each gene, with zero initial protein counts. The results here are from the NADE architecture with natural gradient optimization. Parameters are , , , and , with hyperparameters in Table I. (e) Training performance of different optimizers on the NADE architecture. (f) Comparison of computational cost. NADE reduces overall training time compared to RNN, and both NG and TDVP significantly accelerate training.
Figure 3: NNCME-2 handles large reaction networks efficiently. (a) A schematic of a representative large reaction network (MAPK cascade, 16 species). Bidirectional arrows indicate reversible reactions, and the detailed reactions and rates are in Appendix G.2 . (b) The left panel displays the time evolution of the mean counts of the principal species (MKP3, K, and Kpp), comparing Gillespie simulations (lines), VAN (circles), and ACME (squares). The two panels on the right compare the means and standard deviations obtained from VAN and ACME with Gillespie at time points . (c) Marginal distributions of MKP3, K, and Kpp at and , obtained from the three methods. The inset reports the Hellinger distance between the Gillespie and VAN distributions. (d) Joint distributions of K and Kpp from the VAN at and . The system evolves from the initial distribution to the bimodal state. The Gillespie results are based on trajectories. The hyperparameters of the VAN are in Table I.
Figure 4: NNCME-2 captures rare events in the one-dimensional spatially extended Schlögl model (6 sites). (a) A schematic of the reaction-diffusion system, where each site hosts a well-mixed Schlögl model with fixed species and , and stochastic species diffusing between neighboring sites with rate . (b) Comparison between the vanilla VAN (top) and the one with enhanced sampling (bottom). The results obtained from the VAN (dots) closely match those from Gillespie simulations (lines) for the time evolution of the mean counts at sites , , , and (left), standard deviations at time points (middle), and the marginal distribution of at with the Hellinger distance (right). The inset highlights the low-probability regime ( , probability ). (c) Joint distributions of and from the VAN with enhanced sampling, with probabilities shown by color in a logarithmic scale. White boxes indicate the transition of the probability peak. The initial distribution is a delta peak at . The Gillespie results are based on trajectories. The hyperparameters of the VAN are in Table I. Mixture ES is used in this example, with exploratory samples comprising 15% of each batch.
Figure 5: The scaling of rare probabilities and computational time for the one-dimensional Schlögl model. (a) Probability mass in across lattice size , from vanilla Gillespie simulations (red), NNCME-2 (blue), and Gillespie (SSIT [ 82 ] ) with an optimized SSA implementation (orange). (b) Total computation time (log scale), with a linear-scale view shown in the inset. Both Gillespie implementations used trajectories. The original Gillespie simulations were run on an Intel Xeon Gold 6248 CPU node, SSIT calculations were performed in MATLAB on a single-processor MacBook Pro, and NNCME-2 computations were performed on an NVIDIA Tesla V100 GPU node. Enhanced sampling is applied to all lattice sizes using the mixture ES scheme, with 15% exploratory samples per training batch.
Figure 6: Absolute error of the rare-event probability for the one-dimensional Schlögl model with as a function of computation time. Results from the vanilla Gillespie implementation, Gillespie (SSIT), and NNCME-2 are compared with the extrapolated reference . The Gillespie and SSIT curves with shaded regions show the variation over 20 resampled estimates. The NNCME-2 points correspond to different training settings. The inset shows the same results on linear axes. The computational platforms follow those described for Fig. 5 .
Figure 7: NNCME-2 captures rare events in the two-dimensional spatially extended Schlögl model ( grid). (a) A schematic of the 2d well-mixed Schlögl model. The parameters are the same as Fig. 4 . (b) Comparison between the vanilla VAN (top) and the enhanced VAN (bottom), using mixture enhanced sampling with 15% exploratory samples per training batch. The left and middle panels show the time evolution of the mean counts and the corresponding standard deviations of sites - , respectively. The right panels display the marginal distribution of at , with the Hellinger distance between VAN and Gillespie results indicated. (c) Joint distributions of and from the enhanced VAN at , shown on a logarithmic probability scale.
Figure 8: The scaling of rare probabilities and computational time for the two-dimensional spatially extended Schlögl model. (a) Probability mass obtained using the vanilla Gillespie implementation, NNCME-2, and Gillespie (SSIT). Both Gillespie implementations use trajectories. (b) The total computation times of the three methods on a logarithmic scale, with the inset showing the same data on a linear scale. The computational platforms follow those described for Fig. 5 . All NNCME-2 results use mixture enhanced sampling with 15% exploratory samples per training batch.
Appendix figures & tables15 assets
Supplementary material from the paper’s appendix.
Appendix
| Symbol | Description |
|---|---|
| Learning rate used in optimization. | |
| Discrete time step of CME propagation. | |
| Number of discrete time steps in the evolution. | |
| Physical time, . | |
| Number of samples drawn per training step. | |
| Total number of trainable parameters. |
Table 2: Mathematical symbols used throughout the manuscript.
Figure 9: Effect of network width on the learning dynamics and accuracy of the toggle-switch system. (a) Evolution of the average loss during training of the initial distribution ( ) for different network widths. (b) Evolution of the Hellinger distance over time between the VAN-generated distribution and the Gillespie reference for different network widths. In these experiments, the VAN model employs NADE architecture ( and ). Training utilizes NG and reverse KL, with , = 0.5 at and , = 0.5 for .
Figure 10: Evaluation of KL-divergence loss functions in learning the dynamics of the toggle switch system. (a) Mean of training loss at the initial time step . (b) The accuracy of the VAN-generated distribution over time, quantified by the Hellinger distance from the benchmark Gillespie trajectories. (c) Joint probability heatmap of and for the benchmark Gillespie simulation at (left) and (right). (d) The time evolution of the average count for the genes and proteins, from the VAN (lines) and the Gillespie simulation (dotted line) (left). Standard deviations comparisons (middle). The heatmaps depict the joint distributions of and at time points and (right). In these experiments, the VAN model employs NADE architecture ( , and ). Training utilizes SGD, with , = 0.005 at and , = 0.005 for .
| Loss Type | Loss |
|---|---|
| ForwardKL | |
| ReverseKL | |
| ReverseKL2 |
Table 3: Summary of loss functions.
Figure 11: Validation of the Hellinger distance as a representative evaluation metric for the toggle-switch system. The VAN is trained using different loss functions based on variants of the KL divergence, and model performance is evaluated using multiple distance measures, including Total Variation (TV), Kullback–Leibler (KL), Jensen–Shannon (JS), and Hellinger distance (HD). The reference distribution is based on trajectories generated via the Gillespie algorithm. In these experiments, the VAN model employs NADE architecture ( and ). Training utilizes SGD, with , = 0.005 at and , = 0.005 for .
Figure 12: Effect of the learning rate on NG. (a) Mean of loss at for different values of . (b) Standard deviation of the loss at for different values of . The red dashed line marks a threshold of , and the colored vertical lines indicate the first epoch at which the threshold is reached for the corresponding . (c) Mean of loss at . (d) The accuracy of the VAN-generated distribution for various over time, quantified by the Hellinger distance from the Gillespie trajectories. In these experiments, the VAN model employs NADE architecture ( ) and is optimized using the reverse KL-divergence. Training utilizes NG, with for and for , and .
Figure 13: Comparison of the training efficiency and convergence behavior of different optimization algorithms for the VAN in the toggle-switch system. (a) Means of training loss at the initial time step for SGD (green) and NG (blue) with 1000 epochs. (b) The accuracy of the VAN-generated distribution using three optimization methods over time, quantified by the Hellinger distance from the benchmark Gillespie trajectories. (c) Joint probability heatmap of and for the benchmark Gillespie simulation at (left) and (right). (d) The left panel shows the time evolution of the mean counts of the genes and proteins obtained from the VAN (lines) and from Gillespie simulations (dotted lines). The middle panel presents a comparison of the corresponding standard deviations. The right panel displays heatmaps of the joint distributions of and at time points and . All experiments were trained with reverse KL-divergence, the optimizer parameters: learning rate = 0.005, epoch = 100 at , epoch = 2000 at for SGD; = 0.5, epoch = 5 for , epoch = 50 for for NG, in NADE architecture (width , ).
Figure 14: Comparative analysis of training loss for NG and SGD. (a) Means of training loss at the initial time step . (b) Means of training loss at an early training stage . (c) Standard deviations of training loss at the initial time step , the horizontal dotted lines indicate convergence threshold ( ). (d) Standard deviations of training loss at an early training stage , the convergence threshold is set as . Red arrows quantify the epoch reduction ratio. All experiments were trained with reverse KL-divergence, the optimizer parameters: learning rate = 0.005, epoch = 100 at , epoch = 2000 at for SGD; = 0.5, epoch = 5 for , epoch = 50 for for NG, in NADE architecture (width , ).
| Algorithm | Parameter update |
|---|---|
| Stochastic Gradient Descent | |
| Natural Gradient | |
| Time-Dependent Variational Principle |
Table 4: Comparison of optimization algorithms and their update mechanisms. SGD follows the negative gradient direction scaled by the learning rate ; NG incorporates the FIM for preconditioning; TDVP employs a force term derived from the system’s Hamiltonian. All methods iteratively minimize the reverse KL-divergence loss .
Figure 15: Comparison of probability distributions for the genetic toggle-switch model obtained from Gillespie simulations, VAN, SSIT-FSP, and ACME. (a) Distribution of the promoter state at . (b) Marginal distribution of protein at . (c) Joint distributions of and at and , shown on a logarithmic probability scale.
Figure 16: Joint distributions of K and Kpp for the MAPK cascade at , , and , obtained from Gillespie simulations, ACME, and VAN. The distribution evolves from the delta-initialized state to a bimodal long-time distribution, with close agreement among the three methods.
Figure 17: Effect of hyperparameter on diffusive ES for the Schlögl (2 sites) model. (a) The time evolution of the average count for with different from the VAN (dot) and the Gillespie simulation (gray line). (b) The accuracy of the VAN-generated distribution over time, quantified by the Hellinger distance from the benchmark Gillespie trajectories. (c) Marginal distribution of species from the VAN at for different values (colored) and compared with the Gillespie simulation (gray). The learning parameters for the system are fixed at and . All experiments are performed using the NADE architecture with and . The Gillespie simulation has trajectories. In these experiments, the VAN model employs NADE architecture ( , , and ). Training utilizes reverse KL, with , = 0.5 at using NG and , =1 for using TDVP.
Figure 18: Effect of hyperparameter on ES for the Schlögl (2 sites) model. Smaller values of correspond to stronger exploration in the sampling process. The layout of the figure and the hyperparameters of VAN are the same as those of Fig. 17 .
Figure 19: Results of four ES strategies- Vanilla , Mixture ES, Diffusive ES and ES used in training VAN on the Schlögl (2 sites) system. (a) Time evolution of the means of species estimated by each method. (b) Comparison of the means of between VAN-based estimates and Gillespie simulation. (c) Comparison of the standard deviation of between VAN-based estimates and Gillespie simulation. (d) Marginal distributions of species and at (first and second rows), and joint probability distribution heatmaps of and at (third row). The Gillespie simulation has trajectories. In these experiments, the VAN model employs NADE architecture ( , , and ). Training utilizes reverse KL, with , = 0.5 at using NG and , =1 for using TDVP.
Figure 20: Comparison of computational efficiency and precision between PyTorch and JAX frameworks, framework implementations are distinguished by color (red: PyTorch, blue: JAX) and line style (solid: TDVP, dashed: NG algorithm). Experiments were conducted using the NADE net ( ). Performance metrics were evaluated through reverse KL-divergence measurements. Computational efficiency was evaluated through Hellinger distance measurements relative to Gillespie simulated trajectories. All experiments were executed on an NVIDIA Tesla V100 GPU under identical computational constraints.