Gradient estimators for parameter inference in discrete stochastic kinetic models
Authors: Ludwig Burger, Annalena Kofler, Lukas Heinrich, Ulrich Gerland
Organizations: Physics of Complex Biosystems, School of Natural Sciences, James-Franck-Straße 1, 85748 Garching, Germany · Max Planck Institute for Intelligent Systems, Max-Planck-Ring 4, 72076 Tübingen, Germany · Max Planck Institute for Gravitational Physics (Albert Einstein Institute), Am Mühlenberg 1, 14476 Potsdam, Germany · Data Science in Physics, School of Natural Sciences, James-Franck-Straße 1, 85748 Garching, Germany · Munich Center for Machine Learning (MCML), Munich, Germany
Stochastic kinetic models are ubiquitous in physics, yet inferring their parameters from experimental data remains challenging. For deterministic models, parameter inference often relies on gradients, which can be obtained efficiently through automatic differentiation (AD). However, AD cannot be applied directly to the Gillespie stochastic simulation algorithm (SSA), since sampling from a discrete set of reactions introduces non-differentiable operations. In this work, we adopt three gradient estimators from machine learning for the Gillespie SSA: the Gumbel-Softmax Straight-Through (GS-ST) estimator, the Score Function estimator, and the Alternative Path estimator. We use the estimators to evaluate gradients of steady-state and time-dependent observables, and compare their performance in representative biophysical systems with relaxation dynamics (bimolecular association) and oscillatory dynamics (repressilator). We find that the GS-ST estimator generally yields well-behaved gradient estimates, but exhibits diverging variance in challenging parameter regimes, which can cause parameter inference to fail. In these cases, other estimators provide more robust, lower variance gradients. Our results demonstrate that gradient-based parameter inference can be effectively combined with the Gillespie SSA, with different estimators offering complementary advantages.
Figures & tables
Figure 1: Bimolecular association model. (a) Starting from an initial configuration without bimolecular complexes, the system relaxes toward a steady state with a non-zero number of A–B complexes. The number of complexes in steady-state is governed by the dissociation rate k , the total numbers of molecules A and B, and the system volume (here NAtot=NBtot=200 , V=20 , and kref=5 ). Solid lines and shaded areas show the mean and standard deviation of 104 trajectories. (b) The steady-state distribution of A–B complexes obtained after s=500 Gillespie steps (orange bars) agrees well with the distribution derived from the exact solution of the Chemical Master Equation (dashed line, see Appendix A ). The histogram is obtained from 104 samples, with each bin weighted by the lifetime of the corresponding state to obtain an unbiased steady-state distribution (Eq. 1 ). (c) The loss function, defined as the relative deviation between the average steady-state number of complexes at k and at kref , exhibits a clear minimum at kref . (d–f) The gradient of the loss function, dL/dlogk , evaluated using the GS-ST method (d), the SF method (e), and the AP method (f), agrees well with the analytical reference (dashed lines). Data points in panels (c–f) show the mean and uncertainties obtained by averaging over 10 independent batches, each consisting of 100 stochastic trajectories.
Figure 2: Variance scaling of steady-state gradient estimators in the bimolecular association model. (a) The relative variance of the GS-ST estimator increases sharply as the temperature τ decreases below a critical value. (b) GS-ST variance scales differently with the number of Gillespie steps s depending on τ : it converges for large τ , and grows exponentially for small τ . The dashed line shows the analytical estimate of the asymptotic variance for large τ (Supplementary Material Section ). (c) The Lyapunov exponent indicates parameter regimes with convergent (negative exponent) or divergent variance (positive exponent). The temperature at which the Lyapunov exponent changes sign depends on the dissociation rate k . (d) The minimal temperature τ∗ required to obtain non-diverging variance depends on the parameter regime of k . (e) The variance of the SF and AP gradient estimates grows linearly with s . The dashed line for SF shows the analytical approximation of the variance (Supplementary Material Section ). (f) For SF, the linear variance scaling of the estimated gradient results from the linear variance scaling of the score function; the dashed line shows the analytical estimate. Data points show the variance and its uncertainty obtained by computing the variance across 104 stochastic trajectories and repeating this procedure ten times; in panels (a–b), the error bars indicate the 10th–90th percentile interval instead of the standard deviation. Unless stated otherwise, k=kref=5 , NAtot=NBtot=200 , and V=20 .
Figure 3: Illustration of the Gillespie SSA and gradient estimators for time-dependent observables. (a) The Gillespie SSA generates stochastic trajectories by sampling a molecule update ΔNs and a time increment Δts at each step s based on a previous state and input parameters θ . It is not directly differentiable with respect to θ because both Δts and ΔNs are sampled from parameter-dependent probability distributions, and ΔNs is a discrete random variable. (b) In the GS-ST estimator, gradients are obtained by reparameterizing the sampling of Δts and ΔNs using inverse transform sampling for the time and the Gumbel-max trick for the molecule update. For gradient evaluation, the discrete molecule update is approximated by the continuous, differentiable softmax function, enabling automatic differentiation. (c) In the SF estimator, gradients are obtained by evaluating the product of the simulation output and the score function. At each Gillespie step, score contributions from the molecule and the time update are computed, and accumulated along the trajectory. (d) In the AP estimator, gradients are obtained by evaluating the weighted difference of an alternative and a primal path. The alternative path ΔN~s can either be updated or re-sampled in every Gillespie step. The gradient contribution of the time update Δts is included through a score term, analogous to the SF estimator.
Figure 4: Bimolecular association model evaluated at fixed time. (a) Starting from an initial configuration without bimolecular complexes, the system reaches a nonzero number of complexes by the evaluation time tfinal=0.1 . Solid lines and shaded areas show the mean and standard deviation of 104 trajectories. (b–d) Gradient of the loss function evaluated at tfinal=0.1 using the GS-ST (b), SF (c), and AP (d) estimators. For GS-ST, τ=0.03 and τtime=2.5⋅10−5 . The estimated gradients agree with the analytical reference (dashed lines) only when the contribution of the waiting times to the gradient is explicitly taken into account (with time contribution (TC), dark data points); otherwise the gradient is overestimated (without TC, light data points). Data points show mean values and uncertainties by averaging over ten batches, each consisting of 104 trajectories. In all panels, kref=5 , NAtot=NBtot=200 , and V=20 .
Figure 5: Variance scaling of time-dependent gradient estimators in the bimolecular association model. (a) For small τ , the relative gradient variance of the GS-ST estimator increases exponentially with time, irrespective of whether the waiting-time contribution is accounted for (with TC) or not (without TC). (b) For large τ , including the waiting-time contribution increases the GS-ST variance relative to the time-independent case; the dashed line shows the analytical approximation of the time-independent asymptotic variance (Supplementary Material Section ). (c) The variance of the time-dependent GS-ST estimator scales quadratically with t for large τ . (d) The variance of the SF estimator scales linearly with t (dark green), analogously to the time-independent case (light green). Data points show the variance and its uncertainty obtained by computing the variance across 104 stochastic trajectories and repeating this procedure ten times; in panels (a–c), the error bars indicate the 10th–90th percentile interval instead of the standard deviation. In all panels, τtime=2.5⋅10−5 , kref=5 , NAtot=NBtot=200 , and V=20 .
Figure 6: Parameter inference in the repressilator model. (a) Three protein species are produced and degraded, and mutually repress each other’s production. (b) In certain parameter regimes, the system exhibits self-sustained oscillations under intrinsic noise. Mean and standard deviation are obtained from 104 Gillespie trajectories with kp=100 , Kd=10 , h=3 , and V=1 . (c–h) The production rate constant kp and dissociation constant Kd are inferred from trajectory data by minimizing the loss function ( h=3 , V=1 fixed). Results are shown for 50 optimization runs, corresponding to different reference parameter sets and different initializations of the optimization. Panels (c–d), (e–f), and (g–h) compare inferred and reference parameter values obtained using GS-ST, SF, and AP, respectively. In each optimization step, the gradient of the loss is estimated from 1000 stochastic trajectories; for GS-ST, we use τ=0.3 and τtime=0.05 .
Figure 7: Divergent gradient variance explains failures of parameter inference with the GS-ST estimator. (a–b) For the three parameter sets where GS-ST fails to converge to the true optimum, the gradients with respect to kp and Kd exhibit high variance. (c) The GS-ST gradient variance depends sensitively on the system parameters; in particular, the variance is larger for stronger binding affinity (lower Kd ). In panels (a–c), the variances are evaluated at the parameter values used to initialize the optimization. (d) For GS-ST and SF, the variance of the loss gradients increases sharply with binding affinity, although the dependence on Kd is weaker for SF. (e) The variance of the gradient of the change in molecule count in a single Gillespie step ΔN increases as Kd decreases, with comparable sensitivity to Kd for GS-ST and SF (Supplementary Fig. ). (f) For GS-ST, the variance of the gradients of the molecule counts Ns diverges at low Kd (dark blue), but increases approximately linearly with s for high Kd (light blue). By contrast, SF does not exhibit divergent variance even at strong binding affinity and instead scales approximately linearly with s for sufficiently large s (dark green). For panels (d–f), variance data are obtained at fixed kref=(200,6.67) and kp=200 , while varying Kd . In all panels, variances of the loss gradients are computed by averaging gradients over 103 independent trajectories and evaluating the variance across 50 repetitions. Variances of molecule counts are computed by first evaluating the variance across 104 trajectories and then averaging over 100 repetitions.
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: Analysis of the chemical rate equations of the repressilator model. (a–b) The mean-field repressilator model shows self-sustained oscillations with finite amplitude and period for Hill coefficient h>2 . The amplitude and period are largely insensitive to the choice of Hill coefficient, but depend on the effective production rate keff . (c) The dimensionless amplitude scales approximately linearly with keff . (d) The dimensionless period is proportional to lnkeff .
Figure 9: Temperature-dependent performance of GS-ST for parameter inference in the repressilator model. (a–b) Parameter inference via GS-ST at τ=0.75 succeeds for most parameter sets, including several that failed at τ=0.30 (highlighted by colored boxes), but a new subset of runs fails to converge. (c–d) For one of the newly non-converging parameter sets (black boxes in panels a–b), the GS-ST gradient closely matches the reference gradient from the unbiased SF estimator at low temperature (dark blue), but deviates systematically at high temperature (light blue), indicating strong gradient bias. (e–f) The gradient bias causes the optimization to converge to an incorrect optimum rather than the true parameters. In (c–f), lines and shaded regions show the mean and uncertainty across 20 independent optimization runs; in each run, gradients are estimated from 1000 stochastic trajectories at each optimization step.
Exact stochastic simulation of continuous-time Markov chains (CTMCs) is essential when discreteness and noise drive system behavior, but the hard categorical event selection in Gillespie-type algorithms blocks gradient-based learning. We eliminate this constraint by decoupling forward simulation from backward differentiation, with hard categorical sampling generating exact trajectories and gradients propagating through a continuous massively-parallel Gumbel-Softmax straight-through surrogate. Our approach enables accurate optimization at parameter scales over four orders of magnitude beyond existing simulators. We validate for accuracy, scalability, and reliability on a reversible dimerization model (0.09% error), a genetic oscillator (1.2% error), a 203,796-parameter gene regulatory network achieving 98.4% MNIST accuracy (a prototypical deep-learning multilayer perceptron benchmark), and experimental patch-clamp recordings of ion channel gating (R^2 = 0.987) in the single-channel regime. Our GPU implementation delivers 1.9 billion steps per second, matching the scale of non-differentiable simulators. By making exact stochastic simulation massively parallel and autodiff-compatible, our results enable high-dimensional parameter inference and inverse design across systems biology, chemical kinetics, physics, and related CTMC-governed domains.
Jose M. G. Vilar, Leonor Saiz
Biofisika Institute (CSIC, UPV/EHU), University of the Basque Country (UPV/EHU), P.O. Box 644, 48080 Bilbao, Spain · IKERBASQUE, Basque Foundation for Science, 48011 Bilbao, Spain · Department of Biomedical Engineering, University of California, 451 E. Health Sciences Drive, Davis, CA 95616, USA
Continuous-time Markov chains (CTMCs) provide the backbone for modeling discrete stochastic dynamics across applied, physical, and biological sciences. Their integration with modern gradient-based machine learning, however, is limited by the hard categorical event selection intrinsic to Gillespie-type simulation algorithms. We exploit the affine state update to obtain the exact one-step conditional-mean sensitivity by differentiating normalized reaction propensities. We pair this backward rule with exact forward trajectories to define the propensity straight-through (PST) estimator. At the trajectory level, we show that one-step sensitivities composed across events can depart from the exact multistep sensitivity. We derive the resulting per-step discrepancy in closed form and prove that it vanishes identically for affine downstream dependence. PST matches the accuracy of Gumbel-Softmax straight-through across all benchmarks: reversible dimerization (0.06% error), a genetic oscillator (1.7% error), a 50-task repressilator suite (0.17% median error), and patch-clamp ion-channel recordings (R2 = 0.988). Under matched settings, PST converges 3.0-fold faster on the oscillator and 2.1-fold faster on the ion channel. At deep-learning scale, PST trains a 203,796-parameter stochastic reaction network with hard sampling, reaching 98.22% MNIST digit classification accuracy. By differentiating an exact conditional mean rather than a relaxed sample, PST offers a temperature- and Gumbel-free path to scalable gradient-based learning through exact stochastic trajectories.
Jose M. G. Vilar, Leonor Saiz
Biofisika Institute (CSIC, UPV/EHU), University of the Basque Country, P.O. Box 644, 48080 Bilbao, Spain · IKERBASQUE, Basque Foundation for Science, 48011 Bilbao, Spain · Department of Biomedical Engineering, University of California, 451 East Health Sciences Drive, Davis, CA 95616, USA
Gradient estimation -- the task of computing the gradient of the expected value of a probabilistic program -- has diverse applications in scientific computing, but is notoriously difficult because of issues such as high-dimensional integration, discrete random choices, and complex stochastic dependencies. This article introduces gradient inference, a new approach to developing sound and efficient gradient estimators for probabilistic programs. Gradient inference rests on a formal reduction from a gradient estimation problem to a closely related probabilistic inference problem, whose solution can be differentiated to obtain a gradient estimator. This inference problem is obtained by applying two powerful statistical operations -- coupling and factorization -- to the input probabilistic program. Our reduction lets us leverage the rich toolkit of probabilistic inference algorithms to design novel gradient estimators that extend and improve upon existing methods. We introduce GradInf, a probabilistic programming system that facilitates the sound and automated implementation of gradient inference. GradInf is centered around programmable source-to-source transformations for coupling and factorizing higher-order probabilistic programs, whose soundness is proven in terms of a denotational semantics. Key to our development is the use of information-flow typing to allow random choices in a probabilistic program to be factored out and partially evaluated, which improves our ability to deploy sophisticated probabilistic inference algorithms. The resulting system offers practitioners a principled framework for designing gradient estimators. We apply GradInf to several challenging case studies, showing that it can express prominent gradient estimators from the literature and enables the construction of new state-of-the-art estimators that outperform the best existing baselines.
Gaurav Arya, Mathieu Huot, Moritz Schauer +2
Carnegie Mellon University, Pittsburgh, USA · Massachusetts Institute of Technology, Cambridge, USA · Chalmers University of Technology & University of Gothenburg, Gothenburg, Sweden +1