Source-linked AI summary
Propensity Straight-Through Gradients for Discrete Stochastic Systems
Jose M. G. Vilar, Leonor Saiz
TL;DR
Gradient-based learning through exact discrete stochastic trajectories is limited by nondifferentiable categorical reaction selection. PST keeps hard exact-forward sampling but differentiates normalized propensities, achieving comparable or faster optimization than GS-ST across benchmarks and scaling to a 203,796-parameter classifier.
Problem
Exact Gillespie-style trajectories are difficult to optimize with ordinary pathwise gradients because categorical reaction selection is discrete.
Method
PST retains exact SSA forward sampling and uses the derivative of normalized propensities as the backward rule for the one-step conditional mean.
Results
PST achieves parameter-recovery accuracy comparable to GS-ST while converging 3.0-fold faster on the genetic oscillator and remaining accurate in a 203,796-parameter classifier.
Takeaways & Limitations
Gradient-based learning through exact discrete stochastic simulation does not require relaxing the reaction selected by the simulator.
Takeaways & Limitations
Trajectory-level discrepancies can accumulate, cancel, change sign, or be damped over longer trajectories, and their control depends locally on downstream curvature and the backward extension.
Abstract
from arXiv · showhide
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.
Introduction
Exact discrete-event simulation is difficult to optimize with gradients because categorical reaction selection blocks ordinary pathwise differentiation. PST addresses this by differentiating normalized propensities for the exact one-step conditional mean while retaining hard forward trajectories.
- Motivation: Gradient-based optimization is needed for mechanistic models embedded in inference, inverse design, and machine-learning procedures.These applications can involve many kinetic parameters and parameter spaces beyond parameter-by-parameter search.
- Differentiation challenge: Gillespie-type exact simulation alternates exponential waiting-time sampling with categorical reaction selection, making pathwise differentiation difficult.The waiting time is differentiable, but the categorical reaction realization is piecewise constant in its probabilities.
- PST approach: PST differentiates normalized propensities to obtain the exact derivative of the one-step conditional mean state.This exactness follows from the affine SSA state update and also applies to affine post-reaction observables.
- PST approach: PST retains exact forward reaction selection while replacing relaxed backward sampling with a marginal-based derivative of normalized propensities.Unlike GS-ST, PST removes the backward-pass temperature and Gumbel perturbation.
- Theory and scope: PST’s trajectory-level sensitivities can differ from exact multistep sensitivities even though its one-step conditional-mean sensitivity is exact.The discrepancy is explicitly characterized, while PST complements score-function, finite-difference, unbiased, and pathwise sensitivity methods.
- Evaluation: PST is characterized theoretically and evaluated on biochemical, oscillator, ion-channel, repressilator, and large stochastic-classifier benchmarks.The classifier contains 203,796 trainable parameters, extending evaluation to deep-learning scale.
Results
PST keeps exact SSA trajectories in the forward pass while differentiating normalized propensities in the backward pass, yielding exact one-step conditional-mean sensitivities. Across reaction-network, oscillator, ion-channel, repressilator, and MNIST tests, it achieves accurate inference and competitive or faster optimization than GS-ST.
- PST method: PST leaves hard reaction selection and exponential waiting-time sampling unchanged, modifying only the derivative through the reaction indicator.The forward value remains the exact one-hot reaction indicator, while automatic differentiation uses the normalized propensity derivative.
- PST method: The PST backward pass exactly differentiates the one-step conditional mean and affine post-reaction observables.For a one-step squared-error objective on the population mean, this gives the exact parameter gradient by the chain rule.
- Trajectory-level sensitivity: Multistep PST discrepancies vanish for affine downstream functions and are controlled locally by downstream curvature along stoichiometric directions.The approximation replaces finite differences between reachable states with directional derivatives through the computational graph.
- Benchmark results: 0.060% two-parameter MAPE was achieved by PST across eight reversible-dimerization conditions, compared with 0.069% for GS-ST.Exact trajectories from PST-inferred parameters reproduced target ensemble means for all three species at the representative k2 = 0.32 condition.
- Benchmark results: 3.0-fold faster threshold attainment was observed for PST than GS-ST under equal per-step computational cost in the genetic oscillator.Median epochs to sustained error below 3% were 262.5 for PST and 782.5 for GS-ST.
- Benchmark results: PST matched GS-ST accuracy across discrete and large-scale benchmarks, including R^2 = 0.988 for ion-channel recordings, 0.168% median MAPE on 50 repressilator tasks, and 98.22% MNIST accuracy.The MNIST model used 203,796 trainable parameters with hard categorical sampling and no temperature, annealing, or soft-forward warm-up.
Discussion
PST enables gradient-based learning through exact stochastic trajectories by retaining hard forward sampling and differentiating normalized propensities. It matches GS-ST accuracy across benchmarks, often optimizes faster, and complements existing CTMC sensitivity methods.
- PST retains exact hard categorical reaction selection in the forward pass while using normalized propensities for the backward rule.The rule is exact for the one-step conditional-mean state and affine post-reaction observables.
- The trajectory-level PST discrepancy replaces exact finite differences with directional derivatives and vanishes for affine downstream dependence.Its local behavior is governed by downstream curvature, while longer-trajectory discrepancies may accumulate, cancel, change sign, or be damped.
- PST matches GS-ST parameter-recovery accuracy across inference benchmarks while avoiding a relaxed reaction-selection backward sample.The comparison includes the genetic oscillator, ion channel, repressilator suite, and a 203,796-parameter classifier.
- 3.0-fold fewer epochs were required for the genetic oscillator to reach the sustained 3% recovery threshold under matched settings.
- PST complements unbiased, finite-difference, score-function, pathwise, and exact finite-state chemical-master-equation sensitivity approaches.
Materials and Methods
The methods compare PST and GS-ST under matched stochastic models, objectives, optimization settings, and exact-forward simulation protocols across several inference benchmarks and a large reaction-network classifier.
- Comparison protocol: PST and GS-ST use the same state model, objective, optimizer, learning-rate schedule, ensemble size, initialization protocol, and simulation horizon within each inference benchmark.They differ in the backward treatment of categorical reaction selection.
- Reversible dimerization: The dimerization protocol generates targets from 100,000 exact SSA trajectories per condition across eight reverse-rate conditions and trains with 100,000 model trajectories per epoch for 250 epochs.Training minimizes squared error between target and model ensemble-mean time courses on a common physical-time grid.
- Genetic oscillator: The oscillator study infers five parameters from a nine-species, sixteen-reaction network using exact-SSA target segments and segment-level rate statistics.Each epoch samples 8192 starting states and generates 25 model trajectories per state.
- Ion-channel benchmark: The ion-channel experiment uses 100 NaV1.5 patch-clamp sweeps, 262,144 model trajectories per epoch, 20 SSA event slots per trajectory, and 400 training epochs.The objective compares experimental mean open-channel counts with exact overlap integrals over experimental time bins.
- Repressilator benchmark: The repressilator suite evaluates production rate and dissociation constant inference across 50 reference–initialization tasks using 50,000 exact trajectories for targets, training evaluations, and validation.PST and GS-ST share the exact SSA forward process, objective, optimizer, stopping rule, and refinement stage.
- Deep-learning-scale test: The stochastic gene-regulatory classifier uses 203,796 trainable parameters in a 784–256–10 reaction-network architecture and trains for 40 epochs with the published MNIST setup.The protocol includes categorical cross-entropy, RMSprop, stochastic weight averaging, and post-training Monte Carlo and temporal averaging.