Gradient-based optimization of exact stochastic kinetic models

Abstract

Stochastic kinetic models describe systems across biology, chemistry, and physics where discrete events and small populations render deterministic approximations inadequate. Parameter inference and inverse design in these systems require optimizing over trajectories generated by the Stochastic Simulation Algorithm, but the discrete reaction events involved are inherently non-differentiable. We present an approach based on straight-through Gumbel-Softmax estimation that maintains exact stochastic simulations in the forward pass while approximating gradients through a continuous relaxation applied only in the backward pass. We demonstrate robust performance on parameter inference in stochastic gene expression, first recovering kinetic rates of telegraph promoter models from both moment statistics and full steady-state distributions across diverse and challenging synthetic parameter regimes, then inferring the kinetic parameters of a four-state promoter model from experimental single-molecule RNA timecourse measurements. We further apply the method to inverse design in stochastic thermodynamics, optimizing non-equilibrium currents in an interacting particle system under kinetic resource constraints and recovering known analytical bounds. The ability to efficiently differentiate through exact stochastic simulations provides a foundation for systematic scalable inference and rational design across the many domains governed by continuous-time Markov dynamics.