SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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
(a) Average test log likelihood across seeds
(b) Best test log likelihood across seeds
Figure 5: Visualization of the results for data generated according to Problem C.1, where only Faithful is ablated with respect to the bottleneck size. Left: Average test log likelihood for each of the compared models, with standard error as error bars. Right: Best test log likelihood result achieved across seeds for each of the models.
the relevant information about the variance. We rather conjecture that the amount of information that is propagated is dependent on the initialization of the network, since no gradient actively encourages the network  $f_{\mathbf{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:
$$
y = \sin (x) + (0. 1 + | 0. 5 x |) \varepsilon \quad \text {w i t h} \quad x \sim \operatorname {U n i f} [ 2. 5, 1 2. 5 ] \quad \text {a n d} \quad \varepsilon \sim \mathcal {N} (0, 1). \tag {26}
$$
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 marginal-likelihood optimization algorithm [Immer et al., 2021a], also detailed in Alg. 1. We train the neural network parameter using learning rate 0.01 with Adam [Kingma and Ba, 2014] and optimize the hyperparameters using learning rate 0.1 with the same optimizer every 50 epochs for 50 gradient steps. In this example, we use the efficient Kronecker-factored Laplace approximation and use 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 neural network baselines, MC-dropout and mean-field VI. We find that the proposed heteroscedastic Laplace approximation with natural parameterization provides
more visually appealing in-between epistemic uncertainties and has a more sensible split into aleatoric and epistemic uncertainties. In Figure 6, we display the difference when the generated data were homoscedastic: despite this, our well-regularized heteroscedastic method performs on par or better with a homoscedastic baseline trained in an identical way.
Figure 7: Comparison of heteroscedastic Laplace (right) to Monte-Carlo Dropout [Gal and Ghahramani, 2016], mean-field variational inference, and homoscedastic Laplace on the Skafte [Skafte et al., 2019] regression example. Laplace obtains better "in-between" uncertainties as reported also by Foong et al. [2019].
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.
Objective	Regular-ization	PosteriorPredictive	LL (↑)
boston	concrete	energy	kin8nm	naval	plant	wine	yacht
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)
Deterministic VI	EB	✗	-2.41 (0.02)	-3.06 (0.01)	-1.01 (0.06)	1.13 (0.00)	6.29 (0.04)	-2.80 (0.00)	-0.90 (0.01)	-0.47 (0.03)
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)
✗	-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)
Naive NLL	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.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)
Table 6: Comparison of our models trained with EB regularization, with the work of Wu et al. [2019], that also adopts an EB regularization procedure.
Objective	Regular-ization	MSE (↓)
boston	concrete	energy	kin8nm	naval	plant	wine	yacht
Homoscedastic	EB	0.130 (0.012)	0.096 (0.005)	0.002 (0.000)	0.064 (0.001)	0.001 (0.000)	0.057 (0.002)	0.606 (0.024)	0.002 (0.000)
Naive NLL	GS	0.170 (0.020)	0.130 (0.013)	0.009 (0.005)	0.072 (0.001)	5.673 (5.216)	0.056 (0.002)	0.605 (0.022)	0.003 (0.000)
β-NLL (0.5)	GS	0.131 (0.013)	0.122 (0.007)	0.002 (0.000)	0.069 (0.001)	6.293 (5.367)	0.055 (0.002)	0.608 (0.023)	0.002 (0.001)
β-NLL (1.0)	GS	0.144 (0.012)	0.108 (0.009)	0.003 (0.001)	0.069 (0.001)	0.003 (0.000)	0.057 (0.002)	0.609 (0.022)	0.002 (0.000)
Faithful	GS	0.140 (0.015)	0.097 (0.007)	0.004 (0.001)	0.070 (0.001)	0.003 (0.001)	0.055 (0.002)	0.616 (0.023)	0.002 (0.001)
MC-Dropout	GS	0.131 (0.016)	0.131 (0.006)	0.057 (0.003)	0.117 (0.004)	0.145 (0.011)	0.059 (0.002)	0.606 (0.025)	0.124 (0.038)
VI	GS	0.131 (0.016)	0.134 (0.006)	0.056 (0.004)	0.076 (0.002)	0.146 (0.009)	0.058 (0.002)	0.645 (0.037)	0.042 (0.012)
Naive NLL	EB	0.136 (0.016)	0.100 (0.006)	0.002 (0.000)	0.065 (0.001)	0.027 (0.017)	0.055 (0.002)	0.621 (0.027)	0.003 (0.001)
Natural NLL	GS	0.111 (0.010)	0.081 (0.008)	0.002 (0.000)	0.067 (0.001)	0.001 (0.000)	0.051 (0.002)	0.612 (0.022)	0.003 (0.001)
EB	0.102 (0.011)	0.094 (0.007)	0.002 (0.000)	0.061 (0.001)	0.001 (0.000)	0.052 (0.002)	0.652 (0.046)	0.003 (0.000)
Table 3: Test MSE results for the compared models on the UCI regression datasets, where we report mean and standard error (dataset names in italic).
Objective	Regular-ization	Approximation	Posterior Predictive	LL (↑)
boston	concrete	energy	kin8nm	naval	plant	wine	yacht
Natural NLL	EB	kfac	✓	-2.44 (0.04)	-3.05 (0.02)	-0.94 (0.02)	1.36 (0.01)	6.64 (0.01)	-2.76 (0.02)	-0.94 (0.02)	-0.69 (0.05)
full	✓	-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)
kfac	✗	-2.51 (0.07)	-3.04 (0.03)	-0.83 (0.03)	1.35 (0.01)	6.69 (0.01)	-2.76 (0.02)	-5.41 (3.54)	-0.64 (0.05)
full	✗	-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)
Table 4: Ablation comparing test log likelihood results for the natural parametrization with empirical Bayes (Natural Laplace) adopting the full Hessian Laplace approximation (full) and the Kronecker-factored Laplace approximation (kfac).
Method	Regular-ization	Posterior Predictive Variant	LL (↑)
boston	concrete	energy	kin8nm	naval	plant	wine	yacht
Natural NLL	GS	mean	-2.45 (0.04)	-2.92 (0.04)	-0.73 (0.05)	1.32 (0.01)	6.66 (0.01)	-2.76 (0.02)	-0.94 (0.02)	-0.29 (0.07)
lse	-2.46 (0.04)	-2.92 (0.04)	-0.74 (0.06)	1.32 (0.01)	6.66 (0.01)	-2.76 (0.02)	-0.94 (0.02)	-0.29 (0.07)
Natural NLL	EB	mean	-2.39 (0.04)	-2.92 (0.03)	-0.71 (0.03)	1.36 (0.01)	6.66 (0.01)	-2.76 (0.02)	-0.94 (0.02)	-0.51 (0.04)
lse	-2.36 (0.03)	-2.93 (0.02)	-0.71 (0.03)	1.36 (0.01)	6.66 (0.01)	-2.76 (0.02)	-0.94 (0.02)	-0.51 (0.04)
Table 5: Albation comparing the two variants of posterior predictive for Natural NLL objective with Empirical Bayes regularization: mean and lse.
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.
(c) Replicate 3
Figure 8: Box-plots reporting test log likelihood results on the CRISPR datasets on single experiment replicated. For the survival-screen-A375 dataset, only two replicates are available. Note that the NLL mean-variance parameterization with MAP prediction, as well as the MC-Dropout baseline, markedly underperform in this setting, and are hence not reported in the comparisons.
β-NLL	Homoscedastic GS + PP	Homoscedastic EB + PP	Natural GS + PP	Natural EB + PP
Training	44 min.	44 min.	8 min.	49 min.	9 min.
Inference	0.023 sec.	0.124 sec.	0.124 sec.	0.134 sec.	0.134 sec.
Table 7: Average runtime over 5 trials including hyperparameter optimization.
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 [asd1; Ottawa, 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.
Objective	Regular-ization	PosteriorPredictive	MNIST with MLP	FashionMNIST with CNN
LL (↑)	DKL (↓)	RMSE (↓)	LL (↑)	DKL (↓)	RMSE (↓)
Homoscedastic	EB	✓	-4.02 (0.10)	3.75 (1.92)	11.9 (0.2)	-3.87 (0.02)	1.37 (0.08)	12.5 (0.2)
Naive NLL	GS	X	-4.56 (0.25)	1.39 (0.25)	20.9 (7.7)	-4.06 (0.16)	0.88 (0.15)	12.7 (0.2)
β-NLL (0.5)	GS	X	-4.30 (0.16)	1.13 (0.17)	16.4 (1.9)	-4.01 (0.02)	0.84 (0.02)	15.0 (0.1)
β-NLL (1)	GS	X	-4.21 (0.09)	1.04 (0.09)	15.0 (1.2)	-4.05 (0.04)	0.88 (0.03)	12.3 (0.1)
Faithful	GS	X	-4.17 (0.04)	0.99 (0.04)	14.1 (0.5)	-4.42 (0.05)	1.25 (0.05)	17.8 (1.6)
MC-Dropout	GS	✓	-4.14 (0.05)	0.97 (0.05)	15.9 (0.6)	-3.78 (0.05)	0.61 (0.05)	15.5 (1.1)
VI	GS	✓	-4.15 (0.01)	0.98 (0.01)	13.9 (0.1)	-3.86 (0.03)	0.69 (0.03)	16.3 (0.7)
Naive NLL	EB	✓	-3.85 (0.02)	0.75 (0.01)	12.2 (0.3)	-3.66 (0.01)	0.60 (0.02)	13.1 (0.3)
X	-3.92 (0.02)	0.75 (0.01)	12.2 (0.3)	-3.77 (0.02)	0.60 (0.02)	13.1 (0.3)
Natural NLL	GS	✓	-4.16 (0.01)	0.99 (0.01)	14.5 (0.1)	-3.66 (0.01)	0.61 (0.02)	13.2 (0.4)
X	-4.16 (0.01)	0.99 (0.01)	14.5 (0.1)	-3.78 (0.02)	0.61 (0.02)	13.2 (0.4)
EB	✓	-3.73 (0.01)	0.76 (0.02)	11.0 (0.2)	-3.66 (0.01)	0.61 (0.01)	13.1 (0.2)
X	-3.94 (0.02)	0.76 (0.02)	11.0 (0.2)	-3.78 (0.01)	0.61 (0.01)	13.1 (0.2)
Table 8: Image regression with heteroscedastic noise variance that depends on the rotation applied, i.e.,  $y \sim \mathrm{rot} + \sqrt{|\mathrm{rot}|}\varepsilon$ . Similar to the results with heteroscedastic label-based noise in Table 2, empirical Bayes and the posterior predictive improve the performance consistently.
Objective	Regular-ization	PosteriorPredictive	MNIST with MLP	FashionMNIST with CNN
LL (↑)	DKL (↓)	RMSE (↓)	LL (↑)	DKL (↓)	RMSE (↓)
Homoscedastic	EB	✓	-4.15 (0.02)	0.86 (0.07)	12.8 (0.4)	-4.06 (0.01)	0.64 (0.02)	12.7 (0.2)
Naive NLL	GS	X	-5.38 (0.00)	1.66 (0.00)	51.8 (0.0)	-4.18 (0.03)	0.45 (0.03)	12.8 (0.2)
β-NLL (0.5)	GS	X	-4.34 (0.10)	0.62 (0.10)	16.0 (1.5)	-4.27 (0.02)	0.54 (0.01)	15.5 (0.1)
β-NLL (1)	GS	X	-4.31 (0.05)	0.59 (0.05)	14.8 (0.8)	-4.26 (0.04)	0.53 (0.03)	12.6 (0.3)
Faithful	GS	X	-4.31 (0.03)	0.58 (0.03)	14.6 (0.5)	-4.57 (0.01)	0.84 (0.01)	21.0 (0.4)
MCDO	GS	✓	-4.29 (0.04)	0.56 (0.04)	16.1 (0.9)	-4.08 (0.02)	0.36 (0.02)	15.6 (1.1)
VI	GS	✓	-4.30 (0.01)	0.58 (0.01)	14.5 (0.1)	-4.10 (0.01)	0.38 (0.01)	15.0 (0.3)
Naive NLL	EB	✓	-4.09 (0.01)	0.42 (0.02)	12.7 (0.1)	-4.02 (0.01)	0.33 (0.01)	13.4 (0.2)
X	-4.14 (0.02)	0.42 (0.02)	12.7 (0.1)	-4.06 (0.01)	0.33 (0.01)	13.4 (0.2)
Natural NLL	GS	✓	-4.26 (0.00)	0.54 (0.00)	14.3 (0.1)	-3.99 (0.01)	0.33 (0.01)	13.0 (0.2)
X	-4.26 (0.00)	0.54 (0.00)	14.3 (0.1)	-4.06 (0.01)	0.33 (0.01)	13.0 (0.2)
EB	✓	-4.06 (0.01)	0.38 (0.00)	12.0 (0.1)	-4.00 (0.01)	0.33 (0.01)	13.2 (0.2)
X	-4.12 (0.01)	0.38 (0.00)	12.0 (0.1)	-4.06 (0.01)	0.33 (0.01)	13.2 (0.2)
Table 9: Image regression with homoscedastic noise variance that is fixed at scale of 10. Similar to the results with heteroscedastic label-based noise in Table 2, empirical Bayes and the posterior predictive improve the performance consistently. Also, homoscedastic and  $\beta$ -NLL variants perform relatively well due to the reduced heteroscedasticity.
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,  $\delta$  per layer. The algorithm optimizes neural network parameters  $\theta$  and hyperparameters  $\delta$  jointly during training and further the intermediate marginal likelihood values,  $\log p(\mathcal{D}|\delta)$  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
Require: model  $\mathbf{f}(\mathbf{x};\pmb {\theta})$  (natural parameters), dataset  $\mathcal{D}$  (initial) prior precision  $\delta$  , epochs  $E$  , burn-in   
 $B$  , frequency  $F$  , steps  $S$  , learning rate  $\alpha$  , hyperparameter learning rate  $\gamma$ $\forall l:\delta_l\gets \delta$  Set prior precision (regularizer) for each layer   
for epoch  $e\leq E$  do   
for batch  $\mathcal{B}\subseteq \mathcal{D}$  do Optimize model parameters  $\pmb{\theta}$ $\pmb {\theta}\leftarrow \pmb {\theta} + \alpha \nabla_{\pmb{\theta}}[\frac{|D|}{|B|}\log p(\mathcal{B}|\pmb {\theta}) + \sum_l\log \mathcal{N}(\pmb {\theta}_l|\mathbf{0},\delta_l\mathbf{I})]$  SGD parameter update   
end for   
if  $e$  mod  $F = 0$  and  $e\geq B$  then Optimize regularization  $\delta$    
for step  $s\leq S$  do   
 $\forall l:\delta_l\gets \delta_l + \frac{\partial}{\partial\delta_l}\log p(\mathcal{D}|\pmb {\delta})$  (Equation 12) Hyperparameter update   
end for   
end if   
end for   
if (Bayesian) posterior predictive then   
Compute  $\mu (\mathbf{x}_{*}),\sigma^{2}(\mathbf{x}_{*})$  as in Equation 13 Fast approximate posterior predictive   
else   
 $\mu (\mathbf{x}_{*})\gets -\frac{\mathrm{f}_{1}(\mathbf{x}_{*};\pmb{\theta})}{2\mathrm{f}_{2}(\mathbf{x}_{*};\pmb{\theta})}$  and  $\sigma^2 (\mathbf{x}_*)\gets -\frac{1}{2\mathrm{f}_2(\mathbf{x}_*;\pmb{\theta})}$  Point-wise MAP predictive   
end if   
Predictive  $\mathcal{N}(\mu (\mathbf{x}_*),\sigma^2 (\mathbf{x}_*))$
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:
from math import log, pi  
import torch  
C = -0.5 * log(2 * pi)  
def heteroscedastic_mse_loss(input, target, reduction='mean'):  
    '''Heteroscedastic Normal negative log likelihood.  
Parameters  
-  
    input: torch.Tensor (n, 2)  
    two natural parameters per data point  
    target: torch.Tensor (n, 1)  
    targets  
    '''  
    assert input.ndim == target.ndim == 2  
    assert input.shape[0] == target.shape[0]  
    n, _ = input.shape  
    target = torch.cat([target, target.shape(), dim=1)  
    inner = torch.einsum('nk,nk->n', target, input)  
    log_A = (input[:, 0].square() / (4 * input[:, 1]) + 0.5 * torch.log(-2 * input[:, 1]))  
    log lik = n * C + inner.sum() + log_A.sum()  
if reduction == 'mean':  
    return -log lik / n  
elif reduction == 'sum':  
    return -log lik  
else:  
    raise ValueError('Invalid reduction', reduction)