SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhe Zeng
Computer Science Department  
University of California, Los Angeles  
zhezeng@cs.ucla.edu
Computer Science Department University of California, Los Angeles guyvdb@cs.ucla.edu
Bayesian neural networks (BNNs) provide a formalism to quantify and calibrate uncertainty in deep learning. Current inference approaches for BNNs often resort to few-sample estimation for scalability, which can harm predictive performance, while its alternatives tend to be computationally prohibitively expensive. We tackle this challenge by revealing a previously unseen connection between inference on BNNs and volume computation problems. With this observation, we introduce a novel collapsed inference scheme that performs Bayesian model averaging using collapsed samples. It improves over a Monte-Carlo sample by limiting sampling to a subset of the network weights while pairing it with some closed-form conditional distribution over the rest. A collapsed sample represents uncountably many models drawn from the approximate posterior and thus yields higher sample efficiency. Further, we show that the marginalization of a collapsed sample can be solved analytically and efficiently despite the non-linearity of neural networks by leveraging existing volume computation solvers. Our proposed use of collapsed samples achieves a balance between scalability and accuracy. On various regression and classification tasks, our collapsed Bayesian deep learning approach demonstrates significant improvements over existing methods and sets a new state of the art in terms of uncertainty estimation as well as predictive performance.
Uncertainty estimation is crucial for decision making. Deep learning models, including those in safety-critical domains, tend to estimate uncertainty poorly. To overcome this issue, Bayesian deep learning obtains a posterior distribution over the model parameters hoping to improve predictions and provide reliable uncertainty estimates. Among Bayesian inference procedures with neural networks, Bayesian model averaging (BMA) is particularly compelling (Wasserman, 2000; Fragoso et al., 2018; Maddox et al., 2019). However, computing BMAs is distinctly challenging since it involves marginalizing over posterior parameters, which possess some unusual topological properties such as mode-connectivity (Izmailov et al., 2021). We show that even with simple low-dimensional approximate parameter posteriors as uniform distributions, doing BMA requires integrating over highly non-convex and multi-modal distributions with discontinuities arising from non-linear activations (cf. Figure 1a). Accurately approximating the BMA can achieve significant performance gains (Izmailov et al., 2021). Existing methods mainly focus on general-purpose MCMC, which can fail to converge, or provides inaccurate few-sample predictions (Kristiadi et al., 2022), because running longer sampling chains is computationally expensive, and variational approaches that typically use a mean-field approximation that ignores correlations induced by activations (Jospin et al., 2022).
In this work, we are interested in developing collapsed samplers, also known as cutset or Rao-Blackwellised samplers for BMA. A collapsed sampler improves over classical particle-based methods by limiting sampling to a subset of variables and further pairing each sample with a closed-form
(a)  $p(y\mid \boldsymbol {x},\boldsymbol {w})$  being Gaussian. See Example 3.
(b)  $p(y\mid \boldsymbol {x},\boldsymbol {w})$  being triangular. See Section 3.
Figure 1: The integral surface of (a) the expected prediction in BMA, and (b) our proposed approximation. Both are highly non-convex and multi-modal. The z-axis is the weighted prediction  $y p(y \mid x, w) p(w \mid \mathcal{D})$ . Integration of (a) does not admit a closed-form solution, yet integration of (b) is a close approximation that can be solved exactly and efficiently by WMI solvers.
representation of a conditional distribution over the rest whose marginalization is often tractable. Collapsed samplers are effective at variance reduction in graphical models (Koller & Friedman, 2009), however no collapsed samplers are known for Bayesian deep learning. We believe that this is due to the lack of a closed-form marginalization technique congruous with the non-linearity in deep neural networks. Our aim is to overcome this issue and improve BMA estimation by incorporating exact marginalization over (close approximate) conditional distributions into the inference scheme. Nevertheless, scalability and efficiency are guaranteed by the sampling part of our proposed algorithm.
Marginalization is made possible by our observation that BMA reduces to weighted volume computation. Certain classes of such problems can be solved exactly by so-called weighted model integration (WMI) solvers (Belle et al., 2015a). By closely approximating BMA with WMI, these solvers can provide accurate approximations to marginalization in BMA (cf. Figure 1b). With this observation, we propose CIBER, a collapsed sampler that uses WMI for computing conditional distributions. In the few-sample setting, CIBER delivers more accurate uncertainty estimates than the gold-standard Hamiltonian Monte Carlo (HMC) method (cf. Figure 2). We further evaluate the effectiveness of CIBER on regression and classification benchmarks and show significant improvements over other Bayesian deep learning approaches in terms of both uncertainty estimation and accuracy.
2 Bayesian Model Averaging as Weighted Volume Computation
In Bayesian Neural Networks (BNN), given a neural network  $f_{\boldsymbol{w}}$  parameterized by weights  $\boldsymbol{w}$ , instead of doing inference with deterministic  $\boldsymbol{w}$  that optimize objectives such as cross-entropy or mean squared error, Bayesian learning infers a posterior  $p(\boldsymbol{w} \mid \mathcal{D})$  over parameters  $\boldsymbol{w}$  after observing data  $\mathcal{D}$ . During inference, this posterior distribution is then marginalized to produce final predictions. This process is called Bayesian Model Averaging (BMA). It can be seen as learning an ensemble of an infinite number of neural nets and aggregating their results. Formally, given input  $\boldsymbol{x}$ , the posterior predictive distribution and the expected prediction for a regression task are
$$
p (y \mid \boldsymbol {x}) = \int p (y \mid \boldsymbol {x}, \boldsymbol {w}) p (\boldsymbol {w} \mid \mathcal {D}) d \boldsymbol {w}, \quad \text {a n d} \quad \mathbb {E} _ {p (y | \boldsymbol {x})} [ y ] = \int y p (y \mid \boldsymbol {x}) d y. \tag {1}
$$
For classification, the (most likely) prediction is the class  $\arg \max_y p(y\mid x)$ . BMA is intuitively attractive because it can be risky to base inference on a single neural network model. The marginalization in BMA gets around this issue by averaging over models according to a Bayesian posterior.
BMA requires approximations to compute posterior predictive distributions and expected predictions, as the integrals in Equation 1 are intractable in general. Deriving efficient and accurate approximations remains an active research topic (Izmailov et al., 2021). We approach this problem by observing that the marginalization in BMA with ReLU neural networks can be cast as weighted volume computation (WVC). Later we show that it can be generalized to any neural network when combined with sampling. In WVC, various tools exist for solving certain WVC problem classes (Baldoni et al.,
(a) HMC with 10 samples
Figure 2: Uncertainty estimates for regression. The red line is the ground truth. The dark blue line shows the predictive mean. The shaded region is the  $90\%$  confidence interval of the predictive distribution. For the same number of samples, (b) CIBER is closer than (a) small-sample HMC to (c) a highly accurate but slow HMC with a large number of samples. See the Appendix for details.
(b) CIBER with 10 collapsed samples
(c) HMC with  $2k$  samples
2014; Kolb et al., 2019; Zeng et al., 2020c). This section reveals the connection between BMA and WVC. It opens up a new perspective for developing BMA approximations by leveraging WVC tools.
Definition 1 (WVC). A weighted volume computation (WVC) problem is defined by a pair  $(\boxdot, \phi)$  where a region  $\boxdot$  is a conjunction of arithmetic constraints and weight  $\phi: \boxdot \to \mathbb{R}$  is an integrable function assigning weights to elements in  $\boxdot$ . The task of WVC is to compute the integral  $\int_{\boxdot} \phi(x) dx$ .
2.1 A Warm-Up Example
Consider a simple yet relevant setting where the predictive distribution  $p(y \mid x, w)$  is a Dirac delta distribution with zero mass everywhere except at  $f_{\boldsymbol{w}}(\boldsymbol{x})$ , such that  $\int y p(y \mid x, w) dy = f_{\boldsymbol{w}}(\boldsymbol{x})$ .
Example 2. Assume a model  $f_{\mathbf{w}}(x) = \mathsf{ReL}\mathsf{U}(w\cdot x)$  with a uniform posterior over the parameter:  $p(w\mid \mathcal{D}) = \frac{1}{6}$  with  $w\in [-3,3]$ . Let the input be  $x = 1$ . For parameter  $w\in [-3,0]$ , the model  $f_{\mathbf{w}}$  always predicts 0, and otherwise (i.e.,  $w\in (0,3]$ ), it predicts  $w$ . Thus, the expected prediction (Equation 1) is  $\mathbb{E}_{p(y|\mathbf{x})}[y] = \int_{\mathbb{T}_\perp}0\cdot \frac{1}{6} d\mathbf{w} + \int_{\mathbb{T}_\top}w\cdot \frac{1}{6} d\mathbf{w}$ . That is, a summation of two WVC problems  $(\widehat{\mathbb{Q}}_{\perp},0)$  and  $(\widehat{\mathbb{Q}}_{\top},w / 6)$  with  $\widehat{\mathbb{Q}}_{\perp} = (-3\leq w\leq 0)$  and  $\widehat{\mathbb{Q}}_{\top} = (0\leq w\leq 3)$ . The BMA integral decomposes into WVC problems with different weights due to the ReLU activation.
These WVC problems have easy closed-form solutions. This is no longer the case in the following.
Example 3. Assume a model  $f_{\mathbf{w}}(x)$  and posterior distribution  $p(w \mid \mathcal{D})$  as in Example 2. Let the predictive distribution  $p(y \mid x, w)$  be a Gaussian distribution  $p_{\mathcal{N}}(y; f_{\mathbf{w}}(x), 1)$  with mean  $f_{\mathbf{w}}(x)$  and variance 1. Given input  $x = 1$ , the expected prediction (Equation 1) is
$$
\mathbb {E} _ {p (y | x = 1)} [ y ] = \int_ {\mathbb {D} _ {\perp}} y \cdot p _ {\mathcal {N}} (y \mid 0, 1) \cdot \frac {1}{6} d y d w + \int_ {\mathbb {D} _ {\top}} y \cdot p _ {\mathcal {N}} (y \mid w, 1) \cdot \frac {1}{6} d y d w.
$$
It is a summation of two WVC problems with  $\widehat{\mathbb{D}}_{\perp} = (-3 \leq w \leq 0) \wedge (y \in \mathbb{R})$  and  $\widehat{\mathbb{D}}_{\top} = (0 \leq w \leq 3) \wedge (y \in \mathbb{R})$ , whose joint integral surface is shown in Figure 1a.
These WVC problems do not admit closed-form solutions since they involve truncated Gaussian distributions. Moreover, Figure 1a shows that computing BMA, even in such a low-dimensional parameter space, requires integration over non-convex and multi-modal functions.
2.2 General Reduction of BMA to WVC
Let model  $f_{\mathbf{w}}$  be a ReLU neural net. Denote the set of inputs to its ReLU activations by  $\mathcal{R} = \{r_i\}_{i=1}^R$ , where each  $r_i$  is a linear combination of weights. For a given input  $x$ , the parameter space is partitioned by whether each ReLU activation outputs zero or not. This gives the WVC reduction
$$
p (y \mid \boldsymbol {x}) = \sum_ {\boldsymbol {B} \in \{0, 1 \} ^ {R}} \int_ {\boxdot_ {B}} p (y \mid \boldsymbol {x}, \boldsymbol {w}) p (\boldsymbol {w} \mid \mathcal {D}) d \boldsymbol {w},
$$
where  $\pmb{B}$  is a binary vector. The region  $\widehat{\mathbb{D}}_B$  is defined as  $\wedge_{i=1}^{R} \ell_i$  where arithmetic constraint  $\ell_i$  is  $r_i \geq 0$  if  $B_i = 1$  and  $r_i \leq 0$  otherwise. The expected prediction  $\mathbb{E}_{p(y|x)}[y]$  is analogous but includes an additional factor and variable of integration  $y$  in each WVC problem.
This general reduction, however, is undesirable since it amounts to a brute-force enumeration that implies a complexity exponential in the number of ReLU activations. Moreover, not all WVC
problems resulting from this reduction are amenable to existing solvers. We will therefore appeal to a framework called weighted model integration (WMI) that allows for a compact representation of these WVC problems, and a characterization of their tractability for WMI solvers (Kolb et al., 2019). This inspires us to approximate BMA by first reducing it to WVC problems and further closely approximating those with tractable WMI problems.
3 Approximating BMA by WMI
WMI is a modeling and inference framework that supports integration in the presence of logical and arithmetic constraints (Belle et al., 2015a,b). Various WMI solvers have been proposed in recent years (Kolb et al., 2019), ranging from general-purpose ones to others that assume some problem structures to gain scalability. However, even with the reduction from BMA to WVC from the previous section, WMI solvers are not directly applicable. Existing solvers have two main limitations: (i) feasible regions need to be defined by Boolean combinations of linear arithmetic constraints, and (ii) weight functions need to be polynomials. In this section, we show that these issues can be bypassed using a motivating example of how to form a close approximation to BMA using WMI.
In WMI, the feasible region is defined by satisfiability modulo theories (SMT) constraints (Barrett et al., 2010): an SMT formula is a (typically quantifier-free) expression containing both propositional and theory literals connected with logical connectives; the theory literals are often restricted to linear real arithmetic, where literals are of the form  $(\mathbf{c}^T\mathbf{X}\leq b)$  with variable  $\mathbf{X}$  and constants  $\mathbf{c}^T$  and  $b$ .
Example 4. The ReLU model  $f_{\mathbf{w}}(x)$  of Example 2 can be encoded as an SMT formula (see box).
The curly bracket denotes logical conjunction, the symbol  $\Rightarrow$  is a logical implication, variable  $W$  is the weight, and variable  $Z$  denotes the model output.
$$
\boxed {\Delta_ {\text {R e L U}} = \left\{ \begin{array}{l} W \cdot x > 0 \Rightarrow Z = W \cdot x \\ W \cdot x \leq 0 \Rightarrow Z = 0 \end{array} \right.}
$$
The encoding of ReLU neural networks into SMT formulas is explored in existing work to enable verification of the behavior of neural networks and provide formal guarantees (Katz et al., 2017; Huang et al., 2017; Sivaraman et al., 2020). We propose to use this encoding to define the feasible region of WMI problems. Let  $x \models \Delta$  denote the satisfaction of an SMT formula  $\Delta$  by an assignment  $x$ , and  $[[x \models \Delta]]$  be its corresponding indicator function. We formally introduce WMI next.
Definition 5. (WMI) Let  $X$  be a set of continuous random variables. A weighted model integration problem is a pair  $\mathcal{M} = (\Delta, \Phi)$ , where  $\Delta$  is an SMT formula over  $X$  and  $\Phi$  is a set of per-literal weights defined as  $\Phi = \{\phi_{\ell}\}_{\ell \in \mathcal{L}}$ , where  $\mathcal{L}$  is a set of SMT literals and each  $\phi_{\ell}$  is a function defined over variables in literal  $\ell$ . The task of weighted model integration is to compute
$$
\operatorname {W M I} (\Delta , \Phi) = \int_ {\boldsymbol {x} \vDash \Delta} \prod_ {\ell \in \mathcal {L}} \phi_ {\ell} (\boldsymbol {x}) ^ {[ \boldsymbol {x} \vDash \ell ]} d \boldsymbol {x}.
$$
That is, the task is to integrate over the weighted assignments of  $X$  that satisfy the SMT formula  $\Delta$ .<sup>1</sup>
An approximation to the BMA of Example 3 can be achieved with WMI using the following four steps:
Step 1. Encoding model  $f_{w}(x)$ . This has been shown as the SMT formula  $\Delta_{\mathrm{ReLU}}$  in Example 4.
Step 2. Encoding posterior distribution  $p(w \mid \mathcal{D})$ . The uniform distribution  $p(w \mid \mathcal{D}) = \frac{1}{6}$  with  $w \in [-3,3]$  can be encoded as a WMI problem pair  $(\Delta_{pos}, \Phi_{pos})$  as follows:
$$
\boxed {\Delta_ {p o s} = - 3 \leq W \leq 3 \quad \Phi_ {p o s} = \left\{\phi_ {\ell} (W) = \frac {1}{6} w i t h \ell = \text {t r u e} \right\}}
$$
Step 3. Approximate encoding of predictive distribution  $p(y \mid w, x)$ . Recall that  $p(y \mid w, x) = p_{\mathcal{N}}(y; f_{w}(x), 1)$  is Gaussian, which cannot be handled by existing WMI solvers. To approximate it with polynomial densities, we simply use a triangular distribution encoded as a WMI problem pair:
$$
\Delta_ {p r e d} = \left\{ \begin{array}{l l} Y \leq Z + \alpha \\ Y \geq Z - \alpha \end{array} \right. \quad \Phi_ {p r e d} = \left\{ \begin{array}{l l} \phi_ {\ell_ {1}} (Y, Z) = \frac {1 - Y + Z}{\alpha} w i t h \ell_ {1} = Y \geq Z \\ \phi_ {\ell_ {2}} (Y, Z) = \frac {1 + Y - Z}{\alpha} w i t h \ell_ {2} = Y <   Z \end{array} \right\}
$$
In the literature, WMI is defined on mixed discrete-continuous domains. Since we only work with WMI problems over continuous variables, we ignore the discrete ones in the definition for succinctness.
In this encoding,  $\alpha$  is a constant that defines the shape of the triangular distribution. It is obtained by minimizing the  $L2$  distance between a standard normal distribution and the symmetric triangular distribution. We visualize this approximation in the right figure.
Step 4. Approximating BMA by calling WMI solvers. With the
above encodings, the predictive posterior  $p(y \mid x)$  (Equation 1) can be computed using two calls to a WMI solver. For example, the uncertainty of a prediction  $y = 1$  for input  $x = 1$  is
$$
p (y = 1 \mid x = 1) = \mathsf {W M I} (\Delta \wedge (Y = 1), \Phi) / \mathsf {W M I} (\Delta , \Phi) = 0. 1 6 4 / 1,
$$
where  $\Delta = \Delta_{\mathrm{ReLU}} \wedge \Delta_{pos} \wedge \Delta_{pred}$  and  $\Phi = \Phi_{pos} \cup \Phi_{pred}$ . Similarly, the expected prediction  $\mathbb{E}_{p(y|x=1)}[y]$  (Equation 1) can be computed using two calls to a WMI solver:
$$
\mathbb {E} _ {p (y \mid x = 1)} [ y ] = \operatorname {W M I} (\Delta , \Phi^ {*}) / \operatorname {W M I} (\Delta , \Phi) = 0. 7 5 2 / 1,
$$
where  $\Phi^{*} = \Phi \cup \{\phi_{\ell}(Y) = Y$  with  $\ell = \mathrm{true}\}$ . The above formulations also work for unnormalized distributions since the WMI in the denominator serves to compute the partition function.
We visualize the integral surface of the resulting approximate BMA problem in Figure 1b. It is very close to the integral surface of the original BMA problem in Figure 1a. However, it can be exactly integrated using existing WMI solvers while the original one does not admit such solutions. Next, we show how this process can be generalized to a scalable and accurate approximation of BMA.
4 CIBER: Collapsed Inference for Bayesian Deep Learning via WMI
Given a BNN with a large number of weights, naively approximating it by WMI problems can lead to computational issues, since it involves doing integration over polytopes in arbitrarily high dimensions and this is known to be #P-hard (Valiant, 1979; De Loera et al., 2012; Zeng et al., 2020c). Further, weights involved with non-ReLU activation might not be amenable to the WMI encoding. To tackle these issues, we propose to use collapsed samples to combine the strengths from two worlds: the scalability and flexibility from sampling and the accuracy from WMI solvers.
Definition 6. (Collapsed BMA) Let  $(\mathbf{W}_s, \mathbf{W}_c)$  be a partition of parameters  $\mathbf{W}$ . A collapsed sample is a tuple  $(\mathbf{w}_s, q)$ , where  $\mathbf{w}_s$  is an assignment to the sampled parameters  $\mathbf{W}_s$  and  $q$  is a representation of the conditional posterior  $p(\mathbf{W}_c \mid \mathbf{w}_s, \mathcal{D})$  over the collapsed parameter set  $\mathbf{W}_c$ . Given collapsed samples  $\mathcal{S}$ , collapsed BMA estimates the predictive posterior and expected prediction as
$$
\begin{array}{l} p (y \mid \boldsymbol {x}) \approx \frac {1}{| \mathcal {S} |} \sum_ {\left(\boldsymbol {w} _ {s}, q\right) \in \mathcal {S}} \left[ \int p (y \mid \boldsymbol {x}, \boldsymbol {w}) q \left(\boldsymbol {w} _ {c}\right) d \boldsymbol {w} _ {c} \right], a n d \\ \mathbb {E} _ {p (y | \boldsymbol {x})} [ y ] \approx \frac {1}{| \mathcal {S} |} \sum_ {\left(\boldsymbol {w} _ {s}, q\right) \in \mathcal {S}} \left[ \int y p (y \mid \boldsymbol {x}, \boldsymbol {w}) q \left(\boldsymbol {w} _ {c}\right) d \boldsymbol {w} _ {c} d y \right]. \tag {2} \\ \end{array}
$$
The size of the collapsed set  $W_{c}$  determines the trade-off between scalability and accuracy. The more parameters in the collapsed set, the more accurate the approximation to BMA is. The fewer parameters in  $W_{c}$ , the more efficient the computations of the integrals are since the integration is performed in a lower-dimensional space. Later in our experiments, we choose a subset of weights at the last or second-to-last hidden layer of the neural networks to be the collapsed set. This choice is known to be effective in capturing uncertainty as shown in Kristiadi et al. (2020); Snoek et al. (2015).
To develop an algorithm to compute collapsed BMA, we are faced with two main design choice questions: (Q1) how to sample  $\boldsymbol{w}_s$  from the posterior? (Q2) what should be the representation of the conditional posterior  $q$  such that the integrals in Equation 2 can be computed exactly? Next, we provide our answers to these two questions that together give our proposed algorithm CIBER.
4.1 Approximation to Posteriors
For (Q1), we follow Maddox et al. (2019) and sample from the stochastic gradient descent (SGD) trajectory after convergence and use the information contained in SGD trajectories to efficiently approximate the posterior distribution over the parameters of the neural network, leveraging the interpretation of SGD as approximate Bayesian inference (Mandt et al., 2017; Chen et al., 2020).
Given a set of parameter samples  $\mathcal{W}$  from the SGD trajectory, the sample set is defined as  $\mathcal{W}_s = \{\pmb{w}_s \mid \pmb{w} \in \mathcal{W}\}$ . For each assignment  $\pmb{w}_s$ , an approximation  $q(\pmb{W}_c)$  to the conditional posterior  $p(\pmb{W}_c \mid \pmb{w}_s, \mathcal{D})$  is necessary since the posteriors induced by SGD trajectories are implicit. Next, we discuss the choice of approximation to the conditional posterior that is amenable to WMI.
4.2 Encoding into WMI Problems
As shown in Section 3, if a BNN can be encoded as a WMI problem, the posterior predictive distribution and the expected prediction, which involve marginalization over the parameter space, can be computed exactly using WMI solvers. This inspires us to use the WMI framework as the closed-form representation for the conditional posteriors of parameters. The main challenge is how to approximate the integrand in Equation 2 using an SMT formula and a polynomial weight function in order to obtain a WMI problem amenable to existing solvers.
For the conditional posterior approximation  $q(\mathbf{W}_c)$ , we choose it to be a uniform distribution that can be encoded into a WMI problem as  $\mathcal{M}_{pos} = (\Delta_{pos},\Phi_{pos})$  with the SMT formula being  $\Delta_{pos} = \wedge_{i\in c}\left(l_i\leq W_i\leq u_i\right)$  and weights being  $\Phi_{pos} = \{\phi_{\ell}(\mathbf{W}_c) = 1\mid \ell = \mathrm{true}\}$ , where  $l_{i}$  and  $u_{i}$  are domain lower and upper bounds for the uniform distribution respectively. While seemingly over-simplistic, this choice of approximation to the conditional posterior is sufficient to robustly deliver surprisingly strong empirical performance as shown in Section 6. The intuition is that uniform distributions are better than a few samples. We further illustrate this point by comparing the predictive distributions of CIBER and HMC in a few-sample setting. Figure 2 shows that even with the same 10 samples drawn from the posterior distribution, since CIBER further approximates the 10 samples with a uniform distribution, it yields a predictive distribution closer to the ground truth than HMC, indicating that using a uniform distribution instead of a few samples forms a better approximation.