SlowGuess/NeurIPS_2023_OCR huggingface dataset:
To create the modified datasets, we follow the procedure as outlined in Problem 2.1. An observation  $\mathbf{x}$  is generated by rotating an MNIST (resp. FashionMNIST) image with rotation angle drawn as  $\mathrm{rot} \sim \mathrm{Unif}(-90, 90)$ , and the corresponding target  $y$  is generated as  $y = \mathrm{rot} + (11c + 1)\epsilon$ , where  $c \in \{0, 1, \dots, 9\}$  is the image class of  $\mathbf{x}$  as integer number and  $\epsilon \sim \mathcal{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.  $\beta$ -NLL, MC-Dropout, and VI all give relatively good results but fall slightly short of the proposed methods.
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.
<sup>4</sup>Note, however, that the original images are not completely free of rotation which creates an additional, potentially heteroscedastic, noise source that we cannot quantify.
AI gratefully acknowledges funding by the Max Planck ETH Center for Learning Systems (CLS). EP and AM were supported by a fellowship from the ETH AI Center.
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 Ratsch, 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.
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 pseudo-independent 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.
A Laplace with Mean and Variance Parametrization 15
B Comparison of Gradient Updates 15
C Orthogonal Mean-Variance Sources 17
D Details on Experiments and Additional Results 18
D.1 Experimental Details 18
D.1.1 Skafte Illustrative Example 18
D.1.2 UCI Regression and CRISPR-Cas13 Knockdown Efficacy Experiments 18
D.1.3 Image Regression 19
D.2 Additional Experimental Results 20
D.2.1 Skafte Illustrative Example 20
D.2.2 UCI Regression Datasets 20
D.2.3 CRISPR Gene Knockdown Efficacy Datasets 21
D.2.4 Image Regression 21
E Implementation 22
E.1 Pseudo-Code for Our Method 23
E.2 Implementation of Natural Gaussian Log Likelihood 24
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  $\mathbf{m}(x;\boldsymbol {\theta}) = [\mu (\mathbf{x};\boldsymbol {\theta}),\sigma^2 (\mathbf{x};\boldsymbol {\theta})]^{\top}$  and compute the Hessian of the Gaussian negative log likelihood w.r.t.  $\mathbf{m}$
$$
- \nabla_ {\mathbf {m}} ^ {2} \log p (y | \mathbf {x}, \boldsymbol {\theta}) = \left[ \begin{array}{c c} \frac {1}{\sigma^ {2} (\mathbf {x} ; \boldsymbol {\theta})} & \frac {y - \mu (\mathbf {x} ; \boldsymbol {\theta})}{(\sigma^ {2} (\mathbf {x} ; \boldsymbol {\theta})) ^ {2}} \\ \frac {y - \mu (\mathbf {x} ; \boldsymbol {\theta})}{(\sigma^ {2} (\mathbf {x} ; \boldsymbol {\theta})) ^ {2}} & \frac {(\mu (\mathbf {x} ; \boldsymbol {\theta}) - y) ^ {2} - \frac {\sigma^ {2} (\mathbf {x} ; \boldsymbol {\theta})}{2}}{(\sigma^ {2} (\mathbf {x} ; \boldsymbol {\theta})) ^ {3}} \end{array} \right] \tag {14}
$$
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
$$
\det  \left(- \nabla_ {\mathbf {f}} ^ {2} \log p (y | \mathbf {x}, \boldsymbol {\theta})\right) = - \frac {1}{2 \left(\sigma^ {2} (\mathbf {x} ; \boldsymbol {\theta})\right) ^ {3}} <   0.
$$
A simple way to overcome this issue is to map the mean-variance parameters  $\mathbf{m}$  to the natural parameters  $\pmb{\eta}$  using Equation 3 and apply the approximations as derived in Sec. 4, i.e.,
$$
\eta_ {1} (\mathbf {x}; \boldsymbol {\theta}) = \frac {\mu (\mathbf {x} ; \boldsymbol {\theta})}{\sigma^ {2} (\mathbf {x} ; \boldsymbol {\theta})} \text {a n d} \eta_ {2} (\mathbf {x}; \boldsymbol {\theta}) = - \frac {1}{2 \sigma^ {2} (\mathbf {x} ; \boldsymbol {\theta})}, \tag {15}
$$
and then apply the Gauss-Newton approximation as in Sec. 4 to obtain
$$
- \nabla_ {\mathbf {m}} ^ {2} \log p (y | \mathbf {x}, \boldsymbol {\theta}) \approx - J _ {\boldsymbol {\eta} | \mathbf {m}} ^ {\top} \nabla_ {\boldsymbol {\eta}} ^ {2} \log p (y | \mathbf {x}, \boldsymbol {\theta}) J _ {\boldsymbol {\eta} | \mathbf {m}} = \left[ \begin{array}{c c} \frac {1}{\sigma^ {2} (\mathbf {x} ; \boldsymbol {\theta})} & 0 \\ 0 & \frac {1}{2 \left(\sigma^ {2} (\mathbf {x} ; \boldsymbol {\theta})\right) ^ {2}} \end{array} \right], \tag {16}
$$
where the Jacobians  $\mathbf{J}_{\eta|\mathbf{m}}$  are of  $\eta$  w.r.t.  $\mathbf{m}$ . It is easy to see that the Hessian approximation in Equation 16 is positive definite by evaluating the determinants of all sub-matrices. The implementation is very simple: we simply use the output of the neural network  $\mathbf{m}$ , apply the transformation to the natural parameters  $\eta$ , 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  $\theta$  can be written as
$$
\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 {17}
$$
5 Interestingly, the individual determinants for the  $1\times 1$  and  $2\times 2$  sub-matrices are equal to the absolute value of the corresponding determinants of  $-\nabla_{\mathbf{f}}^{2}\log p(y|\mathbf{x},\pmb {\theta})$  reminiscent of saddle-free methods.
Figure 4: Illustration for the faithful heteroscedastic regression method proposed by Stirn et al. [2023]. Functions parameterized by neural networks are color-coded in blue, backpropagated gradients are color-coded in violet, and the two black arrows only indicate that the values for  $\mu$  and  $\Sigma$  are used as input for loss computation. The two modifications to conventional heteroscedastic regression introduced by Stirn et al. [2023] are color-coded in orange: the introduction of the scaling of the gradient  $\nabla_{\mu}\ell_{\mathrm{FF}}(\pmb{\theta})$  by the covariance, and of a stop gradient operation to avoid  $\nabla_{\Sigma}\ell_{\mathrm{FF}}(\pmb{\theta})$  gradients being propagated to  $f_{\theta_{\mathbf{z}}}$ . In this setup, specific dependencies for  $\Sigma$  on  $\mathbf{x}$  which are independent of the mean function are discarded, especially in the case in which  $\mathbf{z}$  acts as a bottleneck. Note that in this image we make the dependency of neural network functions on their respective parameters explicit. In the rest of the work we omit the  $\theta$  for the sake of brevity and readability.
As Seitzer et al. [2022] point out, the gradient of  $\ell_{\mu,\sigma}(\pmb{\theta})$  with respect to  $\mu(\mathbf{x};\pmb{\theta})$  is scaled by the estimated precision, i.e,
$$
\nabla_ {\mu} \ell_ {\mu , \sigma} (\boldsymbol {\theta}) = \frac {\mu (\mathbf {x} ; \boldsymbol {\theta}) - y}{\sigma^ {2} (\mathbf {x} ; \boldsymbol {\theta})}, \quad \nabla_ {\sigma^ {2}} \ell_ {\mu , \sigma} (\boldsymbol {\theta}) = \frac {\sigma^ {2} (\mathbf {x} ; \boldsymbol {\theta}) - (y - \mu (\mathbf {x} ; \boldsymbol {\theta})) ^ {2}}{2 \left(\sigma^ {2} (\mathbf {x} ; \boldsymbol {\theta})\right) ^ {2}} \tag {18}
$$
which they argue can lead to a poor mean estimate. In particular, they state that if the variance is well-calibrated, i.e.,  $\sigma^2 (\mathbf{x};\pmb {\theta})\approx (\mu (\mathbf{x};\pmb {\theta}) - y)^2$ , the gradient with respect to the mean reduces to
$$
\nabla_ {\mu} \ell_ {\mu , \sigma} (\boldsymbol {\theta}) \approx \frac {\mu (\mathbf {x} ; \boldsymbol {\theta}) - y}{\sigma^ {2} (\mathbf {x} ; \boldsymbol {\theta})} = \frac {1}{\mu (\mathbf {x} ; \boldsymbol {\theta}) - y} \tag {19}
$$
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.
$\beta$ -NLL To improve the mean fit, Seitzer et al. [2022] introduce the  $\beta$ -NLL loss
$$
\ell_ {\beta} (\boldsymbol {\theta}) = \left\lfloor \sigma^ {2 \beta} (\mathbf {x}; \boldsymbol {\theta}) \right\rfloor \left(\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\right) \tag {20}
$$
which introduces a stop-gradient operation denoted as  $\lfloor \cdot \rfloor$  to trade-off the influence of the variance estimate on the gradients. The corresponding gradients for the  $\beta$ -NLL loss are
$$
\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 {21}
$$
Thus, for  $\beta = 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  $\beta = 0$  it is the standard heteroscedastic gradient and for  $0 < \beta < 1$  the approach interpolates between both extremes. A special case appears at  $\beta = 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  $\mu(\mathbf{x})$ ,  $\sigma(\mathbf{x})$  should be faithful to a homoscedastic model that learns  $\mu_0(\mathbf{x})$ . As faithfulness they define that (cf. Stirn et al. [2023, Def. 2])
$$
\mathbb {E} _ {(\mathbf {x}, y) \sim \mathcal {D}} \left[ \left(y - \mu_ {0} (\mathbf {x})\right) ^ {2} \right] \geq \mathbb {E} _ {(\mathbf {x}, y) \sim \mathcal {D}} \left[ \left(y - \mu (\mathbf {x})\right) ^ {2} \right]. \tag {22}
$$
To achieve this property, they propose to decouple the estimation into three networks, as shown in Figure 4: a shared representation learner  $f_{\mathbf{z}}$  computes a representation  $\mathbf{z}$  from  $\mathbf{x}$ , which is passes into two individual networks  $f_{\mu}$  and  $f_{\Sigma}$ , which receive  $\mathbf{z}$  as input and output the mean respectively the covariance matrix. Based on this network architecture, they minimize the following loss
$$
\ell_ {\mathrm {F F}} (\boldsymbol {\theta}) = \frac {(y - \mu (\mathbf {x} ; \boldsymbol {\theta})) ^ {2}}{2} - \log \mathcal {N} (y | \left\lfloor \left(f _ {\mu} \circ f _ {\mathbf {z}}\right) (\mathbf {x}) \right\rfloor , \Sigma \left(\left\lfloor f _ {\theta_ {z}} (\mathbf {x}) \right\rfloor\right)), \tag {23}
$$
which incorporates two stop-gradient operations. They show that the proposal is equivalent to the following two modifications, which are illustrated in Figure 4: First, they stop the gradient of  $f_{\Sigma}$  from propagating to  $f_{\mathbf{z}}$ . Second, they propose to multiply the gradient of  $f_{\mu} \circ f_{\mathbf{z}}$  with  $\Sigma$ . 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  $\beta$ -NLL. In particular, for the sub-networks  $f_{\mathbf{z}}$  and  $f_{\mu}$  the gradients are multiplied with  $\Sigma$  which is the multivariate analogue of the effect of the  $\beta$ -NLL loss on the mean gradient for  $\beta = 1$ .
Proposed: Natural Parametrization The natural parametrization, that we use in this paper, has the following gradients with respect to  $\ell_{\eta}(\theta)$  with respect to  $\eta_1$  and  $\eta_2$  as
$$
\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 {24}
$$
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  $\mu (\mathbf{x};\pmb {\theta})$  from the estimated natural parameters, it is equal to  $-\frac{\eta_1(\mathbf{x};\pmb{\theta})}{2\eta_2(\mathbf{x};\pmb{\theta})}$ . Similarly, note that the variance  $\sigma^2 (\mathbf{x};\pmb {\theta})$  can be computed from the natural parametization as  $-\frac{1}{2\eta_2(\mathbf{x};\pmb{\theta})}$ . If we replace the corresponding quantities in Equation 24, we get that
$$
\nabla_ {\eta_ {1}} \ell_ {\boldsymbol {\eta}} (\boldsymbol {\theta}) = \mu (\mathbf {x}; \boldsymbol {\theta}) - y, \quad \nabla_ {\eta_ {2}} \ell_ {\boldsymbol {\eta}} (\boldsymbol {\theta}) = \sigma^ {2} (\mathbf {x}; \boldsymbol {\theta}) - (y ^ {2} - (\mu (\mathbf {x}; \boldsymbol {\theta})) ^ {2}), \tag {25}
$$
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  $\mathbf{X} \in \mathbb{R}^2$  be a centered multivariate Gaussian distributed with  $\Sigma = I$ . We construct the target variable  $Y$  as
$$
y = \sin (x _ {1}) + 1 / 2 \left(| x _ {2} | + 0. 2\right) \epsilon
$$
where  $\epsilon \sim \mathcal{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 train-validation/test according to the ratio (0.9/0.1). Then the train-validation data are split into training and validation sets (0.9/0.1).
As baselines we consider our empirical Bayes training with the natural homoscedastic and 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_{\mathbf{z}}$  and instantiate  $f_{\mu}$  and  $f_{\Sigma}$  as linear projections from dimension  $|\mathbf{z}|$  to 1. We ablate the size of the hidden dimension  $|\mathbf{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.