Search papers, labs, and topics across Lattice.
The paper introduces a method for exact stochastic simulation of continuous-time Markov chains (CTMCs) that overcomes the gradient-blocking issue of Gillespie-type algorithms by decoupling forward simulation from backward differentiation. This is achieved by using hard categorical sampling for exact trajectory generation and a continuous Gumbel-Softmax straight-through estimator for gradient propagation. The method enables accurate optimization at parameter scales significantly exceeding existing simulators, as demonstrated on various models including a large gene regulatory network achieving high MNIST accuracy.
Unlock gradient-based learning for exact stochastic simulations of continuous-time Markov chains, scaling parameter optimization by four orders of magnitude.
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.
