SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Alexander Immer*,1,2 Emanuele Palumbo*,1,3 Alexander Marx†1,3 Julia E. Vogt†1
$^{1}$ Department of Computer Science, ETH Zurich, Switzerland  
 $^{2}$ Max Planck Institute for Intelligent Systems, Tübingen, Germany  
 $^{3}$ AI Center, ETH Zurich, Switzerland
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.
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  $\mathbf{x} \in \mathbb{R}^D$  and a scalar response  $\mathbf{y} \in \mathbb{R}$ , and model  $\mathbf{y} \mid \mathbf{x} = \mathbf{x}$  as a conditional Gaussian distribution with mean  $\mu(\mathbf{x})$  and standard deviation  $\sigma(\mathbf{x})$  being dependent on input  $\mathbf{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].
*Equal contribution. † Shared last author. Correspondence to: alexander.immer@inf.ethz.ch alexander.marx@inf.ethz.ch. Code at https://github.com/aleximmer/heteroscedastic-nn.
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 to learn  $\mu (\mathbf{x})$  and  $\sigma (\mathbf{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  $\mu (\mathbf{x})$  and  $\sigma (\mathbf{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  $\beta$ -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. 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}$ DVI [Wu et 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 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),
$$
\ell_ {\mu , \sigma} (\boldsymbol {\theta}) = \frac {1}{2} \log \sigma^ {2} (\mathbf {x}; \boldsymbol {\theta}) + \frac {(y - \mu (\mathbf {x} ; \boldsymbol {\theta})) ^ {2}}{2 \sigma^ {2} (\mathbf {x} ; \boldsymbol {\theta})} + c o n s t, \tag {1}
$$
where  $\theta$  denote the neural network parameters used to estimate mean and variance, Seitzer et al. [2022] trace the problem to the fact that  $\nabla_{\mu}\ell_{\mu ,\sigma}(\pmb {\theta}) = \frac{\mu(\mathbf{x};\pmb{\theta}) - y}{\sigma^2(\mathbf{x};\pmb{\theta})}$  heavily depends on the learned variance. To overcome this effect, Seitzer et al. [2022] introduce the  $\beta$  NLL loss  $\ell_{\beta}(\pmb {\theta})$  , which is equal to  $\lfloor \sigma^{2\beta}(\mathbf{x};\pmb {\theta})\rfloor \cdot \ell_{\mu ,\sigma}(\pmb {\theta})$  , where  $\lfloor \cdot \rfloor$  denotes a stop-gradient operation, and  $\beta$  is a hyperparameter controlling the dependency of gradients on the predictive variance. As a result, the gradients for  $\beta$  NLL are equal to
$$
\nabla_ {\mu} \ell_ {\beta} (\boldsymbol {\theta}) = \frac {\mu (\mathbf {x} ; \boldsymbol {\theta}) - y}{\sigma^ {2 - 2 \beta} (\mathbf {x} ; \boldsymbol {\theta})}, \quad \nabla_ {\sigma^ {2}} \ell_ {\beta} (\boldsymbol {\theta}) = \frac {\sigma^ {2} (\mathbf {x} ; \boldsymbol {\theta}) - (y - \mu (\mathbf {x} ; \boldsymbol {\theta})) ^ {2}}{2 \sigma^ {4 - 2 \beta} (\mathbf {x} ; \boldsymbol {\theta})}. \tag {2}
$$
With  $\beta = 0$ ,  $\ell_{\beta}(\pmb{\theta})$  is equivalent to  $\ell_{\mu, \sigma}(\pmb{\theta})$ , whereas for  $\beta = 1$  the gradient with respect to the mean is proportional to the gradient for homoscedastic regression. When setting  $0 < \beta \leq 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_{\mathbf{z}}$  computes a representation  $\mathbf{z}$  from  $\mathbf{x}$ , which is passed into two individual networks  $f_{\mu}$  and  $f_{\Sigma}$ , which receive  $\mathbf{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_{\mathbf{z}}$  and  $f_{\mu}$  as setting  $\beta = 1$  in  $\beta$ -NLL, and the second one stops any gradient from the variance network  $f_{\Sigma}$  from propagating to  $f_{\mathbf{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  $\beta$ -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  $\sigma^2 \leq 1$  and hence the gradient with respect to the variance is amplified stronger compared to the mean. Especially in near-deterministic processes ( $\sigma^2 \to 0$ ) this might be problematic. Further, it is not clear which value of  $\beta$  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
Figure 2: Left: Example heteroscedastic image regression data point with target  $y \sim \mathrm{rot} + 56\epsilon$ . Middle: Test log likelihood for different values of prior precision for  $\beta$ -NLL and Faithful. Right: Downstream accuracy using a linear model on the last layer representation after learning the heteroscedastic image regression. Compromised variance fit leads to worse downstream accuracy.
Method	Accuracy
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  $\mu (\mathbf{x})$  and  $\Sigma (\mathbf{x})$  to share a joint representation that receives no gradient from the variance. If, for example,  $\Sigma (\mathbf{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  $\sim$  Unif(-90,90). We generate the target  $y$  as  $y = \text{rot} + (11c + 1)\epsilon$ , where  $c \in \{0,1,\dots,9\}$  is the image class as integer number and  $\epsilon \sim \mathcal{N}(0,1)$  is an independent noise source. The heteroscedastic regression task of learning the distribution of  $y$  given observations (images)  $\mathbf{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  $\beta$ -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.  $\beta$ -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_{\mathbf{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 = \mathrm{rot} + 10\varepsilon$ . Further, we use a heteroscedastic noise that is based on the rotational magnitude,  $y = \mathrm{rot} + \sqrt{|\mathrm{rot}|}\varepsilon$ . 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  $\mathcal{D} = \{(\mathbf{x}_n,y_n)\}_{n = 1}^N$  with  $N$  pairs of input  $\mathbf{x}_n\in \mathbb{R}^D$  and scalar response  $y_{n}\in \mathbb{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  $\pmb{\eta}$  and the mean  $\mu$  and variance  $\sigma^2$  of the parametrization is
$$
\eta_ {1} = \frac {\mu}{\sigma^ {2}} \text {a n d} \eta_ {2} = - \frac {1}{2 \sigma^ {2}} <   0, \tag {3}
$$
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  $\mathbf{f}(\mathbf{x};\boldsymbol {\theta})\in \mathbb{R}^2$  that is parameterized by weights  $\pmb {\theta}\in \mathbb{R}^{P}$  and acts on the inputs  $\mathbf{x}\in \mathbb{R}^{D}$ . To satisfy the constraint that  $\eta_{2} < 0$ , we link  $\mathbf{f}$  to  $\pmb{\eta}$  using a positive function  $g_{+}:\mathbb{R}\mapsto \mathbb{R}_{+}$  and have the following mapping:
$$
\eta_ {1} (\mathbf {x}; \boldsymbol {\theta}) = f _ {1} (\mathbf {x}; \boldsymbol {\theta}) \text {a n d} \eta_ {2} (\mathbf {x}; \boldsymbol {\theta}) = - g _ {+} \left(f _ {2} (\mathbf {x}; \boldsymbol {\theta})\right). \tag {4}
$$
We use either the exponential  $g_{+}(\cdot) = \frac{1}{2}\exp (\cdot)$  or softplus  $g_{+}(\cdot) = \frac{1}{\beta}\log (1 + \exp (\cdot))$  as typical for heteroscedastic regression with mean-variance parametrization. Mathematically, the heteroscedastic Gaussian log likelihood of our model is given by
$$
\log p (y | \mathbf {x}, \boldsymbol {\theta}) = \left[ \begin{array}{l} \eta_ {1} (\mathbf {x}; \boldsymbol {\theta}) \\ \eta_ {2} (\mathbf {x}; \boldsymbol {\theta}) \end{array} \right] ^ {\top} \left[ \begin{array}{l} y \\ y ^ {2} \end{array} \right] + \frac {\eta_ {1} (\mathbf {x} ; \boldsymbol {\theta}) ^ {2}}{4 \eta_ {2} (\mathbf {x} ; \boldsymbol {\theta})} + \frac {1}{2} \log (- 2 \eta_ {2} (\mathbf {x}; \boldsymbol {\theta})) + c o n s t, \tag {5}
$$
which we also denote by  $-\ell_{\pmb{\eta}}(\pmb{\theta})$ . Assuming the data are i.i.d., we have  $\log p(\mathcal{D}|\pmb{\theta}) = \sum_{n=1}^{N} \log p(y_n | \mathbf{x}_n, \pmb{\theta})$ . 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  $\eta$ , leading to
$$
\nabla_ {\eta_ {1}} \ell_ {\boldsymbol {\eta}} (\boldsymbol {\theta}) = - \frac {\eta_ {1} (\mathbf {x} ; \boldsymbol {\theta})}{2 \eta_ {2} (\mathbf {x} ; \boldsymbol {\theta})} - y, \quad \nabla_ {\eta_ {2}} \ell_ {\boldsymbol {\eta}} (\boldsymbol {\theta}) = \frac {\left(\eta_ {1} (\mathbf {x} ; \boldsymbol {\theta})\right) ^ {2}}{4 \left(\eta_ {2} (\mathbf {x} ; \boldsymbol {\theta})\right) ^ {2}} - \frac {1}{2 \eta_ {2} (\mathbf {x} ; \boldsymbol {\theta})} - y ^ {2}. \tag {6}
$$
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  $\mu (\mathbf{x};\pmb {\theta})$  as  $-\frac{\eta_1(\mathbf{x};\pmb{\theta})}{2\eta_2(\mathbf{x};\pmb{\theta})}$  and  $\sigma^2 (\mathbf{x};\pmb {\theta})$  as  $-\frac{1}{2\eta_2(\mathbf{x};\pmb{\theta})}$ , then  $\nabla_{\eta_1}\ell_\eta (\pmb {\theta})$  reduces to  $\mu (\mathbf{x};\pmb {\theta}) - y$ , and similarly,  $\nabla_{\eta_2}\ell_\eta (\pmb {\theta})$  to  $\sigma^2 (\mathbf{x};\pmb {\theta}) - (y^2 -(\mu (\mathbf{x};\pmb {\theta}))^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  $\ell^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(\boldsymbol{\theta}|\boldsymbol{\delta}) = \prod_{l}\mathcal{N}(\boldsymbol{\theta}_{l};\mathbf{0},\delta_{l}^{-1}\mathbf{I})$ . In the case of the last layer or a simple linear model with the natural heteroscedastic likelihood, it induces a mode on  $\eta_1(\mathbf{x}) = 0$  and  $\eta_2(\mathbf{x}) = -\frac{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(\mathcal{D}|\boldsymbol{\theta})$  (below Equation 5) with the prior  $p(\boldsymbol{\theta}|\boldsymbol{\delta})$ , we have the joint distribution  $p(\mathcal{D},\boldsymbol{\theta}|\boldsymbol{\delta})$ , which corresponds to a regularized objective. According to Bayes' theorem, we have the posterior  $p(\boldsymbol{\theta}|\mathcal{D},\boldsymbol{\delta}) \propto p(\mathcal{D},\boldsymbol{\theta}|\boldsymbol{\delta})$ . The normalization constant, also referred to as marginal likelihood, is given by  $p(\mathcal{D}|\boldsymbol{\delta}) = \int p(\mathcal{D},\boldsymbol{\theta}|\boldsymbol{\delta})\mathrm{d}\boldsymbol{\theta}$  and gives us the Type II maximum likelihood objective to optimize the prior precisions  $\boldsymbol{\delta}$ . 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_{*}|\mathbf{x}_{*},\mathcal{D})$  for a new data point  $\mathbf{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