SlowGuess/NeurIPS_2023_OCR huggingface dataset:
approximating the marginal likelihood, posterior, and predictive in deep learning [Daxberger et al., 2021]. In comparison to other approximate inference methods, it can rely on effective training algorithms developed for deep learning and also offers a differentiable marginal likelihood estimate that enables empirical Bayes (EB). Efficient curvature approximations further make it scalable to deep learning [Ritter et al., 2018]. We extend these to the heteroscedastic regression setting.
Laplace approximates the posterior locally at a mode of the posterior with a Gaussian distribution,  $p(\boldsymbol{\theta}|\mathcal{D},\boldsymbol{\delta}) \approx \mathcal{N}(\boldsymbol{\theta};\boldsymbol{\theta}_*,\boldsymbol{\Sigma})$ . The mean is given by a stationary point of the posterior,  $\boldsymbol{\theta}_{*} = \arg \max_{\boldsymbol{\theta}}\log p(\mathcal{D},\boldsymbol{\theta}|\boldsymbol{\delta})$ , and the covariance by the curvature at that mode,  $\boldsymbol{\Sigma}^{-1} = \nabla_{\boldsymbol{\theta}}^{2}\log p(\mathcal{D},\boldsymbol{\theta}|\boldsymbol{\delta})|_{\boldsymbol{\theta} = \boldsymbol{\theta}_{*}}$ . This is due to a second-order Taylor approximation of the log posterior around the mode. The mean is hence the result of neural network training, however, the covariance requires estimating and inverting a Hessian, which is typically intractable.
4.1 Linearized Laplace for Natural Heteroscedastic Regression
The linearized Laplace approximation [MacKay, 1995, Khan et al., 2019, Foong et al., 2019, Immer et al., 2021b] overcomes issues of the vanilla Laplace approximation. Linearizing the neural network about the parameters at the mode, the Hessian, which is the generalized Gauss-Newton in this case [Martens, 2020], becomes positive semidefinite and offers efficient structured approximations. In particular, we have the following Hessian approximation due to linearization
$$
\boldsymbol {\Sigma} ^ {- 1} \approx \sum_ {n = 1} ^ {N} \mathbf {J} _ {*} (\mathbf {x} _ {n}) ^ {\top} \left[ - \nabla_ {\boldsymbol {\eta}} ^ {2} \log p \left(y _ {n} | \mathbf {x} _ {n}, \boldsymbol {\theta} _ {*}\right) \right] \mathbf {J} _ {*} (\mathbf {x} _ {n}) + \nabla_ {\boldsymbol {\theta}} ^ {2} \log p (\boldsymbol {\theta} | \boldsymbol {\delta}) | _ {\boldsymbol {\theta} = \boldsymbol {\theta} _ {*}}, \tag {7}
$$
where  $[\mathbf{J}_*(\mathbf{x})]_{cp} = \frac{\partial\eta_c(\mathbf{x};\pmb{\theta})}{\partial\theta_p}|_{\pmb {\theta} = \pmb{\theta}_*}$  is the Jacobian of the neural network  $\pmb {\eta}(\mathbf{x};\pmb {\theta})$  at the mode. The first summand is the generalized Gauss-Newton and the second term is the Hessian of the prior, which is simply a diagonal matrix constructed from entries  $\delta_l > 0$ , the layer-wise regularization parameters.
Due to the natural parametrization of the likelihood, the Hessian approximation of the linearized Laplace approximation is guaranteed to be positive definite. Since the log prior Hessian is diagonal with entries  $\delta_l > 0$ , one only has to show that the negative log likelihood Hessian,  $-\nabla_{\eta}^{2}\log p(y|\mathbf{x},\boldsymbol{\theta}_{*})$  is positive semidefinite, which is simply a property of naturally parameterized exponential families [Martens, 2020]. Note that this would not be the case for the mean-variance parametrization<sup>3</sup>. For efficient computation, we can decompose the log likelihood Hessian as
$$
\begin{array}{l} \boldsymbol {\Lambda} _ {*} (\mathbf {x}) \stackrel {\text {d e f}} {=} \nabla_ {\boldsymbol {\eta}} ^ {2} \log p (y | \mathbf {x}, \boldsymbol {\theta} _ {*}) = \left[ \begin{array}{c c} \frac {1}{2 \eta_ {2} (x ; \boldsymbol {\theta} _ {*})} & - \frac {\eta_ {1} (x ; \boldsymbol {\theta} _ {*})}{2 \eta_ {2} (x ; \boldsymbol {\theta} _ {*}) ^ {2}} \\ - \frac {\eta_ {1} (x ; \boldsymbol {\theta} _ {*})}{2 \eta_ {2} (x ; \boldsymbol {\theta} _ {*}) ^ {2}} & \frac {\eta_ {1} (x ; \boldsymbol {\theta} _ {*}) ^ {2}}{2 \eta_ {2} (x ; \boldsymbol {\theta} _ {*}) ^ {3}} - \frac {1}{2 \eta_ {2} (x ; \boldsymbol {\theta} _ {*}) ^ {2}} \end{array} \right] (8) \\ = \left[ - \frac {1}{\sqrt {- 2 \eta_ {2} (x ; \boldsymbol {\theta} _ {*})}} \quad \frac {\eta_ {1} (x ; \boldsymbol {\theta} _ {*})}{\eta_ {2} (x ; \boldsymbol {\theta} _ {*}) \sqrt {- 2 \eta_ {2} (x ; \boldsymbol {\theta} _ {*})}} \right] ^ {2} - \left[ 0 \quad \frac {1}{\sqrt {2} \eta_ {2} (x ; \boldsymbol {\theta} _ {*})} \right] ^ {2} (9) \\ \stackrel {\text {d e f}} {=} \lambda_ {1} (\mathbf {x}) \lambda_ {1} (\mathbf {x}) ^ {\top} - \lambda_ {2} (\mathbf {x}) \lambda_ {2} (\mathbf {x}) ^ {\top}, (10) \\ \end{array}
$$
where the square indicates the outer product. This alleviates the need to compute this matrix and instead allows to work with outer products of Jacobian-vector products to compute the covariance in Equation 7. However, the full Hessian approximation remains quadratic in the number of neural network weights. We tackle this issue with a scalable Kronecker-factored approximation.
4.2 Scalable Kronecker-Factored Hessian Approximation
To enable the application of the Laplace approximation to heteroscedastic deep neural networks and large datasets, we use a layer-wise Kronecker-factored approximation [KFAC; Martens and Grosse, 2015, Botev et al., 2017]. To overcome the quadratic scaling in the number of parameters, KFAC makes two efficient approximations: first, it constructs a block-diagonal approximation to  $\mathbf{H}$  from blocks  $\mathbf{H}_l$  per layer  $l$ . Second, each block  $\mathbf{H}_l$  is approximated as a Kronecker product that enables efficient storage and computation using only the individual factors. In the following, we revise KFAC for linear layers and define it for the specific case of the heteroscedastic natural Gaussian likelihood. The same derivation applies similarly to other layer types [Osawa, 2021].
Using the Hessian decomposition of the natural log likelihood in Equation 10, we derive a KFAC approximation that can efficiently be computed in a closed-form. We can write the Jacobian of a
3In Sec. 4.5 we show how Laplace can be applied for the mean-variance parametrization despite that.
fully connected layer that maps a  $D$  to a  $D'$ -dimensional representation as a Kronecker product  $\mathbf{J}_l(\mathbf{x}_n)^\top = \mathbf{a}_{l,n} \otimes \mathbf{g}_{l,n}$  with  $\mathbf{a}_{l,n} \in \mathbb{R}^{D \times 1}$  as the layer's input and  $\mathbf{g}_{l,n} \in \mathbb{R}^{D' \times 2}$  as transposed Jacobian w.r.t. the output, both for the input  $\mathbf{x}_n$ . Following Martens and Grosse [2015] and Botev et al. [2017], we then have the KFAC approximation
$$
\begin{array}{l} [ \boldsymbol {\Sigma} ^ {- 1} ] _ {l} = \sum_ {n = 1} ^ {N} [ \mathbf {a} _ {l, n} \otimes \mathbf {g} _ {l, n} ] \boldsymbol {\Lambda} _ {n} [ \mathbf {a} _ {l, n} \otimes \mathbf {g} _ {l, n} ] ^ {\intercal} + \delta_ {l} \mathbf {I} = \sum_ {n = 1} ^ {N} [ \mathbf {a} _ {l, n} \mathbf {a} _ {l, n} ^ {\intercal} ] \otimes [ \mathbf {g} _ {l, n} \boldsymbol {\Lambda} _ {n} \mathbf {g} _ {l, n} ^ {\intercal} ] + \delta_ {l} \mathbf {I} \\ \approx \frac {1}{N} \left[ \sum_ {n = 1} ^ {N} \mathbf {a} _ {l, n} \mathbf {a} _ {l, n} ^ {\intercal} \right] \otimes \left[ \sum_ {n = 1} ^ {N} \sum_ {k = 1} ^ {2} \mathbf {g} _ {l, n} \boldsymbol {\lambda} _ {n, k} \boldsymbol {\lambda} _ {n, k} ^ {\intercal} \mathbf {g} _ {l, n} ^ {\intercal} \right] + \delta_ {l} \mathbf {I} \stackrel {\mathrm {d e f}} {=} \mathbf {A} _ {l} \otimes \mathbf {B} _ {l} + \delta_ {l} \mathbf {I}, \tag {11} \\ \end{array}
$$
where  $\lambda_{n,k} \in \mathbb{R}^{2 \times 1}$  is due to the decomposition of  $\Lambda_{n}$  into outer products Equation 10 and the approximation is due to exchanging the sum and product. Conveniently, the terms  $\mathbf{g}_{l,n} \boldsymbol{\lambda}_{n,k}$  can be computed efficiently using two Jacobian-vector products and the Kronecker factors can then be extracted within the second-order framework of Osaka [2021]. To the best of our knowledge, this is the first instantiation of KFAC for heteroscedastic regression. While we derive it for the Laplace approximation, it could also be useful for optimization [Martens and Grosse, 2015].
4.3 Empirical Bayes for Automatic Regularization
To automatically regularize the heteroscedastic neural network, we use an empirical Bayes (EB) procedure that optimizes the layer-wise prior precisions,  $\delta_l$ , during training by maximizing the Laplace approximation to the marginal likelihood [Immer et al., 2021a]. This procedure can exhibit a Bayesian variant of Occam's razor [Rasmussen and Ghahramani, 2000] and trades off model fit and complexity. Although online training violates the stationarity assumption of the Laplace approximation, it has been observed to work well in practice [Immer et al., 2021a, 2023b, Daxberger et al., 2021, Lin et al., 2023]. We use gradient-based optimization of the log marginal likelihood,
$$
\log p (\mathcal {D} | \boldsymbol {\delta}) \approx \log p (\mathcal {D} | \boldsymbol {\theta} _ {*}) + \log p (\boldsymbol {\theta} _ {*} | \boldsymbol {\delta}) + \frac {1}{2} \log | \boldsymbol {\Sigma} | + \frac {P}{2} \log 2 \pi \propto \log p (\boldsymbol {\theta} _ {*} | \boldsymbol {\delta}) + \frac {1}{2} \log | \boldsymbol {\Sigma} |, \tag {12}
$$
which crucially requires differentiating the log-determinant w.r.t.  $\delta$ , which is only tractable for small neural networks. For deep neural networks, it can be done efficiently using the KFAC approximation derived in Sec. 4.2 by eigendecomposition of the individual Kronecker factors [Immer et al., 2021a, 2022]. In practice, we compute the marginal likelihood approximation every few epochs to adapt the regularization, which effectively mitigates overfitting, and use it as an early-stopping criterion. The detailed empirical Bayes training algorithm is also described in Alg. 1.
4.4 Posterior Predictive for Epistemic Uncertainties
We use the linearized posterior predictive that is in line with the linearized Laplace posterior approximation and performs typically better than sampling weights directly [Immer et al., 2021b]. We define the linearized neural network as  $\eta_{*}^{\mathrm{lin}}(\mathbf{x};\boldsymbol {\theta})\stackrel {\mathrm{def}}{=}\boldsymbol {\eta}(\mathbf{x};\boldsymbol {\theta}_{*}) + \mathbf{J}_{*}(\mathbf{x})(\boldsymbol {\theta} - \boldsymbol{\theta}_{*})$ . Due to the Gaussianity of the Laplace approximation,  $\mathcal{N}(\pmb {\theta}_*;\pmb {\Sigma})$ , and the linearization, we can express the function-space posterior as a Gaussian on the natural parameters  $\eta_{*}^{\mathrm{lin}}(\mathbf{x})\sim \mathcal{N}(\boldsymbol {\eta}(\mathbf{x};\boldsymbol {\theta}_*);\mathbf{J}_*(\mathbf{x})\pmb {\Sigma}\mathbf{J}_*(\mathbf{x})^\top)\stackrel {\mathrm{def}}{=}q(\boldsymbol {\eta}|\mathbf{x})$ . A straightforward way to approximate the posterior predictive is then a Monte-Carlo estimate of  $p(y_{*}|\mathbf{x}_{*},\mathcal{D})\approx \int p(y_{*}|\mathbf{x}_{*},\boldsymbol {\eta})q(\boldsymbol {\eta}|\mathbf{x}_{*})\mathrm{d}\boldsymbol{\eta}$  by sampling multiple  $\eta_{*}^{\mathrm{lin}}(\mathbf{x})$
Alternatively, we propose to use an approximation to the posterior predictive that can be computed in a closed-form without sampling, similar to the probit approximation commonly used in the classification setting [Daxberger et al., 2021]. To enable a closed-form posterior predictive, we restrict ourselves to the epistemic uncertainty about the mean as proposed by Le et al. [2005] for heteroscedastic Gaussian process regression. Instead of linearizing the natural parameters, we first transform them to the mean and variance using the inverse mapping of Equation 3, i.e., we have  $\mu (\mathbf{x};\pmb {\theta}) = -\frac{\eta_1(\mathbf{x};\pmb{\theta})}{2\eta_2(\mathbf{x};\pmb{\theta})}$  and  $\sigma^2 (\mathbf{x};\pmb {\theta}_*) = -\frac{1}{2\eta_2(\mathbf{x};\pmb{\theta}_*)}$ . Next, we only linearize the mean function  $\mu_{*}^{\mathrm{lin}}(\mathbf{x};\pmb {\theta}) = \mu (\mathbf{x};\pmb{\theta}_{*}) + \mathbf{J}_{*,\mu}(\mathbf{x})(\pmb {\theta} - \pmb{\theta}_{*})$  where  $\mathbf{J}_{*,\mu}(\mathbf{x})$  is the Jacobian of the mean, and have
$$
\begin{array}{l} p (y _ {*} | \mathbf {x} _ {*}, \mathcal {D}) \approx \int \mathcal {N} (y _ {*} | \mu_ {*} ^ {\mathrm {l i n}} (\mathbf {x} _ {*}), \sigma^ {2} (\mathbf {x}; \boldsymbol {\theta} _ {*})) \mathcal {N} (\boldsymbol {\theta}; \boldsymbol {\theta} _ {*}, \boldsymbol {\Sigma}) \mathrm {d} \boldsymbol {\theta} \\ = \mathcal {N} \left(y _ {*}; \mu (\mathbf {x}; \boldsymbol {\theta} _ {*}), \underbrace {\mathbf {J} _ {*} , \mu (\mathbf {x} _ {*}) \boldsymbol {\Sigma} \mathbf {J} _ {*} , \mu (\mathbf {x} _ {*}) ^ {\top}} _ {\text {e p i s t e m i c}} + \underbrace {\sigma^ {2} (\mathbf {x} ; \boldsymbol {\theta} _ {*})}\right), \tag {13} \\ \end{array}
$$
where the epistemic and aleatoric uncertainty about the mean are clearly split. An example of this posterior predictive approximation is shown in Figure 1.
4.5 Laplace Approximation for Mean-Variance parametrization
Using the natural parameter mapping, it is possible to apply above Laplace approximations to the mean-variance parametrization and profit from the empirical Bayes and posterior predictive procedures. However, because the negative log likelihood Hessian w.r.t. the mean and variance parameters can be indefinite, this does not work naively. By mapping  $\mu, \sigma^2$  to the natural parameters using Equation 3, all above derivations apply and correspond to a separate Gauss-Newton approximation of the log likelihood with Jacobians of the mapping and the natural log likelihood Hessian (see App. A).
4.6 Limitations
As common for heteroscedastic regression, we make the assumption that the conditional distribution of the target given the observations follows a Gaussian distribution. A practitioner should keep in mind that the full Laplace approximation has high computational complexity, and should resort to the KFAC approximation, which we also use for more complex settings in experiments. Further, the true posterior of neural networks is in most cases multimodal while our Laplace approximation only covers a single mode. However in practical settings, the Laplace approximation provides strong results [Daxberger et al., 2021] and can even be ensembled [Eschenhagen et al., 2021].
We evaluate the effectiveness of the natural parameterization compared to the mean-variance (naive) one, and empirical Bayes (EB) to optimizing a single regularization parameter using a grid search on the validation set (GS), and the MAP prediction vs a Bayesian posterior predictive (PP) in comparison to state-of-the-art baselines on three experimental settings: the UCI regression benchmark [Hernandez-Lobato and Adams, 2015], which is also well-established for heteroscedastic regression [Seitzer et al., 2022, Stirn et al., 2023], the recently introduced CRISPR-Cas13 gene expression datasets [Stirn et al., 2023], and our proposed heteroscedastic image-regression dataset (cf. Problem 2.1) in three noise variants. For UCI regression, we use the full Laplace approximation and we resort to the KFAC approximation for the remaining tasks for computational efficiency. Further, we use the proposed closed-form posterior predictive for our Laplace approximations. If applicable, as for our methods, we test the effect of having a Bayesian posterior predictive (PP) in comparison to the point estimate.
Baselines As control to assess the benefit of heteroscedastic aleatoric uncertainty, we include a homoscedastic model with EB and Laplace posterior predictive. In addition, we include the mean-variance parameterization of the negative log likelihood loss for heteroscedastic regression (Naive NLL). As competitive baselines, we include the recently proposed  $\beta$ -NLL [Seitzer et al., 2022] and Faithful [Stirn et al., 2023] heteroscedastic losses. Finally, we include well-established approaches for heteroscedastic regression with Bayesian neural networks, namely mean-field variational inference (VI) [Graves, 2011] and Monte-Carlo Dropout [Gal and Ghahramani, 2016] (MC-Dropout). Note that for our VI baseline we employ Flipout [Wen et al., 2018] to improve gradient estimation.
Details on model architectures, hyperparameter tuning, and additional results are in App. D.
5.1 UCI Regression
In Table 1 we report results on the UCI regression datasets [Hernandez-Lobato and Adams, 2015] for the compared models in terms of test log likelihood. We underline the best performance for each dataset, and bold results which are not statistically distinguishable from best performance. That is, if the mean performance of a model falls within the region of the mean plus minus two standard errors of the best performing model. The two rightmost columns report the total number of wins and ties for each model: wins are the number of datasets in which the given model achieves the best performance, while ties are the number of datasets in which the model achieves results which are not statistically distinguishable from the best performing model. The results validate the effectiveness of EB regularization on different heteroscedastic loss parameterizations. Notably, with the Bayesian posterior predictive (PP) our proposed methods considerably improve in performance in comparison to using a point estimate, and significantly outperform existing state-of-the-art approaches. Finally the results confirm that, particularly when the point prective is used, using the Natural parameterization of the heteroscedastic NLL improves training stability.
Objective	Regular-ization	PosteriorPredictive	LL (↑)
boston	concrete	energy	kin8nm	naval	plant	wine	yacht	Wins Ties
Homoscedastic	EB	✓	-2.51 (0.04)	-3.03 (0.03)	-0.71 (0.04)	1.29 (0.01)	5.58 (0.45)	-2.83 (0.02)	-0.94 (0.02)	-1.05 (0.14)	0	2
Naive NLL	GS	✗	-2.88 (0.2)	-3.96 (0.32)	-1.06 (0.16)	1.31 (0.01)	3.88 (0.27)	-2.85 (0.04)	-0.95 (0.02)	-0.60 (0.18)	0	1
β-NLL (0.5)	GS	✗	-2.61 (0.06)	-3.45 (0.24)	-0.96 (0.13)	1.30 (0.01)	2.60 (2.70)	-2.85 (0.04)	-0.94 (0.02)	-1.39 (0.49)	0	1
β-NLL (1.0)	GS	✗	-2.71 (0.06)	-3.13 (0.04)	-0.88 (0.07)	1.29 (0.01)	6.47 (0.12)	-2.86 (0.03)	-0.94 (0.02)	-1.09 (0.20)	0	1
Faithful	GS	✗	-2.68 (0.07)	-3.12 (0.06)	-1.10 (0.11)	1.29 (0.01)	6.41 (0.14)	-2.80 (0.02)	-0.95 (0.02)	-1.32 (0.31)	0	2
MC-Dropout	GS	✗	-3.01 (0.24)	-3.05 (0.03)	-1.61 (0.03)	1.17 (0.01)	5.64 (0.03)	-2.82 (0.01)	-1.12 (0.07)	-1.01 (0.08)	0	1
VI	GS	✗	-2.62 (0.08)	-3.10 (0.04)	-1.55 (0.03)	1.31 (0.01)	5.70 (0.24)	-2.82 (0.02)	-1.07 (0.07)	-1.16 (0.05)	0	1
Naive NLL	EB	✓	-2.47 (0.04)	-2.88 (0.03)	-0.46 (0.03)	1.37 (0.01)	6.38(0.16)	-2.79 (0.02)	-0.93 (0.02)	-0.05 (0.15)	5	6
✗	-6.03 (1.82)	-3.62 (0.26)	-0.97 (0.15)	1.36 (0.01)	-10.9 (8.01)	-2.83 (0.04)	-2.26 (0.51)	-1.27 (0.51)	0	2
Natural NLL	GS	✓	-2.46 (0.04)	-2.92 (0.04)	-0.74 (0.06)	1.32 (0.01)	6.66 (0.01)	-2.76 (0.02)	-0.94 (0.02)	-0.29 (0.07)	0	4
✗	-2.57 (0.08)	-2.95 (0.05)	-0.69 (0.08)	1.31 (0.01)	6.69 (0.01)	-2.76 (0.02)	-0.95 (0.02)	-0.87 (0.28)	1	3
EB	✓	-2.36 (0.03)	-2.93 (0.02)	-0.71 (0.03)	1.36 (0.01)	6.66 (0.01)	-2.76 (0.02)	-0.94 (0.02)	-0.51 (0.04)	2	5
✗	-2.51 (0.09)	-2.99 (0.04)	-0.65 (0.05)	1.35 (0.01)	6.68 (0.01)	-2.77 (0.02)	-1.00 (0.03)	-0.60 (0.08)	0	3
Table 1: Test log likelihood results for the compared models on the UCI regression datasets, where we report mean and standard error (dataset names in italic). Underlining of results indicates best performance while bolded results are statistically identical to best performance. The bottom half shows the proposed methods, which have the best performance overall.
Figure 3: Box-plots reporting test log likelihood results on the three CRISPR datasets, that vary with respect to cell and screen type. The natural parametrization, and using the posterior predictive improve the performance. Note that NLL mean-variance (MAP) and MC-Dropout markedly underperform in this context and are hence excluded from the figure.
5.2 CRISPR Datasets for Gene Knockdown Efficacy
As Stirn et al. [2023] highlight, heteroscedasticity is an often encountered problem in natural sciences. To test heteroscedastic regression approaches in such a setting, they introduce a set of datasets to model the efficacy of the CRISPR-Cas13 system for gene knockdown. In particular, Cas13 is a system that targets specific RNA transcripts to temporarily decrease gene expression. Decreasing gene expression levels serves multiple purposes in molecular biology, including the understanding of the function of specific genes and for therapeutic aims. In particular, the authors propose three datasets, and for each dataset multiple efficacy scores are reported, resulting from replicated experiments. The response variable is a scalar value measuring knockdown efficacy from a one-hot encoded target sequence. Note that here we report results only on the de-noised averaged data across replicates (see Stirn et al. [2023]), while we leave the results on the raw replicates for App. D. Results on these datasets validate the advantage of using the posterior predictive for our methods, as well as the effectiveness of using the natural parameterization, which results in improved training stability.
5.3 Regression with Image Data
To assess the performance of different approaches in a typical deep learning setting, we create versions of the MNIST and FashionMNIST datasets where the task is to regress the rotational angle under label-dependent heteroscedastic noise (Problem 2.1). Such a setting is interesting to study since most (heteroscedastic) regression datasets are tabular and we intend to benchmark on a more common deep learning task where tabular regression models would not be applicable.
Objective	Regular-ization	PosteriorPredictive	MNIST with MLP	FashionMNIST with CNN
LL (↑)	DKL (↓)	RMSE (↓)	LL (↑)	DKL (↓)	RMSE (↓)
Homoscedastic	EB	✓	-5.55 (0.01)	0.77 (0.02)	25.6 (0.8)	-5.54 (0.00)	0.67 (0.00)	18.7 (0.2)
Naive NLL	GS	✗	-5.56 (0.01)	0.68 (0.00)	19.0 (0.1)	-5.37 (0.01)	0.49 (0.00)	20.1 (0.1)
β-NLL (0.5)	GS	✗	-5.36 (0.01)	0.49 (0.01)	20.7 (0.8)	-5.47 (0.03)	0.59 (0.03)	23.5 (0.3)
β-NLL (1)	GS	✗	-5.38 (0.01)	0.51 (0.01)	21.5 (0.5)	-5.53 (0.03)	0.65 (0.03)	25.5 (0.2)
Faithful	GS	✗	-5.56 (0.01)	0.69 (0.00)	21.1 (0.3)	-5.78 (0.00)	0.91 (0.00)	51.8 (0.0)
MC-Dropout	GS	✓	-5.42 (0.02)	0.55 (0.01)	21.0 (0.5)	-5.37 (0.01)	0.48 (0.00)	20.1 (0.1)
VI	GS	✓	-5.42 (0.04)	0.56 (0.04)	21.9 (0.4)	-5.39 (0.01)	0.51 (0.01)	21.4 (0.5)
Naive NLL	EB	✓	-5.30 (0.00)	0.44 (0.01)	18.2 (0.3)	-5.34 (0.00)	0.47 (0.01)	18.3 (0.2)
✗	-5.31 (0.01)	-5.35 (0.01)
Natural NLL	GS	✓	-5.42 (0.00)	0.55 (0.00)	19.2 (0.1)	-5.40 (0.01)	0.56 (0.01)	24.3 (1.4)
✗	-5.42 (0.00)	-5.44 (0.01)
EB	✓	-5.30 (0.01)	0.45 (0.01)	19.1 (0.2)	-5.34 (0.00)	0.46 (0.00)	18.1 (0.1)
✗	-5.32 (0.01)	-5.35 (0.00)
Table 2: Performance metrics on heteroscedastic image regression (Problem 2.1) on MNIST and FashionMNIST with MLP and CNN architectures, respectively. We report the mean and standard error and bold all numbers that are statistically indistinguishable from the best result. The bottom half shows the proposed methods. We find that the empirical Bayes (EB) procedure leads to better regularized and performing models than using grid search (GS). The posterior predictive also strictly improves the test log likelihood, albeit not as significantly.