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, accurately recovering kinetic rates of telegraph promoter models from both moment statistics and full steady-state distributions across diverse and challenging parameter regimes. We further demonstrate the method's applicability to inverse design problems in stochastic thermodynamics, characterizing Pareto-optimal trade-offs between non-equilibrium currents and entropy production. The ability to efficiently differentiate through exact stochastic simulations provides a foundation for systematic inference and rational design across the many domains governed by continuous-time Markov dynamics.