Original: On diagonal approximations to the extended Kalman filter for online training of Bayesian neural networks
Revised: On diagonal approximations to the extended Kalman filter for online training of Bayesian neural networks

Original: Abstract
Revised: Abstract

Original: We present two approaches to approximate online Bayesian inference for the parameters of
Revised: We present two approaches to approximate online Bayesian inference for the parameters of

Original: DNNs.
Revised: DNNs.

Original: Both are based on diagonal Gaussian approximations and linearize the network at each step to ensure efficient computation.
Revised: Both are based on diagonal Gaussian approximations and linearize the network at each step to ensure efficient computation.

Original: The first approach optimizes the exclusive KL,
Revised: The first approach optimizes the exclusive KL,

Original: D KL ( q ∥ p ); this amounts to matching the marginal mean and precision of p and q .
Revised: D KL ( q ∥ p ); this amounts to matching the marginal mean and precision of p and q .

Original: The second approach optimizes the inclusive KL, D KL ( q ∥ p ), which amounts to matching the marginal mean and variance of p and q .
Revised: The second approach optimizes the inclusive KL, D KL ( q ∥ p ), which amounts to matching the marginal mean and variance of p and q .

Original: The latter approach turns out to be equivalent to the previously proposed “fully decoupled EKF” approach.
Revised: The latter approach turns out to be equivalent to the previously proposed “fully decoupled EKF” approach.

Original: We show experimentally thatinclusive KL yields better results than exclusive KL, and both are more effective than  one-pass SGD.
Revised: We show experimentally that exclusive KL is more effective than both inclusive KL and one-pass SGD.

Original: 1. Introduction
Revised: 1. Introduction

Original: In this short paper, we consider the problem of online Bayesian inference for the parameters of a neural network.
Revised: In this short paper we consider the problem of online Bayesian inference for the parameters of a neural network.

Original: That is, we want to recursively compute p ( θ t | y 1: t , x 1: t ) in an efficient manner, where x t ∈ R N x are the input features, y t ∈ R N y are the output labels, and θ t ∈ R N z are the model parameters.
Revised: That is, we want to recursively compute p ( θ t | y 1: t , x 1: t ) in an efficient manner, where x t ∈ R N x are the input features, y t ∈ R N y are the output labels and θ t ∈ R N z are the model parameters.

Original: We assume the likelihood has the form p ( y t | θ t , x t ) =
Revised: We assume the likelihood has the form p ( y t | θ t , x t ) =

Original: expfam( y t | g − 1 ( h ( x t , θ t ))), where h is a neural network that predicts the natural parameters η of the exponential family output and g is some link function that maps from natural parameters η to moment parameters m .
Revised: expfam( y t | g − 1 ( h ( x t , θ t ))), where h is a neural network that predicts the natural parameters η of the exponential family output, and g is some link function that maps from natural parameters η to moment parameters m .

Original: To handle non-stationary distributions, we allow the parameters to drift according to a Gaussian random walk, so p ( θ t | θ t − 1 ) =
Revised: To handle non-stationary distributions, we allow the parameters to drift according to a Gaussian random walk, so p ( θ t | θ t − 1 ) =

Original: N ( θ t | θ t − 1 , Q t ), where Q t = diag( q t ) is the process noise.
Revised: N ( θ t | θ t − 1 , Q t ), where Q t = diag( q t ) is the process noise.

Original: For continual learning problems where a single task drifts over time, this method tracks the optimal parameters and their uncertainty.
Revised: For continual learning problems where a single task drifts over time, this method tracks the optimal parameters and their uncertainty.

Original: If instead we assume q t = 0, so the parameters are stationary, then the method can also be used to compute p ( θ | x 1: T , y 1: T ) with a single pass over the data.
Revised: If instead we assume q t = 0, so the parameters are stationary, then the method can also be used to compute p ( θ | x 1: T , y 1: T ) with a single pass over the data.

Original: Exact Bayesian inference in this model family is intractable, so we must resort to approximations.
Revised: Exact Bayesian inference in this model family is intractable, so we must resort to approximations.

Original: In the signal processing and engineering communities, a standard approach (e.g., Singhal and Wu, 1989; Puskorius and Feldkamp, 2003) is the extended Kalman filter (EKF), which amounts to linearizing the observation model h at each step, and then performing an exact Bayesian update.
Revised: In the signal processing and engineering communities, a standard approach (e.g., Singhal and Wu, 1989; Puskorius and Feldkamp, 2003) is the extended Kalman filter (EKF), which amounts to linearizing the observation model h at each step, and then performing an exact Bayesian update.

Original: In the machine learning literature, it is more common (e.g., Challis and Barber, 2013; Broderick et al., 2013; Nguyen et al., 2018; Lambert et al., 2021b) to solve the following variational inference problem, using exclusive KL divergence:
Revised: In the machine learning literature, it is more common (e.g., Challis and Barber, 2013; Broderick et al., 2013; Nguyen et al., 2018; Lambert et al., 2021b) to solve the following variational inference problem, using exclusive KL divergence:

Original: [Equation]
Revised: [Equation]

Original: where we omit the conditioning on x t for brevity, and where the one-step-ahead predictive distribution is given by
Revised: where we omit the conditioning on x t for brevity, and where the one-step-ahead predictive distribution is given by

Original: [Equation]
Revised: [Equation]

Original: Some prior work (e.g., Hern´andez-Lobato and Adams, 2015; Soudry et al., 2014; Ghosh et al., 2016) instead uses the assumed density filtering (ADF) approach, which uses the inclusive KL divergence:
Revised: A third approach (e.g., Hern´andez-Lobato and Adams, 2015; Soudry et al., 2014; Ghosh et al., 2016) based on assumed density filtering (ADF), minimizes the inclusive KL divergence:

Original: [Equation]
Revised: [Equation]

Original: In this paper, we focus our attention on (block) diagonal Gaussian approximations of the form q t ( θ t ) =N ( θ t | µ t , Σ t ), where Σ t = diag( Σ iit ).
Revised: In this paper, we discuss efficient and deterministic methods to optimize both of these KL objectives in O ( N z ) time. To do this, we use two approximations: first, we use (block) diagonal Gaussian approximations of the form q t ( θ t ) =

Original: Unfortunately, even with this assumption, optimizing these KL objectives is challenging, and often requires Monte Carlo approximations.

Original: We show how wecan use the EKF approximation to compute the optimum of both (approximate) KL objectives in closed form in linear time.
Revised: N ( θ t | µ t , Σ t ), where Σ t = diag( Σ ii t ); second, we use the EKF approximation to locally linearize the observation model at each step.

Original: When we apply this approach to the exclusive KL objective, the resulting algorithm is(as far as we know) a novel variational approach to DNNs.
Revised: When we apply this approach to the exclusive KL objective, the resulting algorithm is a novel diagonal variant of the mean field VI method (Blundell et al., 2015).

Original: When we apply this approach to the inclusive KL objective, we recover the “decoupled EKF” method (Puskorius and Feldkamp, 1991, 2003; Murtuza and Chorian, 1994).
Revised: When we apply this approach to the inclusive KL objective, we recover the “decoupled EKF” method (Puskorius and Feldkamp, 1991, 2003; Murtuza and Chorian, 1994).

Original: In Section 2, we describe our inference methods in more detail, and in Section 3, we perform an experimental comparison.
Revised: In Section 2, we describe our inference methods in more detail, and in Section 3, we perform an experimental comparison.

Original: We show that the inclusive KL objective gives better results than the exclusive KL, and both are more statistically efficient than  stochastic gradient descent (SGD), which requires multiple passes over the data, making it harder to apply to the streaming setting.
Revised: We show that the exclusive KL objective is more statistically efficient than both inclusive KL and stochastic gradient descent (SGD), which requires multiple passes over the data, making it harder to apply to the streaming setting.

Original: (We leave a comparison to online SGD variants that use replay buffers, such as Hu et al.
Revised: (We leave a comparison to online SGD variants that use replay buffers, such as Hu et al.

Original: (2021), to future work.)
Revised: (2021), to future work.)

Original: We conclude in Section 4.
Revised: We conclude in Section 4.

Revised: Our code is available at https:// github.com/probml/dynamax/tree/main/dynamax/generalized_gaussian_ssm/dekf .

Original: 2. Methods
Revised: 2. Methods

Original: In this section, we briefly describe our methods.
Revised: In this section, we briefly describe our methods.

Original: For the derivation, please see the appendix.
Revised: For the derivation, please see the appendix.

Original: (We use the notation A i and A i to represent the i th row and column of a matrix A , and
Revised: (We use the notation A i and A i to represent the i th row and column of a matrix A , and

Original: A ii to denote the i th (block) diagonal entry, if A is square.)
Revised: A ii to denote the i th (block) diagonal entry, if A is square.)

Revised: Refer to Appendix A.6 for the different ways we can compute the posterior predictive distribution.

Original: 2.1. Extended Kalman filter
Revised: 2.1. Extended Kalman filter

Original: For completeness, we summarize the usual EKF procedure, specialized to our setting.
Revised: For completeness, we summarize the usual EKF procedure, specialized to our setting.

Original: For notational simplicity, we drop the conditioning on inputs x t .
Revised: For notational simplicity, we drop the conditioning on inputs x t .

Original: The predict step is given by Eq.
Revised: The predict step is given by Eq.

Original: If we make a diagonal assumption on the process noise, we have µ t | t − 1 = µ t − 1 and Σ t | t − 1 =diag
Revised: (2), where µ t | t − 1 = µ t − 1 and Σ t | t − 1 =

Original: ( Σ iit − 1) +diag( Q iit ).
Revised: Σ t − 1 + Q t .

Original: If we assume a Gaussian likelihood, p ( y t | θ t ) =
Revised: If we assume a Gaussian likelihood, p ( y t | θ t ) =

Original: N ( y t | h ( θ t ) , R t ), then the update step becomes p ( θ t | y 1: t ) =
Revised: N ( y t | h ( θ t ) , R t ), then the update step becomes p ( θ t | y 1: t ) =

Original: N ( θ t | µ t | t , Σ t | t )
Revised: N ( θ t | µ t | t , Σ t | t )

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: where H t ≡ Jac( h )( µ t | t − 1 ) is a linearization of the observation model.
Revised: where H t ≡ Jac( h )( µ t | t − 1 ) is the N y × N z Jacobian matrix of the observation model and

Revised: K t is the N z × N y Kalman gain matrix.

Original: (The EKF approach can be extended to exponential family likelihoods using the techniques from Ollivier (2018);
Revised: We now derive an alternative update form that we will use below. The Woodbury matrix inversion lemma states the following: [Equation]

Original: Tronarp et al.
Revised: Letting A = Σ − 1 t | t − 1 , U

Original: In block form, we can write the elements of theposterior covariance matrix as follows:
Revised: = H ⊺ t , C = R − 1 t , V = H t and applying the inversion lemma gives the covariance update in information form, as follows:

Original: [Equation]
Revised: [Equation]

Original: 2.2. Approximating the update step
Revised: 2.2. Proposed diagonal approximations

Original: Let p =
Revised: Let p =

Original: p
Revised: p

Original: ( θ t | y 1: t )
Revised: ( θ t | y 1: t )

Original: =
Revised: =

Original: N ( θ t | µ p , Σ p ) be the full covariance computed by EKF, where µ p = µ t | t and Σ p =
Revised: N ( θ t | µ p , Σ p ) be the full covariance computed by EKF, where µ p = µ t | t and Σ p =

Original: Σ t | t , as in Section 2.1.
Revised: Σ t | t , as in Section 2.1.

Original: In general the EKF takes O ( N  2 z ) time to compute.
Revised: In general the EKF takes O ( N 3 y + N 2 z ) time to compute, due to the need to invert S t ∈ R N y × N y and then compute the N z × N z matrix K t H t Σ t | t − 1 .

Original: We seek to approximate this  by  q = N ( θ t | µ q , Σ q ), where Σ q = diag( Σ ii ). {+To derive our method, first note that if the previous posterior Σ t − 1 is (block) diagonal, then the prior predictive covariance for the next state is also diagonal, Σ iit | t − 
 = Σ iit − 1 + Q iit .+}
Revised: We seek to approximate this in O ( N z ) time by computing a diagonal approximation, q = N ( θ t | µ q , Σ q ), where Σ q = diag( Σ ii ). {+To derive our method, first note that if the previous posterior Σ t − 1 is (block) diagonal, then the prior predictive covariance for the next state is also diagonal, Σ iit | t − 
 = Σ iit − 1 + Q iit .+}

Original: We discuss different objectives below, which can be optimized in O ( N z ) time.
Revised: Also, the posterior covariance after the next observation can be efficiently computed using [Equation]

Original: 2.3. Exclusive KL (variational diagonal EKF)
Revised: From Eq. (7), the blocks of the (exact) EKF posterior covariance matrix are as follows: [Equation] The above equations tell us how to compute p .

Original: Suppose we want to  minimize D KL (q ∥ p  ), which is mode seeking, using diagonal q .
Revised: Now suppose we want to approximate this by a diagonal q so as to minimize D KL ( p ∥ q ), which is mode covering.

Original: We call this the “variational diagonal EKF” or “VD-EKF” method. In the appendix we show that the result is given by  µ q = µ p  and 
Revised: In the appendix we show that the result is given by setting µ q = µ p , and by zeroing out the off-diagonal blocks of the posterior covariance, Σ ijt | t for i ̸ = j . From Eq.

Revised: (11) this gives

Original: [Equation]
Revised: [Equation]

Original: This should be compared to the exact update in Eq.
Revised: This is equivalent to the fully decoupled EKF method of Puskorius and Feldkamp (2003); we therefore call this the “FD-EKF” method. Now suppose we want to minimize D KL ( q ∥ p ), which is mode seeking. We call this the “variational diagonal EKF” or “VD-EKF” method.

Original: We see that the variational approximation requires computing the diagonal elements of the posterior precisionmatrix , Σ − 1 t | t .
Revised: In the appendix we show that the result is given by setting µ q = µ p and by zeroing out the off-diagonal elements of the posterior precision , Σ − 1 t | t .

Original: By using the matrix inversion lemma, we can also write the result as follows:
Revised: From Eq. (9) this gives

Original: [Equation]
Revised: [Equation]

Original: This is similar to the diagonal elements in the update of Eq.

Original: (7), but drops some of the terms in S t ; we discuss the implications of this below. 2.4. Inclusive KL (fully decoupled EKF) We now seek to minimize D KL ( p ∥ q ), which is mode covering, usingdiagonal q . In the appendix we show that the result is given by µ q = µ p and
Revised: By using the matrix inversion lemma, we can also write the result as follows:

Original: [Equation]
Revised: [Equation]

Original: This requires computing the diagonal elements of the posterior covariance matrix , Σ t | t .

Original: The resulting update is equivalent to the fully decoupled EKF method of Puskorius and

Original: Feldkamp (2003); we therefore call this the “FD-EKF” method.

Original: 2.5. Discussion

Original: We see that the main difference between the exclusive KL in Eq.
Revised: We see that the main difference between the inclusive KL in Eq.

Original: (10) and the inclusive
Revised: (12) and the exclusive

Original: KL in Eq.
Revised: KL in Eq.

Original: (11) is that the former replaces S t with S it :=
Revised: (14) is that the former uses the full S t matrix, whereas the latter uses S it :=

Original: R t 
Revised: R t +

Original: + H it Σ iit | t − 1 ( H it ) T .
Revised: H it Σ iit | t − 1 ( H it ) T .

Original: Intuitively, using S t should be superior to S it, since S t captures moreof the uncertainty.
Revised: Intuitively, using S t instead of S it captures more uncertainty.

Original: In particular, from the perspective of estimating θ i , the missing terms in S it , namely (cid:80)
Revised: In particular, from the perspective of estimating θ i , the missing terms in S it , namely (cid:80)

Original: j ̸ = i H jt Σ jjt | t − 1 ( H jt ) T , play the same role as the observation covariance R t , but reflect uncertainty in the other parameters. 
Revised: j ̸ = i H jt Σ jjt | t − 1 ( H jt ) T , play the same role as the observation covariance R t , but reflect uncertainty in the other parameters. 2.3. Handling non-Gaussian observations

Revised: To handle non-Gaussian observations, such as discrete labels, we follow the technique of

Original: In Section 3, we will show experimentally that the inclusive KL gives better results.
Revised: Ollivier (2018) and Tronarp et al. Specifically, we assume p ( y t | θ t ) = expfam( y t | m t ) is an exponential family distribution with mean parameter m t = g − 1 ( η t ), where η t = h ( x t , θ t ) are the natural parameters (e.g., for classification, g − 1 = S is the softmax transform and η t are the logits). We define the predicted observation to be ˆ y t =

Original: 3. Experimental results
Revised: E [ m t | y 1: t − 1 ] = h ( µ t | t − 1 ), and the error (innovation) term to be e t = T ( y t ) − ˆ y t where T ( y t ) is the sufficient statistics vector. We define the observation covariance to be R t =

Revised: Cov [ T ( Y ) | ˆ y t ].

Revised: In the case of the categorical distribution with K labels, we have T ( y t ) =

Revised: [ I ( y t = 1) , . . . ,

Revised: I ( y t = K )], which is the one-hot encoding.

Revised: Similarly, we have ˆ y t =

Original: Since the exclusive KL objective neglects some forms of parameter uncertainty, it “trusts” the likelihood more, which is similar to using SGD with a larger learning rate.
Revised: [ p 1 t , . . . , p Kt ], where p kt = e η kt / ( (cid:80) j e η jt ) are the softmax probabilities, and η t are the logits. Finally, we have R t = diag( p t ) − p t p T t .

Original: This is illus- trated in Fig.
Revised: To avoid numerical problems with the sum-to-one constraint on T ( y t ), p t and R t , we can drop the last row/column. Alternatively, we can replace matrix inverses with least squares solvers, which work even if their argument is singular. For example, we can compute R − 1 using jnp.linalg.lstsq(R, jnp.eye(D))[0] .

Original: 1, where we display the posterior predictive distribution from a small MLP fit to a 2d binary classification dataset. We see that one pass of VD-EKF gives similar results to 5 epochs of SGD with a large learning rate of η = 0 .
Revised: In this section, we provide an experimental comparison of VD-EKF, FD-EKF and SGD for fitting a small LeNet CNN model to MNIST. We approximate the posterior predictive distribution by using a Monte Calro approximation (see Appendix A.6 for details).

Original: 45, whereas one pass of FD-EKF gives similarresults to 5 epochs of SGD with a much smaller learning rate of η = 0 .
Revised: The results are shown in Fig. We see that FD-EKF is similar to SGD, but that VD-EKF is consistently better than both.

Original: In Fig.
Revised: This is perhaps surprising given that VD-EKF is arguably a “less accurate” posterior approximation, since it ignores more posterior uncertainty when computing S t .

Original: 2, we show that both VD-EKF and FD-EKF can be much more data efficient thanSGD, in the sense that they need to see less data to converge to a good predictive distribution.
Revised: However, we conjecture that the mode seeking behavior of VD-EKF results in more robust estimates than the mode-covering behavior of FD-EKF. See A.7 for more thorough experimental results.

Original: 4. Conclusion and future work
Revised: 4. Conclusion and future work

Original: We haveseen how optimizing the inclusive KL, as is common in ADF and expectation propagation, can give better results than optimizing the exclusive KL, which is the basis of variational inference. We have also shown how we can efficiently compute both forms of update by linearizing the likelihood model, as in the EKF.
Revised: We have shown how we can efficiently compute a diagonal approximation to the posterior by linearizing the likelihood model, as in the EKF, and then optimizing KL ( q || p ) or KL ( p || q ) deterministically in closed form. Perhaps surprisingly, we find that the former objective does better.

Original: In the future, we would like to extend this approach to include more expressive approximations, such as diagonal plus low rank (e.g., see Mishkin et al., 2018; Lambert et al., 2021a).
Revised: In the future, we would like to extend this approach to include more expressive approximations, such as diagonal plus low rank (e.g., see Mishkin et al., 2018; Lambert et al., 2021a).

Original: We also want to apply it to more complex non-stationary and continual learning problems, which may be achievable with more sophisticated dynamical priors.
Revised: We also want to apply it to more complex non-stationary and continual learning problems, which may be achievable with more sophisticated priors for how the weights evolve over time.

Original: Appendix A. Appendix
Revised: Appendix A. Appendix

Original: In this section we give the derivation of the algorithms.
Revised: In this section, we give the derivation of the algorithms and more experimental results.

Original: A.1. Variational diagonal EKF A.1.1. Covariance
Revised: A.1. Objectives

Original: The covariance update for the EKF corresponds to the following equation: [Equation] The Woodbury matrix inversion lemma states the following:

Original: [Equation] Letting A = Σ − 1 t | t − 1 , U
Revised: Let µ p = µ t | t and Σ p =

Original: = H ⊺ t , C = R − 1 t , V = H t and applying the inversion lemma leads to the following simplified covariance update:
Revised: Σ t | t be the exact Gaussian posterior computed by EKF.

Original: [Equation]
Revised: In the

Original: We make the variational assumption that the covariance is fully diagonal at each time step. The prior is assumed to be:
Revised: VD-EKF, we want to find the approximation q ∼ N ( µ q , Σ q ), where Σ q is diagonal, which minimizes D KL ( q ∥ p ), given by

Original: [Equation]
Revised: [Equation]

Original: and we would like a diagonal posterior Σ t | t that is the closest approximation to  the expression on the right of Eq. (15), i.e., one that minimizesthe KL-divergence between the Gaussians. The KL-divergence between D -dimensional Gaussians p ∼ N (µ p , Σ p ) and q ∼ N ( µ q , Σ q ) can be simplified to the following expression:
Revised: In FD-EKF, we want to find the approximation that minimizes D KL ( p ∥ q ), given by

Original: [Equation]
Revised: [Equation]

Revised: A.2. Solving for the mean

Original: For the means, the optimal solution  is to setˆ µ q = µ p .
Revised: In both cases, the optimal solution for the mean is to set µ q = µ p .

Original: We can verify this by setting the partial derivativeof D KL ( q ∥ p ) with respect to the mean elements µ q,i to zero. 
Revised: We can verify this by setting the partial derivative with respect to the mean elements µ q,i to zero. The optimal µ p is given by [Equation]

Original: Since the covariance and the mean contribute additively to the objective in Eq. 
Revised: where the Kalman gain from Eq. (6) is given by [Equation] We now use the push-through identity

Revised: [Equation]

Revised: which states the following:

Revised: [Equation]

Revised: Letting U = Σ t | t − 1 H ⊺ t R − 1 t and V =

Revised: H t and applying the identity to Eq.

Revised: [Equation]

Revised: In summary

Revised: [Equation]

Revised: the mean update is given by

Original: (17), we can plugin the optimal  µ q  andthen solve the following simplified objective for Σ q :
Revised: [Equation] A.3. Solving for the covariance We next plug the optimal mean µ q = µ p into the KL objectives and solve the resulting simplified objectives for Σ q . The result will be denoted by

Original: [Equation]
Revised: [Equation]

Revised: where the form of σ it | t will be given below.

Revised: A.4. Variational diagonal EKF (reverse KL)

Revised: In this section, we use the reverse (exclusive) KL.

Revised: We plug in the optimal µ q to get the following simplified objective for Σ q :

Original:  Suppose we chose Σ q = diag( σ 2 i ).
Revised: [Equation] Suppose we chose Σ q = diag( σ 2 i ).

Original: Then we get
Revised: Then we get

Original: [Equation]
Revised: [Equation]

Original: Setting ∂J ( q ) ∂σ i = 0, we get 
Revised: Setting ∂J ( q ) ∂σ i = 0, we get [Equation]

Original: [Equation]
Revised: [Equation]

Original: (20) which gives
Revised: which gives

Original: [Equation]
Revised: [Equation]

Original: which says that we should match the marginal precisions of the two distributions.
Revised: which says that we should match the marginal precisions of the two distributions.

Original: Applying Eq.
Revised: Applying Eq.

Original: (21) to Eq.
Revised: (30) to Eq.

Original: (15) above, then:
Revised: (9) we get

Original: [Equation]
Revised: [Equation]

Original: and therefore the posterior covariance has the following form, matching Eq.
Revised: and therefore the posterior covariance has the following form, matching Eq.

Original: [Equation]
Revised: [Equation]

Original: The mean update corresponds to the following equation: [Equation] First
Revised: We now discuss the time complexity of this update.

Original: [Equation] we can rewrite the Kalman gain from Eq. [Equation] We can use the push-through identity
Revised: Naively computing the diagonal matrix (cid:0) H ⊺ t R − 1 t H t (cid:1) ii takes O ( N 2 z

Original: [Equation]
Revised: N y + N z

Original: which states the following: [Equation] Letting U = Σ t | t − 1 H ⊺ t R − 1 t and V = H t and applying the identity to Eq.
Revised: N 2 y + N 3 y ) time.

Original: [Equation] Therefore [Equation] the mean update can be re-expressed in the following form:
Revised: However, we can leverage the following result: if A ∈ R n × m and B ∈ R m × n , then the diagonal elements can be computed in O ( nm ) time using

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Note that the bottleneck at each step is the computation of the diagonal matrix (cid:0) H ⊺ t R − 1 t H t (cid:1) ii in the update step. If we naively compute the full matrix product before taking the diagonal elements, it takes O ( N 2 z N y + N z N 2 y + N 3 y ) per step, where θ t ∈ R N z and y t ∈ R N y . If N z ≫ N y , as is usually the case when training deep neural networks, we can further optimize this using the following lemma.
Revised: Hence we can first compute A

Original: Lemma 1. Let A ∈ R n × m , B ∈ R m × n . Then computing the diagonal elements of the product, ( AB ) ii , takes O( nm ) .
Revised: = H ⊺ t R − 1 t in O ( N z N 2 y + N 3 y ) time, and then compute ( AB ) ii in an additional O

Original: Proof Note that we do not need to evaluate the full matrix product to take the diagonal elements. Computing each diagonal element takes time O ( m ): [Equation] Since there are n diagonal elements, the total operation takes O ( nm ).
Revised: ( N z N y ) time, where B = H t .

Original: Computation [Equation] takes O ( N z N 2 y + N 3 y ). Then since by Lemma 1, computing the diagonal elements of the product ( H ⊺ t R − 1 t ) H t takes O ( N z N y ), overall the computations take O ( N z N 2 y ).
Revised: The total time is therefore linear in N z .

Original: A.2. Fully decoupled EKF 
Revised: A.5. Fully decoupled EKF (forwards KL)

Original: We now consider the inclusive KL objective: D KL ( p ∥ q ): [Equation]
Revised: We now consider the forwards (inclusive) KL objective.

Original: The optimal solution for the mean is to set ˆ µ q = µ p , which gives us the objective:
Revised: Plugging in µ q = µ p we get the simplified objective

Original: [Equation]
Revised: [Equation]

Original: Choosing Σ q = diag( σ 2 i ):
Revised: Choosing Σ q = diag( σ 2 i ):

Original: [Equation]
Revised: [Equation]

Original: Setting ∂J ( q ) ∂σ i = 0:
Revised: Setting ∂J [Equation]

Original: [Equation]
Revised: [Equation]

Original: And therefore the solution is simply taking the diagonal elements of Σ p :σ
Revised: And therefore the solution is simply taking the diagonal elements of Σ p :

Original: [Equation]
Revised: [Equation]

Original: = ( Σ

Original: [Equation]

Original: ) ii (39)

Original: Plugging in the full-covariance EKF posterior covariance Σ p =
Revised: Plugging in the full-covariance EKF posterior covariance Σ p =

Original: Σ t | t :
Revised: Σ t | t :

Original: [Equation] [Equation] 
Revised: [Equation] [Equation] [Equation]

Original: Let A i and A i represent the i th row and column of matrix A , respectively.
Revised: Let A i and A i represent the i th row and column of matrix A , respectively.

Original: Then, constraining the prior to be diagonal Σ t | t − 1 = diag( σ ′ 2 i ): 
Revised: Then, constraining the prior to be diagonal Σ t | t [Equation] = diag( σ ′ 2 i ): [Equation]

Original: [Equation]
Revised: [Equation]

Original: Plugging in the full expression for K t :
Revised: Plugging in the full expression for K t :

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Revised: [Equation]

Revised: [Equation]

Original: In addition, we see that:
Revised: In addition [Equation] we see that:

Original: [Equation]
Revised: [Equation]

Original: Plugging these in, therefore, we recover the covariance update equation for the fullydecoupled EKF, matching Eq. 
Revised: Plugging these in, therefore, we recover the covariance update equation for the fullydecoupled EKF, matching Eq. [Equation]

Revised: A.6. Computing the posterior predictive distribution

Revised: After performing posterior inference for the parameters, we next compute the posterior predictive distribution.

Revised: In the classification case, this becomes

Original: [Equation]
Revised: [Equation]

Revised: The simplest approach to this integral is to replace the Gaussian posterior with a delta function centered at the posterior mean/mode ˆ µ t = µ t | t , which gives the plugin approximation

Revised: [Equation]

Revised: We call this the “plugin predictive”.

Revised: Alternatively we can use a Monte Carlo approximation, in which we sample θ t from the Gaussian posterior and plug into h ; we call this the “MC predictive”.

Revised: However, Immer et al.

Revised: (2021) argues that it is better to plug the samples into the linearized model, ˆ h ( x t , θ t ) =

Revised: H t x t ; we call this the “linearized predictive”.

Revised: A.7. More experimental results