Backpropagation for Effectful Languages I: Finite Probability and Discrete Output Algebraic Effects
This work provides a foundational framework for extending automatic differentiation to probabilistic programming languages with discrete effects, benefiting researchers in probabilistic programming and AD.
The paper extends reverse-mode automatic differentiation to discrete probabilistic programs with finite support and differentiable weights, proving correctness via categorical logical relations. It provides a reusable pattern for differentiating discrete-output algebraic effects like non-determinism, exceptions, and writer-style accumulation.
We analyse reverse-mode automatic differentiation (AD) for discrete probabilistic programs. Our construction is formulated in the framework of Combinatory Homomorphic Automatic Differentiation (CHAD), treating AD as a structure-preserving transformation of programs, guided by a denotational semantics. The main case study is the finite atomic distribution monad, whose computations have finite support and differentiable weights. The key point is that differentiating probabilistic programs requires cotangents to flow backwards not only through deterministic computations, but also through the probabilistic structure itself. We define the corresponding reverse-mode code transformation and prove its correctness, for handled real-output programs, by a categorical logical-relations argument. Although the paper focuses on finite discrete probability, the construction gives a reusable pattern for differentiating discrete-output algebraic effects, including finite multiset non-determinism (e.g., from fork-join parallelism), exceptions, and writer-style accumulation (e.g., for in-place accumulation of high-dimensional vectors). More broadly, we view this work as a foundational step towards extending CHAD to richer probabilistic languages and to other algebraic effects with handlers.