Skip to content

Propensity Straight-Through Gradients for Discrete Stochastic Systems

Aug 2026 · 0 citations · 38 references
Biology Physics Computer Science

TL;DR

This work exploits the affine state update to obtain the exact one-step conditional-mean sensitivity by differentiating normalized reaction propensities, and defines the propensity straight-through (PST) estimator, a temperature- and Gumbel-free path to scalable gradient-based learning through exact stochastic trajectories.

Abstract

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.

View source

Similar papers

Preprint Aug 2026

StocBench: A Benchmark for Generative Modeling of Stochastic Dynamics

While stochastic diffusion samplers such as DDPM better preserve the enstrophy spectrum during rollouts in the stochastic setting, deterministic samplers such as DDIM and DPM-2 show better spectral preservation in the deterministic setting.

S. Pfister, Benjamin J. Holzschuh, Nils Thürey · 0 citations
#machine learning Preprint Sep 2026

Time Without Timesteps: Simulating Coupled Dynamical Systems via Self-Consistency

Numerical simulation of dynamical systems is usually organized as a causal march through time: each state is computed from the previous one. We explore a different formulation for coupled systems. For each subsystem type we train a neural surrogate mapping a full driving trajectory and initial condition directly to a full output trajectory; following classical waveform relaxation, coupled systems are assembled by enforcing self-consistency among these trajectories: simulation becomes a fixed-point problem over complete trajectories rather than a stepwise rollout. On coupled van der Pol oscillators and Hodgkin-Huxley neuron networks, sequential depth becomes the number of solver iterations: 4-10 Newton iterations where the reference integrator takes 1500 steps. The gradient likewise loses its time recursion: it becomes a linear system solved by GMRES at memory independent of solver depth. A single scalar measured from the learned operator, the spectral radius of its Jacobian, predicts in advance where the coupled solve will converge; past that boundary, unrolled backpropagation diverges and a Neumann adjoint fails, while the implicit gradient remains correct to 0.04%. We report where the approach succeeds and where surrogate error degrades it.

Liyu Zerihun, M. Lee · 0 citations
Review Aug 2026

From Continuous Dynamics to Practical Gradient-Based Samplers

Gradient-based Markov chain Monte Carlo methods are often introduced as a catalog of algorithms: Hamiltonian Monte Carlo (HMC), the Metropolis-adjusted Langevin algorithm (MALA), the No-U-Turn Sampler (NUTS), and several underdamped variants. This presentation obscures the common structure of the methods and, more importantly, the reasons why a sampler that is correct in principle may be ineffective in practice. We develop a unified account, beginning with exact continuous-time dynamics that represent idealized sampling methods and for which Metropolis adjustments are not required. Numerical discretization makes the dynamics computationally feasible but introduces bias. Metropolis adjustment removes the asymptotic bias by converting numerical errors into rejection, leading to HMC, MALA, NUTS, and the Metropolis-adjusted kinetic Langevin algorithm (MAKLA). The second half of the paper presents geometric design choices that determine practical performance, namely, although MAKLA and NUTS have nice theoretical properties, their sampling efficiency may be slow in practice. Importantly, a fixed mass matrix can whiten globally anisotropic targets, often fixing sampling inefficiency in Bayesian posteriors with large data. Whereas hierarchical posteriors introduce their own problem, causing state-dependent variation in the Hessian (e.g., Neal's funnel). We explain how a randomized step size can be used effectively to sample from such a distribution. The resulting paper is both a tutorial on the mechanics of gradient-based sampling and a set of practical recipes to improve sampler performance.

J. Chok · 0 citations
Preprint Sep 2026

A Constrained Kuramoto Gradient-Flow System Can Perform High-Accuracy Finite-Time Inference

A central question in physical inference is whether strongly constrained dynamical systems can realize accurate input--output maps through their own finite-time evolution. We study this question in Kuramoto phase networks, whose deterministic dynamics form an input-conditioned gradient flow and whose predictions are read directly from output oscillators. As a constructive training approach, we develop a two-stage teacher--student procedure. A neural teacher is first converted into an explicit phase trajectory whose terminal oscillator activations reproduce the teacher outputs, and the Kuramoto parameters are trained by matching the student vector field along this prescribed path. Because accurate teacher-forced path matching does not ensure accurate autonomous inference, we then differentiate through the autonomous finite-time rollout and directly align its terminal output with the neural target. The resulting oscillator system, with $74$ oscillators, reaches mean test accuracies of $96.711\%$ on MNIST and $86.399\%$ on Fashion-MNIST. This capability persists across neural-teacher architectures, matched system sizes, thermal perturbations, and integration-grid refinement. Together, these results provide a constructive demonstration that a strongly constrained, small-sized Kuramoto gradient-flow system can be trained for high-accuracy finite-time inference through a direct oscillator readout.

Yi Cheng, Zong-Li Lin · 0 citations
Jul 2026

Markov state models revisited: Principles and algorithms for unbiased observables.

Markov state models (MSMs) have become ubiquitous tools for analyzing molecular dynamics (MD) simulations because of their simple, powerful premise: although complete MD sampling may be impossible, the MSM can "stitch together" transition probabilities derived from local sampling to provide a global picture of kinetics and mechanisms. In the standard MSM framework, the available MD data is organized into a single transition matrix, which is then used to estimate all observables at a lag time chosen so the coarse-grained dynamics are approximately Markovian. This approach leads to avoidable model bias and motivates long lag times that obscure short-timescale processes of interest. In contrast, this paper shows how to obtain unbiased coarse-grained observables at any fixed lag time and for any fixed coarse-graining in the limit of infinite, properly weighted data. The central idea is to replace the single-matrix framework with two transition matrices -- one representing equilibrium dynamics and another representing source-sink recycling dynamics -- and use the correct matrix or matrices to estimate the matched dynamical observables.

D. Aristoff, Robert J Webber, D. M. Zuckerman · 0 citations
Preprint Aug 2026

Conditional-path Monte Carlo for rare stochastic dynamics on networks: Details and derivations

This paper provides the rigorous mathematical foundations and algorithmic details underlying the CPMC framework, and details a dynamic programming scheme to exactly implement complex boundary conditions - including patient-zero and macroscopic outbreak-size constraints - enabling the rejection-free generation of valid trajectories.

Thomas Barthel, Jiazheng Sun, Jhao-Hong Peng · 1 citation

Related blog posts

We use cookies to run the site and, with your consent, for analytics and to show ads. See our Cookie Policy.