Propensity Straight-Through Gradients for Discrete Stochastic Systems

Jose M. G. Vilar; Leonor Saiz

q-bio.QMAug 26, 2026

Propensity Straight-Through Gradients for Discrete Stochastic Systems

Authors: Jose M. G. Vilar, Leonor Saiz

Organizations: 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

Abstract

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 ( $R^2$ = 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.

Explore similar work

Feb 23, 2026

Exact Discrete Stochastic Simulation with Deep-Learning-Scale Gradient Optimization

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

Differentiable Physics

Ask PDF

Apr 2, 2026

Gradient estimators for parameter inference in discrete stochastic kinetic models

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.

Ludwig Burger, Annalena Kofler, Lukas Heinrich +1

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 +2

Parameter Estimation Simulation-Based Inference

Ask PDF

Oct 2, 2026

Tangent Schrödinger Bridge Matching: Learning Stochastic Transport with Mechanistic Sensitivities

Predicting how stochastic systems respond to changes in viscosity, reaction rates, or external forces requires costly simulations, motivating reusable learned models. Yet matching observed outcome distributions does not ensure accurate intervention responses. We introduce Tangent Schrödinger Bridge Matching (Tangent-SBM), which learns stochastic transports from endpoint observations and mechanistic sensitivities. It propagates parameter derivatives alongside trajectories and supervises them against supplied targets. For average-response targets, single-rollout squared error also penalizes response variability; our objective uses two independent rollouts to match the mean without this additional penalty. We establish conditions under which sensitivity accuracy bounds finite-change prediction error and decision regret. Across Gaussian, stochastic double-well, PDEBench reaction--diffusion, and stochastic Navier--Stokes systems, Tangent-SBM improves sensitivity and finite-change prediction over matched conditional-bridge baselines while maintaining comparable endpoint and distributional accuracy. Controls examine target correctness, response objectives, and simulator-budget allocation. To test decision usefulness, we evaluate calibrated viscosity selection in Navier--Stokes: Tangent-SBM reduces tracking error relative to taking no action on every evaluated task.

Jowaria Khan, Elizabeth Bondi-Kelly

University of Michigan

Schrödinger Bridges Optimal Transport Stochastic Differential Equations PDE Surrogate Modeling

Ask PDF