SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 1: Hyperparameters used for the experiments
Dataset
Model	CIFAR10
ResNet20	MLP	MNIST
MLP
Data Aug.	False	False	False
Batch size	500	500	500
Temperature	1.0	1.0	1.0
Test samples	128	128	128
Train samples	1	1	1
VI std. init	0.01	0.01	0.01
Base features	16	512	512
Prior var	0.01	0.0025	0.01
Learning rate	0.000001	0.000001	0.000001
Train epochs	1000	1000	1000
et al. [2], we don't use data augmentation. A possible protocol for handling data augmentation in BNNs is presented by Osawa et al. [77] and involves carefully tuning the likelihood temperature to correctly counting the number of data points.
Table 2: MLP
Layer	Dimensions
Linear-ReLU	512 × Din
Linear-ReLU	512 × 512
Linear-ReLU	512 × 512
Linear-Softmax	Dout × 512
Table 3: ResNet20
Layer	Dimensions
Conv2D	16 × 3 × 3 × Din
Residual Block	3 × 3,16	× 3
3 × 3,16
Residual Block	3 × 3,32	× 3
3 × 3,32
Residual Block	3 × 3,64	× 3
3 × 3,64
AvgPool	8 × 8
Linear-Softmax	Dout × 64
C.1 Computing platform
The experiments have been performed using JAX [15] and run on two AWS p4d.24xlarge instances with 8 NVIDIA A100 GPUs. Experiments were conducted using in the eu-west-1 region, which has a carbon efficiency of  $0.62\mathrm{kgCO_2eq / kWh}$ . A cumulative of 6500 hours of computation was performed on GPUs and it includes interactive sessions as well as small experiments with very low GPU usage, providing a pessimistic estimation of the true utilization. Total emissions are estimated to be 1007.5  $\mathrm{kgCO_2eq}$  of which 100 percents were directly offset by AWS.
We present timings obtained by profiling the time needed to solve the SOLAP with the Wasserstein cost, as well as the time for the deterministic case [2]. In Fig. A14 we show the results for MLP and ResNet20 architectures, varying the model width. It is evident that, in the majority of cases, the algorithm completes within a minute. Moreover, as anticipated, in case of VI solving our distribution alignment problem for wide neural networks is more computationally demanding compared to merely matching weights from SGD solutions.
Finally, we also test our setup on the CIFAR100 dataset [56]. Surprisingly, we were not able to replicate the same level of performance as in the other cases. In Fig. A15, we see that, despite converging well, we fall short to find zero-barrier solutions. Similarly to the comments of Ainsworth et al. [2], we also stress that the failure to align distributions does not rule out the existence of a proper permutation map that the algorithm couldn't find. Nonetheless, this raises a number of questions: the Bayesian posterior is the product of two ingredients, the prior and the likelihood, conditioned to observing a dataset.
Figure A14: Timings. Profile of the algorithms to align the VI solutions and to match weights from SGD solutions.
Figure A15: Alignment failure. The method proposed fails to recover zero-barrier solutions for CIFAR100.
Finally, as an additional check, we analyze the log-posterior by projecting the density into two dimensional slices, following the setup in [47, 36]. We study the two dimensional subspace of the parameter space supported by the hyperplane  $H$  of the form
$$
H = \left\{\boldsymbol {\theta} \in \mathbb {R} ^ {d} \mid \boldsymbol {\theta} = a \boldsymbol {\theta} _ {a} + b \boldsymbol {\theta} _ {b} + (1 - a - b) \boldsymbol {\theta} _ {c} \right\},
$$
where  $a, b \in \mathbb{R}$  and  $\theta_{a}, \theta_{b}$  and  $\theta_{c}$  are the means of  $q_{0}$ ,  $q_{1}$  and  $P_{\#} q_{1}$ . With this configuration, all three solutions lie on this hyper-plane. In Fig. 6, we present the visualization of ResNet20 trained on CIFAR10. We see that the distributions  $q_{0}$  and  $P_{\#} q_{1}$  are connected by higher density regions than the ones between  $q_{0}$  and  $q_{1}$ . Also, as expected the symmetries arise from the form of the likelihood, and the prior has a comparable strength with respect to the three posteriors. Later, we also study in more details the effect of the prior's variance in finding low-barrier solutions.
Log-posterior
Figure A16: Posterior density visualization. All three solutions are local approximation of posterior, but  $q_{0}$  and  $P_{\#}q_{1}$  are connected by lower density regions.
E A primer on variational inference for Bayesian neural networks
VI is a classic tool to tackle intractable Bayesian inference [50, 13]. VI casts the inference problem into an optimization-based procedure to compute a tractable approximation of the true posterior. Assume a generic parametric model  $f$  parameterized by some unknown parameters  $\pmb{\theta}$  (i.e.  $f(\cdot, \pmb{\theta})$ ) and a collection of data  $\pmb{y} \in \mathbb{R}^N$  corresponding to some input points  $\pmb{X} = \{\pmb{x}_i \mid \pmb{x}_i \in \mathbb{R}^{D_{\mathrm{in}}} \}_{i=1,\dots,N}$ . In our setting, we have a probabilistic model  $p(\pmb{y} \mid f(\pmb{X}; \pmb{\theta}))$  with parameters  $\pmb{\theta}$ , a prior distributions on them  $p(\pmb{\theta})$  and a set of observations  $\{\pmb{X}, \pmb{y}\}$ . In a nutshell, the general recipe of VI consists of (i) introducing a set  $\mathcal{Q}$  of distributions; (ii) defining a tractable objective that "measure" the distance between any arbitrary distribution  $q(\pmb{\theta}) \in \mathcal{Q}$  and the true posterior  $p(\pmb{\theta} \mid \pmb{y})$ ; and finally (iii) providing a programmatic way to find the distribution  $\widetilde{q}(\pmb{\theta})$  that minimizes such distance. In practice,  $q(\pmb{\theta})$  has some free parameters  $\nu$  (also known as variational parameters), which are optimized such that
the approximating distribution  $q(\pmb{\theta}; \pmb{\nu})$  is as closer as possible to the true posterior  $p(\pmb{\theta} \mid \pmb{y})$ . We can derive the variational objective starting from the definition of the KL,
$$
\begin{array}{l} \operatorname {K L} [ q (\boldsymbol {\theta}; \boldsymbol {\nu}) \| p (\boldsymbol {\theta} \mid \boldsymbol {y}) ] = \mathbb {E} _ {q (\boldsymbol {\theta}; \boldsymbol {\nu})} [ \log q (\boldsymbol {\theta}; \boldsymbol {\nu}) - \log p (\boldsymbol {\theta} \mid \boldsymbol {y}) ] = \tag {22} \\ = \mathbb {E} _ {q (\boldsymbol {\theta}; \boldsymbol {\nu})} [ \log q (\boldsymbol {\theta}; \boldsymbol {\nu}) - \log p (\boldsymbol {y} | \boldsymbol {\theta}) - \log p (\boldsymbol {\theta}) ] + \log p (\boldsymbol {y}) \\ \end{array}
$$
Rearranging we have that
$$
\log p (\boldsymbol {y}) - \mathrm {K L} [ q (\boldsymbol {\theta}; \boldsymbol {\nu}) \| p (\boldsymbol {\theta} \mid \boldsymbol {y}) ] = \mathbb {E} _ {q (\boldsymbol {\theta}; \boldsymbol {\nu})} [ \log q (\boldsymbol {\theta}; \boldsymbol {\nu}) - \log p (\boldsymbol {y} \mid \boldsymbol {\theta}) - \log p (\boldsymbol {\theta}) ] \tag {23}
$$
The r.h.s. of the equation defines our variational objective, also known as ELBO, that can be arranged as follows,
$$
\mathcal {L} _ {\mathrm {E L B O}} (\boldsymbol {\nu}) = \underbrace {\mathbb {E} _ {q (\boldsymbol {\theta} ; \boldsymbol {\nu})} \log p (\boldsymbol {y} \mid \boldsymbol {\theta})} _ {\text {M o d e l f i t t i n g t e r m}} - \underbrace {\operatorname {K L} [ q (\boldsymbol {\theta} ; \boldsymbol {\nu}) \| p (\boldsymbol {\theta}) ]} _ {\text {R e g u l a r i z a t i o n t e r m}}. \tag {24}
$$
This formulation highlights the property of this objective, which is made of two components: the first one is the expected log-likelihood under the approximate posterior  $q$  and measures how the model fits the data. The second term, on the other hand, has the regularization effect of penalizing posteriors that are far from the prior as measured by the KL. Before diving into the challenges of optimization of the ELBO, we shall spend a brief moment discussing the form of the approximating distribution  $q$ . One of the simplest and easier choice is the mean field approximation [42], where each variable  $\theta_{i}$  is taken to be independent with respect to the remaining  $\theta_{-i}$ . Effectively, this imposes a factorization of the posterior,
$$
q (\boldsymbol {\theta}; \boldsymbol {\nu}) = \prod_ {i = 1} ^ {K} q \left(\theta_ {i}; \boldsymbol {\nu} _ {i}\right) \tag {25}
$$
where  $\pmb{\nu}_i$  is the set of variational parameters for the parameter  $\theta_i$ . On top of this approximation,  $q(\theta_i)$  is often chosen to be Gaussian,
$$
q \left(\theta_ {i}\right) = \mathcal {N} \left(\mu_ {i}, \sigma_ {i} ^ {2}\right) \tag {26}
$$
Now, the collection of all means and variances  $\{\mu_i,\sigma_i^2\}_{i = 1}^K$  defines the set of variational parameters to optimize.
For BNNs the analytic evaluation of the ELBO (and its gradients) is always untractable due the nonlinear nature of the expectation of the log-likelihood under the variational distribution. Nonetheless, this can be easily estimated via Monte Carlo integration [72], by sampling  $N_{\mathrm{MC}}$  times from  $q_{\nu}$ ,
$$
\mathbb {E} _ {q (\boldsymbol {\theta}; \boldsymbol {\nu})} \log p (\boldsymbol {y} \mid \boldsymbol {\theta}) \approx \frac {1}{N _ {\mathrm {M C}}} \sum_ {j = 1} ^ {N _ {\mathrm {M C}}} \log p (\boldsymbol {y} \mid \widetilde {\boldsymbol {\theta}} _ {j}), \quad \text {w i t h} \quad \widetilde {\boldsymbol {\theta}} _ {j} \sim q (\boldsymbol {\theta}; \boldsymbol {\nu}) \tag {27}
$$
In practice, this is as simple as re-sampling the weights and the biases for all the layers  $N_{\mathrm{MC}}$  times and computing the output for each new sample.
We now have a tractable objective that needs to be optimized with respect to the variational parameters  $\pmb{\nu}$ . Very often the KL term is known, making its differentiation trivial. On the other hand the expectation of the likelihood is not available, making the computation of its gradients more challenging. This problem can be solved using the so-called reparameterization trick [92, 54]. The reparameterization trick aims at constructing  $\theta$  as an invertible function  $\mathcal{T}$  of the variational parameters  $\pmb{\nu}$  and of another random variable  $\varepsilon$ , so that  $\theta = \mathcal{T}(\varepsilon; \pmb{\nu})$ . Generally, a  $\mathcal{T}$  that suits this constraint might not exist; Ruiz et al. [90] discuss how to build "weakly" dependent transformation  $\mathcal{T}$  for distributions like Gamma, Beta and Log-normal. For discrete distributions, instead, one could use a continuous relaxation, like the Concrete [69].  $\varepsilon$  is chosen such that its marginal  $p(\varepsilon)$  does not depend on the variational parameters. With this parameterization,  $\mathcal{T}$  separates the deterministic components of  $q$  from the stochastic ones, making the computation of its gradient straightforward. For a Gaussian distribution with mean  $\mu$  and variance  $\sigma^2$ ,  $\mathcal{T}$  corresponds to as simple scale-location transformation of an isotropic Gaussian noise,
$$
\theta \sim \mathcal {N} (\mu , \sigma^ {2}) \Longleftrightarrow \theta = \mu + \sigma \varepsilon \quad \text {w i t h} \quad \varepsilon \sim \mathcal {N} (0, 1). \tag {28}
$$
This simple transformation ensures that  $p(\varepsilon) = \mathcal{N}(0,1)$  does not depend on the variational parameters  $\pmb{\nu} = \{\mu, \sigma^2\}$ . The gradients of the ELBO can be therefore computed as
$$
\nabla_ {\boldsymbol {\nu}} \mathcal {L} _ {\mathrm {E L B O}} = \mathbb {E} _ {p (\varepsilon)} \left[ \nabla_ {\boldsymbol {\theta}} \log p (\boldsymbol {y} \mid \boldsymbol {\theta}) \mid_ {\boldsymbol {\theta} = \mathcal {T} (\varepsilon ; \boldsymbol {\nu})} \nabla_ {\boldsymbol {\nu}} \mathcal {T} (\varepsilon ; \boldsymbol {\nu}) \right] - \nabla_ {\boldsymbol {\nu}} \mathrm {K L} [ q (\boldsymbol {\theta}; \boldsymbol {\nu}) \| p (\boldsymbol {\theta}) ]. \tag {29}
$$
The gradient  $\nabla_{\theta}\log p(\pmb {y}|\pmb {\theta})$  depends on the model and it can be derived with automatic differentiation tools [1, 78], while  $\nabla_{\nu}\mathcal{T}(\varepsilon ;\nu)$  doesn't have any stochastic components and therefore can be known deterministically. Note that the reparameterization trick can be also used when the KL is not analytically available. In that case, we would end up with,
$$
\nabla_ {\boldsymbol {\nu}} \mathcal {L} _ {\mathrm {E L B O}} = \mathbb {E} _ {p (\varepsilon)} \left[ \nabla_ {\boldsymbol {\theta}} \log p (\boldsymbol {y} \mid \boldsymbol {\theta}) + \log q (\boldsymbol {\theta}; \boldsymbol {\nu}) - \log p (\boldsymbol {\theta}) \right] _ {\boldsymbol {\theta} = \mathcal {T} (\varepsilon ; \boldsymbol {\nu})} \nabla_ {\boldsymbol {\nu}} \mathcal {T} (\varepsilon ; \boldsymbol {\nu}) \tag {30}
$$
Roeder et al. [86] argue that when we believe that  $q(\pmb{\theta}; \pmb{\nu}) \approx p(\pmb{y} \mid \pmb{\theta})$ , Eq. (30) should be preferred over Eq. (29) even if computing analytically the KL is possible. Note that this case is very unlikely for BNN posteriors, and that the additional randomness introduced by the Monte Carlo estimation of the KL could be harmful.
In case of large datasets and complex models, the formulation summarized in Eq. (29) can be computationally challenging, due to the evaluation of the likelihood and its gradients  $N_{\mathrm{MC}}$  times. Assuming factorization of the likelihood,
$$
p (\boldsymbol {y} \mid \boldsymbol {\theta}) = p (\boldsymbol {y} \mid f (\boldsymbol {X}; \boldsymbol {\theta})) = \prod_ {i = 1} ^ {N} p \left(y _ {i} \mid f \left(\boldsymbol {x} _ {i}; \boldsymbol {\theta}\right)\right) \tag {31}
$$
this quantity can be approximated using mini-batching [39, 43]. Recalling  $\mathbf{y}$  as the set of labels of our dataset with  $N$  examples, by taking  $\mathcal{B} \subset \mathbf{y}$  as a random subset of  $\mathbf{y}$ , the likelihood term can be estimated in an unbiased way as
$$
\log p _ {\boldsymbol {\theta}} (\boldsymbol {y} \mid \boldsymbol {\theta}) \approx \frac {N}{M} \sum_ {y _ {i} \sim \mathcal {B}} \log p \left(y _ {i} \mid \boldsymbol {\theta}\right). \tag {32}
$$
where  $M$  is the number of points in the minibatch. At the cost of increase "randomness", we can use Eq. (29) to compute the gradients of the ELBO with the minibatch formulation in Eq. (32). Stochastic optimization, e.g. any version of SGD, will converge to a local optimum provided with a decreasing learning rate and sufficient gradient updates [85].