# Automatic Differentiation of Programs with Discrete Randomness

arXiv: https://arxiv.org/abs/2210.08572 · HTML: https://mlsanity.com/paper/2210.08572/

## Key facts

- Cited by 29 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 60 references that are in the map.
- First posted on arXiv 16 Oct 2022.

## Builds on (most cited first, 18 of 18)

- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Automatic differentiation in machine learning: a survey](https://mlsanity.com/paper/1502.05767/) (2015)
- [Monte Carlo Gradient Estimation in Machine Learning](https://mlsanity.com/paper/1906.10652/) (2019)
- [Gradient Estimation Using Stochastic Computation Graphs](https://mlsanity.com/paper/1506.05254/) (2015)
- [Backpropagation through the Void: Optimizing control variates for black-box gradient estimation](https://mlsanity.com/paper/1711.00123/) (2017)
- [REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models](https://mlsanity.com/paper/1703.07370/) (2017)
- [Differentiable Particle Filters: End-to-End Learning with Algorithmic Priors](https://mlsanity.com/paper/1805.11122/) (2018)
- [Differentiable Particle Filtering via Entropy-Regularized Optimal Transport](https://mlsanity.com/paper/2102.07850/) (2021)
- [Don't Unroll Adjoint: Differentiating SSA-Form Programs](https://mlsanity.com/paper/1810.07951/) (2018)
- [Differentiable Agent-based Epidemiology](https://mlsanity.com/paper/2207.09714/) (2022)
- [Differentiable Particle Filtering without Modifying the Forward Pass](https://mlsanity.com/paper/2106.10314/) (2021)
- [Storchastic: A Framework for General Stochastic Automatic Differentiation](https://mlsanity.com/paper/2104.00428/) (2021)
- [Efficient Learning of the Parameters of Non-Linear Models using Differentiable Resampling in Particle Filters](https://mlsanity.com/paper/2111.01409/) (2021)
- [Plots.jl -- a user extendable plotting API for the julia programming language](https://mlsanity.com/paper/2204.08775/) (2022)
- [Direct Policy Gradients: Direct Optimization of Policies in Discrete Action Spaces](https://mlsanity.com/paper/1906.06062/) (2019)
- [Situating Agent-Based Modelling in Population Health Research](https://mlsanity.com/paper/2002.02345/) (2020)

## Built on by (most cited first, 29 of 29)

- [On the limits of agency in agent-based models](https://mlsanity.com/paper/2409.10568/) (2024)
- [Smoothing Methods for Automatic Differentiation Across Conditional Branches](https://mlsanity.com/paper/2310.03585/) (2023)
- [Differentiating Metropolis-Hastings to Optimize Intractable Densities](https://mlsanity.com/paper/2306.07961/) (2023)
- [Branches of a Tree: Taking Derivatives of Programs with Discrete and Branching Randomness in High Energy Physics](https://mlsanity.com/paper/2308.16680/) (2023)
- [An efficient Wasserstein-distance approach for reconstructing jump-diffusion processes using parameterized neural networks](https://mlsanity.com/paper/2406.01653/) (2024)
- [On Feynman--Kac training of partial Bayesian neural networks](https://mlsanity.com/paper/2310.19608/) (2023)
- [Probabilistic Programming with Programmable Variational Inference](https://mlsanity.com/paper/2406.15742/) (2024)
- [Bayesian calibration of stochastic agent based model via random forest](https://mlsanity.com/paper/2406.19524/) (2024)
- [Enhanced SMC2: Leveraging Gradient Information from Differentiable Particle Filters Within Langevin Proposals](https://mlsanity.com/paper/2407.17296/) (2024)
- [Generalizing Stochastic Smoothing for Differentiation and Gradient Estimation](https://mlsanity.com/paper/2410.08125/) (2024)
- [Learning Individual Behavior in Agent-Based Models with Graph Diffusion Networks](https://mlsanity.com/paper/2505.21426/) (2025)
- [Some challenges of calibrating differentiable agent-based models](https://mlsanity.com/paper/2307.01085/) (2023)
- [Diffusion differentiable resampling](https://mlsanity.com/paper/2512.10401/) (2025)
- [A Blueprint for Equilibrium-Based Differentiable Continuous-Variable Thermodynamic Computing](https://mlsanity.com/paper/2607.16183/) (2026)
- [Auto-Differentiation of Relational Computations for Very Large Scale Machine Learning](https://mlsanity.com/paper/2306.00088/) (2023)
- [Automatic Gradient Estimation for Calibrating Crowd Models with Discrete Decision Making](https://mlsanity.com/paper/2404.04678/) (2024)
- [Towards Learning Stochastic Population Models by Gradient Descent](https://mlsanity.com/paper/2404.07049/) (2024)
- [Fully Data-driven but Interpretable Human Behavioural Modelling with Differentiable Discrete Choice Model](https://mlsanity.com/paper/2412.19403/) (2024)
- [Gradient estimators for parameter inference in discrete stochastic kinetic models](https://mlsanity.com/paper/2604.02121/) (2026)
- [DynaLay: An Introspective Approach to Dynamic Layer Selection for Deep Networks](https://mlsanity.com/paper/2312.12781/) (2023)
- [A Principled Bayesian Framework for Training Binary and Spiking Neural Networks](https://mlsanity.com/paper/2505.17962/) (2025)
- [A generalized Wasserstein-2 distance approach for efficient reconstruction of random field models using stochastic neural networks](https://mlsanity.com/paper/2507.05143/) (2025)
- [Control of dynamical systems with neural networks](https://mlsanity.com/paper/2510.12810/) (2025)
- [Refined Gradient-Based Temperature Optimization for the Replica-Exchange Monte-Carlo Method](https://mlsanity.com/paper/2601.13542/) (2026)
- [Dimensional Peeking for Low-Variance Gradients in Zeroth-Order Discrete Optimization via Simulation](https://mlsanity.com/paper/2602.00075/) (2026)
- [partialCBDs: Differentiable Causal Block Diagrams](https://mlsanity.com/paper/2602.07581/) (2026)
- [GradInf: Gradient Estimation as Probabilistic Inference](https://mlsanity.com/paper/2607.07840/) (2026)
- [Scalable Training of Continuous-Time Spiking Neural Networks with Differentiable Spike-Time Discretization](https://mlsanity.com/paper/2607.14672/) (2026)
- [Learning Graph Energies for Finite-Time Generation](https://mlsanity.com/paper/2610.19946/) (2026)

## Abstract (arXiv)

Automatic differentiation (AD), a technique for constructing new programs which compute the derivative of an original program, has become ubiquitous throughout scientific computing and deep learning due to the improved performance afforded by gradient-based optimization. However, AD systems have been restricted to the subset of programs that have a continuous dependence on parameters. Programs that have discrete stochastic behaviors governed by distribution parameters, such as flipping a coin with probability p of being heads, pose a challenge to these systems because the connection between the result (heads vs tails) and the parameters (p) is fundamentally discrete. In this paper we develop a new reparameterization-based methodology that allows for generating programs whose expectation is the derivative of the expectation of the original program. We showcase how this method gives an unbiased and low-variance estimator which is as automated as traditional AD mechanisms. We demonstrate unbiased forward-mode AD of discrete-time Markov chains, agent-based models such as Conway's Game of Life, and unbiased reverse-mode AD of a particle filter. Our code package is available at https://github.com/gaurav-arya/StochasticAD.jl.
