Gradient estimators for parameter inference in discrete stochastic kinetic models

TL;DR

This paper explores the use of machine learning gradient estimators to enable parameter inference in discrete stochastic kinetic models, specifically using the Gillespie algorithm, and compares their effectiveness.

Contribution

It introduces and compares three gradient estimators for the Gillespie SSA, demonstrating their advantages and limitations in different dynamic regimes.

Findings

GS-ST estimator often produces well-behaved gradients but can have high variance in difficult regimes.

Score Function and Alternative Path estimators offer more robust, lower variance gradients in challenging scenarios.

Gradient-based inference can be effectively integrated with Gillespie SSA using these estimators.

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.