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 that exclusive KL is more effective than both inclusive KL and 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: 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:
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 discuss efficient and deterministic methods to optimize both of these
Revised: In this paper, we discuss efficient and deterministic methods to optimize both of these

Original: KL objectives in O ( N z ) time.
Revised: KL objectives in O ( N z ) time.

Original: To do this, we use two approximations: first, we use (block) diagonal Gaussian approximations of the form q t ( θ t ) =
Revised: To do this, we use two approximations: first, we use (block) diagonal Gaussian approximations of the form q t ( θ t ) =

Original: N ( θ t | µ t , Σ t ), where Σ t = diag( Σ ii t ); second, we use the EKF approximation to locally linearize the observation model at each step.
Revised: N ( θ t | µ t , Σ t ), where Σ t = diag( Σ iit ); 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 a novel diagonal variant of the mean field VI method (Blundell et al., 2015).
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 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.
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.

Original: (2021), to future work.)

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

Original: Our code is available at https:// github.com/probml/dynamax/tree/main/dynamax/generalized_gaussian_ssm/dekf .
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.)

Original: Refer to Appendix A.6 for the different ways we can compute the posterior predictive distribution.
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: (2), where µ t | t − 1 = µ t − 1 and Σ t | t − 1 =
Revised: (2), where µ t | t − 1 = µ t − 1 and Σ t | t − 1 =

Original: Σ t − 1 + Q t .
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 the N y × N z Jacobian matrix of the observation model and
Revised: where H t ≡ Jac( h )( µ t | t − 1 ) is the N y × N z Jacobian matrix of the observation model and

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

Original: We now derive an alternative update form that we will use below.
Revised: We now derive an alternative update form that we will use below.

Original: The Woodbury matrix inversion lemma states the following:
Revised: The Woodbury matrix inversion lemma states the following:

Original: [Equation]
Revised: [Equation]

Original: Letting A = Σ − 1 t | t − 1 , U
Revised: Letting A = Σ − 1 t | t − 1 , U

Original: = H ⊺ t , C
Revised: = H ⊺ t , C

Original: =
Revised: =

Original: R − 1 t , V = H t and applying the inversion lemma gives the covariance update in information form, as follows:
Revised: 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. Proposed diagonal approximations
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 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 .
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 in O ( N z ) time by computing a diagonal approximation, q =
Revised: We seek to approximate this in O ( N z ) time by computing a diagonal approximation, q =

Original: N ( θ t | µ q , Σ q ), where Σ q = diag( Σ ii ).
Revised: N ( θ t | µ q , Σ q ), where Σ q = diag( Σ ii ).

Original: 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 − 

Revised: 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 −

Original: = Σ iit − 1 + Q iit .
Revised: = Σ iit − 1 + Q iit .

Original: Also, the posterior covariance after the next observation can be efficiently computed using
Revised: Also, the posterior covariance after the next observation can be efficiently computed using

Original: [Equation]
Revised: [Equation]

Original: From Eq.
Revised: From Eq.

Original: (7), the blocks of the (exact) EKF posterior covariance matrix are as follows:
Revised: (7), the blocks of the (exact) EKF posterior covariance matrix are as follows:

Original: [Equation]
Revised: [Equation]

Original: The above equations tell us how to compute p .
Revised: The above equations tell us how to compute p .

Original: Now suppose we want to approximate this by a diagonal q so as to minimize D KL ( p ∥ q ), which is mode covering.
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: 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 ̸
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 ̸

Original: = j .
Revised: = j .

Original: From Eq.
Revised: From Eq.

Original: (11) this gives
Revised: (11) this gives

Original: [Equation]
Revised: [Equation]

Original: This is equivalent to the fully decoupled EKF method of Puskorius and Feldkamp (2003); we therefore call this the “FD-EKF” method.
Revised: This is equivalent to the fully decoupled EKF method of Puskorius and Feldkamp (2003); we therefore call this the “FD-EKF” method.

Original: Now suppose we want to minimize D KL ( q ∥ p ), which is mode seeking.
Revised: Now suppose we want to minimize D KL ( q ∥ p ), which is mode seeking.

Original: We call this the “variational diagonal EKF” or “VD-EKF” method.
Revised: We call this the “variational diagonal EKF” or “VD-EKF” method.

Original: 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 .
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: From Eq.
Revised: From Eq.

Original: (9) this gives
Revised: (9) this gives

Original: [Equation]
Revised: [Equation]

Original: By using the matrix inversion lemma, we can also write the result as follows:
Revised: By using the matrix inversion lemma, we can also write the result as follows:

Original: [Equation]
Revised: [Equation]

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

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

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

Original: (14) is that the former uses the full S t matrix, whereas the latter uses 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 instead of S it captures more 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.

Original: 2.3. Handling non-Gaussian observations
Revised: 2.3. Handling non-Gaussian observations

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

Original: Ollivier (2018) and Tronarp et al.
Revised: Ollivier (2018) and Tronarp et al.

Original: Specifically, we assume p ( y t | θ t ) =
Revised: Specifically, we assume p ( y t | θ t ) =

Original: 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).
Revised: 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).

Original: We define the predicted observation to be ˆ y t =
Revised: We define the predicted observation to be ˆ y t =

Original: E [ m t | y 1: t − 1 ]
Revised: E [ m t | y 1: t − 1 ]

Original: = h ( µ t | t − 1 ), and the error (innovation) term to be e t =
Revised: = h ( µ t | t − 1 ), and the error (innovation) term to be e t =

Original: T ( y t ) − ˆ y t where T ( y t ) is the sufficient statistics vector.
Revised: T ( y t ) − ˆ y t where T ( y t ) is the sufficient statistics vector.

Original: We define the observation covariance to be R t =
Revised: We define the observation covariance to be R t =

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

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

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

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

Original: Similarly, we have ˆ y t =
Revised: Similarly, we have ˆ y t =

Original: [ p 1 t , . .
Revised: [ p 1 t , . .

Original: . , p Kt ], where p kt = e η kt / ( (cid:80) j e η jt ) are the softmax probabilities, and η t are the logits.
Revised: . , p Kt ], where p kt = e η kt / ( (cid:80) j e η jt ) are the softmax probabilities, and η t are the logits.

Original: Finally, we have
Revised: Finally, we have

Original: R t = diag( p t ) − p t p T t .
Revised: R t = diag( p t ) − p t p T t .

Original: To avoid numerical problems with the sum-to-one constraint on
Revised: To avoid numerical problems with the sum-to-one constraint on

Original: T ( y t ), p t and R t , we can drop the last row/column.
Revised: T ( y t ), p t and R t , we can drop the last row/column.

Original: Alternatively, we can replace matrix inverses with least squares solvers, which work even if their argument is singular.
Revised: Alternatively, we can replace matrix inverses with least squares solvers, which work even if their argument is singular.

Original: For example, we can compute R − 1 using jnp.linalg.lstsq(R, jnp.eye(D))[0] .
Revised: For example, we can compute R − 1 using jnp.linalg.lstsq(R, jnp.eye(D))[0] .

Original: 3. Experimental results
Revised: 3. Experimental results

Original: In this section, we provide an experimental comparison of VD-EKF, FD-EKF and SGD for fitting a small LeNet CNN model to MNIST.
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.

Original: We approximate the posterior predictive distribution by using a Monte Calro approximation (see Appendix A.6 for details).
Revised: We approximate the posterior predictive distribution by using a Monte Carlo approximation (see Appendix A.6 for details).

Original: The results are shown in Fig.
Revised: The results are shown in Fig.

Original: We see that FD-EKF is similar to SGD, but that VD-EKF is consistently better than both.
Revised: We see that FD-EKF is similar to SGD, but that VD-EKF is consistently better than both.

Original: 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 .
Revised: This is perhaps surprising given that VD-EKF is arguably a “less accurate” posterior approximation, since it ignores more outcome variance when computing S t .

Original: However, we conjecture that the mode seeking behavior of VD-EKF results in more robust estimates than the mode-covering behavior of FD-EKF.
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.

Original: See A.7 for more thorough experimental results.
Revised: See A.7 for more thorough experimental results.

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

Original: 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.
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 D KL ( q ∥ p ) or D KL ( p ∥ q ) deterministically in closed form.

Original: Perhaps surprisingly, we find that the former objective does better.
Revised: 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 priors for how the weights evolve over time.
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.

Revised: Finally we want to compare to online SGD variants that use replay buffers, such as Hu et al.

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

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

Original: A.1. Objectives
Revised: A.1. Objectives

Original: Let µ p = µ t | t and Σ p =
Revised: Let µ p = µ t | t and Σ p =

Original: Σ t | t be the exact Gaussian posterior computed by EKF.
Revised: Σ t | t be the exact Gaussian posterior computed by EKF.

Original: In the
Revised: In the

Original: VD-EKF, we want to find the approximation q ∼ N ( µ q , Σ q ), where Σ q is diagonal, which minimizes D KL ( q ∥ p ), given by
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: In FD-EKF, we want to find the approximation that minimizes D KL ( p ∥ q ), given by
Revised: In FD-EKF, we want to find the approximation that minimizes D KL ( p ∥ q ), given by

Original: [Equation]
Revised: [Equation]

Original: A.2. Solving for the mean
Revised: A.2. Solving for the mean

Original: In both cases, the optimal solution for the mean 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 derivative 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.

Original: The optimal µ p is given by
Revised: The optimal µ p is given by

Original: [Equation]
Revised: [Equation]

Original: where the Kalman gain from Eq.
Revised: where the Kalman gain from Eq.

Original: (6) is given by
Revised: (6) is given by

Original: [Equation]
Revised: [Equation]

Original: We now use the push-through identity
Revised: We now use the push-through identity

Original: [Equation]
Revised: [Equation]

Original: which states the following:
Revised: which states the following:

Original: [Equation]
Revised: [Equation]

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

Original: H t and applying the identity to Eq.
Revised: H t and applying the identity to Eq.

Original: [Equation]
Revised: [Equation]

Original: In summary
Revised: In summary

Original: [Equation]
Revised: [Equation]

Original: the mean update is given by
Revised: the mean update is given by

Original: [Equation]
Revised: [Equation]

Original: A.3. Solving for the covariance
Revised: A.3. Solving for the covariance

Original: We next plug the optimal mean µ q = µ p into the KL objectives and solve the resulting simplified objectives for Σ q .
Revised: We next plug the optimal mean µ q = µ p into the KL objectives and solve the resulting simplified objectives for Σ q .

Original: The result will be denoted by
Revised: The result will be denoted by

Original: [Equation]
Revised: [Equation]

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

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

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

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

Original: [Equation]
Revised: [Equation]

Original: Suppose we chose Σ q = diag( σ 2 i ).
Revised: 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

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

Original: 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: (30) to Eq.
Revised: (30) to Eq.

Original: (9) we get
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: We now discuss the time complexity of this update.
Revised: We now discuss the time complexity of this update.

Original: Naively computing the diagonal matrix (cid:0) H ⊺ t R − 1 t H t (cid:1) ii takes O ( N 2 z
Revised: Naively computing the diagonal matrix (cid:0) H ⊺ t R − 1 t H t (cid:1) ii takes O ( N 2 z

Original: N y + N z
Revised: N y + N z

Original: N 2 y + N 3 y ) time.
Revised: N 2 y + N 3 y ) time.

Original: 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
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: Hence we can first compute A
Revised: Hence we can first compute A

Original: = 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
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: ( N z N y ) time, where B = H t .
Revised: ( N z N y ) time, where B = H t .

Original: The total time is therefore linear in N z .
Revised: The total time is therefore linear in N z .

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

Original: We now consider the forwards (inclusive) KL objective.
Revised: We now consider the forwards (inclusive) KL objective.

Original: Plugging in µ q = µ p we get the simplified 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
Revised: Setting ∂J

Original: [Equation] [Equation]
Revised: [Equation] [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: 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]

Original: [Equation]
Revised: [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
Revised: Then, constraining the prior to be diagonal Σ t | t

Original: [Equation]
Revised: [Equation]

Original: = diag( σ ′ 2 i ):
Revised: = diag( σ ′ 2 i ):

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

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

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

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

Original: In addition
Revised: In addition

Original: [Equation]
Revised: [Equation]

Original: we see that:
Revised: 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.

Original: [Equation]
Revised: [Equation]

Original: A.6. Computing the posterior predictive distribution
Revised: A.6. Computing the posterior predictive distribution

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

Original: In the classification case, this becomes
Revised: In the classification case, this becomes

Original: [Equation]
Revised: [Equation]

Original: 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: 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

Original: [Equation]
Revised: [Equation]

Original: We call this the “plugin predictive”.
Revised: We call this the “plugin predictive”.

Original: 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: 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”.

Original: However, Immer et al.
Revised: However, Immer et al.

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

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

Original: A.7. More experimental results
Revised: A.7. More experimental results