SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Fabian Zaiser
University of Oxford, UK
fabian.zaiser@cs.ox.ac.uk
Andrzej S. Murawski
andrzej.murawski@cs.ox.ac.uk
C.-H. Luke Ong
Nanyang Technological University, Singapore
luke.ong@ntu.edu.sg
We present an exact Bayesian inference method for discrete statistical models, which can find exact solutions to a large class of discrete inference problems, even with infinite support and continuous priors. To express such models, we introduce a probabilistic programming language that supports discrete and continuous sampling, discrete observations, affine functions, (stochastic) branching, and conditioning on discrete events. Our key tool is probability generating functions: they provide a compact closed-form representation of distributions that are definable by programs, thus enabling the exact computation of posterior probabilities, expectation, variance, and higher moments. Our inference method is provably correct and fully automated in a tool called Genfer, which uses automatic differentiation (specifically, Taylor polynomials), but does not require computer algebra. Our experiments show that Genfer is often faster than the existing exact inference tools PSI, Dice, and Prodigy. On a range of real-world inference problems that none of these exact tools can solve, Genfer's performance is competitive with approximate Monte Carlo methods, while avoiding approximation errors.
Bayesian statistics is a highly successful framework for reasoning under uncertainty that has found widespread use in a variety of fields, such as AI/machine learning, medicine and healthcare, finance and risk management, social sciences, climate science, astrophysics, and many other disciplines. At its core is the idea of representing uncertainty as probability, and updating prior beliefs based on observed data via Bayes' law, to arrive at posterior beliefs. A key challenge in Bayesian statistics is computing this posterior distribution: analytical solutions are usually impossible or intractable, which necessitates the use of approximate methods, such as Markov-chain Monte Carlo (MCMC) or variational inference. In this work, we identify a large class of discrete models for which exact inference is in fact possible, in particular time series models of count data, such as autoregressive models, hidden Markov models, switchpoint models. We achieve this by leveraging probability generating functions (GFs) as a representation of distributions. The GF of a random variable  $X$  is defined to be the function  $G(x) \coloneqq \mathbb{E}[x^X]$ . In probability theory, it is a well-known tool to study random variables and their distributions, related to the moment generating and the characteristic functions [14, Chapter 4]. In computer science, GFs have previously been used for the analysis of probabilistic programs [17, 4] and for exact inference on certain classes of graphical models [26, 27].
Here we apply them uniformly in a much more general context via probabilistic programming, enabling exact inference on more expressive Bayesian models.
We characterize the class of supported models with the help of a probabilistic programming language. Probabilistic programming [25] has recently emerged as a powerful tool in Bayesian inference. Probabilistic programming systems allow users to specify complex statistical models as programs in a precise yet intuitive way, and automate the Bayesian inference task. This allows practitioners to focus on modelling, leaving the development of general-purpose inference algorithms to experts. Consequently, probabilistic programming systems such as Stan [3] enjoy increasing popularity among statisticians and data scientists. We describe a programming language, called SGCL (statistical guarded command language), that extends pGCL (probabilistic GCL) [17]. This language is carefully designed to be simple yet expressive, and just restrictive enough to enable exact Bayesian inference on all programs that can be written in it.
Contributions We provide a new framework for exact inference on discrete Bayesian models.
(a) Our method is applicable to a large class of discrete models with infinite support, in particular time series models of count data, such as autoregressive models for population dynamics, Bayesian switchpoint models, mixture models, and hidden Markov models. To our knowledge, no exact inference method for all but the simplest population models was known before.
(b) The models are specified in a probabilistic programming language (PPL), which provides flexibility in model definition, thus facilitating model construction. Our PPL supports stochastic branching, continuous and discrete priors, discrete observations, and conditioning on events involving discrete variables.
(c) Every program written in the language can be translated to a generating function that represents the posterior distribution in an automatic and provably correct way. From this generating function, one can extract posterior mean, variance, and higher moments, as well as the posterior probability masses (for a discrete distribution).
(d) We have built an optimized tool, called Genfer ("GENerating Functions for inFERence"), that takes a probabilistic program as input and automatically computes the aforementioned set of descriptive statistics about the posterior distribution.
(e) We demonstrate that (1) on benchmarks with finite support, Genfer's performance is often better than existing exact inference tools, and (2) on a range of real-world examples that no existing exact tool supports, Genfer is competitive with approximate Monte Carlo methods, while achieving zero approximation error.
Related Work Computing the exact posterior distribution of probabilistic programs is intractable in general as it requires analytical solutions to integrals [9]. For this reason, existing systems either restrict the programming language to allow only tractable constructs (this is our approach) or cannot guarantee successful inference. In the former category are Dice [16], which only supports finite discrete distributions, and SPPL [22], which supports some infinite-support distributions but requires finite discrete priors. Both are based on probabilistic circuits [5] or extensions thereof, which allow for efficient and exact inference. In the latter category are the systems PSI [9] and Hakaru [21], which do not impose syntactic restrictions. They rely on computer algebra techniques to find a closed-form solution for the posterior. Such a form need not exist in general and, even if it does, the system may fail to find it and the running time is unpredictable and unscalable [22]. None of the case studies featured in our evaluation (Section 5.2) can be handled by the aforementioned tools.
Probability generating functions are a useful tool in probability theory to study random variables with infinite support, e.g. in the context of branching processes [7, Chapter 12]. In the context of Bayesian inference, probabilistic generating circuits (PGCs) leverage them to boost the expressiveness of probabilistic circuits [28] and enable efficient inference [15]. However, PGCs cannot handle random variables with infinite support, which is the focus of our approach. Generating functions have also been applied to probabilistic graphical models: Winner and Sheldon [26] find a symbolic representation of the generating function for a Poisson autoregressive model and extract posterior probabilities from it. Subsequently, Winner et al. [27] extend this model to latent variable distributions other than the Poisson distribution, where symbolic manipulation is no longer tractable. Instead they evaluate generating functions and their derivatives using automatic differentiation, which enables exact inference for graphical models. Probabilistic programming is an elegant way of generalizing graphical models, allowing a much richer representation of models [10, 25, 12]. Our contribution here is a new framework for exact inference on Bayesian models via probabilistic programming.
In the context of discrete probabilistic programs without conditioning, Klinkenberg et al. [17] use generating functions to (manually) analyze loop invariants and determine termination probabilities. In follow-up work, Chen et al. [4] extend these techniques to automatically check (under certain restrictions) whether a looping program generates a specified distribution. Their analysis tool Prodigy recently gained the ability to perform Bayesian inference as well [18]. It supports discrete distributions with infinite support (but no continuous priors) and is less scalable than our automatic-differentiation approach (see Section 5) since it relies on computer algebra.
**Limitations** Exact posterior inference is already PSPACE-hard for probabilistic programs involving only finite discrete distributions [16, Section 6]. It follows that our method cannot always be performant and has to restrict the supported class of probabilistic programs. Indeed, our programming language forbids certain constructs, such as nonlinear transformations and observations of continuous variables, in order to preserve a closed-form representation of the generating functions (Section 2). For the same reason, our method cannot compute probability density functions for continuous parameters (but probability masses for discrete parameters and exact moments for all parameters are fine). Regarding performance, the running time of inference is polynomial in the numbers observed in the program and exponential in the number of program variables (Section 4). Despite these limitations, our approach shows that exact inference is possible in the first place, and our evaluation (Section 5) demonstrates support for many real-world models and efficient exact inference in practice.
2 Bayesian Probabilistic Programming
Probabilistic programming languages extend ordinary programming languages with two additional constructs: one for sampling from probability distributions and one for conditioning on observed values. We first discuss a simple program that can be written in our language using a simplified example based on population ecology (cf. [26]).
Suppose you're a biologist trying to estimate the size of an animal population migrating into a new habitat. The immigration of animals is often modeled using a Poisson distribution. You cannot count the animals exhaustively (otherwise we wouldn't need estimation techniques), so we assume that each individual is observed independently with a certain probability; in other words, the count is binomially distributed. For simplicity, we assume that the rate of the Poisson and the probability of the binomial distribution are known, say 20 and 0.1. (For more realistic population models, see Section 5.) As a generative model, this would be written as  $X \sim \text{Poisson}(20); Y \sim \text{Binomial}(X, 0.1)$ .
Suppose you observe  $Y = 2$  animals. The Bayesian inference problem is to compute the posterior distribution  $\mathbb{P}(X = x \mid Y = 2)$  given by Bayes' rule as  $\mathbb{P}(X = x \mid Y = 2) = \frac{\mathbb{P}(X = x) \mathbb{P}(Y = 2 \mid X = x)}{\mathbb{P}(Y = 2)}$ , where  $\mathbb{P}(X = x)$  is called the prior probability,  $\mathbb{P}(Y = 2 \mid X = x)$  the likelihood and  $\mathbb{P}(Y = 2)$  the evidence or normalization constant.
Example 2.1. In our probabilistic programming language, this simplified population model would be expressed as:
$$
X \sim \text {P o i s s o n} (2 0); Y \sim \text {B i n o m i a l} (X, 0. 1); \text {o b s e r v e} Y = 2;
$$
The syntax looks similar to generative model notation except that the observations are expressible as a command. Such a program denotes a joint probability distribution of its program variables, which are viewed as random variables. The program statements manipulate this joint distribution. After the first two sampling statements, the distribution has the probability mass function (PMF)  $p_1(x,y) = \mathbb{P}[X = x, Y = y] = \text{Poisson}(x;20) \cdot \text{Binomial}(y; x, 0.1)$ . Observing  $Y = 2$  restricts this to the PMF  $p_2(x,y) = \mathbb{P}[X = x, Y = y = 2]$ , which equals  $\mathbb{P}[X = x, Y = 2] = p_1(x, 2)$  if  $y = 2$  and 0 otherwise. Note that this is not a probability, but a subprobability, distribution because the total mass is less than 1. So, as a final step, we need to normalize, i.e. rescale the subprobability distribution back to a probability distribution. This corresponds to the division by the evidence in Bayes' rule, yielding the PMF  $p_3(x,y) = \frac{p_2(x,y)}{\sum_{x,y} p_2(x,y)}$ . To obtain the posterior  $\mathbb{P}[X = x \mid Y = 2]$ , which is a distribution of the single variable  $X$ , not the joint of  $X$  and  $Y$ , we need to marginalize  $p_3$  and find  $\sum_{y} p_3(x,y) = \frac{\sum_{y} p_2(x,y)}{\sum_{x,y} p_2(x,y)} = \frac{p_1(x,2)}{\sum_{x} p_1(x,2)} = \frac{\mathbb{P}[X = x, Y = 2]}{\mathbb{P}[Y = 2]} = \mathbb{P}[X = x \mid Y = 2]$ , as desired.
Programming constructs Next we describe our probabilistic programming language more formally. It is based on the probabilistic guarded command language (pGCL) from [17] but augments
it with statistical features like conditioning on events and normalization, which is why we call it Statistical GCL (SGCL). Each program operates on a fixed set of variables  $\mathbf{X} = \{X_{1},\ldots ,X_{n}\}$  taking values in  $\mathbb{R}_{\geq 0}$ . A program consists of a list of statements  $P_{1};P_{2};\dots ;P_{m}$ . The simplest statement is skip, which does nothing and is useful in conditionals (like pass in Python). Variables can be transformed using affine maps, e.g.  $X_{2} := 2X_{1} + 7X_{3} + 2$  (note that the coefficients must be nonnegative to preserve nonnegativity of the variables). Programs can branch on the value of a (discrete) variable (e.g. if  $X_{k}\in \{1,3,7\} \{\dots \}$  else  $\{\dots \}$ ), and sample new values for variables from distributions (e.g.  $X_{k}\sim \mathrm{Binomial}(10,0.5)$  or  $X_{k}\sim \mathrm{Binomial}(X_{j},0.5)$ ). The supported distributions are Bernoulli, Categorical, Binomial, Uniform (both discrete and continuous), NegBinomial, Geometric, Poisson, Exponential, and Gamma. They need to have constant parameters except for the compound distributions Binomial(X,p), NegBinomial(X,p), Poisson(λ·X), and Bernoulli(X), where  $X$  can be a variable. One can observe events involving discrete variables (e.g. observe  $X_{k}\in \{3,5\}$ ) and values from (discrete) distributions directly (e.g. observe  $3\sim \mathrm{Binomial}(X_j,0.5)$ ). Note that observe  $m\sim D$  can be seen as a convenient abbreviation of  $Y\sim D$ ; observe  $Y = m$  with a fresh variable  $Y$ . After an observation, the variable distribution is usually not a probability distribution anymore, but a subprobability distribution ("the numerator of Bayes' rule").
Syntax In summary, the syntax of programs  $P$  has the following BNF grammar:
$$
\begin{array}{l} P := \operatorname {s k i p} \left| P _ {1}; P _ {2} \right| X _ {k} := a _ {1} X _ {1} + \dots + a _ {n} X _ {n} + c | \text {i f} X _ {k} \in A \left\{P _ {1} \right\} \text {e l s e} \left\{P _ {2} \right\} \\ \mid X _ {k} \sim D \mid X _ {k} \sim D (X _ {j}) \mid \text {o b s e r v e} X _ {k} \in A \mid \text {o b s e r v e} m \sim D \mid \text {o b s e r v e} m \sim D (X _ {j}) \\ \end{array}
$$
where  $P, P_{1}$  and  $P_{2}$  are subprograms;  $a_{1}, \ldots, a_{n}, c \in \mathbb{R}_{\geq 0}; m \in \mathbb{N}$ ;  $A$  is a finite subset of the naturals; and  $D$  is a supported distribution. To reduce the need for temporary variables, the abbreviations  $+\sim$ ,  $+=$ , and if  $m \sim D\{\dots\}$  else  $\{\dots\}$  are available (see Section 4). A fully formal description of the language constructs can be found in Appendix A.
Restrictions Our language imposes several syntactic restrictions to guarantee that the generating function of any program admits a closed form, which enables exact inference (see Section 3.1): (a) only affine functions are supported (e.g. no  $X^2$  or  $\exp(X)$ ), (b) only comparisons between variables and constants are supported (e.g. no test for equality  $X = Y$ ), (c) only observations from discrete distributions on  $\mathbb{N}$  and only comparisons of such random variables are supported (e.g. no observe  $1.5 \sim \text{Normal}(0,1)$ ), (d) a particular choice of distributions and their composites is supported, (e) loops or recursion are not supported. A discussion of possible language extensions and relaxations of these restrictions can be found in Appendix B.3.
3 Generating Functions
Consider Example 2.1. Even though the example is an elementary exercise in probability, it is challenging to compute the normalizing constant, because one needs to evaluate an infinite sum:  $\mathbb{P}[Y = 2] = \sum_{x\in \mathbb{N}}\mathbb{P}[Y = 2\mid X = x]$ $\mathbb{P}[X = x] = \sum_{x\in \mathbb{N}}\binom{x}{2}0.1^2 0.9^{x - 2}\cdot \frac{20^x e^{-20}}{x!}$ . It turns out to be  $2e^{-2}$  (see Example 3.1), but it is unclear how to arrive at this result in an automated way. If  $X$  had been a continuous variable, we would even have to evaluate an integral. We will present a technique to compute such posteriors mechanically. It relies on probability generating functions, whose definition includes an infinite sum or integral, but which often admit a closed form, thus enabling the exact computation of posteriors.
Definition Probability generating functions are a well-known tool in probability theory to study random variables and their distributions, especially discrete ones. The probability generating function of a random variable  $X$  is defined to be the function  $G(x) \coloneqq \mathbb{E}[x^X]$ . For a discrete random variable supported on  $\mathbb{N}$ , this can also be written as a power series  $G(x) = \sum_{n \in \mathbb{N}} \mathbb{P}[X = n] \cdot x^n$ . Since we often deal with subprobability distributions, we omit "probability" from the name and refer to  $G(x) \coloneqq \mathbb{E}_{X \sim \mu}[x^X]$ , where  $\mu$  is a subprobability measure, simply as a generating function (GF). For continuous variables, it is often called factorial moment generating function, but we stick with the former name in all contexts. We will use the notation  $\mathrm{gf}(\mu)$  or  $\mathrm{gf}(X)$  for the GF of  $\mu$  or  $X$ . Note that for discrete random variables supported on  $\mathbb{N}$ , the GF is always defined on  $[-1, 1]$ , and for continuous ones at  $x = 1$ , but it need not be defined at other  $x$ .
Probability masses and moments Many common distributions admit a closed form for their GF (see Table 1). In such cases, GFs are a compact representation of a distribution, even if it has infinite or continuous support like the Poisson or Exponential distributions, respectively. Crucially, one can
Table 1: GFs for common distributions with constant (left) and random variable parameters (right)
Distribution D	gf(X ~ D)(x)	Distribution D(Y)	gf(X ~ D(Y))(x)
Binomial(n,p)	(px + 1 - p)n	Binomial(Y,p)	gf(Y)(1 - p + px)
Geometric(p)	p/1-(1-p)x	NegBinomial(Y,p)	gf(Y)(p/1-(1-p)x)
Poisson(λ)	eλ(x-1)	Poisson(λ·Y)	gf(Y)(eλ(x-1))
Exponential(λ)	λ/λ-log x	Bernoulli(Y)	1 + (x - 1) · (gf(Y))'(1)
extract probability masses and moments of a distribution from its generating function. For discrete random variables  $X$ ,  $P[X = n]$  is the  $n$ -th coefficient in the power series representation of  $G$ , so can be computed as the Taylor coefficient at 0:  $\frac{1}{n!} G^{(n)}(0)$  (hence the name "probability generating function"). For a discrete or continuous random variable  $X$ , its expected value is  $\mathbb{E}[X] = G'(1)$  and more generally, its  $n$ -th factorial moment is  $\mathbb{E}[X(X - 1)\dots (X - n + 1)] = \mathbb{E}[\frac{\mathrm{d}^n}{\mathrm{d}x^n} x^X]|_{x = 1} = G^{(n)}(1)$  (hence the name "factorial moment generating function"). The raw and central moments can easily be computed from the factorial moments. For instance, the variance is  $\mathbb{V}[X] = G''(1) + G'(1) - G'(1)^2$ . We will exploit these properties of generating functions through automatic differentiation.
Multivariate case The definition of GFs is extended to multidimensional distributions in a straightforward way: for random variables  $\mathbf{X} = (X_{1},\ldots ,X_{n})$ , their GF is the function  $G(\pmb {x})\coloneqq \mathbb{E}[\pmb{x}^{\pmb{X}}]$  where we write  $\pmb {x}\coloneqq (x_{1},\dots,x_{n})$  and  $\pmb{x}^{\pmb{X}}\coloneqq x_1^{X_1}\dots x_n^{X_n}$ . We generally follow the convention of using uppercase letters for random and program variables, and lowercase letters for the corresponding parameters of the generating function. Marginalization can also be expressed in terms of generating functions: to obtain the GF  $\tilde{G}$  of the joint distribution of  $(X_{1},\ldots ,X_{n - 1})$ , i.e. to marginalize out  $X_{n}$ , one simply substitutes 1 for  $x_{n}$  in  $G$ :  $\tilde{G} (x_{1},\ldots ,x_{n - 1}) = G(x_{1},\ldots ,x_{n - 1},1)$ . This allows us to compute probability masses and moments of a random variable in a joint distribution: we marginalize out all the other variables and then use the previous properties of the derivatives.
3.1 Translating programs to generating functions
The standard way of describing the meaning of probabilistic programs is assigning (sub-)probability distributions to them. An influential example is Kozen's distribution transformer semantics [19], where each program statement transforms the joint distribution of all the variables  $X = (X_{1},\ldots ,X_{n})$ . Since Kozen's language does not include observations, we present the full semantics in Appendix A. We call this the standard semantics of a probabilistic program and write  $[[P]]_{\mathrm{std}}(\mu)$  for the transformation of  $\mu$  by the program  $P$ . As a last step, the subprobability distribution  $\mu$  has to be normalized, which we write  $\mathrm{normalize}(\mu)\coloneqq \frac{\mu}{\int\mathrm{d}\mu}$ . This reduces the Bayesian inference problem for a given program to computing its semantics, starting from the joint distribution  $\mathrm{Dirac}(\mathbf{0}_n)$  in which all  $n$  variables are initialized to 0 with probability 1. While mathematically useful, the distribution transformer semantics is hardly amenable to computation as it involves integrals and infinite sums.
Instead, we shall compute the generating function of the posterior distribution represented by the probabilistic program. Then we can extract posterior probability masses and moments using automatic differentiation. Each statement in the programming language transforms the generating function of the distribution of program states, i.e. the joint distribution of the values of the variables  $X = (X_{1},\ldots ,X_{n})$ . Initially, we start with the constant function  $G = 1$ , which corresponds to all variables being initialized with 0 since  $\mathbb{E}[x^0 ] = 1$
The generating function semantics of a program  $\llbracket P\rrbracket_{\mathrm{gf}}$  describes how to transform  $G$  to the generating function  $\llbracket P\rrbracket_{\mathrm{gf}}(G)$  for the distribution at the end. It is defined in Table 2, where the update notation  $\pmb{x}[i\mapsto a]$  denotes  $(x_{1},\ldots ,x_{i - 1},a,x_{i + 1},\ldots ,x_{n})$  and  $\mathbf{1}_n$  means the  $n$ -tuple  $(1,\dots ,1)$ . The first five rules were already described in [4], so we only explain them briefly: skip leaves everything unchanged and  $P_{1};P_{2}$  chains two statements by transforming with  $\llbracket P_1\rrbracket_{\mathrm{gf}}$  and then  $\llbracket P_2\rrbracket_{\mathrm{gf}}$ . To explain linear assignments, consider the case of only two variables:  $X_{1}:= 2X_{1} + 3X_{2} + 5$ . Then
$$
\llbracket P \rrbracket_ {\mathbf {g f}} (G) (\boldsymbol {x}) = \mathbb {E} \left[ x _ {1} ^ {2 X _ {1} + 3 X _ {2} + 5} x _ {2} ^ {X _ {2}} \right] = x _ {1} ^ {5} \mathbb {E} \left[ \left(x _ {1} ^ {2}\right) ^ {X _ {1}} \left(x _ {1} ^ {3} x _ {2}\right) ^ {X _ {2}} \right] = x _ {1} ^ {5} \cdot G \left(x _ {1} ^ {2}, x _ {1} ^ {3} x _ {2}\right).
$$
For conditionals if  $X_{k} \in A\{P_{1}\}$  else  $\{P_2\}$ , we split the generating function  $G$  into two parts: one where the condition is satisfied  $(G_{X_k \in A})$  and its complement  $(G - G_{X_k \in A})$ . The former is transformed