Effective Bayesian Heteroscedastic Regression with Deep Neural Networks
Authors: Alexander Immer
Emanuele Palumbo
Alexander Marx
Julia E Vogt
Venue: NeurIPS 2023 poster (2023)
Keywords: Heteroscedastic Regression
Marginal Likelihood
Bayesian Neural Networks
Uncertainty Estimaton
Model Selection
Laplace Approximation
TL;DR: We propose the first efficient Laplace approximation for heteroscedastic neural networks, show it scales to image data, and derive its posterior predictive.

Abstract: Flexibly quantifying both irreducible aleatoric and model-dependent epistemic uncertainties plays an important role for complex regression problems. While deep neural networks in principle can provide this flexibility and learn heteroscedastic aleatoric uncertainties through non-linear functions, recent works highlight that maximizing the log likelihood objective parameterized by mean and variance can lead to compromised mean fits since the gradient are scaled by the predictive variance, and propose adjustments in line with this premise. 
We instead propose to use the natural parametrization of the Gaussian, which has been shown to be more stable for heteroscedastic regression based on non-linear feature maps and Gaussian processes. Further, we emphasize the significance of principled regularization of the network parameters and prediction. We therefore propose an efficient Laplace approximation for heteroscedastic neural networks that allows automatic regularization through empirical Bayes and provides epistemic uncertainties, both of which improve generalization.
We showcase on a range of regression problems—including a new heteroscedastic image regression benchmark—that our methods are scalable, improve over previous approaches for heteroscedastic regression, and provide epistemic uncertainty without requiring hyperparameter tuning.

Effective Bayesian Heteroscedastic Regression with Deep Neural Networks
Alexander Immer
Emanuele Palumbo 1.3 Alexander Marx
Julia E. Vogt"
Department of Computer Science, ETH Zurich, Switzerland Max Planck Institute for Intelligent Systems, Tübingen, Germany 3 AI Center, ETH Zurich, Switzerland
Abstract
Flexibly quantifying both irreducible aleatoric and model-dependent epistemic uncertainties plays an important role for complex regression problems. While deep neural networks in principle can provide this flexibility and learn heteroscedastic aleatoric uncertainties through non-linear functions, recent works highlight that maximizing the log likelihood objective parameterized by mean and variance can lead to compromised mean fits since the gradient are scaled by the predictive variance, and propose adjustments in line with this premise. We instead propose to use the natural parametrization of the Gaussian, which has been shown to be more stable for heteroscedastic regression based on non-linear feature maps and Gaussian processes. Further, emphasize the significance of principled regular- ization of the network parameters and prediction. We therefore propose an efficient Laplace approximation for heteroscedastic neural networks that allows automatic regularization through empirical Bayes and provides epistemic uncertainties, both of which improve generalization. We showcase including a new heteroscedastic image regression benchmark -thrt our methods are scalable, improve over previous approaches for heteroscedastic regression, and provide epistemic uncertainty without requiring hyperparameter tuning.
range of regression problems-
Introduction
Capturing the epistemic (model uncertainty) and aleatoric uncertainty (observation noise) allows for computing the predictive variance of a model, which is crucial for areas such as active learn- ing [Houlsby et al., 2011, Kirsch, 2023], reinforcement learning [Osband al., 2016, Yu al., 2020] and decision making. Bayesian neural networks allow for modelling both epistemic and alcatoric uncertainty. as such, they lend themselves naturally this task. Typically, they are applied under the assumption of homoscedasticityy i.c., constant noise [MacKay, 1995, Foong et al., 2019, Kirsch, 2023, and others), but also adaptations of variational inference (VI) to model heteroscedasticity; such as mean-field VI [Graves, 2011]. deterministic VI [Wu et al., 2019, DVIJ, and Monte- Monte-Carlo Dropout [Gal and Ghahramani, 2016], have been studied. In this paper, we are interested in learning the epistemic and aleatoric uncertainty in heteroscedastic regression for potentially complex tasks such as image regression, where we are given inputs x E RRD and scalar response y E IR, and model y
conditional Gaussian distribution with mean ((x) and standard deviation a ax() being dependent on input x. Inherently, the problem involves robustly modelling the aleatoric uncertainty, corresponding to the variance, which is relevant problem in econometrics, statistics [Harvey, 1976, Amemiya, 1985], and causal discovery [Guyon et al., 2019, Xu et al., 2022].
'Equal contribution. t Shared last author. Correspondence alexander ifmerQinf ethz, ch alexander marx@inf ethz. ch. Code httpp::/ggithie http:///.w.....aaa////e/oooaeaaa
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Homoscedaastic MAP
Homoscedastic Laplace
Heterosoedastic MAP
Hetersscedastic Laplace
2.5
-2.
-0
5
-G
-5
0
5
0
ground
epistemic
aleostoric
Figure 1: Illustration of the proposed training and posterior predictive of a heteroscedastic Bayesian neural network (right) in comparison to a homoscedastic one (left).
Classical approaches based on linear models and non-linear feature maps [Cawley et al., 2004] attempt learn μ(x) and (x) directly by first fitting the mean function and subsequently fitting the standard deviation from the log-residuals, e.g.. feasible generalized least squares (FGLS) [Wooldridge, 2015]. The corresponding objective is, however, only convex one of the parameters kept constant, but is not jointly convex in both parameters [Cawley et al., 2004, Yuan and Wahba, 2004]. To alleviate this problem in the context Gaussian processes, [Le et al., 2005] propose to reparametrize the objective in terms of the natural parameters, which induces jointly convex objective. Advantages of the natural parametrization compared to standard FGLS have also been demonstrated for non-linear features maps [[mmer et al., 2023a].
In the neural network literature, the standard approach is to model p(x) and (xx as outputs of the network and maximize the corresponding Gaussian log likelihood via stochastic gradient descent [Nix and Weigend, 1994, Lakshninarayanan et al., 2017, Kendall and Gal, 2017], or use two separate neural networks to model mean and standard deviation [Skafte et al., 2019]. As in classical estimators, this parametrization might not be ideal and can lead to overconfident variance estimates [Skafte et al., 2019, Stirn and Knowles, 2020] and possibly compromised mean fits. Two recent strategies aim to adjust for this problem by reducing the influence of the predictive variance on the gradient of the mean. In particular, Seitzer et al. [2022] introduce surrogate loss, the 3-NLL loss, which regulates the influence of the variance on the gradients of the loss by introducing a stop-gradicn operation. As an alternative solution, Stirn et al. [2023] propose architectural constraints coupled with two stop gradient operations, regularizing the heteroscedastic model such that its mean fit is not compromised compared to a homoscedastic baseline. We provide more details to both approaches in Sec. 2.
Our Contributions In comparison previous work heteroscedastic regression with neural networks, we take a different perspective postulating that current estimators lack principled regu- larization. While recent work aims at regularizing the influence of the predictive variance, as we discuss in Sec. 2, we show that this focus can shift the problem to compromising capacity for the variance estimate. Instead, we propose three major modifications to tackle the problem of fitting heteroscedastic neural networks: akin to previous work on GPs [Le et al., 2005] and linear mod- els [Immer al., 2023a], we propose to re- re-parameterize the loss using the natural parametrization (cf. Sec. 3) which is known to be jointly concave in both parameters. Empirically, we find that this parametrization can be more stable during optimization. Further, we derive an efficient Laplace approximation to the marginal likelihood for heteroscedastic regression that can automatically learn regularization via empirical Bayes and provide an early- signal to prevent overfitting without requiring grid search based on a validation set in Sec. 4, Additionally, the Laplace approximation provides epistemic uncertainty through the Bayesian posterior predictive, which generally improves predictive performance. We provide a fast closed-form approximation to the posterior predictive that also provides a simple split into aleatoric and epistemic uncertainties. This predictive illustrated on the right in Figure 1 and shown in comparison to the prediction with only aleatoric uncertainty (MAP) as well as the corresponding homoscedastie regression solution on the left.
Besides showing that our approach performs favorably on commonly used UCI benchmarks and the CRISPR-Care knockdown efficacy datasets [Stirn al., 2023], we notice the lack of more complex (heteroscedastic) regression benchmark datasets. To that end, we propose new image-regression benchmarks based on image classification datasets. The input images are randomly rotated and the targets are the random rotations with heteroscedastic noise that depends on the label.
PDVI [Wu al., 2019] also employs empirical Bayes regularization, but does only apply to MLPs.

---

2 Heteroscedastic Regression with Neural Networks
Due to their universal approximation guarantees [Hornik al., 1989], deep neural networks have the capacity to solve complex regression problems. As in any regression task, however, unregularized function approximators also have the tendeney to overfit to the data. Due to the additional degree of freedom granted by learning both mean and variance in heteroscedastic regression, this problem amplified, and persists when doing a classical grid search over regularization hyperparameters. In the following, we review prior work that proposes to regularize the influence of the variance and highlight some limitations these approaches by proposing a new image regression task.
2.1 Regularizing the Influence of the Variance
Skafte et al. [2019] note that naively minimizing the Gaussian log likelihood by learning the mean and variance as output of neural network can lead overconfident variance estimates that compromise the mean fit. Taking into account the gradients of the negative log likelihood (NLL),
E (0) 2 log o o²(x; 0) +
(9- ()-m(x:01)² 2n (x:0)
+ const,
where O denote the neural network parameters used estimate mean and variance, Seitzer al. [2022] trace the problem to the fact that Valu,o (0) = -(x:0)= heavily depends on the learned variance. To overcome this effect, Seitzer et al. [2022] introduce the 3-NLL loss (²(0). which is equal to [x; 0)] (0), where L-] denotes stop stop-gradient operation, and (3 is hyperparameter controlling the dependency of gradients on the predictive variance. As a result, the gradients for 3-NLL are equal to
d²1x;0)
V (0) -(x:0)- (x:0) Volls(0) = (2(x:0) - p(x:0)22 2 (x:0)
(2)
With 3 O, ()) is equivalent to (0), whereas for 3 1 the gradient with respect to the mean is proportional to the gradient for homoscedastic regression. When setting 0 < 3 < 1 we interpolate between both settings. As an alternative approach, Stirn et al. [2023] propose to decouple the estimation into three networks: shared representation learner computes representation from x, which is passed into two individual networks fa and fz, which receive z as input and output the mean and covariance matrix, respectively. To ensure that the gradient with respect to the mean is equal to the gradient of the homoscedastic model, they introduce two stop-gradient operations: the first one has equivalent effect on the mean gradient for fa and fa setting B = 1 in B-NLL, and the second one stops any gradient from the variance network fr from propergating to fz [Stirn et al., 2023]. We provide more details in App. B.
Taking a different perspective, one can view both proposals as implicit regularization techniques for the variance parameter. The surrogate score of Seitzer et al. [2022] regularizes the influence of the predictive variance on the gradient, while the network architecture and sto-ggradient operations introduced by Stirn et al. [2023] have similar, yet stronger regularization effect for the variance- i... stopping its influence on the joint representation Iearner. An additional hurdle presents itself due to the fact that need to tune the regularization to calibrate the models for certain dataset, which can have a strong influence on the result, as we show below. Concurrent work [Wong-Toi et al., 2023] also highlights the necessity for regularization in heteroscedastic regression when considering individual mean and variance networks. They find that the networks relative regularization critical.
2.2 Learning Complex Variance Dependencies
Despite alleviating the problem of compromised mean fits, the question arises if such regularization limits the capabilities for learning the alcatoric uncertainty.
In the 3-NLL objective, we need to select the influence of the estimated variance on both gradients, while the effect of scaling is different as can be seen from Equation 2. Due to the typically applied standardization of the data, we expect that o² < 1 and hence the gradient with respect to the variance amplified stronger compared to the mean. Especially this might be problematic. Further, it is not clear which value of 3 is best suited for problem priori introducing another hyperparameter that has to be tuned in addition to tuning regularization hyperparameters. For the approach by Stir et al. [2023, Faithful] the regularization of the variance is
near-deterministic processes (a" -> 0)
3

---

original
rotated
Method
Accuracy
Test NLL (0)
5
8-LLL (0.5)
Homoscedastic
52.2 #0.7
Faithful
Faithful
51.3 #0.8
8-NLL (0.5)
71.6 #0.8
20
3-NLL (1.0)
69.3 #0.8
Natural Laplace (ours)
76,8 #0.4
class
rot 52
10
10a
10a
10*
regularization
Figure 2: Left: Example heteroscedastic image regression data point with target y ~ rot + 56e. Middle: Test log likelihood for different values of prior precision for 3-NLL and Faithful. Right: Downstream accuracy using a linear model on the last layer representation after learning the het- eroscedastic image regression. Compromised variance fit leads worse downstream accuracy.
more severe due to forcing (x() and [xx) to share joint representation that receives no gradient from the variance. If, for example, E(x) depends on set of variables that is independent of those influencing the mean, no information about the variance might be contained in the joint representation. To illustrate that only regularizing the variance can be suboptimal, consider the following example.
Problem 2.1 (Heteroscedastic Regression from Image Data) Consider version of rotated Uni,(-90, 90). We generate the target y as 9} is the image class integer number and ~ NO, 1) is an independent noise source. The heteroscedastic regression task of learning the distribution of y given observations (images) involves learning complex non-linear mapping for the variance, while learming the mean only requires learning the rotation angle.
MNIST with rotation angle drawn as rot rot + (11c+ 1)e, where E {0, .)
We train simple 3-layer MLP of width 500 using the baseline objectives and our approach (Natural Laplace introduced later in Secs. 3 and 4) on data generated as described above. First, we note that the test-NLL of 3-NLL and Faithful strongly depends the regularization strength (cf. Figure 2), which emphasizes the importance of additional regularization. Further, we report how much information about the image classification task is contained in the last layer of the MLP in Figure 2hich provides a proxy on the information about the variance that has been picked up by each approach. The homoscedastic model serves control that does not require knowledge about the label. As expected, Faithful achieves a low accuracy on this downstream task since it limits the capability for learning the variance. 3-NLL shows better performance but is still significantly outperformed by our proposal. The difference becomes even more evident when using CNN architecture on rotated FashionMNIST as we demonstrate Sec. 5.3. We also provide minimal example illustrating the limitations of Faithffll in App. C, which shows that the bottleneck size of the joint network fa has strong influence of the estimation quality of the variance.
As alternatives to Problem 2.1, we consider two generative processes for the rotated image regression task: first, use homoscedastic heteroscedastic noise that is based on the rotational magnitude, y rot V|rot|e. In this case, both mean and variance depend on the same feature, which is a setting where Faithful [Stirn et al., 2023] can theoretically work.
as ablation with y
10e. Further, we
3 Naturally Parameterized Heteroscedastic Regression
model a dataset D {(x, ) with pairs of input X, E IRP and scalar response yn. E R using the natural form of the Gaussian likelihood with unknown variance, first introduced for the case of Gaussian processes [Le et al., 2005] and recently applied in the context of causal discovery with maximum likelihood [Immer et al., 2023a]. The relationship between the natural parameters n and the mean and variance o² of the parametrization is
71 = - and 7/2 = 2 0,
(3)
which can be understood as the ggnal-to-varianee ratio and the negative precision (inverse variance). We model these natural parameters using neural network f(x; 0) E R² that is parameterized by weights 0 E R and acts on the inputs x E RD. To satisfy the constraint that 722 < we link to m using a positive function 9+ ER + R4 and have the following mapping:
Ti (x: 0) fi 0) and n2(x; 0) 9+ (x2 0)).
(4)

---

We use either the exponential (+) = 2 exp(-) or softplus g+(+) 3 log(1 + exp(-)) as typical for heteroscedastic regression with mean-variance parametrization, Mathematically, the heteroscedastic Gaussian log likelihood of our model is given by
logp(u|x, 0)
m (x: 8) (x: 0)
y (x:0)) (x)(x:0] + = log( (010--22x;x0) const,
(5)
y
which we also denote by (0). Assuming the data are ii.d. we have logo(D =
En= log p(yn Xn. 0). Immer et al. [2023a] used this maximum likelihood objective for bivari- ate causal discovery using linear models and small neural networks. Further, the gradient and Hessian take on simple forms due to the properties of the natural parametrization.
3.1 Gradients of the Natural parametrization
Similar to Seitzer al. [2022] and Stirn et al. [2023], also inspect the gradients of the correspond- ing negative log likelihood with respect to 17. leading to
Vn (0) = m (x:0) [x.
y: V (0)) (x:0)))
(6)
(x:0)) [x;
The gradients by themselves cannot be directly linked mean and variance updates. We note, however, that if we relate the natural parameters to mean and variance, i.e., we compute μ(x; 0) as (x:0)
(x;(x;0) and o² (x: 0) (x;)(x;0) then Vn, (0) reduces to p(x; 0) --y, y,, and similarly, (,
to 10o2(x: (x; (v² (μ(x; ²)) which would be desired because these are simply separate residuals for mean and variance [Seitzer al., 2022, Stirn al., 2023]. Empirically, we observe that this parametrization can be more stable to train and less prone to insufficient regularization also
observed previously for Gaussian process and ridge regression [Le al., 2005, Immer et al., 2023a].
3.2 Regularization using Bayesian Inference
Due to the expressiveness of heteroscedastic regression, regularization is crucial to obtain well- generalizing models (ef. Figure 2). We achieve regularization in two ways: first, regularize parameters towards low norm using classical g² regularization, which corresponds to Gaussian prior on the parameters. Further, we use Bayesian posterior predictive, which additionally accounts for uncertainties of the model as depicted in the illustration (Figure 1, right).
For effective regularization of deep neural networks, parameters given by p(0)8) (0;;; 0, "T). In the case of the last layer simple linear model with the natural heteroscedastic likelihood, it induces mode on 77 (x) = 0 and (2(x) 3 which corresponds to zero mean and unit variance and is reasonable for standardized response variables. The layer-wise prior further allows to differently regularize parts of the neural network and has been observed to improve generalization in image classification [Immer et al., 202 la, Daxberger et al., 2021, Antorán et al., 2022]. Concurrent work by also suggests that this might be the case for heteroscedastic regression [Wong-Tee Toi et al., 2023]. However, optimizing a regularization parameter per layer is intractable using validation- based grid search.
use a layer-wise Gaussian prior the
To optimize wiser-wise prior precisions and obtain a posterior predictive, we make use Bayesian inference. Combining the natural Gaussian likelihood p (DD))) ()) (below Equation 5) with the prior p()))), we have the joint distribution pDD, 0(5), which corresponds to regularized objective. According to Bayes' theorem, we have the posterior p(0)D, 8) pDD, 0|)) The normalization constant, also referred to as marginal likelihood, is given by (D10) = f(D 018)d0 and gives us the Type II maximum likelihood objective to optimize the prior precisions ö. This procedure referred to as empirical Bayes (EB). Inferring the posterior distribution of the neural network parameters, further have access to the posterior predictive P(y, [x, D) for a new data point x.. By averaging over multiple hypotheses from the posterior, the predictive can be better regularized than single model [Wilson, 2020]. Unfortunately, inference is intractable for deep neural networks.
4 Approximate Inference with Laplace Approximation
We develop scalable Laplace approximation for the posterior in heteroscedastic regression with deep neural networks. The Laplace approximation [MacKay, 1995] is effective method for
5

---

approximating the marginal likelihood, posterior, and predictive in deep learning [Daxberger et al., 2021]. In comparison other approximate inference methods, 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 scalable to deep learning [Ritter et al., 2018]. We extend these to the heteroscedastic regression setting.
Laplace approximates the posterior locally at mode of the posterior with a Gaussian dis- tribution, p(0 8) & NN 0., E). The mean is given by stationary point of the poste- rior, 0. arg max log log 0)). and the covariance by the curvature at that mode, Σ- V po (D, 0 d))o= This is due to a second-orde 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 Lincarized Laplace for Natural Heteroscedastic Regression
The linearized Laplace approximation [MacKay, 1995, Khan 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 linearization
E 2 J. VI logp(un|>,, Xx., 0.)] (x,)) + V log p ( (01o)l==a.
where [x](x)]
le-0. is the Jacobian of the neural network (); 0) 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 5, > O, the layer-wise regularization parameters.
0n.(x:0) 88
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 diagonal with entries & > 0, one only has to show that the negative log likelihood Hessian, - V log p(y|x, O.), positive semidefinite, which is simply a property of naturally parameterized exponential fami- lies [Martens, 2020]. Note that this would not be the case for the mean-variance parametrization' For efficient computation, we can decompose the log likelihood Hessian as
m1 (w:
A. (x) V2 logp(s/xx00)
2na (m:
(8)
Mi (.:
Mi (2:0
20z (m:00_
2 (ar000 )
2va(=:00)
Tn (v:0.)
(9)
2m2
(m:00))
2m:22m2:::: (m:0.)
2mp (1: 0.:
def (x(x)xx -xx())x)))
(10)
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-vectes products to compute the covariance in Equation 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 mameeckr-facctred approximation [KFAC; Martens and Grosse, 2015, Botev 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 H from blocks H; per layer Second, each block H, 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 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 KFAC approximation that can efficiently be computed in a closed-form. We can write the Jacobian of
In Sec. 4.5 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 JI(x,)' = at, & gt., with at,. E IR Dx.1 as the layer's input and gr,a E R x2 as transposed Jacobian w.r.t. the output, both for the input xn. Following Martens and Grosse [2015] and Botev et al. [2017], we then have the KFAC approximation
Xn.
E 'I
En Jav,n @ gI. An @ gl.)]
+ + En==Enilanne
lat
@ [g.. An [+ 6,1
1 N
Ex= at an,
Zk=1 Bi, BL. N Ai, 8
Jet A
B
+ 6,1,
(11)
where An, 2 is due to the decomposition of A,, into outer products Equation 10 and the approximation is due to exchanging the sum and product. Conveniently, the terms Br. can be computed efficiently using two Jacobian-vectes vector products and the Kronecker factors can then be extracted within the second-order framework of Osawa [2021]. To the best of our knowledge, this 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, procedure that optimizes the layer-wise prior precisions, 6, 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 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., 202 la, 2023b, Daxberger al., 2021, Lin et al., 2023]. We use gradient-
empirical Bayes (EB)
optimization of the log marginal likelihood,
logp(D))) ss logp(D)(..) + log p (0.16) + 1 log + E I log 2T x log plp.p(O00 (8) + 1 log [>], (12)
which crucially requires differentiating the log- dee-deetemiant w.r.t. 6, which is only tractable for small neural networks. For deep neural networks, can be done efficiently using the KFAC approximation derived in Sec. 4.2 by eigendecomposition of the individual Kronecker factors [Immer et al., 2021 la, 2022]. In practice, regularization, which effectively mitigates overfitting, and use it as an early-stopping criterion. The detailed empirical Bayes training algorithm is also described in Alg.
compute the marginal likelihood approximation every few epochs to adapt the
4.4 Posterior Predictive for Epistemic Uncertainties
We use the linearized posterior predictive that is in line with the linearized Laplace posterior approxi- mation and performs typically better than sampling weights directly [Immer et al., 2021b]. We define the linearized neural network as 17. Jin (x; 0) der (x; 0.) + J. (x) 0.). Due to the Gaussianity of the Laplace approximation, (0,; E). and the lincarization, we can express the function-space posterior as a Gaussian on the natural parameters (x) ~ (n)(xx; 0. J. (x) d 4(n x). A straightforward way approximate the posterior predictive is then a Monte-Carle estimate of P(y. Xx, D) = p(y. 7) ())(())x, dn by sampling multiple ry.inn (x).
Alternativel,, we propose to use 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 closed-form posterior predictive, 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, transform them to the mean and variance using the inverse mapping of Equation 3, i.c., we have
first
(x;; 0) = ni (x:0)
(x)(x:0) and ox(x;0 0.) = 2na(x:0.)' Next, we only linearize the mean function Lin (x: 0) = p(x; 0.) + J J., (x) O.), where J (x) is the Jacobian of the mean, and have p(u.1x.. |x., D) [N 1n!! (xx,, =N(y 0.), (x. EJ (x.)) + (x: 0.))
(x.,), o² (x: (.)) 0., E) de
(13)
epistemic
aleatoric
where the epistemic and aleatoric uncertainty about the mean are clearly split. example of this posterior predictive approximation is shown in Figure 1.
7

---

4.5 Laplace Approximation for Mean- Varinnnr parametrization
Using the natural parameter mapping, is possible to apply above Laplace approximations to the mean-variance parametrization and profit from the empirical Bayes and posterior predictive proce- dures. 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 μt, o² to the natural parameters using Equation 3, all above derivations apply and correspond to separate Gauss-Newton approximation of the log likelihood with Jacobians of the mapping and the natural log likelihood Hessian (sce 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 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 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].
5 Experiments
We evaluate the effectiveness of the natural parameterization compared to the mean-variame vernnnar (naive) one, and empirical Bayes (EB) 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 hetcroscedastic regression [Seitzer et al., 2022, Stirn et al., 2023], the recently introduced CRISPR C3 gene expression datasets [Stirn et al., 2023], and our proposed heteroscedastic imagc-regression dataset (cf. Problem 2.1) in three noise variants. For UCI regression,
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 assess the benefit of heteroscedastic alcatoric 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 3-NLL [Seitzer et al., 2022] and Faithfill [Stirn et al., 2023] heteroscedastic losses. Finally we include well-estale eslablstall 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 we report results on the UCI regression datasets (lernandezzLLoatt 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 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 comparison to using point estimate, and significantly outperform existing -fate-of-the-art approaches. Finally the results confirm that, particularly when the point prective used, using the Natural parameterization of the heteroscedastic NLL improves training stability.
8

---

(L
Regular- Posterior ization Predictive
Objective
concrete
energy
naval
plant
winte
yachr
Wins
Homoscedattic
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.55 (0.14)
0
Naive NLJ
GS
K
-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
A-NLL (0.5)
(0.06) (0.24)
(2.70)
(0.49)
(L.0)
GS
X
-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
Farthful
(0.14)
-1.32 (0.31)
MC-Dropout
GS
X
-3.01 (0.24) 3305 (0.03) -1661 (0.03) 1.17 (0.01) 5.64 (0.03) -2.82 (0.01)
-1.12 (0.07) -1.01 (0.08)
0
VI
GS
-2.62 (0.08) -3.10 (0.04) -155 (0.03) 1,31 (0.01) 5,70 (0.24) -2.82 (0.02)
-1.07 (0.07) -1.16 (0.05)
0
/
-2.47 (0.04) -2.88 (0.03) 40,46 (0.03) 1,37 (0,01) 6.38(0.16)
-2.79 (0,02)
-0.93 (0M12) +0,05 (0.15)
5
6
Naive NLL
-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
G8
-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
Natural NIJ.
x
-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)
V .2.36 (0.03) -2.93 (0.02) 40.71 (0.03) 1,36 (0.01) 6.66 (0.01)
.2.76 (0.02)
-0.94 (0.02) -0.51 (0.04)
2
X -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
Table 1: Test log likelihood results for the compared models on the UCI regression datasets, where we report mean and standard error (dataset names 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.
Hhw-eetoomety-HHBe HEX229
survival-screen-A33.
survial-scree--HEX320
Homoscadaetic GS PP-
+
Naive NLLL-
+
(ANLL (0.5)-
+
(.NL.. (1.0)
[|]
()
PP
Nake PP-
Naturl NLJ. GS PP-
H
Natural NLL (GS-
-
Natual NLL EB PP-
+
Natural NLL
(1)
LL
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 NL.L 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 approches in such a setting, they introduce set of datasets model the efficacy of the CRISPR-Cas 13 system for gene knockdown. In particular, Cas 13 is system that targets specific RNA transcripts 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 scalar value measuring knockdown efficacy from 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 typical deep learning setting, we create versions of the MNIST and FashionMNIST datasets where the task is regress the rotational angle under label-dependent heteroscedastic noise (Problem 2. Such 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.

---

MNIST with MLP
FashionMNIST with CNN
Regular- Lzalion
Posserior Predictive
Objective
IT.(t)
DKL (4)
RMSE
1I-(t)
PKL (4)
RMSE (L)
Homoscedastic
EB
V
5.55 (0.01)
0.77 (0.02)
(0.6
-5.54 (0.00)
0.67 (0.00)
18.7 (0.2)
Naive NLL
x
(0.51)
0.68 (0.00)
(9.0
-5.37 (0.01)
0.49 (0.00)
20.1 (0.1)
G-NLL (0.5)
GS
x
-5.36
0,49 (0.01)
(0.8
-5.47 (0.03)
0.59 (0.03)
23.5 (0.3)
B-NLLL()
GS
X
-5.38 (0.01)
0.51 (0.01)
21.1 (0.5
-5.53 (0.03)
0.65 (0.03)
25.5 (0.2)
Faithful
GS
-5.56 (0.01)
0.69 (0.00)
-5.78
0.91 (0.00)
51.8
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)
-5.42 (0.04)
0.56 (0.04)
21.9 (0.4)
(0.01)
(0.01)
21.4 (0.5)
Naive NLL
EB
V
X
-5.31 (0.01)
.5.30 (0.00)
0.44 (0.01)
18.2 (0.3)
-5.35 (0.01)
.5.34 (0.00)
0.47 (0.01)
18.3 (0.2)
/
-5.42 (0.00)
(0.01)
GS
0.55 (0.00)
(0.1
0.56 (0.01)
24.3 (1.4)
Natural NLL
X
-5.42 (0.00)
-5.44 (0.01)
/ .5.30 (0.01)
.5.34 (0.00)
0.45 (0.01)
19.1 (0.2)
0.46 (0.00)
18.1 (0.1)
X -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.
To create the modified datasets, we follow the procedure as outlined in Problem 2.1. An observation x is generated by rotating an MNIST (resp. FashionMNIST) image with rotation angle drawn as rot- ~ Unif(- -90, 90), and the corresponding target y is generated as y = rot + (1lc 1)e, where C {0, 1.. ... 9} is the image class of x as integer number and (0 1) is an independent noise source. To solve the problem, a heteroscedastic regression model needs to 1) learn the mapping of rotated image to angle and 2) learn the label of the image to conrectly fit the heteroscedasticity of the observation noise, both of which require learning complex functions. We evaluate the methods based on the test log likelihood and RMSE. In addition, we compute the KL-divergence of the predicted distribution from the ground truth.* We train a 3-layer 500-neuron wide MLP with ReLU activation for the MNIST-based dataset and a CNN conv. and linear) for the task based on FaShionMNIST. In App. D.2.4, we provide additional details and results on two alternative tasks.
The results on image regression with heteroscedastic label noise indicate that the proposed empirical Bayes (EB) algorithm leads to better generalizing models and the posterior predictive (PP) strictly improves performance for both parameterizations. It is also notable that Faithful fails to fit the model due to the restrictions on its variance function, and therefore performs worse than a homoscedastic model FashionMNIST. B-NILL, MC-Dropout, and VI all give relatively good results but fall slightly short of the proposed methods.
Conclusions
We tackled the problem of heteroscedastic regression using neural networks with focus on regular- ization. To this end we proposed three individual improvements: 1) we use the natural form of the Gaussian likelihood; we derive a scalable Laplace approximation that 2) can be used for automatic principled regularization without validation data, and 3) enables epistemic uncertainty estimation through its posterior predictive. We show on image data that our method scalable despite the Bayesian procedures and benchmark our approach on UCI regression as well as CRISPR-CasS datasets achieving ofate-of-the-ar performance. For future work, would be interesting to relax the assumption about Gaussianity and apply our approach to other real-world datasets.
'Note, however, that the original images are not completely free of rotation which creates an additional, potentially heteroscedastic, noise source that cannot quantify.
10

---

Acknowledgements
Al gratefully acknowledges funding by the Max Planck ETH Center for Learning Systems (CLS). EP and AM were supported by a fellowship from the ETH Al Center.
References
Takeshi Amemiya. Advanced econometrics. Harvard university press, 1985.
Javier Antorán, David Janz, James U Allingham, Erik Daxberger, Riccardo Rb Barbano, Eric Nalisnick, and José Miguel Herández-Lobato. Adapting the linearised laplace model evidence for modern deep learning. In International Conference on Machine Learning, pages 796-821, PMLR, 2022.
Alcksandar Botev, Hippolyt Ritter, and David Barber. Practical gauss-newton optimisation for deep learning. In International Conference Machine Learning, pages 557-565. PMLR, 2017.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http://github: www. coom/google/jr
Gavin Cawley Nicola LC Talbot, Robert Foxall, Stephen R Dorling, and Danilo Mandic. Heteroscedastic kernel ridge regression. Neurocomputingg 57:105-124, 2004.
Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace redux-effortlen bayesian deep learning. Advances in Neural Information Processing Systems, 34:20089 20103, 2021.
Runa Eschenhagen, Erik Daxberger, Philipp Hennig, and Agustinus Kristiadi. Mixtures of laplace ap- proximations for improved post-hoc uncertainty in deep learning. arXiv preprint arXiv:2111.03577,
2021.
Andrew YK Foong, Yingzhen Li, José Miguel Hernández-Lobato. and Richard E Turner. 'in- between uncertainty in bayesian neural networks. arXiv preprint arXiv:1906. 11537, 2019.
Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In International Conference on Machine Learning. 2016.
Alex Graves. Practical variational inference for neural networks. In Advances Neural Information
Processing Systems, 2011.
Isabelle Guyon, Alexander Statnikov, and Berna Bakir Batu. Cause effect Pairs in machine learning. Springer, 2019.
Andrew Harvey. Estimating regression models with htt.tppccccttttttttooooccctctttt Ecooootticccty.cooooomeitccccaa Econometrica: Journal of the Econometric Society, pages 461-465 1976.
Jose Miguel Hernandez-Lobato and Ryan Adams. Probabilistic backpropagation for scalable learning of bayesian neural networks. In International Conference on Machine Learning. PMLR, 2015.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural networks, 2(5):359-366, 1989.
Neil Houlsby, Ferenc Huszár, Zoubin Ghahramani, and Máté Lengyel. Bayesian active learning for classification and preference learning, arXiv preprint arXiv: 2011.
Alexander Immer, Matthias Bauer, Vincent Fortuin, Gunnar Rätsch, and Khan Mohammad Emtiyaz. Scalable marginal likelihood estimation for model selection in deep learning. In International Conference on Machine Learning, pages 4563-4573. PMLR, 2021a.
Alexander Immer, Maciej Korzepa, and Mathias Bauer. Improving predictions of bayesian neural nets via local linearization. International Conference on Artificial Intelligence and Statistics, pages 703-711. PMLR, 202

---

Alexander Immer, Tycho van der Ouderaa, Gunnar Rätsch, Vincent Fortuin, and Mark van der Wilk. Invariance learning in deep neural networks with differentiable laplace approximations. Advances in Neural Information Processing Systems, 35:1 1122441444649 2022.
Alexander Immer, Christoph Schultheiss, Julia E Vogt, Bernhard Schölkopf, Peter Bühlmann, and Alexander Marx. On the identifiability and estimation of causal location-scale noise models. In International Conference Machine Learning, pages 1221661-33.... PMLR, 2023a.
Alexander Immer, Tycho FA Van Der Ouderaa, Mark Van Der Wilk, Gunnar Ratsch, and Bernhard Schölkopf. Stochastic marginal likelihood gradients using neural tangent kernels. In International Conference on Machine Learning, pages 11333-14352 PMLR, 2023b.
Alex Kendall and Yarin Gal. What uncertainties do we need bayesian deep learning for computer vision? Advances in neural information processing systems, 30, 2017.
Mohammad Emtiyaz E Khan, Alexander Immer, Ehsan Abedi, and Maciej Korzepa. Approximate inference turns deep networks into gaussian processes. Advances in neural information processing systems, 32, 2019.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:/412.6998 2014.
Andreas Kirsch. Black-box batch active learning for regression. arXiv preprint arXiv:3302.08981, 2023.
Ranganath Krishnan, Pi Esposito, and Mahesh Subedar. Bayesian-torchh Bayesian neural network lay- ers for uncertainty estimation. attps://gitee com/ http//ntelaabbbaaeeeeee bbayssiaaassaaaa torch, January 2022.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems, 30, 2017.
Quoc V Le. Alex J Smola, and Stéphane Canu. Heteroscedastie gaussian process regression. In International Conference on Machine Learning, pages 489-496. PMLR, 2005.
Jihao Andreas Lin, Javier Antorán, and José Miguel Hemández--Lobato. Online laplace model selection revisited. arXiv preprint (/XXv:2307.06093 2023.
David JC MacKay. Probable networks and plausible predictions-a review of practical bayesian methods for supervised neural networks. Network: computation in neural systems, 6(3):469, 1995.
James Martens. New insights and perspectives on the natural gradient method. The Journal of
Machine Learning Research, 21(1):5776-881) 2020.
James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In International conference on machine learning, pages 2408-2417 PMLR, 2015.
David A Nix and Andreas Weigend. Estimating the mean and variance of the target probability distribution. In IEEE international conference on neural networks, volume pages 55-60. IEBE, 1994.
Kazuki Osawa. Automatic Second-Order Differentiation Library (ASDL), 2021. URL http:
//github. won//kazukiosawa//aaa asdl.
Ian Osband, Charles Blundell, Alexander Pritzel, and Benjamin Van Roy. Deep exploration via bootstrapped dqn. Advances in neural information processing systems, 29, 2016.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang. Zachary Vito,o, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differenniaaion in pytorch. In NewrIPS Workshop Autodiff, 2017.
Carl Rasmussen and Zoubin Ghahramani. Occam's razor. Advances in neural information processing systems, 13, 2000.
12

---

Hippolyt Ritter, Aleksandar Botev, and David Barber. A scalable laplace upproximation for neural networks. In International Conference on Learning Representations, volume 2018.
Maximilian Seitzer, Arash Tavakoli, Dimitrije Antic, and Georg Martius. On the pitfalls het- eroscedastic uncertainty estimation with probabilistic neural networks. In International Conference on Learning Representations. 2022.
Nicki Skafte, Martin Jø rgensen, and So ren Hauberg. Reliable training and estimation of variance networks. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Andrew Stirn and David A Knowles. Variational variance: Simple, reliable, calibrated heteroscedastic noise variance parameterization. arXiv preprint MXXv:22006.04910. 2020.
Andrew Stirn, Harm Wessels, Megan Schertzer, Laura Pereira, Neville Sanjana, and David Knowles. Faithful heteroscedastic regression with neural networks. In International Conference on Artificial Intelligence and Statistics, pages 5593-5613. PMLR, 2023.
Marcin B Tomezak, Siddharth Swaroop, and Richard E Turner. Neural network ensembles and variational inference revisited. In 1st Symposium on Advances in Approximate Bayesian Inference, pages 1-11, 2018.
Yeming Wen, Paul Vicol, Jimmy Ba, Dustin Tran, and Roger Grosse. Flipout:
Efficient pseudo- independent weight perturbations on mini-batches. In International Conference Learning Representations, 2018.
Andrew Gordon Wilson. The case for bayesian deep learning. arXiv preprint arXi::2001 2020.
Eliot Wong-Toi, Alex Boyd, Vincent Fortuin, and Stephan Mandt, Understanding pathologies deep heteroskedastic regression. arXiv preprint arXiv:2306.1677 2023.
Jeffrey M Wooldridge. Introductory econometrics: A modern approach. Cengage learning, 2015.
Anqi Wu, Sebastian Nowozin, Edward Meeds, Richard E. Turner, Jose Miguel Hernandez-Lobee and Alexander L. Gaunt. Deterministic variational inference for robust bayesian neural networks. In International Conference on Learning Representations, 2019.
Sascha Xu, Osman Mian, Alexander Marx, and Jilles Vreeken. Inferring cause and effect in the
presence of heteroscedastic noise. In International Conference on Machine Learning, 2022.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. Advances in Neural Information Processing Systems, 33: 53:14129-14142 2020.
Ming Yuan and Grace Wahba. Doubly penalized likelihood estimator in heteroscedastic regression. Statistics & probability letters, 69(1):11 1-20, 2004.
13

---

Appendix
Table of Contents
A Laplace with Mean and Variance Parametrization
15
Comparison Gradient Updates
15
C Orthogonal Mean- Variance Sources
Mean-Variance
17
Details on Experiments and Additional Results
18
D.I
Experimental Details
18
D.1.1
Skafte Illustrative Example
18
D.1.2 UCI Regression and CRISPR-Cas13 Knockdown Efficacy Experiments
18
D.1.3
Image Regression...
19
D.2
Additional Experimental Results
20
D.2.1
Skafte Illustrative Example
20
D.2.2
UCI Regression Datasets
20
D.2.3
CRISPR Gene Knoekdown Efficacy Datasets
21
D.2.4
Image Regression
21
Implementation
22
E.1
Pseudo-Code for Our Method
23
E.2
Implementation of Natural Gaussian Log Likelihood
24
14

---

A Laplace with Mean and Variance Parametrization
In Sec. we introduced the Laplace approximation for the naturally parameterized Gaussian like- lihood. The natural parametrization allows an efficient generalized Gauss-Newer approximation to the Hessian due to the properties the second derivative [Martens, 2020], i.e., that is positive definite w.r.t. the natural parameters. The positive definiteness is necessary that the generalized Gauss-Newton, which pre and post-multiplies the Jacobians to this Hessian, can only be positive semidefinite. This is necessary requirement for the Laplace approximation.
The Laplace approximation provides different useful capabilities for heteroscedastic regression: 1) automatic regularization through marginal likelihood optimization (empirical Bayes), 2) an approx- imation the posterior predictive, and 3) epistemic uncertainties on the functional output. While we derived these in the context of the natural parametrization in the main text, it is also possible to obtain these benefits for the common mean-vant varrivae parametrization of a neural network.
Since the Hessian of the Gaussian likelihood under the mean-variance parametrization is not positive definite, cannot directly use this parametrization for the Gauss-Newton approximation. To verify
this statement, define m(;; 0) = [u[x; 0), (x; 0)] and compute the Hessian of the Gaussian
negative log likelihood w.r.t. m
1
y- y=u(x::)
(x:00)
(o² (x;(x))))3 (u(x::)-v)) (x.00)
V logp(u/x, 0)
(14)
w-w(x:0) (x:)))
(x²(xx0))²
By Silvester's criterion, Hermitian matrix is positive definite if and only if the determinants of all sub-matrices are positive. For the Hessian Equation 14, this criterion is not fulfilled as
det (- ? log p(ylx, 0)) 2(="(x:0)))
simple way overcome this issue is to map the mean-variance parameters the natural parameters using Equation 3 and apply the approximations as derived in Sec. 4, i.c.,
m) (x: 0)
=(x:0) and (2 0) (x.0) (x:0)
(15)
and then apply the Gauss-Newton approximation as in Sec. 4 to obtain
22 logp(ulxx 0) = -J V log p) = of(x:0)
(16)
(x:00)))
where the Jacobians J, Jnjm are of 77 w.r.t. m. It is easy to see that the Hessian approximation in Equation 16 is positive definite by evaluating the determinants of all sub-matrices:" The implemen- tation is very simple: we simply use the output of the neural network m, apply the transformation to the natural parameters 77. and then compute the quantities derived for the natural Laplace approximation.
B Comparison of Gradient Updates
In the following, we first review the problem of the mean-variance parametrization identified by Seitzer et al. [2022] and present their proposed solution. Then we recap the proposal by Stirn et al. [2023], and subsequently explain how the natural parametrization avoids the gradient scaling problem.
The Problem Seitzer et al. [2022] identify that problem of previous neural estimators for heteroscedastic regression is due to the gradient of the negative log likelihood (NLL) which are parametrized to estimate the mean and standard deviation. First, recap that the corresponding NLL with respect to neural network parameters O can be written as
l (0) 2 log o (x; 0) +
(g) (g==(x:0))² 20 (x:00
+ consi.
(17)
'Interestingly, the individual determinants for the x and 2 2 sub-matrin matrices are equal to the absolute value of the corresponding determinants of - v log p x, 0) reminiscient of saddle-fre free methods.
15

---

J0
06 2
ou oz
06 pr
dz
dz ae
02 de
00
f
11 dj
of Fr 02
Oy
00,
O2 08g
FF
Figure 4: Illustration for the faithful heteroscedastic regression method proposed by Stirn et al. [2023]. Functions parameterized by neural networks are color-coded in blue, backpropagated gradients are color-coded in violet, and the two black arrows only indicate that the values for y and E are used as input for loss computation. The two modifications to conventional heteroscedastic regression introduced by Stirn et al. [2023] are color-coded in orange: the introduction of the scaling of the gradient V, (O) ) by the covariance, and of stop gradient operation to avoid FF:(FF(0) gradients being propagated Jo.. In this setup, specific dependencies for Σ x which are independent of the mean function are discarded, especially in the case in which z acts as a bottleneck. Note that in this image we make the dependency of neural network functions on their respective parameters explicit. In the rest of the work we omit the for the sake of brevity and readability.
As Seitzer et al. [2022] point out, the gradient of fu (0) with respect to μ(x; 0) is scaled by the estimated precision, i.e,
( (0)
((x:0) [xx00 Vo (x(00)) - (xx)))) 2 (x;0)=
(18)
which they argue can lead to a poor mean estimate. In particular, they state that the variance well-calibrated, i.c., o² o²(x; 0) z (u(x: 0) y)² the gradient with respect to the mean reduces to
Vplno
p(x:0)y xx
(19)
P
which rewards points with already accurate fit and undersamples points for which the mean fit is inaccurate [Seitzer et al., 2022]. We would like to note, however, that the same holds the irreducible aleatoric uncertainty low resp. high in certain region, which does not necessarily imply poor mean fit.
3-NLL
To improve the mean fit, Seitzer al. [2022] introduce the B-NLL loss
() = [o (x: 0)J (t log (x;; (y-H(x:0))² 2x-(x.0)
+
(20)
which introduces a noo-ggradeett operation denoted as [-] to trade-off the influence the variance estimate the gradients. The corresponding gradients for the B-NLL loss are
V [O) p(x:0)= (x:0)
V vo=lz(0) <x(x:0) (p pxx0()))))
(21)
(x:0)
Thus, for 3 = 1, the gradient for the mean is proportional to the gradient for the homoscedastic objective (which does not hold for the gradient of the variance), for B it is the standard heteroscedastic gradient and for < 3 < the approach interpolates between both extremes. A special case appears at B 0.5, where the gradient of the mean is weighted by the standard deviation.
Proposal of [Stirn et al., 2023] To approach the gradient scaling, Stirn et al. [2023] argue that a heteroscedastic regression model that learns u(x), axx should be faithful homoscedastic model that learns (x()). As faithfulness they define that (cf. Stirn et al. [2023, Def. 2J)
E (x.x)> (yy (=(x)) 2 E(x,)~D(((r {1
(22)
16

---

To achieve this property, they propose to decouple the estimation into three networks, as shown in Figure 4: shared representation learner fs computes a representation z from x, which is passes into two individual networks fn. and Sx, which receive as input and output the mean respectively the covariance matrix. Based on this network architecture, they minimize the following loss
4F 0) (v- log. (uI [x)(x),, (x],,(())))
(23)
which incorporates two stop-gradient operations. They show that the proposal is equivalent to the following two modifications, which are illustrated in Figure 4: First, they stop the gradient of fx from propergating to fz. Second, they propose to multiply the gradient of f o with E. In combination, both modifications imply that the gradient with respect to the mean is equivalent to the homoscedastic model [Stirn et al., 2023]. As discussed in Sec. 2 and App. C. these modifications restrict the capacity of the joint network such that it cannot pick-up complex dependencies for the variance that do not affect the mean.
Last, the procedure has some commonalities to 3-NLL. In particular, for the sub-networks fz and fu the gradients are multiplied with which is the multivariate analogue of the effect of the 33NLL loss on the mean gradient for / 1.
Proposed: Natural Parametrization
The natural parametrization, that we use in this paper, has the following gradients with respect to (0) with respect to and 172 as
Vn (0) =
ni (x:0)
y. (n) =
(m (x:0))²
2np(x;00 y²
(24)
(xx(0=
(x;]))=
Although these gradients do not have direct implication on updates with respect to mean and variance, note that if we compute the mean estimate μxx; 0) from the estimated natural parameters, it (x:0)
equal to (x: Similarly, note that the variance o>; 0) can be computed from the natural parametization (x:0 If replace the corresponding quantities in Equation 24, we get that o²(x:²2(0p C, = ,(x: 0) y. Vm (0)) = (x: 0) - (y² (n(x; 0))
(25)
which desired according to Seitzer et al. [2022] and Stirn et al. [2023] since these are simply separate residuals for mean and variance. Empirically, we observe that this parametrization can be more stable to train and less prone to insufficient regularization, which is in line with previous work
on Gaussian processes and ridge regression [Le et al., 2005, Immer et al., 2023a].
C Orthogonal
Sources
To illustrate the limitations using a joint encoder network for mean and variance, which receives no gradients from the variance estimate (cf. [Stirn et al., 2023, Faithful)) we consider toy setting in which mean and variance are generated from independent random variables:
Problem C.1 Let X E R²2 be a centered multivariate Gaussian distributed with Σ I. We construct the target variable Y as
y = sin(x1) + 1/2 (112) + +
where NO 1) is an independent noise variable.
From the mechanism described in Problem C.I, we construct 20 datasets consisting of 1000 i.i.d. sam- ples, where each dataset is generated using different random seed. We split the data into train- validation/test according to the ratio (0.9/0.1). Then the train-validation data are split into training and validation sets (0.9/0.1).
As baselines we consider our empirical Bayes training with the natural homoscedastic and het- croscedastic objective using only point predictive.both trained with the full Laplace approximation, where we use an MLP with hidden-layee consisting of 100 hidden units. For Faithful, we use the same MLP with 1 hidden layer for the joint encoder fa and instantiate ff and fz linear projections from dimension zz| to 1. We ablate the size of the hidden dimension zz] from to 128 neurons to create bottleneck. In theory, size of 2 should be sufficient encode both mean and variance.
As shown in Figure 5,the bottleneck size has an influence on how much information about the variance is propagated, however, large bottleneck size does not guarantee that the network learns
17

---

-64
Toxt TeelangsBeeleee
1.0
-12
-12
16
32
64
128
16
32
128
Bottlescck sloo
Bottlenec size
(a) Average test log likelihood across seeds
(b) Best test log likelihood across seeds
Figure 5: Visualization of the results for data generated according to Problem C.I, where only Faithful is ablated with respect to the bottleneck size. Left: Average test log likelihood for each the compared models, with standard error as error bars. Right: Best test log likelihood result achieved across seeds for each of the models.
the relevant information about the variance. We rather conjecture that the amount of information that propagated is dependent on the initialization of the network, since no gradient actively encourages the network fz to pick up information about X2, which influences the variance. Since for a larger bottleneck size the network can encode much more information than just X1, observe that encodes information about X₂ more consistently, and for some seeds (see Figure 5b) it performs comparably well with our heteroscedastic variant. On average, however, its performance lies, as expected, between the homoscedastic and heteroscedastic models with natural parametrization trained with the full Laplace approximation.
D Details on Experiments and Additional Results
In this section, we first provide technical details for our experimental comparisons, and then show some additional results for our experiments.
D.1
Experimental Details
D.1.1 Skafte Illustrative Example
In Figure 1, we compare homoscedastic to heteroscedastic likelihood and the point-wise prediction the applac--appooximate posterior predictive. The dataset is created similar Skafte al. [2019] with the following generative model:
sin(x) + (0.1 0..zz))e with x~ Uniff2.5, 12.5) and E 100,
(26)
For training, we center the input data by subtracting the mean 7.5 from x and standardize the observations zero mean and unit variance common for regression problems. We create a dataset with 1000 samples.
We model the data with neural network with 100 hidden units on a single hidden layer using tanh activation. Both homoscedastic and heteroscedastic models are trained using the marginal- likelihood optimization algorithm [Immer et al., 2021a], also detailed in Alg. 1. We train the neural network parameter using learning rate 0.01 with Adam [Kingma and Ba, 2014] and optimize the hyperparameters using learning rate 0.1 with the same optimizer every 50 epochs for 50 gradient steps. In this example, we use the efficient Kronecker-factored Laplace approximation and use model that obtained the best marginal likelihood approximation during training, akin to early stopping.
D.1.2 UCI Regression and CRISPR-Cas13 Knockdown Efficacy Experiments
For both UCI and CRISPR regression datasets all compared models use an MLP with hidden layer, with width of 50 units. We use the GeLU activation function for models trained on the UCI regression datasets, while we adopt the ReLU activation function for models trained on the CRISPR datasets. We train all models, except for the VI and MC-Dropou baselines, with Adam optimizer using
18

---

batch size of 256 for 5000 epochs and an initial learning rate of 10 that is decayed to 10 using a cosine schedule. As we observe that gives overall better performance, MC-Dropou and VI on UCI datasets are trained for 1000 epochs, with learning rate of 10-3 that is not decayed. In the models using classic grid-search (GS) for hyperparameter tuning the grid-search procedure is used to tune the prior precision, i.e., weight decay, hyperparameter. On the CRISPR datasets, VI is trained for 500 epochs with a learning rate of 10- -3,3, For all methods using grid-search, we first split the training data into a 90/10 train-validation split. We then train the model for each order of magnitude of the prior precision from 10 3 to 10 The best prior precision is selected on the validation data and the model is retrained on the entire combined training dataset and assessed on the test data. The observation noise for the homoscedastic model is chosen on the training data as maximum likelihood solution. With the Empirical Bayes (EB) training approach, training data and optimize the prior precision values per layer of the neural network during training. We use the algorithm in Immer et al. [2021a] with hyperparameter learning rate of 0.01 decayed to 0.001 using a cosine schedule, 100 burn-in epochs, every 50 epochs for 50 hyperparameter steps. The pseudocode for the optimization is provided in App. F.1. For the UCI regression datasets, use the full Hessian Laplace approximation, while for CRISPR dataset, due to higher computational cost of training, we resort to the Kronecker-factored Laplace approximation. Finally for MC-Dropout, we set the dropout probability for dropout layers to 0.005. For all compared models, results are averaged over 20 independent runs on the UCI regression datasets and 10 independent runs on the CRISPR datasets, where independent runs use different seeds. We report standard errors. Runs were executed on computing cluster, but without the need for GPU support. Note that in our experiments to implement the VI baseline we make use of the PyTorch implementation from Krishnan et al. [2022].
only need to fit the model once on the entire
D.1.3 Image Regression
We use the same hyperparameters for MNIST and FashionMNIST but two different architectures. The architecture for MNIST is an MLP with hidden layers, each with width of 500 units and using ReLU activation function. The architecture for FashionMNIST is a LeNet with three convolutional layers with 6, 32, 256 channels, respectively, with max- poollnn in-between followed by two hidden layers of width 128 and 64. In all cases except for VI and MC-Dropout, we train the models with SGD using a batch size of 512 for 500 epochs and initial learning rate of 0.1 that is decayed to 10-4 using a cosine schedule. For MC-Dropout, we use a learning rate of 0.01 and train the model for 50 epochs with the Adam optimizer, with dropout probability for Dropout layers of 0.005. For VI, we use learning rate of 10 that is decayed to 10-4 using cosine schedule when training on the FashionMNIST dataset, while we use a non-decayed learning rate of 10 for the MNIST dataset. In both cases, the VI baseline is trained for 500 epochs with the Adam optimizer. In practice, we observe that the VI and MC-Dropout baselines are more sensitive to choices of hyperparameters than the other compared models, and we adopted the settings that would lead to the best results according to our experiments. In the models using classic grid-search (GS) for hyperparameter tuning (GS) the grid-search procedure is used to tune the prior precision, i.e., weight decay, hyperparameter. For all methods using grid-search, we first split the training data into a 90/10 train-validation split. We then train the model for each order of magnitude of the prior precision from 10-2 to 10a. 'he best prior precision is selected on the validation data and the model is retrained on the entire combined training dataset and assessed on the test data. The observation noise for the homoscedastic model is chosen on the training data as maximum likelihood solution. With the Empirical Bayes (EB) training approach, we only need to fit the model once on the entire training data and optimize the prior precision values per layer of the neural network during training. We use the algorithm in Immer et al. [2021a] with hyperparameter learning rate of 0.1, 50 burn-in epochs, every 25 epochs for 50 hyperparameter steps. These choices are in line with their suggestions. Due to the large models on the image datasets, we use the cronecker-factoree and inference time between Laplace and baselines are almost identical. The training was done 5 times (different seeds) per model-dataser pair to estimate mean and standard error and were run on a computing cluster with V100 and A100 NVIDIA GPUs.
Laplace approximation so that the training
19

---

D.2 Additional Experimental Results
D.2.1 Skafte Illustrative Example
In Figure 7. we additionally compare the proposed approach to Bayesian neu- ral network baselines, MC-dropout and mean-field VI. We find that the pro- posed heteroscedastic Laplace approximation with natural parameterization provides more visually appealing in-between epistemic uncertainties and has a more sensible split into aleatoric and epistemic uncertainties. In Fig- ure 6, we display the difference when the gen- erated data were homoscedastic: despite this, our well-regularized heteroscedastic method per- forms on par or better with a homoscedastic baseline trained in an identical way.
Homoscodastic Laplace
Heteroscedastic Laplace
2.5
2.5
-5
MC Dropout
MFVI
Homosccdastic Laplace
Heteroscodastic Laplace
2.5
0.0
-2.5
-5
-5
5
5
5
-5
5
22
ground truth
epistemic
alestoric
Figure 7: Comparison of heteroscedastic Laplace (right) to Monte-Carlo Dropout [Gal and Ghahra- mani, 2016], mean-field variational inference, and homoscedastic Laplace on the Skafte [Skafte et al., 2019] regression example. Laplace obtains better "in-between" uncertainties reported also by Foong et al. [2019].
D.2.2 UCI Regression Datasets
In this section, report additional results the we show in Sec. 5.1 for UCI regression datasets. In Table 3 we report test Mean Squared Error (MSE) results for the models compared in Table while Table 4 reports ablation where for our proposed natural parameterized NLL heteroscedastic loss model EB regularization, comparing results obtained with the full Hessian Laplace upproximation (full) and results obtained with the Kromecker-factored Laplace approximation (kfac). Results show that the full approximation can improve results, and thus should be preferred when feasible. At the same time, the kfac approximation leads comparable results in most settings, justifyyng its use for larger datasets. In Table Table 5, we compare the results obtained with two different variants of posterior predictive for the Natural NLL (EB) loss on the UCI regression datasets, namely the closed-form posterior predictive proposed in Sec. 4.4 (mean) and the asymptotically exact Monte- Monte-Carlo estimate the NLL ['Tomezak et al., 2018] (lse). Finally Table shows that the results from our work for models trained with EB regularization favourably compare results from the work of Wu al. [2019] on a deterministic variational inference approach for BNNs that also adopts an EB procedure for regularization. Note however, that this last approach does not scale well to more complex architectures than an MLP and hence is not reported the other experimental settings.
LL (1) kinSrmm
Regular- Posterior ization Predictive
Objective
boston
concrete
naval
plant
wine
yacht
Homoscedastic
EB
-2.51 (0.04) 3303 (0.03) -0.71 (0.04) 1.29 (0.01) 5.58 (0.45) -283 (0.02) 0.94 (0.02) -1.05 (0.14)
Deterministi VI
x
-3.06 (0.01) -1.01
1.13
6.29 -2.80 (0.00)
(0.01) -0.47
V
-2.47 (0,04) -.888 (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)
Naive NLI.
EB
x
-6.03 (1.82) -3.62 (0.26) -0.97 (0.15) 1.36 11000011-10000 (8.01) -283 (0.04) -2.26 (0.51) -1.27 (0.51)
V -2.36 (0.03) -2.93 (0.02) -0.71 (0.05) 1.36 (0.01) 6.66 (0.01) -2.76 (0.02) -0.94 (0.02) -0.51 (0.04)
NLL
X -2.51 (0.09) -2.99 (0.04) -0.65 (0.05) 1.35 (0.01) 6.68 (0.01) -2.77 (0.02) (0.03) -0.60 (0.08)
Table 6: Comparison of our models trained with EB regularization, with the work of Wu al. [2019). that also adopts an EB regularization procedure.
20

---

MSE (L)
Regular-
Objective
Ization
boston
concrrete
ewerryy
kinkium
Navil
plant
wilee
yachi
Homosedustic
EB
0.130 (0.012) 0.096 (0.005) 0.002 (0.000) 0.064 (0.001) 0.001 (0.000) 0.057 (0.002) 0.606 (0.024) 0.002 (0.000)
Naive NLL
GS
0.170 (0.020)
(0.033)) 0.009 (0.005)
0.072 (0.001)
5.673 (5.216) 0.056 (0.002) 0.605 (0.022) 0.003 (0,000)
8-NLL(0.5)
GS
0.131 (0.013) 0.122 (0.007) 0.002 (0.000)
0.069 (0.001) 6.293 (5.367) 0.055 (0.002) 0.608 (0.023) 0.002 (0.001)
8-NLL.
GS
0.144 (0.012) 0.108 (0.029) 0.003 (0.001) 0.069 (0.001) 0.003 (0.000) 0.057 (0.002) 0.609 (0.022) 0.002 (0,000)
Faithful
GS
(0.015) (0.007)
0.070
(0.001)
(0.023)
MC-Dropout
GS
0.131 (0.016) 0.131 (0.006) 0.057 (0.003)
0.117 ( (0.004) 0.145 (0.011) 0.059 (0.002) 0.606 (0.025) 0.124 (0.038)
VI
GS
0.131 (0.016) (0.00) 0.056 (0.004)
0.076 (0.002)
(0.009)
(0.037)
Naive NLL
0.136 (0.016) (0.066)
0.065 (0.001) 0.027 (0.017) (0.002) (0.027) 0.003 (0.001)
GS
0.111 0.081 (0.088) 0.002 (0.000) 0.067 (0.001) 0.001 (0.000) (0.002) 0.612 (0.02) 0.003 (0,00))
Natural NLL
HB
102 (0.011) 0.084 (0.001) 0.002 (0.000) 0.061 (0.001) 0.001 (0.000) 0.052 (0.002) 0.652 (0.046) 0.003 (0.000)
Table 3: Test MSE results for the compared models the UCI regression datasets, where we report mean and standard error (dataset names italic).
LL (f)
Objective
Regular- ization
Apgroximatiin
Predictive
Posterior
bostow
concrete
kinSom
navval
plant
wätte
yacht
kfac
-2.44 (0.04) -3.05 (0.02) -0.44 (0.02) 1,36 (0.01) 6.64 (0.01) -276 (0.02) -0.94 (0.02) -0:69 (0.05)
Natural NLL
EB
full
V -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)
kíacc
X -2.51 (0.07) -3.04 (0.03) -0.83 (0.03) 1.35 (0.01) 6.69 (0.0)) -2.76 (0.02) -5.41 (3.54) -0.64 (0.05)
-2.99 (0.04) (0.05) 1.35 (0.01) 6.68 (0.91) (0.02) -1.00 (0.03)
(0.08)
Table 4: Ablation comparing test log likelihood results for the natural parametrization with empirical Bayes (Natural Laplace) adopting the full Hessian Laplace approximation (full) and the Kronecker- factored Laplace approximation (kfac).
LL (t)
Posterior Predictive
Regular- ization
Method
boston
energy
kin8inm
naval
plant
wine
yacht
Variant
Natural NLL
GS
mean
(0.04) (0.04) (0.05) 1.32 (0.01) 6.66 (0.01) 2.76 (0.02) 0,94 (0.02) 00,29 (0.07)
Ise
-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)
Natural NLL
EB
mean
-2,39 (0.04) -2,92 (0.03) -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.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)
Table 5: Alblation comparing the two variants of posterior predictive for Natural NLL objective with Empirical Bayes regularization: mean and Ise.
D.2.3 CRISPR Gene Knockdown Efficacy Datasets
Here we report additional results for the compared models in Sec. 5.2 on the CRISPR datasets for gene knockdown efficacy, reporting results for models trained on single experiment replicates. Therefore, we report test log likelihood performance for models trained on the available replicates for each of the three CRISPR datasets in Figure 8. Again, for each model in each dataset run 10 independent seeds. As it is expected, these results are in line with, and complement, the ones shown in Figure 3.
D.2.4 Image Regression
In Table 9, we additionally provide results with the exact same training setting as for Table but with homoscedastic generative process. Table 8, we provide results for the same task but with heteroscedastic noise that depends on the mean output, i.e., the rotational angle. Overall, the pattern the same: empirical Bayes and the posterior predictive strictly improve generalization performance and the natural parameterization seems to overall perform best and most consistently. In Table 7, we report the runtimes for training (EB vs GS) and testing (MAP vs. PP). Training with EB is much faster than using grid search but the posterior predictive requires runtime per sample.
21

---

Bow-cytometry -1168299
servivalscreen HEX291
Betoosedassti PP-
Naie NLL
(-NLL (0.5)
Farnnul
Nave
Natural NLL. PP
Naturll NLL FP
Natural
-20
-15
-1.0
-0.5
-06
-05
-14
-0.
-02
11.11
11.
(a) Replicate
survival-screau
Hommnsslatiic] GS PPp-
Nave
(-NLL (0.5)
(-N11.
H)
Faihful
V1 PP
Natura NIJ.] 1P0-
Natural NLJ.
Natrral PP
Nateal NLL. EB
-1.
-0.0
-00
-0.6
Lf
(b) Replicate 2
Bow-crometry
5 HEX295
survival screen-A375
survival screen HEK229
Homnnsslaaiic pP-
Nave NLL.
+
J-NLL (0.5)
(..L
Fainfal
Naive NLL. PP-
Natrral NLJ. US5
Natual NLL. EB
L.
L.
LL
(c) Replicate 3
Figure 8: Box-plots reporting test log likelihood results on the CRISPR datasets on single experiment replicated. For the Servivalscrree-A331 dataset, only two replicates are available. Note that the NLL mean-variance parameterization with MAP prediction, as well as the MC-Dropout baseline, markedly underperform in this setting, and are hence not reported in the comparisons.
B-NII.
Homoscedastic GS PP
Homoscedastic EB PP
Natural GS PP
Natural EB PP
Training
44 min.
44 min.
8 min.
49 min.
min.
Inference
0.023 sec.
0.124 sec.
0.124 sec.
0.134 sec.
0.134 sec.
Table 7: Average runtime over 5 trials including hyperparameter optimization.
E
Implementation
In the following, we detail the interleaved Laplace marginal likelihood optimization [Immer et al., 2021a] but adapted the heteroscedastie setting in App. E.I. Further, we provide an efficient and stable implementation of the natural parametrization of the heteroscedasttc Gaussian likelihood in App. E.2. We implement our method the peplace-toreed package [Daxberger al., 2021] and the KFAC for heteroscedastic regression in the automatic second-order differentiation library [asdl; Osawa, 2021]. Due to anonymity reasons, we implemented these changes as private forks and will only propose them to the open source packages upon publication of the paper. This ensures that the proposed method can used with any supported network by these packages and different scalable Laplace approximation variants.
22

---

MNIST with MLP
FashiomMNIST with CNN
Regular- Lzalion
Posserior Predictive
Objective
IT.(f)
PKL (4)
(M)
1I.(t)
PKL (4)
RMSE (
Homoscedastic
EB
V
4.02 (0.10)
(1.92
(1.9
-387 (0.02)
1.37 (0.08)
12.5 (0.2)
Naive
x
(0.25)
1.39 (0.25)
(7.7
(0.16)
(0.15)
12.7 (0.2)
G-NLL (0.5)
GS
x
-4.30
1.13 (0.17)
(6.9
4.01 (0.02)
0.84 (0.02)
(0.1)
B-NLLOO)
G8
X
.4.21 (0.09)
1.04 (0.09)
(5.0
.4.05 (0.04)
0.88 (0.03)
12.3 (0.1)
Faithful
GS
-4.17
0,99 (0.04)
(0.5)
1,25 (0.05)
17.8
MC-Dropout
GS
-4.14 (0.05)
0,97 (0.U5)
15.9 (0.6)
-3.78 (0.05)
0.61 (0.05)
(1.1(
.4.15 (0.01)
(0.01)
(0.1)
(0.03)
(0.03)
16.3 (0.7)
Naive NLL
EB
V
X
-3.92 (0.02)
-3.85 (0.02)
0.75 (0.01)
0.75 (0.01)
(2.3
(2.3
-3.77 (0.02)
.366 (0.01)
0.60 (0.02)
0.60 (0.02)
13.1
13.1 (0.3)
(0.01)
0.99 (0.01)
(0.1)
(0.01)
(0.02)
13.2 (0.4)
GS
Natural NLL
x
-4.16 (0.01)
0.99 (0.01)
[4.5 (4.1
-3.78 (0.02)
0.61 (0.02)
13.2 (0.4)
/ .3.73 (0.01)
0.76 (0.02)
11.0 (0.2)
.3.66 (0.01)
0.61 (0.01)
13.1 (0.2)
X -3.94 (0.02)
0.76 (0.02)
11.0 (0.2)
-3.78 (0.01)
0.61 (0.01)
13.1 (0.2)
Table 8: Image regression with heteroscedastic noise variance that depends on the rotation applied, ie.. y rot + /rot|s. Similar to the results with heteroscedastic label-based noise in Table 2, empiricial Bayes and the posterior predictive improve the performance consistentlly.
MNIST with MLP
FashionMNIST with CNN
Regular- ization
Posterior Predictive
Objective
(L
DEL (1)
RMSE (L)
(L(f)
DEL
RMSE
Homoscedastic
EB
V
-.115
0.86 (0.07)
12.8 (0.4)
.4.0 (0.01)
0.64 (0.02)
12.7 (0.2)
-5.38
1.66 (0.00)
(0.03)
(0.03)
J-NLL (0.5)
GS
X
(0.14(0.10)
0.62 (0.10)
(1.5
.4.27 (0.02)
0,54 (0.01)
15.5 (0.5(
(INLL(()
x
-4.31 (0.05)
0.59 (0.05)
(0.8)
.4.26 (0.04)
0.53 (0.03)
12.6 (0.3)
Faihhul
GS
X
-4.31 (0.03)
0.58 (0.03)
14.6 (0.5)
-4.57 (0.01)
0.84 (0.01)
21.0 (0.4)
MCDO
GS
..29 (0.04)
0.56 (0.04)
16.1 (0.9)
.408 (0.02)
0.36 (0.02)
(1.6
GS
.4300 (0.01)
0.58 (0.01)
(0.1)
4.100
0.38 (0.01)
(5.3)
/
-4.09 (0.01)
0.42 (0.02)
12.7
.4002 (0.01)
0.33 (0.01)
13.4 (0.2)
Naive NLL
EB
X
-4.14 (0.02)
0.42 (0.02)
12.7 (0.1)
-4.06 (0.01)
0.33 (0.01)
13.4 (0.2)
/
-4.26
0.54 (0.00)
(0.1)
-3,99 (0.01)
0.33 (0.01)
(0.2)
x
-4,26 (0.00)
0.54 (0.00)
14.3 (0.1)
-4.06 (0.01)
0.33 (0.01)
13.0 (0.2)
Natural NLL
/ -4006 (0.01)
0.38 (0.00)
12.0 (0.1)
-4.00 (0.01)
0.33 (0.01)
13.2 (0.2)
x
..12 (0.01)
0.38 (0.00)
12.0
4.06 (0.01)
0.33 (0.01)
13.2 (0.2)
Table 9: Image regression with homoscedastic noise variance that fixed at scale of 10. Similar to the results with heteroscedastic label-hased noise in Table empiricial Bayes and the posterior predictive improve the performance consistently. Also, homoscedastic and 3-NLL variants perform relatively well due the reduced heteroscedassicityy
E.1 Pseudo-Code for Our Method
Alg. is the heteroscedastic regression variant of the algorithm proposed by Immer et al. [2021a] and includes the empirical Bayes hyperparameter optimization of the prior precision, equivalently weight-decay, per layer. The algorithm optimizes neural network parameters and hyperparameters jointly during training and further the intermediate marginal likelihood values, log p(D)( can be used for early stopping. After training, either the Bayesian or point-wise posterior predictive can be used. The posterior predictive is efficient compute and does not require sampling due the simplification proposed in Sec. 4.4. Experimentally, we find that the Bayesian posterior predictive improves over the point-wiee wise predictive almost all cases.
23

---

Algorithm Optimization of Heteroscedastic Regression Models
Require: model f(x; 0) (natural parameters), dataset DD (initial) prior precision d, epochs E, burn-in B, frequency F, steps S, learning rate a, hyperparameter learning rate 7 Set prior precision (regularizer) for each layer
vl 67 <-
for epoch e I E do for batch B _ D do
Optimize model parameters O
0 + aVol
D
log (B))) + E log.V 6,1)]
SGD parameter update
end for
if e mod / = 0 and B then
Optimize regularization
for step $ v S do
:l: dt 6 + 36
log (D))) (Equation 12)
Hyperparameter update
end for end if
end
if (Bayesian) posterior predictive then
Compute (xxx),, ox.) as in Equation
Fast approximate posterior predictive
else
fx)x:00) μ(x.) f- (x) and o²(xx) + (x.(x.))
Point-wise MAP predictive
end if'
Predictive o((((x),,"(xx.)
E.2 Implementation of Natural Gaussian Log Likelihood
Common automatic differentiation packages like pytorch [Paszke et al., 2017] and jax [Bradbury et al., 2018] do include typical losses and distributions but not the natural parametrization of Gaussian. Below, we provide an implementation for using pytorch:
from math import log, pi import torch
0.5 12g pi)
def heteroscedastic_mee loss (input, target, reduction-'mean """Heteroscedastcc Normal negative log likelihood.
Parameters
input torch. Tensor (n, 2)
two natural parameters per data point target torch, Tensor (n, 1)
targets
assert input.ndi ndim == target ndim == target. shape [0]
assert
n, input.shape
target torch. cat ((targeet, target (quaree dim=1) inner
einseeinknnnnknne nk- target, input) log_A (input 0]. square () (4 input 1))
0.5 torch.log(
input 1)))
log lik -n C sumer.sum() log_A, sum
if reduction == 'mean':
return log_lik elif reduction == 'aum"
return log_lik
else:
raise ValueError ( Invalid reduction', reduction)
24


Reviews:
[{"review_id": "glXZAHnlTh", "reviewer": "Reviewer_2cQi", "summary": "This paper introduces a new method for fitting heteroscedastic regression models with neural networks. In contrast to previous work that highlights the deficiencies of MLE and indirectly hint at regularization, this paper suggests using natural gradients during training and combines Bayesian ideas for regularization and to handle uncertainty. This modeling approach accounts for both aleatoric and epistemic uncertainty through the posterior predictive distribution. Bayesian inference can be computationally expensive, so they utilize a Laplace approximation to the posterior and simplifying factorizations on the Hessian. Natural gradients avoid some of the issues with gradients that occur when training with other parameterizations ($\\mathcal{N}(\\mu, \\sigma^2)$). Empirically this method achieves strong results on real-world data for regression tasks as well as on image data.", "questions": "- How does this method perform on data that is homoscedastic (ie the model is misspecified for the task)? Does the variance function flatten out, or will it overfit the variances (or $\\eta_2$)?\n\n- Does empirical Bayes put this method at risk of overfitting?\n\n- Is $\\eta_2(\\cdot)$ necessary if it gets omitted when estimating the epistemic uncertainty (line 250) and what was the reason for this choice if it could have been done at low cost?\n\n- How does this type of \"Bayesian-ness\" compare to the ideas presented in Stirn and Knowles (2020) or Detlefsen et al (2019) in terms of interpretation?\n\n", "limitations": "There was not much discussion of this, but that is appropriate for the nature of this work.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- Problem is well-motivated, and the distinctions between this method and existing ones are clear\n- Experiments show solid results and extend the method to a setting with image data as opposed to only tabular data\n- Combines ideas from approximate Bayesian inference with a parameterization commonly used for Gaussian processes ", "weaknesses": "- typo line 98: \"the mean respectively the covariance matrix\" --> \"the mean and covariance matrix, respectively\"\n- Structurally, it was odd to see Problem 2.1 presented so high up in the paper\n- I would appreciate more discussion on the differences between the natural MAP vs natural Laplace methods. In particular, the differences in the interpretations of the two models (what sorts of uncertainties they account for) and when to use which", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive evaluation of our work and your constructive comments and questions.\n\n**Weaknesses**\n\n> It was odd to see Problem 2.1 presented so high up in the paper.\n\nCurrently, Section 2 serves as a theoretical comparison to the most related works. Within this context, we introduce the problem for the case-study presented in Figure 2 to highlight the conceptual difference to our approach. Given your feedback and the feedback of the other reviewers, we plan to make the discussion of related work more direct in this section complementing our discussion in the introduction, Section 2 and Appendix B and C. We can defer the definition to the experiments. We also welcome other suggestions.\n\n> I would appreciate more discussion on the differences between the natural MAP vs natural Laplace methods.\n\nThese only differ in how the regularization is optimized: for \"MAP\", we optimize the regularization on a validation set, while for \"Laplace\", we use empirical Bayes (ML-II; see Sec. 4.2) to optimize the regularization jointly on the training data. We provide the pseudocode for the optimization in Appendix E1: for \"MAP\" the optimization of $\\delta$ and the corresponding hyperparameter update are omitted. We will make this clearer by adding replacing Laplace with \"empirical Bayes (EB)\".\n\n**Questions**\n\n> How does this method perform on data that is homoscedastic (ie the model is misspecified for the task)? Does the variance function flatten out, or will it overfit the variances (or $\\eta_2$)?\n\nThanks for the interesting question. We added results for this on the Skafte and image regression tasks with modified homoscedastic noise to the rebuttal pdf in Figure 8. We indeed find that our method successfully regularizes towards a homoscedastic aleatoric uncertainty.\n\n> Does empirical Bayes put this method at risk of overfitting?\n\nSince we only optimize the regularization strength, the potential to overfit is quite limited. Our experiments confirm this: we do not find it to overfit compared to validation-based selection of the regularization.\n\n> Is $\\eta_2$ necessary if it gets omitted when estimating the epistemic uncertainty (line 250) and what was the reason for this choice if it could have been done at low cost?\n\nYes, $\\eta_2$ is necessary since the Jacobian $J_\\mu(x)$ depends on it and is required to approximate the epistemic uncertainty. This is a low-cost estimate since it is in a closed form.\n\n> How does this type of \"Bayesian-ness\" compare to the ideas presented in Stirn and Knowles (2020) or Detlefsen et al (2019) in terms of interpretation?\n\nStirn & Knowles (2020) and Detlefsen et al. (2019) propose non-Bayesian methods and therefore infer only heteroscedastic aleatoric uncertainties $\\sigma^2(x)$. They have no epistemic uncertainty, which would require uncertainty about the model or its parameters. For a visual depiction of the additional epistemic uncertainty, see \"Heteroscedastic Laplace\" on the right in Fig. 1.\n\nWe hope that we could clarify your questions and if so, would appreaciate if you considered revising your score."}, {"review_id": "Q96MqUQWXt", "reviewer": "Reviewer_PsPP", "summary": "This article focuses on refining the techniques used to manage two types of uncertainties in complex regression tasks utilizing deep neural networks. The uncertainties, known as aleatoric (arising from inherent randomness in the data) and epistemic (originating from the model's limitations), are critical to address for robust and accurate predictions. The authors propose an innovative approach using the natural parameterization of the Gaussian likelihood to overcome the gradient scaling issue commonly encountered in traditional methodologies. Additionally, they introduce an efficient Laplace approximation that enhances heteroscedastic neural networks, providing epistemic uncertainties, and facilitating automatic regularization through empirical Bayes. This method outperforms earlier strategies in heteroscedastic regression, demonstrating scalability and obviating the need for hyperparameter tuning. Empirical validation on diverse datasets, including a new image dataset, UCI regression, and CRISPR-Cas13, yielded superior performance, suggesting potential applicability to other real-world datasets in the future.", "questions": "- How large is the computation cost and time of the proposed method in comparison to the rival methods? In case it is significantly heavier, it should be properly discussed in the limitation section.", "limitations": "The limitations of this work are discussed properly.  ", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- They are addressing very important shortcomings in one of the best algorithms for quantifying uncertainty using neural networks.\n- Their method seems mathematically solid and is supported by extensive experiments.\n- The code is available, and this work seems to be reproducible.\n- Their method is thoroughly compared with the rival algorithms.\n- Compared to rival methods, this method has less hyperparameters. \n\n", "weaknesses": "I cannot spot any weaknesses in this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive remarks! Below we would like to clarify your concerns regarding computational cost of our method.\n\n> How large is the computation cost and time of the proposed method in comparison to the rival methods?\n\n1. Using the natural parameterization of the likelihood does not change the computational cost or complexity.\n2. Using the Laplace posterior and predictive incurs a one-time cost of computing the approximation (Sec. 4.1) and the posterior predictive is slightly more expensive (Sec. 4.3) with $\\mathcal{O}(P\\sqrt{P})$ instead of $\\mathcal{O}(P)$ for an MLP with fixed hidden layer sizes using KFAC. Empirically, the predictives are very fast and this is not a bottleneck. In comparison to the standard predictive, Table 9 in the rebuttal pdf shows that the posterior predictive is only 5x slower than a single forward pass without epistemic uncertainty.\n3. Using empirical Bayes even **reduces** the computational cost because it only requires a single training run in comparison to validation-based selection of the regularization. This is apparent from Table 9 as well where the runtime of empirical Bayes is roughly 5 times faster when the hyperparameter tuning is included.\n\nOverall, the proposed natural Laplace method leads to faster training than rival methods when including the regularization strength selection and a slight increase in runtime of the posterior predictive.\n\nWe will clarify the above points in Sec. 4.4 and hope that our answer alleviates your concerns regarding the computational cost and time of our method. If so, we would appreciate if you considered revising your score."}, {"review_id": "zCRBvodeqy", "reviewer": "Reviewer_HwYi", "summary": "This paper extends the technique of the Laplace Approximation to incorporate heteroscedastic aleatoric uncertainty. The proposed technique achieves this by exploiting the natural parameters of the Gaussian likelihood. Main motivation is the coupling of the mean and the input dependent variance in the Gaussian likelihood. Using the natural parameters, the paper suggests the heteroscedastic Gaussian log-likelihood, linearized Laplace Approximation, and marginal likelihood for obtaining the prior precision term. Two standard benchmarks including UCI and one self-made benchmarks, improvements over the previous methods are demonstrated.", "questions": "(minor comments/questions)\n\n- in lines 20-21: isnt active learning and reinforcement learning part of the decision making problems?\n\n- in lines 22-24: there are lots of work on making data noise as function of inputs. It might be an outdated statement.\n\n- lines 48-54 should be made easier to read.\n\n- Can you comment on the practical relevance of the problem 2.1?\n\n- Is it possible to refer to all the derivations? (if they are present in the appendix?)\n\n- line 198: is it a typo the jacobian term?\n\n- is it possible to comment on the complexity of the overall pipeline, when compared to the standard linearized Laplace Approximation?", "limitations": "The limitations sections exist.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "In my opinion, the followings are the strengths of this paper:\n\n- Presentation of the paper is generally clear to see the contributions of the paper. I appreciated especially the section 2, which carefully addresses the problem at hand. \n\n- The contribution of the paper is relevant to the current Bayesian Deep Learning community. For example, extension of the Laplace Approximation to consider input dependent aleatoric uncertainty is meaningful. To do so, the paper provides long stretch from finding the problem in standard formulation of the Gaussian likelihood, to showing the adaptations on linearized Laplace Approximation, Gaussian log likelihood, predictive distribution and marginal log likelihood.\n\n- Experimental results show good performance.", "weaknesses": "The followings might be the weaknesses of the paper (all major):\n\n- The experiments only focus on regression problem, which is limited in scope for a paper on Bayesian Deep Learning.\n\nThe paper mainly focuses on the regression problems, e.g., all the experiments as well as the methodologies (section 2 and section 5). One way to address this point could be certain changes in the title. Would it be possible to include the term \"regression\" in the title?\n\n- Certain parts of the paper need revision in presentation.\n\n1. Many technical terms are introduced without explaining them well. Major examples are in the introduction 2nd paragraph, e.g., feasible generalized least squares, natural parameters, standard FGLS, etc. \n\n2. The logic flow in the introduction may not be very kind to the reader. In particular, I had difficult time reading the 3rd paragraph. After reading the paper, I could grasp the concept well, but my feeling is that it is diving too much into the details early on.\n\n3. I left minor comments later in the review.\n\n- The baselines are restricted to the area of heteroscedastic aleatoric uncertainty estimation, which may be limited for broader audience. The scale of the experiments are also limited to toy-ish benchmarks. One option would be to use the uncertainty-baselines benchmarks.\n\nIn all the experiments, the baselines are rather restricted. While these choices validate the main point of the paper, I think the paper could improve by also competing against generic state of the art in uncertainty estimation. Would it be possible to include atleast MC-dropout, deep ensemble, and maybe stochastic HMC? (with homoscedastic aleatoric terms?) In this way, one could examine the importance of heteroscedastic aleatoric uncertainty estimation.\n\n- Related work section is missing in the paper, which makes it harder to locate this work within the state-of-the-art.\n\nTo back up, there is no related work section in the main body of the paper. This makes it difficult to locate this work within the current state of the art. Would it be possible to include them? (one could shorten other parts of the paper). Some important areas are the usage of natural parameters of the Gaussian distribution, and more detailed treatment on aleatoric uncertainty, e.g., calibration methods, combination of both model and data uncertainty, etc. To name only few, highly relevant works seem to be: (1) Natural-Parameter Networks: A Class of Probabilistic Neural Networks (one of the works that uses natural parameters in more general context) and (2) estimating model uncertainty of neural networks in sparse information form (one of the works that uses natural parameters of Gaussian distribution for Laplace Approximation). Why the authors build on regularization based aleatoric uncertainty estimation method should also be mentioned.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed feedback and constructive questions which will help inprove our manuscript.\n\n> The paper mainly focuses on the regression problems. [...] Would it be possible to include the term \"regression\" in the title?\n\nThank you for pointing out that heteroscedastic does not imply the focus on regression. We are happy to adjust the title to include \"regression\", for example, \"Effective Bayesian Heteroscedastic Regression: Model Selection and Epistemic Uncertainties\".\n\n> The baselines are restricted to the area of heteroscedastic aleatoric uncertainty estimation, which may be limited for broader audience. [...] Would it be possible to use the uncertainty-baselines benchmarks?\n\nWe believe that the problem of regression, as opposed to classification, is similarly important in practice and for the broader audience but relatively underexplored in deep learning. In fact, the uncertainty-baselines benchmarks illustrate this issue since they exclusively include classification problems.\n\n> In all the experiments, the baselines are rather restricted. [...] Would it be possible to include atleast MC-dropout, deep ensemble, and maybe stochastic HMC with homoscedastic aleatoric terms?\n\nThanks for these suggestions. We are happy to include these as baselines. For UCI regression datasets, we include the results provided by (Wu et al, 2019; Gal et al 2016). These works use the same setup as we do, so results are comparable and already provide evidence that our approach outperforms these additional baselines, speaking to the significance of our work. However, we want to note that deep ensembles are applicable to all the baselines and our method, see Eschenhagen et al. (2021). We would therefore add an ablation that studies the effect of additional ensembling for ours and other methods. We are happy to extend these baselines to all our experiments in the camera-ready version and will assess the feasibility of an HMC baseline.\n\nReferences:\n\n- Wu A., et al. Deterministic Variationa Infeference for Robust Bayesian Neural Networks. In ICLR, 2019. \n- Gal Y., and Ghahramani Z. Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning. In ICML, 2016.\n- Eschenhagen R., et al. Mixtures of Laplace approximations for improved post-hoc uncertainty in deep learning. Bayesian Deep Learning Workshop 2021.\n\n\n> Certain parts of the paper need revision in presentation. [...] undefined technical terms, confusing introduction.\n\nThank you for the clear description of what was unclear. We will improve the text accordingly.\n\n> Related work section is missing in the paper, which makes it harder to locate this work within the state-of-the-art. [...]\n\nWe understand that while we discuss the most recent related work in great detail (Sec. 2, Appendix B and C), we should improve the discussion of related work in a broader context. Thus, we will extend the discussion of related work in Sec. 2, making sure that we cover the Baselines that we now added, and your suggestions discussed below.\n\n> relevant works seem to be: (1) Natural-Parameter Networks: A Class of Probabilistic Neural Networks and (2) estimating model uncertainty of neural networks in sparse information form.\n\nThank you for the additional references, we will refer to them but want to point out that they are rather complementary to our work: (1) uses natural exponential families for posterior approximations in Bayesian neural networks and (2) could be used as alternative to KFAC for the Laplace posterior approximation. Both only apply to homoscedastic regression but could very well be combined with our approach to be applicable for heteroscedastic regression. \n\n> Isn't active learning and reinforcement learning part of the decision making problems?\n\nThanks for pointing out this error, we will fix it.\n\n> Suggestions made in the \"Questions\" section\n\nWe will make the corresponding changes and clarifications: 1) active and reinforcement learning part of decision making problems; 2) correct lines 22-24; 3) clarify lines 48-54; 4) refer to derivations in appendix; 5) clarify Jacobian term.\n\n> Can you comment on the practical relevance of the problem 2.1?\n\nRecent works on heteroscedastic regression focus on ad-hoc regularizations that seem to work for certain problems but fall short when applied to more complex problems. Problem 2.1 is a prototype of a complex heteroscedastic regression problem that a good algorithm should be able to solve.\n\n> Is it possible to comment on the complexity of the overall pipeline, when compared to the standard linearized Laplace Approximation?\n\nThe standard linearized Laplace approximation is not applicable to the heteroscedastic case due to potentially negative definite Hessian, which our approach fixes using natural parameters. The complexity of our approximation is the same as for homoscedastic regression. This also shows in the empirical measurements in Table 9 of the rebuttal pdf where homoscedastic and heteroscedastic inference have the same runtime. We will further clarify the computational complexities in the appendix.\n\nWe hope that we could address your questions and remarks in our response and the attached pdf with additional results. If so, we would appreciate if you considered revising your score."}, {"review_id": "o2TuvR1log", "reviewer": "Reviewer_7fHN", "summary": "The authors consider the task of properly modeling heteroscedastic observation noise with Bayesian Neural Networks (BNNs).\nThey identify several shortcomings in prior work concerning variance modeling and propose to rely on parameterization\nvia natural gradients instead. Together with a Laplace approximation to model epistemic uncertainty and a closed-form\nposterior predictive likelihood approximation the proposed approach is evaluated on three sets of experiments. ", "questions": "- I am not sure how to parse the sentence in l27/28. Hasn't epistemic uncertainty modeling been the main focus for now\n    with aleatoric uncertainty modeling being of secondary interest? Can the authors provide some intuition on what they\n    mean here? (I might just be missing the obvious...)\n- Figure 2: How does the corresponding test NLL look like for the proposed approach on this task?\n- On the VI setups:  \n    - Can the authors speculate on the performance of VI in the naval and plant experiments? In both data sets the \n            performance is not just a drastic outlier compared to the other models, but also with respect to prior results \n            reported in these setups. See, e.g., the results reported in the original paper introducing this set of experiments (Hernández-Lobato & Adams, 2015), as well as numerous other BNN and GP papers on this setup who report variational inference results as part of their baselines (e.g., Bui et al., 2015; Wu et al., 2019; Haussmann et al., 2020,...).\n    - A parallel question applies to VI on MNIST where it performs exceptionally bad, whereas I would expect a similar performance at least to the \n        homoscedastic setup. \n    - Wrt the runtime. VI should in practice be only slightly slower than the other methods. Especially given the\n        statement in the appendix that for all methods in the non-image data, CPUs were sufficient, the argument on why\n        the hyperparameter selection was limited is not completely clear to me. Can the authors elaborate?\n    - Have the authors experimented with a simpler approach than Blundell et al.'s mixture prior, relying instead on a \n        simple mean-field normal prior with local-reparameterization (Kingma et al., 2015) to stabilize the gradients?\n        Then a single weight sample per forward pass should be sufficient to train the models without stability issues\n        on both UCI as well as the image data sets without too many problems increasing the runtime cost to only twice\n        as much as a deterministic net would cost. \n- In UCI and CRISPR, the proposed MAP approach performed almost always similarly to/better than the Laplace approach. \n    This behavior drastically changes for the synthetic image data experiments, especially FashionMNIST. Can the \n    authors speculate/provide a discussion on why the additional structure provided provides this huge difference?\n\n\n\n_____\nBui et al., 2016: Deep Gaussian Processes for Regression using Approximate Expectation Propagation  \nHaussmann et al., 2020: Sampling-Free Variational Inference of Bayesian Neural Networks by Variance Backpropagation  \nKingma et al., 2015: Variational Dropout and the Local Reparameterization Trick  \nWu et al., 2019: Deterministic Variational Inference for Robust Bayesian Neural Networks \n", "limitations": "Limitations of the model are discussed, but the societal impact is not. Given the theoretical nature of the paper,\nI do not consider this lack of discussion to be a problem.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper considers a well-defined task, identifies prior weaknesses, and offers a principled solution.\n- The paper is well written and can be understood easily by a reader\n- The experimental setups cover a wide range of domains", "weaknesses": "- The individual steps, i.e., switching to natural parameters, reliance on Laplace approximations, and approximate\n    closed-form marginalizations for the posterior predictive are mostly minor adaptations from prior work\n- Section 5.2 reads more like an afterthought that is added to the paper to have an additional experiment. \n    - The setup description is mostly left to the prior work which introduced the experiment.\n    - The evaluation of the results is similarly vague. E.g., Heteroscedastic VI (PP) has a huge variance in Figure 1, Exp 1,\nwhereas it suffers no such problems in the individual replications. The Laplace approach of the proposed method \nnot only performs a little bit worse than the others in setting two (survival-screen-A375), but struggles a lot,\nwhereas the MAP approach does not suffer from this problem. \n\n\n## Minor \n- None of the tables do follow the NeurIPS style guide. Captions should be placed above a table, and tables should not contain vertical lines (fig 2)\n\n\n## Typos\n- Fig 2 right and Table 1/2 switch the se notation (from $\\pm$ to brackets)\n- Between Table 1 and 2 the formatting of the method column switches from right-aligned to left-aligned", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the overall positive evaluation of our paper and your detailed and constructive comments and questions!\n\n> Section 5.2 reads more like an afterthought that is added to the paper to have an additional experiment.\n\nThank you for your comment. This is mostly due to lack of space. We will give the section more prominence, and elaborate more on its nature and relevance.\n\n> Evaluation of the results is vague. E.g., Heteroscedastic VI (PP) has a huge variance in Figure 1, Exp 1, whereas it suffers no such problems in the individual replications. The Laplace approach of the proposed method not only performs a little bit worse than the others in setting two (survival-screen-A375), but struggles a lot, whereas the MAP approach does not suffer from this problem.\n\nRegarding VI: during training on the mean response, some runs diverge to outlier values despite exhaustive tuning of hyperparameters and the use of common VI libraries and their recommended settings. Hence those lead to a large error (while the majority of runs achieves good results). We provide more details regarding the VI baseline below.\n\nRegarding Laplace: this is indeed curious and we will investigate further. We want to emphasize that these are challenging real-world datasets and we don't expect any single method to perform best on these. Despite that, we find it encouraging that our method performs best in two out of three settings.\n\n> I am not sure how to parse the sentence in l27/28. Hasn’t epistemic uncertainty modeling been the main focus for now with aleatoric uncertainty modeling being of secondary interest? Can the authors provide some intuition on what they mean here?\n\nIn the case of heteroscedastic regression, recent approaches, e.g. the faithful and $\\beta$-NLL methods, are (variations of) maximum likelihood estimation and therefore do not provide epistemic uncertainties. This observation has motivated our statement. On the other hand, the Bayesian community indeed focused more on epistemic uncertainties instead of aleatoric uncertainties. We will clarify this.\n\n> Figure 2: How does the corresponding test NLL look like for the proposed approach on this task?\n\nThe NLL is reported in Table 2 on the left. Note that the homoscedastic baseline in Fig. 2 is different since we use the empirical Bayesian homoscedastic method in Table 2 for its better performance.\n\n> In UCI and CRISPR, the proposed MAP approach performed almost always similarly to/better than the Laplace approach. This behavior drastically changes for the synthetic image data experiments.\n\nWe hypothesize that the proposed empirical Bayesian approach with Laplace particularly helps for deeper neural networks because it can adjust the prior precision/weight decay individually per layer, which is not tractable with cross-validation. On UCI and CRISPR, the networks only have one hidden layer. We therefore believe that the proposed benchmark is interesting for heteroscedastic neural network regression because it requires more complex networks than prior benchmarks.\n\n> Can the authors speculate on the performance of VI in the naval and plant experiments? In both data sets the performance is not just a drastic outlier compared to the other models, but also with respect to prior results reported in these setups. [...]. A parallel question applies to VI on MNIST where it performs exceptionally bad, whereas I would expect a similar performance at least to the homoscedastic setup.\n\nThank you for drawing our attention to these matters. As stated above, we had training stability issues despite extensive manual tuning for VI and some runs end up diverging while others give good performance. As per your suggestion, we will improve our VI baselines with local reparameterization and simpler prior. \n\n> VI should in practice be only slightly slower than the other methods. Especially given the statement in the appendix that for all methods in the non-image data, CPUs were sufficient, the argument on why the hyperparameter selection was limited is not completely clear to me. Can the authors elaborate?\n\nVI runs 10 times slower due to 10 MC samples. Less MC samples during training showed even more convergence issues. Therefore, we tested hyperparameters manually to obtain a suitable grid for validation. Nonetheless, some outlier runs reduce the performance. As stated above, we hope to alleviate these issues with flipout or local reparameterization. \n\n> Have the authors experimented with a simpler approach than Blundell et al.'s mixture prior, relying instead on a simple mean-field normal prior with local-reparameterization (Kingma et al., 2015) to stabilize the gradients?\n\nThank you for the suggestion. As stated above, we will experiment with these. However, we want to point out that we used a standard implementation for VI in neural networks (`blitz-bayesian-pytorch`) and followed suggested defaults (see also Appendix D.1.2 ll. 599-607). For now, we added the MFVI performance from the suggested DVI paper for comparison in Table 8 (see rebuttal PDF).\n\n> Minor comments and Typos\n\nThanks for pointing these out. We will make the necessary changes.\n\nWe will include the above clarifications, as well as the improvements mentioned in the general response into our manuscript. We hope that these improvements positively influence your assessment of our paper."}]

Author rebuttal:
We thank the reviewers for their time and constructive feedback on our manuscript. It is encouraging to see that the reviewers rated our paper overall positively pointing out that our paper *"identifies prior weaknesses, and offers a principled solution"* (Reviewer 7fHN), that Reviewer HwYi *"appreciated especially the section 2, which carefully addresses the problem at hand"*, that our *"method seems mathematically solid and is supported by extensive experiments"* (Reviewer PsPP) and lastly that Reviewer 2cQi acknowledges our dedication to *"extend the method to a setting with image data as opposed to only tabular data"*. 

Further, as asked by the reviewers, we investigated the following additional aspects (cf. attached pdf): 

* **Baselines (HwYi, 7fHN)**: After confirming that the papers of Wu et al., (2019), and Gal and Ghahramani (2016) follow the same experimental setup as we do for the UCI datasets, we included their results for MFVI (Graves, 2011), DVI (Wu et al., 2019), and MC Dropout (Gal and Ghahramani, 2016) as shown in Table 8. We find that our proposed approaches still reach the overall best performance.
* **Homoscedastic Data (2cQi)**: We performed additional experiments to confirm that our approach also performs well on homoscedastic data: a) we plot the mean and uncertainty estimates for a homoscedastic variant of the data generated for Figure 1 and b) compare both on a homoscedastic variant of the image regression task in Table. The results are in Figure 8 of the rebuttal pdf. They show that the proposed approach does not overfit and instead recovers the homoscedastic aleatoric uncertainty.
* **Runtime (PsPP)** In general, we would like to note that the Laplace methods do not require cross-validation, which is necessary for MAP and the remaining baselines and therefore have an advantage in terms of runtime. In Table 9 of the rebuttal, we show runtimes for the image regression task for training (including hyperparameters) and inference. 
* **Visualizations** In addition to showing the differences between MAP and Laplace, we also plot the mean and uncertainty estimates for MC Dropout and MFVI for the example provided in Figure 1 (cf. Figure 7, rebuttal pdf)---showing a favourable fit for our method.

Beyond these experiments, we will also add MC Dropout and an improved VI baseline to the remaining experiments, and add an ablation in which we instantiate our methods and baselines as deep ensembles (Reviewer HwYi). Further, we will extend Section 2 to discuss more related work---especially from the Bayesian learning literature as suggested by the reviewers, and lastly, include all clarifications provided during the rebuttal into our manuscript.

- Wu A., et al. Deterministic Variational Infeference for Robust Bayesian Neural Networks. In ICLR, 2019. 
- Gal Y., and Ghahramani Z. Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning. In ICML, 2016.
- Graves A. Practical variational inference for neural networks. In NeurIPS, 2011

Meta-review: This paper addresses an important problem: inappropriate homoskedasticity assumptions in Bayesian neural nets (or any model, really) make it impossible for Bayes's rule to yield reasonable uncertainty estimates (and can also distort point estimates). The reviewers agree that this paper offers a principled, effective recipe for addressing this problem, and are all in favor of acceptance. I look forward to seeing the final version with the reviewers' suggestions incorporated.
Decision: Accept (poster)