A Framework for Stochastic Differentiable Programming
2026-08-03 • Emerging Technologies
Emerging Technologies
AI summaryⓘ
The authors introduce Parametrized Stochastic Circuits (PSCs), a new way to represent and work with random processes using programmable parts called stochastic kernels connected by different types of wires. They also created torx, a software tool that helps build and test these circuits efficiently. Their approach matches well with new types of probabilistic computer hardware, helping save energy by avoiding complicated translations. The authors show PSCs can handle various examples like random walks and sampling methods, and they tested it on real hardware where it performed similarly to traditional software methods.
stochastic circuitsprobabilistic hardwarestochastic kernelsJAX frameworkrandom walksMetropolis-Hastingsimportance samplingIsing modelbinary/categorical/continuous data typesCMOS
Authors
Guillaume Verdon, Leo Tyrpak, Owen Lockwood, Seth Morton, Alexander Neagoe, Anton Sugolov, Ian MacCormack, Mirko Amico
Abstract
We introduce Parametrized Stochastic Circuits (PSCs), a gate-based intermediate representation for programmable stochastic dynamics in which typed local stochastic kernels with tunable parameters compose over explicit binary, categorical, and continuous wires, and \texttt{torx}, an open-source JAX framework for constructing, executing, and differentiating them. PSCs' data types and stochastic kernels are chosen to align closely with the native operations exposed by emerging probabilistic hardware. In this way, stochastic algorithms can be designed directly in terms of the operations the hardware executes natively, so that the energy advantage arising at this level is not lost on mappings that introduce substantial decomposition, communication, or control overhead. We demonstrate the framework on a variety of example applications such as random walks on graphs, discrete diffusion, stochastic graph networks, jump diffusion and Ising sampling. We also report a hardware experiment in which probabilistic bits on the X0 subthreshold CMOS test chip, hosted by the XTR-0 desktop platform, provide physical randomness for Metropolis-Hastings and importance-sampling estimators, yielding estimates consistent with a software pseudorandom baseline.