A model of stochastic memoization and name generation in probabilistic programming: categorical semantics via monads on presheaf categories
Younesse Kaddar, Sam Staton
TL;DR
This work addresses the semantic challenge of stochastic memoization in higher-order probabilistic programming, especially for diffuse (non-enumerable) domains. It develops both an operational semantics and a denotational model based on presheaf categories equipped with a probabilistic local-state monad that is affine and commutative, ensuring the dataflow property and allowing memoization of constant Bernoulli functions. The model proves sound with respect to the operational semantics and demonstrates that memoization respects program transformations that preserve dataflow, bridging traditional measure-theoretic probability and higher-order probabilistic programming. A practical Haskell implementation complements the theory, illustrating the feasibility and guiding future generalizations to broader domain classes and function spaces.
Abstract
Stochastic memoization is a higher-order construct of probabilistic programming languages that is key in Bayesian nonparametrics, a modular approach that allows us to extend models beyond their parametric limitations and compose them in an elegant and principled manner. Stochastic memoization is simple and useful in practice, but semantically elusive, particularly regarding dataflow transformations. As the naive implementation resorts to the state monad, which is not commutative, it is not clear if stochastic memoization preserves the dataflow property -- i.e., whether we can reorder the lines of a program without changing its semantics, provided the dataflow graph is preserved. In this paper, we give an operational and categorical semantics to stochastic memoization and name generation in the context of a minimal probabilistic programming language, for a restricted class of functions. Our contribution is a first model of stochastic memoization of constant Bernoulli functions with a non-enumerable type, which validates data flow transformations, bridging the gap between traditional probability theory and higher-order probability models. Our model uses a presheaf category and a novel probability monad on it.