Title: Effective Bayesian Heteroscedastic Regression with Deep Neural Networks
Venue: NeurIPS 2023
Accepted: yes
Keywords: Heteroscedastic Regression

Marginal Likelihood

Bayesian Neural Networks

Uncertainty Estimaton

Model Selection

Laplace Approximation
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

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 problemsincluding 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.

1 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 learning [Houlsby et al., 2011, Kirsch, 2023], reinforcement learning [Osband et al., 2016, Yu et al., 2020] and decision making. Bayesian neural networks allow for modelling both epistemic and aleatoric uncertainty, as such, they lend themselves naturally to this task. Typically, they are applied under the assumption of homoscedasticity, i.e., 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, DVI], and 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 ∈ R D and a scalar response y ∈ R , and model y | x = x as a conditional Gaussian distribution with mean µ ( x ) and standard deviation σ ( x ) being dependent on input x . Inherently, the problem involves robustly modelling the aleatoric uncertainty, corresponding to the variance, which is a relevant problem in econometrics, statistics [Harvey, 1976, Amemiya, 1985], and causal discovery [Guyon et al., 2019, Xu et al., 2022].

-

5

0

5

x

-

2

.

5

0

.

0

2

.

5

y

Homoscedastic MAP

-

5

0

5

x

Homoscedastic Laplace

-

5

0

5

x

Heteroscedastic MAP

-

5

0

5

x

Heteroscedastic Laplace

ground truth

mean

epistemic

aleatoric

Classical approaches based on linear models and non-linear feature maps [Cawley et al., 2004] attempt to 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 if one of the parameters is kept constant, but it is not jointly convex in both parameters [Cawley et al., 2004, Yuan and Wahba, 2004]. To alleviate this problem in the context of Gaussian processes, [Le et al., 2005] propose to reparametrize the objective in terms of the natural parameters, which induces a jointly convex objective. Advantages of the natural parametrization compared to standard FGLS have also been demonstrated for non-linear features maps [Immer et al., 2023a].

In the neural network literature, the standard approach is to model µ ( x ) and σ ( x ) as outputs of the network and maximize the corresponding Gaussian log likelihood via stochastic gradient descent [Nix and Weigend, 1994, Lakshminarayanan 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 a surrogate loss, the β -NLL loss, which regulates the influence of the variance on the gradients of the loss by introducing a stop-gradient 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 to previous work on heteroscedastic regression with neural networks, we take a different perspective postulating that current estimators lack principled regularization. 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 models [Immer et al., 2023a], we propose to 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-stopping signal to prevent overfitting without requiring a grid search based on a validation set in Sec. 4. 2 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 is illustrated on the right in Figure 1 and shown in comparison to the prediction with only aleatoric uncertainty (MAP) as well as the corresponding homoscedastic regression solution on the left.

Besides showing that our approach performs favorably on commonly used UCI benchmarks and the CRISPR-Cas13 knockdown efficacy datasets [Stirn et 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.

2 Heteroscedastic Regression with Neural Networks

Due to their universal approximation guarantees [Hornik et 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 tendency to overfit to the data. Due to the additional degree of freedom granted by learning both mean and variance in heteroscedastic regression, this problem is 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 of 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 a neural network can lead to overconfident variance estimates that compromise the mean fit. Taking into account the gradients of the negative log likelihood (NLL),

where θ denote the neural network parameters used to estimate mean and variance, Seitzer et al. [2022] trace the problem to the fact that ∇ µ ℓ µ,σ ( θ ) = µ ( x ; θ ) -y σ 2 ( x ; θ ) heavily depends on the learned variance. To overcome this effect, Seitzer et al. [2022] introduce the β -NLL loss ℓ β ( θ ) , which is equal to ⌊ σ 2 β ( x ; θ ) ⌋· ℓ µ,σ ( θ ) , where ⌊·⌋ denotes a stop-gradient operation, and β is a hyperparameter controlling the dependency of gradients on the predictive variance. As a result, the gradients for β -NLL are equal to

With β = 0 , ℓ β ( θ ) is equivalent to ℓ µ,σ ( θ ) , whereas for β = 1 the gradient with respect to the mean is proportional to the gradient for homoscedastic regression. When setting 0 < β ≤ 1 we interpolate between both settings. As an alternative approach, Stirn et al. [2023] propose to decouple the estimation into three networks: a shared representation learner f z computes a representation z from x , which is passed into two individual networks f µ and f Σ , 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 an equivalent effect on the mean gradient for f z and f µ as setting β = 1 in β -NLL, and the second one stops any gradient from the variance network f Σ from propergating to f z [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 stop-gradient operations introduced by Stirn et al. [2023] have a similar, yet stronger regularization effect for the variance-i.e., stopping its influence on the joint representation learner. An additional hurdle presents itself due to the fact that we need to tune the regularization to calibrate the models for a 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 is 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 aleatoric uncertainty.

In the β -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 σ 2 ≤ 1 and hence the gradient with respect to the variance is amplified stronger compared to the mean. Especially in near-deterministic processes ( σ 2 → 0 ) this might be problematic. Further, it is not clear which value of β is best suited for a problem a priori introducing another hyperparameter that has to be tuned in addition to tuning regularization hyperparameters. For the approach by Stirn et al. [2023, Faithful ] the regularization of the variance is

class = 5

original

rot = 52

◦

rotated

10

1

10

2

10

3

10

4

regularization

20

40

Test NLL (

↓

)

β

-NLL (0.5)

Faithful

Method

Homoscedastic

52.2

±

0

.

7

Faithful

51.3

±

0

.

8

β

-NLL (0.5)

71.6

±

0

.

8

β

-NLL (1.0)

69.3

±

0

.

8

Natural Laplace (ours)

76.8

±

0

.

4

more severe due to forcing µ ( x ) and Σ( x ) to share a joint representation that receives no gradient from the variance. If, for example, Σ( x ) depends on a 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 a version of rotated MNIST with rotation angle drawn as rot ∼ Unif ( -90 , 90) . We generate the target y as y = rot +(11 c +1) ϵ, where c ∈ { 0 , 1 , . . . , 9 } is the image class as integer number and ϵ ∼ N (0 , 1) is an independent noise source. The heteroscedastic regression task of learning the distribution of y given observations (images) x involves learning a complex non-linear mapping for the variance, while learning the mean only requires learning the rotation angle.

We train a 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 β -NLL and Faithful strongly depends on 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 2-which provides a proxy on the information about the variance that has been picked up by each approach. The homoscedastic model serves as a 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. β -NLL shows a better performance but is still significantly outperformed by our proposal. The difference becomes even more evident when using a CNN architecture on rotated FashionMNIST as we demonstrate in Sec. 5.3. We also provide a minimal example illustrating the limitations of Faithful in App. C, which shows that the bottleneck size of the joint network f z has a 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, we use a homoscedastic case as ablation with y = rot + 10 ε . Further, we use a heteroscedastic noise that is based on the rotational magnitude, y = rot + √ | rot | ε . 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.

3 Naturally Parameterized Heteroscedastic Regression

We model a dataset D = { ( x n , y n ) } N n =1 with N pairs of input x n ∈ R D and scalar response y n ∈ R using the natural form of the Gaussian likelihood with unknown variance, as 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 η and the mean µ and variance σ 2 of the parametrization is

which can be understood as the signal-to-variance ratio and the negative precision (inverse variance). We model these natural parameters using a neural network f ( x ; θ ) ∈ R 2 that is parameterized by weights θ ∈ R P and acts on the inputs x ∈ R D . To satisfy the constraint that η 2 < 0 , we link f to η using a positive function g + : R ↦→ R + and have the following mapping:

We use either the exponential g + ( · ) = 1 2 exp( · ) or softplus g + ( · ) = 1 β log(1 + exp( · )) as typical for heteroscedastic regression with mean-variance parametrization. Mathematically, the heteroscedastic Gaussian log likelihood of our model is given by

which we also denote by -ℓ η ( θ ) . Assuming the data are i.i.d. , we have log p ( D| θ ) = ∑ N n =1 log p ( y n | x n , θ ) . Immer et al. [2023a] used this maximum likelihood objective for bivariate 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 et al. [2022] and Stirn et al. [2023], we also inspect the gradients of the corresponding negative log likelihood with respect to η , leading to

The gradients by themselves cannot be directly linked to mean and variance updates. We note, however, that if we relate the natural parameters to mean and variance, i.e., we compute µ ( x ; θ ) as -η 1 ( x ; θ ) 2 η 2 ( x ; θ ) and σ 2 ( x ; θ ) as -1 2 η 2 ( x ; θ ) , then ∇ η 1 ℓ η ( θ ) reduces to µ ( x ; θ ) -y , and similarly, ∇ η 2 ℓ η ( θ ) to σ 2 ( x ; θ ) -( y 2 -( µ ( x ; θ )) 2 ) , which would be desired because these are simply separate residuals for mean and variance [Seitzer et al., 2022, Stirn et al., 2023]. Empirically, we observe that this parametrization can be more stable to train and less prone to insufficient regularization as also observed previously for Gaussian process and ridge regression [Le et al., 2005, Immer et al., 2023a].

3.2 Regularization using Bayesian Inference

Due to the expressiveness of heteroscedastic regression, regularization is crucial to obtain wellgeneralizing models (cf. Figure 2). We achieve regularization in two ways: first, we regularize parameters towards low norm using classical ℓ 2 regularization, which corresponds to a Gaussian prior on the parameters. Further, we use a 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, we use a layer-wise Gaussian prior on the parameters given by p ( θ | δ ) = ∏ l N ( θ l ; 0 , δ -1 l I ) . In the case of the last layer or a simple linear model with the natural heteroscedastic likelihood, it induces a mode on η 1 ( x ) = 0 and η 2 ( x ) = -1 2 , 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., 2021a, Daxberger et al., 2021, Antorán et al., 2022]. Concurrent work by also suggests that this might be the case for heteroscedastic regression [Wong-Toi et al., 2023]. However, optimizing a regularization parameter per layer is intractable using a validation-based grid search.

To optimize layer-wise prior precisions and obtain a posterior predictive, we make use of Bayesian inference. Combining the natural Gaussian likelihood p ( D| θ ) (below Equation 5) with the prior p ( θ | δ ) , we have the joint distribution p ( D , θ | δ ) , which corresponds to a regularized objective. According to Bayes' theorem, we have the posterior p ( θ |D , δ ) ∝ p ( D , θ | δ ) . The normalization constant, also referred to as marginal likelihood , is given by p ( D| δ ) = ∫ p ( D , θ | δ )d θ and gives us the Type II maximum likelihood objective to optimize the prior precisions δ . This procedure is referred to as empirical Bayes (EB). Inferring the posterior distribution of the neural network parameters, we 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 a single model [Wilson, 2020]. Unfortunately, inference is intractable for deep neural networks.

4 Approximate Inference with a Laplace Approximation

We develop a scalable Laplace approximation for the posterior in heteroscedastic regression with deep neural networks. The Laplace approximation [MacKay, 1995] is an effective method for

approximating the marginal likelihood, posterior, and predictive in deep learning [Daxberger et al., 2021]. In comparison to other approximate inference methods, it can rely on effective training algorithms developed for deep learning and also offers a differentiable marginal likelihood estimate that enables empirical Bayes (EB). Efficient curvature approximations further make it scalable to deep learning [Ritter et al., 2018]. We extend these to the heteroscedastic regression setting.

Laplace approximates the posterior locally at a mode of the posterior with a Gaussian distribution, p ( θ |D , δ ) ≈ N ( θ ; θ ∗ , Σ ) . The mean is given by a stationary point of the posterior, θ ∗ = arg max θ log p ( D , θ | δ ) , and the covariance by the curvature at that mode, Σ -1 = ∇ 2 θ log p ( D , θ | δ ) | θ = θ ∗ . This is due to a second-order Taylor approximation of the log posterior around the mode. The mean is hence the result of neural network training, however, the covariance requires estimating and inverting a Hessian, which is typically intractable.

4.1 Linearized Laplace for Natural Heteroscedastic Regression

The linearized Laplace approximation [MacKay, 1995, Khan et al., 2019, Foong et al., 2019, Immer et al., 2021b] overcomes issues of the vanilla Laplace approximation. Linearizing the neural network about the parameters at the mode, the Hessian, which is the generalized Gauss-Newton in this case [Martens, 2020], becomes positive semidefinite and offers efficient structured approximations. In particular, we have the following Hessian approximation due to linearization

where [ J ∗ ( x )] cp = ∂η c ( x ; θ ) ∂θ p | θ = θ ∗ is the Jacobian of the neural network η ( x ; θ ) 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 δ l > 0 , the layer-wise regularization parameters.

Due to the natural parametrization of the likelihood, the Hessian approximation of the linearized Laplace approximation is guaranteed to be positive definite. Since the log prior Hessian is diagonal with entries δ l > 0 , one only has to show that the negative log likelihood Hessian, -∇ 2 η log p ( y | x , θ ∗ ) , is positive semidefinite, which is simply a property of naturally parameterized exponential families [Martens, 2020]. Note that this would not be the case for the mean-variance parametrization 3 . For efficient computation, we can decompose the log likelihood Hessian as

where the square indicates the outer product. This alleviates the need to compute this matrix and instead allows to work with outer products of Jacobian-vector products to compute the covariance in Equation 7. However, the full Hessian approximation remains quadratic in the number of neural network weights. We tackle this issue with a scalable Kronecker-factored approximation.

4.2 Scalable Kronecker-Factored Hessian Approximation

To enable the application of the Laplace approximation to heteroscedastic deep neural networks and large datasets, we use a layer-wise Kronecker-factored approximation [KFAC; Martens and Grosse, 2015, Botev et al., 2017]. To overcome the quadratic scaling in the number of parameters, KFAC makes two efficient approximations: first, it constructs a block-diagonal approximation to H from blocks H l per layer l . Second, each block H l is approximated as a Kronecker product that enables efficient storage and computation using only the individual factors. In the following, we revise KFAC for linear layers and define it for the specific case of the heteroscedastic natural Gaussian likelihood. The same derivation applies similarly to other layer types [Osawa, 2021].

Using the Hessian decomposition of the natural log likelihood in Equation 10, we derive a KFAC approximation that can efficiently be computed in a closed-form. We can write the Jacobian of a

fully connected layer that maps a D to a D ′ -dimensional representation as a Kronecker product J l ( x n ) T = a l,n ⊗ g l,n with a l,n ∈ R D × , 1 as the layer's input and g l,n ∈ R D ′ × 2 as transposed Jacobian w.r.t. the output, both for the input x n . Following Martens and Grosse [2015] and Botev et al. [2017], we then have the KFAC approximation

where λ n,k ∈ R 2 × 1 is due to to the decomposition of Λ n into outer products Equation 10 and the approximation is due to exchanging the sum and product. Conveniently, the terms g l,n λ n,k can be computed efficiently using two Jacobian-vector products and the Kronecker factors can then be extracted within the second-order framework of Osawa [2021]. To the best of our knowledge, this is the first instantiation of KFAC for heteroscedastic regression. While we derive it for the Laplace approximation, it could also be useful for optimization [Martens and Grosse, 2015].

4.3 Empirical Bayes for Automatic Regularization

To automatically regularize the heteroscedastic neural network, we use an empirical Bayes (EB) procedure that optimizes the layer-wise prior precisions, δ l , during training by maximizing the Laplace approximation to the marginal likelihood [Immer et al., 2021a]. This procedure can exhibit a Bayesian variant of Occam's razor [Rasmussen and Ghahramani, 2000] and trades off model fit and complexity. Although online training violates the stationarity assumption of the Laplace approximation, it has been observed to work well in practice [Immer et al., 2021a, 2023b, Daxberger et al., 2021, Lin et al., 2023]. We use gradient-based optimization of the log marginal likelihood,

which crucially requires differentiating the log-determinant w.r.t. δ , which is only tractable for small neural networks. For deep neural networks, it can be done efficiently using the KFAC approximation derived in Sec. 4.2 by eigendecomposition of the individual Kronecker factors [Immer et al., 2021a, 2022]. In practice, we compute the marginal likelihood approximation every few epochs to adapt the regularization, which effectively mitigates overfitting, and use it as an early-stopping criterion. The detailed empirical Bayes training algorithm is also described in Alg. 1.

4.4 Posterior Predictive for Epistemic Uncertainties

We use the linearized posterior predictive that is in line with the linearized Laplace posterior approximation and performs typically better than sampling weights directly [Immer et al., 2021b]. We define the linearized neural network as η lin ∗ ( x ; θ ) def = η ( x ; θ ∗ ) + J ∗ ( x )( θ -θ ∗ ) . Due to the Gaussianity of the Laplace approximation, N ( θ ∗ ; Σ ) , and the linearization, we can express the function-space posterior as a Gaussian on the natural parameters η lin ∗ ( x ) ∼ N ( η ( x ; θ ∗ ); J ∗ ( x ) ΣJ ∗ ( x ) T ) def = q ( η | x ) . A straightforward way to approximate the posterior predictive is then a Monte-Carlo estimate of p ( y ∗ | x ∗ , D ) ≈ ∫ p ( y ∗ | x ∗ , η ) q ( η | x ∗ ) d η by sampling multiple η lin ∗ ( x ) .

in a closed-form without sampling, similar to the probit approximation commonly used in the classification setting [Daxberger et al., 2021]. To enable a closed-form posterior predictive, we restrict ourselves to the epistemic uncertainty about the mean as proposed by Le et al. [2005] for heteroscedastic Gaussian process regression. Instead of linearizing the natural parameters, we first transform them to the mean and variance using the inverse mapping of Equation 3, i.e., we have µ ( x ; θ ) = -η 1 ( x ; θ ) 2 η 2 ( x ; θ ) and σ 2 ( x ; θ ∗ ) = -1 2 η 2 ( x ; θ ∗ ) . Next, we only linearize the mean function µ lin ∗ ( x ; θ ) = µ ( x ; θ ∗ ) + J ∗ ,µ ( x )( θ -θ ∗ ) , where J ∗ ,µ ( x ) is the Jacobian of the mean, and have

Alternatively, we propose to use an approximation to the posterior predictive that can be computed

where the epistemic and aleatoric uncertainty about the mean are clearly split. An example of this posterior predictive approximation is shown in Figure 1.

4.5 Laplace Approximation for Mean-Variance parametrization

Using the natural parameter mapping, it is possible to apply above Laplace approximations to the mean-variance parametrization and profit from the empirical Bayes and posterior predictive procedures. However, because the negative log likelihood Hessian w.r.t. the mean and variance parameters can be indefinite, this does not work naively. By mapping µ, σ 2 to the natural parameters using Equation 3, all above derivations apply and correspond to a separate Gauss-Newton approximation of the log likelihood with Jacobians of the mapping and the natural log likelihood Hessian (see App. A).

4.6 Limitations

As common for heteroscedastic regression, we make the assumption that the conditional distribution of the target given the observations follows a Gaussian distribution. A practitioner should keep in mind that the full Laplace approximation has high computational complexity, and should resort to the KFAC approximation, which we also use for more complex settings in experiments. Further, the true posterior of neural networks is in most cases multimodal while our Laplace approximation only covers a single mode. However in practical settings, the Laplace approximation provides strong results [Daxberger et al., 2021] and can even be ensembled [Eschenhagen et al., 2021].

5 Experiments

We evaluate the effectiveness of the natural parameterization compared to the mean-variance (naive) one, and empirical Bayes (EB) to optimizing a single regularization parameter using a grid search on the validation set (GS), and the MAP prediction vs a Bayesian posterior predictive (PP) in comparison to state-of-the-art baselines on three experimental settings: the UCI regression benchmark [HernandezLobato and Adams, 2015], which is also well-established for heteroscedastic regression [Seitzer et al., 2022, Stirn et al., 2023], the recently introduced CRISPR-Cas13 gene expression datasets [Stirn et al., 2023], and our proposed heteroscedastic image-regression dataset (cf. Problem 2.1) in three noise variants. For UCI regression, we use the full Laplace approximation and we resort to the KFAC approximation for the remaining tasks for computational efficiency. Further, we use the proposed closed-form posterior predictive for our Laplace approximations. If applicable, as for our methods, we test the effect of having a Bayesian posterior predictive (PP) in comparison to the point estimate.

Baselines As control to assess the benefit of heteroscedastic aleatoric uncertainty, we include a homoscedastic model with EB and Laplace posterior predictive. In addition, we include the meanvariance parameterization of the negative log likelihood loss for heteroscedastic regression ( Naive NLL ). As competitive baselines, we include the recently proposed β -NLL [Seitzer et al., 2022] and Faithful [Stirn et al., 2023] heteroscedastic losses. Finally we include well-established approaches for heteroscedastic regression with Bayesian neural networks, namely mean-field variational inference ( VI ) [Graves, 2011] and Monte-Carlo Dropout [Gal and Ghahramani, 2016] ( MC-Dropout ). Note that for our VI baseline we employ Flipout [Wen et al., 2018] to improve gradient estimation.

Details on model architectures, hyperparameter tuning, and additional results are in App. D.

5.1 UCI Regression

In Table 1 we report results on the UCI regression datasets [Hernandez-Lobato and Adams, 2015] for the compared models in terms of test log likelihood. We underline the best performance for each dataset, and bold results which are not statistically distinguishable from best performance. That is, if the mean performance of a model falls within the region of the mean plus minus two standard errors of the best performing model. The two rightmost columns report the total number of wins and ties for each model: wins are the number of datasets in which the given model achieves the best performance, while ties are the number of datasets in which the model achieves results which are not statistically distinguishable from the best performing model. The results validate the effectiveness of EB regularization on different heteroscedastic loss parameterizations. Notably, with the Bayesian posterior predictive (PP) our proposed methods considerably improve in performance in comparison to using a point estimate, and significantly outperform existing state-of-the-art approaches. Finally the results confirm that, particularly when the point prective is used, using the Natural parameterization of the heteroscedastic NLL improves training stability.

LL (

↑

)

Objective

Regular-

ization

Posterior

Predictive

boston

concrete

energy

kin8nm

naval

plant

wine

yacht

Wins Ties

Homoscedastic

EB

✓

-2.51 (0.04) -3.03 (0.03) -0.71 (0.04)

1.29 (0.01)

5.58 (0.45)

-2.83 (0.02)

-0.94 (0.02)

-1.05 (0.14)

0

2

Naive NLL

GS

✗

-2.88 (0.2)

-3.96 (0.32) -1.06 (0.16)

1.31 (0.01)

3.88 (0.27)

-2.85 (0.04)

-0.95 (0.02)

-0.60 (0.18)

0

1

β

-NLL (0.5)

GS

✗

-2.61 (0.06) -3.45 (0.24) -0.96 (0.13)

1.30 (0.01)

2.60 (2.70)

-2.85 (0.04)

-0.94 (0.02)

-1.39 (0.49)

0

1

β

-NLL (1.0)

GS

✗

-2.71 (0.06) -3.13 (0.04) -0.88 (0.07)

1.29 (0.01)

6.47 (0.12)

-2.86 (0.03)

-0.94 (0.02)

-1.09 (0.20)

0

1

Faithful

GS

✗

-2.68 (0.07) -3.12 (0.06) -1.10 (0.11)

1.29 (0.01)

6.41 (0.14)

-2.80 (0.02)

-0.95 (0.02)

-1.32 (0.31)

0

2

MC-Dropout

GS

✗

-3.01 (0.24) -3.05 (0.03) -1.61 (0.03)

1.17 (0.01)

5.64 (0.03)

-2.82 (0.01)

-1.12 (0.07) -1.01 (0.08)

0

1

VI

GS

✗

-2.62 (0.08) -3.10 (0.04) -1.55 (0.03)

1.31 (0.01)

5.70 (0.24)

-2.82 (0.02)

-1.07 (0.07) -1.16 (0.05)

0

1

Naive NLL

EB

✓

-2.47 (0.04)

-2.88 (0.03) -0.46 (0.03)

1.37 (0.01)

6.38(0.16)

-2.79 (0.02)

-0.93 (0.02) -0.05 (0.15)

5

6

✗

-6.03 (1.82) -3.62 (0.26) -0.97 (0.15)

1.36 (0.01)

-10.9 (8.01)

-2.83 (0.04)

-2.26 (0.51) -1.27 (0.51)

0

2

Natural NLL

GS

✓

-2.46 (0.04)

-2.92 (0.04)

-0.74 (0.06)

1.32 (0.01)

6.66 (0.01)

-2.76 (0.02)

-0.94 (0.02) -0.29 (0.07)

0

4

✗

-2.57 (0.08) -2.95 (0.05) -0.69 (0.08)

1.31 (0.01)

6.69 (0.01)

-2.76 (0.02)

-0.95 (0.02)

-0.87 (0.28)

1

3

EB

✓

-2.36 (0.03) -2.93 (0.02)

-0.71 (0.03)

1.36 (0.01)

6.66 (0.01)

-2.76 (0.02)

-0.94 (0.02)

-0.51 (0.04)

2

5

✗

-2.51 (0.09) -2.99 (0.04) -0.65 (0.05)

1.35 (0.01)

6.68 (0.01)

-2.77 (0.02)

-1.00 (0.03) -0.60 (0.08)

0

3

-

1

.

0

-

0

.

5

LL

(

)

Natural NLL

|

EB

Natural NLL

|

EB

|

PP

Natural NLL

|

GS

Natural NLL

|

GS

|

PP

Naive NLL

|

EB

|

PP

VI

|

PP

Faithful

β

-NLL (1.0)

β

-NLL (0.5)

Naive NLL

Homoscedastic

|

GS

|

PP

flow-cytometry-HEK293

-

0

.

6

-

0

.

4

-

0

.

2

LL

(

)

survival-screen-A375

-

1

.

00

-

0

.

75

-

0

.

50

-

0

.

25

0

.

00

LL

(

)

survival-screen-HEK293

↑

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 a set of datasets to model the efficacy of the CRISPR-Cas13 system for gene knockdown. In particular, Cas13 is a system that targets specific RNA transcripts to temporarily decrease gene expression. Decreasing gene expression levels serves multiple purposes in molecular biology, including the understanding of the function of specific genes and for therapeutic aims. In particular, the authors propose three datasets, and for each dataset multiple efficacy scores are reported, resulting from replicated experiments. The response variable is a scalar value measuring knockdown efficacy from a one-hot encoded target sequence. Note that here we report results only on the de-noised averaged data across replicates (see Stirn et al. [2023]), while we leave the results on the raw replicates for App. D. Results on these datasets validate the advantage of using the posterior predictive for our methods, as well as the effectiveness of using the natural parameterization, which results in improved training stability.

5.3 Regression with Image Data

To assess the performance of different approaches in a typical deep learning setting, we create versions of the MNIST and FashionMNIST datasets where the task is to regress the rotational angle under label-dependent heteroscedastic noise (Problem 2.1). Such a setting is interesting to study since most (heteroscedastic) regression datasets are tabular and we intend to benchmark on a more common deep learning task where tabular regression models would not be applicable.

↑

↑

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 +(11 c +1) ϵ, where c ∈ { 0 , 1 , . . . , 9 } is the image class of x as integer number and ϵ ∼ N (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 correctly 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. 4 We train a 3-layer 500-neuron wide MLP with ReLU activation for the MNIST-based dataset and a CNN (3 conv. and 2 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 on FashionMNIST. β -NLL, MC-Dropout, and VI all give relatively good results but fall slightly short of the proposed methods.

6 Conclusions

We tackled the problem of heteroscedastic regression using neural networks with a focus on regularization. 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 is scalable despite the Bayesian procedures and benchmark our approach on UCI regression as well as CRISPR-Cas13 datasets achieving state-of-the-art performance. For future work, it would be interesting to relax the assumption about Gaussianity and apply our approach to other real-world datasets.

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 Hernández-Lobato. Adapting the linearised laplace model evidence for modern deep learning. In International Conference on Machine Learning , pages 796-821. PMLR, 2022.

Aleksandar Botev, Hippolyt Ritter, and David Barber. Practical gauss-newton optimisation for deep learning. In International Conference on 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.com/google/jax .

Gavin C Cawley, Nicola LC Talbot, Robert J Foxall, Stephen R Dorling, and Danilo P Mandic. Heteroscedastic kernel ridge regression. Neurocomputing , 57:105-124, 2004.

Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace redux-effortless 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 approximations 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. 'inbetween'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 in Neural Information Processing Systems , 2011.

Isabelle Guyon, Alexander Statnikov, and Berna Bakir Batu. Cause effect Pairs in machine learning . Springer, 2019.

Andrew C Harvey. Estimating regression models with multiplicative heteroscedasticity. 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:1112.5745 , 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 Matthias Bauer. Improving predictions of bayesian neural nets via local linearization. In International Conference on Artificial Intelligence and Statistics , pages 703-711. PMLR, 2021b.

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:12449-12463, 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 on Machine Learning , pages 14316-14332. 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 14333-14352. PMLR, 2023b.

Alex Kendall and Yarin Gal. What uncertainties do we need in 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: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Andreas Kirsch. Black-box batch active learning for regression. arXiv preprint arXiv:2302.08981 , 2023.

Ranganath Krishnan, Pi Esposito, and Mahesh Subedar. Bayesian-torch: Bayesian neural network layers for uncertainty estimation. https://github.com/IntelLabs/bayesian-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. Heteroscedastic gaussian process regression. In International Conference on Machine Learning , pages 489-496. PMLR, 2005.

Jihao Andreas Lin, Javier Antorán, and José Miguel Hernández-Lobato. Online laplace model selection revisited. arXiv preprint arXiv: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-5851, 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 S Weigend. Estimating the mean and variance of the target probability distribution. In IEEE international conference on neural networks , volume 1, pages 55-60. IEEE, 1994.

Kazuki Osawa. Automatic Second-Order Differentiation Library (ASDL), 2021. URL http: //github.com/kazukiosawa/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 DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. In NeurIPS Workshop Autodiff , 2017.

Carl Rasmussen and Zoubin Ghahramani. Occam's razor. Advances in neural information processing systems , 13, 2000.

Hippolyt Ritter, Aleksandar Botev, and David Barber. A scalable laplace approximation for neural networks. In International Conference on Learning Representations , volume 6, 2018.

Maximilian Seitzer, Arash Tavakoli, Dimitrije Antic, and Georg Martius. On the pitfalls of heteroscedastic uncertainty estimation with probabilistic neural networks. In International Conference on Learning Representations , 2022.

Nicki Skafte, Martin Jø rgensen, and Sø 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 arXiv:2006.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 Tomczak, 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 pseudoindependent weight perturbations on mini-batches. In International Conference on Learning Representations , 2018.

Andrew Gordon Wilson. The case for bayesian deep learning. arXiv preprint arXiv:2001.10995 , 2020.

Eliot Wong-Toi, Alex Boyd, Vincent Fortuin, and Stephan Mandt. Understanding pathologies of deep heteroskedastic regression. arXiv preprint arXiv:2306.16717 , 2023.

Jeffrey M Wooldridge. Introductory econometrics: A modern approach . Cengage learning, 2015.

Anqi Wu, Sebastian Nowozin, Edward Meeds, Richard E. Turner, Jose Miguel Hernandez-Lobato, 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:14129-14142, 2020.

Ming Yuan and Grace Wahba. Doubly penalized likelihood estimator in heteroscedastic regression. Statistics & probability letters , 69(1):11-20, 2004.

Table of Contents

Appendix

A Laplace with Mean and Variance Parametrization

In Sec. 4, we introduced the Laplace approximation for the naturally parameterized Gaussian likelihood. The natural parametrization allows an efficient generalized Gauss-Newton approximation to the Hessian due to the properties of the second derivative [Martens, 2020], i.e., that is positive definite w.r.t. the natural parameters. The positive definiteness is necessary so that the generalized Gauss-Newton, which pre and post-multiplies the Jacobians to this Hessian, can only be positive semidefinite. This is a 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 approximation to 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-variance parametrization of a neural network.

Since the Hessian of the Gaussian likelihood under the mean-variance parametrization is not positive definite, we cannot directly use this parametrization for the Gauss-Newton approximation. To verify this statement, define m ( x ; θ ) = [ µ ( x ; θ ) , σ 2 ( x ; θ ) ] ⊤ and compute the Hessian of the Gaussian negative log likelihood w.r.t. m

By Silvester's criterion, a Hermitian matrix is positive definite if and only if the determinants of all sub-matrices are positive. For the Hessian in Equation 14, this criterion is not fulfilled as

A simple way to overcome this issue is to map the mean-variance parameters m to the natural parameters η using Equation 3 and apply the approximations as derived in Sec. 4, i.e.,

and then apply the Gauss-Newton approximation as in Sec. 4 to obtain

where the Jacobians J η | m are of η 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. 5 The implementation is very simple: we simply use the output of the neural network m , apply the transformation to the natural parameters η , and then compute the quantities as 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 a 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 θ can be written as

x

dtFF

μz

比FF ∑z

:=0

u

²00 2e

0z00

Z

0

dtFFE

FF OE

re

200

n

FF

As Seitzer et al. [2022] point out, the gradient of ℓ µ,σ ( θ ) with respect to µ ( x ; θ ) is scaled by the estimated precision, i.e,

which they argue can lead to a poor mean estimate. In particular, they state that if the variance is well-calibrated, i.e., σ 2 ( x ; θ ) ≈ ( µ ( x ; θ ) -y ) 2 , the gradient with respect to the mean reduces to

which rewards points with an already accurate mean 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 if the irreducible aleatoric uncertainty is low resp. high in a certain region, which does not necessarily imply a poor mean fit.

β -NLL To improve the mean fit, Seitzer et al. [2022] introduce the β -NLL loss

which introduces a stop-gradient operation denoted as ⌊·⌋ to trade-off the influence of the variance estimate on the gradients. The corresponding gradients for the β -NLL loss are

Thus, for β = 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 β = 0 it is the standard heteroscedastic gradient and for 0 < β < 1 the approach interpolates between both extremes. A special case appears at β = 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 µ ( x ) , σ ( x ) should be faithful to a homoscedastic model that learns µ 0 ( x ) . As faithfulness they define that (cf. Stirn et al. [2023, Def. 2])

To achieve this property, they propose to decouple the estimation into three networks, as shown in Figure 4: a shared representation learner f z computes a representation z from x , which is passes into two individual networks f µ and f Σ , which receive z as input and output the mean respectively the covariance matrix. Based on this network architecture, they minimize the following loss

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 f Σ from propergating to f z . Second, they propose to multiply the gradient of f µ ◦ f z with Σ . 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 β -NLL. In particular, for the sub-networks f z and f µ the gradients are multiplied with Σ which is the multivariate analogue of the effect of the β -NLL 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 ℓ η ( θ ) with respect to η 1 and η 2 as

Although these gradients do not have a direct implication on updates with respect to mean and variance, note that if we compute the mean estimate µ ( x ; θ ) from the estimated natural parameters, it is equal to -η 1 ( x ; θ ) 2 η 2 ( x ; θ ) . Similarly, note that the variance σ 2 ( x ; θ ) can be computed from the natural parametization as -1 2 η 2 ( x ; θ ) . If we replace the corresponding quantities in Equation 24, we get that

which is 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 Mean-Variance Sources

To illustrate the limitations of 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 a toy setting in which mean and variance are generated from independent random variables:

Problem C.1 Let X ∈ R 2 be a centered multivariate Gaussian distributed with Σ = I . We construct the target variable Y as

where ϵ ∼ N (0 , 1) is an independent noise variable.

From the mechanism described in Problem C.1, we construct 20 datasets consisting of 1000 i.i.d. samples, where each dataset is generated using a different random seed. We split the data into trainvalidation/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 heteroscedastic objective using only a point predictive,both trained with the full Laplace approximation, where we use an MLP with 1 hidden-layer consisting of 100 hidden units. For Faithful , we use the same MLP with 1 hidden layer for the joint encoder f z and instantiate f µ and f Σ as linear projections from dimension | z | to 1. We ablate the size of the hidden dimension | z | from 1 to 128 neurons to create a bottleneck. In theory, a size of 2 should be sufficient to 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, a large bottleneck size does not guarantee that the network learns

Homoscedastic | EB

0.4

Faithful

Natural NLL | EB

Test Log-likelihood

0.6

0.8

1.0

1.2

1

2

4

8

16

32

64

128

Bottleneck size

(a) Average test log likelihood across seeds

0.4

-0.6

0.8

1.0

Homoscedastic | EB

Faithful

1.2

Natural NLL | EB

1

2

4

8

16

32

64

128

Bottleneck size

(b) Best test log likelihood across seeds

the relevant information about the variance. We rather conjecture that the amount of information that is propagated is dependent on the initialization of the network, since no gradient actively encourages the network f z to pick up information about X 2 , which influences the variance. Since for a larger bottleneck size the network can encode much more information than just X 1 , we observe that it encodes information about X 2 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 a homoscedastic to a heteroscedastic likelihood and the point-wise prediction to the Laplace-approximated posterior predictive. The dataset is created similar to Skafte et al. [2019] with the following generative model:

For training, we center the input data by subtracting the mean 7 . 5 from x and standardize the observations to zero mean and unit variance as common for regression problems. We create a dataset with 1000 samples.

We model the data with a neural network with 100 hidden units on a single hidden layer using tanh activation. Both homoscedastic and heteroscedastic models are trained using the marginallikelihood 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 the 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 1 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-Dropout baselines, with Adam optimizer using a

batch size of 256 for 5000 epochs and an initial learning rate of 10 -2 that is decayed to 10 -5 using a cosine schedule. As we observe that gives overall better performance, MC-Dropout and VI on UCI datasets are trained for 1000 epochs, with a 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 . 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 5 . 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, 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 a 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. E.1. For the UCI regression datasets, we 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 a 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].

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 3 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-pooling 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 a learning rate of 10 -2 that is decayed to 10 -4 using a cosine schedule when training on the FashionMNIST dataset, while we use a non-decayed learning rate of 10 -2 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 10 4 . 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, 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 a 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 Kronecker-factored Laplace approximation so that the training and inference time between Laplace and baselines are almost identical. The training was done 5 times (different seeds) per model-dataset pair to estimate mean and standard error and were run on a computing cluster with V100 and A100 NVIDIA GPUs.

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

-

5

0

5

x

-

2

.

5

0

.

0

2

.

5

Homoscedastic Laplace

-

5

0

5

x

Heteroscedastic Laplace

Figure 6:

Heteroscedastic proposed method on

homoscedastic ground truth data still performs well

and falls back to homoscedastic predictive.

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.

-

5

0

5

x

-

2

.

5

0

.

0

2

.

5

y

MC Dropout

-

5

0

5

x

MFVI

-

5

0

5

x

Homoscedastic Laplace

-

5

0

5

x

Heteroscedastic Laplace

ground truth

mean

epistemic

aleatoric

D.2.2 UCI Regression Datasets

In this section, we report additional results to the ones 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 1, while Table 4 reports an ablation where for our proposed natural parameterized NLL heteroscedastic loss model EB regularization, comparing results obtained with the full Hessian Laplace approximation ( full ) and results obtained with the Kronecker-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 to comparable results in most settings, justifying 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-Carlo estimate of the NLL [Tomczak et al., 2018] ( lse ). Finally Table 6 shows that the results from our work for models trained with EB regularization favourably compare to results from the work of Wu et 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 in the other experimental settings.

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 we 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 2 but with a homoscedastic generative process. In 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 is 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 a grid search but the posterior predictive requires more runtime per sample.

-

2

.

0

-

1

.

5

-

1

.

0

-

0

.

5

LL

(

↑

)

Natural NLL

|

EB

Natural NLL

|

EB

|

PP

Natural NLL

|

GS

Natural NLL

|

GS

|

PP

Naive NLL

|

EB

|

PP

VI

|

PP

Faithful

β

-NLL (1.0)

β

-NLL (0.5)

Naive NLL

Homoscedastic

|

GS

|

PP

flow-cytometry-HEK293

-

0

.

6

-

0

.

5

-

0

.

4

-

0

.

3

-

0

.

2

LL

(

↑

)

survival-screen-HEK293

(a) Replicate 1

-

1

.

5

-

1

.

0

-

0

.

5

LL

(

↑

)

Natural NLL

|

EB

Natural NLL

|

EB

|

PP

Natural NLL

|

GS

Natural NLL

|

GS

|

PP

Naive NLL

|

EB

|

PP

VI

|

PP

Faithful

β

-NLL (1.0)

β

-NLL (0.5)

Naive NLL

Homoscedastic

|

GS

|

PP

flow-cytometry-HEK293

-

0

.

8

-

0

.

7

-

0

.

6

-

0

.

5

-

0

.

4

LL

(

↑

)

survival-screen-A375

-

0

.

8

-

0

.

7

-

0

.

6

-

0

.

5

-

0

.

4

-

0

.

3

LL

(

↑

)

survival-screen-HEK293

(b) Replicate 2

-

0

.

8

-

0

.

6

-

0

.

4

LL

(

↑

)

Natural NLL

|

EB

Natural NLL

|

EB

|

PP

Natural NLL

|

GS

Natural NLL

|

GS

|

PP

Naive NLL

|

EB

|

PP

VI

|

PP

Faithful

β

-NLL (1.0)

β

-NLL (0.5)

Naive NLL

Homoscedastic

|

GS

|

PP

flow-cytometry-HEK293

-

0

.

8

-

0

.

7

-

0

.

6

-

0

.

5

-

0

.

4

LL

(

↑

)

survival-screen-A375

-

1

.

0

-

0

.

8

-

0

.

6

-

0

.

4

LL

(

↑

)

survival-screen-HEK293

(c) Replicate 3

E Implementation

In the following, we detail the interleaved Laplace marginal likelihood optimization [Immer et al., 2021a] but adapted to the heteroscedastic setting in App. E.1. Further, we provide an efficient and stable implementation of the natural parametrization of the heteroscedastic Gaussian likelihood in App. E.2. We implement our method in the laplace-torch package [Daxberger et 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 be used with any supported network by these packages and different scalable Laplace approximation variants.

E.1 Pseudo-Code for Our Method

Alg. 1 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, or 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 to compute and does not require sampling due to the simplification proposed in Sec. 4.4. Experimentally, we find that the Bayesian posterior predictive improves over the point-wise predictive in almost all cases.

Algorithm 1 Optimization of Heteroscedastic Regression Models

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 a Gaussian. Below, we provide an implementation for it using pytorch:

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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\n\nWEAKNESSES:\n- 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\n\nQUESTIONS:\n- 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\nLIMITATIONS:\nLimitations 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nIn 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\n(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?\n\nLIMITATIONS:\nThe limitations sections exist.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\nI cannot spot any weaknesses in this paper.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nThe limitations of this work are discussed properly.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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\n\nWEAKNESSES:\n- 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\n\nQUESTIONS:\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\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\nLIMITATIONS:\nThere was not much discussion of this, but that is appropriate for the nature of this work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}