A Framework for Stochastic Differentiable Programming

2026-08-02

AI Summary

This work addresses the abstraction gap between differentiable programming and emerging probabilistic hardware by introducing Parameterized Stochastic Circuits (PSCs) as a gate-level intermediate representation that closely mirrors native hardware operations. PSCs unify explicit binary, categorical, and continuous signals with localized stochastic kernels. Building on this foundation, the authors develop torx, an open-source JAX framework that enables, for the first time, direct alignment between differentiable stochastic computation and probabilistic hardware, substantially reducing mapping overhead. Experiments on the X0 subthreshold CMOS probabilistic bit chip demonstrate that PSCs efficiently harness physical randomness in tasks such as graph random walks, discrete diffusion, and stochastic graph networks, achieving results consistent with pseudorandom software baselines while confirming the approach’s validity and hardware energy efficiency.

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.

Problem

• stochastic differentiable programming
• probabilistic hardware
• stochastic algorithms
• hardware-native operations
• energy efficiency

Innovation

• Parametrized Stochastic Circuits
• Stochastic Differentiable Programming
• Probabilistic Hardware
• torx
• Hardware-Aware Stochastic Algorithms