Variational Stochastic Gradient Descent for Deep Neural Networks
Authors: Anna Kuzina
Haotian Chen
Babak Esmaeili
Jakub M. Tomczak
Venue: Accepted by TMLR ()

Abstract: Optimizing deep neural networks is one of the main tasks in successful deep learning. 
Current state-of-the-art optimizers are adaptive gradient-based optimization methods such as Adam. 
Recently, there has been an increasing interest in formulating gradient-based optimizers in a probabilistic framework for better modeling the uncertainty of the gradients. 
Here, we propose to combine both approaches, resulting in the Variational Stochastic Gradient Descent (VSGD) optimizer. We model gradient updates as a probabilistic model and utilize stochastic variational inference (SVI) to derive an efficient and effective update rule. Further, we show how our VSGD method relates to other adaptive gradient-based optimizers like Adam. 
Lastly, we carry out experiments on two image classification datasets and four deep neural network architectures, where we show that  VSGD outperforms Adam and SGD.

Published in Transactions on Machine Learning Research (04/2025)
Variational Stochastic Gradient Descent for Deep Neural Networks
Haotian Chen"
chenhaotian.jit@@maal
Deportment of Mathematics and Computer Science, Eindhoven University of Technology, Netherlands.
Anna Kuzina"
arma..anaa@maiilcommm com
Department Computer Science,
Vrije Universiteit Amsterdom, Netherlands.
Babak Esmaeili
b. b.csmaili@uue...
Department of Mathematics and Computer Science, Findhoven University of Technology, Netherlands.
Jakub M. Tomezak
j. j.m. mmmtomczakBtue...
Deportment of Mathematics and Computer Science, Eindhoven University of Technology, Netherlands.
Reviewed on OpenReview: https: // openreview. net/ forum? id= wudATN cdy
Abstract
Optimizing deep neural networks one the main lasks in successful deep learning. Current seate-of-the-ars optimizers are adaptive gradient- based optimization methods such as ADAM. Recently, there has been incrcasing interest in formulating gradient-basce optimizers in probabilistic framework for better modeling the uncertainty of the gradients. Here, we propose to combine both approaches, resulting in the Variational Stochastic Gradient Descent (VSGD) optimizer. We model gradient updates as probabilistic model and utilize stochastic variational inference (SVI) to derive an efficient and effective update rule. Further, we show how our VSGD method relates to other adaptive gradient-based optimizers like ADAM. Lastly, we carry out experiments on two image classification datasets and four deep neural network architectures, where we show that VSGD outperforms ADAM and SGD.
1 Introduction
The driving force for deep learning success is efficient and effective optimization [Bottou 2012 Sun 2020
Deep Neural Networks (DNNs) introduce multiple optimization challenges due to their complexity, size, and loss landscape with multiple local minima, platcaus, and saddle points. Since the first attempt train DNNs with stochastic gradient descent. (SGD), there have been multiple approaches speed up the learning process and improve the final performance of DNNs.
Nowadays, adaptive mandient-based methods are the leading optimizers in deep learning. The first advance- ments arrived with the idea of adding momentum (Polyak] 1964) to the SGD update rule (SGDM) which demostrated impressive performance in training DNNs (Sutskever et al. 2013). The next major breakthrough was ADAM (Kingma Ba 2015), which uses the first and sccond momenta adapt the learning rate and gradients. Currently, ADAM probably the most widely used optimizer in deep learning due its relative insensitivity to hyperparameter values and much faster initial progress in training compared to SGD (Sun) 2020
Equal contribution

---

Published in Transactions on Machine Learning Research (04/2025)
Recently, there has been an increasing interest in formulating SGD in a probabilistic framework. An example is Mandt et al. (2017) in which SGD is used as an approximate Bayesian inference algorithm analyzed by relating it to the Ornstein- Uhlbeck process, A different approach utilizes a Bayesian perspective on SGD to model the uncertainty for datastreams (Liu et al. 2024).
In this paper, we propose to combine these two approaches, and, as a result, we obtain a probabilistic framework for adaptive madient-based optimizers. First, we formulate a probabilistic model of SGD. Second, employ stochastic variational inference SVI (Hoffman et al. 2013] to derive an efficient adaptive gradient- based update rule. We refer to our approach Variational Stochastic Gradient Descent (VSGD). We apply our method to optimize overparameterized DNNs on image classification. Compared to ADAM and SGD, we obtain very promising results.
Our contributions can be summarized as follows:
We propose Variational Stochastic Gradient Descent (VSGD), a novel optimizer that adopts a probabilistic approach. In VSGD, we model the true gradient and the noisy gradient as latent and observed random variables, respectively, within probabilistic model. Additionally, this unique perspective allows us to manage gradient noise more effectively by adopting distinct noise models for the true and noisy gradients, enhancing the optimizer's robustness and efficiency.
2. We draw connections between VSGD and several established non-probabilistic optimizers, including NORMALIZED-GGD, ADAM, and SGDM. Our analysis reveals that many of these methods can be viewed as specific instances or simplified adaptations of VSGD when a constant noise model is assumed for both the true and observed gradients.
3. Lastly, we carry out an empirical evaluation of VSGD by comparing its performance against the most popular optimizers, namely ADAM and SGD in the context of large-scale image classification tasks using overparameterized deep neural networks. Our results indicate that VSGD not only achieves lower generalization errors, but also converges at a competative rate compared to traditional methods such ADAM and SGD. We believe that these findings underscore the practical advantages of VSGD, making it compelling choice for complex deep learning challenges.
2 Background knowledge
Training Neural Nets with Gradient Descent Let us consider a supervised sctup for deep neural networks (DNNs) in which a DNN with weights 0 E RR predicts target variable y E y for a given object X' y. Here, we treat 0 as a vector to keep the notation uncluttered. Theoretically, finding best values of 0 can be achieved by risk minimization: R0)) = Ep(x.w) [e(fo(x), v)], where is a loss function. In practice, we do not have access p(x. y). but we are given a dataset D. Then, training of fo corresponds empirical risk minimization (Bottou 2012 Sun 2020 : L(0; D) = N E ((fo(xn) Un). We assume that 1.
x XE fo
the loss function is k-differentiable, C* at least
Many optimization algorithms for DNNs are based on gradient descent (GD). Due to computational restrictions calculating the full gradient (that is, over all training datapoints) for DNN, the stochastic version of GD, stochastic gradient descent (SGD), is employed. SGD results in the following update rule:
0: = 0 1
(1)
where The is learning rate, 9: VoL(O;M) is noisy version of the full gradient calculated over mini-batch of M data points, L(O; M) TM ΣnEm E( ((fo(x)) Bn); and we assume that E [ōx] 9t.
Related work Optimization of DNNs a crucial component in deep learning research and, especially, in practical applications. In fact, the success of deep learning models greatly dependent on the effectiveness of optimizers. The loss landscapes of DNNs typically suffer from multiple local minima, plateaus, and saddle points (Du & Lee. 2018), making the optimization process challenging.
2

---

Published in Transactions on Machine Learning Research (04/2025)
Since the introduction of backpropagation (Rumelhart et al. 1986 Schmidhuber 2015 which is based on
SGD updates, multiple techniques have been proposed to improve SCD. Some improvements correspond to
a better initialization of layers that results in a warm start of an optimizer (e.g., (Glorot & Bengio, 2010 He
ct al.
2015 Sutskever et al. 2013) ). or pro-defined learning schedules (c.g., (Loshchilov & Hutter 2016
Another active and important line of research focuses on adaptive gradient- optimizers. Multiple methods
are proposed to calculate adaptive learning rates for DNN, such as RMSPROP (Graves 2013 Tieleman &
Hinton 2012 ; ADADELTA (Zeiler) 2012). and ADAM (Kingma & Baj 2015). ADAM is probably the most
popular optimizer due to its relative insensitivity to hyperparameter values and rapid initial progress during
training (Keskar & Socher 2017 Sun 2020). However, it was pointed out that ADAM may not converge (eddi et a 2018), leading to an improvement of ADAM called ADMSGRAD (Reddi et al. 2018). There are multiple attempts to understand ADAM that lead to new theoretical results and new variants of ADAM (o.g.: (Barakat & Bianchi 2019 Zou et al. 2019)). ) In this paper, we propose a new probabilistic framework that treats gradients as random variables and utilizes stochastic variational inference (SVI) (Hoffman et al.) 2013) to derive an estimate of a gradient. We indicate that our method is closely related to ADAM and is an adaptive gradient-based optimizer.
A different line of research focuses on 2nd-order optimization methods that require the calculation of Hessian matrices. In the case of deep learning, it is infeasible to calculate Hessian matrices for all layers. Therefore, many researchers focus on approximated 2rd-order optimizers. One of the most popular methods is KFAC (Martens & Grosse 2015). Recently, there have been new developments in this direction (e.g., (Eschenhagen lt al. 2023 Lin et al., 2021; 2023)). However, extending K-FAC introduces memory overhead that could be prohibitive, and numerical instabilities might arise in low-precision settings due to the need for matrix inversions or decompositions. Here, focus on the 1 1 optimizer. However, see great potential for extending our framework to the 2nd optimization.
There are various attempts to treat SGD as a probabilistic framework. For instance, Mandt et al. (017) presented an interesting perspective on how SGD can be used for approximate Bayesian posterior inference. Moreover, they showed how SGD relates to stochastic differential equation (SDE), specifically, the Ornstein- Uhlbeck process. The idea of perceiving SGD as an SDE was further discussed in, e.g., (Yokoi & Sato 2019) with interesting conncctions to stochastic gradient Langevin dynamics (Welling & Tch 2011]. Another perspective was on the phrasing SGD as Bayesian method either for sampling (Mingard et al.)
2021
modeling parameter uncertainty for streaming data (Liu et al.) 2024) or generalization (Smith & Le
2017
A different approach uses linear Gaussian models to formulate the generation process of the noisy gradients and infer the hidden true gradient with the Kalman filter or Bayesian filter, e.g. (Bittner & Pronzato 2004
Vuckovie 2018 Yang 2021). In this paper, we propose to model dependencies between true gradients and noisy gradients through the Bayesian perspective using novel framework. Our framework treats both quantities random variables and allows efficient inference and incorporation of prior knowledge into the optimization process. A similar line of work is the natural gradient- gradient-based methods for Bayesian Inference, where the optimizer uses natural gradients instead of Euclidean gradients (Khan et al.] 2018 Osawa et al. 2019). Our work is different from this line of work, as they apply ADAM to learn the variational posteriors of the model. VSGD, on the other hand, learns a variational posterior for the gradients themselves and may not necessarily be applied to estimate the posterior of the model parameters.
3 Methodology
3.1 Our approach: VSGD
A probabilistic model of SGD There are several known issues with applying SGD to the training of DNNs (Bottou 2012 Sun 2020), such as slow convergence due to gradient explosion/ vaniihhing or too noisy gradient, poor initialization of weights, and the need for a hyperparameter search for each architecture and datasot. Here, we would like to focus on finding a better estimate of the true gradient. To achiove that, we propose to perceive SGD through a Bayesian perspective, in which we treat noisy gradients as observed variables and want to infer the posterior distríbutions over true gradients. Subsequently, we then the
3

---

Published in Transactions on Machine Learning Research (04/2025)
Wy
U
2
WT
91
91
gr
gt
Wa
Figure A probabilistic graphical model of our proposed VSGD optimizer. White circles represent latent variables, gray circles correspond to observable variables, filled small black circles denote control variables.
expected value of the posterior distribution over the true gradient gr and plug it into the SGD update rule (Eq. D In the following, we present how we can achieve that.
In considerations, we treat all quantities random variables. Hence, the true gradient 9r. is latent. variable, and observe the noisy gradient 9,- We propose model Ir and ف using Gaussian distributions with precision variables Wy and wg. representing the information of systematic and obscrvation noiscs, respectively. Here, we use 'observation noise' to indicate the sampling noise of the observed gradient and "systematic noise' to indicate the introduced error of trying to approximate the real, local area of the gradient surface with a simplified model. Both precision variables are treated as The advantage of treating precision variables as latent variables instead of hyperparameters is that this way, can incorporate prior knowledge over the systematic and observation noise as well having uncertainty over them.
latent variables.
We also introduce control variate to serve as the prior mean for 9e and is essential for our approach. We define = h1( as a deterministic message passed from previously observed noisy gradients to time t. Later (Eq. 7 we will discuss the specific formulation of the function h(-). In the case of 9,, we propose using It. as the mean.
We assume that the gradients across each dimension are independent by modeling them independently. This approach is adopted for computational efficiency. Moving forward, unless specified otherwise, all symbols should be considered scalars.
For the maximum number of iterations T,, the joint distribution is then defined as follows (Figure
p Wg.9r.r: Iuert u1:r) p(w,)p((w,) II
plo un ())p wg).
(2)
with the following distributions:
p(g.lugi ue) N(uu, wo
(3)
p(0: 9t, wg) = No u. p(w,) I'(77))
(4)
55)
p(wg) = I'(7, Kg7),
where the Gamma distributions I'(-) are parameterized by shape and rate. K, is hyperparameter designed to reflect our prior belief about the variance ratio between 9 and Ji, i.c., K, E " /E [w A recommended value for Kg is greater than 1, e.g. K, = 30, suggesting that the majority observed variance stems from observation noise w rather than systematic noise wT1 y shared hyperparamcter that indicates the strength of the prior, essentially equating the prior knowledge with 7 observations. A smaller value for 7, such as 1e-8, is advised to represent our prior ignorance. The remaining key question is the
4

---

Published in Transactions on Machine Learning Research (04/2025)
choice of the control variate uy way to carry information. In this paper, aggregate the information over by setting the prior mean for the current It as the posterior mean of the previous ones:
Ut = Epis -))
(7)
where uo could be set to 0.
An alternative approach to using to summarize the knowledge before t is to instead let Is explicitly depend on g1 However, incorporating the dependency of on 91a-1 explicitly would lead to a non-scalable, fully connected graphical model. A common alternative solution is to adopt a Markov structure, in which 9: depends only on the most recent observations -t-Lt-1- This would align the model with frameworks similar to Kalman filters or Bayesian filters (Kalman 1960). However, these models either require known precision terms beforehand, which is often impractical or necessitate online learning of precisions through computationally intensive methods such as nested loops in variational inference Monte Carlo methods, rendering them non-scalable for training deep neural networks. By introducing 2 and treating it as an additional observation, we can utilize SVT to avoid nested loops and still guarantee convergence to the estimated procisions (if any). As demonstrated in this paper, this stratogy proves to be both scalable and efficient.
Stochastic Variational Inference Since the model defined by Eqs. 8 requires calculating intractable integral over latent, variables, to find good estimate of gr, we aim to approximate the distribution p(Wg, Wig. I.., f1.r; u1r)) with a variational distribution q(wgg Wg. Iu:r T, @ur)), parameterized by the global and local variational parameters I and {dr]²=1 respectively. We achieve this by maximizing the evidence lower-bound (ELBO) defined by p and q
arg max Eg log p(wy, Wlg 9u. 01:r; u1:r) 7:01
(8)
q(wg, wg. 91.7; T, Φ1.))
We will demonstrate that this objective can be achieved concurrently with the original optimization objective 0* = (rg D), allowing each SGD iteration of 0 serve sinultaneously as an SVI iteration for Eq. S
Following SVT Hoffman et al. 2013 we employ the mean-field approach, assuming the variational posterior is factorized as follows:
((Wgs Wg. 9IT) = g[ (4 wg) II (g:)).
(9)
We define the marginal variational distributions as:
((wg) = T [(ag: bg):
(10)
((wg) = T(ag, bg);
(11)
((ot) = NNHug 2ggg)
(12)
for 1 T, where the global and local variational parameters are = (Ag, ag by bg) and di (Nt.gg oZ.0)² respectively.
We now outline the complete SVI update process to meet the objective set in Eq. 8 For detailed derivations, please see Appendix A.1 Let 7--1 = (at-1.g 01-1, i b l and di-1 = (#t-1.g; ( "2.1,) represent. the global and local variational parameters updated just before the t-th SVI iteration, and we set the control variate:
= Pt-1.91
(13)
as the variational version of u, as in Eq. 7 and let up At iteration t, we start by randomly drawing a noisy gradient ÿe with the sampling algorithm of your choice (e.g. using mini-batch of samples). Then,
5

---

Published in Transactions on Machine Learning Research (04/2025)
conditioned on 91, T--1; and uy, we update the local parameter Φi as follows:
b 01-1.6 1.9
b 1.g gt.
Mt.g
Me-1,9
(14)
1--1 + b1-19
1-- + 1
01-1.g +
11-1,0
Fig
(15)
1,9 b=-1.5
Note that Pt,g represents dynamic weighted average of the previous estimate Pt-1.9 and the noisy observation 9t- Weights are dynamically determined based the learned relative precision between Eq. 3 and Eq. 4
Ptg then acts as control variate for the subsequent iteration. Recall that ais big [wg]. Therefore, the
variance parameter of is equal to the inverse sum the two precisions characterizing the systematic and
observation noise. For an alternative kernel smoothing view on Eq. 14 please refer to Appendix A.2
Given local parameters 6t, update equations for the global parameters 7 have the form:
0t,g ot, +
(16)
6/ 7 + 0.5 (oï.9 + (Ur - pc-1.g)
(17)
b/g = Ky + 0.5 (o,.g (Uu. - 92)
(18)
In Eq. T6 at,g and øt.g remain constant throughout the iterations and depend purely on the prior strength y. For the sake of consistency, they are included in the update equations. Thus, the posterior belief of the precision parameters Wy and Wig are determined purely by the rate parameters (V/.g b.g)). Intuitively, K/ and V/.8 characterize the amount of systematic noise and observation noise accordingly. In both cases, the prior value and the estimated variance of the true gradient additively contribute to the estimation. The main difference ween the two comes the third term, which depends on how much the gradient mean changed in time for the systematic noise and the difference between the observed gradient and the mean estimate for the observation noise.
Finally, the SVI updates for the global parameters take the form:
bt.g (1 (t))(b) + Pt, (t)bk/
(19)
big (1 - (t) Pt, (t)
(20)
where (1) and (t) denote the learning rates of SVI across iterations L. A typical selection for p could be as follows:
(t) = t
(21)
(t.) = (t. t22
(22)
In this formulation, and K2 are hyperparameters that influence the behavior of the SVI learning rales. We typically limit Hi and K2 in the range (0.5, 1] to let ft1 and Pt.2 meet the Robbins-Monro conditions (Robbins & Monro 1951), and we select N2 > Ki to allow for a slower or larger decaying learning rate for Wy. For example, setting 11 0.8 and K2 = 0.9 allows wy to adapt more rapidly to new developments compared to wg-
Note that when using mini-batch of samples, Se is calculated as the average of the sampled gradients in the batch. While in VSGD it is possible to treat the samples in the mini-batch separately, see Appendix A.3 for details.
Having concluded the 1-th SVI iteration, we then incorporate the local parameter de to finalize the 1-th SGD iteration:
0, 0z-1
of Mt.g1
(23)
where p6. + oo.g = VELg represents an estimation the local Lipschitz constant. A summary of the steps of the VSGD algorithm can be seen in Algorithm 1

---

Published in Transactions on Machine Learning Research (04/2025)
Algorithm VSGD
Input: SVI learning rate parameter {Ki, N2), learning rate 17, prior strength Y prior variance ratio Kg. Initialize:
lo, 40.9 = 60.0 = Yi bo.s = bo.s = Ho, = 0
for t= to T do
Compute It for -)
Pt, K1
Pr.2 t2
Update Gigg He.g
Eq.
14
15
Update ar,g. at.
Eq.
16
Update br.g. bbgg
Eq.
20
Update Oz
b Exq.
23
end for
3.2 VSGD, Constant VSGD and their connections other gradient-based deep learning optimizers
Our probabilistic framework offers a way of estimating gradients. natural question to raise here is whether this framework is related to other 1st order optimizers used in deep learning. Similarly well-known methods like ADAM, VSGD maintains cache of the first and second momenta of the gradient, calculated weighted moving averages. However, unlike other methods, where these variables remain constant, VSGD adjusts its weights across iterations.
VSGD and SGDD First, we can show how our VSGD method relates to a version of SGD called NORMALIZEDSGGD (Murray et al. 2019 that updates the gradients as follows:
0, θ.-
7 90.
(24)
0z
If we take 7 + 00, Kg - 0, such that -om,, +o Kg = O, then VSGD approximately equivalent to NORMALIZED-SGD (Murray et al. 2019.. We can show it by taking bo.9 = Y, bo.9 = KGY. Then, for 7 + oc and K, ++0, 0, Eq. 14 yields the following:
Kgy + 91 KG + K,Y+ +
7
N1. = 0
91.
(25)
In subsequent iteration, y dominates all other values, we have: ar 7, bbgg 2 7, byg ~ K2y. As result, br.g 7, big RJ Kyy. By induction, the variational parameter for de al iteration 1 is the following:
KgT²
Kg 0,
ofg 7(K77 +7)
(26)
Kg +1
Kgy + di
? Ky +
M1.9 M--1 K t7
= 9.-
(27)
As a result, the update rule for approximately the update rule of NORMALIZE-SGD, namely:
0, 0,1
7 9t-
(28)
Constant VSGD Before we look at further connections between our VSGD and other optimizers, we first introduce simplified version of VSGD with a constant variance ratio shared between gr and 9r. This results in the following model:
p(w, grer Iiert unrr) p(w) Il plock: us ((0rlll,
(29)

---

Published in Transactions on Machine Learning Research (04/2025)
with the following distributions:
p(gel wi ue) Nui,, w
(30)
pl (e, = (1
(31)
p(w) - ((),)).
(32)
This approach is tantamount to imposing strong prior Eq. 2 dictating that the observation noise is the systematic noise scaled by Kg. This strong prior alfects only the ratio between systematic and observation noises. Regarding the prior knowledge of the observation noise w itself, we can maintain an acknowledgment of our ignorance by setting y to a small value such as 10-
Given its characteristic of maintaining a constant variance ratio, we refer to the algorithm CONSTANT VSGD. Since the derivation of CONSTANT VSGD aligns with that of other variants of VSGD, omit the detailed derivation here and directly present the update equations.
Algorithm CONSTANT VSGD
Input: SVI learning rate parameter K, SGD learning rate 7) prior strength y, variance ratio Kg.
Initialize:
lo, 40.a 7; bo.9 75 N0,, 0
for t = to T do
Draw Ht
= Update og.g. Mt.g
Eq. 33
34
Update az.g. b
Eq. 35
Update 0.
Exq
end for
K
1tg -4 1,9 Kg
Jt
(33)
Kg
oig
(34)
Kg + 11-1,0
= +1
(35)
beg =(1 11-p (p:))b Pe [7 0.5 (.,, +Meg - -.) 0.5Kg (ož,g + (Me, -1-1g9
(36)
0. -0
V Mr. + oî., 11.0
(37)
The complete CONSTANT VSGD algorithm is outlined in Algorithm the updates are notably simpler compared to VSGD, as outlined in Algorithm | yet they still offer sufficient flexibility to estimate the uncertainties the system. Sunmarizing the discussion above, we enumerate the key distinctions between CONSTANT VSGD and VSGD as follows:
1. In VSGD, two global hidden variables, Wy and Wgg employed to represent systematic and obscrvation noise, respectively. On the contrary, CONSTANT VSGD utilizes w for observation noise and constrains the systematic noise to fraction 1 of the observation noise. Due to these reduced degrees of freedom, all the update equations in CONSTANT VSGD simplified.
K,
In the first momentum update Eq. 14 of VSGD, the weights are determined by the ratio between systematic and observation noise. However, Eq. 33 of CONSTANT VSGD, these weights are fixed at K ki+1 } facilitating a direct comparison with other optimizers such as ADAM.
Rg
To sum up, now know how simplify VSGD CONSTANT VSGD. the following, we will indicate how CONSTANT VSGD is connected to other optimizers, namely: ADAM, SGD with momentum, and A AMSRRAD.

---

Published in Transactions on Machine Learning Research (04/2025)
3.2.1 Constant VSGD and Adam
First, we compare CONSTANT VSGD with one of the strongest and most popular optimizers for DNNs, ADAM (Kingma & Ba 2015). At the 1-th iteration, ADAM updates the first two momenta of the gradient. through the following procedure:
m, m + 3i)
(38)
= + -B2)
(39)
where B1,02 are hyperparameters. Then, 0 is updated as follows:
0, = 0,-1
me:
(40)
v
where my and ün are the bias-corrected versions of and vt, namely:
my
(41)
Ut
(42)
B2
Now, the counterparts of my and in CONSTANT VSGD, i.e. the first and second momenta:
9 -
(43)
E
| = 14 + o?.g
(44)
where Pt,g is defined Eq. 33 we set K,
91
then the first momentum updatos in Eq. 33 and Eq.
38
1-31
become identical.
Next, we examine the relationship between Eq. 39 and Eq.
44 Intuitively, Vt (Eq. 39 is an approximation of
the expected sccond momentum of the noisy gradient E
[ While Eq. 44 attempts to quantify the expected
second momentum of the actual gradient [g]].
By plugging to Eq.
33 and Eq.
to Eq.
obtain the following:
K2
E [
P 1.9 (K, + 1)2
+ Di (Kg + 1)2
(45)
2Kg
b
(K, + 1)2 1. + K2 +1 a
(46)
For clarily, we divide the aforementioned equation into two components: Eq. 45 and Eq. 46 The first
component aligns with Eq. 39 as both represent weighted sums an estimated second momentum and d The distinction arises in the sccond component, where CONSTANT VSGD introduces two additional factors 41-1. Specifically, 1t-1,g applies penalty on Eq 45 when 14-1, into the weighted sum: Pt-1,9 and K,+1 a:-1
1 and 9e have opposing signs. The factor K,+I a:-1 Le represents the K, fraction of observation noise learned from the data, while higher observation noise implics greater uncertainty in Jr. leading to higher expected value E [-]]-
In summary, CONSTANT VSGD described in Algorithm 2 can be regarded as a variant of ADAM. The first momentum update in both algorithms is equivalent when K, to 81 1- However, the second momentum update in CONSTANT VSGD incorporates an additional data-informed component that is dynamically adjusted based the observation noise learned from the data.
3.2.2 Constant VSGD and SGD with momentum (SGDM)
The second optimizer to which we compare CONSTANT VSGD is SGD with the momentum term (SGDM). SGDM is quite often used for DNN training due to its simplicity and improved performance compared
9

---

Published in Transactions on Machine Learning Research (04/2025)
to SGD (Luu et al. 2020b). At the t-th iteration, SGDM calculates increment to 8t-1, denoted by un expressed as a weighted sum between the previous increment f and the latest noisy gradient St, namely:
0 A1 + 79t.
(47)
0, 0,- Ut:
(48)
where À is a hyperparameter. After rearranging Eqq 47 we get the following:
11 11 + 79. (A + 7)
A P1-
X 0t
(49)
That is, the increment V in SGDM is weighted average of e4-1 and H., the counterpart of in CONSTANT VSGD is Me,9 updated by Eq. 33 Therefore, by choosing Kg such that Kg = " the update equations 49 and become identical except for constant scaling factor.
33
3.2.3
Constant VSGD and AMSGrad
Another optimizer related to CONSTANT VSGD is AMSGRAD (Reddi et al. 2018). variant of ADAM. AMSGRAD was introduced to address a specific issue of ADAM in estimating the second momentum. To correct the problem, a factor contributing to ADAM's suboptimal generalization performance was introduced. AMSGRAD rectifies it by incorporating the long-term memory of the past gradients.
Specifically, the momentum accumulation process AMSGRAD remains identical to ADAM, as defined in Eq. 38 and Eq. 39 However, for the parameter update step, AMSGRAD differs from AMSGRAD (as
in Eq. 40 whether bias-corrected not) by incorporating the maximum of the past second momentum estimates rather than relying solely on the most recent value, that is:
ê max (0--1, Ut),
(50)
"
0: 0
m..
(51)
Vo:
Here, the long-term memory is carried out by calculating the maximum of all historical second momentum estimations.
Similarly to AMSGRAD, CONSTANT VSGD incorporates form of long-term memory regarding past gradient information in its second um update equations
45 and 46 However, this is not achieved by computing the maximum of historical values. Notably, the factor K,+1 Lik in Eq. 46 represents the expected observation
noise, estimated using data up to the (t - 1)-th iteration. In this expression, while Kg and a remain constant for 1. 2 2, the variable component is b which is determined by Eq. 36 for clarity we restate it as:
big =(1 pc)bi-1.8 + PLS,
(52)
where:
s= 8 0.5 + (Mr.g - -.)))) 0.5K, (0ž.g + (Mr.g - 1))
(53)
It represents the combined squared systematic and observation noises after considering 9.. Thus, Eq suggests that 1y-1, is a balanced summation the historical estimate of by and the newly squared disturbancos. The forgelting rale decreases over l, Pt 1"; promoting longer memory retention in successive iterations. CONSTANT VSGD, the magnitude of the memory retention can be modulated through the hyperparameter K.
52
4 Experiments
In this Section, evaluate the performance of the VSGD compared to other optimizers on high-dimensional optimization problems The task we focus on in our work is training DNNs to solve classification tasks.
Code available github.coom hetpub/.cooggeeratiiiiiitteeeeee
10

---

Published in Transactions on Machine Learning Research (04/2025)
4.1
Setup
Data We used three benchmark datasets: CIFAR100 (Krizhevsky et al. 2009 TINYIMAGENET-2N
Deng et al. 2009a and IMAGENET-1K (Deng et al. 2009b). The CTFAR100 dataset contains 60000 small (32 x 32) RGB images labeled into 100 different classes, 50000 images are used for training, and 10000 are left for testing. In the case of TINYIMAGENET-22 200, the models are trained on 100000 images from 200 different classes and tested on 10000 images. For a large-scale experiment, we use the IMAGENET- dataset which contains 1,281,167 images from 1000 classes.
Architectures We evaluate VSGD on various optimization tasks and use three different neural network architectures: VGG (Sionnyan & Zisserman 2015). RESNEXT (Xie et al. 2017), and CoNvMIXER (Trockmani & Kolter 2023 We keep the model hyperparamelers fixed for all the optimizers in experiments. We provide detailed description of each architecture in Appendix B.1
We conducted a grid search over the following hyperparameters: Learning rate (all optimizers); Weight decay (ADAMW, VSGD); Momentum coefficient (SGD). For each set of hyperparamcters, we trained the models with three different random seeds and chose the best one based on the validation dataset. The complete set of hyperparameters used in all experiments is reported in Table 3
Hyperparameters
Furthermore, we apply the learning rate scheduler, which halves the learning rate after each 10000 training iterations for CIFAR100 and every 20000 iterations for JINYMMGENNTT-200 We train VGG and CoNv using batch size 256 for CIFAR100 and batch size 128 for TINYLMAGENET-200, We use smaller batch size (128 for C1FAR100 and 64 for TINYIMAGENET-200) with the RESNEXT architecture to fit training on a single GPU.
4.2 Results
In Table 5 we compare the top-1 test accuracy for ADAM, ADAMW, SGD (with momentum), VSGD, VSGD with weight decay for the different architectures and datasets. The accuracy is averaged across three random seeds. We observe that VSGD almost always converges to better solution compared ADAM and SGD, outperforming ADAM by average of 2.6% for CIFAR100 and 0.9% for TINYIMAGENRI-200 Additionally, show how the top-1 test accuracy progresses during training. We plot the average values (solid line) minimum cattinummmaximuu maximum values (shaded area). We observe that VSGD's convergence is often the same rate or faster compared to ADAMW. We present training curves for all models in Appendix
In Figures 2 and
B44
Table 1: Final Average test accuracy, over three random seeds.
Table 2: Average training time on GeForce RTX 2080 Ti (scconds per training itera- tion) on CIFAR100 dataset.
VSGD (w/ 122)
VSGD (w/o 1.2)
ADAM (w/o 1.2)
ADAMW (w/ 12)
SGD (w/ mom)
CIFAR100
MODEL
TraInIng TIME
VGG16
70.1
70.0
66.8
66.6
67.9
CoNvMIXER
69.8
69.1
66.5
67.0
65.4
NAME
PARAMS (x10*)
ADAM
VSGD
RESNEXT-18
71.4
71.2
68.2
69.7
68.5
TINYTMAGENET-
VGG16
14.8
0.38
0.41
CONVMINER
0.6
0.36
0.38
VGG19
51.2
52.0
47.6
49.0
50.9
RESNEXT-18
11.1
0.84
0.87
CONvMIXER ResNeXT-18
53.1
52.6
51.9
52.4
52.4
48.7
47.2
48.8
48.9
47.0
Runtime Stability As can be observed in Amgorithm [ VSGD requires additional operation each gradient update step compared to the ADAM optimizer. However, we did not observe large computational overhead compared to ADAM when training on the GPU. To illustrate this, provide the average training time (per iteration) in Table
We Tiny ISaafoet CS231N) provided on inage-net.org
11

---

Published in Transactions on Machine Learning Research (04/2025)
VGG16-bn
Conv 256/8
ResNeXt-118
0.7
0.7
0.7
0.6
0.6
0.6
AlaaW
AdamW
Adaun
SGD (mom)
SGD (mom)
SGD (mon)
VSGD
VSGD
VSGD
0.5
0.5
0.5
10000
20000
30000
10000
20000
30000
10000
20000
30000
Iteration
Iteration
Iteration
Figure 2: Top-1 1 test accuracy on CIFAR100 dataset.
VGG19-bn
ConvMixer-256/
ResNeXt-11
0.5
0.5
0.5
0.4
0.4
AdamW
0.4
AdamW
Adaan
SGD (mom)
SGD (mom)
SGD (monn)
VSGD
VSGD
VSGD
0.3
0.3
0.3
20000
40000
60000
20000
40000
60000
20000
40000
60000
Iteration
Iteration
Iteration
Figure 3: Top-1 test accuracy TINYIMAGENET-200 dataset.
ADAM is known to perform well on deep learning tasks without the need for extensive hyperparamcter luning. We observed that our approach also demonstrates stable performance without the need to lune the hyperparameters for different architectures and datasets separately. Therefore, we believe that VSGD is generally applicable to various DNN architectures and tasks, as the case with ADAM.
Imagenet In order to evaluate VSGD's performance in overparameterizee domains, we performed an experiment on the Imagenet-1k dataset (Deng et al. 2009b) where we trained a ResNet-50 with both VSGD and ADAM. Due to limited computational resources, we not perform any hyperparameter tuning for this experiment. For both optimizers, we used a learning rate of 0.1 with a ONECYCLELR learning rate scheduler (Smith & Topin 2019) and a batch size of 256. Same as our previous experiments, we set y to 5e-8 for VSGD. Both models were trained for 50 epochs. The top-1 accuracy of the validation set for VSGD was 72.7%, exceeding ADAM's 69.8% by 2.9%. We also notice slight improvement in convergence speed (Figure 6).
5 Conclusion
In this work, we introduced VSGD, novel optimization framework that combines gradient descent with probabilistic modeling by treating the true gradients as a hidden random variable. This approach allows for a more principled approach to model noise over gradients. Not only does this framework open a new path to optimizing DNNs, but one can establish links with other popular optimization methods by certain specific modeling choices in our method. Lastly, we evaluated VSGD on largo-scale image classification tasks using a variety of deep learning architectures, where we demonstrated that VSGD consistently outperoreeed the baselines and achieved competitive convergence rate.
In conclusion, we outlinc two main directions for extending our framework. The first dircction is to make stronger dependency assumptions in VSGD. For example, we can model the dependencies between different. gradients by introducing covariates between gradients different parameters in Eq. [3] or considering the second-order momentum of gradients in the VSGD update rules. Second, we advocate for the application of VSGD to a broader spectrum of machine learning challenges beyond classification tasks such as deep generative modeling, representation learning, and reinforcement learning.
12

---

Published in Transactions on Machine Learning Research (04/2025)
Acknowledgements
This work was carried out on the Dutch national e-infrastructure with the support of SURF Cooperativ. Anna Kuzina is funded by the Hybrid Intelligence Center, a 10-year programme funded by the Dutch Ministry of Education, Culture and Science through the Netherlands Organisation for Scientific Research, bybrid--imeellligeeeeeeetrrrr
References
Anas Barakat and Pascal Bianchi. Convergence analysis of a momentum algorithm with adaptive step size for non convex optimization. arViv preprint orXiv:1911.07596, 2019.
Barbara Bittner and Lue Pronzato. Kalman filtering in stochastic gradient algorithms: construction of a stopping rule. In 2004 IEEE International Conference on Acoustics, Speech, and Signal Processing, volume 2, pp. 709. IEEE, 2004.
Léon Bottou. Stochastic Gradient Descent Tricks, pp. 421-436. Springer Berlin Haidelberg, Berlin, Heidelberg, 2012. ISBN 978-3-641-332889-8 doi: 10.1007 111007/978-3--12-3522555 25. URL http::///abi.crr//1.000777 978-3-642-35289-8
J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and L, Fei-Fei. ImageNet: A Large-Scale Hierarchical Image Database. In CVPR09, 2009a.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEFE Conference on Computer Vision and Pattern Recognition, pp. 248 255, 2009b. doi:
Simon Du and Jason Lee. On the power of over- per-paramettization in neural networks with quadratic activation. In International conference on machine learning, pp. 1329-1338. PMLR, 2018.
Runa Eschenhagen, Alexander Immer, Richard E Turner, Frank Schneider, and Philipp Hennig. Kronecker- Factored Approximate Curvature for Modern Neural Network Architectures. In Thirly--eventth Conference on Neural Information Processing Systems, 2023.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference artificial intelligence and statistics, pp. 249-256 JMLR Workshop and Conference Proceedings, 2010.
Alex Graves. Generating sequences with recurrent neural networks. arXie preprind orXio:1308. 300:008000 2013.
Kaiming Hc, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human- level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, pp. 1024-1034, 2015.
Matthew D Hoffman, David M Blei, Chong Wang, and John Paisley. Stochastic variational inference. Journal of Machine Learning Research, 2013.
Rudolph Emii Kalman. A new approach to linear filtering and prediction problems. 1960.
Nitish Shirish Keskar and Richard Socher. Improving generalization performance by switching from adam to sgd. arXiv preprint Xiv: 1 076287 2017.
Mohammad Khan, Didrik Nielsen, Vool Tangkaratt, Wu Lin, Yarin Gal, and Akash Srivastava. Fast and scalable Bayesian deep learning by sright--peruurbation in Adam. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 2611 2620. PMLR, 10 15 Jul 2018. URL https://proceedings.oo ppwss/vv0/khaneeeeee kwtp800/ha118a.... htmll
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference Learning Representations (ICLR), San Diega, CA, USA, 2015.
13

---

Published in Transactions on Machine Learning Research (04/2025)
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Wu Lin, Frank Nielsen, Mohammad Emtiyaz Khan, and Mark Schmidt. Structured second-order methods via natural gradient descent. 2021.
Wu Lin, Valentin Duruisseaux, Melvin Leok, Frank Nielsen, Mohammad Emtiyaz Khan, and Mark Schmidt. Simplifying momentum-based positive-definite submanifold optimization with applications to deep learning. In International Conference on Machine Learning, pp. 21026 21050. PMLR, 2023.
Haitao Liu, Yew-Soon Ong, Xiaobo Shen, and Jianfei Cai. When gaussian process moets big data: A review of scalable gps. IEEE transactions on neural nelworks and learning systems, 31(1)):4405 4423, 2020a.
Tianyi Liu, Yifan Lin, and Enlu Zhou. Bayesian stochastic gradient descent for stochastic optimization with streaming input data. STAM Journal on Oplimization, 34(1):389-411) 2024.
Yanli Liu, Yuan Gao, and Wotao Yin. An improved analysis of stochastic gradient descent with momentum. Advances in Neural Information Processing Systems, 33:18261-18271, 2020b.
Nya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.0390 2016.
Stephan Mandt, Matthew D Hoffman, David M Blei, et al. Stochastic gradient descent approximate bayesian inference. Journal of Machine Learning Research, 18(134):1 35, 2017.
James Martens and Roger Grosse. Optimizing neural networks with tronecker-factore approximate curvature. Tn International conference on machine learning, pp. 2408 2417. PMT.R, 2015.
Chris Mingard, Guillermo Valle-Pérez, Joar Skalse, and Ard A Louis. Is sgd bayesian sampler? well, almost. The Journal of Machine Learning Research, 22(1):3579-3664 3642, 2021.
Ryan Murray, Brian Swenson, and Soummya Kar. Revisiting normalized gradient descent: Fast evasion of saddle points. IEEE Transactions on Automatic Control, 64(111)4418-44844 2019.
Kazuki Osawa, Siddharth Swaroop, Mohammad Emtiyaz E Khan, Anirudh Jain, Runa Eschenhagen, Richard E Turner, and Rio Yokota. Practical deep learning with bayesian principles. Advances in neural information processing systems, 32, 2019.
Boris T Polyak. Some methods of speeding up the convergence of iteration methods. Ussr compulational mathematics and mathematical physics, 4(5):1-17, 17, 1964.
Joaquin Quinonero-Candela and Carl Edward Rasmussen. A unifying view of sparse approximate gaussian process regression. The Journal of Machine Learning Research, 6:1939 1959, 2005.
Sashank J Roddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. In International Conference on Learning Representations, 2018.
Herbert Robbins and Sutton Monro. stochastic approximation method. The annals of mathematical slotisties, pp. 400-407, 1951.
David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. Learning representations by back-propagating errors. nature, (2308) 1986.
Jürgen Schmidhuber. Deep learning in neural networks: An overview. Neural networks, 61:85-117, 2015.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. In International Conference on Learning Representations, 2015.
Leslic N Smith and Nicholay Topin. Super-convergence: Very fast training of neural networks using large learning rates, In Artificial intelligence and machine learning for multi-domain operations applications, volume 11006, pp. 369 386. SPIE, 2019.
14

---

Published in Transactions on Machine Learning Research (04/2025)
Samuel L Smith and Quoc V Le. A bayesian perspective generalization and stochastic gradient descent. arvivv preprint arxi:: vrXxim171000044 06451, 2017.
Ruo- Yu Sun. Optimization for deep learning: An overview. Journal of the Operations Reseorch Society of
China, 8(2):249 299:2499444 2020.
Πya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum in deep learning. In International conference on machine learning, pp. 1139 1147. PMLR, 2013.
Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5-rmsprop, coursera: Neural networks for machine learning. Universily of Toronlo, Technical Report, 6, 2012.
Asher Trockman and J Zico Kolter. Patches are all you need? Transactions Machine Learning Research, 2023.
James Vuckovie. Kalman gradient descent: Adaptive variance reduction in stochastic oplimization. ar preprint Xrxi::1810. 12273, 2018.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pp. 181-688. Citeseer, 2011.
Saining Xie, Ross Girshick, Piotr Dollár, Zhuowen Tu, and Kaiming He. Aggregated residual transformations for deep neural networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1492 1500, 2017.
Xingyi Yang. Kalman optimizer for consistent gradient descent. In ICASSP 2021-2021 IEEE International Conference Acoustics, Speech and Signal Processing (ICASSP), pp. 3900-3904. IEEE, 2021.
Soma Yokoi and Issei Sato. Bayesian interpretation of SGD as Ito process. arXiv preprint arXic:1911.090 2019.
Matthew D Zeiler. Adadelta: an adaptive learning rate method. arXiv preprint wxXiv:1212.500.1 2012.
Fangyu Zou, Li Shen, Zequn Jie, Weizhong Zhang, and Wei Liu. A sufficient condition for convergences of adam and rmsprop. In Proceedings of the IEEE/CVF Conference computer vision and pattern recognition, pp. 11127-11135, 2019.
15

---

Published in Transactions on Machine Learning Research (04/2025)
A Derivations
A.1 Derivation of VSGD
Here, we derive Eq. 14 to Eq. 18 following the SVI technique introduced by (Hoffman et al. 2013 3. First, we
provide a brief overview of the SVI procedure before diving into the details.
In the 1-th SVT iteration, we realize the following steps:
Sample from Our, where the sample is denoted by dc-
2. Update the local parameters P1.1: with 7-1 and dr, such that
plga wegs us ( Ih, wg))
LL arg max B (gg log algx: Øk)
== ==1:T,
(54)
dx
where p(gx ( uz) and p Ik, wa) are defined in Eq. and Eq. respectively. Note that here we use the same Ir for each & E 1 T, as if the same is observed T times.
Calculate the intermediate global parameters o/ with 0r.I.T, such that:
T|- arg max ml log p(wy)p(wg) I p(gN (wgi (o)lle wa) (( Wgit)
(55)
where p(w,) and p(wa) are defined in Eq. and Eq. respectively. Update the global parameters Ti using:
To p)T= pTós
(56)
where Pt ((,,,, defined in Eq. 21 and Eq.
22
5. Pick any 6t,* (Me,k. ,T, E Ir,ITT and set u++1 Pe,k, where u4+1 serves as the control variate for the next iteration. We can choose any òi because by the definition (Eq. 54h, di and Φt.j are identical for any i,j E.:
Since all 6r.k are identical for k = 1 T, for notational ease, we will use o, = (Hr.gy of )z to represent the unique local parameter updated in the f-th SVI iteration. Distributions defined Eq. 2 are in the conjugate exponential family, therefore, we can derive the closed-form update rule for di according to Eq. 54
log qlyni pt oiyy) = (Bg [log (gg u.) + log p(O:l wg))))
(57)
Next, by matching the coefficients of the natural paramcters, we get:
b 1.
b
P4.g =
+ 9: b
(58)
b, b 1, bi-19
bb-19
02. = b b1-1, b= 11
(59)
0 0-
(b.- 1,9
+ b
where Eq. 59 corresponds Eq. 15 Putting the control variate un = P4-1.9 . in Eq. 58
results in Eq. Note
that the shape parameters dt-1,9 and 0z-1, are cancelled out while deriving Eq. 58 reasons being that the shapes all equal to the same constant over the iterations, as later we will show in Eq. 64
16

---

Published in Transactions on Machine Learning Research (04/2025)
Similarly for 7/ (o a bg b( according to Eq. 53
log q og (/ = =B glnn
logp(wg) + log plga wes wg)
(60)
((gg(W))) + log p(g lu,, Wgyl logp(w3) + log p I9k. wa)
(61)
log ((ug) dgg. (6.9) = (an
(62)
= ((rersh)) [log + log p(0: 91, ug))
(63)
Again, by matching the coefficients of the natural parameters:
d/ dgg 7 7 0.57
(64)
b/,0 = 7 0.5T Mi. -/ + 11-1.9"
(65)
b/.9 = K,T + 0.57 + 12. - 2u + 97:)
(66)
In Eq. 64 dg and 0/.g equal the same constant over the iterations. As a result, ar,y and ar are also equal constants, according to Eq 56
Eq. 64 to Eq. 66 reveals that during the posterior updates of Wi and wg. choosing a larger value for T increases the emphasis on the information conveyed by Jt, consequently diminishing the relative influence of the priors. In VSGD, we have already made the prior strongth, Y a free paramcter, then we can replace T with a constant and still perlain the ability to adjust the strength between prior and new observation through values of 7.
That means
constant value, such 1, to replace T in Eq.
64
Eq.
65
and Eq. 66 which result in Eq. 16 is derived from the fact
Eq. 17 and Eq. 18 Eq. 19 and Eq. 20 are derived by simply following Eq.
56
that Eq. 04 is a constant, and interpolation between two equal constants results in the same constant.
A.2 Kernel Smoothing View on VSGD
If we look al VSGD from a kernel smoothing perspective, making use of ut serves as a cheap approximation to conditioning on all of the historical data. Namely, VSGD seeks to smooth the local gradient function in a similar way with the following kernel smoothing task.
There are two noisy observations of the gradient function 8 8, They are (0,, ue) and (0,, (t). We can define two Gaussian kernels for each of the noisy observations, governed by wz and w By denoting the Gaussian kernels Kд and 62, the smoothed gradient value at Or should be as follows:
(1)
R2(0,,, 02)
9t K2 (0,, + (2 0,) U4
Ki (0, 0.) (2( 0,) dc-
(67)
After laking the expectation step of SVI (expect out wy and wg), the above equation simplifies to Eq (15).
Note that uy in VSGD is introduced as a learned summary of the noisy vaservations before / and we only need smooth on u, as cheap approximation to smoothing on all of the historical data. Similar simplification technique is commonly seen in scalable Gaussian Process solutions, as discussed in Quinonero-Candela & Rasmussen (2005): Liu et al. (2020a). In these references, the learned summaries are called the 'inducing points'.
If we look al VSGD as an approximation to the local areas of the gradient function, then the precision term Wy can be seen as a smoothness measure the gradient surface, In VSGD we assume a global Wy. which is equivalent assuming that the smoothness of the whole gradient function governed by constant. This assumption is not necessarily idcal, but we still treat it a global parameter for simplicity and practical reasons.
17

---

Published in Transactions on Machine Learning Research (04/2025)
A.3 Treating Samples Separately in Mini-bach
Assuming the size of the mini-batch is M, denote the sample gradients {9r ji=I:M- According Hoffman
jet al. (2013), we would need to calculate Eq. 14 M times and average them in Eq. 17 and Eq. 18. Specifically,
would need to replace Eq. 14 Eq.
17
and Eq.
18 with:
b 1.0
bz 1.g
H4g =
1,9
9t
i=1: 1=1:M M
(68)
1,
bt-19
1
bl-19
E/ 7 0.5
(oh, + (ui
Me- 119)
b/g
(69)
M
EM Kg + 0.5
o + (M
- g (2)
b/.8
(70)
M
And introduce an additional step
( Σ i==1 Hr.g
Pt.g
(71)
M
before Eq. 23
Based on the equations presented, the computational cost grows linearly with the batch size M. Given the importance of efficiency in contemporary optimization techniques, above equations were omitted from the main paper. Nonetheless, is conceivable that the increased computational expense could be offset by enhanced convergence rates, hypothesis reserved for future investigation.
18

---

Published in Transactions on Machine Learning Research (04/2025)
B Experimental Details
B.1 Architecture
We have used open-source implementations of all three models.
VGG
We use VGG16 and VGGI with batch normalization. We add adaptive average pooling before the final linear layer to make the architecture suitable for both 32 x 32 and 64 x 64 images.
ConvMixer We use ConvMixe with 256 channels and set the depth to
ResNeXt ResNeXt
we set the cardinality (or number of convolution groups) and the depth 18 layers. The widen factor is set to 4, resulting in channels being [64, 256, 512, 1024]. We add adaptive average pooling before the final linear layer to make the architecture suitable for both 32 x 32 and 64 x 64 images.
B.2 Hyperparameters
The full set of hyperparameters used in all experiments is reported in Table
Table Hyperparameter values used in all the experiments.
ADAM
VSGD
(B1, 82)
(0.9, 0.999)
(Ni1 K2)
(0.9, 0.81)
7 Kg
1e-8
30
learning rate
{0.001, 0.005, 0.01, 0.02]
weight deacy momentum coefficient.
{0, 0.01}
{0.9, 0.99)
B.3 Data Augmentations
We use data augmentations during training to improve the gencralizability of the models. These augmentations were the same for the optimizers and only differ between datasets.
Table 4: Data augmentations.
CIFAR100
TINY IMAGENET-200
RandomCrop (32, padding=4) RandumiorrizontalF1ip
RandomCrop padding=4)
Randomforrizonta1FUl
Landomnffinneegggre Translate-(0.11 0.1), scale=(0.9, 1.1))
ColorJitter (breghttter((rigghtnesss.... 2, contrast=0.2. saturation=0.2)
Sgithub. Satthhacccccaaaacceeeeeeeeaaaaae /tttiaa.ww./ceessa///cceeeeeerre gttthbb.cc/ppttr7//heeltterrrrrr
19

---

Published in Transactions on Machine Learning Research (04/2025)
B.4 Training Loss
VGG16-bn
ConyMixer-2550 2566/8
ResNeXt-18
2.0
2.0
2.0
AdamW
Adamn
A
1.5
SGD (mom)
1.5
SGD (mom)
1.5
SGD (mom)
VSGD
VSGD
VSGD
1.0
1.0
1.0
0.5
0.5
0.5
0.0
0.0
0.0
10000
20000
30000
10000
20000
30000
10000
20000
30000
Iteration
Iteration
Iteration
Figure 4: Training loss CIFAR100 dataset.
VGG19-bn
ConvMixer-256/8
ResNcXt-18
Adamt
Adam
Adamn
SGD (mom)
4
SGD (mom)
SGD (mon)
VSGD
VSGD
VSGD
2
2
20000
40000
60000
20000
40000
G0000
20000
40000
60000
Iteration
Iteration
Iteration
Figure 5: Training loss on Timylmagenet2200 dataset.
75
50
25
AdamW VSGD
0
0
50000
100000
150000
200000
250000
Itcration
Figure Top-1 Accuracy on the Imagenet dataset trained with ResNet-50.
Table 5: Final Average train loss, over three random seeds.
VSGD (w/ L2)
VSGD (w/o L2)
ADAM (w/o L2)
ADAMW (w/ L2)
SGD
(w/ mom)
CIFAR100
VGG16
0.0016
0.0022
0.0284
0.0157
0.0018
CONvMIXER
0.0019
0.0022
0.0094
0.0094
0.0027
REsNEXT-18
0.0027
0.0015
0.0081
0.0088
0.0151
TINYIMAGENET- 200
VGG19
0.8783
0.1900
1.2528
1.4722
0.1954
MIXEMIXER
0.8823
0.8513
0.7347
0.6882
0.7581
REsNEXT-18
1.3233
1.2680
1.4096
1.4089
1.8158
20

---

Published in Transactions on Machine Learning Research (04/2025)
C Hyperparameter Sensitivity Analysis
In this Section, we investigate the shape parameter y's sensitivity to VSGD's performance. Higher implics a stronger weight on the prior while a lower implies a stronger weight on the current observation. We train VSGD on a wide range of 7 values on C10AR100 using the CONvMIXER architecture. The results can be and Table Inspecting Figure (Right) more closely, we observe that while there is statistically significant difference between accuracies obtained by different ys, they all perform relatively well long 7 is not too large.
observed in Figures
0.70
1,000
0.68
Train Accuracy
Test Actuuac
0.66
0.996
0.64
0.994
10-0
Figure 7: Analysis of the hyperparameter ssssssiiiiity the VSGD's performance in terms of accuracy.
Train
Test
4
Lo55
2
2
0
0
1.0
0.700
0.675
16:06
Sc.00
0.9
0.650
0.625
0.8
0.600
5000
10000
15000
20000
25000
30000
5000
10000
15000
20000
25000
30000
Figure 8: CIFAR100 learning curves trained with CONVMIXER architecture for different y values.
Accuracy
1e-9
67.75
5e-9
68.59
1c-8
69.03
5e-8
69.77
1e-7
69.71
lc-3
60.97
Table 6: Test accuracy on CIFAR100 trained with CONVMINER for different values.
D
Constant VSGD
In this Section, we offer a comparison between VSGD and CONSTANT VSGD. From the modeling point of view, VSGD has more free parameters that allow model to learn complicated distributions. While
21

---

Published in Transactions on Machine Learning Research (04/2025)
CONSTANT VSGD imposes constraint VSGD, which is less flexible but may lead to better out-of-sample performance. To investigate this, we ran additional experiments on CTFAR100 with CONVMIXER architecture. The final accuracy for VSGD was 69.04 = 0.4, while the accuracy for CONSTANT VSGD was 68.49 + 0.19. The learning curves can be observed in Figure
Train
Test
LO55
2
0
0
1.0
0.700
Optim
Accuracy
0.9
0.650
0.8
0.7
0.600
25000
30000
20000
25000
30000
Figure 9: Comparison of CONSTANT-VSGD and VSGD learning curves trained on CIFAR100 with CONVMIXER architecture. The final accuracy for VSGD is 69.04 0.4, while the accuracy for CONSTANT-
VSGD is 68.49 = 0.19.
22


Reviews:
[{"review_id": "7kDnQzfS9a", "reviewer": "Reviewer_DBpR", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Sorry for the formatting below, but Openreview is refusing to correctly render the equations. I give you directly the raw markdown (it works in vscode).\n\n```md\n### Strengths\n\n#### Interesting approach to apply Bayesian principles to gradient estimation\n\nThe paper introduces Variational Stochastic Gradient Descent (VSGD), an interesting optimization framework that incorporates Bayesian principles into gradient estimation. By treating true gradients as latent variables and noisy gradients as observed variables, VSGD models the dependencies between them using Gaussian distributions with learnable noise parameters. The approach leverages stochastic variational inference (SVI) to estimate the posterior distributions of the true gradients, allowing for more robust and noise-resilient updates. This Bayesian framework also facilitates the incorporation of prior knowledge about gradient noise variance, which can improve the efficiency and adaptability of the optimization process.\n\n#### Empirical results suggest good performance, at least on the CIFAR100 benchmark\n\nThe experimental results highlight VSGD’s good empirical performance, particularly on the CIFAR100 benchmark. The optimizer outperformed Adam and SGD in terms of top-1 test accuracy across various neural network architectures, including VGG16, ResNeXt, and ConvMixer. For instance, VSGD achieved an average improvement of 2.6% over Adam and a 3.2% improvement over SGD on CIFAR100. These results indicate that the proposed approach can consistently achieve lower generalization errors and maintain competitive or superior convergence rates. While the results are promising, it is worth noting that the gains on other datasets like TinyImagenet-200 and Imagenet-1k, though positive, are comparatively smaller, suggesting further evaluation are necessary to generalize the findings fully.\n\n### Weaknesses\n\n#### Modeling assumption is too convoluted and unnecessary\n\nWhile the Bayesian framework effectively models $g_t$ as the mean of the Gaussian distribution over $\\hat{g}_t$, the introduction of separate latent precision variables ($w_g$ and $w_{\\hat{g}}$) for systematic and observation noise adds unnecessary complexity. Specifically, the current approach relies on a factorized joint distribution (Equation 2):  \n\n$$\np(w_g, w_{\\hat{g}}, g_{1:T}, \\hat{g}_{1:T}; u_{1:T}) = p(w_g)p(w_{\\hat{g}}) \\prod_{t=1}^T p(g_t | w_g; u_t)p(\\hat{g}_t | g_t, w_{\\hat{g}}).\n$$  \n\nHowever, in cases where $p(g_t) = \\mathcal{N}(u_t, w_g^{-1})$ and $p(\\hat{g}_t | g_t) = \\mathcal{N}(g_t, w_{\\hat{g}}^{-1})$, the marginal distribution $p(\\hat{g}_t)$ can be computed directly:  \n\n$$\np(\\hat{g}_t) = \\mathcal{N}(u_t, w_g^{-1} + w_{\\hat{g}}^{-1}),\n$$  \n\navoiding the need to explicitly model the latent $g_t$. In the context of the paper, $\\hat{g}_t$ corresponds to the noisy gradient, $g_t$ to the true gradient, and $u_t$ to the prior mean of $g_t$. This raises some questions about the overall justification for the modeling choices.\n\nOn the other hand, this simplification relies on the assumption that the true gradient $g_t$ and the noisy gradient $\\hat{g}_t$ are Gaussian distributed. The explicit Gaussian assumptions _before introducing any variational framework_ for both $g_t$ and $\\hat{g}_t$ do not necessarily align with the practical distribution of gradients in real-world deep learning problems, which are often non-Gaussian and exhibit heavy-tailed behavior.\n\n#### The variational objective in Equation (8) is flawed\n\nThe evidence lower bound (ELBO) defined in Equation (8) aims to approximate the joint distribution over four sets of random variables $(w_g, w_{\\hat{g}}, g_{1:T}, \\hat{g}_{1:T})$ using a variational distribution that only includes three random variables $(w_g, w_{\\hat{g}}, g_{1:T})$. Unless I'm mistaken this discrepancy in dimensionality makes the objective technically invalid as an ELBO since it does not properly bound the marginal log-likelihood of the observed gradients.  \n\nHowever, the redundancy between $g_{1:T}$ and $\\hat{g}_{1:T}$, as noted above, might explain why things work out in the end. Specifically, $\\hat{g}_{1:T}$ can be expressed as noisy observations of $g_{1:T}$ via Equation (4), and the posterior distribution over $g_{1:T}$ implicitly determines the distribution of $\\hat{g}_{1:T}$.\n\nNonetheless, the inconsistency in the formulation highlights a lack of clarity in the probabilistic modeling, that needs to be fixed. \n\n#### Limited experimental campaign\n\nThe experimental evaluation primarily focuses only on image classification tasks. While these experiments demonstrate that VSGD can perform well against Adam and SGD in this context, the narrow scope of the evaluation limits the generalizability of the findings.\n\nOptimization methods like VSGD would be intended to be broadly applicable across a variety of machine learning tasks, yet the paper does not explore other important domains. For instance:\n\n- Generative Models: VSGD could be tested on Variational Autoencoders (VAEs) to assess its ability to optimize latent variable models where gradient dynamics can differ significantly from classification tasks.\n- Text Classification: Transformers and/or RNNs for text classification would provide insights into how VSGD performs in NLP tasks, where optimization challenges such as vanishing gradients and long-term dependencies are prevalent.\n- Reinforcement Learning: Testing on reinforcement learning tasks with a simple policy gradient method could further validate VSGD’s capacity to manage possibly non-stationary rewards and high-variance gradients.\n```", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "- The paper introduces Variational Stochastic Gradient Descent (VSGD), a novel probabilistic optimizer for deep neural networks that models gradients as random variables and uses stochastic variational inference to improve gradient estimation and handle noise effectively.\n- VSGD is shown to generalize and unify existing optimizers like Adam and Normalized-SGD by offering a probabilistic framework, adapting weights dynamically, and incorporating noise models.\n- Experimental results on image classification tasks using CIFAR100, TinyImagenet-200, and Imagenet-1k datasets demonstrate that VSGD achieves higher accuracy and faster convergence compared to Adam and SGD.\n- The computational overhead of VSGD is minimal, and the optimizer exhibits stable performance across different architectures without requiring extensive hyperparameter tuning.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Simplify the modeling framework or, at least, convince me by motivating the necessity of your assumptions more clearly.  \n\n2. Correct the variational objective in Equation (8) and ensure the ELBO derivation aligns with the actual model structure.  \n\n3. Expand the experimental campaign to include tasks beyond image classification, such as generative models (e.g., VAEs or GANs), text classification with small transformer models, and reinforcement learning tasks. These additions would better demonstrate VSGD’s generalizability and practical utility.  \n\n4. Include a sensitivity analysis for other key hyperparameters, and add an empirical study to analyze the impact of batch size on VSGD’s performance and investigate how hyperparameters such as learning rate, $K_g$, and $\\gamma$ should be adjusted for larger batch sizes. This would provide practical guidance for users in scaling the optimizer to different hardware configurations.  \n\n### Minor Changes\n\nIn the second paragraph of the introduction (2nd sentence), the paper implies that momentum combined with SGD was proposed by Sutskever et al. (2013). However, momentum was well-known and discussed in the neural network community long before, as early as Rumelhart et al. (1986) at least. Please revise this statement.", "broader_impact_concerns": "Nothing", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "apFj88Ax96", "reviewer": "Reviewer_FMuq", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "> Strengths\n\n- The method proposes an elegant probabilistic model over gradient, while remaining cognizant of the memory and computational constraints that are dealt with when optimizing deep neural network. As a result, the method needs to make some strong (independence) assumptions to remain tractable (see weaknesses), but the fact that it is very clear what the assumptions are should be regarded as positive.\n- The paper compares with other well-known optimizers, such as Adam, and gives an in-depth comparison under which settings individual update rules are equivalent or in what way they are different.\n- Although experiments on a wider range are needed to make more general claims, the model outperforms strong existing baselines such as Adam on a range of well-known vision models.\n\n> Weaknesses and questions\n\nQ1. Understanding the way in which this method is probabilistic.\n\nThe abstract states that the method may lead to improved \"modeling uncertainties\". Although it is not explicitly mentioned, the reader may think this means the method may lead to better uncertainty estimates in the output. If this is meant by the authors, I do not understand why this would be necessarily be the case since, although a probabilistic model over the gradients is employed, the method ultimately still targets the empirical risk (as described in Sec. 2) and thus no form of calibration. The paper also mentions that the method is \"orthogonal\" to [Khan et al., 2018; Osawa et al.] as VSGD assumes a Bayesian model over the gradients and not applied to a Bayesian inference problem. To my understanding, the proposed method is Bayesian about the gradients, but not the model itself. I am not completely satisfied by the \"orthogonal\" in this sentence and it would be helpful to better understand the distinction between these lines of work. Could the authors elaborate?\n\nQ2. Gamma parameter\n\nI am a bit concerned about the sensitive analysis of the gamma parameters. Figure 8 seems to suggest that a monotonically lower test loss for lower values of gamma. To get a complete picture, it would be nice to include train losses in Table 5 and Figure 7. Have you tried gamma=0, and/or is this case equivalent to an existing optimization algorithm?\n\nQ3. Train loss missing\n\nIt would be helpful to also provide final training losses (the thing we are optimizing!). To me, it is important to distinghuish possible regularizing effects from being a good optimizer. Could training losses be provided? (e.g. train losses in Table 1, Table 5, Figure 7)\n\nPossibly related to this, it seems that lower values of gamma always result in a better loss, but not necesarily test accuracy. Do the authors have an idea what could cause this?\n\nQ4. Memory\n\nIt seems that due to the strong independence assumptions, there is no need to perform any inverses (no correlations between weights or need to do filtering etc..). Table 2 is helpful to see that this indeed leads to negligible increase in runtime.\n\nIn addition to runtime, it would be helpful to provide an overview of total memory (e.g. as a function of parameters) that is required to maintain all (hyper)-parameters, as well as existing methods such as Adam that are discussed.\n\nExperimental validation\n\nThe experimental validation is limited to computer vision tasks. To assess whether the method could be generally better than optimizers such as Adam a wider set of architectures and tasks would be regarded. Yet, the paper is clear about this and does not seem to claim the method to be generally better than Adam/etc. Therefore, I think the experiments to be sufficient for acceptance in PMLR.\n\n\n> Conclusion\n\nThe work proposes an interesting novel method to improve optimization of deep neural networks while incorporating uncertainty over gradient information. The work demonstrates improved performance and relates the work to other existing methods.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "> Summary\n\nThe paper proposes to use a probabilistic model over the gradients using variational inference, named Variational Stochastic Gradient Descent (VSGD). This yields a general and novel way of optimizing deep neural networks. Relationships to other popular optimizers, such as Adam are discussed. The paper presents empirical evicence in favor of using VSGD.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See weaknesses and questions above.\n\nThe most critical requested change mentioned would be providing training losses.", "broader_impact_concerns": "No ethical implications.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "FqkH4s3OTq", "reviewer": "Reviewer_moMe", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths of the submission:\n- The main ideas are pretty clear, and the probabilistic framework does introduce a degree of flexibility which could be useful in some problems.\n- The effort to compare and contrast with other accelerated descent algorithms is appreciated, and strengthens the work.\n- I appreciated that the number of added hyperparameters remained small, making comparison with other algorithms reasonable.\n- The numerical experiments, while modest, still do tell a favorable story.\n\nWeaknesses:\n- In optimization I think many of the assumptions (especially about independence, for example of the gradient across coordinates) are not always realistic. I understand that this makes it tractable, and in some contexts it's useful, but I'm not sure if the Bayesian framework is really the \"right\" way to model an optimization process.\n- The notation (which doesn't differentiate between vectors and scalars) makes some parts of the paper hard to read and evaluate. I'm assuming the authors mean that the square/square root of a vector is meant in a pointwise sense? Depending on how this is intended some of the comparison with normalized descent may need to be modified.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors consider a probabilistic model for describing the relationship between noisy and true gradients during an optimization process. By introducing various priors and parametric approximations, they derive a family of new optimization algorithms. They compare this family of algorithms mathematically to several established accelerated descent algorithms, showing an equivalence with a version of normalized gradient descent in a degenerate prior limit, and highlighting both similarities and differences between the optimization rules utilized in ADAM and AMSGRAD. The authors then compare the performance of their algorithm with ADAM and SGD on some test data sets, showing favorable (modest) performance improvements without a significant increase in computation time.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Spell check the paper.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The three reviewers were all knowledgeable and the authors engaged in appropriate response discussions. While two reviewers were positive, one reviewer raised concerns about the modelling assumptions, the ELBO derivations, and the significance of results. The authors responded, but the issues were only partially concluded between the authors and reviewer. Looking at the issue, I am leaning on the positive reviewer's and authors side.
Decision: Accept (poster)