SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Variational Inference. For VI, the bulk is corrected through the ELBO optimization objective  $E_{q}\log \frac{p(X)}{q(X)}\approx \frac{1}{N}\sum_{i = 1}^{N}\log \frac{p(x_{i})}{q(x_{i})},\quad x_{i}\sim q$ . Since the density  $p$  must also be evaluated, for simplicity, experiments in Table 4 use closed-form marginalized densities for targets. The overall trends also show that GGA yields consistent improvements; the  $\hat{k}$  diagnostic [60] indicates VI succeeds  $(\hat{k}\leq 0.2)$  when a GGA with appropriately matched tails is used and fails  $(\hat{k} >1)$  when Gaussian tails are erroneously imposed.
Table 4: Pareto  $\hat{k}$  diagnostic ([60]) to assess goodness of fit for VI (mean across 100 trials, standard deviation in parenthesis) on targets of varying tail index (smaller  $\alpha =$  heavier tails). A value  $> {0.2}$  is interpreted as potentially problematic so only values not exceeding it are bolded.
Target	α	Normal Affine	Normal Flow	GGA Affine	GGA Flow
Cauchy	α = 2	0.62 (0.26)	0.22 (0.059)	0.68 (0.038)	0.091 (0.04)
IG	α = 2	8.6 (1.8)	8.2 (2.3)	2.0 (0.4)	2.9 (0.71)
StudentT	α = 3	1.2 (0.16)	1.0 (0.43)	1.5 (0.082)	1.3 (0.097)
Chi2	α = ∞	0.57 (0.081)	0.61 (0.067)	0.0093 (0.0067)	0.089 (0.044)
Normal	α = ∞	0.53 (0.17)	0.21 (0.067)	0.4 (0.086)	0.2 (0.089)
Bayesian linear regression. As a practical example of VI applied to posterior distributions, we consider the setting of one-dimensional Bayesian linear regression (BLR) with conjugate priors, defined by the likelihood  $y|X,\beta ,\sigma \sim \mathcal{N}(X\beta ,\sigma^2)$  with a Gaussian prior  $\beta |\sigma^2\sim \mathcal{N}(0,\sigma^2)$  on the coefficients, and an inverse-Gamma prior with parameters  $a_0$  and  $b_{0}$  on the residual variance  $\sigma^2$ . The posterior distribution for  $\beta$  conditioned on  $\sigma^2$  and  $X,y$  is Gaussian. However, conditional on the pair  $(X,y),\sigma^2$  is inverse-Gamma distributed with parameters  $a_0 + \frac{n}{2}$  and  $b_{0} + \frac{1}{2} (y^{\top}y - \mu^{\top}\Sigma \mu))$ , where  $\mu = \Sigma^{-1}X^{\top}X\hat{\beta}$  for  $\hat{\beta}$  the least-squares estimator, and  $\Sigma = X^{\top}X + I$ . Since  $\sigma^2$  is positive, it is typical for PPL implementations to apply an exponential transformation. Hence, a Lipschitz normalising flow starting from a Gaussian initialization will inappropriately approximate the inverse Gamma distributed  $p(\sigma^2 |X,y)$  with log-normal tails. On the other hand, Lipschitz flows starting from a GGA reference distribution will exhibit the correct tails. We assess this discrepancy in Figure 4 under an affine transformation on four subsampled datasets: super (superconductor critical temperature prediction dataset [23] with  $n = 256$  and  $d = 154$ ); who (life expectancy data from the World Health Organisation in the year 2013 [41] with  $n = 130$ ,  $d = 18$ ); air (air quality data [14] with  $n = 6941$ ,  $d = 11$ ); and blog (blog feedback prediction dataset [7] with  $n = 1024$ ,  $d = 280$ ). In Figure 4(i), the GGA-based method seems to perfectly fit to the targets, while in Figure 4(ii), the standard Gaussian approach fails to capture the tail behaviour.
(ii)
Figure 4: Estimated densities for the posterior distribution of  $\sigma^2$  in Bayesian linear regression under optimised exponential + affine transformations from (i) GGA reference, and (ii) Gaussian reference.
Invariant distribution of SGD. For inputs  $X$  and labels  $Y$  from a dataset  $\mathcal{D}$ , the least squares estimator for linear regression satisfies  $\hat{\beta} = \min_{\beta} \frac{1}{2} \mathbb{E}_{X,Y \sim \mathcal{D}} (Y - X\beta)^2$ . To solve for this estimator, one can apply stochastic gradient descent (SGD) sampling over independent  $X_k, Y_k \sim \mathcal{D}$  to obtain the sequence of iterations
$$
\beta_ {k + 1} = \left(I - \delta X _ {k} X _ {k} ^ {\top}\right) \beta_ {k} + \delta Y _ {k} X _ {k}
$$
for a step size  $\delta > 0$ . For large  $\delta$ , the iterates  $\beta_{k}$  typically exhibit heavy-tailed fluctuations [24]. In this regard, this sequence of iterates has been used as a simple model for more general stochastic optimization dynamics [22, 24]. In particular, generalization performance has been tied to the heaviness of the tails in the iterates [47]. Here, we use our algebra to predict the tail behaviour in a simple one-dimensional setting where  $X_{k} \sim \mathcal{N}(0, \sigma^{2})$  and  $Y_{k} \sim \mathcal{N}(0, 1)$ . From classical theory [5], it is known that  $X_{k}$  converges in distribution to a power law with tail exponent  $\alpha > 0$  satisfying  $\mathbb{E}|1 - \delta X_k^2|^{\alpha} = 1$ . In Figure 5, we plot the density of the representative for  $\beta_{10^4}$  obtained using our
algebra against a kernel density estimate using  $10^{6}$  samples when  $\sigma \in \{0.4, 0.5, 0.6\}$  and  $\delta = 2$ . In all cases, the density obtained from the algebra provides a surprisingly close fit.
Figure 5: Kernel density estimate of iterates of SGD (blue) vs. GGA predicted tail behaviour (orange)
Heavy tails and probabilistic machine learning. For studying heavy tails, methods based on subexponential distributions [18] and generalized Pareto distributions (GPD) (or equivalently, regularly varying distributions [49]) have received significant attention historically. For example, [35] presents closure theorems for regularly varying distributions which are special cases of Proposition 1 and Theorem 2. Heavy tails often have a profound impact on probabilistic machine learning methods: in particular, the observation that density ratios  $\frac{p(x)}{q(x)}$  tend to be heavy tailed has resulted in new methods for smoothing importance sampling [53], adaptively modifying divergences [55], and diagnosing VI through the Pareto  $\hat{k}$  diagnostic [60]. These works are complementary to our paper, and our reported results include  $\hat{k}$  diagnostics for VI and  $\hat{\alpha}$  tail index estimates based on GPD.
Our work considers heavy-tailed targets  $p(x)$  which is the same setting as [25, 33]. Whereas those respective works lump the tail parameter in as another variational parameter and may be more generally applicable, the GGA may be applied before samples are drawn and leads to perfectly calibrated tails when applicable.
Probabilistic programming. PPLs can be broadly characterized by the inference algorithms they support, such as: Gibbs sampling over Bayes nets [13, 48], stochastic control flow [19, 58], deep stochastic VI [4, 52], or Hamiltonian Monte-Carlo [8, 59]. Our implementation target beanmachine [50] is a declarative PPL selected due to availability of a PPL compiler and support for static analysis plugins. Similar to [4, 46], it uses PyTorch [40] for GPU tensors and automatic differentiation. Synthesizing an approximating distribution during PPL compilation (Section 3) is also performed in the Stan language by [30] and normalizing flow extensions in [57]. We compare directly against these related density approximators in Section 4.
Static analysis. There is a long history of formal methods and probabilistic programming in the literature [26, 29], with much of the research [10] concerned with defining formal semantics and establishing invariants [54] (see [3] for a recent review). Static analysis uses the abstract syntax tree (AST) representation of a program in order to compute invariants (e.g., the return type of a function, the number of classes implementing a trait) without executing the underlying program. It has traditionally been applied in the context of formalizing semantics [29], and has been used to verify probabilistic programs by ensuring termination, bounding random values values [44]. As dynamic analysis in a PPL is less reliable due to non-determinism, static analysis techniques for PPLs become essential. As recent examples, [32] proposes a static analyzer for the Pyro PPL [4] to verify distribution supports and avoid -Inf log probabilities. More relevant to our work are applications of static analysis to improve inference. [38] and [12] both employ static analysis to inform choice of inference method. However, both works do not account for heavy tails whereas the primary goal of GGA-based analysis is to ensure tails are properly modelled.
In this work, we have proposed a novel systematic approach for conducting tail inferential static PPL analysis. We have done this by defining a heavy-tailed algebra, and by implementing a three-parameter generalized Gamma algebra into a PPL compiler. Initial results are promising, showing that improved inference with simpler approximation families is possible when combined with tail
metadata. While already useful, the generalized Gamma algebra and its implementation currently has some notable limitations:
- The most significant omission to the algebra is classification of log-normal tails. Addition may be treated using [20], but multiplication with log-normal tails remains elusive.
- Since the algebra assumes independence, handling of dependencies between defined random variables must be conducted externally. This can be addressed using a symbolic package to decompose complex expressions into operations on independent random variables.
- Scale coefficients  $\sigma$  for conditional distributions may often be inexact, as exact marginalization in general is NP-hard [28]. Treatment of disintegration using symbolic manipulations is a significant open problem, with some basic developments [9, 45].
- Compile-time static analysis is only applicable to fixed model structures. Control flow, open-universe models [36], and PPLs to support them [4] are an important future research direction.
[1] Søren Asmussen and Hansjorg Albrecher. *Ruin probabilities*, volume 14. World scientific, 2010.
[2] Søren Asmussen, Enkelejd Hashorva, Patrick J Laub, and Thomas Taimre. Tail asymptotics of light-tailed Weibull-like sums. Probability and Mathematical Statistics, 37(2):235-256, 2017.
[3] Ryan Bernstein. Static analysis for probabilistic programs. arXiv preprint arXiv:1909.05076, 2019.
[4] Eli Bingham, Jonathan P Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul Szerlip, Paul Horsfall, and Noah D Goodman. Pyro: Deep universal probabilistic programming. The Journal of Machine Learning Research, 20(1):973-978, 2019.
[5] Dariusz Buraczewski, Ewa Damek, Thomas Mikosch, et al. Stochastic models with power-law tails. Springer Ser. Oper. Res. Financ. Eng., Springer, Cham, 10:978-3, 2016.
[6] Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. arXiv preprint arXiv:1509.00519, 2015.
[7] Krisztian Buza. Feedback prediction for blogs. In Data analysis, machine learning and knowledge discovery, pages 145-152. Springer, 2013.
[8] Bob Carpenter, Andrew Gelman, Matthew D Hoffman, Daniel Lee, Ben Goodrich, Michael Betancourt, Marcus Brubaker, Jiqiang Guo, Peter Li, and Allen Riddell. Stan: A probabilistic programming language. Journal of statistical software, 76(1), 2017.
[9] Kenta Cho and Bart Jacobs. Disintegration and Bayesian inversion via string diagrams. Mathematical Structures in Computer Science, 29(7):938-971, 2019.
[10] Guillaume Claret, Sriram K Rajamani, Aditya V Nori, Andrew D Gordon, and Johannes Borgström. Bayesian inference using data flow analysis. In Proceedings of the 2013 9th Joint Meeting on Foundations of Software Engineering, pages 92-102, 2013.
[11] Aaron Clauset, Cosma Rohilla Shalizi, and Mark EJ Newman. Power-law distributions in empirical data. SIAM review, 51(4):661-703, 2009.
[12] Marco F Cusumano-Towner, Feras A Saad, Alexander K Lew, and Vikash K Mansinghka. Gen: a general-purpose probabilistic programming system with programmable inference. In Proceedings of the 40th acm sigplan conference on programming language design and implementation, pages 221-236, 2019.
[13] Perry de Valpine, Daniel Turek, Christopher J Paciorek, Clifford Anderson-Bergman, Duncan Temple Lang, and Rastislav Bodik. Programming with models: writing statistical algorithms for general model structures with nimble. Journal of Computational and Graphical Statistics, 26(2):403-413, 2017.
[51] David Tolpin, Jan-Willem van de Meent, Hongseok Yang, and Frank Wood. Design and implementation of probabilistic programming language anglican. In Proceedings of the 28th Symposium on the Implementation and Application of Functional programming Languages, pages 1-12, 2016.
[52] Dustin Tran, Matthew W Hoffman, Dave Moore, Christopher Suter, Srinivas Vasudevan, and Alexey Radul. Simple, distributed, and accelerated probabilistic programming. Advances in Neural Information Processing Systems, 31, 2018.
[53] Aki Vehtari, Daniel Simpson, Andrew Gelman, Yuling Yao, and Jonah Gabry. Pareto smoothed importance sampling. arXiv preprint arXiv:1507.02646, 2015.
[54] Di Wang, Jan Hoffmann, and Thomas Reps. Pmaf: an algebraic framework for static analysis of probabilistic programs. ACM SIGPLAN Notices, 53(4):513-528, 2018.
[55] Dilin Wang, Hao Liu, and Qiang Liu. Variational inference with tail-adaptive f-divergence. Advances in Neural Information Processing Systems, 31, 2018.
[56] George Neville Watson. A treatise on the theory of Bessel functions. Cambridge university press, 1995.
[57] Stefan Webb, Jonathan P. Chen, Matrin Jankowiak, and Noah Goodman. Improving automated variational inference with normalizing flows. In ICML Workshop on Automated Machine Learning, 2019.
[58] David Wingate, Andreas Stuhlmüller, and Noah Goodman. Lightweight implementations of probabilistic programming languages via transformational compilation. In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, pages 770-778. JMLR Workshop and Conference Proceedings, 2011.
[59] Kai Xu, Hong Ge, Will Tebbutt, Mohamed Tarek, Martin Trapp, and Zoubin Ghahramani. Advancedhmc. jl: A robust, modular and efficient implementation of advanced hmc algorithms. In Symposium on Advances in Approximate Bayesian Inference, pages 1-10. PMLR, 2020.
[60] Yuling Yao, Aki Vehtari, Daniel Simpson, and Andrew Gelman. Yes, but did it work?: Evaluating variational inference. In International Conference on Machine Learning, pages 5581-5590. PMLR, 2018.