Gradient estimators for parameter inference in discrete stochastic kinetic models

Abstract: Stochastic kinetic models are ubiquitous in physics, yet inferring their parameters from experimental data remains challenging. In deterministic models, parameter inference often relies on gradients, as they can be obtained efficiently through automatic differentiation. However, these tools cannot be directly applied to stochastic simulation algorithms (SSA) such as the Gillespie algorithm, 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 compare the properties of all estimators in two representative systems exhibiting relaxation or oscillatory dynamics, where the latter requires gradient estimation of time-dependent objective functions. We find that the GS-ST estimator mostly yields well-behaved gradient estimates, but exhibits diverging variance in challenging parameter regimes, resulting in unsuccessful parameter inference. In these cases, the other estimators provide more robust, lower variance gradients. Our results demonstrate that gradient-based parameter inference can be integrated effectively with the Gillespie SSA, with different estimators offering complementary advantages.

• The paper demonstrates that GS-ST, SF, and AP yield distinct trade-offs between bias and variance in gradient estimation for Gillespie SSA-based parameter inference.
• Methodologically, GS-ST uses softmax relaxation while SF employs likelihood ratios, with SF offering more robust performance in long trajectory regimes.
• Empirical studies reveal that SF achieves lower variance than AP, and GS-ST can suffer from exponential gradient variance in extreme parameter regions.

Gradient Estimators for Parameter Inference in Discrete Stochastic Kinetic Models

Introduction

Stochastic kinetic models, particularly those describing biochemical reaction networks with intrinsic fluctuations, pose fundamental challenges for parameter inference. Unlike deterministic models where analytical gradients enable efficient optimization via automatic differentiation, stochastic descriptions—especially those operationalized via the Gillespie Stochastic Simulation Algorithm (SSA)—involve non-differentiable discrete sampling steps, obstructing direct gradient-based inference. This work benchmarks three prominent gradient estimators from machine learning—Gumbel-Softmax Straight-Through (GS-ST), the score function (SF) estimator, and the alternative path (AP) estimator—for their suitability in parameter inference within Gillespie SSA, elucidating their statistical properties and operational regimes.

Figure 1: Overview of Gillespie SSA and the three gradient estimators: GS-ST (b) uses softmax relaxation for differentiability; SF (c) accumulates score terms along the trajectory; AP (d) evaluates weighted differences using primal and alternative paths.

Methodological Framework

The Gillespie SSA models reaction networks as continuous-time Markov chains, iteratively sampling both discrete reaction events and inter-reaction waiting times from parameterized probability distributions. Inference necessitates gradients of expected observables (e.g., $\nabla_\theta \mathbb{E}[f(N(t))]$), yet the discrete nature of event sampling precludes naive differentiation.

The three estimators adapted for such discrete stochastic processes are:

• GS-ST: Utilizes the Gumbel-max trick to relax discrete samples via softmax, enabling (biased) differentiable surrogates for gradients during the backward pass but retaining true discrete trajectories in the forward simulation. Control via temperature $\tau$ introduces a bias-variance trade-off—low $\tau$ reduces bias but can induce gradient variance explosion.
• SF: Provides unbiased estimators via the likelihood ratio (reinforce) method, accumulating parameter-gradient terms (the score) scaled by deviations of observed quantities from their expected values. While unbiased, its variance grows linearly with the number of reaction events.
• AP: Computes gradients by generating a paired alternative trajectory with an infinitesimal parameter perturbation, leveraging boundary-shift weights. Like SF, AP is unbiased but typically incurs even larger variance.

Empirical Studies: Relaxation and Oscillatory Dynamics

Simple Relaxation: Bimolecular Association

The benchmarking commences with a reversible bimolecular association system ($A + B \leftrightarrow AB$), where the dissociation rate $k$ is the key parameter. The system admits an exact chemical master equation solution, facilitating unbiased comparison of estimators.

Empirical loss surfaces and gradients computed by all three estimators align well on average with ground truth but reveal key distinctions in uncertainty:

• GS-ST achieves low-variance estimates at optimal temperature, but at low $\tau$ or extreme $k$, experiences exponential gradient variance growth, which is analytically tracked via a Lyapunov exponent formalism.
• SF and AP maintain linear-in-trajectory-length variance; SF achieves consistently lower variance than AP, especially as the required trajectory lengths increase with large $k$.

(Figure 2)

Figure 2: Relaxation dynamics and gradient estimates for the bimolecular association system across estimators and parameter regimes.

Oscillatory Dynamics: The Repressilator

A more complex test uses the repressilator, a model genetic oscillator with feedback capable of sustained stochastic oscillations. Target parameters are the production rate $k_{\rm p}$ and dissociation constant $K_{\rm d}.$ The loss is defined on the logarithm of species copy number trajectories over one oscillation period.

Parameter inference via stochastic gradient descent is evaluated for 50 randomly sampled ground truths and initialization points per estimator:

• Both GS-ST and SF recover true parameters accurately in most cases, but GS-ST can fail in regimes of high binding affinity ($\tau$0), where the gradient variance becomes catastrophic.
• AP exhibits prohibitively high variance and consistently yields inferior inferences.

(Figure 3)

Figure 3: Inference results comparing true and learned parameters for all three estimators in the repressilator system.

Gradient Variance Analysis

A critical contribution is the identification of estimator regimes through variance scaling analysis:

• GS-ST: Exhibits mode switch between bounded and exponentially divergent variance, governed by the interplay of $\tau$1 and model parameters. In oscillatory and high-propensity regimes (large $\tau$2), a single estimator temperature cannot reliably balance bias and variance, leading to unacceptably noisy gradients in subset regions of parameter space.
• SF and AP: Variance grows linearly with the number of events for both estimators; for AP, the scaling constant is substantially larger. Thus, SF is preferred whenever computational budget is at a premium or reaction bursts generate long event sequences.

(Figure 4)

Figure 4: Gradient variance scaling versus Gillespie steps and parameters for GS-ST and SF, demonstrating exponential versus linear variance escalation.

Implications and Outlook

This systematic comparison highlights the nuanced operational regimes and limitations of gradient-based inference for stochastic kinetic models under SSA:

• GS-ST is efficient in tractable regimes (modest event counts, moderate parameter values) but can become intractable without self-adaptive temperature control or variance reduction. Arbitrarily increasing $\tau$3 introduces gradient bias, impeding convergence.
• SF is robust but limited by linear variance growth; inference with long trajectories or large populations will require subsampling or additional variance-reduction strategies (e.g., control variates, Rao-Blackwellization).
• AP is rarely competitive due to its even higher intrinsic variance.

Practical ramifications include the choice of estimator for systems biology and biophysics applications: SF is preferred for robustness across parameter regimes and as a basis for future improvements. GS-ST may be useful for rapid inference in low to moderate complexity settings or with careful adaptive tuning. Large gradient variance is the primary obstacle preventing broad adoption of gradient-based inference in stochastic simulation—variance-reduction remains an open, critical research direction.

Theoretically, robust score estimators can integrate gradient-based Bayesian inference methods such as Hamiltonian Monte Carlo, linking likelihood-free inference techniques with scalable simulation-based models. Extensions to models with dynamic reaction network topologies (non-static species, reactions appearing/disappearing) represent another frontier.

Conclusion

Gradient-based parameter inference in discrete stochastic kinetic models simulated with Gillespie SSA is feasible via advanced gradient estimators from machine learning. The GS-ST, SF, and AP estimators each present unique operational advantages and challenges: GS-ST achieves low-variance gradients in select regimes at the cost of potential divergence elsewhere, while SF offers robust, unbiased inference with manageable linear variance scaling. AP is generally dominated by SF in terms of variance. Future research should focus on estimator variance reduction and further integration with Bayesian inference frameworks to enable practical, scalable inference in more complex biochemical networks.