Backpropagation for Effectful Languages I: Finite Probability and Discrete Output Algebraic Effects

2026-07-15

AI Summary

This work addresses the challenge of implementing reverse-mode automatic differentiation for programs featuring algebraic effects such as finite discrete probabilistic choice. Building upon the Compositional Homomorphic Automatic Differentiation (CHAD) framework, it formulates automatic differentiation as a semantics-preserving program transformation and constructs a backward-pass mechanism tailored to the finite atomic distribution monad. The correctness of this construction is established using logical relations from category theory. This study presents the first systematic extension of reverse-mode automatic differentiation to effectful languages with discrete outputs, introducing a reusable differentiation scheme applicable to a broad class of algebraic effects—including nondeterminism, exceptions, and writer effects. The approach not only enables correct reverse differentiation of programs with finite discrete probabilistic structure but also lays a foundational theoretical groundwork for differentiating more general effectful languages.

Abstract

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.

Problem

• automatic differentiation
• probabilistic programs
• algebraic effects
• reverse-mode
• discrete probability

Innovation

• reverse-mode automatic differentiation
• algebraic effects
• probabilistic programming
• CHAD
• cotangent propagation