Gradient-based optimization of exact stochastic kinetic models

TL;DR

Stochastic kinetic models involve non-differentiable discrete events, complicating gradient-based inference and design. The authors propose a reparameterized pathwise gradient using straight-through Gumbel-Softmax to differentiate through reaction choices while keeping SSA forward trajectories exact, together with a variance-reduction strategy using baseline simulations. They validate the method on stochastic gene expression (telegraph promoter) for moment- and distribution-based parameter inference and on a three-state ring for Pareto optimization of current versus entropy production under a kinetic budget, obtaining accurate parameter recovery and tight Pareto fronts. The framework enables efficient gradient-based optimization for objectives computed from exact stochastic trajectories and offers a scalable route toward systematic inference and inverse design in domains governed by continuous-time Markov dynamics, with key result $J^*(\mathcal{A},K_{tot}) = \frac{K_{tot}}{9}\tanh\left(\frac{\mathcal{A}}{6}\right)$ in the thermodynamic setting.

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.

Figures (23)

• Figure 1: Modeling and Optimization Framework. SSA trajectories are generated exactly in the forward pass by sampling waiting times and discrete reaction events (Left, blue shaded box). In the backward pass, gradients are computed by directly reparameterizing waiting times and applying a continuous Gumbel-Softmax relaxation to reaction selection (Right, yellow shaded box), enabling efficient gradient-based optimization without approximating the stochastic dynamics.
• Figure 2: Moment matching inferences from synthetic data. Inference of parameters from mean and variance of the true underlying generative model at steady state, for 25 different parameter sets. Gray markers indicate the random initializations. (a) Two parameters of the telegraph model ($k_{\text{on}}$, $k_{\text{tx}}$) are simultaneously estimated. (b) Mean and standard deviation of the RNA levels at steady state for the true and optimized models. (c) Comparison between the estimated and true rates.
• Figure 3: Distribution matching inferences from synthetic data. Inference of parameters from the RNA histogram of the true underlying generative model at steady state, for 25 different parameter sets. Gray markers indicate the random initializations. (a) Three parameters of the telegraph model ($k_{\text{on}}$, $k_{\text{tx}}$, $k_{\text{mdeg}}$) are simultaneously estimated. (b) Comparison between the estimated and true transcription rate. (c) Relative errors in the estimation of the degradation rate $k_{\text{mdeg}}$ and the on-rate $k_{\text{on}}$ illustrate sloppy parameters directions. (d) Two representative examples comparing the target data distribution (black) and the distribution generated by the optimized model (red), showing agreement over several orders of magnitude.
• Figure 4: Optimal current-dissipation trade-off in a three-state ring. Optimization of cycle current $J$ under constraints on entropy production rate $\sigma$ and total kinetic budget $K_\text{tot}$. (a) Schematic of the three-state ring with forward ($k_{ij}^+$) and backward ($k_{ij}^-$) transition rates. (b) Optimized cycle current (blue circles) compared to the theoretical Pareto front (dashed line). Dotted line: asymptotic near-equilibrium scaling. Inset: Pareto front on linear axes. (c) Relative errors in $J$ (versus theory), $\sigma$, and $K_\text{tot}$ (versus target). Shaded band: $\pm 1\%$. (d) Optimized forward rates $k_{12}^+$, $k_{23}^+$, $k_{31}^+$ (light green symbols), their mean $\langle k^+ \rangle$ (green diamonds), and backward rate $k^-$ (squares), compared to theoretical predictions (dashed and dotted lines).
• Figure S1: Moment matching optimization results.(Left) Loss trajectories for all 25 parameter sets of the telegraph promoter model. Light curves show individual runs; dark curve shows the mean. (Center) Parameter recovery accuracy: ratio of estimated to true parameter values for $k_\text{on}$ and $k_\text{tx}$. Gray points show initial guesses; colored points show optimized values, with color indicating final loss. Optimized parameters cluster near the true values (intersection of gray lines at ratio 1). (Right) Same data as center panel, with color indicating the simulated trajectory duration.