Stable Optimization of Gaussian Likelihoods
Authors: Denis Megerle
Fabian Otto
Michael Volpp
Gerhard Neumann
Venue: Submitted to ICLR 2023 (2023)
Keywords: Gaussian
Likelihood
Optimization
Uncertainty
Stabilization

Abstract: Uncertainty-aware modeling has emerged as a key component in modern machine learning frameworks. The de-facto standard approach adopts heteroscedastic Gaussian distributions and minimizes the negative log-likelihood (NLL) under observed data. However, optimizing this objective turns out to be surprisingly intricate, and the current state-of-the-art reports several instabilities. This work breaks down the optimization problem, initially focusing on non-contextual settings where convergence can be analyzed analytically. We show that (1) in this learning scheme, the eigenvalues of the predictive covariance define stability in learning, and (2) coupling of gradients and predictions build up errors in both mean and covariance if either is poorly approximated. Building on these insights, we propose Trustable, a novel optimizer that overcomes instabilities methodically by combining systematic update restrictions in the form of trust regions with structured, tractable natural gradients. We demonstrate in several challenging experiments that Trustable outperforms current optimizers in regression with neural networks in terms of the NLL, MSE, and further performance metrics. Unlike other optimizers, Trustable yields an improved and more stable fit and can also be applied to multivariate outputs with full covariance matrices.

STABLE OPTIMIZATION OF GAUSSIAN LIKELIHOODS
Anonymous authors
Paper under double-blind review
ABSTRACT
Uncertainty-aware modeling has emerged as a key component in modern ma- chine learning frameworks. The de-facto standard approach adopts heteroscedas- tic Gaussian distributions and minimizes the negative log-likelihen oog-likelhhod (NLL) un- der observed data. However, optimizing this objective turns out to be surpris- ingly intricate, and the current state-of-the-art reports several instabilities. This work breaks down the optimization problem, initially focusing on non-contextual settings where convergence can be analyzed analytically. We show that (1) in this learning scheme, the eigenvalues of the predictive covariance define stabil- ity in learning, and (2) coupling of gradients and predictions build up errors in both mean and covariance if either is poorly approximated. Building on these insights, we propose Trustable, novel optimizer that overcomes instabilities me- thodically by combining systematic update restrictions in the form of trust regions with structured, tractable natural gradients. We demonstrate in several challenging experiments that Trustable outperforms current optimizers in regression with neu- ral networks in terms of the NLL, MSE, and further performance metrics. Unlike other optimizers, Trustable yields an improved and more stable fit and can also be applied multivariate outputs with full covariance matrices.
INTRODUCTION
Generating models capable of quantifying uncertainty is crucial in modern machine learning appli- cations. These uncertainties arise, amongst others, due to underlying stochastic processes, inherent fuzziness of the data, partial visibility (Hüllermeier (02T). Efficient and reliable strategies for predicting uncertainties allows understanding the underlying data better, interpreting how well a model fares, guiding decisions we make based on model.
Conventionally, uncertainty is modeled using a probabilistic approach. That is, a model induces directly or indirectly a probability distribution describing uncertainty with its parameters. well- established technique of estimating alcatoric uncertainty, i. c., uncertainty that accounts for inherent randomness in an experiment, predicts distributions via neural networks (NNs). NNs are extensively addressed methods (Seitzer et al.) 2022] Detlefsen et al.) [2199] Abdar et al.) 2021; Hüllermeier, 202T]. primarily due to their ability learn non-lineas and complex relationships. Under the as- sumption of Gaussian distributed targets, the de-facto standard for various use cases, these NNs infer Gaussian mean and covariance conditioned on the respective input context.
A common objective for training such networks is the negative log-likelihood (NLL) objective, which is given by
0* arg min Σ log po (2t, vi),
where 0 are the parameters of the NN and D = {24, 3i is the training data. Typically, optimiza- tion involves computing some form of gradient VoLNLL 0, D). Yet. these Gaussian NLL gradient updates were shown be often unstable and hard use(Guo al. 2017 Hüllermeter 2021). This
for example typically the case if large portion of training data indicates highly certain subregions gradients explode (Takahashi et al. 2018) or fit only sub-par (Seitzer et al. 2022).
Hence, implementations across all domains rely on pragmatic techniques (Seitzer et al. 2022 Gar-
nelo et al. 2018: Schulman al. 2017al for stabilizing uncertainty output and gradient training for

---

Moldy E Fold
NNO
Me A
PDTRL
11
f Σ
LIL
A,
a,
Moldy E oold
Φ
va, VM
Rnew neww
PDTRL
nnew
Reg
Σ new
schematic of Trustable for heteroscedastic prediction of Gaussians. We retrieve the exact natural gradient of the target distributions via local parametrizations (Lin al. 2021 and regress on the closed-form trust region projected new distributions (Otto et al.; 2021] The local loss manifold aware natural gradient cancels destructive gradient scalings while trust regions stabilize the sensitive update allowing for stable learning of full covariances.
Figure 1:
NLLL-baeed learning. Unfortunately, these techniques also impair model expressivity and degrade prediction quality. For example, full covariance matrices in higher dimensions are primarily avoided (Gopinath] 1998 Fiszeder & Orzeszko 2021) due to their inherent instability during training. A gra- dient must keep all predicted covariances positive-definite at all times and ensure numerical stability. Most implementations, therefore, do not model correlations but only predict diagonal Gaussian distribution. Furthermore, clamping the variance into empirically performing regions (Raffin et al.) 2021 Garnelo et al., 2018) is classic makeshift patch to avoid small or large variances in which gradient descent with the Gaussian NLL becomes unstable (Seitzer et al. 2022). In this case, how- ever, the target distribution's 's true variance cannot be learned if it lies outside the clamped region. Other approaches for compensating destructive gradients include different variance parametriza- tion (Raffin et al. 2021), objective adjustments (Seitzer et al.) 2022) or a change to the optimization strategy (Detlefsen et al. 2019 Choi et al. 2004 Takahashi al. (2018).
Albeit all these efforts, reliably predicting large ranges of variances and correlations in data remains challenging. The noted shortcomings impact model expressivity and prediction quality. In this work, we examine the underlying cause of those shortcomings and open up new ways to further improve learning for wide variety of models based the NLLL.
Summary of Contributions. We first demonstrate that in an update step with the NLL,, the covari- ance plays distinct role in the convergence behavior. We extend the idea from Seitzer et al. (2022) and argue that the eigenvalues of the predictive covariance define stability in learning. Furthermore, coupling of gradients and predictions detrimental as it destabilizes training and can build up errors in both optimization parameters if either is poorly approximated. We then introduce an optimization strategy dubbed Trustable (see Figure D that methodically overcomes coupling and instabilities. Through structured, tractable natural gradients (NGs) (Lin et al., (021)) we efficiently obtain more informative, decoupled gradients and systematically constrain the update step via Trust Region Pro- jection Layers (TRPLs) (Otto et al.) (202I) for stability. Finally, we empirically find that Trustable displays an improved and less volatile fit in terms of the NLL, mean squared error (MSE), and the Wasserstein L2 distance (W2) for non-contextual and contextual regression with full covariance matrices in higher dimensions.
PRELIMINARIES
In supervised learning, the optimization algorithm receives a labeled set of training data D = {14, y. JeEooNN C X' ),
where the instance x; and target yi are distributed according to a complex, usually intractable distri- bution p(x,3) y). Our goal is to approximate the conditional distribution p(y|æ)) over some data includ- ing its aleatoric uncertainty. some Gaussian distribution, which is in the multivariate case given by N Ez). where μ Bd represents the mean and Σ Raxd the covariance. The covariance is positive semi-definite matrix
common assumption is that each outcome is distributed according
2

---

defining the individual variances in its diagonal and correlations in its off-diagonal elements. The eigendeconposition of the covariance matrix is given by Σ = QTAQ with Q as eigenvector and A strong function approximators such as NNs to map
diagonal eigenvalue matrix. Typically the input values to mean and covariance of the Gaussian distribution, i.e., f + Px, Ex.
For each sample (æ, y), the d-dimensional multivariate Gaussian defines the NLL via Equation 1 with its density function as
1 (d (d 2n + log Ee + (y (--)))) [x (y pz))
LLL
log PGauss (y (1z Eny)
(2)
where Piz Rd. Ex E Rdxd are sesitive-definit Standard differentiation yields the gradients
1 - (w)-))) px) T- Ex)_,T (3)
NLL
Ey Mx)
Vx, NN.L
Gauss
2
2.1 OPTIMIZATION STRATEGIES
Gradient descent is the method choice for optimizing NNs. Among SGD (Ruder) 2016) and
Adam (Kingma & Ba 2015), discuss assorted strategies to stabilize and improve gradients for NLL optimization..
Pitfalls. Seitzer et al. 2022 indicate that the variance scaling of the gradients introduces importance weights for each sample. As solution, they propose the univariate (3-LLL loss, multiplying the variance out.
Trust Regions. Trust regions are a prominent method in stabilization of non-linear optimization processes (Sun Yuan 2006). They produce robust updates by restricting the change output dis- tributions per iteration which is enforced by bounding the dissimilarity of the new distribution to the old distribution for each update. While many implementations resort approximations (Schulman iet al. 2017aab) TRPLs (Otto et al., (021) allow to impose trust regions for each input individually. Formally, the layer solves the constrained optimization problem
arg min denean (f ((x)),
s.t. Ameen ((x))(()))>))))))))}{{{))))) S cn
(4)
Az
arg min deow (E, E(x));
deor E(x)) S <3
for each input to find projected distribution parameters N( Sixr adhering to the trust region. The employed dissimilarity measures are dmean and dews with trust region bounds " and 6x, re- spectively. For example, we could employ the W2 distance Kullback-Leit divergence. The trust region layers can be made fully differentiable due to the use convex optimization Otto et al. 2021.
Natural Gradients. NGs improve convergence by incorporating information about the geometry of the loss manifold through the Fisher information matrix. Instead of expensively computing this matrix per sample (Kakade) 2002 Schulman et al. 2017a), Lin et al. (202I) leverage specific lo- cal, auxiliary and global parametrizations parametrization ensures non-singularity of the Fisher, while the auxiliary parametrization connects global and local parametrization via mappings ub, d. In local space, the Fisher information matrix collapses into an identity matrix. We employ their mappings introduced for Gaussian distributions
the prior distribution over the parameters. The local
= (A) (A): AAT) (#)= - Φ2 ( =
+ A.5 exp( (
and retrieve the NG for the global space by leveraging the chain rule through the mappings (Lin jet al. (02T). This approach alleviates some of the computations from previous NG methods such as computing the FIM. The only overhead is the parameter mapping.
papters matrix derivatives (specifically 0log X feex X if X asymmetric vs. 0 XXl/ax = 2XX -X ool if X
'The dichotomy for these derivations in various papers is the adhering to additional constraints of symmetric
symmetric). Indeed, the gradients here correct; for detailed discussion of this issue, we refer Srinivasan
Panda 2020.
3

---

2 8 8 8
[[Grndient NLL
0.8
Smalled Eigenvalue
10"
20
0.
-10
1.0
-05
0.4
0
0
20
100
0.5
10
0.0
0.2
det Z
(a) Univariant variance gradient
(b) Bivariate gradient magnitudes
Figure 2:
The covariance step progression and gradient magnitudes. dap The univariate variance gradient with a learning rate of a = 0.01 on òz,μ = (a ((t-)) (4) and o² experiences various undesirable traits. Values plateau towards larger variances and explode towards smaller ones. The dip in the middle indicates negative (out of parameter space) variance updates. (b) Covariance NLL gradient magnitude on the predictive covariance determinant. Standard Gaussian sampled eigenvalues gen- erate the predictive, bivariate covariance (with randomly rotated eigenvectors). We also sample 6x. through a standard Gaussian distribution. Overall, small eigenvalues indicate large gradient magni- tudes and small ò ب values indicate high sensitivity to the covariance eigenvalues.
LIKELIHOOD LEARNING WITH GAUSSIAN DISTRIBUTIONS
3.1 ANALYSIS OF NON-CONTEXTUAL GAUSSIANS
While the non-contextual setting allows closed form solution for mean and variance of multivari- ate Gaussian distribution, it is informative to look at the properties of gradient-barl based optimization already in this setup as we can draw conclusions for the contextual case with NNs as well.
Gradient Scaling, Much like Seitzer et al. (2022): Takahashi et al. (2018). we regard the gradient scaling with the covariance as main culprit for instabilities. However, we want to extend their perspective for univariate Gaussians to the multivariate case. In the univariate case, the gradient scales with o 2 for the mean and o -4 for the variance. Figure 2a displays this behavior for the univariate NLL gradient around both o² 0 and òz,a 0 and suggests several difficulties. First, cliff-like like structures close to o² O indicate abrupt variance changes, such that even values for sample error 8z, = (x - (4) corresponding to small variances still yield large destructive gradients. In these regions, the gradient is sensitive to noise in 82,, (cf. zero-variance problem in Takahashi ct al. (018)). This behavior is counter-intuitive, as we would expect progressively smaller gradients, especially when the true variance is small. Second, the predictive variance o² becomes larger, a plateau emerges, leading to a slow convergence as large steps are necessary for large variance domains reach the true variance.
For the multivariate case, we inspect the covariance update through eigendecomposition, which results for one step with learning rate in
Enxxt
Y a V
Gauss NLL
2
0 5
a
-5-1(x
1)(x 1=
2
2
( = QAQ"))))
Q
a
+ a
A
-QQ (x - p) w\ 11 QT.
MC estimate
The applied gradient interpolates between the current predictive covariance Σ and a rescaled ver- sion of the Monte-Carlo (MC) estimate (x- (x 4²) given the current predictive mean. We now examine the second term from the inside (the inner-most part being the MC estimate) to the outside. Suppose the outer product (x (x 14) defines an endomorphism. Then, the inner- most transformation QT #x(x = change of basis, expresses the covariance estimate
4

---

10
10/44
segllegCovar
SGD
Tractable
TAPLKL
10
SGb
Tractabee
10
Trustable
10a4
10P4
10
10"1
16-1
101
10
10"4
M Www.yu 1rr
10
10a4
30
15
20
25
120
200
300
400
300
200
300
400
mezn
Rerations
Iborat
Fitting initially small variance
(h) Fitting initially large (left) and small (right) eigenvalues
Figure 3: Fitting Gaussian distributions with varying parametrizations and optimizers. Training progression of univariate Gaussian from an initially small variance (o) for 400 iterations (until x) with SGD. Small variances during training lead to exploding mean and covariance gradients. The true distribution (red lines) cannot be recovered regardless of parametrization. (b) The Kullback- Leibler divergence between true and predictive distribution for 10D Gaussian with different initial covariance eigenvalues. Standard optimizers struggle, TRPL successfully compensates large, but does not converge for small gradients. While Tractable converges, initial extreme updates can lead to instabilities. In contrast, Trustable shows a fast and stable convergence.
via eigenvectors of the predictive covariance Σ. After this basis change, the estimate's new space scaled with the squared inverse eigenvalues the predictive covariances. In that case, the in- verse eigenvalues of the predictive covariance define the stability and convergence behavior along the corresponding eigenvectors. Along eigenvectors with small eigenvalues, large gradients harm the update. Along eigenvectors with large eigenvalues, small gradients slow down convergence. We empirically observe this concept when visualizing the (euclidean) magnitudes of randomly gener- ated two-dimensional covariance gradients in Figure 2b
Coupling & Joint Optimization. The mean and covariance gradients both depend on the predic- tive mean and covariance. From Equation (3, we can that the inverse of the covariance matrix scales the mean gradient. Thus, an incorrect covariance influences the mean through eigenvalues and eigenvectors of the matrix and destabilizes training. Vice versa, the predictive mean alters the covariance gradient in Equation B). To begin with, the MC estimate could be far from the true covariance. This discrepancy arises from potentially incorrect mean prediction and the stochasti- cally sampled which allow for possible outliers. Furthermore, the predictive covariance could be arbitrarily wrong, in which case the rescaling is faulty, In other words, if the MC estimate of the covariance for one sample given the predictive mean is far away from the predictive covariance, end up with significant, potentially destructive changes the predictive covariance disproportional to the true difference in distributions. This coupling is disadvantageous because it destabilizes the training and can lead to errors in both optimization parameters if either is in unfavorable region.
Structural Integrity. A dip is observed in the middle of Figure 2a the predictive variance takes on negative values outside the allowed parameter space. In other words, the NI.L gradient itself cannot sustain structural integrity for the predictive variance. Analytically, any function approxima- tor requires to predict positive-definite matrix in all training iterations. However, the gradient does ensure positive definiteness for any step. Especially in sensitive regions, the covariance quickly drifts out of the parameter space of positive definite matrices. Numerically, due to the volatility of training for small covariance eigenvalues, large region of the parameter space must be rep- resentable. However, the uneven distribution of floating point numbers and limited representation options lead to fewer actual representable, positive-definite matrices. In particular, for matrices that contain many numbers of sparse float regions (like close to or far away from zero), rounding emrors might occur.
Practical Implications. For the non-contextual setting, examine various optimizer and parametrization combinations, progressively increasing the complexity from low- to high- dimensional predictions (see Appendix A.I and Figure 3). The results confirm theoretical findings and expose eigenvalues and gradient coupling as the main culprits inducing a vicious cycle of harm- ful, ineffective updates.
5

---

Parametrizations can be divided into square-root-basee and log-based. The gradient is sensitized by square-root parametrizations, which improves convergence in previously plateau- areas while exacerbating the small eigenvalue issue. Therefore, the Cholesky parameterization for multivari- ate predictions experiences instabilities during our experiments. On the other hand, logarithm parametrizations highlight the regions around covariance of the standard Gauss distribution and improve convergence in the broader range, while degrading gradients further away.
Standard optimizers like SGD or Adam do not compensate for the exploding gradients. The Pitfalls adjustment 3-NLL (Seitzer et al. 2022) with high & values canceling out the covariance or update restricting optimizers like TRPL absorb the impact of destructively large gradients. Restricting the update step with trust regions stabilizes the training but slows down convergence, while not compensating for small gradients in regions with large variances. The natural gradient variants converge, but exhibit comparably large gradients in the beginning. Overall, we argue that these gradient explosions harm learning for neural networks. Hence, optimizers that lessen the cycle of damaging updates by decoupling the gradients, dampening volatile regions, and inciting plateau regions - outperform the competition.
As Figure suggesss, none of the above optimization strategies parametrizations could efficiently mitigate all influences of gradient scaling and coupling. Hence, one must choose an appropriate parametrization Pr.E to enforce the parameter space integrity, utilize more informative gradients, and restrict the update step for stabilization.
3.2 ANALYSIS OF CONTEXTUAL GAUSSIANS
This section extends the above study to the contextual case based on probabilistic NNs. We examine the additional difficulty adhering to parameter space constraints and sharing model parameters between contexts for prediction.
Network & Batching Gradients. Suppose a model defined by fo x + (Maz Ex)) is trained via batches. The gradient corresponds to expectation over various samples
0L1LL 00
T.NLL M., Opz, TaLNLL 0pz Opue Opz Opu- 00 σc, Ops. 00
where Pu: Pr are the parametrizations of the Gaussian mean and covariance, respectively. our case, both the mean and covariance gradient affect the same parameters they shared. If the parameters are shared between the distributions, harmful gradients can spread from each gradient to each distribution parameter. If the size of the network and computational cost allow it, this can be mitigated by splitting the network, i. e., providing network for the mean and a separate network for the covariance. Similarly, if the parameters are shared between contexts, harmful gradients spread between contexts. This concept of error spreading through the whole network due to a single or few volatile samples destabilizes the training process.
At the same time, batching reduces variance (Bottou 2012) for the gradient and stabilizes it. For the NLL, it dilates the effect of subset of harmful gradients. As a result, a NN trains stably as long as a significant portion of the training batch (for all batches) produces meaningful gradients. However, the resulting model is a subpar but stable fit as we either ignore or dampen some training samples.
Normalization. Typically, before using a data set for NNs, the data is normalized via mean and covariance. The idea center the data and normalize the spread in each dimension. This technique helps to train various types of networks (Bottou] 2012) but does not solve the NILL eigenvalue scaling issue. That is, normalization with global mean and variance does not correspond to normalization of each context individually.
Subpar Fitting. As already introduced by Seitzer et al. (2022), when we use batch gradient descent and average over the training batch, each sample defines its own importance (factor) through its gradient magnitude.
a multivariate setting, this weighting results in samples with large gradients, e., low predictive covariance eigenvalues, to be responsible for how the learning process is adapt- ing and behaving. Contrarily, samples with small gradients barely influence learning. As a result, suboptimal fits emerge (Seitzer al. 2022

---

Table 1: NLL, MSE, and W2 to the ground truth (# standard deviation) for the synthetic datasets.
NNCmmcii
Publls
Tartetfon NLL
Hioden
Miskll
NLL
MSE
ADAM
2.33 1.73
0,09
4.26
0.00
0.0161
28.60 C.91
1,54 2.42
Palalls
1:001 0.45
0.41 0:111
1.90 00.50
-6.28 1
0.01 0.01
0.0119 0.0000
2.90 0.00
36. 19:199 0.22
0.15 0.008
30.50
TRPL-W2
103 0.20
6.23 0.65
0.67=00.42
0.55 0:335
DJS DJIl
0,2319 = 01012A
2008 0.30
28.70 =
12,5T = 22.11
Thittbli
1.52
0.25
6.19-2.0:0 0.003
0.33
-6.45.1-0.15
0.012:000 0.00
0.0859 : 0.001
2.85 0.00
58.45 =2 0.06
0.01 2: 0.03
TrustaMt--y W22
0.75 0.00
6.21 0.06
0.00
-0.69 0.00
0.03 0.01
0/0047 =: 0.0012
2:00 (100
28.64 0.33
0.03
6638
Pitulls
6:29 0.25
0.0070
0,0015 0.0017
2.90=
0.24
0U05
S0s90a.0
TW2
100:2:0.15
0.45 0:33
0.35 135.2...0 0.30
-6.51
1 0.52
0J0322 0.0211
0.6062 5 0.0000
2:92 0.002
56.86 6.55
0.09 0.06
Thatable
2.02 0.0.0
6.13.3 0.00
0.56-2.0.2
-4-78 1:0.3
0.0111 1 0.0044
0.0015 2: 0.00011
2,900 0..01
38.39 6.20
0.01 0.00
TruszaMe-W W2
0.06
0.33
6.19 0.04
0.21::0.06
0.94 0.01
0.0089 A: 0.0033
0.0014:: 0.0012
2.00 0.001
28.55 =z 0.38
0.03
4 TRUST REGION TRACTABLE NATURAL GRADIENT
To recap, we require an optimization strategy that incites low gradient regions, dampens high gra- dient regions, decouples gradients, and stabilizes training. In this section, we propose Trustable, an optimization procedure that aims to satisfy all these. We combine tractable, structured NG (Lin al.
202T for improved convergence and a systematic restriction in the form of TRPL. (Otto et al. to stabilize the resulting distributions.
202
Non-contextual Trustable Update. The tractable structured NG update provides an efficient way to retrieve a more informative gradient for the Gaussian NLL loss. It grants information about the local loss manifold to cancel destructive gradient scalings. Additionally, preserves structural integrity, even with covariance priors or the desitive-definitenesss constraint of the covariance matrix. The update naturally decouples the mean gradient from the covariance altogether and descales the covariance accordingly (Lin et al., 202T). Since this NG is only a local approximation, we stabilize it by restricting the Gaussian parameter updates via closed-form projection based on TRPLs.
For the forward pass, first compute the TRPL projected distributions given predictions (Mt, Ad)- Then, we introduce zero-initialized local parameters (6, .M) and map them via Equation (6). In the backward pass, we backpropagate through the projections (6, M). According to Lin et al. (202 the global NG computes through chain-rule to the local parametrization, while the Fisher informa- tion matrix reduces to the identity matrix. As result, the distribution N (Pnew: Doow) the NG adapted version with the computational overhead of evaluating the mappings twice. The tractable update is a local approximation, we must ensure that the distance from new natural parameters to the previous iteration is restricted. Hence, we again project the new natural parameters via TRPL.
We found that this combination improves convergence, as the TRPLs prevent extreme updates of tractable NG while stabilizing the training. In several sample cases (see Appendix [A.T] this update scheme converges quicker and provides more stable targets than other tested strategies. Further, across the Gaussian parameter space, the updates keep delivering favorable gradients (see Figure 3b).
Contextual Trustable Update. When predicting a non-contextual Gaussian, we have access to its parameters and can update them exactly. For example, the tractable update indeed is the parameters' true NG. Additionally, the model update issue of TRPL (Otto et al. 2021) does not exist; the dis- tribution parameters update directly to the projection. In contrast, for contextual distributions, the Gaussian parameters are defined indirectly via the network. Hence, both the tractable NG and the TRPL cannot be updated by setting the new distributions precisely.
Figure I provides a schematic overview for this setting. The NN predicts the mean and cholesky (14, Ar) of Gaussian distribution. The TRPL is applied for each output by carrying on old distri- bution parameters. For each prediction, the TRPL projects its predictive distribution to adhere to a trust region respective the old distribution and the bounds E,s e,r Eg, solving Equations (4p and (5). We
then introduce iirtiinitialized local parameters (8, M). apply the mappings from Equation (6 to
the projected distribution (pe, Ao) and retrieve our actual predictions (p, 2). Note that, in this case, the mappings do not change the projected predictions. They are purely required to compute the NG towards the distribution parameters.
To this end, we compute the NLL loss and backpropagate to the local parameters. As per Lin et al. (2021), the new natural parameters (Mnew. Enew) can be computed directly through applying the mappings from Equation (6) with Vō and VΣ. The non-contextual experiments showed that these new parameters can be extreme (e g. case III of Figure 8. Hence, regressing on them might

---

Fixt nn connstantt
Fod nn mn_constamt
AD44
ADAM
4.03
F4t3i5 Tesitare-
Pite
2.90
6.22
6.02
1.00
2.02
6.0
0.0
1003
2000
2533
3500
500
1590
2000
3000
1003
1500
2000
2200
3200
<pects
(a) NLL
(b) MSE
(c) W2
Figure 4: Training progression for various optimizers on NNConstant. From left to right, the NLL, MSE between ground truth and predictive mean and W2 distance between ground truth variance and the predictive covariance. The plots represent the median and 85% confidence interval between 10 runs with the best hyperparameters obtained through grid search.
Table 2: UCI regression (multivariate) sets over 10 runs after hyperparameter tuning.
NLLL
MSE
Optimizer
carbon
energy
carbon
energy
ADAM
-1.18
8.00 + 0.74
62.71 50.11
Pitfalls
-12.24
7.19 2.74
9.79e-06 3.1ee006 06
2.54 1.18
Pitfalls
3
6.16 4.61
48.82 47.75
1.03e 7.33e-07
0.89 + 0.26
Trustable- W2
-12.41 + 0.19
3.06 0.83
1.30e 1.50e-00
1.28
destabilize our NN. We utilize the TRPL again to retrieve projected targets (Anew, Enew) reasonably close to old distributions and regress on the projected natural parameters train the network.
For this regression, any distribution metric, such as the W2 or the Kullback-Leebllr divergence (KL), can be employed. our case, we use the same metric that is used in the TRPL project the distributions.
EXPERIMENTS
We refer Appendices rameters. In the experiments we project via W2 metric, as the non-contextual experiments showed that these deliver good regression targets. Hence, we employ the NLL to evaluate the general model likelihood as well as MSE and W2 to the ground truth for the quality of predictions in terms of mean and covariance.
A.3 for detailed information about datasets, training, and hyperpa-
5.1 SYNTHETIC DATASETS
Step Function with Step Noise (NNConstant). We first implement a step-wise constant function with segments of varying variances named NNConstant. The challenge in this dataset comes from some segments with minuscule variances, leading to instability or slow learning. Figure 4 presents the confidence over the validation over 10 model runs after hyperparameter optimization. We notice more volatility in the variance prediction in both the NLL and W2 for other optimizers. In contrast, Trustable fits the function and reduces instability in all metrics. We conjecture that through the NG, our optimizer can better trade off the function approximator's allocation to fit the variance better. As a result, the overall likelihood increases. Furthermore, we test for both a two and three layer model (cf. Table D and find that only Trustable fit the data with smaller model size.
Sinusoidal without Heteroscedastic Noise (Pitfalls). We investigate Seitzer et al. 2022)'s 's illus- trative sinusoidal with homoscedastic noise: 0.4 sin Ç. with (0, 0.001). As the authors noted, an optimizer requires symmetry breaking to escape locally stable mean fits for such a function. Table display the best models after hyperparameter tuning over 10 trials w.r.t. NLL and W2. While we confirm the results from Pitfalls, where the basic Adam optimizer has issues fitting the function, both the Pitfalls and Trustable converge better fit.
8

---

Tho sweryy
2.4
AD4M TrataalerO
ADAM
X00M 7Mpal 16.51
3.0
0.20
TholudteW2
Tuasbba 992
23.4
2300
2000
3030
(a) 3DSpiral (I: 5). (left) NLL, (right) MSE
(b) UCI Multivariate NLL, (left) carbon, (right) energy
Figure 5: Training progression for various optimizers on multivariate data. a NLL and W2 between ground truth and predictive covariance for the 3DSpiral with & 5. (b) NLL UCI sets carbon (3 targets) and energy (2 targets) trained with full covariance. All plots represent the median and 85% confidence interval between 10 runs with the best hyperparameters obtained through grid search.
Sinusoidal with Heteroscedastic Noise (Detlefsen). The sine from Detlefsen al. (2019) with additive and increasing noise = x sin(2)) + æς1 + {2 with Cı, (O, 0.3) is relatively easy to fit for all methods (see Table TT Each optimizer learns it reasonably well according to all metries. We observe that the Trustable optimizer without mean bound converges slightly faster than the Adam and Pitfalls optimizer and, once converged, is stable.
3D Spiral with Contextual Full Covariance. This function is based on a diree-dimensional spiral of the form sx)) cos(t) (). y() = rsin(t), (t) at. We set 0.5 and 0.05 retrieve multiple windings for E [0; 20] describing our mean. For the covariance, two eigenvalues are fixed 0.1 and 0.2 and clip(-0.01
10.0)² + 0.1, 0.1,10 10-10 oo) describes the last eigenvalue. Finally, the covariance rotates on the z-axis based on t. Scaling the covariance with 10 * I increases the challenge of learning small eigenvalues. Figure 10 visualizes & O. We compare Adam, Pitfalls, and Trustable with respect to the NLL and the normalized W2 (Wang et al., 2022) in Figure 5a Adam does not converge, especially the covariance cannot be fit. Pitfalls cancels out portion the destabilizing gradients but remains volatile during optimization in both metries. Trustable reaches the target distribution in less iterations, typically experiencing less volatility.
5.2 UCI REGRESSION DATASETS
We examine the UCI datasets (Hernánezz-Lobato & Adams, 2015) check basic convergence. Table presents an overview for Adam, Pitfalls, and Trustable. While Trustable stays on par with other optimizers for the NLL, it improves the MSE on some datasets. On the two datasets with multiple outputs carbon (3 targets) and energy (2 targets), we learn a full covariance matrix. An overview over best NLL and MSE values are shown Table 2 Adam does not fit the function properly, achieving worse scores on all metrics. For Pitfalls, convergence depends on 3 but generally improves the fit, but high have hard time fitting a full covariance. The training progress in Figure suggests that, across training, Trustable performs smoother and fits to lower NLL.
6 DISCUSSION
In this work, report on the instabilities of NLL optimization and highlight the role of the predic- tive covariance for stability. Along eigenvectors with corresponding small eigenvalues, large gradi- ents destroy the update. Along eigenvectors with corresponding large eigenvalues, small gradients slow down convergence. Furthermore, we argue that coupling between the predictive parameters and gradients causes poor update steps in both optimization parameters if either one is poorly estimated.
We propose Trustable, combining TRPLs and tractable, structured NGs to methodically tackle these issues. The combination proves to be consistently par better than comparable optimizers on various regression tasks while stabilizing full covariance matrix prediction. The main limitation of our method lies in carrying old distribution parameters over for projection of the current predictions. While static datasets can benefit from Trustable directly, dynamic data like in reinforcement learning requires more approximations. In future work, we could address this issue by either regressing on the final network or adding an auxiliary loss to incorporate projection information, like Otto et al. 2021.

---

REPRODUCIBILITY STATEMENT
Hyperparameters and settings are provided in Appendices http:///hhhbboomm/////////222222
A.2 and A.3 Code is available at
REFERENCES
Moloud Abdar, Farhad Pourpanah, Sadiq Hussain, Dana Rezazadegan, Li Liu, Mohammad Ghavamzadeh, Paul Fieguth, Xiaochun Cao, Abbas Khosravi, U. Rajendra Acharya, Vladimir Makarenkov, and Saeid Nahavandi. Review of Uncertainty Quantification in Deep Learning: Techniques, Applications and Challenges. Information Fusion, 76:243 December 2021. ISSN 15662535, doi: 10.1016/j.inffus 2021.05.008 i000 06225
URL http: http:://xxxvv.rrr.ooo atwpv///gggbbb//21111111
Xiv: 1011.066 2011.06225.
Léon Bottou. Stochastic Gradient Descent Tricks. In Grégoire Montavon, Geneviève B. Orr, and Klaus-Robert Müller (eds.), Neural Networks: Tricks of the Trade: Second Edition, Lecture
Notes in Computer Science, pp. 421-436 Springer, Berlin, Heidelberg, 2012. ISBN 978-3- 642-35289 doi: 10.1007/978- 2-788-3-642 011179978-3-42-35288-88 978 3- 142233529-8822
URL httph:ttt://a/c /cci:: wttp:/.i...r..1100000011010....0 htt.:/te///....rr/111.o.1.. 1007/
Stephen P. Boyd and Lieven Vandenberghe. Convex optinization. Cambridge University Press, Cambridge, UK New York, 2004. ISBN 978-0-521-83378-3.
Heeyoul Heeyoul"Henry Choi, Sookjeong Kim, and Seungjin Choi. Trust-region learning for ICA. volume pp. 46, August 2004. ISBN 978-0-7803-8359-3 doi: 10. 10.11010CNN.2004.117888888
Nicki S. Detlefsen, Martin Jørgensen, and Søren Hauberg. Reliable training and estimation of vari- ance networks. arXiv: VrXiv:/906.033 9326003260 [cs, stat), November 2019. URL http::/arxxv.orrg
abs/1906. 03260
arXiv: 1906.03260.
Piotr Fiszeder and Witold Orzeszko.
Covariance matrix forecasting using support vector regres- sion. Applied Intelligence, (1(10):7029-7042 October 2021. ISSN 1573-7497. doi: 10. S10489-021-02211-5.
Marta Garnelo, Dan Rosenbaum, Chris J. Maddison, Tiago Ramalho, David Saxton, Murray Shana- han, Yee Whye Teh, Danilo Rezende, and M. Ali Eslami. Conditional Neural Processes, July 2018. URL /
1807 (1613 arXiv:1807.011 .06613 [cs, stat].
R.A. Gopinath. Maximum likelihood modeling with Gaussian distributions for classification. In Proceedings of the 1998 IEEE International Conference on Acoustics, Speech and Signal Pro- cessing, ICASSP '98 (Cat. No. No.98CH36181) volume pp. 661-664 vol.2, May 1998. doi: ISSN: 1520-6149.
10.1109/ICASSP.1999 6753511
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q. Weinberger. On Calibration of Modern Neural Networks. Proceedings of the 34th International Conference on Machine Learning, pp. 1321- 1330. PMLR, July 2017. URL http:::/ttp//oeeeeo sttps::/prroeeccnngs mllr wtt.ppppes//77//uuull 7 html
ISSN: 2640-3498.
José Miguel Hernández- and Ryan Adams. Probabilistic backpropagation for scalable learn- ing of bayesian neural networks. In International conference on machine learning, pp. 1861- 1869. PMLR, 2015.
Eyke Hüllermeier. Alcatoric and epistemic uncertainty in machine learning: introduction to concepts and methods. Machine Learning, pp. 50, 2021.
Sham M Kakade. A Natural Policy Gradient. In Advances in Neural Information Processing Sys- tems, volume 14. MIT Press, 2002. URL http:///tttt///oeeeiisss.../ www.eeiiisiinn..n....e.........e cC/ paper/
2001/hash 68 6abe4 8d35 Becf 9etccff944cc 94 Abst.eetb.ttttttttt......
D. P. Kingma and L. J. Ba. Adam: A Method for Stochastic Optimization. arXiv preprint arXiv:1412.6988 2015. URL httt::tt uva.arreeea..... i a20791
Pubiss
Publisher: Ithaca, NYarXiv.org.
10

---

Wu Lin, Frank Nielsen, Mohammad Emtiyaz Khan, and Mark Schmidt. Tractable structured natural gradient descent using local parameterizations. arXiv:2102.07405 [cs, stat], November 2021.
/tlptpp:///xxiv.rrr abs/ 07405
arXiv: 2102.07405.
Fabian Otto, Philipp Becker, Ngo Anh Vien, Hanna Carolin Ziesche, and Gerhard Neumann. Dif- ferentiable Trust Region Layers for Deep Reinforcement Learning, arXiv:2101.099 09000 [cs, math), March 2021. URL /ttp:////a.w..../////aaaaaa/
09207 arXiv: 2101.09207.
Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dor- mann. Stable-Baselines3: Reliable Reinforcement Learning Implementations. Journal of Ma- chine Learning Reseurch, 22(268):1-8. 2021. ISSN 1533-7928. URL http://tp/// jtmr..ml.rrr
Sebastian Ruder. An overview of gradient descent optimization algorithms. arXiv preprint arXiv:1609.04447, 2016.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust Region Policy Optimization. In Proceedings of the 32nd International Conference on Machine Learn-
ing. pp. 1889-1897. PMLR. June 2015. URL httttss:: httpps://prroceelnggs.o wwt.s..rpp../31/
schulman15.hta ISSN: 1938-7228.
John Schulman, Sergey Levine, Philipp Moritz, Michael I. Jordan, and Picter Abbeel. Trust Re- gion Policy Optimization, April 2017a. URL http:://axxxv.orr//a////////3 arXiv:1502.05477 [cs].
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov, Proximal Policy Optimization Algorithms, August 2017b. URL http:///araiv.orrgaaa////111 arXiv:1707.06347 [es].
Maximilian Seitzer, Arash Tavakoli, Dimitrije Antic, and Georg Martius. On the Pitfalls of Heteroscedastic Uncertainty Estimation with Probabilistic Neural Networks. Technical Re-
port arXiv:2203.09166 arXiv, April 2022. /tLtttp:///arxt wwwpiiii.rr./ooeeo http://aaav..or/aaa///0..... http://bbb/2200016688 arXiv:2203.09 [cs, stat] type: article.
Shriram Srinivasan and Nishant Panda. What is the gradient of a scalar function of a symmetric matrix 0rXiv:1911.066 [cs, math], May 2020. URL htt.tttpp//aaxv..oo////////1111 06491
arXiv: 1911.06491.
Wenyu Sun and Ya-Xiang Yuan. Optimization Theory and Methods: Nonlinear Programming. Springer Science & Business Media, August 2006. ISBN 978-0-387-24976-66 Google-Books ID: oOB YHL.
Hiroshi Takahashi, Tomoharu Iwata, Yuki Yamanaka, Masanori Yamada, and Satoshi Yagi. Student- Variational Autoencoder for Robust Density Estimation. pp. 2696-2702, 2018. URI. https:
/www. aaaai wwweaa..rgpprooceeeoo
Jinwang Wang, Chang Xu, Wen Yang, and Lei Yu. A Normalized Gaussian Wasserstein Dis- tance for Tiny Object Detection, June 2022. URL BRLhttp:/// /tt//vxiivv///////////////////aa http:///awxxvv.oo/aa//////////// 13389 MxXiv:2110.11 13399 [cs].

---

A APPENDIX
A.I NON-CONTEXTUAL IMPLICATIONS
We investigate the convergence of the NLL gradient for various optimizers and well-established parametric alternatives a non-contextual setting (see Figure(6p. We sample batches data directly
y: PuE - Ra X
R Raxa ++
L
P
Σ
Apuz
optimizer
Figure 6:
the non-contextual setting, one Gaussian distribution is predicted and parametrized directly by P..Z. We sample from the true distribution and compute the NLL loss, adapting the parameters as the respective optimizer defines.
from the true distribution and iteratively update predictive mean and covariance according to the optimizer. Therefore, we employ commonly known optimizers, trust region based optimizers to restrict destructive updates and more informative gradients such as the natural gradient:
Stochastic gradient descent (SGD)
Adaptive moment estimation (Adam)
Pitfalls adjusted SGD: We dub the B-NLL loss from Seitzer et al. 2022 Pitfalls. In con- trast to their work, expand the authors' idea to higher dimensions by rescaling the eigenvalues along each covariance eigenvector, analogously to one dimension, and retrieve the gradients
V VaC3-NLL E NI.
Gauss
VILe-NLL E Gauss Σ
NLL Caus
Natural gradient (NG): For the NG, we estimate the Fisher information matrix (FIM) per batch of training samples (Kakade (002) and utilize Trust Region Policy Optimization (TRPO)-like update (Schulman al. 2015) without line search
25 FTF-10
ÒNG 1 F VoLLLL
with as the trust region bound in terms of the Kullback-Leiberr divergence (KL) between the current and updated Gaussian distribution.
Tractable Structured NG: We optimize via tractable structured NGs (Lin et al. 2021 utiliz- ing Gauss parametrization with factorized covariance structure (Cholesky or Squareroot). The gradient propagates local parameters and computes the exact FIM in the local pa- rameter space, updating the distribution accordingly.
Trust Region Projection Layer (TRPL): In this algorithm dubbed Trp in the plots, the TRPLs (Otto et al. (02T) restrict the update step. We save one set old distributions (Hold, Eola) for the projections and evaluate all three metries Frobenius distance (FD), W2 and KL.
Combination of Tractable and TRPL: Combining the tractable NG with trust regions might lead to improved convergence. We dub this combination Trustable or TrpTrac and test the Tractable update with all three trust region projections.
12

---

Gauss Newstn:: Gauss-Newton defines baseline for these experiments. The NLL is not jointly optimizable (Boyd & Vandenberghe, 2004). Hence, we optimize via block- diagonal
update. To retrieve stable updates, the optimizer solves least- square optimization problem
H - with H VIL.
VE x:
(9)
computing the second order gradient = [u Veclāz]]T. We use Tensorflow's built-in least-square solver with L2 regularization of 10-7
The following univariate cases (see Table simulate real-world issues that might arise during opti- mization:
Learning minuscule variances Initially, a standard variance of o = 1.0 must be turned into the rather small true variance of o = 0.01. We would expect proper training at the beginning, and the smaller our variance becomes, the more inconsistent and unstable.
2. Baseline Training from mit 1.0 oe 1.0.
3. Unfavorable initialization Starting at of 0.01, expect the training be unstable with large leaps in variance and mean immediately. The target is of 1.0.
Another set of multivariate cases reflect various challenges for the covariance (see Table 4 We draw optimal covariances to learn as true covariances and the following initial covariances per case:
1. Optimal eigenvalues an optimal optimization problem, all true covariances are centered around eigenvalues of Hence, we expect rapid and stable convergence.
2. Mostly small eigenvalues The initial covariance exhibits mostly minuscule eigenvalues. In this case, the learning process should have large gradients and frequent instabilities, usually not converging the true covariance.
3. Mostly large eigenvalues Initially, random large eigenvalues in the covariance. When mostly large eigenvalues lead to small gradients, we expect slow convergence, likely not reaching the optimal covariance.
4. "A mix of II and III" This most difficult case combines II and III.
We draw the covariance eigenvalues randomly from the distributions depicted Figure 7 and ran- domly rotate the matrix in each span direction. Tables and summarize hyperparameters.
A, dotted line 60
20
100
3.0
1.0
optinal snall
2.5
0.8
(A), solid line
P(A), dotted line
1.5
0.4
1.0
80
0.2
0.5
0.0
0.0
0.0
0.2
0.4
0.6
0.8
1.0
1.2
1.4
X, solid line
Figure 7: The distribution of eigenvalues per non-contextual, multivariate test case. All case distri- butions are essentially d-distributions over the eigenvalues. For Case IV, Bernoulli sample from # (II) and 2 (III) with p = 0.5.
13

---

Table 3: Non-contextual experiments optimizer hyperparameters
Symbol
Optimizer
Description
Value
Default
Number of iterations
400
B
Batch size per update step
128
M
Number of mini batches
8
Learning rate
5e - 2
NG
Number of iterations
50
3
Pitfalls GaussNewton
Interpolation coefficient
{0.5, 1.0)
Number of iterations
50
Experiments 1D
NG
Natural gradient KL trust region
0.001
6p
TRPL
Trust region mean bound Trust region covariance bound
0.1**
TRPL
0.001
Experiments 10D
NG
Natural gradient KL trust region
0.1
Ey
TRPL
Trust region mean bound
1.0**
E
TRPL
Trust region covariance bound
1.0**
6p
TRPL- Force
Trust region mean bound
0,1**
6
TRPL Force
Trust region covariance bound
0,1**
6p
Trustable/TrpTTrc
Trust region mean bound
1.0**
C
Trustable/TppTra
Trust region covariance bound
1.0**
Default parameters are used if not specified otherwise ## Bound for all projections Frob/W2/KL.
Table 4: Non-contextual experiments cases breakdown
Case
Description
Initial Distribution
Target Distribution
Experiments ID
Minit 0.0
Mire 5.0
Learning minuscule variances
o 2
1.0
Juue 0.1
Minit = 0.0
Mque 5.0
II Bascline
o mit 1.0
2
22 1.0
Uinit 0.0
Minue 5.0
Unfavorable initialization
o init 0.1
2 1.0
Experiments 10D*
#4 51-5.0.5.0
14 4[-500.5.0
I
Optimal eigenvalues
X; Beta 2) + 0.5
X Beta 2) + 0.5
#4 2[-5.0.5.00
μ4 ~ U[-5.00.500
II Mostly small cigenvalues
X ~ Beta(0.5, 8) + 0.01
A ~ Bcta(2, + 0.5
#4 24[-5.0.5.0]
14 -21-5.0.5.0)
III Mostly large eigenvalues
A Beta(8, 0.5) 100
A; 2eta(2, + 0.5
#4 U1-5.0.5.0)
μ4 ~ U1-5.0,5.0)
IV
A mix of II and III
X ~ Beta(2, 0.5
A ~ Beta 0.5
* We sample Eigenvalues according to the distributions and compute > QAQT for random Q
14

---

Table 5: M-AUC and mean convergence iteration and time (ITER/TIME) of selected optimizers in the multivariate, non-contextual experiments for 10 runs per algorithm, case and parametrization on random seed.
Case ITER
Case
M-AUC
ITER
TIME
M-AUC
TIME
M-AUC
ITER
TIME
M-AUC
ITER
TIME
Vanalaa Cholesky
3.Hec
269.70
33.51
4-1lcee5
500e+2
4,500e43
SGD
4,38c+1
270.20
35.59
2.09ce6
1.22e+2
255.50
43.95
3,09cr6
SuntCovariacee
4.13c+1
270.20
35.04
2.6lem6
1.22++2
255.50
42.80
3.03c+6
2.54c40
157.80
2.08+42
Adam
Cholesky
109.60 80
[.98c=]
52.10
9.50c+1
151.40
21.39
2.43++1
$4.61
SgnCovariance
6.05e40
117.40
16.54
1
100.00
52.09
1.35e+2
162.0
22.70
2.20c=1
71.00
55.00
Vinilla
1.270+1
7.33
1.9c+1
300++
NG
Cholesky
1.26++
7.33c+1
1.97c1
N.28c+1
SovCovariance
1.26++1
7.33e=1
1,97++1
6.28c+
Vanilla
1.23:40
112.50
2570
5. TTT=2
4Mio02
5,20042
0.5
Cholesky
4.50c+0
41.90
9.14
2. 1.96+=
73.70
71.99
1.5te+1
48.50
10.71
3.51ee
79.20
78.18
SguCovariance
4.89c40
41.90
9.40
2.79ea4
71.50
74,90
1.5css1
48.50
3.36014
115.10
75.71
3.19+40
49.20
2.53e40
17.58
TRP W2
Cholesky
7.79c+0
122.30
2.800e+0
31.80
9,53
1,22e+2
255.50
59:67
3.26c+0
10.32
SqutCivrrianee
7.95c40
123.10
2.79c40
33.80
9.77
1.22e+2
255.50
3.17c40
38.30
10.97
Tracsable
TractablsCholeuky
ApprssimseTraciiee
8.85c-1
6.56e-1
22.80
24.60
12:52
3,86
1.67643
56.80
9.24
6.90c+0
6.50c+0
31.70
31.70
16.11
5.27
2.32c+3
59.70
9.62
Trussale-W2
Approsimate Traciable
Tintahle hilovky
8.270-1
8.27e-1
22.30
22.20
1578
7.19
1.68--0
22,80
23.00
14.01
7.29
6,00010
6.90e+0
31,80
31.30
17.86
9.56
1,93++0
1.93540
23.30
23.30
12344
7.33
Instability. We visualize the training progress in the univariate case with the KL between true and predicted distribution in Figure For the multivariate case, we evaluate speed and stability for each optimization run. Therefore, we compute the mean area under curve (AUC) over all optimization runs per optimizer and parametrization, as well as time and iterations to convergence (KI. < 10 and obtain Table 5 Runs suffering from numerical instability are marked as (-). The two pri- mary concerns causing instability are: Mean and covariance gradients in each eigenvector direction are scaled with the inverse eigenvalues squared. This scaling results tiny predictive covariance eigenvalues scaling the gradient up, potentially failing the update owing to huge, non-representable dositive-defime or out-of-paramen ance eigenvalues provide such tiny gradients that there is no convergence. Second, the mean and covariance gradients are linked with the predictions themselves. A poor mean prediction results in poor covariance gradient, and a poor covariance prediction results in poor mean gradient. This results in a vicious cycle of harmful updates.
space matrices. On the other hand, sizeable predictive covari-
Parametrizations. Parametrizations essentially adjust the scaling of the gradients. Two groups crystallize out: Square root parametrizations sensitize the gradient, improving convergence in pre- viously plateau-like areas but increasing the small eigenvalue problem. Logarithm parametrizations emphasize the regions around a covariance of / and improve convergence in the broader range while degrading gradients further away.
Algorithms. Converging combinations restrict the gradient directly (trust regions) or indirectly (by rescaling it proportionally to the variance, Pitfalls). Alternatively, they decouple the mean gradient from the variance such that even if one of the two parameters is far off its optimum, the other can still converge (NG, Newttn, Pitfalls). Canceling out the covariance to some extent is a reoccurring theme: The Pitfalls update does it directly. It turns out that true second-order methods like Gauss- Newton or similar methods like NG also comply with this idea. For example, the Pitfalls mean update with 3 = 1.0 is equivalent to the tractable NG mean update (ef. Equation (8) and Equation (12) of Lin et al. (2021)) Furthermore, this update is also equivalent to the Gauss-Newton mean update (cf. Equations (8) and (9). As result, can interpret exponentially smoothed maximum likelihood estimation (MLE) as a second-ordd optimization for the mean with respect to the NLL.
Unfortunately, none of the investigated optimizer parametrization combos solves all provided test cases efficiently. Simply restricting the gradient through various metrics, as via TRPL, stabilizes learning. Of course, then the tiny gradients are not influenced. On the other hand, the tractable NG update indeed converges even for low eigenvalue cases. However, its intermediate distributions can be extreme, hence unstable.
Non-contextual Trustable Finally, the Trustable update combines tractable NGs, which converge albeit with extreme intermediate values, with TRPLs for stabilization. As TRPLs prevent immediate volatility, the distributions never reach extreme values if not supported by the data and do not need recover from there, which saves time.
15

---

Variance
124
300
toffee
3345
33"
1.33
1.53
1.000
Precision
20
20
14-
15"1
15'
1.43
1.00
Sqrt Variance
334
1000
1044
10"
13:4
190
133
122
Log Variance
722
trghte
1.07
300
toffee
3
10""
12"
11-1
12
12:4
12"
0.31
1.30
1.00
Log Variance
22
tpfics
1.9
3
8
100
120
04
200
22
64
330
230
TractableCholesky
10"
3p
100
0.31
110
Case
Case II
Case
Figure 8: KL between true and predicted distribution for each non-contextual ID experiment case, parametrization and optimizer on time
16

---

15
0.-
02-
-5-
(a) NNConstant
(b) Pitfalls
(c) Detlefsen
Figure 9: Each plot depicts one function. The ground truth mean and standard deviation are orange, while a typically drawn training set can observed through the blue samples. All functions map from R. to I and exhibit contextual variances.
Figure 10: Visualization of mean (red), covariance confidence ellipsoids (grey) and samples (green) of 3DSpiral with & = O.
A.2 DATASETS AND FUNCTIONS
Step Function with Step Noise. This mapping captures complicated predict discontinu- ities and in the predictive means and variances through multiple segments. Each segment i can be described by y = C₁,1 + {c,2 with Ç ~ (0, We let C,1 increase equidistantly from
1-0,0 1.0] such that the function resides within y [-3.0, 3.0] and set c,2 [0.01, 0.02, 1.0, 5.0, 0.01, 0.02, 1.0, 5.0]. The result is NNConstant, visualized in Figure 9a Some segments experience small variances and, thus, influence the whole model through the zero-variance problem and sample importance weightings.
Sinusoidal without Heteroscedastic Noise. We utilize the function y 0.4 sin + with Gaussian (a = 0.01, see Figure 9b) distributed (Seitzer et al., 2022). For training, we evaluate 2000 samples across [0; 12] evenly spaced and keep the variance as ground truth for evaluation.
Sinusoidal with Heteroscedastie Noise. We utilize the function = r sin(x) + æ + C2 with
Ç1. S2 (0, 0.3) (see Figure 9ch from Detlefsen al. (2019). For training, we again evaluate 2000 samples across x E [0; 10] evenly spaced and keep the function to evaluate the ground truth variance.
spiral We set 0.5 and 0.05 retrieve two and half windings for [0; 20] describing our mean. For the covariance, two eigenvalues are fixed at and 0.2. The equation clip(-0.01
3D Spiral with Contextual Full Covariance. The function is based on
of the form a(t) cos(t), (t) = sin(t), (t)
- 10.0) + 0.1, 10-4, .2) describes the last eigenvalue. Finally, the covariance rotates on the z-axis based on t. To increase difficulty given analysis stating that small eigenvalues generate learning issues, we scale the covariance by 10 10-k J. Figure 10 visualizes the spiral for Io 0. We draw 5000 samples across [0: 20] evenly spaced through the respective means and covariances.
UCI Regression Datasets. We choose subset of the UCI datasets employed in HernándezzLooaar & Adams (2013). The data is zero-mean unit-vare miitvaaiance normalized on the training set, while the reported NI.. is provided in the original scale of the data.
17

---

Table 6: Hyperparameters for the contextual experiments
Symbol
Hyperparameter
Value
Default
learning rate
{le - 4,50 - le 3, 5e 3,, le 2}
batch size
512
B
Pitfalls Beta
0.5
Tractable Regression Metrics
MSE (mean) W2 (covariance)
B
Tractable Beta
le 4, 5e - 4, le - 3, 5e - 3, le - 2}
TRPL Type
W2
Detlefsen, Sinusoidals
mean {lel, oc}
Eus 6x
TRPL. Bounds
covariance {1e2, lel, 1e0, le 1}
3DSpiral
mean {0o}
Ejs eg
TRPL Bounds
covariance (1c1)
UCI Univariate
batch size
256
mean {le lel, 00 covariance {le le 3, 1}
Eyy TRPL Bounds
UCI Multivariate
mean {le - 1,e-1,,70 1e,el,00 covariance {le le 3, le 1}
Eps EΣ
TRPL Bounds
A.3 HYPERPARAMETER SETTINGS AND TRAINING
Step Function, Sinusoidals and 3DSpiral. We train a multi-layer perceptron with ReLU activations output parameter, i. e., for mean and covariance separately, predicting the covariance cholesky clamped between 1e-8 and 1e3 to extendable for higher dimensions. For the univariate functions, use either two hidden layers with 50 neurons each and for the 3DSpiral three hidden layers with 50 neurons each.
We take 70% of the samples as train set and the rest for testing. We split the train set further 70 30 and hypertune (over the parameters provided in Table 6 for 1000 epochs with an early stopping patience of 200 monitoring log-likelihood the validation set. We then train 10 models for 3000 epochs on the train set and evaluate on the test set.
UCI Regression Datasets. Like in Hernández-Lobate & Adams (2015), before but with only one hidden layer with 50 neurons for the univariate case. For the multivariate case increase to layers with 50 neurons each.
the same NN as
We follow the same training procedure as for the functions above, but split the data by 20% test and again 20% evaluation. Furthermore, we adjust 5000 epochs in both hypertuning and the train set with an early patience of 100 epochs for the univariate and 50 epochs for the multivariate case.
It should be noted that the performance provide is not comparable to that of other publications, as it is known to vary among data splits.
A.4 ADDITIONAL RESULTS
UCT Regression Datasets Univariate.
18

---

Table 7: UCI regression (univariate) sets over 10 runs after hyperparameter tuning.
DRAEZY
power
Paiill
0.07
0.22
2.03
-1.1810000 0.05
144.0 0.331
2.59 0.03
Pafilk
0.22
2:90
-1.10
-14.35-0.22
0.00
ThenshloW3
0.12
-1.15+00.07
-1396
0.05
carloon
contecte
eaagg
lonsing
katkmn
soval
powes
yadei
ADAM
0.0025
0.0044
24.25:t:2:44
2.75 2::::00 1.12
25.50 = 1.27
0.01 0.000
0.002 : 0.0066
11:24:2:030
0.59 :: 0.19
Pabii
8 0.5
0.0023 + 0.0008
23.01 +-2.0
0:60-1 0.14
18:46 + 0.53
0.01 0.0.0
0.0003 0.80072
10.70-t00.0 0.50
0.96 + 0.00
Pitalls
8 1.0
0.0028
0.0006
20.00 3,16
0.64 1: 0.10
18:23 1: 0.54
1.01.10000
0.0003
0.0001
10.61 1.
0147 1: 0.12
Trusibble W2
0.0029
0.0007
21.93 1. 1.67
0.58
0.15
17:05-1155
0.000
0.0003
0.0002
11.26.1.0.7
0.33....11
19


Reviews:
[{"review_id": "x6Nt1gthf-", "reviewer": "Reviewer_YtD7", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes Trustable, a novel optimizer that overcomes instabilities methodically by combining systematic update restrictions in the form of trust regions with structured, tractable natural gradients. It is demonstrated in several challenging experiments that Trustable outperforms current optimizers in regression with neural networks in terms of the NLL, MSE, and further performance metrics.", "summary_of_the_review": "This paper is well-done except the experiments could be extended to larger-scale and more settings.The violation of anonymity is also found in the github repo.", "strength_and_weaknesses": "Strengths:\n1. Well motivated. This is indeed an critical problem for uncertainty learning. \n2. Proposes good methedology from the perspective of eigenvalues of covariance and decoupled gradients. And also the good bulding of contextual gradients from the optimization process. All above brings the stability. \n3. promising experimental results\n\nWeakness:\n1. Somewhat small scale experiment and also only regression is tested in the paper. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "uSa6ntBSrGE", "reviewer": "Reviewer_gGLg", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes an optimization method for the Gaussian likelihood. The authors argue that optimization of the Gaussian log-likelihood is prone to instabilities. The method is a combination of two approaches: using structured natural gradients with local parameterizations, and constraining them with trust region layers. The authors provide experimental evidence of the instabilities of the Gaussian NLL optimization and performance of their proposed method. The new method is more stable and performs on par or better than other methods such as ADAM or Pitfalls.", "summary_of_the_review": "The paper provides a new strategy to optimize Gaussian NLL. The approach combines previously known methods. The authors should discuss more clearly the limitations of the method. My initial judgement is the paper is borderline.\n\nEDIT: Given the lack of authors' response and in light of other reviews, I'd like to lower my score to 3: Reject.", "strength_and_weaknesses": "*Strength*. The paper provides analysis of the source of the instability of the Gaussian NLL optimization. The authors did a multitude of experiments to provide evidence of the said instability. They provided necessary background information required to understand their approach. The method provides a more stable and in many cases better-performing optimization routine for Gaussian NLLs.\n\n*Weakness*. The new optimization strategy is a combination of previous work. The new method seems to achieve superior results on synthetic data while on UCI datasets, results are quite close for different optimizers (table 7). Although Trustable does seem to produce a more stable descent (Figure 5).\nLimitations of the method are not discussed clearly.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "gF6h-_Kx4h", "reviewer": "Reviewer_WXLj", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors consider the problem of Gaussian likelihood maximization. Estimating the conditional mean and the conditional covariance of a Gaussian likelihood model under neural modeling assumptions on the mean and the covariance leads to difficult optimization problems. The paper proposes a method to solve such optimization problem", "summary_of_the_review": "Interesting algorithmic contribution, but not very scalable.", "strength_and_weaknesses": "Pros:\nA fairly novel method to solve the ensuing Gaussian likelihood optimization problem. \nThe experimental results do seem to indicate better performance than other competing methods.\n\nCons:\nBy very nature the methods seems to be unscalable (as it estimates a full covariance matrix) to higher dimension.s", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "3Htb_B5rdX", "reviewer": "Reviewer_zKvQ", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper focuses on the problem of optimising multivariate Gaussian likelihoods, specifically wrt to its mean and covariance matrix. They identify the root causes of various failure modes in Gaussian likelihood optimisation through the lens of spectral analysis. Specifically, they characterise these modes based on the covariance matrix's eigenvalues and propose combining numerous techniques into a pipeline to mitigate these issues. In particular, the methods consist of structured natural gradients (NG) and trust region projection layers (TRPL).", "summary_of_the_review": "I recommend Rejection at this time. There are a number of significant issues in this paper that make it not yet ready for publication. In particular, as outlined above, there are gaps in the analysis of the proposed method and in the empirical results. Furthermore, there is room for improvement in the writing and presentation. That being said, this paper seeks to tackle an important problem, and could become a great contribution once the aforementioned issues are addressed. \n\nEDIT: I am not inclined to update my score. The authors have not provided a response, and I have not found supportive arguments from other more positive reviews compelling enough to change my views.", "strength_and_weaknesses": "\n## Strength\n\nThis paper seeks to tackle a timely and important problem, the solution for which would have a significant impact.\n\n## Weakness\n\nThere is no strong theoretical/analytical justification for the proposed approach. It is well-known that small eigenvalues can cause the matrix to be non-PSD in practice due to the limitations of finite-precision floating-point arithmetic. We can usually reason about this precisely or construct failure examples by first assuming some unit round-off error u, constructing a PSD matrix with some properties in terms of u, and showing that finite-precision arithmetic can lead to small negative eigenvalues (and thus non-PSD-ness). For the proposed method to be convincing, one would ideally like to see some analysis and guarantees that it is robust or can recover, regardless of how large the round-off error is. This may be true by construction, but is not immediately obvious, and needs to be discussed in the manuscript. There is some discussion on Pg. 5, under \"Structural Integrity\", but the argument is a little hand-wavy. I am not entirely convinced that the proposed optimization method can actually always mitigate these issues. Instead, they may simply be less likely to encounter these problematic regions.\n\nConcerning the experimental results, it appears the authors did tune the hyperparameters of their proposed approach on an intermediate validation/development set (according to A.3), but it is unclear whether the baselines against which they benchmarked their method benefited from any kind of similar tuning? Looking at Table 6, there appears to be a significant number of hyperparameters. \n\nSecond, I found the results themselves to be a bit underwhelming. In particular, the final test MSE often appears to be worse lower than the best-performing baseline? (e.g. Table 1-2)\n\nFinally, there is no indication of performance in terms of runtime. Instead of merely comparing convergence in terms of the number of epochs (e.g. Figure 5), it is important to get a comparison in terms of at least one of wall-clock time, throughput (samples per sec), or FLOPs.\n\n### Misc questions and issues\n\nHere is a (non-exhaustive) list of some miscellaneous issues and questions:\n\nPage 1:\n- \"fit only subpar\"\nPage 2:\n- \"empirically performing regions\": meaning ambiguous\n- \"destructive gradients\": meaning ambiguous\n- \"_Albeit_ all these efforts\": \"despite\"\n- \"improved and less volatile fit in terms of [...]\": vague\nPage 3:\n- \"allow to\"\nPage 4:\n- Figure 2 (a) \"Univariant variance\": univariate\nPage 5:\n- Figure 3 (a): this figure is quite difficult to parse. Also, there is nothing in the text dedicated to explaining what each item in the legend represents exactly. Most of these can be inferred but 'SqrtCovariance' -- does this just represent the matrix square root (Cholesky factorization)?\n- \"compensates large, but does not converge\"\n- \"Vice versa\"\n- \"rescaling is faulty\"\n- \"potentially destructive changes\"\n- \"disproportional\": \"disproportionate\"\n- \"cannot sustain structure integrity\": meaning ambiguous. Also, \"structural integrity\" is an odd choice of words.\n- \"requires to\"\n- \"the covariance quickly drifts out of the parameter space of PSD matrices\": please be more specific here. Presumably, you are referring to parameterizations that result, by construction, in a PSD covariance matrix such as the low-rank (Cholesky) or diagonal parameterization, and the reason they are not PSD is due to floating-point precision issues?\nPage 6:\n- \"the gradient is _sensitised_ by\"\n- \"degrading gradients further away\": meaning ambiguous\n- \"we argue that these gradient explosions harm learning for neural networks\": sentences like this are examples of where this manuscript is unnecessarily verbose. This \"argument\" can be unsaid because nobody would ever argue with this.\n- \"_inciting_ plateau regions\"\n- \"parameter space integrity\": the use of the word \"integrity\" here is non-standard and carries ambiguous meaning. I understand this to mean numerical stability and PSDness but this is not clear from the writing\n- Normalization: I would use terms like \"whitened\", \"standardized\", or \"de-correlated\" here, because normalization could mean so many other things (e.g. projected onto unit hypersphere, affine transformed to be in unit hypercube, etc.)\nPage 7:\n- \"_delivering_ favorable gradients\"\n- \"cholesky\": Upper case\n- \"respective the old distribution\"\nPage 8:\n- \"trade off the function approximator's _allocation_\": meaning ambiguous - I understand you are trying to say its \"representational capacity\" or something along these lines but this is difficult to decipher.\nPage 9: \n- \"without mean bound\"\n- \"once converged, is stable\": this sentence is tautological - if it's not stable, then it can't have converged?\n- \"multiple _windings_\": ?\n- \"_especially_ the covariance cannot be fit\"\n- \"but beta have a hard time fitting a full covariance\"", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject