Exact Discrete Stochastic Simulation with Deep-Learning-Scale Gradient Optimization

TL;DR

This paper introduces a novel deep-learning-based method for exact stochastic simulation of continuous-time Markov chains, enabling scalable, accurate, and differentiable simulations suitable for high-dimensional systems and inverse design.

Contribution

It decouples simulation from differentiation using Gumbel-Softmax, allowing gradient-based optimization of complex stochastic models at unprecedented scales.

Findings

Achieves less than 0.1% error on a reversible dimerization model

Reaches 98.4% accuracy on a gene regulatory network benchmark

Performs 1.9 billion simulation steps per second on GPU

Abstract

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.