Bayesian Online Natural Gradient (BONG)
Authors: Matt Jones
Peter G. Chang
Kevin Patrick Murphy
Venue: NeurIPS 2024 poster (2024)
Keywords: online learning
Bayesian neural networks
variational inference
natural gradient descent
Primary area: probabilistic_methods
TL;DR: We improve on online variational Bayes using natural gradient descent on expected log-likelihood.

Abstract: We propose a novel approach to  sequential Bayesian inference based on variational Bayes (VB).
  The key insight is that,
  in the online setting,
  we do not need to add the KL term to regularize to the prior (which comes from the posterior at the previous timestep);
  instead we can optimize just the expected log-likelihood,
  performing a single step of natural gradient descent
  starting at the prior predictive.
  We prove this method
 recovers exact Bayesian inference 
  if the model is conjugate.
 We also  show how to compute an
 efficient deterministic 
 approximation to the VB objective,
  as well as our simplified objective,
 when the variational distribution is
 Gaussian or a sub-family, including the case of
 a diagonal plus low-rank
precision matrix.
We show empirically that our
method outperforms other online VB methods
  in the  non-conjugate setting,
  such as online learning for neural networks,
  especially when controlling for computational costs.

Bayesian Online Natural Gradient (BONG)
Matt Jones University of Colorado oecjone@coolrado.o edu
Peter Chang
Kevin Murphy Google DeepMind
MIT gyuyoungâmit. edu
kpmurphy@google.cm
Abstract
We propose novel approach to sequential Bayesian inference based on variational Bayes (VB). The key insight is that, in the online setting, we not need to add the KL term to regularize to the prior (which comes from the posterior at the previous timestep); instead we can optimize just the expected fog-likelihood, performing single step of natural gradient descent starting at the prior predictive. We prove this method recovers exact Bayesian inference if the model is conjugate. We also show how to compute an efficient deterministic approximation to the VB objective, as well as our simplified objective, when the variational distribution is Gaussian sub-family, including the of diagonal plus low-rank precision matrix. We show empirically that our method outperforms other online VB methods in the non-conjugate setting, such online learning for neural networks, especially when controlling for computational costs.
Introduction
Bayesian methods for neural network (NN) training aim to minimize the Kullback-Leibler divergence between true and estimated posterior distributions. This is equivalent to minimizing the variational loss (or negative ELBO)
C(bb) -gg log p (D + Duc(gelpo)
Here O are the network parameters, ?/2 are the variational parameters of the approximate posterior ge (0). D the training dataset, and (0) is the prior. The two terms the variational loss correspond to data fit and regularization the prior, the latter being analogous to regularizer r(0) = log Po(O in traditional point estimation methods like SGD.
An important set of approaches learns the variational parameters by gradient descent L( ab) BBlundell al. 2015). More recently Khan and colleagues [Khan et al. (018b) Khan and Rue 2023 Shen et 2024 have proposed using the natural gradient V where F& is the Fisher information matrix of the variational family evaluated 7Φ- Natural gradient descent (NGD) is often more efficient than vanilla GD because accounts for the intrinsic geometry of the variational family Amari 1998). Khan and Rue 2023) call this approach the "Bayesian Learning Rule" or BLR. Using various choices for the variational distribution, generalized losses replacing negative log- logllikelihood, and other approximations, they reproduce many standard optimization methods such as Adam, and derive new ones.
We study Bayesian NN optimization in online learning, where the data observed in sequence, D. = {(2+, ys) 1s)f== J. and the algorithm maintains an approximate posterior qu. (0,) ze p(0, Dt). which it updates each step. Fast updates (in terms of both computational speed and statistical efficiency) are critical for many online learning applications [Zhang et al. 2024). To allow for nonstationarity in the datastream, may change over time, as is standard for approaches based on state-space models and the extended
include time index O,, to represent that the parameters
Kalman filter (see e.g., [Sarkka and Svensson 2023) ). The belief state is updated recursively using
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

the prior 4 derived from the previous time step so that the variational loss becomes
[(eb)) = -Eg,- [log p zt. 0.)l + -EL (9 (wh)e-))
(2)
One option for this online learning problem is to apply NGD on L((br) at each time step, iterating until ub converges before consuming the next observation. Our first contribution is a proposal for skipping this inner loop by (a) performing a single natural gradient step with unit learning rate and (b) omitting the DxL, term in Eq. (2) so that learning is based only on expected loglikelihood:
a = 1//--- F 4b
Valui -o.-gy..., [log p (ye I 0x)]
(3)
These two modifications work together: instead of regularizing toward the prior explicitly using DKL (9,, () we do so implicitly by using - as the starting point of our single natural gradient step. This may appear as a heuristic but we prove in Proposition 4.1 that it yields exact Bayesian inference when O& and (3 x, 0) are conjugate and 9 is an exponential family with natural parameter sb. Thus our proposed update can be viewed as relaxation of the Bayesian update to the non-conjugate variational case. As is common in work on variational inference, we view the result for the conjugate case as a motivating foundation that ensures our method exact in certain simple settings. The experiments reported in Section and Appendix [ complement the theory by showing our method also works well in more general settings. We call Eq. (3) the Bayesian online natural gradient (BONG).
Our second contribution concerns ways of computing the expectation in Eqs. T (3). This is intractable for NNs. even for variational distributions that are easy to compute, since the likelihood takes the form p(y xt, 0,) = pf 0z)) with f(z,, 0t), representing the function computed by the network, is a complex, nonlinear function of 0,. Many previous approaches have approximated the expected loglikelihood by sampling methods which add variance and computation time depending
on the number of samples [Blundell et al., 2015 Shen et al.) 2024). We propose a deterministic, closed-form update that applies when the variational distribution is Gaussian (or sub-family) and the likelihood is an exponential family with natural parameter f( 0²) and mean parameter Mx,,, 0,) (e.g., for classification, f returns the vector of class logits, returns class probabilities, and h softmax(J)) (f). This update can be derived in two equivalent ways. First, we use local linear approximation of the network h (z; 02) == (O.) [Immer et al. 2021 and Gaussian approximation of the likelihood (2y (Ot). R,) Ollliveer 2018 Tronarp al. 2018). Under these
assumptions the expectation in Eq. (3) can be calculated malytically. Alternatively, we use a different linear approximation f(x,, 0.) (O.) and a delta approximation - (0,) x (eeic-(O.)) where Mzjr-1 = Egw,,.., [0c] is the prior mean, that the expectation in Eq. (3) is replaced by
plugin prediction. The
Gaussian approximation is previously known but the linear( -deeta approximation is new, and we prove in Proposition 4.2) that they yield the same update, which we call linearized BONG, or BONG-LIN. Finally, we discuss different ways of approximating the Hessian of the objective, which is needed for NGD.
Our BONG framework unifies several existing methods for Bayesian online learning, and it offers new algorithms based on alternative variational families or parameterizations. We define a large space of methods by combining 4 different update rules with 4 different ways of computing the relevant expected gradients and Hessians and 3 different variational families (Gaussians with full, diagonal, and Lingonal-plus-lor rank precision matrices). We conduct experiments syst systematically testing how these factors affect performance. We find support for all three principles of our approach- NGD, implicit regularization to the prior, and linearization- in terms of both statistical and computational efficiency. Code for our experiments is available at
2 Related work
Variational inference approximates the Bayesian posterior from within some suitable family in way that bypasses the normalization term [Zellner 1988 Dordan al. 1999]. A common choice for the variational family is Gaussian. For online learning, the exact update equations for Gaussian variational filtering are given by the RVGA method of [Lambert et al., 2021]. This update implicit but can be approximated by explicit RVGA update which we show arises as a special case of BONG. Most applications of Gaussian VI use a mean-field approximation defined by diagonal covariance, which scales linearly with model size. More expressive but still linear in the model size are methods

---

that express the covariance [Tomczak et al. 2020] or precision Mishhkin et al. 2018 Lambert et al. 2023 Chang et al., 2023] as a sum of diagonal and low rank matrices (DLR). In this paper, we consider variational families defined by full covariance, diagonal covariance, and DLR covariance.
For NNs and other complicated models, even the variational approximation can be intractable, so methods have been developed to approximately minimize the VI loss. Bayes by backprop (BBB) [Blundell al. 2015 learns a variational distribution on NN weights by iterated GD on the VI loss of Eq. They focus on mean-field Gaussian approximations but the approach also applies to other variational families. Here we adapt BBB to online learning to compare to our methods.
The Bayesian learning rule (BLR) replaces BBB's GD with NGD [Khan and Rue 2023). Several variants of BLR have been developed including VON and VOGN for mean field Gaussian prior (Khan 2018b) and SIANNG for a DIR Gaussian [Mishkin et al. 2018). BL. has also been used to derive versions of many classic optimizers including SGD, RMSprop and Adam (Khan al., 20T8a
et al.
Khan and Rue, 2023 Lin et al. 2024| Shen et al. 2024). Although BLR has been applied to online learning, we are particularly interested in Bayesian filtering including in nonstationary environments, where observations must be processed one at time and updates are based on the posterior from the previous step, often in conjunction with parameter dynamics. We therefore develop filtering versions of BLR to compare to BONG, some of which reduce to VON, VOGN and SLANG in the batch setting, while others are novel. We also note BLR a mature theory including several elever tricks we have not yet incorporated into our framework.
Khan and Rue 2023 observe that conjugate updating is equivalent one step of BLR with learning rate 1. This is similar to our Proposition 4.1 except that BLR retains the KI. term in the variational loss. BL. and BONG agree in this case because the gradient of the KI. is zero on BLR'S first iteration: V Dxu (g (gwl - = O. Therefore BONG can be seen as a special case of BLR with one update step per observation and learning rate Our contribution is to recognize that doing a single update step allows the KL term to be dropped entirely, yielding a substantially simpler algorithm which our experiments show also performs better.
While BLR allows alternative losses in place of the NLL in Eq. (2. we can also replacing the KL term with other divergences (Knoblauch et al. 2022 Our approach fits within that "generalized VB" framework in that it drops the divergence altogether. Our approach of implicitly regularizing to the prior using a single NGD step is also similar to the implicit MAP filter of [Bencomo et al., 2023] which performs truncated GD from the prior mode. The principal difference is they perform GD on model parameters (0,) while we do NGD the variational parameters (14,). Thus BONG maintains a full prior and posterior while IMAP is more concerned with how the choice of optimizer can substitute for explicit tracking of covariance.
We show two other ways to derive the BONG update in Appendix D one of which is to replace the expected NLL in Eq. (2) with a linear approximation and solve the resulting equation exactly. Several past works have taken this approach, arriving at updates similar to ours. ACéelee-Abdellatf et al. 2019 study streaming variational Bayes and propose solving Eq. 2 with linearized expected NLL. When the variational family is an exponential family their update becomes NGD [Khan and
Lin 2017] and matches the BONG update. Hoeven al. 2018 show how mirror descent can be derived as a special case of Exponential Weights [Littlestone and Warmuth| 1994). which is closely
related to Bayesian updating. The resulting algorithm is similar to BONG and follows from linearizing
the NLL instead of expected NLL, with an additional delta assumption at the prior mean. Lyu and
Tsang | study relaxed block-box optimization where the objective is arg mineo Ex-qu for some target function f. They use mirror descent formulation with linearized expected loss and KL regularizer and show the resulting update is NGD on expected loss, formally equivalent to our BONG update. From the perspective of this prior work, our contribution is to express the BONG update simply as NGD on the expected NLL., motivated by replacing the KI. with implicit regularization, and to show how this yields variety of known and novel algorithms for Bayesian filtering.
EKF applications to NNs apply Bayesian filtering using a local linear approximation of the network, leading simple closed form updates (Singhal and Wu. 1989 Puskorius and Feldkamp) 1991).
The classic EKF assumes a Gaussian observation distribution but it has been extended to other
exponential families (e.g. for classification) by matching the mean and covariance in what we call
the conditional moments EKF (CM-EKF) [Ollivier] 2018 Tronarp et al. 2018). Applying a KL
projection diagonal covariance yields the variational diagonal EKF (VD-EKF) [Chang et 2022
3

---

Alternatively, projecting to diagonal plus low rank precision using SVD gives LO-FI [Chang et al. 2023]. We derive all these methods as special cases of BONG-LIN. Further developments this direction include the method of [Titsias et al., 2024] which does Bayesian filtering on only the final weight layer, and WoL.F [Duran-Martin et al. 2024 which achieves robustness outliers through data-dependent weighting of the loglikelihood.
3 Background
We study online supervised learning where the agent receives input 2t RD and observation y: E EC on each time step, which it aims to model with function f((O,) = f(x,, 0,) such as a NN with weights 0: E R.P. The predictions for y are given by some observation distribution (O.)). For example, If may compute the mean for regression or the class logits for classification.
We work in a Bayesian framework where the agent maintains an approximate posterior distribution over O, after observing data D. {( )=))==)- The filtering posterior Toh (0.) R p (0, (Dr) is approximated within some parametric family indexed by the variational parameter ab;- We allow for nonstationarity by assuming 0 changes over time according some dynamic model (.-)) By pushing the posterior from step t through the dynamics we obtain a prior for step t given by - (0,) R p(0, (D.-1). For example suppose the variational posterior from the previous step is Gaussian, 9:h-- (0,-))
1 ) and the dynamics model is an Ornstein Omstein-UUllenbeck process, as proposed in prior work (Kurle et al. 2020 Titsias et al. 2024 to handle non-stationarity, 140 Qt). where Q: = (1 > Eo the covariance of the noise process, 0 S S 1 is the degree of drift, and p(0n) = =N(uo, En) is the prior. In this case, the parameters of the prior predictive distribution are M1jt-1 = 7//4++ 7%
i.e. the dynamics model has the form 0. ~ .>, (xx00,111
7) (1-7%
and Enjr-1 = % Σ₁- + Q.. In general the predict step may require approximation to stay in the variational family (e.g., if the dynamics
nonlinear). In this paper, from - to 2ht upon observing (It, Mt), so for simplicity we assume constant (static) parameters, ie., p(0, 0-1) = 800 0 (equivalently 7/- = 1), so 1 = t however, our method can trivially handle non-stationary parameters.
focus is the update step
Variational inference seeks an approximate posterior that minimizes the KL divergence from the exact Bayesian update from the prior. In the online setting this becomes
5 arg min DKL. (1+ (0.)
p(v.)fe(0.)) (
arg min Le(bb)
(4)
where L, is the online VI loss defined in Eq. 2 and the normalization term Z, (which depends on æx)) drops out an additive constant. Our goal is an efficient approximate solution this variational optimization problem.
We will sometimes assume the variational posterior 96 is exponential family distribution with natural parameter 2b so that 90, (0,) = exp (T).T(0,) - 2 ((b))) with log pog-partiiinn function & and sufficient statisties T T(,)). Assuming & is strictly convex (which holds in the cases study) there a bijection between 2b, and the dual (or expectation) parameter Pt = Eo.- [T)).)] Classical
thermodynamic identities imply that the Fisher information matrix has the form F Fs,. 0p:/ Prpb.
This has important implications for NGD exponential families [Khan and Rue 2023 because implies that for any function defined on the variational parameter space the natural gradient wrt
natural parameters 4/2 the regular gradient wrt the dual parameters Pt, i.e., F vb. V Veree f= V.ll.
Methods
We propose to approximate the variational optimization problem in Eq. using the BONG update in Eq. (3). When qu an exponential family, the fact that the natural gradient wrt the natural parameters tb, is the regular gradient wrt the dual parameters P: implies an equivalent mirror descent form (see Appendix [ for further analysis of BONG from the MD perspective):
a/- 1/-1-1 + Peit-1 F
Eo Mogg((y)an (x, 0,)
(5)
This is NGD with unit learning rate on the variational loss in Eq. ② but ignoring the DRI. (gpp (gpllgpste term. In this section we first prove this method is optimal when the model conjugate and then describe extensions to more complex cases of practical interest.

---

4.1
Conjugate case
Our approach is motivated by the following result which states that BONG matches exact Bayesian inference when the variational distribution and the likelihood are conjugate exponential families: Proposition 4.1. Let the observation distribution (likelihood) be an exponential family with natural parameter Or (where Ti(y:) Yt is the sufficient statistics for the likelihood and A(O,) is the log-partition function)
Py O.) = exp (0) y: -((,)) - -(y)))
(6)
and let the prior be the conjugate exponential family
Tou (0.) = exp (())) - ((b1(111)
with T(0)) - A(0)). Then the exact Bayesian update agrees with Eq. (5).
The proof is in Appendix Writing the natural parameters of the prior as d = X show the Bayesian update and BONG both yield X Xxjt-1 3/ and in Mije-1 + L. Intuitively, we are just accumulating a sum of the observed sufficient statistics, and counter of the sample size (number of observations seen far).
4.2 Variational case
In practical settings the conjugacy assumption of Proposition 4.1 will only approximate the Bayesian update. In this paper we restrict to Gaussian variational families. We refer the unrestricted case as FC (full covariance), defined by the variational distribution
will not be met, so Eqs. (3) and (5)
Q (0:) = N (0, Mej-11 Eq(-1)
where Σnjt-1 can be any positive semi-definite (PSD) matrix. The natural and dual parameters
are v/- = E μ. (recE-1)) and p (H,, vec ( + E)). Appendix E.1.1 shows that Eq. 5 translated back to (ju, Σ) gives the following BONG update for the FC case:
μr Mije-1 E, Bo
Ve. log ((:)
(9)
9
Σ E 1
EB
Ve
(o
(10)
G, which matches the explicit update in the RVGA method Lambert et al. 2021
4.3 Monte Carlo approximation
The integrals over the prior We in Eqs. approximated. One option is Monte Carlo, what we call BONG-MC. Given M independent
(9) and (00) are generally intractable and must be
samples 0m ~ 1 we estimate the expected gradient 9: E0 Vo, log (()().((((.)))
and expected Hessian G. = Eg Eor-
[V. log (yelfe (0.)) as the empirical means
M
gº
Σ
(m)
(m)
V o, - aims log
(11)
M
m=1
G
MC-HESS
(m)
G (m) v2 o, - p(mi log pl ((
(12)
M
We G HESS only for small models. Otherwise empirical Fisher (Section 4.5
4.4 Lincarized BONG
As an alternative BONG-MC, propose linear approximation we call BONG-LIN that yields a deterministic and closed-form update. Assume the likelihood is an exponential family as in Proposition 4.1 but with natural parameter predicted by some function fO. = (x,, 0.):
p Ft 0,) = exp (f. ye A(f.(e)) - b(y:))
(13)
5

---

We also define the dual (moment) parameter of the likelihood as h,)0,) = ly f((0.). In a NN, ft and h, are related by the final response layer. For example in classification ft and h, give the class logits and probabilities, with hig (0.) = (oftmax(f{(O.)), with Yt being the one-hot encoding.
We now define two methods for approximating the expected gradient Jt and expected Hessian Gr, based on linearizing the predictive model at the prior mean Maje-1 in terms of either ft or (2)),, and then prove their equivalence.
The (acarsi) approximation Ollivier 2018
Tronarp et al. 2018 linearizes h4 (0,)
The (0.) U, (0,
14j(1)
(14)
y: h4 (H
(15)
The
H. 00, 10.0010
(16)
and approximates the likelihood by Gaussian with variance based at p
p. (v))0)) = (, (0:), Rt), R₂ = V [y: = [xj:-1]]
(17)
The (nnear()-dehee approximation linearizes f((),) and maintains the original exponential family likelihood distribution in Eq. (13)
(..)) = fe( + F. peje-1)
(18)
F. = Of
(19)
00: 10.=111-11
I.D (3) x exp (fi(0:) - ( A(f((e))) - b(yr)) P
(20)
also uses a plug-in approximation that replaces 4e (0.) with a point mass 6 (0,) so that the expected gradient and Hessian approximated by their values at the prior mean, i.e.,
Vol=mue log 11.15 (x)) and V V log. log 1. ():)))), rather than being sampled.
Proposition 4.2. Under Gaussian variational distribution, the limear(h)- Gaaussia and linear( J)- delta approximations yield the same values for the expected gradient and Hessian
geth LIN H. (y/ -9))
(21)
G
LEN-
H R,
H,
(22)
See Appendix Clor the proof. The main idea for the nN part is that the Gaussian assumptions make the gradient linear in 0g so the expected gradient equals the gradicnt at the mean. The main idea for the G HESS part is that eliminating the Hessian of the NN requires different linearizing assumptions for the Gaussian and delta approximations, and the remaining nonlinear terms (from the log-like log-likelihood in Eq. (13) agree because of the property of exponential families that the Hessian of the log- oog-partitinn A equals the conditional variance R₁-
Applying Proposition 4. 2 to Eqs. 9 and 10 gives the BONG-LIN update for FC Gaussian prior
|r: = P Ke(y: - (:)
(23)
E; Emj--1 K,H,E,ft-
(24)
K, H (R, + H. H!)
(25)
where K, is the Kalman gain matrix (see Appendix E.1.2 This matches the CM-EKF 2018 Ollivier 2018).
Tronarp et al.
4.5 Empirical Fisher
Name
Eqs.
The methods in Sections and 4.4 require explicitly computing the Hessian of the loss (MC-HESS) or the Jacobian of the network (LIN-HESS). These are too expensive for large models or high-dimensional observations. Instead we can use an empirical Fisher approximation that replaces the Hessian with the outer product of the gradient (see e.g, [Martens] 2020)..
MC-HESS IN-HESS
MC-EF LIN-EF
26
28
29
Table 1: The 4 Hessian approximations.

---

For the MC-EF variant, we make the following approximation:
GN
G (RL M
1
(1:M0 G(:MM)
(26)
where G ((2:M) = 19t 9 ( is the P x M matrix of gradients from the MC samples.
We can also consider similar approach for the LIN-EF variant that is Jacobian-free and sampling- free. Note that if Dt were the true value of E 2 (i.c., if the model were correct) then we would have -.INN-ISS This suggests using
E (1y+ (v: Ú.) R2, implying E [g] IN (gl [
g Vo.=mie-1
[- (3/ he h,,),)) RT " - hh(0,))]
(27)
(O))(O))
R h,((4t)-11)) gh
(28)
00, 0:=MI1-1
G
GLIN-EF
- IN (g!
(29)
where Eq. 29 is the EF approximation to Eq.
more accurate EF approximation is possible by sampling virtual observations ÿr from p(-1 p( (m))
p(- ((()e=)) and using them for the gradients in Eq. (26 Eq. (29) (respectively) [Martens]
2020
Kunstner et al. 2020]. However, in our experiments we use the actual observations which IS
faster and follows previous work (e.g., [Khan al. 2018b]).
4.6 Update rules
Name
LOSS
Update
In addition to the four ways of approximating the expected Hessian (summarized in T.. we also consider four vari- ants of BONG, based on what kind of loss we optimize and what kind of update we perform, as we describe below. See Table 2 for summary.
BONG
E
NLL
NGD(L = 1)
BOG
E
NNL
GD(T = 1)
VI
NGD(T 2
BBR
GD(T 21)
Table 2: The update algorithms.
BONG (Bayesian online natural gradient) performs one step of NGD on the expected log- logelioid We set learning rate to = since this is optimal for conjugate models. The update (for an exponential variational family) is as Eq. (5:
J w + V Pajt-1 Ea, Ea [log p(yll(y...... 0c)]
(30)
BOG (Bayesian online gradient) performs one step of GD (instead of NGD) on the expected log- likelihood. We include learning rate because GD does not have the ceale-invrriance of NGD:
tb w/b + ar V Eo, [log (1g)(()))((00)))
(31)
BLR (Bayesian learning rule, Khan and Rue 2023 uses NGD (like BONG) but optimizes the VI loss using multiple iterations, instead of optimizing the expected NLL with a single step. When modified to the online setting, BLR starts an inner loop at each time step with ahr,o = T6pp-1 and iterates
4/i.i = th1,i-1 + ar F.F -e Fe
V (Bo,
(lgg(( - DxL. (ox.(9.... 19.....) ( (32)
For exponential variational family this can be written in mirror descent form
1/, = 1 + ar V pu.-1 Eg [log p(y (e(())) - DEL. (1wr, 140m)
(33)
BBB (Bayes By Backprop, Blundell et al. 2015 is like BLR but uses GD instead of NGD. When adapted online learning, starts each time step at 4/1.0 = d/ and iterates with GD:
Vi,i = 1 + ar Ven. (Eo.nqo..., (ogp((el..(())) DEL (goi. (lw)))))
(34)
7

---

Variational Family
Update
Hessian
Full
Diag
DLR
BONG MC-EF
O(MP))
[RVGA]
O(MP)
-((R+M) - M)
BLR
MC-EF
O
O(IMP) [VON] O(MP)
O (I(R+MM)PP) P) [SLANG]
BOG
MC-EF
O(Pa)
O(RMP)
BBB
MC-EF
O
O(IMP) [BBB]
O M)P)
BONG LIN-HESS O(CP²)
[CM-EKF]
O(C*P) [VD-EKF]
((P+
BLR LIN-HESS
((1P))
P(
0 C)
BOG LIN-HESS
O(PS)
O(C*P)
O(C(C+ R)P)
BBB LIN-HESS
O(TP²)
O(TC²P)
O
BONG
LIN-EF
O(P4)
O(P)
O(R*P)
BLR
LIN-EF
O(IP²)
O(IP)
O(TRPP)
BOG
LIN-EF
O(Pa)
O(P)
O(RP)
BBB LIN-EF
O(IP))
O(IP)
O(RRP)
Table 3: Time complexity of the algorithms. P: params, C: observation dim, M: MC samples, I: iterations, R: DLR rank. We assume P > {I, R.) so display only the terms of leading order in P. Time complexities for MC-HESS algorithms (not shown) are always at least great as for the corresponding MC-EF. Full (full covariance) and Diag (diagonal covariance) columns indicate natural parameters; corresponding algorithms using moment parameters have the same complexities except BOG-FC_MON which is (MM for MC-EF, O(CP²) for LIN-HESS, and O(P2) for LIN-EF. [Names] correspond to the following existing methods (or variants thereof) in the literature: RVGA: Lambert et al. [021] (explicit update version); VON: [Khan et al. 2018b] (modified for online);
SLANG: [Mishkin et al. 2018 (modified for online); BBB: [Blundell al. 2015 (modified for online
and uses moment parameters); CM-EKF: [Ollivier 2018
Tronarp et al.
2018); VD-EKF:
[Chang
al. 2022 LO-FI:
(Chang al. 2023].
4.7 Variational families and their parameterizations
We investigate five variational families for the posterior distribution: (1) FC Gaussian using natural parameters //- = (Σ μe. - "), (2) FC Gaussian using central moment parameters ab- = (c. E), (3) diagonal Gaussian using natural parameters tb (a- μ4, o -2) (using elementwise exponents and products), (4) diagonal Gaussian using central moment parameters rb (H, o²), and (5) DLR Gaussian with parameters v/2 = (1, Y, W) and precision Σ-1 = + WWT where E RR diagonal and W E Px R with R & P. The moment parameterizations
included to test the importance of using natural parameters per Proposition[4. [. The diagonal family allows learning of large models because it scales linearly in the model size P. DLR also scales linearly but is more expressive than diagonal, maintaining some of the correlation information between parameters that is lost in the mean field (diagonal) approximation [Lambert et al. 2023 Mishkin et al., 2018 Chang et al. 2023]
Optimizing the BONG objective wrt (n, Y, W) using NGD methods is challenging because the Fisher information matrix in this parameterization cannot be efficiently inverted. Instead we first derive the update wrt the FC natural parameters (leveraging the fact that the prior - DLR to make this efficient), and then use SVD to project the posterior precision back to low-rank form, following our prior LO-FI work [Chang et al. 2023 However, if we omit the Fisher preconditioning matrix and use GD as BOG and BBB, we can directly optimize the objective wrt (H, Y, W) (see Appendix E.5
4.8 Overall space of methods
Crossing the four algorithms in Table the four methods of approximating the Hessian in Table and the five variational families yields 80 algorithms. Table shows the 36 based on the three tractable Hessian approximations and the three variational families that use natural parameters. Update equations for all the algorithms are derived in Appendix Pscudocode is given in Appendix
3hows

---

5 Experiments
This section presents our primary experimental results. These are based on MNIST (D = 784, Niran 60k, New 10k, C == 10 classes) [LeCun et al., 20T0]. See Appendix for more details on these experiments, and more results on MNIST and other datasets. We focus on training on prefix of the first T 2000 examples from each dataset, since our main interest is in online learning from potentially nonstationary distributions, where rapid adaptation of model in response to a small number of new data points critical.
Our primary evaluation objective is the negative log predictive density (NLPD) of the test set as function of the number of training points observed so far"] It defined as NLPD, Num DiED log [J (x,, 0.)) (O.)dee.. We approximate this integral in two main ways: (1) using Monte Carlo sampling or (2) using plugin approximation, where we replace the posterior qv. (6,) with delta function centered at the mean, 6(0; pe).
I
For methods that require learning rate (i.e., all methods except BONG), we optimize it wrt mid-way or final performance on a holdout validation set, using Bayesian optimization on NLL. All methods require specifying the prior belief state, p(0g) N (mo: Σg ra We optimize over of and sample pág from a standard NN initializer. As Hessian approximations, we use MC-EF with M 100 samples, as well as the deterministic approximations LIN-HESS and LIN-EF. (In the appendix, also study MC-IIESS but find that it works very poorly, even with M 1000.)
In Fig. T we compare the 4 main algorithms using DLR family with rank R 10. We apply these to a CNN with two convolutional layers (each with 16 features and a (5,5) kernel), followed by two linear layers (one with 64 features and the final one with 10 features), for total of 57,722 parameters. Shaded regions in these and all other plots indicate +1 SE, based on independent trials randomly varying in the prior mean Moo data ordering, and MC sampling. From this figure (and additional
results in Appendix B) we conclude
Linearization helps: LIN-HESS and LIN-EF both outperform the MC variants.
NGD helps: BONG outperforms BOG.
Implicit regularization helps: BONG outperforms BLR.
The LIN-HESS approximation outperforms LIN-EF, at least for BONG.
BBB generally does poorly.
The BONG posterior predictive (using LIN-HESS) is slightly better calibrated than BOG, and both
are much better than BLR and BBB, at least for small sample sizes, shown in Fig. 7
The plugin posterior predictive is similar to the Lin-MC predictive (see Footnote 2 and both are generally much better than the simple MC predictive, as shown in Fig.
In Fig. we compare BONG using different variational families, and conclude
DLR-10 outperforms DLR-1, which is similar to diagonal (except when using BONG-LIN-EF, where DLR-1 is worse than diagonal). Also, find (in results not reported here) that rank 5-10 often gives results as good as FC, but is much cheaper.
Both natural and moment parameterizations for the diagonal representation perform comparably to each other, although with LIN-EF, the moment parameterization can be numerically unstable
Finally, in Fig. 3 report the runtimes from these experiments and conclude
One-step methods (BONG and BOG) are faster than iterative methods (BLR and BBB), as expected. Linearized methods HEIS and LIN-EF) are faster than MC methods (MC-EF).
WWe assume the training and test sets are drawn from the same static distribution. Alternatively, if there is only one stream of data coming from a potential notstationary source, we can use the prequential or one-step-ahead log predictive density Gama et al. 20133. We leave studying the non-stationary case to future work.
"That is, we compute
100 posterior samples 0, p(0, (Dimin) and then use p(y|a)
ES p(z (x, 0.). ). For efficiently sampling from DLR Gaussian we follow Mishkin al. (2018). Following Immer et al. 2021b). we find it better to linearize the model (i.e., replacing h( O,) with h(2,, 0.) defined in Eq. (14) before pushing posterior samples through. We call this the MC approximation.
TTis is likely due to the fact that in the moment update in Eq. (249). we add diag to the variance, whereas in the natural update in Eq. (213), we subtract diag (G.) from the precision; the latter is more stable since G negative semi-definite for all Hessian approximations.

---

EF methods (LIN-EF) are bit faster than methods that compute the Hessian exactly (LIN-HESS) HESS). especially for diagonal family. (This speedup is larger when the output dimensionality C is big.)
Enn cive poitt Cav did
aopaatooeee lecoall
bangy
Figure 1: Performance on MNIST using Lin-MC posterior predictive, where the posterior is computed using BONG, BOG, BBB and BLR and the tractable Hessian approximations with DLR- 10 variational family.
hong deskin DI
1411 bing sec3l1
Figure Performance on MNIST using Lin-MC posterior predictive, where the posterior is computed using BONG with different variational families, namely diagonal (natural and moment), DLR-I, DLR- 10.
6 Conclusions, limitations and future work
Our experiment results show benefits of BONG's three main principles: NGD, implicit regularization to the prior, and linearization. The clear winner across datasets and variational families is BONG-LIN- HESS, which embodies all three principles. BLR-LIN-HESS nearly matches its performance but is much slower. Several of the best-performing algorithms previously known (notably CM-EKF and LO-FI) but we explain these results within systematic theory that also offers new methods (including BLR-L.N-HESS).
BONG motivated by Proposition 4.1 which applics only in the idcalized setting of a conjugate prior. Nevertheless find it performs well in non-conjugate settings. On the other hand our experiments are based on relatively small models and datasets. It will be important to test how our methods scale up, especially using the promising DLR representation.
10

---

Acknowledgments and Disclosure of Funding
Thanks to Gerardo Durån-Martin and Alex Shestopaloff for helpful input. MJ was supported by NSF grant 2020-906.
References
S Amari. Natural gradient works efficiently in learning. Neural Comput., 10(2):251-776, 1998. URL htttp:/dae doi. org/10../... 11622/00.11/11111111 116200899766983000077777
Gianluca M Bencomo, Jake C Snell, and Thomas L Griffiths. Implicit maximum a posteriori filtering via adaptive optimization. arXiv preprint arXiv:2311.10580, 2023.
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wienstra. Weight uncertainty in neural networks. In ICML, 2015. Utt.httpp//arxiivorr////s/ wwt.://rrraabs/55...oooo 05424
Georges Bonnet. Transformations des signaux aléatoires a travers les systemes non linéaires sans mémoire. In Annales des Télécommunications, volume 19, pages 203-220, Springer, 1964.
Peter Chang, Kevin Patrick Murphy, and Matt Jones. On diagonal approximations the ex- tended kalman filter for online training of bayesian neural networks. In Continual Lifelong Learning Workshop at ACMI. 2022, December 2022. BtLLhtttss///pperreeet http:/pprreviee..ee/fffruur net/forum? id=asgeEt25kk
Peter G Chang, Gerardo Durán-Martín, Alexander Y Shestopaloff, Matt Jones, and Kevin Murphy, Low-rank extended Kalman filtering for online learning of neural networks from streaming data. In COLLAS, May 2023. URL httptt/////xxxv http://aaaiiv..orrra/////2//0/ 19535
Badr-Eddine
Pierre Alquier, and Mohammad Emtiyaz Khan. A generalization bound for online variational inference. In Asian conference machine learning, pages 662-677.
PMLR, 2019.
Gerardo Duran-Mann Martin, Matias Altamirano, Alexander Y Shestopaloff, Leandro Jeremias Knoblauch, Matt Jones, François-Xavis Briol, and Kevin Murphy. Outlier-robust kalman filtering through generalised bayes. arXiv preprint 240i::20000006-46 2024.
João Gama, Raquel Sebastião, and Pedro Pereira Rodrigues. On evaluating stream learning algorithms. MLJ, 90(3):317-346, March 2013. URL
www/ homm/mxfk44w
Nathan Halko, Per-Gunnar Martinsson, and Joel A Tropp. Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions. SIAM review, 53(2): 217-288, 2011.
Dirk Hoeven, Tim Erven, and Wojeiech Kotlowski. The many faces of exponential weights online learning. In Conference On Learning Theory, pages 2067-2092. PMLR, 2018.
Michael F Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Statistics-Simularte and Computation, 1853):0099-107, 1989.
Alexander Immer, Maciej Korzepa, and Matthias Bauer. Improving predictions of bayesian neural via local linearization. In Arindam Banerjee and Kenji Fukumizu, editors. AISTATS, volume 130 of Proceedings of Machine Learning Research, pages 703-711. PMLR, R, 2021a. URI.httpa
//proceedings.co .mlr ppwps//13//eeieppppppppe htwp:////..mmmmmmmmmmm..mmm html
Alexander Immer, Macicj Korzepa, and Matthias Bauer, Improving predictions of bayesian neural nets via local lincarization. In Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, pages 703-711. PMLR, 2021b. URL wtt.ttts////ooeeeeiiiiiei mlr press/ 10.00//mm222a...m html
Michael Jordan, Zoubin Ghahramani, Tommi Jaakkola, and Lawrence Saul. An introduction to variational methods for graphical models. Machine learning, 37:183-233, 1999.
11

---

Mohammad Khan and Wu Lin. senjjnaatecomputaatie variational inference: Converting variational inference in non-conjugate models inferences in conjugate models. In Artificial Intelligence and Statistics, pages 878 887. PMLR, 2017.
Mohammad Khan, Didrik Nielsen, Voot Tangkaratt, Wu Lin, Yarin Gal, and Akash Srivastava. Fast and scalable bayesian deep learning by weight-perturbation in adam. International conference machine learning. pages 2611-2666 PMLR, 2018a.
Mohammad Emtiyaz Khan and Håvard Rue. The bayesian learning rule. J. Mach. Learn. Res., 2023. URL http:://axiii org/ abs/ 2107 .04662
Mohammad Emtiyaz Khan, Didrik Nielsen, Voot Tangkaratt, Wu Lin, Yarin Gal, and Akash Srivastava. Fast and scalable bayesian deep learning by Weight- Peight-Perturbation adam. In ICML, 2018b. URL
http::/arxiv..om org/ abs/ 1806. .4854
Jeremias Knoblauch, Jack Jewson, and Theodoros Damoulas. An optimization-centre view bayes' rule: Reviewing and generalizing variational inference. Journal of Machine Learning Research, 23 (132):1-109, 2022.
Frederik Kunstner, Lukas Balles, and Philipp Hennig. Limitations of the empirical fisher approxima- tion for natural gradient descent, 2020.
Richard Kurle, Botond Cseke, Alexej Klushyn, Patrick van der Smagt, and Stephan Günnemann. Continual learming with bayesian neural networks for non-stationary data. ICLR, March 2020.
URL https::/ wttp://wwww...oe..oooooooosda//s
Marc Lambert, Silvère Bonnabel, and Francis Bach. The recursive variational gaussian approxi- mation (R-V (GA). Stal. Comput., 32(1):10, December 2021. URL https://hal inriaa.iirra.... fr/
hal-03086627/doccee document)
Marc Lambert, Silvère Bonnabel, and Francis Bach. The limited-memory recursive variational gaussian approximation (I-rvga). Statistics and Computing, 33(3):70, 2023.
Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. 2, 2010.
Available: http://yann.lecu
Wu Lin, Felix Dangel, Runa Eschenhagen, Juhan Bae, Richard Turner, and Alireza Makhzani. Can we remove the square-root in adaptive gradient methods? second-order perspective. arXiv preprint ar 2xr.i:2240..00 240.0499000 2024.
Nick Littlestone and Manfred K Warmuth. The weighted majority algorithm. Information and computation, 108(2):212- 261, 1994.
Yueming Lyu and Ivor W Tsang. Black-box optimizer with stochastic implicit natural gradient. In Machine Learning and Knowledge Discovery in Databases. Research Track: European Conference, ECML. PKDD) 2021, Bilbuo, Spain, September 13-17, 2021, Proceedings, Part III 21, pages 217-232. Springer, 2021.
James Martens. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 21(146):1-76 2020.
Aaron Mishkin, Frederik Kunstner, Didrik Nielsen, Mark Schmidt, and Mohammad Emtiyaz Khan. SLANG: Fast structured covariance approximations for bayesian deep learning with natural gradient. NIPS, pages 6245-6255. Curran Associates, Inc., 2018.
Yann Ollivier. Online natural gradient as a kalman filter. Electron. J. Stat., 12(2):2930-2911, 2018. URL https::/pprojecteuclldd http:://ww..........d///////////
Robert Price. A useful theorem for nonlinear devices having gaussian inputs. IRE Transactions on Information Theory, 4(2):69-72, 1958.
G V Puskorius and L. A Feldkamp. Decoupled extended kalman filter training of feedforward layered networks. In International Joint Conference on Neural Networks, volume pages 771-777 vol.1, 1991. URL http: //xx doi. org/10. 10. 1111/11909// 1991. 155276
12

---

Carl Edward Rasmussen and Christopher K. I. Williams. Gaussian Processes for Machine Learning. MIT Press, 2006,
Simo Sarkka and Lennart Svensson. Bayesian Filtering and Smoothing (2nd edition). Cambridge University Press, 2023.
Yuesong Shen, Nico Daheim, Bai Cong, Peter Nickl, Gian Maria Marconi, Clement Bazan, Rio Yokota, Iryna Gurevych, Daniel Cremers, Mohammad Emtiyaz Khan, and Thomas Möllenhoff. Variational learning is effective for large deep networks. arXiv preprint 240. 2024.
Sharad Singhal and Lance Wu. Training multilayer perceptrons with the extended kalman algorithm. In NIPS, volume 1, 1989.
Michalis K Titsias, Alexandre Galashov. Amal Rannen- Triki, Razvan Pascanu, Yee Whye Teh, and Jorg Bornschein. Kalman filter for online classification of non-stationary data. In ICIR, 2024.
Marcin Tomezak, Siddharth Swaroop, and Richard Turner. Efficient low rank gaussian variational nwwrinurrpp... cc/ Paper .ppp
inference for neural networks. In NIPS, 2020. URL https:
Filip Tronarp, Ángel F Garcia-Fernánderr and Simo Särkkä. Iterative filtering and smoothing in nonlinear and Non- Gaussia systems using conditional moments. IEEE Signal Process. Lett., 25 (3):408-412, 2018. URL httpss////a aalto. cm cm_parapub.pdl
Zhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa, Kurt Keutzer, and Michael W Mahoney, ADAHESSIAN: An adaptive second order optimizer for machine learning. In AAAI, 2021. URL
Amold Zellner. Optimal information processing and bayes's theorem. The American Statistician, 42 (4):278-2)) 280, 1988.
Wenxuan Zhang, Youssef Mohamed, Bernard Ghanem, Philip Torr, Adel Bibi, and Mohamed Elhoseiny. Continual learning a diet: Learning from sparsely labeled streams under constrained computation, In ICLR, 2024, httpttttt//// wttp:eeew....edd//eeeee/eotaa.a
13

---

A Abstract pseudocode
Algorithms to give pseudocode for applying the methods we study. For the predict step. we assume the dynamics model has the form 0, ~~ -N 00 b, Q²). The update step in Algorithm calls one of four grad functions (Algorithmss 1oo77 that estimate the expected gradient and Hessian using either MC or linearization combined with either the direct Hessian (or Jacobian and observation covariance) or EF. The update step also calls an inner step function that implements BONG, BLR, BOG or BBB on some variational family corresponding to the encoding of s/ (not shown). In practice the grad-fn fad--fn and inner-step-fe fn are not cleanly separated because the full matrix G2 is not passed between them except when the variational family is FC. When the family diagonal, grad-fn only needs to return diag( Gr,)). When grad-fn uses EF, it only needs to return G ((EM) (grad-MC-FF) or gi. (grad-LIN-EF) and inner-step-fn will implicitly use the outer product of this output as Gr,1- Finally, note the expressions for Di.: in Algorithms 6 and are equivalent ways of computing the same quantity, as explained after Eq. 28
for / = 1 xo do
-//r- 1 predict(3/:-1)
4.
y1)
end
Algorithm 1: Main loop.
def predict( 1/ = (11-1, E(-1))): Melee-1 = 1 + bt
Eqj-1 F, FFE Q
Return 1/1/1-1 = (11)/-1: 1
Algorithm 2: Predict step.
def ( Jz, N/: A0, grad-fn, inner-step-fn, a, I, M):
2 1/|1--1
fo,O,) f(x,, 0z)
he(0, = Elyr [r, 0r] = (==(((0.)A(n)
V. (e().) = Cov y xt. 0,] = V2 100, A(n))
(0,)) log p(u x 0,) = log
for 1: do
( G.,,)
l,, hr V,, M) 2/, inner-step-fin 1, 11,1-11 G.
end
Return Φh.,
Algorithm 3: Update step. The inner-step-fn is BONG, BLR, BOG or BBB (not shown).
B Additional experiment results
In this section, we give more thorough set of experimental results.
B.1 Running time measures
The running times the methods for the experiments in Figs. Tand where we fit a CNN to MNIST, shown in Fig. 3
The running times of the methods for the FC and DLR case, where we fit an MLP to a synthetic
regression dataset, are shown in Fig. 4 The slower speed of BLR (even with = 1) relative to BONG is at least partly attributable to the fact that BLR must compute the SVD of larger matrix (see Appendices E.5.3 and E.5.4).
14

---

def grad-MC-HERS 4 1 h, = D. V. [)
for m = 1: M do
((m)
Uprii (0)
9t,1
(m)
0s (m) l (1))
G(m) V2
0 ((m) lo (0.)
end
g
M m== 91,
Σ
(m)
G
GMC-HESS
M Σm 2m=1
(m) 1,
Return (9.,
GMC-HESS 1,1
Algorithm 4: MC gradient/Hessian estimator
def grad-MC-EF
= D. V, = M):
for M do
ab (0)
8 glon) Le
end
MC
M Σ m=1 1 It
(m)
(1:M)
t
19.,
(1)
gt.i
(M
GMC-EF
(1:M) (L:M)T 37
f,
Return (ghhr G
Algorithm MC gradient/Hessian estimator with Empirical Fisher
def grad-LIN- HES.LINN-HSSH..----S- = 11 hr Vr, M = D):
Mr,i-1 El0.lubr.-111
Dr.i
hr (14
H 30, 10-pii
Bhe
R, Ve(M((ii1)
gIN = H., R, t. (y: )1,))
GLIN-MESS -H R H G
Return (g G
-LIS-HESS)
Algorithm 6: Linearized gradient/Hessian estimator
def ENad-L 11. - 0, V,, M = =[):
, = Fl (2pp///111111
R = 11(
gl. = Vo.=pe,.-1
(yr - ho,)) R (ye he.)))
GLIN-Br D LIN T
Return (g LEI
Algorithm 7: Linearized gradient/Hessian estimator with empirical Fisher
15

---

140
120
100
10.
G4l
Junime
40
20
10c of0-- of mccn0
ber- micito
bogs dein-et mosoo
boop aeD of ecmcco
d8b aer--iit---os... An
ne-dirlo-nes, Jon
bug dorpdrr----sss Jo
bomp homp--ttzo- No
1a6 dirto-t Jo
bong beng- peng-drl0ue...
Bong- bong--lap-rt Sapplllpeettae me3300
sump dlgg Jhh
boog dag. of (mncto Bong- dig nom namahos. In
coupelag, nommal
bomp Stte me3300
bong oo O0 mmctt
bmpgalrr heos Jh
bovo-dtry sinp-dtra..o out
bong-atr hoos
bono if
bong slup
Figure 3: Runtimes for methods MNIST. Left: Corresponding to Fig. Dssnng different algorithms
on DIRR-10 family. Right: Corresponding to Fig. 2
using BONG on different variational families.
122--- (22-in
1C-MC305-C-L10 022:3F2-Lnd
de 1C-CC0543-14 (21-5F2-1))
6.5
W.B-PCLIB-2D- 20- coL-EroLLD0
N-PCESS4E-L1 012-275-int
5.03
Moo E--------------IEE STL-MCC00-22-8 100
Moo-ant--ii:33aa 2012
T--CCEO2-BB-4E L5O, S0L4Fo-in)
0:04
(CPEROTA1D LII LHe
0.3
0.03
6.2
6.0
Mum.
(b)
Figure 4: Running time (seconds) number of parameters P (size of state space) synthetic regression problem. For BBB and BLR, we show results using = and I = 10 iterations per step. Hessian approximations are denoted as follows: EFO-LinO = MC-Hess, EFI-LinO = EF-Hess, EFO-Linl Lin-Hess. (a) Full Covariance representation. (b) DLR representation. The BLR plot is truncated due to out of memory problem.
B.2 Detailed results for CNN MNIST
We show 3 approximations to the NLPD: plugin, MC, and Linearized MC each of these approximations to the posterior predictive, we also measure the corresponding misclassification rate based on picking the most probable predicted class. Results are shown in Figs. 5 and 6 We see that the plugin and lin-MC approximations are similar, and both are generally much better than standard MC.
Here we report further metrics for the experiments in Figs. and 2
Finally, in Fig. 7 and Fig. 8 we report the test-set expected calibration error (ECE) at time steps [250, 500. 1,000, 2,000], computed using 20 bins. Note that among the LIN-HESS variants, BONG-DLR-10
$ *Lin-MC defined in Footnote The motivation this approximation (from Immer et al. 2021b is the following: If we push posterior samples through nonlinear predictive model, the results can be poor if the samples are far from the mean, but if we linearize the predictive model, extrapolations away from the mean are more sensible. This is true if the posterior was not explicitly computed using linear approximation.
16

---

120
UnB
110
1200
100
1330
110
1003
136
212
1150
100
0220
170
120
80
000-011 af mxt100
Figure 5: MNIST results for methods using DLR family. Left column shows misclassification rate, right column showns NL... First row uses plugin approximation to the posterior predictive, second row uses linearized MC approximation, and third row uses standard MC approximation.
method is the most well-calibrated (in addition exhibiting the strongest plugin and linearized-Me NLPD results), when compared to other DLR-10 10 methods well other BONG variants.
17

---

3353
peppartin INSCCE
20
Prenu
by Premin Tep-deee
Soru Sap
Figure 6: MNIST results for BONG variants. Left column shows misclassification rate, right column shows NLL. First row uses plugin approximation to the posterior predictive, second row uses linearized MC approximation, and third row uses standard MC approximation.
18

---

Gccected laation anar (AcO wED.Nee S00M
920
5x2
630
120
300
920
532
500
Figure 7: MNIST expected calibration error (ECE) results at selected timesteps for methods using DLR family.
3.0
5.3
020
512
665
3.2
020
512
505
AB
605
Figure 8: MNIST expected calibration error (ECE) results at selected timesteps for BONG variants. We see that BONG and BOG are both well calibrated, as least when combined with LIN-HESS, with BONG have slight edge, especially for small sample sizes, where there is more posterior uncertainty.
B.3 SARCOS dataset
In addition to MNIST, we report experiments on the SARCOS regression dataset (D = 22, Nusai 44,484, Nie = 4449, C = 1). This dataset derives from inverse dynamics problem for seven degrees-of-freede SARCOS anthropomorphic robot arm. The task is to map from a 21 2i-dimensional input space (7 joint positions, 7 joint velocities. joint accelerations) the corresponding 7 joint torques. Following Rasmussen and Williams [2006), we pick single target output dimension, so org/ gpml/datat.
C = The data is from http::////aausssapoooeeeeeee
We use a small MI.P of size 21 21-20-20-1, 1, so there are 881 parameters. For optimizing learning grid search NLPD-PI. We fix the variance of the prior belief state
rates for SARCOS,
19

---

to o = 1.0, which represents mild degree of regularization We fix the observation variance to R. 0.1R, where Rt Var(y1:)) is the maximum likelihood estimate based the training sequence; we can think of this as simple empirical Bayes approximation, and the factor of 0.1 accounts for the fact that the variance of the residuals from the learned model will be smaller than from the unconditional baseline. We focus on DLR approximation of rank 10. This gives similar results to full covariance, but is much faster to compute. We also focus the plugin approximation to NLPD, since the MC approximation gives much worse results (not shown).
B.3.1 Comparison of BONG, BLR, BBB and BOG
In Fig. we show the results of using the LIN-HESS approximation. For iteration per step, we see that BONG and BLR are indistinguishable in performance, and BBB and BOG are also indistinguishable. but much worse. For 10 iterations per step, we see that BBB improves significantly, and approaches BONG and BLR. However, BLR and BBB are less than 10, due to constant factors of the implementation.) (Note that BONG and BOG always use a single iteration, their performance does not change.)
about 10 times slower. (In practice, the slowdown is
In Fig. TO we show the results of using the MC-EF approximation with MC 100 samples. The trends are similar to the LINHHES case. In particular, for / = 1, BONG and BL.R are similar, with BONG having a slight edge; and for = 10, BBB catches up with both BONG and BLR, with BOG always in last place. Finally, see that the performance of MC-EF is slightly worse than LIN-HESS when I = 1, but catches up with = 10. However, in larger scale experiments, we usually find that LIN-IIESS is significantly better than MC-EF, even with / 10.
Mozec
tept: buteo
Maddi
Jihwdo3) Deter bug: iahholl--30
bos untuunt aL und 25 12 21-822 ripap t-ataa:3.a0i 0.13)
loo r12 untheas 813 un0 (ee:246 riss sigriwzoo 3 2.33 Ln-d--10 u20 uno boxzz 21-022 12pdd y ef 2000:-231 (0.4) hoo 41 leen jeec
trrdi-10 Lintiss us unu looc 1-1-00000 vaa sag w2000 201-2 1-2221
11.3 so feec 0.0 por
pat 2000 tie
longy Liiiittss boach.2 o.1-
bony di Lintitte bocch.t
a.2
2-
10.3
2000
13:0
2002
720 num. training observations
2002
training observations
(a) iteration.
(b) 10 iterations.
Figure 9: Predictive performance SARCOS using MLP 21-20-20-1 with DLR rank 10. Error bars represent +1 standard deviation computed from 3 random trials, randomizing over data order and initial state Moo (a) We show all 4 algorithms combined with LIN-HESS approximation and 1. (b) Same as (a) but with 10.
B.3.2 Learning rate sensitivity
In Fig. [ we show the test set performance for BLR (with LIN-HESS approximation) for 5 different learning rates (namely x 10 10-3, 1 10 2, 10 10-2, 2, x 10 and 10-1).
When using 1 iteration per step, the best learning rate is 0.5, which is also the value chosen based on validation set performance. With this value, BLR matches BONG. For other learning rates, BLR performance is much worse. When using 10 iterations per step, there are several learning rates all of which give performance as good as BONG.
*This value was based on small amount of trial and error. Using a smaller value of do results in underfitting relative to a linear least squares baseline, and using a much larger value results in unstable posterior covariances, causing the NLPD-MC samples result in NaNs after few hundred steps.
20

---

Mozel -01+0011;
Eept: bateus
Model JIPw0223, Cata: Egg: sakaoth
bio ria MccPSco LRC 25 5,3 midd g 007-2000555000003
t00-5 cessscoo-oto---a L30, 65 Jsec B-------- rise pigh-2000 7.70 0.10) Meddino 1o0 (Mt_]] Isec (iile-335 net wosee 45m 0.05)
14
tendi10 acaF UNI Isocili reed aetwoon 4.s5
boy arteeoe kay Usec S.te-li radt 601220066.
rist #T-22000 w0)
12
bony di Moceaze bec 01 wisd:i@th20002 2.233.
bony di sicerize bec 12.3--33-000 1ipe-1gge00000
Abp
rip3
2-
2-
520
1000
1000
3752
2002
250
730
2000
1550
2002
num, training observations
num. training obsorvations
(a) iteration.
(b) 10 iterations.
Figure 10: Same Fig.
except use MC-EF approximation with MC 100.
Model iioo 20, 20,22.2 1/111011 Dasa: sarowa beck bi/2 ritngg
Model mpp 20.00.00 (r1) Date serces tagt- ooop----ong
12
hong dir Liieees jparis
rip 2000 3224-3
Iono di Lintees joac &
2000-2200
br-d-100 Unites UR: baech 022 repdpp d 2020:8.13--003
Lt-d-10 Unttls 110 une 325 joe 0.1 sigm
bi-dt-to--iiter 2R.2 (C---00
2000 1..55--3:11
bi-drro-Uhiller, Unilerr 200-RR Jese2t Se widd 2016:1.36- 6.14) bi-d-10 Uniter L2o Jene22 de 2010:0.30-- 6.17)
bi--d-20 Uniteen URZ 0.5 2107 3.221
0104
ntod
observations
observations
(a) iteration.
(b) 10 iterations.
Figure 11: Same setup as Fig. 9 except now plot performance for BLR for 5 different learning rates. We also show BONG as a baseline, which uses a fixed learning rate step size of 1.0.
Model 20 3P 0213 Cabs:
newetbbler
Mode: 20 0013 Cata:
now bot10
bring dirt- intte-tto---ieeeeeeee------- (nt-iieeeeeeeeeete Joann +-6. mix -ix--lom---20 rlomi-200n 3 3214-3211
hrnn iiioeii Jonc 6. d:11 mtps =igT-2200n 3234--332
boo unthor ano pc5 Seecd.5 3--00 midd a 20tt 151--0 2212
bbo 4 Unfhree 812 uno, 209 las - ntodi 73200025.73 0.02)
Mot-o 102---110--m unthoo ki--ne 31 breca
repdlr 20000005 20 u.)))
Moo-ss 1i0-0112-0e AiPhorthoostt E1P uno Seccia ei -.1-bo hol----stt m2200 300 (220)
hoz-c rib--ree-ue
InecS. de 6.1 mepip 3-00000: 30-
boo-s urb--niiee J8o (sec744 23 pigTi 30
fom-s Lnteess L10 Lseczs wed
untrurresss RXD (N0 (De01ig
2-
250
300
720
2000
:220
1300
2752
2002
250
520
720
2000
1220
1500
2752
2002
pum. training observationn
num. training obsersations
(a) iteration.
(b) 10 iterations.
Figure 12: Same setup as Fig. 9 except now we plot performance for BBB for 5 different learning rates. We also show BONG as a baseline.
In Fig. 12
we show the analogous plot for BBB. When using iteration per step, all learning rates result in poor performance, with many resulting in NaNs. When using 10 iterations per step, there are some learning rates that enable BBB get close (but still not match) the performance of BONG.
21

---

teat:
Hode: 8811
bong deto-untlees (sac6.4 2.1.::66----00000 midd =igmT=2000 3.22 3.32
bong 5730 MCCET22 Tec: -pp--1g--200--
12- 0.533
bop Lia
20,265 (sec 5--0..
ricdindd 22002
boy wo.drtamaze (L0 180, 500 Jsic de
-2000 8.17--0
bog 112 lasc rightpitttttt1 pigi=2000
2.33
110 lasc a:0t- 3.00
bug 1122
1es 011 2 2-24
Alps
rip3
3-
320
1000
1200
3752
2002
250
750
2000
15000
2002
training observations
num. training obsorvations
(a) BOG with LIN-HESS.
(b) BOG with MC-EF.
Figure 13: We plot performance for BOG for 5 different learning rates. We also show BONG
baseline. (a) LIN-IIESS approximation. (b) MC-EF approximation.
Finally, in Fig. 13a where results are much worse.
we show the analogous plot for BOG with LIN-HESS, and in Fig. 13b with MC-EF,
Overall we conclude that all the methods (except BONG) are quite sensitive to the learning rate. In experiments, pick value based performance validation set, but in the truly online setting, where there is just single data stream, picking an optimal learning rate is difficult, which is an additional advantage of BONG.
C Proof of Propositions 4.1 and 4.2
Proposition 4.7 To case notation we write the natural parameters the prior as 3/1/1-1 =
Xije-ii Mej which can interpreted as the prior sufficient statistics and prior sample size. Note that ø; can be omitted as a constant. Based on the prior qua. (0:) the exact posterior is
p(O:(D;)
900,,-, Pr.
(35)
x exp (x Vaje A((i) exp (0) v. A(0.)
(36)
qsp, (0:)
(37)
Xt 1-111 y/: V/tt-1 +1
wb.
(38)
For BONG, we first note the dual parameter is given by E0 [(())]
Pitt-1
(39)
0, A((,,)
(40)
Therefore the natural gradient in Eq. (5) is
V Puje Ee [log p (yel [.)]] V Pup-1 Ee, Bor
( y - (0,)
(41)
Veeu .P.
yt:
(42)
1
yt
(43)
Therefore the BONG update yields vh [xp-1 + Yti Vije-111 1] in agreement with Eq.
38
Proposition 4.2 The intuition behind this proof is as follows. For the mean update in Eq.
l
linear in O, so the expectation equals the value at the mean. For the covariance update in Eq.
10
22

---

Ve, 4 is independent of 0, so we can drop the expectation operator. The tricky part is why we need different linearizations for the Hessians to agree. It has do with making the Hessian of the NN disappear (as in GGN). In the Gaussian approximation this happens when the predicted mean ( fi (()) linear 0r. In the plugin approximation happens when the outcome-dependen part of the loglikelihood (h, (0,) (T yt) is linear in 0, In the latter case the only nonlinear term remaining in the log-likelihood is the log-partition A, and the two methods end up agrecing because of the property that the Hessian of the log-partition equals the conditional variance R.
Formally, under the linear(h)) Gaussian approximation in Eqs. (14) and 17 the expected gradient and Hessian can be calculated directly:
Eg, E8 Boa Vo, log P (ya [e))
H (y) ú. H.(0. Mclc-1)
(44)
H,R," (yy - 91))
(45)
Eg. Eg.-
Vg log (y: 0.)]
Eg, Bg R, H.
(46)
- H,
(47)
For the linear( approximation we use the properties of exponential families that (1) the gradient the log-partition A with respect to the natural parameter f+(0,) equals the expectation parameter h4 (0.). (2) the Hessian of the log-partition with respect to the natural parameter equals the conditional variance, and consequently (3) the Jacobian of the expectation parameter with respect to the natural parameter equals the conditional variance R₁:
7= - A(n) he(Me((11) 9
(48)
V pr-fem () A(n) -V [yrr000 - Me|c-1] = Rs
(49)
(O.)
R,
(50)
(O.) 01=M1
The last of these implies F, = R, H, Therefore the expected gradient and Hessian can be calculated
Eg Vo, logp, D. (2)(0)) = Vo.=mu log P.. (3.
(51)
Flyr - F.. -
(52)
FT -is)
(53)
H.R, "(ve (r)
(54)
Eg
V2 log o. (v.10.))
V8. V8= log LLD (y. (.))
(55)
=-F (V2 ( A(r) F.
(56)
=-F.R,F,
(57)
=-H,R, R, 'H.
(58)
D Mirror descent formulation
In this section we give more detailed derivation of BONG as mirror descent and use this to give two alternative interpretations of how BONG approximates exact VB: (1) by approximating the expected
NLL as linear in the expectation parameter p. or (2) by replacing an implicit update with explicit one.
Assume the variational family is an exponential one introduced at the end of Section 3 with natural and dual parameters ap and p. sufficient statistics T(Ø), and log-pau parrtttttionn (<bb):
(0) exp (Nb ((() + (bb))
(59)
p Ea-ay [T())]
(60)
23

---

We first review how NGD on an exponential family is a special case of mirror descent (Khan and
Rue 2023 Martens 2020 The mirror map is the gradient of the log- partition, which satisfies the
thermodynamic identity
(
(61)
This is bijection when + is convex (which includes the cases we study), so we can implicitly treat a/ and p as functions of each other. Given a loss function L(ab), MD iteratively solves the local optimization problem
4/+ = arg min (V,: L(bb)), p) + Da ub)
(62)
a
The first term is linear (in p) approximation of about the previous iteration t/ and the second term is the Bregman divergence
Do 14+1) = ((b.) (+b++1) (vb: + Pi+1
(63)
The Bregman divergence acts as a regularizer toward t/: and captures the intrinsic geometry of the parameter space because of its equivalence with the (reverse) KL divergence
DKL. (gp. 11 (w:)) 1eo- [(bb+1 *b.)) =b..TT(T0 T(0) + +(bb)) ((bs+11)
(64)
(1), +b++1)
(65)
Importantly, this recursive regularizer is not part of the loss and serves only to define an iterated algorithm that converges local minimum of L. Solving Eq. (62) by differentiating by yields the
MD update
s 1++
3 a L(b.)
(66)
Because the Fisher matrix for an exponential family is F = 0p/0 this is equivalent to NGD with respect to sb. Khan and Rue [2023] offer this as derivation of the BLR, when L(bb) is taken to be the variational loss from Eq.
By applying this analysis to the online setting, approach can be seen to follow from two insights. First, the online variational loss Eq. (2) already includes KL divergence from the previous step, so we do not need the artificial regularizer in Eq. (62). That is, if we start from the online variational problem in Eq. 44 and define (zb) as the expected NLL.
(( = Eg [log p(v: f(())))]
(67)
then replacing Le(1) with a linear approximation based 2/s/t-1 and applying Eq. (65) leads
ab.- arg min P (o( D (1b 1b)
(68)
62 see this defines MD algorithm with unit learning rate that works in single step rather than by iterating. Paralleling the derivation of Eq. (66) from Eq. (62) we get
By comparing to Eq.
+ - V Le
(69)
which matches the BONG update in Eq. 55) Thus BONG be seen approximate solution of
the online variational problem in Eq. (4) based on lincarizing the expected NLL wrt p. (Note this is different from the assumption underlying BONG-LIN that (0,) or hG.))) is linear in 0..)
Second, in the conjugate case, this linearity assumption is true: L, is linear p (see proof of Proposition 4.1). Therefore (8 is equivalent to solving Eq. (4) exactly:
v arg min (z(1)) Do (4/4)-1,
(70)
This recapitulates Proposition 4. that BONG is Bayes optimal in the conjugate case. In general the exact solution Eq. 70 is
tb
- - Vp, Le(bt)
(71)
This is an implicit update because the gradient is evaluated at the (unknown) posterior, whereas Eq. (69)) an explicit update because it evaluates the gradient at the prior. (In the Gaussian case these can be shown to match the implicit and explicit RVGA updates of [Lambert et al. 2021 Therefore
BONG can also interpreted as an approximation of exact VB that replaces the implicit update,
Eq. (71. with explicit update, Eq. 69).
24

---

E
Derivations
This section derives the update equations for all 80 algorithms investigate (Table 3 plus the MC-HESS and LIN-EF variants), In Appendix E.6 we also translate the BLR algorithms from our online setting back to the batch setting used in Khan and Rue [2023]
For an exponential variational family with natural parameters s and dual parameters p. we can derive all 16 methods (BONG, BLR, BOG, BBB under all four Hessian approximations) from four quantities:
V pr. Ee, 00,, Dlgg (yelft (0x))
(72)
V P.... DEL 1 19
(73)
V - -v. log (ye fi (0.)
(74)
V Ve DxL PexI(QO..... (o (1)))
(75)
The NGD methods (BONG and BLR) use gradients with respect to Pt.c-i while the GD methods (BOG and BBB) gradients with respect to 2 For BONG and BOG the DEL term is not relevant, and there is no inner loop so 1/t,i-1 = f and gt, = 9t. G1, = G2.
When 1a is not the natural parameter of an exponential family we must explicitly compute the inverse-
Fisher preconditioner for the NGD methods. Therefore the updates can be derived from these three quantities:
F P V V Ro, /o log (Brlft (0:))
(76)
(77)
V V DEL (gor.. POKa(ggr....
-..)
(78)
We will frequently use Bonnet's and Price's theorems [Bonnet 1964 Price 1958
V p. EvT 1.E. 1) flogp (yelfa (0.)) ENi( E.. Ve. log (Velfe (0.))
(79)
gt,i
(80)
E EN (yelfa (0,)) -[N(m))- E.
Vo, log (y Ii (0,))
(81)
_Gr.ii
(82)
2
For diagonal Gaussians with covariance Diag (a²). Price's 's theorem also implies off NN E.E 1) [log p (3/ ft (0:)) diag (Gr.i)
(83)
Update equations for MC-HESS, MC-EF and LIN-EF methods are displayed together in the subsections that follow, because for the most part they differ only in the choice of G GMC--HLSs G-M--EP GLIN-Dr to approximate G; and gºd or gi to approximate 9t- We note cases where decomposing G.C-EE
((: ( (1:M)) -M Gt
G GLIN-ET -! (N) allows more efficient update.
We derive updates for BONG-LIN-HESS and BOG-LIN- HESS from the corresponding BONG-MC-HESS and BOG-MC-HESS updates using Proposition 4.2 which entails substituting
9: H (y) ÿt)
(84)
G + H H,
(85)
For the algorithms with inner loops (BLR and BBB) we adapt the notation of Section 4.4 as follows:
Yr. = h1 (Hr,i-1
(86)
Thu
H D0, 10=
(87)
R₁,: = V [yr 00 = i1.
(88)
gt. = H. R (yr (r.))
(89)
G. = H R H,
(90)
'We use diag(A) to denote the vector of diagonal elements of matrix A and Diag(v) to denote the matrix whose diagonal entries and off-diagonal entries O.
25

---

This corresponds to basing the linear( Gaussian and linear(h) deeta approximations at pt,i-1 instead of M1jt-1- Thus the updates for BLR-LIN-HEST and BBB-LIN-HIE HESS are obtained by substituting
g - H R (ye- (c.))
(91)
G + H R H
(92)
E.1 Full covariance Gaussian, natural parameters
The natural and dual parameters for a general Gaussian are given by
41,1
(1)
- M
P = M
(93)
-1,-1 -
(2)
Pi = -Mii-114 114, + E1,-11
(2)
(94)
Inverting these relationships gives
Mr.i-1 = - -/0 - 1i,6-1 = =Pi
(95)
Σ Σt,i- = - (2)
= =P -Pii-IPui-
(2) (1) (1)T
(96)
The KL divergence in the VI loss is
DRI. (1 M82ge 1) = 2 (Mei11 11((111 1 2/12- (Me -1
+ Tr (I Ev.-11) - 2 log 10, + const
(97)
with gradients
V piii DKL (Aor.. (1))))1))
2 (Me,s-1 M(((-1)
(98)
V Exai DKLL (go (4bbee-1) 1) (iiji-1 - -
(99)
Following Appendix C of Khan et al.
2018a| for any scalar function e the chain rule gives
3pt
1 V V 0E,,i-1 V...-1
(100)
P
0p.i-1
8pth11
= Vumni (VE,.-1)) Mrii-1
(101)
0p1 V 1l
P
(2)
Opt.1-1
1+
8pp.111
(2)
VE....
(102)
VV L
(103)
Therefore
V (1) Eg. llop (Belff (0,)) gr, Gr.iHri-1
(104)
P
V (2) Ee. Ee [log p (3e (0:)) } GG,,i
(105)
and
Dier. Dier. 110))--)) = E 1 14,ii-1 - - 2//2-
(106)
Pr
DEL - = Σ (2it11 -E
(107)
P
Following the same approach for ab gives
0
2Σ, 0E,,i-1
V
Vm.._,
e+
+
(108)
6,1-1
=-2 -1))()))11
(109)
= Er,i-1 1. V.n.i
(110)
0p,1i-11 1
V..._,e e+ 0E
V
(111)
25 G1-1
(1)
0
(2)
(3)-1 (V.... () 1 1 + 1/ (V (2)-)
(1)
2/ (2) 1 (2)-1
(112)
--E,5-1 (V Vpn.- 1) 14,1111 + 22,,-11 (VI,,-1)) E1,1-1
(113)
26

---

Therefore
Ea logp (3olft (0.)) 19,5-19ti
(114)
Es E [logp (yelft (0.)) -2 + Gr,
(115)
1i
and
DKI. (T [ = -1,i-1 Σ-1
f|t-1 1
(14 14|c-1)
(116)
c,(-1
DKL. (aa,(99 Vobese 1) 2-11i-11 211111 (Kt 1o/c-1) 16
B
(2)
E { E Ec1-1
(117)
E.1.1 BONG FC (explicit RVGA)
Substituting Eqs. 100 and
105 into Eq. 5 gives
2b 2 (1)
f|t-1 + gt G.
(118)
4,
(2)
+ G.
(119)
Translating to (m: Ex) gives the RONG-FC update
#t = Mijt-1 + Engt
(120)
Σ, = Eー1 q|-1 Gr
(121)
This is equivalent the explicit update form of RVGA [Lambert al. 2021]. Using G GNC-HESS this update takes O(P)) because of the matrix inversion. Using G and the Woodbury matrix identity can write the update in form that takes O(M + +M*;
Maj-1 + K,1M
(122)
Σ, Σmj-1 - -KGG, ((:MM) Efjt11
(123)
K, G (MINs + G "( G(1:.))) -1
(1:M))
(124)
Likewise using G. GN takes O(P2):
fe pt|t-1 + K,
(125)
Σ, En 1-1 K, (n Edp-
(126)
Edje-1gh1N
K, 1 (N) E1ee-19 IN
(127)
E.1.2 BONG-LIN-HESS FC (CM-EKF)
Applying Proposition 4.2
to Eqs.
120 and 121 gives the BONG-LIN-HESS-FE update
H,H,R (yr (.) P1j-1
(128)
Σ,
H.R,'H. 'H. Σ 1|-
(129)
This is equivalent CM-EKF TTonarp al. Woodbury identity in form that takes O(CP* +))):
2018 Ollivier 2018 and can be rewritten using the
μr P Ki(y: U.)
(130)
Σ, Enj-11 K,H,E.Je-.
(131)
K, Eqp-HH (R, + H, H1
(132)
27

---

E.1.3
BLR
Substituting Eqs.
104
(07
into Eq. (33 gives
v 1, v - + a gr. Gr. - - 1 + (eje-))
(1)
(133)
2/, 4i.311 + GG. + 2 - - 2(((11
(2)
(134)
Translating to (Hu,i Et,i) gives the BLR-FC update
Pt; = P1,i-
a E1, - 1-1 (Mit-11-Het- 1 (1,1-1) + Ci,
(135)
2 (1 a)) - a o aGr,
(136)
This update takes O(P3) per iteration because of the matrix inversion. In Appendix E.1.1 we were able use the Woodbury identity to exploit the low rank of G.C-FF and G, G. and obtain BONG updates with complexity quadratic in P. This does not appear possible with Eq. (136) because of the extra precision term on the RHS (applying Woodbury would require inverting (1 (1-)) + (
Therefore unlike BONG-FC, BLR-FC requires time cubic in the model size, for reasons that can be traced back the KL term in Eq. (33).
E.1.4 BLR-LIN-HESS FC
Applying Proposition 4.2 to Eqs. 135 and (166) gives the BLR-LIN-HESS-FC update
Me Mt 1 + a E 16- (Mie-1 M H (v. -
(137)
D (1 a) DT - + a -i + aH R H
(138)
This update takes ((P3) per iteration because of the matrix inversion.
E.1.5 BoG FC
Substituting Eqs.
114 and
115 into Eq. 31 gives
wb 2b i|t-1 + 0 11-19 19t
(139)
w12) * w + 2a 2n -- 19. + E
(140)
Translating to (1t, E,) gives the BOG-FC update
Σ,> 1 - M a
(141)
E, Σj 4a
- 1
(142)
This update takes O(P3) because of the matrix inversion. The greater cost of the BOG-FC update relative to BONG-FC can be traced to the difference between GD and NGD: the NLL gradients wrt tb 1-1 in Eqs. (114) and (115) more complicated than the gradients wrt Pac-1 in Eqs. (104)
and (055.
E. BOG-LIN-HESS FC
Applying Proposition 4.2 to Eqs. 141 and 142 gives the BOG-LIN-HESS-FC update
#: = Σ,E 1|1 M + a E, H R, (y) ÿt)
(143)
Σ Σ 4a E H
(y: úi) phr-1 + 2a H R, EHpE11
(144)
This update takes 0(p3) because of the matrix inversion.
28

---

E.1.7 BBB FC
Substituting Eqs.
114 to (117) into Eq. 34 gives
26 16,3-1 + a a 190 aZir -- E (1-11 i|(-1 -
(1)
(145)
2 1 + -2a 19, i-1 + aZitt-11 G
-2a 1/111 (Mt,i-11 - 11 / - aDi,i-1 2-11--
-2 2
E (146)
Translating to (uu,s E1,i) gives the BBB-FC update
Mr. -1 + a2 (g.s + Eqj(11 (M1) (1--1))
(147)
2 (Hi|e-1 ((,ii-1) 14 + Ip
Σ - - 20
(148)
+2g1, + (G.s - 2 Er,i-1
This update takes O(P3) per iteration because of the matrix inversion.
E.18 BBB-LIN-HESS FC
Applying Proposition 4.2 to Eqs. 147 and 148
gives the BRB-L.N-HESS-FC update
,1, = E,, E 1,1 M1
+ Ec 1
H.. RY (y) (1.)) + E (Hite-1 ((,(-1))
(149)
2Σ #/1 1-1 (Mije-1 - M1,i-1) 14 + Ip
Eq. -2 20.-1
+ +2H t, R,y, 1 (y) ÿr,e) 1,,
(150)
H R, H + 1|/-11 Djji-1
This update takes O(P3) per iteration because of the matrix inversion.
E.2 Full covariance Gaussian, Moment parameters
The Bonnet and Price theorems give
V 1 Eo aw. pogg (yr [ft (0.)) gt,i
(151)
V En.: E E [logp (y:ls: [i (0,)) GG,,s
(152)
From Appendix E.1 we have
V Prss- DEL 1 110))-.)) = ETpL1 (M4,5-1 - -1((-1)
(153)
E.,.. DEL (Awe..- (40))111) = | ETi-1))
(154)
We write the Fisher with respect to the moment parameters 2/ = (1r, veol ()) as a block matrix:
F pnu
F,.E
F
(155)
F
FI
FEx
F
29

---

The blocks can be calculated by the second-order Fisher formula
F p.
Buy V
pape log 4vp (0)
(156)
=I-1
(157)
F Fx = -Bg V Vu
log qv (0)
(158)
= -Bey
(Vy (0- 14]
(159)
= 0
(160)
Fx Euy
[ V x log (0)
(161)
Egy.
Vx 2
Σ (0 1 -0-1)00-1)-> - )
(162)
(VyΣ (0 m)(0 -000-110 -1)1--1
E
--
(1) (0- 14) (V=I--)) - -(VxΣ-1)
(163)
(164)
2 E
(165)
In the final line we used
Vr (M-1) ij = E - Σ - je
(166)
=- E
1 ij
(167)
with ij and kl treated as composite indices in [P Therefore the preconditioner in the NGD methods
Σ1,-1
F 1-
(168)
2 8 E
E.2.1 BONG FC, Moment
Substituting Eqs.
151
152
and (688) into Eq. gives the BONG-FC.MOM update
14 Mr|t-1 Σ 1
(169)
Σ, = E + - G
(170)
This update takes ((33) using G. GMC-HESS P2) using G.N-EF and O(P2) using GLI-EE The
update is similar to the BONG-FC update except that Eq. 170 ignores the G. (1:M)* E ((:M)
(1:M)
term in Eq. (124) the (N) Eqp-1g term in Eq. 127 which estimate the epistemic part of predictive uncertainty.
E.2.2 BONG-LIN-HESS FC, Moment
Applying Proposition 4.2 to Eqs. [69 and 170 gives the HEILLIN--ES-..OOO update
11, Pat-1 Enp- H R, (yr 9:)
(171)
Σ, E
E H. 'HE
(172)
This update takes O
E.2.3 BLR FC, Moment
Substituting Eqs.
[15] 154 and (168) into Eq. (32
gives the BLR-FC. update
Dr.i Pr.i + a a.E Etit f/t-1 (M M1,i-1) + a...i-19t.
(173)
E, (1 a1) E,,-1 1 + a GG.- Σ Eg.i-1
(174)
This update takes O(P3) per iteration because of the matrix inversion. Using G (1:M) G (1:M)) or GUI--E - (g! allows the BONG-FC MOM covariance update
M in Eq. 170 to scale quadratically, but this does not help here. Instead the BLR-FC update
scales cubically because of the additional - term that comes from the KL divergence in the VI objective.
30

---

E.2.4 BLR-LIN-HESS FC, Moment
Applying Proposition 4.2 to Eqs.
173 and 174 gives the BLR-LIN-HESS MOM update
Mt, = Ptúi- a 1|--1 (Hije-1 (L,-1) + H R,,! (ye yr.))
(175)
E, (1+ a)) Et aL
E 2/c-1 HJ, R He. Eiii-
(176)
This update takes ( per iteration because of the matrix inversion in the 2 term that comes from the KI. divergence in the VI objective.
E.2.5
BOG FC, Moment
Substituting Eqs.
[151) and 152 into Eq. 31 gives the BOG-FC.MOM update
14 M +
ogt
(177)
Σ, > 1|111 G.
(178)
Note the mean update is vanilla online gradient descent (OGD) and does not depend on the covariance. This update takes O P2) using G MC--EESS G.c-Er and O (P2) using G
E.2.6
BOG-LIN-HESS FC, Moment
Applying Proposition 4.2
to Eqs. 177 and 178
gives the
update
μ1 Mijt11 + R, (y 9:)
(179)
E, E 2
" H;
(180)
This update takes O(CP2).
E.2.7
BBB FC, Moment
Substituting Eqs. 151 to 154 into Eq. (34 gives the BBB-FC.MOM
H1s = Mt,i-1 + a a. i|t-1 (Heje-1 - M1,5-1) + agg,
(181)
Σ, Σi,-1
- + G1,a.
(182)
This update takes O(P*) per iteration because of the matrix inversion, which traces back to the VI objective. Comparing to the BOG-FC MOM update in Eqs. (177) and (178) (which has quadratic complexity in P), the extra terms here come from the KI. part of Eq. (34).
E.2.8 BBB-LIN-HESS FC, Moment
Applying Proposition 4.2 to Eqs.
181 and 182) gives the BBB-LIN-HESS-FCMOO update
Mr, = Mt.i + a a>
1/t-1 (Mite-1 M,,i-1 + aH! R,, (y: - (t.))
(183)
Σ, Et,-1
E-1 -
E-1 f||=1 H., H.,
(184)
This update takes O [3) per iteration because of the matrix inversion. Comparing to the BOG-LIN- HESS-FC.MOM update Eqs. (179) and (180) (which has quadratic complexity P), the extra terms here come from the KL part of Eq.
E.3 Diagonal Gaussian, Natural parameters
Throughout this subsection, vector multiplication and exponents are elementwise.
The natural and dual parameters for a diagonal Gaussian are given by
(1) -2
1 0 (2)
p Pt.i-1
(185)
p Pi,-1 = -M15-1144 + o² fi--
(186)
111-1 0
31

---

Inverting these relationships gives
|
1,3-11
(2)
1/,i-1 = -iii-1
(1)
(187)
2 0i-1 2
(
(2)
= P - (P(_-1)
(188)
The KL divergence in the VI loss is
DRL. (1 Mube 1) 2 (Ma-i-1 Mo|t-1) 1 0111- + Σ
-2
- (qt--1 -i - loggo
+costt
(189)
with gradients
V Me 1 DKL PPKL(Qgs... (9 (1)))--)) = 1/t-1 (Mr, - Mele-1)
2
(190)
DKI.( (geb..- (4pat-)) = 2 (//1-1 - -
(191)
For any scalar function l the chain rule gives
0pt V
00,,- 1,1-1
V ( L =
p 1+
V
l
(192)
0pt.11
-pi.-1
Vp l 2411,i-1 V Vo L
(193)
0p1,i-1 V
00
2
P1,-1
(2)
8
(2)
,e+
0p
(2)
1,1-1
l
(194)
Vol..
(195)
Therefore
V Eg, E0,~au.. - Dlogp (yelft (0.)) gr, diag (Gr,s priil
(196)
Pi.i-]
(2) Eo Eo flogp (yel.ft (0.))] = = (G.,,)
(197)
Pra--
and
V DEL (agp.,---1 190 .) = 0 -2 0, - -
1- Milt-
(198)
P
V DRL (1 1(bbe)--))== = 1 -2
(1/(-1 -07.0-1
(199)
Pr,-1
Following the same approach for ab gives
f 0pp-i1
d 0 V
(200)
awll)
Jubl!! i-
= = (w (111-11
V l
(201)
= oi.-1 V Ve
(202)
1/ l 0p1,i-1 V
00 0 Vo o
e+
(203)
0up-
2 (w10))
-2
(1)
|, l + 2
(1)11)
-2
(204)
= -20 11,5111 V Vpr.,, l + 20 1
(,1-1 Vo2 Vo
(205)
Therefore
Eg Ea Dlogp (v. ff (0x)) = ,i.--191.
(206)
Ea- Dogg (v.1 J1 (0.)) = 20 Mr. 19t, + - (Grs)
(207)
32

---

and
DEL (100.,.--19000000 = : l|--1 (Mri-11 Moje-1)
(208)
DEL (es 20 (
20,11-10
11: 111-1 (He,i-1 pelje-1)
-2
o
(o
#|t-1 0
-0,,
(209)
Our implementations often make of the following trick: Suppose A Ruxm and B Rmxx. Then we can efficiently compute diag(AB) in ((m)) time using (AB);: Σ/= Ayp
For MC-HESS methods, we approximate the diagonal of the Hessian for each MC sample 0,m)) using Hutchinson's trace estimation method [Hutchinson 1989 which has been used in other DNN
optimization papers such adahessian Yao al. [002 This involves an extra inner loop with size denoted N.
E.3.1 BONG Diag
Substituting Eqs. (196 and
197 into Eq. 5 gives
wb 4b f|t-1 + diag (G:) Paje-1
(210)
4 (2)
wb (2) diag (G,)
(211)
Translating to (Me oi) gives the BONG-DIAG update
f: Pc|t- + 0
(212)
0, 0/1-1 diag (Gt)
(213)
This update takes (M P) estimate G, using G G. O PP) using G. G.C-HESS and Hutchinson's method, and O(P) using GL
E.3.2 BONG-LIN-HESS Diag (VD-EKF)
Applying Proposition 4.2
Eqs. 210 and gives the BONG-LIN-IIES--DIAO update
[r: Mtlt 1 + 2 (H.R," (y: 9.)) σ i|t-1 + diag RH.R,"" "H:)
(214)
(215)
This update is equivalent VD-EKF (Chang al. 2022 and takes OCC2))
E.3.3 BLR Diag (VON)
Substituting Eqs. 196 to 199 into Eq. 33 gives
2 = 11211 + a gr.i diag (G.,)) Me.i-1 - 07--1M1 +
to -2
(216)
10 = 1 + a diag (G.,,) + o 2
01i-1 f/|-11
(217)
2
Translating to (Ht,i o of)) gives the BIRRDIAG update
He Pr - ao. - (|-1 (Mije1 - Mr.i-1) + ao? 1
(218)
0 2 = (1 a) - ao 1|11 a diag (Gr,s)
(219)
This update takes OMP)) per iteration to estimate G, using G.IC-EF O(NMP) per iteration using G G-MC-HES and Hutchinson's method, and O(P) per iteration using G GIIN-EF
The MC-HESS and MC-EF versions of this update are respectively equivalent to VON and VOGN Khan al. 2018b) the batch setting where we replace -p..- with a spherical prior A( - (see
Appendix E.6
33

---

E.3.4 BLR-LIN-HESS Diag
Applying Proposition 4.2 to Eqs. 218 and 219 gives the BLR-LIN- BLR-LIN-HESS-DIE update
Pt, = 1 + t0 i d/t [11p11 Mt,i-1) + ao (H), R (ye (c.;)) (220)
o - = (1 (1- 0 No f|t-1 + a diag (H) H.:)
(221)
This update takes P) per iteration.
E.3.5 BOG Diag
Substituting Eqs. 206 and 207 into Eq. 31 gives
(/ w/ f|t-1 + 00
11t-19t
(222)
v/ (2) 4 (2) 200
1 + ao 1t diag (G;)
(223)
Translating to (14, of) gives the BOG-DIAG update
14 o. 1 -
Mi ||- ao of 11--19
(224)
9
2 = 0 ||-1
2 40 2/1 t/t- 1 - 20 (G,)
(225)
This update takes (M to estimate G2 using G.c-nrr O(NMP) using G. G.c--uss and Hutchinson's method, and O(P) using G GUIN-Er
E.3.6 BOG-LIN-HESS Diag
Applying Proposition 4.2 to Eqs.
224 and 22.5
gives the BOG-LIN-HESS-DIAC update
14 = 0..0 11- 1 + ao; 1 H
R. (ye ÿt)
(226)
σ f|1-1 - 4 i M (H.R, (y. .i)) + 200 ddaag
diag (H.R," H.)
(227)
This update takes P).
E.3.7 BBB Diag
Substituting Eqs. (006) (209) into Eq. gives
t.1 ( + oo 19-1 - ao 1,1- ( (14 - M()(-1)
(228)
w/ [b + 200 2.0-1111144 11-1M1,1-150 + to (Gr,s)
200 2, 1 P4i-1 (M1111 M1j--1 a0 (0/11-1 - -0 -0
(229)
Translating to (14,i E.,)) gives the BBR-DIAG update
Pt. 0, 0
1,i 14 + aσ
2
2, .. 2 19
i a 1. -2 q|i-1 (Hip--1 11,1-11
(230)
-2
2 4uo 2
1r. 200 1, 2, diag G.,z)
1
1.:-
+ 4a0
2 c|t-1 - Mt,i-1 (Mr,i-1 ( + 200 200
2
1,1-1
i-i-1
0
f/t-1
-
o
(231)
This update takes OMMP) per iteration to estimate G, using G, G. O(NMP) per iteration using GMC-HEES and Hutchinson's method, and O(P) per iteration using GLIN-EP
E.3.8 BBB-LIN-HESS Diag
Applying Proposition 4.2 to Eqs. 230 and 23 gives the BBB-LIN-HESS-DIA. update
it, 0
+ ao
1 [H] R (ye (r,))
0 i 1
r|c-1 ( 11,i-1)
(232)
2 o 2 4
HH,, (y: (r.))) + 200 2000 diag (H H..)H..))
o,,
-i- Mr.i-
400 - 1 Mن:1 (H 1qj-1) 200 1 6j(-1 220
(233)
This update takes ( per iteration.
34

---

E.4
Diagonal Gaussian, Moment parameters
Throughout this subsection, vector multiplication and exponents are elementwise.
The Bonnet and Price theorems give
Pt. Bo [log (3/ ft (0.)) 9t,
(234)
? Ea Ea Dogg (y: Se (0t)) 7dliag (Gr,))
(235)
From Appendix E.3 we have
V Me.s DEL (ppri [1w.)=))) = o 2
1|t-1 (Mi,i-1 - Mi((-1)
(236)
V Voi.- Der.
Mapie .) = 2|t-1 - -1,3-11
(237)
We write the Fisher with respect to the moment parameters 2b (H, σ²) as block matrix:
F pip
F Fu,0²
F =
(238)
o2
FoR.02
The blocks be calculated by the second-order Fisher formula
F Fr. = - Egg ( VM log (0))
(239)
Diag (o
(240)
F Fu.0² -Bgs (V (Vn.zz log 1/ (0(()) (0)
(241)
-gg [Diag ((0 11)0r ")]
(242)
=0
(243)
F F -Egs 02 log (0))
(244)
-Egs (Diag (- (1 (-H
0)²
2
(245)
2 .iDiag (o 4)
(246)
Therefore the preconditioner for the NGD methods is
Diag (oi,--)
F
(247)
0
2Diag (ot,,-1)
E.4.1
BONG Diag, Moment
Substituting Eqs.
234
(25
and (247 into Eq. 3 gives the BONG-DIAG, update
μ: Mtlt-1 + 19t
(248)
of o /|| + o , diag (G.)
(249)
This update takes O(MP) to estimate G, using G.M-EE O(NMP) using G.C-HESS and Hutchinson' s method, and (PP) using G ILI--PP
E.4.2 BONG-LIN-HESS Diag, Momemt
Applying Proposition 4.2 to Eqs. 248 and 249
gives the LONG-LIN-HESS-DOE update
14 Pitt- + f||-1 H (y: 9.))
(250)
o: o ip-11 ft diag H.R," (H,
(251)
This update takes P).
E.4.3
BLR Diag, Moment
Substituting Eqs. 234
237 and 247 into Eq.
(32
gives the BLR-DIAG.MOX update
M1, Mt + ao
1,1-101 0 (Myz-1 (e,i-1) + aot 119t
2
(252)
0 = - + a (o
o
(|--1) + aot (Gr,i)
2
(253)
This update takes (MPP) per iteration to estimate G, using G O(NMP) per iteration using GMC--ESS and Hutchinson's 's method, and O(P) per iteration using GLI--Lr
35

---

E.4.4 BLR-LIN-HESS Diag, Moment
Applying Proposition 4.2 to Eqs.
252 and 253 gives the HESS-IGGSS-- MOM update
N1, Mr,i-1 + ao a0 1 11-1 ( ((,i-1) + aoi
2
HH R ( vt,)))
(254)
o = o + a0 1 ( - 1je-1) n ao diag
-
-2
(255)
This update takes O P) per iteration.
E.4.5 BOG Diag, Moment
Substituting Eqs.
234 and 235
into Eq. 31 gives the BOG-DIAG.MOM update
14 M + ag:
(256)
o; i 2
a
diag (G.)
(257)
This update takes ((MP) to estimate G, using G. CO using G.C-HISS and Hutchinson's method, and ((P) using GLIN-rr.
BOG-LIN-HESS Diag, Moment E.4.6
Applying Proposition 4.2 to Eqs. 256 and 257 gives the BOG-LIN-HESS-DIAG... MOM update
Maj- + aH, (y: Un) o, o a diag H. H,
(258)
(2559
This update takes PC.
E.4.7
BBB Diag, Moment
Substituting Eqs.
234 to
237 into Eq.
34
gives the BBB-DIAG MOM
Mt M4 a0 (Mijo-1 (1,i-1) + uggt a
(260)
o -i.-1 2 a 0((-1)
a
diag (Gr,s)
(261)
2
This update takes O(MP) iteration to estimate G, using G.N-Er O(NMP) per iteration using GIN--IES and Hutchinson's 's method, and O(P) per iteration using GUIN-EP
This is similar to the original diagonal Gaussian method in Blundell et al. 2015 except (1) they reparameterize o log( exp(p)) (elementwise) and do GD on (μ,, p) instead of (μt, σ²), and (2) they use the reparameterization trick instead of Price's 's theorem for calculating the gradient with respect p (via σ))
E.4.8 BEB-LIN-HESS Diag, Moment
Applying Proposition
4.2 to Eqs. 260 and 26
gives the BBB-LIN-HESS-DIAN MOM update
Mr Mt-i + a0 1|-1 (Mije1 ((,i-1) aH! R (y: "r.i)
(262)
o; -i + a
2 ( - -2 a
d 2
diag H HH,,R,
R R H.,)
(263)
This update takes O P) per iteration.
E.5 Diagonal plus low rank
Assume the prior is given by
P (0.) = =N (0.lMe (Teit-1 + Wii-1 (o--
(264)
36

---

with W E RRPR and diagonal T ft- E RPPP. We sometimes abuse notation by writing Y 1|-11 for the vector diag (T when the meaning is clear from context.
Substituting the DLR form in the gradients for the KL divergence derived in Appendix E.2 gives
V p DKL ( (1 1 (T (Tac-1 + Witr- W/(-1) (142211 - 11:-1)
(265)
V DEL (W ( = (aal--1 - M + Weft-1 W Wirp-1 - Wos-IWE.
(266)
For any function é the chain rule gives
Vr,,-- -diag ((III-1+ + Weje- W Wi(-1) "7(-=)"((')))) (Vey)-)) (Tell-1 + W e/c-1
(267)
( Wij- W (VEnc-1 (Tqje11 W (ill--W
W =je-1
(268)
Therefore the gradients we need are
V Fu EE,-,... plogp (yelh, (0.)) = I
(269)
(Tr.i- Wt.- WT (1-1
G
Vr... Eg. plog (,e (0,)) - dddaggg
x (T W. 1,1-1
(270)
Vw.,, Eo..eo... plogp (Yelh, (0.)) = - (M + WT, G.,
x (M + Wii-i. W
(271)
and
V DEL (1oberr 11w))r-))
(T + 111-1+Wei-- WT, -1
r T1j-11 Yr,i-1 + Wif-1 WT 1|-1 - W1,-1 W )
diaa
(272)
(Tr,1-1 + W7 (--1
Ww. DKI.(au.
(M W W WT 1-11
x (T([-11 - T + W Wip-
WT M|-11 W WT,
x (T + Wt,51 WT,-) W
(273)
The Fisher matrix can be decomposed as a block-diagona with blocks for μt. and for (Y..s. W..). but (in contrast to the FC Gaussian case in Eq. (168) we have not found an efficient way to analytically invert the latter block, which has size P + RP. To avoid the O (R* cost of brute force inversion, use different strategy for BONG and BLR of performing the update on the natural parameters of the FC Gaussian as in Appendix E.T and projecting the result back to rank R using SVD. Specifically, assume the updated precision from applying Eq. (12)) for BONG or Eq. (366 for BLR be written as
E T.s + W. WT,
(274)
and let the SVD of We, f. be
W = U,,A, VL, with U J, Vt, orthogonal and A rectangular-diagonall Following Mishkin et al. 2018 Chang
(275)
et al. 2023 we define the update
W U. ::: R A |: R.: R]
(276)
Y Ye. + diag WT, - W.,W...)
(277)
37

---

so that Wi contains the top R singular vectors and values of the FC posterior and the diagonal is preserved: diag (Tr, W. diag [ This approach works for MC-EF and LIN-EF methods but not MC-HESS, which we omit.
Finally, in the MC-EF methods we sample from the DLR prior using the routine in Mishkin et al.
2018
Lambert al. 2023 which takes O M)P).
E.5.1
BONG DLR
Substituting the DLR prior from Eq. (264) into the FC precision update from Eq. 121 and using the MC-EF approximation yields the posterior precision
--1= -u--- + Wij-1 WT
G.
GM--EF
(278)
i|t-1
=r,+ W.W
(279)
r. -Y11-1
(280)
W, WaR-1
G
(1:M)
(281)
VM
Note W. (P Using this posterior precision in the mean update from Eq. 120 yields
(T (T
w.W?) -
F4 - Mij-1++ +
g:
(282)
= -iilt1
r 1
M W. (IR-M+ WT 1/-1 (.) w.)
w Hjt-1
gr
(283)
where the second line comes from the Woodbury identity and can be computed in O(R+ M)P (R+ - ()), Applying the SVD projection gives
W, = U, l,, 18 A. R;; RR]
(284)
w.W!)
r. = Teje-1 + diag
(285)
(U, A,, _ = SVD (w.)
(286)
which takes O(RR 2) for the SVD. Therefore the BONG-MC-RF-DLR update defined by Eqs. (81) and (283 to 286) and takes O(RR My² + (R M))).
The BONG-LIN-EF-DLR update comes from replacing - G (1:M) with gi in Eq. 281 and replacing / IR+M with LR+1 in Eq. 288 This update takes O( 1)2 (R 1)3).
E.5.2 BONG-LIN-HESS DLR (LO-FI)
Applying Proposition 4.2 to Eqs. 281 and 283
gives the BONG-LIN-HESS-DDE update:
He = 1 + r - We + W/ w/r 1
x H,R," -1 (y,, v.)
(287)
W.- le R A₁ R.: R]
(288)
Y.- = T:p-1 + diag (W.WT- w.WW!)
(289)
(U,, A,, ) = SVD (W:)
(290)
W. = [W 11 H[A
(291)
A, chol (R.
(292)
This is equivalent to LO-FI Chang et al. 2023 and takes 0( (R (R+C)3)
38

---

E.5.3 BLR DLR (SLANG)
Substituting DLR forms for 1 and - into the FC precision update from Eq. (136) and using the MC-EF approximation yields the posterior precision
Σ = -1 a1 (Tr,-1 Wow-1W +a (xpe-1 Wej- W
oG
t,
MC-FF
(293)
= -T WL:
(294)
Yus = (1 x (1-a)T + ah
(295)
Wi, VI- W. 1- va d V
(1:M)
(296)
M 2,
Note W. R Mx Using this posterior precision in the mean update from Eq. 135 yields
Nr |4 + (r,, W.,WT,
((Te)-1 W W W (Haee-1 -1,,-1), gr.s)
(297)
-M + - W.. [12R+M + WT Y W1,)
(WW,, (Tett-1+ Wii-i WT (Hie-11Mi -
WTT
gt.i
(298)
where the second line comes from the Woodbury identity and can be computed in 0 (2R + 3)) Applying the SVD projection gives
P
Wi, = U, : R] A [: R,: R
(299)
Y Yr, + diag ( WT, WT,
(300)
(U,,, Ar,, = SVD (Wr.)
(301)
which takes O(2R + 2 P) for the SVD. Therefore the BLR-MC-EF-DLR update defined by Eqs. (25). (96) and 2998) to (30} and takes O((2R + M) (2R + M) per iteration. Notice
W, has larger rank for BLR than for BONG (2R M R + M) because of the extra Va fjt-1 term that originates in the KL part of the VI loss. This difference will slow BLR-MC-EF-DLR relative BONG-MC-EF-DLE especially when R is not small relative to M.
This method closely resembles SLANG Mishkin et al. 2018 in the batch setting where we replace - with spherical prior N (0, X Ip) (see Appendix E.6
We can also define a BIR-LIN EF- EF-DIR method by replacing / M G (1:M) with vag! in Eq. (96
and 12R+M with 128+1 in Eq. 298 This update takes O(2R + 2 + (2R + 1)3) per iteration.
E.5.4 BER-LIN-HESS DLR Applying Proposition 4.2 to Eqs. 296 and 298 gives the BLR-LIN-HESS--LR update:
M1. = 11.51 + a |. r W.. (IeRto + WT, Y. W..)
WT, f. r.
x ((re)1-1+
W (111--1-11 + HT R (y: 01.))
(302)
W., U2, [:: Ar |: R,; R]
(303)
Ye, Te, diag Wi, WT. WT,)
(304)
(U,,, A )) SVD (W..)
(305)
W [V1 a W W VaW /HH
(306)
Ar, chol (R
(307)
f (1 a) Y + 1
(308)
39

---

This update takes O + C) + (2R+ ()*) per iteration. As with the EF versions of BLR-DLR, W, has larger rank for BLR-LIN-HESS-DL.R than for BONG-LIN-HESS-DLR (2R + C vs. R + C). This difference will slow BLR-LIN-IIESS-DLR relative to BONG-LIN-IIESS-DLR especially when not small relative to C.
E.5.5 BOG DLR
Substituting Eqs. 269
into Eq.
gives the BOG-DLR update
Majt-1 1 + tag:
(309)
Y, Y cj1-1 2
0t <diag (Ti/-11 + Waj-1 (ic--TTe- W
GG G. (Y -
W We WT W ejt-1
(310)
W, = Wit1- - (T WT G.
Wij--1
(311)
Using the EF approximation and Woodbury, the BOG-MC-EF-DLR update can be rewritten
Y. T1/=-1 2M
diag (B,B,)
(312)
W, Wije-1
B.B.W B.B.War. W 6/--1
(313)
M
B. rx_ - TT-1 d/--1 Wii-1 (IR + WT
f/--1 wij-1 - f/r-1 W(sc-)) Wee-1 x
-1
G (I:M)
1/6-1
e/e-1
(314)
which takes ORRMP+ MR + R²))
The BOG-LIN-EF-DLR update comes from replacing G(LMM) with g, in Eq. (314) and dropping the M - factors in Eqs. 312) and (113). This update takes ORRP R2).
E.5.6 BOG-LIN-HESS DLR
Applying Proposition 4.2 to Eqs. 309
and 312) to 314 gives the BOG-LIN-HESS-DIE HESS-DLR update
F: = Md + R, (y: 9.)
(315)
r. T: + diag (B.BT)
(316)
W,- Wil-1 + a.. B. W :|c-1
(317)
B: = (T r T W W IR + WT W W.TTWeee
-1 WT W x 1/c-1
H
(318)
This update takes OCC(C++ R)P + CR + R²))
E.5.7 BEB DLR
Substituting Eqs. 265 and 269
to (233 into Eq. 334 gives the BBB-DLR update
H4. 11,5-1 +
T T W f/|-- WT 1 (Hij--1 ((,i-1 + ag:
(319)
Ye.: T1.i11
= diag Ex,i-1 ( Yr.i- Wii-1 WT - W WT,
G1.e
Er,i-1
(320)
Wt, W1.-1
0 (T11-11 Y + Welo- WT ffr-1 W..- WT W.,11 -G.s) E W. t,1-1
WT
(321)
40

---

The previous covariance can be written using Woodbury as
E = 1--1111-1 - Wi.i- (In
+ W M W 4i-1
(322)
The BRB-MC-RF-DL.R update be computed efficiently by expanding terms in Eqs. For example the terms involving G can be calculated as
(320) to
(22)
r t,i (1:M) (1:M)
diag (-Ec,i-1
Gt,
GMC-ER E(ii-1) -e
1 M
diag
r - (
(1:M)) B
(323)
B, GBa
Y T Li-1 + B
GCCCER -E,,-1G It,i-1 W
1
r
G (1:A))
Tú W
M
r,i-
(1:M) G. B W
M
ti
1
B
T (1-1 W
B,, W (324)
M
M
B T7 W (IR + WT 6, Y 77,1 W (w,-1)
(1:M) - - 1
(325)
Using this strategy the update takes O(RR + MP + MR + R*).
The BBB-LIN-EF-DLR update comes from replacing G (1:M)
with n. and dropping the M- factors
in Eqs. 323 to 325
This update takes O P R2).
E.5.8
BBB-LIN-HESS DLR
Applying Proposition 4.2 to Eqs. 319 to 321
gives the BBB-LIN-HESS-DLR update
H4, Pr.i-1 + a (qp-1++ W (|- WT (1-) (M - + aH! R (v: - Ūt.i)
(326)
Y [|| - Y - + WIj- WT i1-1
)
Te, = Y diag Et,i11
(327)
2
W W WT,-1 + R H
T1/1-1 - T , W Wij-1 4:-1
W,, W
E W
(328)
W1,1-1 WT i,i-1 H., H
This can be computed in O(CC + 2 CR + R²5) using Eq. (322) and following a computational
approach similar to Eqs. 323
335
E.6 Batch BLR
It is interesting translate the BLR updates derived here back the batch setting where BLR was developed [Khan and Rue 2023 by replacing / (Haje- E1j-11 with a centered spherical prior A/ (O.)-Ipp).
The batch BLR-FC update becomes
#4 = 11-1 aE, (g: -1p1-1)
(329)
- (1- - + a (AIp G;)
(330)
The batch BLR-LIN-HESS-FC update becomes
#4: p:-1 + aΣ, H (y: 4.)) \p-=11)
(331)
Σ = (1- (AIr + H. R, HH,)
(332)
The batch BLR-FC MOM update becomes
[g \p-=1) μ: = Pi-1 + a a
(333)
Σ, (1 (1+ aΣ (G; >Ip) E,-1
(334)
41

---

The batch BLR-LIN- update becomes
N4 pi + aZ R,, (y/- ÿ:) A
(335)
Σ. (1) aZ,-1 (MI H,R," (R,, Ei-1
(336)
The batch BLR-DIAG update becomes
μ4 Pi-1 + oo (g: Ap==1)
(337)
o, 2 = (1 o) + a diag (AIp - G.)
(338)
The MC-HESS version of this update equivalent to vON [Khan et al. 2018b] we use MC
approximation with M = 1. The MC-EF version with M = is equivalent to VOGN [Khan et al. 2018b]
The batch BIR-LIN-HESS--DAAG update becomes
μ: = p4-1 ao? [H (3: - -..) Apt-1)
(339)
2 (1 al0 + diag (AIp H.R, H.)) 0,
(340)
The batch LAR-DIAGGMOM. update becomes
μ4; Pi-1 ao (g. Api-1)
(341)
o (1 [ aaoo ao diag (G; AIp)
(342)
The batch BIR-LIN-HESS-DIAS MOM update becomes
μ4 P41 + ao? R.,R, (y: (.)) \pp--1)
(343)
of (1+a)o diag (AIp H. R2 H.))
(344)
The batch BLR-DLR update becomes
14-1 + 1-11-r-11 (IRMW TeeM-W/Y,, "
x (g: - Appi-1)
(345)
W; U, R2] A, |: R, RR]
(346)
Y: Y, diag
W. W.W. w.w!)
(347)
(U,, A, ) SVD (W.)
(348)
W, VI aW
i-1, aW...
((:M)
(349)
A, = chol (R=1)
(350)
r, (1 r-)
(351)
This is equivalent to SLANG except for the following differences. SLANG processes a minibatch of M examples at each iteration, using single sample Ô ~ 4w.., for each minibatch. It uses a different SVD routine which slightly faster but stochastie, taken from [Halko et al. [001]] Most significantly, SLANG applies the SVD before the mean update, meaning μt; is calculated using the rank- R W, and Y; instead of the rank-(R R + M) W; and Y; thus ignoring the non-diagonal information in the M discarded singular vectors from W..
42

---

The batch BLR-LIN-HESS-DLR update becomes
- 14-1 + (r-)) - -r. -r,,,, (IR+C WTT. "
x (H,R, -1
; "y1 Úi) Ap--1)
(352)
W, = U, [:: R A; R,: R
(353)
Y; Y; + diag
WW.W W.W!)
(354)
(U, A _ SVD (W:)
(355)
W; [VI aW 1-1, V
(356)
A, = chol RR,')
(357)
Y, (1 a T + a
(358)
This algorithm could be called SLANG-LIN-HESS and would be deterministic and faster than SLANG since it does not need MC sampling. We can also define SLANG-LIN-BE which would be even faster, by replacing 15 G M
(I:M) with Vag! in Eq. 349 and IR+ with L in Eq. (35).
43

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Section D states two main theoretical contributions (proved in Propositions 4.
and and summarizes the experiment findings reported in Section S
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss limitations in Section 6 We report asymptotic efficiency in Table and report actual running time in Fig.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model cell-specification. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
44

---

Answer:
[Yes]
Justification: We state two formal theorems and provide complete proofs in Appendix C which sketch in the main text in Sections 4. and 4.4
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We give pseudocode for the main methods in Appendix and complete update equations for all algorithms in Appendix D We give full details of the experiment methods
in Section
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
45

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We will post our code on github after the blind review period, including scripts for exactly reproducing all experiments.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://niis.cc)
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/nips. wttpcc/pcppcppccccceeeeooop /eddeeeeeeoemmiiiccccccc cetisubimsssiinnio seriissssoopoiiii
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: p links to the data and reports hyperparameters and architec- ture details.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification: Plots show 1 SE shading based on empirical SD across trials. Trials vary randomly in network initialization, data ordering and MC samples. This is stated in Section
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
46

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for caleulating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (c.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We report running times. Experiments all on any standard GPU/TPU so further details are not necessary.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
We have read the Code of Ethics and affirm our research adheres to all points
therein.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: We believe there are no direct societal impacts (positive or negative) beyond the generic impacts of foundational work training neural networks.
Guidelines:
47

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g.. gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
We provide citations and URLs for the two datasets we use Section
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
48

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
batasets.papersstt
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors encouraged to reach out the asset's s creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: There new data. As noted above, will post model code after the blind review period.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
49

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
50


Reviews:
[{"review_id": "SFa7ySWiMb", "reviewer": "Reviewer_ch8o", "summary": "This paper proposes a generalized framework for variational sequential inference of Bayesian neural networks with Gaussian prior distributions by further approximating the KL-divergence and expectations. It presents a very elegant unification of a wide range of (approximate) variational inference algorithms for Bayesian neural networks developed over the past 10 years and contains an extensive empirical evaluation on subsets of the MNIST dataset.", "questions": "Page 4: The abbreviation PSD as \"positive-semi-definite\" has never been introduced.\nPage 5: jac has not been introduced\nPage 7: In line 239, the component-wise product of \\sigma^{-2} and \\mu should be indicated somewhere", "limitations": "This is not applicable.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The strongest point of the paper is the unification of many previously presented approximations as well as the methodological evaluation and comparison among them.", "weaknesses": "The space of 9 pages is way to short to do this full justice and I often had to read into the 30 pages (!) appendix. I recommend the authors to further considering a journal submission.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback and positive evaluation. We will consider submitting to a journal in the future, but we feel that NeurIPS has higher visibility. We agree the unifying framework makes a good contribution on its own (and we have more theoretical work in progress on these lines) but for NeurIPS we want to highlight the proposed BONG methods and their empirical success.\n\nWe will explain PSD, jac (Jacobian), and the componentwise products in the revision. Thanks for catching those.", "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": "", "extra_scores": {}}, {"review_id": "fAWyISwxTC", "reviewer": "Reviewer_6NTx", "summary": "This paper proposes a new approximate Bayesian technique specifically for online learning whereby posterior distributions over modeling parameters at time $t$ can be achieved with a single natural gradient step evaluated on a model using the previous posterior distribution at time $t-1$ as the prior distribution. This method is rigorously derived in a conjugate setting with an exponential family and extended approximately with various ablations to more expressive black-box settings. Extensive empirical studies are conducted validating the performance of the method.", "questions": "No other questions, please address the main weaknesses. Please let me know if I missed or misunderstood anything; I am more than happy to be convinced otherwise.", "limitations": "The authors adequately addressed the limitations of the proposed method.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "An interesting connection is made for online learning in the conjugate setting allowing for dropping the KL regularization term and producing posterior updates with a single natural gradient step. As mentioned in the summary, various ablations are explored which trade-off precision and computation speed. Finally, many extensive experiments are conducted comparing this method to other Bayesian online learning approaches.", "weaknesses": "The main weakness of the method is the lack of experiments evaluating the resulting model uncertainty achieved after training. From what I read, it seems to only be predictive performance via misclassification and negative log predictive density. Granted, the latter does have some ablations where it is integrated over posterior samples; however, the presentation of these results are a bit difficult to interpret other than just methods with lower values being better. It would be interesting to investigate the predictive uncertainty via expected calibration error, or some other calibration metric.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful comments.\nAs requested, we have attached some figures that plot the ECE (expected calibration error) vs time for the various methods, when evaluated on MNIST using a CNN. Fig 1 shows that the BONG method is (slightly) better calibrated than the other\nmethods, for small sample sizes, but not surprisingly, all methods converge to similar results once they have seen enough data.\nWe will add these to the appendix in our revision.", "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": "", "extra_scores": {}}, {"review_id": "8NZiq46hhr", "reviewer": "Reviewer_ARk4", "summary": "The paper proposed a novel parameter update rule. The rule is natural gradient descent on the variational inference loss, where the prior term is dropped. Such rule is motivated in theory under the assumption of likelihood and variational distribution being conjugate, and in practice though experiments on MNIST.\nSeveral variants are considered: various loss, various hessian approximations and various way of dealing with the intractable expectation.", "questions": "The notation $f_t$ that appears in Equations (9) and (10) was never introduced. I guess it refers to the evaluation of $f$ on $x_t$, is it correct?", "limitations": "Limitations are well discussed. Good job on this point", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 1, "strengths": "The paper is well written and clear.\n\nThe setting is explained properly and the related work are exhaustive.\n\nA lot of variants of the algorithm are considered. These are well explained and the previously known methods are highlighted. This schematic approach is very valuable for two reasons: \\\n-it helps in better understanding the proposed algorithms \\\n-it creates a clear picture of the field of approximated bayesian inference, putting in perspective every method with each other\n\nIn my opinion this is a contribution that may be worth acceptance, not the proposed update rule.", "weaknesses": "The theoretical motivation is quite weak:\\\n-Proposition 4.1 is nice, but the assumptions are overly restrictive. It feels like this is such a special case that is has no implication on real case scenarios \\\n-The claimed contribution (in Line 44) is a very weak contribution. The closed from expression with linearization+gaussian likelihood was already observed by Immer, and I don't think it is fair to claim it as a novel contribution of the paper. Specifically because there is no motivation to such approximations.\n\nOn the other hand, as the authors also correctly point out in Line 307, the empirical evaluation is based on MNIST only. A very small dataset which doesn't really motivate the superiority of the proposed algorithm.\n\nOverall, I don't think the authors provided enough evidence (either theoretical or empirical) of the good performance of the proposed update rule.\n\nTYPOS: \\\n-Line 126. I think the $\\theta_{-1}$ was supposed to be a $\\theta{t-1}$ \\\n-There is a duplicate reference: Line 336 and Line 340 are the same", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback. We have a few responses below.\n\n“Proposition 4.1 is nice, but the assumptions are overly restrictive. It feels like this is such a special case that is has no implication on real case scenarios.”\n\nMost efficient VI methods exploit tricks from conjugate Bayesian inference. Thus we believe that Proposition 4.1 provides a useful motivating foundation for our approach, and ensures that our method is exact in certain simple cases. The experiments complement the theory by showing our method also works well in more general settings.\n\n“The closed form expression with linearization+gaussian likelihood was already observed by Immer”\n\nWe agree line 44 should be revised to better explain our second contribution. (Please note our primary theoretical contribution is replacing the KL in the ELBO with implicit regularization from one-step updating, which is definitely new.)  We cite Immer for the linear(f) approximation when we first introduce it (line 51). Our contribution here is that the update rule for CM-EKF (which we extand to other BONG variants) follows from two different approximations. One is known (and we credit Immer, Ollivier, and Tronarp et al.) while the other is new. Informally, showing the same algorithm arises in different ways from different approximations adds motivation for the approach. Further motivations for the linear(f)-Gaussian approximation are in Immer and Tronarp et al., which we could explain briefly in our paper.\n\n“the empirical evaluation is based on MNIST only”\n\nWe also report experiments using the SARCOS dataset (appendix B4-B10). We agree it will be important to scale up as we continue this line of research, but we believe the current experiments are sufficient for a primarily theoretical paper.\n\n“I don't think the authors provided enough evidence (either theoretical or empirical) of the good performance of the proposed update rule.” \n\nWe believe the experiments show a clear advantage of the proposed update rule, especially when taking compute costs into account. We agree that MNIST is a simple dataset, but we did not have time for more experiments, since our main contribution is the theoretical framework. \n\n“$f\\_t$ in Equations (9) and (10) refers to the evaluation of $f$ on $x\\_t$?”\n\nYes we will add this.", "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": "", "extra_scores": {}}, {"review_id": "HzPiBhB9nF", "reviewer": "Reviewer_uSab", "summary": "The authors explore sequential Bayesian inference using variational inference (VI). They propose to remove the regularizing KL term, performing a single natural gradient step using the expected log-likelihood only, and provide relevant ways to approximate the hard-to-compute expectations. They support their approach with detailed experiments.", "questions": "NA", "limitations": "NA", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper is sound, and the detailed experiments provide strong support for the method proposed in the paper. I also appreciated the nice motivation involving conjugacy, as well as the algorithmic considerations.", "weaknesses": "My main concern is about the claimed novelty of the approach (\"removing the KL\" and \"implicit regularization towards the prior\"). The idea of using a one-step natural gradient descent on the expected log-loss directly (without the KL term) has already been proposed in the literature and analyzed; for example, [Lyu & Tsang, 2021] studied this algorithm in a model-free setting where a general function f is considered instead of the log-loss. This one-step natural gradient update is in fact simply the solution of Equation (2) of the present paper, but with a linearized expected log-loss. This dual interpretation resonates with the dual interpretation of standard Euclidean gradient descent (single update or minimization of a first-order development of the loss to be minimized). \n\nY. Lyu and I. W. Tsang: Black-box optimizer with stochastic implicit natural gradient. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pages 217–232, 2021.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback. \n\nWe were not aware of the Lyu & Tsang paper. We will add discussion of it to Section 2 but also want to point out some important differences: Their goal is optimization, not inference or Bayesian updating. Thus the KL regularizer is not part of their objective whereas it is for us. Instead their regularizer is an auxiliary device to obtain an incremental algorithm. (Also note that iteration number in their optimization setting is not the same as time step in our filtering setting.) Then they use the following two well-known facts (e.g., Raskutti & Mukherjee 2015) to obtain an NGD update:\n\nA) Mirror descent (MD) can be derived by optimizing linearized loss plus Bregman divergence (we believe this is the dual interpretation you refer to). That is, given parameters $\\\\boldsymbol{w}$, loss $\\\\mathcal{L}$, and Bregman function $\\\\Phi$, the MD update is equivalent to $\\\\boldsymbol{w}\\_{i+1} = \\\\arg\\\\min\\_{\\\\boldsymbol{w}} \\\\langle\\\\boldsymbol{w}, \\\\nabla\\\\mathcal{L}(\\\\boldsymbol{w}\\_i)\\\\rangle + D\\_{\\\\Phi}(\\\\boldsymbol{w}\\_i,\\\\boldsymbol{w})$.\n\nB) NGD on the mean parameter of an exponential family is equivalent to MD where the Bregman function is the log-partition.\n\nNevertheless their approach is closely related to an alternative MD-based derivation of BONG that we plan to add (this has already been written for an arXiv version after our NeurIPS submission). Using facts A and B above we show BONG can be obtained by linearizing the expected NLL wrt the variational mean parameter ($\\\\rho$). This is the same observation you make about our Eq (2) and linearized expected log-loss. But to be clear, connecting all these ideas to the ELBO is new in our paper and leads to a novel approach to variational inference.\n\nWe believe our paper has multiple contributions, including a unified framework, a novel update, and extensive experimental comparison of different methods. Even if the BONG update equation is similar to algorithms used in other domains (such as Lyu & Tsang’s) its derivation and application to variational inference are novel, and moreover we believe it is just one piece of our overall contribution. We therefore respectfully ask you to consider raising your score.\n\nRaskutti, G., Mukherjee, S.: The information geometry of mirror descent. IEEE Transactions on Information Theory 61(3), 1451–1457 (2015)", "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": "", "extra_scores": {}}, {"review_id": "SCmeJISydl", "reviewer": "Reviewer_zVLu", "summary": "The paper introduces a framework for variational online learning, in which a method is defined by three ingredients: (1) implicit or explicit regularisation (2) type of approximation of the expected gradient and hessian (3) geometry in which the gradient is taken (ie employing natural gradients or standard gradients). In implicit regularisation, one learning step is taken for each datapoint, with no regularisation; explicit method is the more standard approach of optimising with multiple steps and regularisation to the posterior at the previous step. A case study on different combinations of the ingredients is conducted; some of the methods appeared in previous papers. The study finds the best performing method is the one with implicit regularisation, linearisation of the hessian, and natural gradient.", "questions": "* Given the well documented efficiency of natural gradient steps, is there a danger of the method with no regularisation being more prone to \"forgetting\"?", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is a valuable contribution to the field, offering a framework for disjointed methods. It reads like a review paper in variational online learning, with a novel contribution based on the review. The list of methods of approximation is thorough and the comparison exhaustive.", "weaknesses": "Given that the main strength of the paper is the unifying recipe, and the comparison of using different ingredients, a more exhaustive case study would strengthen the work (and the less essential comments, like the equivalence to the full Bayesian update in the conjugate case, in my opinion can be placed in the appendix and referenced). Nevertheless, I find the contribution of the paper warrants it being accepted.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive feedback.\n\nWe appreciate the suggestion to report a more exhaustive case study and give less space to the motivation from the conjugate case. The reason we prefer the present structure is that it puts more weight on the new algorithms. We’re glad you find the organizing framework and review-like aspects of the paper to be a good contribution (we agree), but we also want to highlight the BONG method and its empirical success.\n\n_Can NGD promote forgetting because of its efficiency?_\n\nThis is an interesting question. We think the right way to think about it is by comparison to the exact Bayesian update, which optimally maintains past information (i.e., forgets just the right amount). Our methods are approximations of the exact update. So even if the NGD step went all the way to the minimizer of the variational loss (as it does in the conjugate case) it should not overlearn or excessively forget past observations.", "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": "", "extra_scores": {}}]

Author rebuttal:
We thank all the reviewers for their useful feedback. We give individual responses  below.
As requested by reviewer 6NTx, we are also attaching some figures that plot the ECE (expected calibration error) vs time for the various methods, when evaluated on MNIST using a CNN. Fig 1  shows that the BONG method is (slightly) better calibrated than the other
methods, for small sample sizes, but not surprisingly, all methods converge to similar results once they have seen enough data.
We will add these to the appendix in our revision.

Meta-review: This paper studies the problem of online Bayesian inference in Bayesian neural networks using variational inference. The paper considers an online natural gradient update rule given by only considering the expected log-likelihood term and using a unit learning rate. This Bayesian online natural gradient update rule is studied theoretically under a conjugate prior setting. Next, the paper proposes a deterministic way of efficiently computing the expectations in the update rule for Gaussian variational families. The method is then evaluated on a number of data sets, including the MNIST data set. 

Overall, the reviewers found the contribution to be interesting to the NeurIPS community. Reviewer 6NTx commented on a lack of experiments evaluating the resulting model uncertainty; these concerns were addressed during the rebuttal and discussion phase. Several reviewers brought up issues of novelty and wanted to see additional discussion that much more explicitly discusses the contributions of this paper in the context of related update rules in the literature. Please revise the paper to include these discussions in the final version.

Based on the reviewer comments, I recommend accepting the paper.
Decision: Accept (poster)