Automatic Differentiation of Programs with Discrete Randomness (2210.08572v3)

Abstract: 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.

• The paper introduces a novel automatic differentiation method that computes unbiased derivatives for programs with inherent discrete randomness.
• It leverages stochastic derivatives to generalize pathwise gradient estimators, reducing variance compared to finite-difference methods.
• The approach extends AD capabilities to simulation models, enabling efficient optimization in discrete stochastic systems like Markov chains and agent-based models.

Automatic Differentiation of Programs with Discrete Randomness: A Technical Analysis

This paper introduces a novel approach to automatic differentiation (AD) of programs involving discrete randomness, addressing a critical gap in AD systems which have traditionally been limited to handling continuous parameter dependencies. By leveraging a reparameterization-based method, the authors develop an approach that can derive programs for computing the expectation of a derivative, even when the underlying program involves discrete stochastic behaviors, such as those found in discrete-time Markov chains or agent-based models like Conway's Game of Life.

Summary of Methodology and Contributions

The paper outlines a methodology centered around stochastic derivatives, which are designed to estimate the derivative of statistical quantities in a low-variance and unbiased manner. Key components of this methodology include:

1. Stochastic Derivatives: These derivatives consider the proportional probability of different event outcomes due to infinitesimal changes in parameters, effectively generalizing pathwise gradient estimators, commonly used in continuous randomness, to discrete contexts.
2. Composability and Unbiasedness: The authors highlight the method's ability to automatically differentiate stochastic programs composed of elementary samples from distributions (e.g., Bernoulli, Poisson, Geometric). It maintains unbiasedness by preserving discrete structure and avoiding bias-introducing continuous relaxations.
3. Variance Reduction: By correlating paths through computation using minimal perturbations, the method achieves variance reduction, providing a significant advantage over traditional finite differences, which struggle with unbounded variance as step size decreases.

The practical utility of this approach is demonstrated through its application to forward-mode AD in stochastic simulations and reverse-mode AD via smoothing techniques, which are shown to produce unbiased differentiable estimates for particle filters.

Numerical Results and Applications

The numerical results presented are compelling. The authors report that their approach delivers unbiased estimates with competitive computational complexity for various synthetic examples. Notably, the method exhibits variance improvements compared to score function estimators (common in handling discrete randomness). Practical applications of the method are demonstrated through implementations in StochasticAD.jl, showcasing its effectiveness in differentiating discrete-time Markov processes and the stochastic Game of Life.

Theoretical and Practical Implications

Theoretically, the work represents a considerable advancement in extending AD to discrete stochastic programs, a domain previously not fully addressed by existing AD frameworks. Practically, this implies enhanced capability in optimizing models where discrete decisions are integral, such as in reinforcement learning, agent-based modeling, or any scientific computing task involving discrete probability distributions.

Speculations on Future Developments in AI

The seamless extension of AD to discrete randomness opens the door to novel applications of AI in domains where decision-making involves stochastic processes. For instance, in areas like automated reasoning or probabilistic programming, where discrete random variables are prevalent, this research can substantially improve optimization routines. Furthermore, as more complex systems are modeled stochastically in machine learning, the ability to automatically differentiate across hybrid (discrete-continuous) systems will become increasingly crucial, making this contribution timely.

In summary, this paper provides a comprehensive methodological advancement in AD, enabling unbiased differentiation of programs interlaced with discrete randomness. The proposed stochastic derivatives pave the way for more robust and efficient AD implementations in domains with inherent stochastic behaviors. The work, while theoretical, firmly sets a groundwork for applications with significant implications across various AI and scientific computing tasks.