arXiv:2206.15078v1 [cs.LG] 30 Jun 2022

Laplacian Autoencoders for
Learning Stochastic Representations
Marco Miani1 , Frederik Warburg1 ,Pablo Moreno-Muñoz, Nicke Skafte Detlefsen, Søren Hauberg{mmia, frwa, pabmo, nsde, sohau}@dtu.dk
Technical University of Denmark
https://github.com/FrederikWarburg/LaplaceAE

Abstract
Established methods for unsupervised representation learning such as variational
autoencoders produce none or poorly calibrated uncertainty estimates making it
difficult to evaluate if learned representations are stable and reliable. In this work,
we present a Bayesian autoencoder for unsupervised representation learning, which
is trained using a novel variational lower-bound of the autoencoder evidence. Thisis maximized using Monte Carlo EM with a variational distribution that takes
the shape of a Laplace approximation. We develop a new Hessian approximationthat scales linearly with data size allowing us to model high-dimensional data.
Empirically, we show that our Laplacian autoencoder estimates well-calibrated
uncertainties in both latent and output space. We demonstrate that this results inimproved performance across a multitude of downstream tasks.

1

Introduction

Unsupervised representation learning is a brittle matter. Consider the classic autoencoder (AE)
(Rumelhart et al., 1986), which compresses data x to a low-dimensional representation z from whichdata is approximately reconstructed. The nonlinearity of the model implies that sometimes smallchanges to data x gives a large change in the latent representation z (and sometimes not). Likewise,for some data, reconstructions are of low quality, while for others it is near perfect. Unfortunately,
the model does not have a built-in quantification of its uncertainty, and we cannot easily answer when
the representation is reliable and accurately reflects data.
The celebrated variational autoencoder (VAE) (Kingma and Welling, 2014; Rezende et al., 2014)
address this concern directly through an explicit likelihood model p(x|z) and a variational approximation of the representation posterior p(z|x). Both these distributions have parameters predicted byneural networks that act similarly to the encoder–decoder pair of the classic autoencoder.
But is the VAEs quantification of reliability reliable? To investigate, we fit a VAE with a twodimensional latent representation to the MNIST dataset (Lecun et al., 1998), and illustrate the predicteduncertainty of p(x|z) in Fig. 1a. The model learns to assign high uncertainty to low-level image
features such as edges, but predicts its smallest values far away from the data distribution. Not
only is such behavior counter-intuitive, it is also suboptimal in terms of data likelihood (Sec. 1.1).
Retrospectively, this should not be surprising as the uncertainty levels away from the data are governedby the extrapolatory behavior of the neural network determining p(x|z). This suggests that perhaps
uncertainty should be a derived quantity rather than a predicted one.
1

Denotes equal contribution.

Preprint. Under review.

µ<latexit sha1_base64="I+0UOx9m5Mg8vkDHEjZJgZ5mbJ4=">AAAB9XicbVBNS0JBFJ1nX2ZfVss2gxIYgbwnWc+dFUFLg/wAn8m8cdTBeR/M3FfJw3/Rok2LItr2X9r5bxo1orIDFw7n3Mu997ih4ApMc2wkFhaXlleSq6m19Y3NrfT2Tk0FkaSsSgMRyIZLFBPcZ1XgIFgjlIx4rmB1d3A+8eu3TCoe+NcwDFnLIz2fdzkloKUbB9g9xDlygGunF6N2OmvmS6Uju2hjTexj2ypgK29O8U2y5Yxz+DAuDyvt9IfTCWjkMR+oIEo1LTOEVkwkcCrYKOVEioWEDkiPNTX1icdUK55ePcL7WungbiB1+YCn6s+JmHhKDT1Xd3oE+uqvNxH/85oRdO1WzP0wAubT2aJuJDAEeBIB7nDJKIihJoRKrm/FtE8koaCDSukQ5l6eJ7VC3irmzSudxhmaIYn2UAblkIVOUBldogqqIookekTP6MW4M56MV+Nt1powvmZ20S8Y75/+rZUq</latexit>(a) VAE<latexit sha1_base64="KnTu97m0JcXB3xIDlkbq5VJDXf4=">AAAB7nicbVBNS8NAEJ34WetX1aOXxUbwFJKiNd6KXjxWsB/QhrLZbtulu0nY3Qgl9Ed48aCIV3+PN/+Nm7aIWh8MPN6bYWZemHCmtOt+Wiura+sbm4Wt4vbO7t5+6eCwqeJUEtogMY9lO8SKchbRhmaa03YiKRYhp61wfJP7rQcqFYujez1JaCDwMGIDRrA2Usu2uyK17V6p7DpXvudVKygn1XPfR57jzvBNyrBAvVf66PZjkgoaacKxUh3PTXSQYakZ4XRa7KaKJpiM8ZB2DI2woCrIZudO0alR+mgQS1ORRjP150SGhVITEZpOgfVI/fVy8T+vk+qBH2QsSlJNIzJfNEg50jHKf0d9JinRfGIIJpKZWxEZYYmJNgkVTQhLLy+TZsXxLhz3rlKuXS/iKMAxnMAZeHAJNbiFOjSAwBge4RlerMR6sl6tt3nrirWYOYJfsN6/ADgOjtc=</latexit><latexit sha1_base64="2xtfArd+6jiIrrxb8xG9C8PUOiw=">AAAB83icbVDLSsNAFJ3UV62vqks3g43gKiRBa9wV3bisYB/QxDKZTtqhM0mYmQgl9DfcuFDErT/jzr9x0hZR64ELh3Pu5d57wpRRqWz70yitrK6tb5Q3K1vbO7t71f2DtkwygUkLJywR3RBJwmhMWooqRrqpIIiHjHTC8XXhdx6IkDSJ79QkJQFHw5hGFCOlJd80fUmHHN27ptmv1mzr0nOcugsLUj/zPOhY9gzfpAYWaParH/4gwRknscIMSdlz7FQFORKKYkamFT+TJEV4jIakp2mMOJFBPrt5Ck+0MoBRInTFCs7UnxM54lJOeKg7OVIj+dcrxP+8XqYiL8hpnGaKxHi+KMoYVAksAoADKghWbKIJwoLqWyEeIYGw0jFVdAhLLy+Ttms555Z969YaV4s4yuAIHINT4IAL0AA3oAlaAIMUPIJn8GJkxpPxarzNW0vGYuYQ/ILx/gWlgpDI</latexit><latexit sha1_base64="IGuruciO9i8L1qsaALt5I7n1+Js=">AAACAnicbVDLSgNBEJyNrxhfUU/iZUgQImLYFR+bW1QEDx4UjApJCLOT2WTI7M4y0yuGJXjxE/wFLx4U8epXePNvnCQiaixoKKq66e7yIsE12PaHlRobn5icSk9nZmbn5heyi0sXWsaKsgqVQqorj2gmeMgqwEGwq0gxEniCXXqdw75/ec2U5jI8h27E6gFphdznlICRGtmVGrAbSAreOj7ZP8KFSGrYbEu63mtk83axVNp2d1xsiLvrOlvYKdoDfJN8OVfbuP8od08b2fdaU9I4YCFQQbSuOnYE9YQo4FSwXqYWaxYR2iEtVjU0JAHT9WTwQg+vGaWJfalMhYAH6s+JhARadwPPdAYE2vqv1xf/86ox+G494WEUAwvpcJEfCwwS9/PATa4YBdE1hFDFza2YtokiFExqGRPCyMuj5GKr6OwU7TOTxgEaIo1WUQ4VkIP2UBkdo1NUQRTdogf0hJ6tO+vRerFeh60p62tmGf2C9fYJwOGZXg==</latexit>(b) LAE (post-hoc)2µ<latexit sha1_base64="I3lyBSIdfE5Wl35PdctE+++DNmc=">AAACAHicbVDLSgMxFM3Ud31VXbhwE1qEFqHMiI9xVxXBhQsF+4C2lEya1mAmMyR3xGHoxm/wD9y4UMStn+HOvzHTiqj1wIXDOfcm9x4vFFyDbX9YmYnJqemZ2bns/MLi0nJuZbWmg0hRVqWBCFTDI5oJLlkVOAjWCBUjvidY3bs+Tv36DVOaB/IS4pC1fdKXvMcpASN1custYLeQFGkJnx2e4GIg05dKg06uYJcPDnbcXRcb4u65zjZ2yvYQ36RQybe27j8q8Xkn997qBjTymQQqiNZNxw6hnRAFnAo2yLYizUJCr0mfNQ2VxGe6nQwPGOBNo3RxL1CmJOCh+nMiIb7Wse+ZTp/Alf7rpeJ/XjOCnttOuAwjYJKOPupFAkOA0zRwlytGQcSGEKq42RXTK6IIBZNZ1oQwdvI4qW2Xnd2yfWHSOEIjzKINlEdF5KB9VEGn6BxVEUUD9ICe0LN1Zz1aL9brqDVjfc2soV+w3j4BaeuYow==</latexit>(c) LAE (online)<latexit sha1_base64="KnTu97m0JcXB3xIDlkbq5VJDXf4=">AAAB7nicbVBNS8NAEJ34WetX1aOXxUbwFJKiNd6KXjxWsB/QhrLZbtulu0nY3Qgl9Ed48aCIV3+PN/+Nm7aIWh8MPN6bYWZemHCmtOt+Wiura+sbm4Wt4vbO7t5+6eCwqeJUEtogMY9lO8SKchbRhmaa03YiKRYhp61wfJP7rQcqFYujez1JaCDwMGIDRrA2Usu2uyK17V6p7DpXvudVKygn1XPfR57jzvBNyrBAvVf66PZjkgoaacKxUh3PTXSQYakZ4XRa7KaKJpiM8ZB2DI2woCrIZudO0alR+mgQS1ORRjP150SGhVITEZpOgfVI/fVy8T+vk+qBH2QsSlJNIzJfNEg50jHKf0d9JinRfGIIJpKZWxEZYYmJNgkVTQhLLy+TZsXxLhz3rlKuXS/iKMAxnMAZeHAJNbiFOjSAwBge4RlerMR6sl6tt3nrirWYOYJfsN6/ADgOjtc=</latexit><latexit sha1_base64="2xtfArd+6jiIrrxb8xG9C8PUOiw=">AAAB83icbVDLSsNAFJ3UV62vqks3g43gKiRBa9wV3bisYB/QxDKZTtqhM0mYmQgl9DfcuFDErT/jzr9x0hZR64ELh3Pu5d57wpRRqWz70yitrK6tb5Q3K1vbO7t71f2DtkwygUkLJywR3RBJwmhMWooqRrqpIIiHjHTC8XXhdx6IkDSJ79QkJQFHw5hGFCOlJd80fUmHHN27ptmv1mzr0nOcugsLUj/zPOhY9gzfpAYWaParH/4gwRknscIMSdlz7FQFORKKYkamFT+TJEV4jIakp2mMOJFBPrt5Ck+0MoBRInTFCs7UnxM54lJOeKg7OVIj+dcrxP+8XqYiL8hpnGaKxHi+KMoYVAksAoADKghWbKIJwoLqWyEeIYGw0jFVdAhLLy+Ttms555Z969YaV4s4yuAIHINT4IAL0AA3oAlaAIMUPIJn8GJkxpPxarzNW0vGYuYQ/ILx/gWlgpDI</latexit>2µ<latexit sha1_base64="KnTu97m0JcXB3xIDlkbq5VJDXf4=">AAAB7nicbVBNS8NAEJ34WetX1aOXxUbwFJKiNd6KXjxWsB/QhrLZbtulu0nY3Qgl9Ed48aCIV3+PN/+Nm7aIWh8MPN6bYWZemHCmtOt+Wiura+sbm4Wt4vbO7t5+6eCwqeJUEtogMY9lO8SKchbRhmaa03YiKRYhp61wfJP7rQcqFYujez1JaCDwMGIDRrA2Usu2uyK17V6p7DpXvudVKygn1XPfR57jzvBNyrBAvVf66PZjkgoaacKxUh3PTXSQYakZ4XRa7KaKJpiM8ZB2DI2woCrIZudO0alR+mgQS1ORRjP150SGhVITEZpOgfVI/fVy8T+vk+qBH2QsSlJNIzJfNEg50jHKf0d9JinRfGIIJpKZWxEZYYmJNgkVTQhLLy+TZsXxLhz3rlKuXS/iKMAxnMAZeHAJNbiFOjSAwBge4RlerMR6sl6tt3nrirWYOYJfsN6/ADgOjtc=</latexit><latexit sha1_base64="2xtfArd+6jiIrrxb8xG9C8PUOiw=">AAAB83icbVDLSsNAFJ3UV62vqks3g43gKiRBa9wV3bisYB/QxDKZTtqhM0mYmQgl9DfcuFDErT/jzr9x0hZR64ELh3Pu5d57wpRRqWz70yitrK6tb5Q3K1vbO7t71f2DtkwygUkLJywR3RBJwmhMWooqRrqpIIiHjHTC8XXhdx6IkDSJ79QkJQFHw5hGFCOlJd80fUmHHN27ptmv1mzr0nOcugsLUj/zPOhY9gzfpAYWaParH/4gwRknscIMSdlz7FQFORKKYkamFT+TJEV4jIakp2mMOJFBPrt5Ck+0MoBRInTFCs7UnxM54lJOeKg7OVIj+dcrxP+8XqYiL8hpnGaKxHi+KMoYVAksAoADKghWbKIJwoLqWyEeIYGw0jFVdAhLLy+Ttms555Z969YaV4s4yuAIHINT4IAL0AA3oAlaAIMUPIJn8GJkxpPxarzNW0vGYuYQ/ILx/gWlgpDI</latexit>

2

Figure 1: (a) The VAE learns to estimate high variance for low level image features such as edges, but
fails at extrapolating uncertainties away from training data. (b) Applying post-hoc laplace to the AE
setup shows much better extrapolating capabilities, but fails in estimating calibrated uncertainties in
output space. (c) Our online, sampling-based optimization of a Laplacian autoencoder (LAE) gives
well-behaved uncertainties in both latent and output space.
From a Bayesian perspective the natural solution is to form an (approximate) posterior over theweights of the neural networks. To investigate, we consider a state-of-the-art implementation of a posthoc Laplace approximation (Daxberger et al., 2021) of the weight posterior. This amounts to traininga regular autoencoder, and thereafter approximating the weight uncertainty with the Hessian of theloss (Sec. 1.1). Fig. 1b shows that uncertainty now grows, as intuitively expected, with the distance tothe data distribution, but there seems to be little semantic structure in the uncertainty in output space.
This suggests that while the post-hoc procedure is computationally attractive it is too simplistic.
In this paper we introduce a new framework for Bayesian autoencoders in unsupervised representation learning. Our method takes inspiration from the Laplace approximation to build avariational distribution on the neural network weights. We first apply existing post-hoc methods
to the autoencoder domain, showcasing their lack of proper calibrated uncertainty. Secondly, wedevelop a new fast and memory-efficient Hessian approximation, which allows us to maximize avariational lower bound using Monte Carlo EM, such that model uncertainty is a key part of modeltraining rather than estimated post-hoc. Fig. 1c gives an example of the corresponding uncertainty,
which exhibits a natural and semantically meaningful behavior.
1.1

BackgroundR
The VAE is a latent variable model that parametrize the data density p(x) = p(x|z)p(z)dz
using a latent variable (representation) z. Here p(z) is a, usually standard normal, prior over the
representation, and p(x|z) is a likelihood with parameters predicted by a neural network. The
nonlinearity of the likelihood parameters render the marginalization of z intractable, and a variational
lower bound of p(x) is considered instead. To arrive at this, one first introduces a variational
approximation q(z|x) ≈ p(z|x) and write p(x) = Eq(z|x) [p(x|z)p(z)/q(z|x)]. A lower bound onp(x) then follows by direct application of Jensen’s inequality,p(x) ≥ LVAE (x) = Eq(z|x) [log p(x|z)] − KL(q(z|x)kp(z)).

(1)

If we momentarily assume that p(x|z) = N (x|µ(z), σ 2 (z)), we see that optimally σ 2 (z) should be
as large as possible away from training data in order to increase p(x) on the training data. Yet thisis not the observed empirical behavior in Fig. 1a. Since the σ 2 network is left untrained away fromtraining data, its predictions depends on extrapolation. In practice, σ 2 take fairly small values neartraining data (assuming the mean µ provides a reasonable data fit), and σ 2 likewise extrapolates tosmall values even if this is suboptimal in terms of data likelihood. Similar remarks hold for otherlikelihood models p(x|z) and encoder distributions q(z|x): relying on neural network extrapolation
to predict uncertainty does not work.
The Laplace approximation (Laplace, 1774; MacKay, 1992) is an integral part of our proposedsolution. In the context of Bayesian neural networks, we seek the weight-posterior p(θ|D) ∝exp(−L(D; θ)), where θ are network weights, D is the training data, and L is the applied loss
function interpreted as an unnormalized log-posterior. This is generally intractable and Laplace’s
approximation (LA) amounts to a second-order Taylor expansion around a chosen weight vector θ∗1
log p(θ|D) = L∗ + (θ − θ ∗ )> ∇L∗ + (θ − θ ∗ )> ∇2 L∗ (θ − θ ∗ ) + O(kθ − θ ∗ k3 )
2
2

(2)

where we use the short-hand L∗ = L(D; θ∗ ). The approximation, thus, assumes that p(θ|D) is
Gaussian. Note that when θ∗ is a MAP estimate, the first order term vanishes and the second orderterm is negative semi-definite. We provide more details on the Laplace approximation in Appendix B.
In practice, computing the full Hessian is too taxing both in terms of computation and memory, andvarious approximations are applied (Sec. 3).

Figure 2: Model overview. We learn a distribution over parameters such that we can sample encoderseψ and decoders dφ . This allow us to compute the empirical mean and variance in both the latent
space z and the output space fθ (xn ) = dφ (eψ (xn )).

2

Laplacian Autoencoders

We consider unsupervised representation learning from i.i.d. data D = {xn }N
n=1 consisting ofobservations xn ∈ RD . We also define a continuous latent space such that representations zn ∈ RK .
Similar to AEs (Hinton and Salakhutdinov, 2006), we consider two neural networks eψ : RD →RK
and dφ : RK →RD , widely known as the encoder and decoder. These have parameters θ = {ψ, φ}.
We refer to the composition of encoder and decoder as fθ = dψ ◦ eφ .
Model overview. The autoencoder network structure implies that we model the data as being
distributed on a K-dimensional manifold parametrized by θ. We then seek the distribution of the reconstruction xrec = fθ (x) given observation x, where the uncertainty comes from θ being unknown,p(xrec |x, f ) = Eθ∼p(θ|x,f ) [p(xrec |θ, x, f )].

(3)

Notice that we explicitly condition on f , which is the operator θ 7→ fθ , even if this is not stochastic;
this choice will simplify notation at a later stage. Mimicking the standard autoencoder reconstruction
loss, we set p(xrec |θ, x, f ) = N (xrec |fθ (x), I). Since p(θ|x, f ) is unknown, the reconstruction likelihood (3) is intractable, and approximations are in order. Similar to Blundell et al. (2015), we resort toa Gaussian approximation, but rather than learning the variance variationally, we opt for LA. This willallow us to sample NNs and deduce uncertainties in both latent and output space as illustrated in Fig. 2.
Intractable joint distribution. Any meaningful approximate posterior over θ should be similar to
the marginal of the joint distribution p(θ, xrec |x, f ). This marginal isp(θ|x, f ) = Exrec ∼p(xrec |x,f ) [p(θ|xrec , x, f )](4)

which can be bounded on a log-scale using Jensen’s inequality,
log p(θ|x, f ) ≥ Lθ = Exrec ∼p(xrec |x,f ) [log p(θ|xrec , x, f )].

(5)

Our first approximation is a LA of p(θ|x, f ) ≈ q t (θ|x, f ) = N (θ|θ (t) , H−1t ), where we postpone(t)the details on how to acquire θ and Ht . These will eventually be iteratively computed fromthe lower bound (5); hence the t index. Fig. 3 (a) illustrates the situation thus far: p(θ|x, f ) is
approximately Gaussian, but the non-linearity of f gives p(xrec |x, f ) a non-trivial density.

Linearization for gradient updates. Standard gradient-based learning can be viewed as a linearization of f in θ, i.e. for a loss L = l(fθ ), the gradient is ∇θ L = Jθ fθ ∇f l(f ), where Jθ fθ is theJacobian of f . In a similar spirit, we linearize f in θ in order to arrive at a tractable approximation ofp(xrec |x, f ). Specifically, we perform a Taylor expansion around θ (t)fθ (x) = fθ(t) (x) + Jθ fθ(t) (x)(θ − θ (t) ) + O kθ − θ (t) k2(6)3

<latexit sha1_base64="2KvYP7NsFo+5DzNrCEbA/Xa8hsI=">AAACFHicbVC7SgNBFJ31GeNr1dJmMAgJStiVoJZBG8sI5gHZGGYns8mQ2Qczd8Ww7kfY+Cs2ForYWtj5N84mKTTxwjCHc+7lnnvcSHAFlvVtLCwuLa+s5tby6xubW9vmzm5DhbGkrE5DEcqWSxQTPGB14CBYK5KM+K5gTXd4menNOyYVD4MbGEWs45N+wD1OCWiqax5FxcRxQ9FTI19/2IEBA5LiB8cnMHC95D499m6TIpTSUtcsWGVrXHge2FNQQNOqdc0vpxfS2GcBUEGUattWBJ2ESOBUsDTvxIpFhA5Jn7U1DIjPVCcZH5XiQ830sBdK/QLAY/b3REJ8lZnWnZlVNatl5H9aOwbvvJPwIIqBBXSyyIsFhhBnCeEel4yCGGlAqOTaK6YDIgkFnWNeh2DPnjwPGidl+7Rcua4UqhfTOHJoHx2gIrLRGaqiK1RDdUTRI3pGr+jNeDJejHfjY9K6YExn9tCfMj5/AP3Mnsc=</latexit>p(✓|x, f (t) )<latexit sha1_base64="WqZzBqPF57lFjPH3lovcJQZsrVs=">AAACDnicbVDLSsNAFJ34rPUVdelmsBQqSEmkqMuiG5cV7AOaUCbTSTt08mDmRiyxX+DGX3HjQhG3rt35N07aLLT1wjCHc+7lnnu8WHAFlvVtLC2vrK6tFzaKm1vbO7vm3n5LRYmkrEkjEcmORxQTPGRN4CBYJ5aMBJ5gbW90lentOyYVj8JbGMfMDcgg5D6nBDTVM8txJXW8SPTVONAfdmDIgEzwgxMQGHp+ej858Y97ZsmqWtPCi8DOQQnl1eiZX04/oknAQqCCKNW1rRjclEjgVLBJ0UkUiwkdkQHrahiSgCk3nZ4zwWXN9LEfSf1CwFP290RKApXZ1Z2ZSTWvZeR/WjcB/8JNeRgnwEI6W+QnAkOEs2xwn0tGQYw1IFRy7RXTIZGEgk6wqEOw509eBK3Tqn1Wrd3USvXLPI4COkRHqIJsdI7q6Bo1UBNR9Iie0St6M56MF+Pd+Ji1Lhn5zAH6U8bnD5CYnHA=</latexit>p(✓|x, f )

xrec

xrec

<latexit sha1_base64="2cxQXRfe85+mCFDpmaMPI13leKk=">AAAB/3icbVDLSsNAFJ3UV62vqODGzWARXJVEirosunFZwT6gDWEynbRDJw9mbqQlZuGvuHGhiFt/w51/46TNQlsPDBzOuZc593ix4Aos69sorayurW+UNytb2zu7e+b+QVtFiaSsRSMRya5HFBM8ZC3gIFg3lowEnmAdb3yT+50HJhWPwnuYxswJyDDkPqcEtOSaR2k/IDDyfDzJ3D6wCaSS0cw1q1bNmgEvE7sgVVSg6Zpf/UFEk4CFQAVRqmdbMTgpkcCpYFmlnygWEzomQ9bTNCQBU046y5/hU60MsB9J/ULAM/X3RkoCpaaBpyfzsGrRy8X/vF4C/pWT8jBOgIV0/pGfCAwRzsvAA66PBTHVhFDJdVZMR0QSCrqyii7BXjx5mbTPa/ZFrX5XrzauizrK6BidoDNko0vUQLeoiVqIokf0jF7Rm/FkvBjvxsd8tGQUO4foD4zPH86ZlqA=</latexit><latexit sha1_base64="2cxQXRfe85+mCFDpmaMPI13leKk=">AAAB/3icbVDLSsNAFJ3UV62vqODGzWARXJVEirosunFZwT6gDWEynbRDJw9mbqQlZuGvuHGhiFt/w51/46TNQlsPDBzOuZc593ix4Aos69sorayurW+UNytb2zu7e+b+QVtFiaSsRSMRya5HFBM8ZC3gIFg3lowEnmAdb3yT+50HJhWPwnuYxswJyDDkPqcEtOSaR2k/IDDyfDzJ3D6wCaSS0cw1q1bNmgEvE7sgVVSg6Zpf/UFEk4CFQAVRqmdbMTgpkcCpYFmlnygWEzomQ9bTNCQBU046y5/hU60MsB9J/ULAM/X3RkoCpaaBpyfzsGrRy8X/vF4C/pWT8jBOgIV0/pGfCAwRzsvAA66PBTHVhFDJdVZMR0QSCrqyii7BXjx5mbTPa/ZFrX5XrzauizrK6BidoDNko0vUQLeoiVqIokf0jF7Rm/FkvBjvxsd8tGQUO4foD4zPH86ZlqA=</latexit><latexit sha1_base64="PjGGKJAagzJfRE18Pt7JUfWBE3o=">AAAB7nicbVDLSgNBEOyNrxhfUY9eBoMQL2FXgnoMevEYwTwgiWF2MpsMmZ1dZnqFsOQjvHhQxKvf482/cZLsQRMLGoqqbrq7/FgKg6777eTW1jc2t/LbhZ3dvf2D4uFR00SJZrzBIhnptk8Nl0LxBgqUvB1rTkNf8pY/vp35rSeujYjUA05i3gvpUIlAMIpWagWPaRnPp/1iya24c5BV4mWkBBnq/eJXdxCxJOQKmaTGdDw3xl5KNQom+bTQTQyPKRvTIe9YqmjITS+dnzslZ1YZkCDSthSSufp7IqWhMZPQt50hxZFZ9mbif14nweC6lwoVJ8gVWywKEkkwIrPfyUBozlBOLKFMC3srYSOqKUObUMGG4C2/vEqaFxXvslK9r5ZqN1kceTiBUyiDB1dQgzuoQwMYjOEZXuHNiZ0X5935WLTmnGzmGP7A+fwB5CyPSg==</latexit>

f (t)
f

<latexit sha1_base64="0df+aMQQZMIi+SmF2E5H9EPmtQg=">AAACFXicbVDJSgNBEO1xjXGLevTSGIQEQpiRoB6DXjxGMAtk4tDT6Uma9Cx010jCOD/hxV/x4kERr4I3/8bOAmrig4LHe1VU1XMjwRWY5pextLyyurae2chubm3v7Ob29hsqjCVldRqKULZcopjgAasDB8FakWTEdwVruoPLsd+8Y1LxMLiBUcQ6PukF3OOUgJacXCkq2D6Bvuslw9SxgQ0hkYym+P5HLnm3SQGKadHJ5c2yOQFeJNaM5NEMNSf3aXdDGvssACqIUm3LjKCTEAmcCpZm7VixiNAB6bG2pgHxmeokk69SfKyVLvZCqSsAPFF/TyTEV2rku7pzfKqa98bif147Bu+8k/AgioEFdLrIiwWGEI8jwl2uEwAx0oRQyfWtmPaJJBR0kFkdgjX/8iJpnJSt03LlupKvXsziyKBDdIQKyEJnqIquUA3VEUUP6Am9oFfj0Xg23oz3aeuSMZs5QH9gfHwDfyGfpA==</latexit>p(xrec |x, f (t) )<latexit sha1_base64="mXTTqJLVGzQHkYptyoikscqKUlE=">AAAB6HicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkqMeiF48t2FpoQ9lsJ+3azSbsboQS+gu8eFDEqz/Jm//GbZuDtj4YeLw3w8y8IBFcG9f9dgpr6xubW8Xt0s7u3v5B+fCoreNUMWyxWMSqE1CNgktsGW4EdhKFNAoEPgTj25n/8IRK81jem0mCfkSHkoecUWOlZtgvV9yqOwdZJV5OKpCj0S9/9QYxSyOUhgmqdddzE+NnVBnOBE5LvVRjQtmYDrFrqaQRaj+bHzolZ1YZkDBWtqQhc/X3REYjrSdRYDsjakZ62ZuJ/3nd1ITXfsZlkhqUbLEoTAUxMZl9TQZcITNiYgllittbCRtRRZmx2ZRsCN7yy6ukfVH1Lqu1Zq1Sv8njKMIJnMI5eHAFdbiDBrSAAcIzvMKb8+i8OO/Ox6K14OQzx/AHzucPzd+M8w==</latexit>f<latexit sha1_base64="mXTTqJLVGzQHkYptyoikscqKUlE=">AAAB6HicbVBNS8NAEJ3Ur1q/qh69LBbBU0mkqMeiF48t2FpoQ9lsJ+3azSbsboQS+gu8eFDEqz/Jm//GbZuDtj4YeLw3w8y8IBFcG9f9dgpr6xubW8Xt0s7u3v5B+fCoreNUMWyxWMSqE1CNgktsGW4EdhKFNAoEPgTj25n/8IRK81jem0mCfkSHkoecUWOlZtgvV9yqOwdZJV5OKpCj0S9/9QYxSyOUhgmqdddzE+NnVBnOBE5LvVRjQtmYDrFrqaQRaj+bHzolZ1YZkDBWtqQhc/X3REYjrSdRYDsjakZ62ZuJ/3nd1ITXfsZlkhqUbLEoTAUxMZl9TQZcITNiYgllittbCRtRRZmx2ZRsCN7yy6ukfVH1Lqu1Zq1Sv8njKMIJnMI5eHAFdbiDBrSAAcIzvMKb8+i8OO/Ox6K14OQzx/AHzucPzd+M8w==</latexit><latexit sha1_base64="X1VmISPnNOhxK2dc5jLSlQfK+dM=">AAACD3icbVDLSgMxFM34rPVVdekmWJQKUmakqMuiG5cV7AM6ZcikmTY0kxmSO9Iy9g/c+CtuXCji1q07/8b0AWrrgcDhnHvJPcePBddg21/WwuLS8spqZi27vrG5tZ3b2a3pKFGUVWkkItXwiWaCS1YFDoI1YsVI6AtW93tXI79+x5TmkbyFQcxaIelIHnBKwEhe7iguuCGBrh+k/aHnAutDqhgd4vsf+SQ49nJ5u2iPgeeJMyV5NEXFy3267YgmIZNABdG66dgxtFKigFPBhlk30SwmtEc6rGmoJCHTrXScZ4gPjdLGQaTMk4DH6u+NlIRaD0LfTI6O1LPeSPzPayYQXLRSLuMEmKSTj4JEYIjwqBzc5iY7iIEhhCpubsW0SxShYCrMmhKc2cjzpHZadM6KpZtSvnw5rSOD9tEBKiAHnaMyukYVVEUUPaAn9IJerUfr2Xqz3iejC9Z0Zw/9gfXxDQzvnU0=</latexit>p(xrec |x, f )<latexit sha1_base64="R16uFHf8w6wU28czfzZ5uIpx/rA=">AAACL3icbVDLSgNBEJz1GeMr6tHLYBAihLAroh6DgnhUMImQjWF20psMzj6Y6RXDun/kxV/xIqKIV//C2STgs2CYoqqb7i4vlkKjbT9bU9Mzs3PzhYXi4tLyymppbb2po0RxaPBIRurSYxqkCKGBAiVcxgpY4EloedfHud+6AaVFFF7gMIZOwPqh8AVnaKRu6SSupK4XyZ4eBuajLg4AWXbnBgwHnp/eZl0X4RZTBTy7Siu4k1W/vKo/lna6pbJds0egf4kzIWUywVm39Oj2Ip4EECKXTOu2Y8fYSZlCwSVkRTfREDN+zfrQNjRkAehOOro3o9tG6VE/UuaFSEfq946UBTq/x1Tmq+rfXi7+57UT9A87qQjjBCHk40F+IilGNA+P9oSJAeXQEMaVMLtSPmCKcTQRF00Izu+T/5Lmbs3Zr+2d75XrR5M4CmSTbJEKccgBqZNTckYahJN78kheyKv1YD1Zb9b7uHTKmvRskB+wPj4BKeKq7A==</latexit><latexit sha1_base64="BjhehX1ZMznaJipFpBJaFpM0KNA=">AAACKXicbVDJSgNBEO1xjXGLevTSGIQIEmZE1GPQi8cIZoFMDD2dGtPYs9BdI4ZxfseLv+JFQVGv/og9ScD1QdOP96qoqufFUmi07Tdranpmdm6+sFBcXFpeWS2trTd1lCgODR7JSLU9pkGKEBooUEI7VsACT0LLuzrJ/dY1KC2i8ByHMXQDdhkKX3CGRuqVanEldb1I9vUwMB91cQDIsls3YDjw/PQm67kIN5gq4NlFWsGdbPfL2/V3eqWyXbVHoH+JMyFlMkG9V3py+xFPAgiRS6Z1x7Fj7KZMoeASsqKbaIgZv2KX0DE0ZAHobjq6NKPbRulTP1LmhUhH6veOlAU6v8RU5kvq314u/ud1EvSPuqkI4wQh5ONBfiIpRjSPjfaFCQDl0BDGlTC7Uj5ginE04RZNCM7vk/+S5l7VOajun+2Xa8eTOApkk2yRCnHIIamRU1InDcLJHXkgz+TFurcerVfrfVw6ZU16NsgPWB+feNGolQ==</latexit><latexit sha1_base64="ZBu5/Q8eb7SOA9Y75IEsSHrYVWE=">AAACBHicbVDLSsNAFJ34rPUVddlNsAh1UxIp6rLoxmUF+4Amhsl00g6dPJi5kZaQhRt/xY0LRdz6Ee78GydtFtp6YOBwzr3MPceLOZNgmt/ayura+sZmaau8vbO7t68fHHZklAhC2yTikeh5WFLOQtoGBpz2YkFx4HHa9cbXud99oEKyKLyDaUydAA9D5jOCQUmuXrEDDCPPTyeZawOdQCooye7TGpxmrl416+YMxjKxClJFBVqu/mUPIpIENATCsZR9y4zBSbEARjjNynYiaYzJGA9pX9EQB1Q66SxEZpwoZWD4kVAvBGOm/t5IcSDlNPDUZH6yXPRy8T+vn4B/6aQsjBOgIZl/5CfcgMjIGzEGTEUGPlUEE8HUrQYZYYEJqN7KqgRrMfIy6ZzVrfN647ZRbV4VdZRQBR2jGrLQBWqiG9RCbUTQI3pGr+hNe9JetHftYz66ohU7R+gPtM8fvqmYzQ==</latexit>

x(t)
rec

p(✓|x(t)rec , x, f )<latexit sha1_base64="ZBu5/Q8eb7SOA9Y75IEsSHrYVWE=">AAACBHicbVDLSsNAFJ34rPUVddlNsAh1UxIp6rLoxmUF+4Amhsl00g6dPJi5kZaQhRt/xY0LRdz6Ee78GydtFtp6YOBwzr3MPceLOZNgmt/ayura+sZmaau8vbO7t68fHHZklAhC2yTikeh5WFLOQtoGBpz2YkFx4HHa9cbXud99oEKyKLyDaUydAA9D5jOCQUmuXrEDDCPPTyeZawOdQCooye7TGpxmrl416+YMxjKxClJFBVqu/mUPIpIENATCsZR9y4zBSbEARjjNynYiaYzJGA9pX9EQB1Q66SxEZpwoZWD4kVAvBGOm/t5IcSDlNPDUZH6yXPRy8T+vn4B/6aQsjBOgIZl/5CfcgMjIGzEGTEUGPlUEE8HUrQYZYYEJqN7KqgRrMfIy6ZzVrfN647ZRbV4VdZRQBR2jGrLQBWqiG9RCbUTQI3pGr+hNe9JetHftYz66ohU7R+gPtM8fvqmYzQ==</latexit>

x(t)
rec

(t)
p(✓|x(t))
rec , x, f<latexit sha1_base64="Tufr5DbNYu8EtBQx4Af+1O49IMI=">AAACMXicbVDLSgNBEJz1GeMr6tHLYBAiSNiVoB6DXnJUMImQjWF20msGZx/M9ErCur/kxT8RLx4U8epPOJsEfDYMU1R30dXlxVJotO1na2Z2bn5hsbBUXF5ZXVsvbWy2dJQoDk0eyUhdekyDFCE0UaCEy1gBCzwJbe/mNO+3b0FpEYUXOIqhG7DrUPiCMzRUr9SIK27AcOD56TDruQhDTBXwjN7R1PUi2dejwHzUxQEgy67SCu5l+1+SfX9C7fVKZbtqj4v+Bc4UlMm0znqlR7cf8SSAELlkWnccO8ZuyhQKLiEruomGmPEbdg0dA0MWgO6m44szumuYPvUjZV6IdMx+V6Qs0LlzM5lb1b97Oflfr5Ogf9xNRRgnCCGfLPITSTGieXy0L0w6KEcGMK6E8Ur5gCnG0YRcNCE4v0/+C1oHVeewWjuvlesn0zgKZJvskApxyBGpkwY5I03CyT15Ii/k1Xqwnq03630yOmNNNVvkR1kfn/4Aq0A=</latexit>p(xrec |✓ (t) , x, f (t) )<latexit sha1_base64="AXgQpWVX/UwETRTFmerAC0uozT0=">AAACK3icbVDJSgNBEO1xjXGLevTSGIQIEmYkqEeJF48KxgQyMfR0akxjz0J3jRjG+R8v/ooHPbjg1f+wJwm4xIKmH6/qUa+eF0uh0bbfrKnpmdm5+cJCcXFpeWW1tLZ+oaNEcWjwSEaq5TENUoTQQIESWrECFngSmt71cd5v3oDSIgrPcRBDJ2BXofAFZ2iobqkeV9yAYd/z09us6yLcYqqAZ/SOpq4XyZ4eBOajLvYBWXaZVnAn2/2W7Po73VLZrtrDopPAGYMyGddpt/Tk9iKeBBAil0zrtmPH2EmZQsElZEU30RAzfs2uoG1gyALQnXR4a0a3DdOjfqTMC5EO2Z+KlAU692wmc5P6by8n/+u1E/QPO6kI4wQh5KNFfiIpRjQPjvaEyQXlwADGlTBeKe8zxTiaeIsmBOfvyZPgYq/q7FdrZ7XyUX0cR4Fski1SIQ45IEfkhJySBuHknjySF/JqPVjP1rv1MRqdssaaDfKrrM8vS2ao6Q==</latexit>p(xrec |✓ (t) , x, f )<latexit sha1_base64="378pGGggjM8Q53SlrPI1a49ctwc=">AAAB83icbVBNS8NAEJ3Ur1q/qh69BItQLyWRoh6LXjxWsB/QxLLZbtqlm03YnQgl9G948aCIV/+MN/+N2zYHbX0w8Hhvhpl5QSK4Rsf5tgpr6xubW8Xt0s7u3v5B+fCoreNUUdaisYhVNyCaCS5ZCzkK1k0UI1EgWCcY3878zhNTmsfyAScJ8yMylDzklKCRPA9HDMljVsXzab9ccWrOHPYqcXNSgRzNfvnLG8Q0jZhEKojWPddJ0M+IQk4Fm5a8VLOE0DEZsp6hkkRM+9n85ql9ZpSBHcbKlER7rv6eyEik9SQKTGdEcKSXvZn4n9dLMbz2My6TFJmki0VhKmyM7VkA9oArRlFMDCFUcXOrTUdEEYomppIJwV1+eZW0L2ruZa1+X680bvI4inACp1AFF66gAXfQhBZQSOAZXuHNSq0X6936WLQWrHzmGP7A+vwBybaRiA==</latexit>

✓ (t)
✓(t)
(a) Non linear f<latexit sha1_base64="F/w4Kniq7SYtkdXe3Ilv2qX75uE=">AAACA3icbVDLSsNAFJ34rPUVdaebwSLUTUmkqMuiG5cV7AOaWCbTSTt08mDmRigh4MZfceNCEbf+hDv/xkmbhbYeGOZwzr3ce48XC67Asr6NpeWV1bX10kZ5c2t7Z9fc22+rKJGUtWgkItn1iGKCh6wFHATrxpKRwBOs442vc7/zwKTiUXgHk5i5ARmG3OeUgJb65mHqeJEYqEmgP+zAiAHJ7tMqnGZ9s2LVrCnwIrELUkEFmn3zyxlENAlYCFQQpXq2FYObEgmcCpaVnUSxmNAxGbKepiEJmHLT6Q0ZPtHKAPuR1C8EPFV/d6QkUPmWujIgMFLzXi7+5/US8C/dlIdxAiyks0F+IjBEOA8ED7hkFMREE0Il17tiOiKSUNCxlXUI9vzJi6R9VrPPa/XbeqVxVcRRQkfoGFWRjS5QA92gJmohih7RM3pFb8aT8WK8Gx+z0iWj6DlAf2B8/gBKepfw</latexit><latexit sha1_base64="KuvvGSHVOL0Fo0vYGZksqTvaZYs=">AAACAnicbVDLSgMxFM3UV62vUVfiJliEuikzUtRl0Y0rqWAf0JaSSTNtaCYzJHfEMhQ3/oobF4q49Svc+Tdm2llo64ELh3PuTe49XiS4Bsf5tnJLyyura/n1wsbm1vaOvbvX0GGsKKvTUISq5RHNBJesDhwEa0WKkcATrOmNrlK/ec+U5qG8g3HEugEZSO5zSsBIPfugA+wBkhI5wTehxOkzROEJ9gs9u+iUnSnwInEzUkQZaj37q9MPaRwwCVQQrduuE0E3IQo4FWxS6MSaRYSOyIC1DZUkYLqbTE+Y4GOj9LEfKlMS8FT9PZGQQOtx4JnOgMBQz3up+J/XjsG/6CZcRjEwSWcf+bHAEOI0D9znilEQY0MIVdzsiumQKELBpJaG4M6fvEgap2X3rFy5rRSrl1kceXSIjlAJuegcVdE1qqE6ougRPaNX9GY9WS/Wu/Uxa81Z2cw++gPr8wdC2ZYP</latexit>✓<latexit sha1_base64="wfHxuECdHRyEThgDlsUPohDkDuA=">AAAB/XicbVDLSgMxFM3UV62v8bFzEyyCqzIjRV0W3bisYB/QGUomk7ahmcyQ3BHqUPwVNy4Ucet/uPNvzLSz0NYDIYdz7iUnJ0gE1+A431ZpZXVtfaO8Wdna3tnds/cP2jpOFWUtGotYdQOimeCStYCDYN1EMRIFgnWC8U3udx6Y0jyW9zBJmB+RoeQDTgkYqW8fZV4Qi1BPInNhD0YMyLRvV52aMwNeJm5BqqhAs29/eWFM04hJoIJo3XOdBPyMKOBUsGnFSzVLCB2TIesZKknEtJ/N0k/xqVFCPIiVORLwTP29kZFI5/nMZERgpBe9XPzP66UwuPIzLpMUmKTzhwapwBDjvAoccsUoiIkhhCpusmI6IopQMIVVTAnu4peXSfu85l7U6nf1auO6qKOMjtEJOkMuukQNdIuaqIUoekTP6BW9WU/Wi/VufcxHS1axc4j+wPr8AQRFlZk=</latexit><latexit sha1_base64="F/w4Kniq7SYtkdXe3Ilv2qX75uE=">AAACA3icbVDLSsNAFJ34rPUVdaebwSLUTUmkqMuiG5cV7AOaWCbTSTt08mDmRigh4MZfceNCEbf+hDv/xkmbhbYeGOZwzr3ce48XC67Asr6NpeWV1bX10kZ5c2t7Z9fc22+rKJGUtWgkItn1iGKCh6wFHATrxpKRwBOs442vc7/zwKTiUXgHk5i5ARmG3OeUgJb65mHqeJEYqEmgP+zAiAHJ7tMqnGZ9s2LVrCnwIrELUkEFmn3zyxlENAlYCFQQpXq2FYObEgmcCpaVnUSxmNAxGbKepiEJmHLT6Q0ZPtHKAPuR1C8EPFV/d6QkUPmWujIgMFLzXi7+5/US8C/dlIdxAiyks0F+IjBEOA8ED7hkFMREE0Il17tiOiKSUNCxlXUI9vzJi6R9VrPPa/XbeqVxVcRRQkfoGFWRjS5QA92gJmohih7RM3pFb8aT8WK8Gx+z0iWj6DlAf2B8/gBKepfw</latexit>✓ (t)(b) Linear f✓<latexit sha1_base64="wfHxuECdHRyEThgDlsUPohDkDuA=">AAAB/XicbVDLSgMxFM3UV62v8bFzEyyCqzIjRV0W3bisYB/QGUomk7ahmcyQ3BHqUPwVNy4Ucet/uPNvzLSz0NYDIYdz7iUnJ0gE1+A431ZpZXVtfaO8Wdna3tnds/cP2jpOFWUtGotYdQOimeCStYCDYN1EMRIFgnWC8U3udx6Y0jyW9zBJmB+RoeQDTgkYqW8fZV4Qi1BPInNhD0YMyLRvV52aMwNeJm5BqqhAs29/eWFM04hJoIJo3XOdBPyMKOBUsGnFSzVLCB2TIesZKknEtJ/N0k/xqVFCPIiVORLwTP29kZFI5/nMZERgpBe9XPzP66UwuPIzLpMUmKTzhwapwBDjvAoccsUoiIkhhCpusmI6IopQMIVVTAnu4peXSfu85l7U6nf1auO6qKOMjtEJOkMuukQNdIuaqIUoekTP6BW9WU/Wi/VufcxHS1axc4j+wPr8AQRFlZk=</latexit><latexit sha1_base64="S2x9aUIVoFa/gptVKX71tgwgNm4=">AAACBHicbVC7TgJBFJ3FF+Jr1ZJmIjGBhuwaopZEGwsLTOSRwEpmh1mYMPvIzF0j2VDY+Cs2Fhpj60fY+TfOwhYKnuQmJ+fcO3PvcSPBFVjWt5FbWV1b38hvFra2d3b3zP2DlgpjSVmThiKUHZcoJnjAmsBBsE4kGfFdwdru+DL12/dMKh4GtzCJmOOTYcA9TgloqW8We8AeICm7FXytnyAST7F3l5ShMi30zZJVtWbAy8TOSAllaPTNr94gpLHPAqCCKNW1rQichEjgVLBpoRcrFhE6JkPW1TQgPlNOMjtiio+1MsBeKHUFgGfq74mE+EpNfFd3+gRGatFLxf+8bgzeuZPwIIqBBXT+kRcLDCFOE8EDLhkFMdGEUMn1rpiOiCQUdG5pCPbiycukdVK1T6u1m1qpfpHFkUdFdITKyEZnqI6uUAM1EUWP6Bm9ojfjyXgx3o2PeWvOyGYO0R8Ynz+rLpbU</latexit>(t)

Figure 3: Illustrative example for fixed x. The likelihood for a fixed θ (t) , shown by the columns, areassumed Gaussian p(xrec |θ (t) , x, f ) = N (fθ(t) (x), I). We model the marginalised density p(θ|x, f )(first axis) over parameters θ with Gaussians. With the additional assumption of linear f thenp(xrec |x, f ) (second axis) is Gaussian. This makes the joint distribution tractable.
where f (t) denote the associated first-order approximation. Under this approximation,
 the joint
distribution p(θ, xrec |x, f ) becomes Gaussian, p(θ, xrec |x, f (t) ) ≈ N θ, xrec |µ(t) , Σ(t) , with
!


H−1
Jθ fθ(t) (x)>
θ (t)
t
(t)
(t)
−1
µ =
, and Σ =. (7)(t)
Jθ fθ(t) (x) Jθ fθ(t) (x)> Ht Jθ fθ(t) (x)+Ifθ(t) (x)This approximation is illustrated in Fig. 3(b).

Iterative learning. With these approximations we can readily develop an iterative learning scheme
which updates q(θ|x, f ). The mean of this approximate posterior can be updated according to a
standard variational gradient step, θ (t+1) = θ (t) + λ∇θ Lxrec , where
Lxrec = Eθ∼qt (θ|x,f ) [log p(xrec |θ, x, f (t) )],(8)

is a lower bound on Eq. 3, which we evaluate with a single Monte Carlo sample. The last missing pieceis the covariance of q t+1 (θ|x, f ), which ideally should be the inverse of the Hessian of log p(θ|x, f )at θ (t+1) . Since this is intractable, we instead compute the Hessian of the lower bound (5)hi
Ht+1 = −∇2θ Lθ θ(t+1) = Ep(xrec |x,f (t) ) −∇2θ log p(xrec |θ, x, f (t) ) −∇2θ log q t (θ|x, f ) (t+1) (9)θThe= −Ht +
 last term can be approximated by noting that
(t+1)
(t)
O kθ
− θ k . To efficiently cope with the O-term, we introduce a parameter α, suchthat the final approximation ishi
Ht+1 ≈ (1 − α)Ht + Ep(xrec |x,f ) −∇2θ log p xrec θ, x, f (t)(10)θ (t+1)


>
(t)
(t)
= (1 − α)Ht − Jθ fθ ∇2xrec log p xrec θ (t+1) , x, f (t) Jθ fθ .
(11)∇2θ log q t (θ|x, f )|θ=θ(t+1)

The practical intuition is that when α is small, the network uncertainty decreases faster. The overall
training procedure is summarized in Fig. 4. We initialize q 0 as a Gaussian with θ (0) = 0 and H0 = I.
We provide more details on the model, the linearization, and iterative learning in Appendix D.

Why not just. . . ? The proposed training procedure may at first appear non-trivial, and it isreasonable to wonder if existing methods could be applied to similar results. Variational inference
often achieves similar results to Laplace approximations, so could we use ‘Bayes by Backprop’(Blundell et al., 2015) to get an alternative Gaussian approximate posterior over θ? Similar to thesupervised experiences of Jospin et al. (2020), we unfortunately found this approach too brittle toallow for practical model fitting. But then perhaps a post-hoc LA as proposed by Daxberger et al.
4t+=1<latexit sha1_base64="h8Pt7eIf+S1zbJhwPYfujaumzrQ=">AAAB73icbVBNS8NAEJ34WetX1aOXYCMIQkkKoheh6MVjBfsBbSib7aZdutnE3YlQQv+EFw+KePXvePPfuG1z0NYHA4/3ZpiZFySCa3Tdb2tldW19Y7OwVdze2d3bLx0cNnWcKsoaNBaxagdEM8ElayBHwdqJYiQKBGsFo9up33piSvNYPuA4YX5EBpKHnBI0UttB5/za8ZxeqexW3BnsZeLlpAw56r3SV7cf0zRiEqkgWnc8N0E/Iwo5FWxS7KaaJYSOyIB1DJUkYtrPZvdO7FOj9O0wVqYk2jP190RGIq3HUWA6I4JDvehNxf+8TorhlZ9xmaTIJJ0vClNhY2xPn7f7XDGKYmwIoYqbW206JIpQNBEVTQje4svLpFmteBcV975art3kcRTgGE7gDDy4hBrcQR0aQEHAM7zCm/VovVjv1se8dcXKZ47gD6zPH6vWjmk=</latexit><latexit sha1_base64="06KnXmEKdfwDSkvj0AQGFGg+QkM=">AAACEHicbVC7SgNBFJ2NrxhfUUubxSAmIGFXFG2EoI1lBPOAbAyzk7vJkNkHM3fFZckn2PgrNhaK2Fra+TdOHoUmHhg4nHMPd+5xI8EVWta3kVlYXFpeya7m1tY3Nrfy2zt1FcaSQY2FIpRNlyoQPIAachTQjCRQ3xXQcAdXI79xD1LxMLjFJIK2T3sB9zijqKVO/tC7S4tYGl44CA+oWIo00Zlh0TtysA9IJ3apky9YZWsMc57YU1IgU1Q7+S+nG7LYhwCZoEq1bCvCdkolciZgmHNiBRFlA9qDlqYB9UG10/FBQ/NAK13TC6V+AZpj9Xcipb5Sie/qSZ9iX816I/E/rxWjd95OeRDFCAGbLPJiYWJojtoxu1wCQ5FoQpnk+q8m61NJGeoOc7oEe/bkeVI/LtunZevmpFC5nNaRJXtknxSJTc5IhVyTKqkRRh7JM3klb8aT8WK8Gx+T0YwxzeySPzA+fwCZrpzt</latexit>f (t) = taylor(f, ✓(t) )<latexit sha1_base64="cb3tJBjMq0SnGYZ0RlEPCLzJJPk=">AAACwXicdVFdaxQxFM1M/ajr16qPvgQXYUvLMlOU+qKUFsHHCm5b2JkOmUxmNzQf0+SOdEnzJ/uk/8bMdkRt1wshh3vOyb25t2wEt5AkP6N44979Bw83Hw0eP3n67Pnwxctjq1tD2ZRqoc1pSSwTXLEpcBDstDGMyFKwk/L8sONPvjNjuVbfYNmwXJK54jWnBEKqGP5wWalFZZcyXDiDBQPiz9wYttMtjz/i/9GB3MaZCIUqgjNJYFGW7rMv3DpDZrnsRbW78GfgZ5kipSDr5T48rOe4wePfnktfZMAuwRlG/dU6zw7+o93B9VZeDEfJJFkFvgvSHoxQH0fF8DqrNG0lU0AFsXaWJg3kjhjgVDA/yFrLGkLPyZzNAlREMpu71QY8fhsyFa61CUcBXmX/djgibddwUHZt2ttcl1zHzVqoP+SOq6YFpuhNoboVGDTu1okrHkYCYhkAoYaHXjFdEEMohKUPwhDS21++C453J+n7SfL13Wj/oB/HJnqN3qAxStEe2kdf0BGaIhp9iqpIRio+jHncxOZGGke95xX6J2L3C2eG3ww=</latexit>

✓

(t+1)

=✓

(t)

+ E✓⇠qt [r✓ log p(xrec |✓, x, f )]<latexit sha1_base64="nN9ydiRBc7synEi9rg21xfZLybU=">AAAB7nicbVDLSgNBEOyNrxhfUY9eBoMQL2FXFD0GvXiMYB6QxDA7mU2GzM6uM71CWPIRXjwo4tXv8ebfOEn2oIkFDUVVN91dfiyFQdf9dnIrq2vrG/nNwtb2zu5ecf+gYaJEM15nkYx0y6eGS6F4HQVK3oo1p6EvedMf3Uz95hPXRkTqHscx74Z0oEQgGEUrNR8f0jKeTnrFkltxZyDLxMtICTLUesWvTj9iScgVMkmNaXtujN2UahRM8kmhkxgeUzaiA962VNGQm246O3dCTqzSJ0GkbSkkM/X3REpDY8ahbztDikOz6E3F/7x2gsFVNxUqTpArNl8UJJJgRKa/k77QnKEcW0KZFvZWwoZUU4Y2oYINwVt8eZk0zireRcW9Oy9Vr7M48nAEx1AGDy6hCrdQgzowGMEzvMKbEzsvzrvzMW/NOdnMIfyB8/kD836PUA==</latexit>q (t)<latexit sha1_base64="Hlr8pNcIwIuOJ+TeetO8yvxjghM=">AAAB83icbVDJSgNBEO2JW4xb1KOXxiDEyzATEp3cgl48RjALZMbQ0+lJmvQsdNcIYchvePGgiFd/xpt/Y2dB1Pig4PFeFVX1/ERwBZb1aeTW1jc2t/LbhZ3dvf2D4uFRW8WppKxFYxHLrk8UEzxiLeAgWDeRjIS+YB1/fD3zOw9MKh5HdzBJmBeSYcQDTgloyXVhxIDcZ2U4n/aLJcus16tOzcGaOBeOXcG2ac3xTUpoiWa/+OEOYpqGLAIqiFI920rAy4gETgWbFtxUsYTQMRmynqYRCZnysvnNU3ymlQEOYqkrAjxXf05kJFRqEvq6MyQwUn+9mfif10shcLyMR0kKLKKLRUEqMMR4FgAecMkoiIkmhEqub8V0RCShoGMq6BBWXl4l7Ypp10zrtlpqXC3jyKMTdIrKyEaXqIFuUBO1EEUJekTP6MVIjSfj1XhbtOaM5cwx+gXj/Qs6rJHR</latexit>✓(t)<latexit sha1_base64="EORz1PPkHhv/aAkh5TV1mZLqHZc=">AAAB8XicbVDLSgNBEJz1GeMr6tHLYBC8GHZDoptb0EuOEcwDkzXMTmaTIbOzy0yvEJb8hRcPinj1b7z5N04eiBoLGoqqbrq7/FhwDbb9aa2srq1vbGa2sts7u3v7uYPDpo4SRVmDRiJSbZ9oJrhkDeAgWDtWjIS+YC1/dD31Ww9MaR7JWxjHzAvJQPKAUwJGuqv1Upjcp+fOpJfL24VKpeSWXWyIe+E6RewU7Bm+SR4tUO/lPrr9iCYhk0AF0brj2DF4KVHAqWCTbDfRLCZ0RAasY6gkIdNeOrt4gk+N0sdBpExJwDP150RKQq3HoW86QwJD/debiv95nQQC10u5jBNgks4XBYnAEOHp+7jPFaMgxoYQqri5FdMhUYSCCSlrQlh6eZk0iwWnXLBvSvnq1SKODDpGJ+gMOegSVVEN1VEDUSTRI3pGL5a2nqxX623eumItZo7QL1jvX7v5kPc=</latexit>

Ht<latexit sha1_base64="OrmFwE+RoR4/bKLrExIobaZF8yA=">AAAC73iclVLLbhMxFPUMFEp4NClLNldUSIloo5lKiG4qVWVTsSpS01aKJyOP40mseh6M71SNjH+iGxYgxJbfYcePsMaTRKi02XAly0f3nPu0k1JJjUHwy/Pv3V978HD9Uevxk6fPNtqdzVNd1BUXA16oojpPmBZK5mKAEpU4LyvBskSJs+TiXcOfXYpKyyI/wVkpooxNcplKztC54o63RjOG0yQ1RzY2+Dq0sA/dcIcyVU5Z7waJsAPmfWxoUqixnmXuAopTgcymI9PFnl3J2dEJ0G2gOUsUG+06zTLllY0piis0leDWAlXFBErorqQ/rcrskv7VbsOyh15T7P+7jNtbQT+YG9wF4RJsHfRoJ/o9oMdx+ycdF7zORI5cMa2HYVBiZFiFkithW7TWomT8gk3E0MGcZUJHZv5eFl45zxjSonInR5h7b0YYlummP6dsJtS3uca5ihvWmO5FRuZljSLni0JprQALaB4fxtJtE9XMAcYr6XoFPmUV4+i+SMstIbw98l1wutsP3/SDD24bh2Rh6+QFeUm6JCRvyQE5IsdkQLinvGvvi/fV/+h/9r/53xdS31vGPCf/mP/jD4rP8mw=</latexit>

Ht+1 = (1

↵)Ht

1

(t) T

J✓ f ✓(t)r2xrec log p(xrec |✓, x, f (t) ) J✓ f✓t+=1<latexit sha1_base64="h8Pt7eIf+S1zbJhwPYfujaumzrQ=">AAAB73icbVBNS8NAEJ34WetX1aOXYCMIQkkKoheh6MVjBfsBbSib7aZdutnE3YlQQv+EFw+KePXvePPfuG1z0NYHA4/3ZpiZFySCa3Tdb2tldW19Y7OwVdze2d3bLx0cNnWcKsoaNBaxagdEM8ElayBHwdqJYiQKBGsFo9up33piSvNYPuA4YX5EBpKHnBI0UttB5/za8ZxeqexW3BnsZeLlpAw56r3SV7cf0zRiEqkgWnc8N0E/Iwo5FWxS7KaaJYSOyIB1DJUkYtrPZvdO7FOj9O0wVqYk2jP190RGIq3HUWA6I4JDvehNxf+8TorhlZ9xmaTIJJ0vClNhY2xPn7f7XDGKYmwIoYqbW206JIpQNBEVTQje4svLpFmteBcV975art3kcRTgGE7gDDy4hBrcQR0aQEHAM7zCm/VovVjv1se8dcXKZ47gD6zPH6vWjmk=</latexit>µ✓ =0Figure 4: Recursive training procedure. Givena distribution qt over parameters, and a linearizedfunction f (t) , compute first and second order derivatives to update the distribution on parameters.
<latexit sha1_base64="qVw4aSXTNs8FzZLhRtfIOAr69mg=">AAAB/HicbVDLSgNBEJz1GeMrmqMig0HwFHbFx+YgBL14VDAmkA1hdtIxQ2YfzPSKYYk3v8OLB0W8it/hzW/wJ5wkImosaCiquunu8mMpNNr2uzUxOTU9M5uZy84vLC4t51ZWL3SUKA4VHslI1XymQYoQKihQQi1WwAJfQtXvHg/86hUoLaLwHHsxNAJ2GYq24AyN1MzlvSBpetgBZB7CNaaHfbuZK9jFUmnX3XOpIe6+6+xQp2gP8U0K5fXXs4/bjdfTZu7Na0U8CSBELpnWdceOsZEyhYJL6Ge9REPMeJddQt3QkAWgG+nw+D7dMkqLtiNlKkQ6VH9OpCzQuhf4pjNg2NF/vYH4n1dPsO02UhHGCULIR4vaiaQY0UEStCUUcJQ9QxhXwtxKeYcpxtHklTUhjL08Ti52is5e0T4zaRyRETJkjWySbeKQA1ImJ+SUVAgnPXJHHsijdWPdW0/W86h1wvqayZNfsF4+AWbxmPc=</latexit>(2021) for supervised learning? Empirically, we found it to be important to center the approximate
posterior around a point where the Hessian provides useful uncertainty estimates. Our online training
moves in this direction as the Hessian is part of the procedure, but this is not true for the post-hoc LA.
f^{(t)}=\textsc{taylor}(f,\theta^{(t)})

Conceptually, we argue that our approach, while novel, is not entirely separate from existing methods.
Our reliance on lower bounds make the method an instance of variational inference (Jordan et al.,
1999; Opper and Archambeau, 2009b), and we maximize the bounds using Monte Carlo EM (Cappé,2009). We rely on a LA as our choice of variational distribution, which has also been explored byPark et al. (2019). Finally, we note that our linearization trick (6) bares great similarities to classic
extended Kalman filtering (Gelb et al., 1974).

3

Scaling the Hessian to Large Images

The largest obstacle to apply LA in practice stem from the Hessian matrix. This matrix has a quadratic
memory complexity in the number of network parameters, which very quickly exceeds the capabilities
of available hardware. To counter this issue, several approximations have been proposed (Ritter et al.,2018; Botev et al., 2017; Martens and Grosse, 2015b) that improve the scaling w.r.t. to the number
parameters. The currently most efficient Hessian implementations (Dangel et al., 2020; Daxberger
et al., 2021) builds on the generalized Gauss-Newton (GGN) approximation of the Hessian>∇2θl L(fθ (x)) ≈ Jθl fθ (x) · ∇2xrec L (xrec ) · Jθl fθ (x) ,(12)which neglects second order derivatives of f w.r.t. the parameters. Besides, the computational benefitsof this approximations, previous works on LA (Daxberger et al., 2021) relies on GGN to ensure thatthe Hessian is always semi-negative definite. In contrast, the model presented in Sec. 2 implies that
GGN is no longer a practical and unprincipled trick, but rather the exact Hessian.
Albeit relying on first order derivates, the layer-block-diagonal GGN, which assume that layers are✓l scales quadratically with the output dimension of the considered neural
independent to each other,
{x
}
{x }
network f . This lack xof=scaling
is particularly
detrimental for convolutional layers as these have low
f xparameter counts, but potentially very high output dimensions.
{M
}
{M }
✓l
J f theM chainJ f rule, one realises that
Expanding Jθl fθM(x)=with⇣⌘
x(l) = f✓ x(l 1)x(l 1)x(l)the Jacobian can be computed as a function of the Jacobian of the next layer.
that an intermediate
H ⇡ Fig.
J f 5Millustrate
J f
(l 1)
= JxT M (l) Jx
quantity M , which is initialised as ∇2xrec L(xrec ), can be
M (l)M (l 1) M
efficiently backpropagated through multiplication with theJacobian w.r.t. input of each layer. This process leads to
H (l) ⇡ J✓T M (l) J✓a block diagonal approximation of the Hessian as illustrated in Fig. 6(a). However, diagonal blocks are generally
too large to store and invert. To combat this, each block can Figure 5: Forward pass of feature map x
be further approximated by its exact diagonal (Daxberger for layer l with parameters θ l and exet al., 2021) as depicted in Fig. 6(b). This scales linearly tended backward pass in which M is
w.r.t. parameters, but still scales quadratically w.r.t. the backpropagated to previous layers. Via
the chain rule and M the Hessian of each
output resolution (Tab. 1).
layer can be computed efficiently.
In order to scale our Laplacian autoencoder to highdimensional data, we propose to approximate the diagonal of the Hessian rather than relying on
exact computations. This is achieved by only backpropagating a diagonal form of M as illustrated
in Fig. 6(c). This assumes that features from the same layer are uncorrelated, and consequently<latexit sha1_base64="f4PnUuGIw+bSuQRFBwHCC4ZqL+E=">AAAB73icbVBNS8NAEJ3Ur1q/qh69BIvgqSSi6LHoxWMFWwttKJvtpF262cTdiVBK/4QXD4p49e9489+4bXPQ1gcDj/dmmJkXplIY8rxvp7Cyura+UdwsbW3v7O6V9w+aJsk0xwZPZKJbITMohcIGCZLYSjWyOJT4EA5vpv7DE2ojEnVPoxSDmPWViARnZKVWhwZIrCu75YpX9WZwl4mfkwrkqHfLX51ewrMYFXHJjGn7XkrBmGkSXOKk1MkMpowPWR/blioWownGs3sn7olVem6UaFuK3Jn6e2LMYmNGcWg7Y0YDs+hNxf+8dkbRVTAWKs0IFZ8vijLpUuJOn3d7QiMnObKEcS3srS4fMM042YhKNgR/8eVl0jyr+hdV7+68UrvO4yjCERzDKfhwCTW4hTo0gIOEZ3iFN+fReXHenY95a8HJZw7hD5zPHyg0kAs=</latexit><latexit sha1_base64="PhUk9zVCJtGlg9ebMlArNnMx7OU=">AAACAHicbVDLSsNAFJ3UV62vqAsXbgaLUBeWRBTdCEU3rqSCfUCThsl00g6dTMLMRCwhG3/FjQtF3PoZ7vwbp20W2nrgwuGce7n3Hj9mVCrL+jYKC4tLyyvF1dLa+sbmlrm905RRIjBp4IhFou0jSRjlpKGoYqQdC4JCn5GWP7we+60HIiSN+L0axcQNUZ/TgGKktOSZe0762E0r7Ng+yjzuZF7KL+2se+uZZatqTQDniZ2TMshR98wvpxfhJCRcYYak7NhWrNwUCUUxI1nJSSSJER6iPuloylFIpJtOHsjgoVZ6MIiELq7gRP09kaJQylHo684QqYGc9cbif14nUcGFm1IeJ4pwPF0UJAyqCI7TgD0qCFZspAnCgupbIR4ggbDSmZV0CPbsy/OkeVK1z6rW3Wm5dpXHUQT74ABUgA3OQQ3cgDpoAAwy8AxewZvxZLwY78bHtLVg5DO74A+Mzx+88JXa</latexit>(l 1) Nnn=1

<latexit sha1_base64="AFR+QBglecMruAL91aI9PFwm5nM=">AAAB/nicbVDLSsNAFJ34rPUVFVduBotQNyURRTdC0Y0rqWAf0KRhMp22QyeTMDMRyxDwV9y4UMSt3+HOv3HaZqGtBy4czrmXe+8JE0alcpxva2FxaXlltbBWXN/Y3Nq2d3YbMk4FJnUcs1i0QiQJo5zUFVWMtBJBUBQy0gyH12O/+UCEpDG/V6OE+BHqc9qjGCkjBfa+px87usyOs4B7WaD5pZt1bgO75FScCeA8cXNSAjlqgf3ldWOcRoQrzJCUbddJlK+RUBQzkhW9VJIE4SHqk7ahHEVE+npyfgaPjNKFvViY4gpO1N8TGkVSjqLQdEZIDeSsNxb/89qp6l34mvIkVYTj6aJeyqCK4TgL2KWCYMVGhiAsqLkV4gESCCuTWNGE4M6+PE8aJxX3rOLcnZaqV3kcBXAADkEZuOAcVMENqIE6wECDZ/AK3qwn68V6tz6mrQtWPrMH/sD6/AHZSJVo</latexit>

(l) N
n n=1

<latexit sha1_base64="0J5S3oTDHViRKLV3dNr31aMBAnk=">AAACFnicbVDLSsNAFJ34tr6iLt0MFqFdtCSi6EYQ3bhUsK3Q1DCZTtqhk0mYuRFLyFe48VfcuFDErbjzb5zULHwduHDmnHuZe0+QCK7BcT6sqemZ2bn5hcXK0vLK6pq9vtHWcaooa9FYxOoqIJoJLlkLOAh2lShGokCwTjA6LfzODVOax/ISxgnrRWQgecgpASP5duPWl9dZTdRzfIRDP/NgyID4IvcCPqjdFlbDree+LN513646TWcC/Je4JamiEue+/e71Y5pGTAIVROuu6yTQy4gCTgXLK16qWULoiAxY11BJIqZ72eSsHO8YpY/DWJmSgCfq94mMRFqPo8B0RgSG+rdXiP953RTCw17GZZICk/TrozAVGGJcZIT7XDEKYmwIoYqbXTEdEkUomCQrJgT398l/SXu36e43nYu96vFJGccC2kLbqIZcdICO0Rk6Ry1E0R16QE/o2bq3Hq0X6/WrdcoqZzbRD1hvn5tanmM=</latexit>

(l)
n

(l 1)
n

✓l

<latexit sha1_base64="f4PnUuGIw+bSuQRFBwHCC4ZqL+E=">AAAB73icbVBNS8NAEJ3Ur1q/qh69BIvgqSSi6LHoxWMFWwttKJvtpF262cTdiVBK/4QXD4p49e9489+4bXPQ1gcDj/dmmJkXplIY8rxvp7Cyura+UdwsbW3v7O6V9w+aJsk0xwZPZKJbITMohcIGCZLYSjWyOJT4EA5vpv7DE2ojEnVPoxSDmPWViARnZKVWhwZIrCu75YpX9WZwl4mfkwrkqHfLX51ewrMYFXHJjGn7XkrBmGkSXOKk1MkMpowPWR/blioWownGs3sn7olVem6UaFuK3Jn6e2LMYmNGcWg7Y0YDs+hNxf+8dkbRVTAWKs0IFZ8vijLpUuJOn3d7QiMnObKEcS3srS4fMM042YhKNgR/8eVl0jyr+hdV7+68UrvO4yjCERzDKfhwCTW4hTo0gIOEZ3iFN+fReXHenY95a8HJZw7hD5zPHyg0kAs=</latexit><latexit sha1_base64="JROU3DIP4PoWJ+YdBbHHPQDoxM0=">AAACAHicbVDLSsNAFJ3UV62vqAsXbgaLUBeWRBTdCEU3bpQK9gFNGibTSTt0MgkzE6GEbPwVNy4UcetnuPNvnLZZaOuBC4dz7uXee/yYUaks69soLCwuLa8UV0tr6xubW+b2TlNGicCkgSMWibaPJGGUk4aiipF2LAgKfUZa/vB67LceiZA04g9qFBM3RH1OA4qR0pJn7jnpbTetsGP7KPO4k3kpv7Sz7p1nlq2qNQGcJ3ZOyiBH3TO/nF6Ek5BwhRmSsmNbsXJTJBTFjGQlJ5EkRniI+qSjKUchkW46eSCDh1rpwSASuriCE/X3RIpCKUehrztDpAZy1huL/3mdRAUXbkp5nCjC8XRRkDCoIjhOA/aoIFixkSYIC6pvhXiABMJKZ1bSIdizL8+T5knVPqta96fl2lUeRxHsgwNQATY4BzVwA+qgATDIwDN4BW/Gk/FivBsf09aCkc/sgj8wPn8AeGiVrw==</latexit><latexit sha1_base64="SKD4gJPomsW3b5dc17IO6xydT8A=">AAAB/nicbVDLSsNAFJ34rPUVFVduBotQNyURRTdC0Y0bpYJ9QJOGyXTaDp1MwsxEKEPAX3HjQhG3foc7/8Zpm4W2HrhwOOde7r0nTBiVynG+rYXFpeWV1cJacX1jc2vb3tltyDgVmNRxzGLRCpEkjHJSV1Qx0koEQVHISDMcXo/95iMRksb8QY0S4keoz2mPYqSMFNj7nr7t6DI7zgLuZYHml27WuQvsklNxJoDzxM1JCeSoBfaX141xGhGuMENStl0nUb5GQlHMSFb0UkkShIeoT9qGchQR6evJ+Rk8MkoX9mJhiis4UX9PaBRJOYpC0xkhNZCz3lj8z2unqnfha8qTVBGOp4t6KYMqhuMsYJcKghUbGYKwoOZWiAdIIKxMYkUTgjv78jxpnFTcs4pzf1qqXuVxFMABOARl4IJzUAU3oAbqAAMNnsEreLOerBfr3fqYti5Y+cwe+APr8weVFpU9</latexit>(l 1) N
n
n=1

(l) N
n n=1

T

<latexit sha1_base64="Eod4uIUlAlqYQ0U4YYDXVLaln6g=">AAACTHicjVBNSwMxFMxWrVq/qh69BIvQHiy7ouhFEL2IUKhgtdCPJZtm29BsdkneimXZH+jFgzd/hRcPigimH4JaDw48GGbekLzxIsE12PaTlZmZncvOLyzmlpZXVtfy6xvXOowVZTUailDVPaKZ4JLVgINg9UgxEniC3Xj9s6F/c8uU5qG8gkHEWgHpSu5zSsBIbp5WXNlOimLXKaX4GDc93i1euMndWC2lKfbdpAk9BsQV6dAuta9w5cv+V8DNF+yyPQKeJs6EFNAEVTf/2OyENA6YBCqI1g3HjqCVEAWcCpbmmrFmEaF90mUNQyUJmG4lozJSvGOUDvZDZUYCHqnfEwkJtB4EntkMCPT0b28o/uU1YvCPWgmXUQxM0vFDfiwwhHjYLO5wxSiIgSGEKm7+immPKELB9J8zJTi/T54m13tl56BsX+4XTk4ndSygLbSNishBh+gEnaMqqiGK7tEzekVv1oP1Yr1bH+PVjDXJbKIfyGQ/Ac4csrU=</latexit>

(l 1)
n

(l)
xn

✓l

(l)
n

(l)
xn

✓l

<latexit sha1_base64="mS6wHAYMLBqUoyYIapvqWgqzaqY=">AAACFXicbVBNS8NAEN34bf2qevQSLEILWhJR9CKIXjxWsLXQxLDZTtqlmw92J2IJ+RNe/CtePCjiVfDmv3Hb5uDXg4HHezPMzPMTwRVa1qcxNT0zOze/sFhaWl5ZXSuvb7RUnEoGTRaLWLZ9qkDwCJrIUUA7kUBDX8C1Pzgf+de3IBWPoyscJuCGtBfxgDOKWvLKu3c3WVXU8pPAyxzsA1IvE3nuCAiwOvb27FruSN7rY80rV6y6NYb5l9gFqZACDa/84XRjloYQIRNUqY5tJehmVCJnAvKSkypIKBvQHnQ0jWgIys3GX+Xmjla6ZhBLXRGaY/X7REZDpYahrztDin312xuJ/3mdFINjN+NRkiJEbLIoSIWJsTmKyOxyCQzFUBPKJNe3mqxPJWWogyzpEOzfL/8lrf26fVi3Lg8qp2dFHAtki2yTKrHJETklF6RBmoSRe/JInsmL8WA8Ga/G26R1yihmNskPGO9fQvKe3A==</latexit><latexit sha1_base64="Z3udPZCLxSYNr4dJmspyP5BL+x0=">AAAB8HicbVBNSwMxEJ2tX7V+VT16CRahHiy7ouix6MVjBfsh7VqyabYNTbJLkhXL0l/hxYMiXv053vw3pu0etPXBwOO9GWbmBTFn2rjut5NbWl5ZXcuvFzY2t7Z3irt7DR0litA6iXikWgHWlDNJ64YZTluxolgEnDaD4fXEbz5SpVkk78wopr7AfclCRrCx0v3TQ1rmJ97xuFssuRV3CrRIvIyUIEOtW/zq9CKSCCoN4VjrtufGxk+xMoxwOi50Ek1jTIa4T9uWSiyo9tPpwWN0ZJUeCiNlSxo0VX9PpFhoPRKB7RTYDPS8NxH/89qJCS/9lMk4MVSS2aIw4chEaPI96jFFieEjSzBRzN6KyAArTIzNqGBD8OZfXiSN04p3XnFvz0rVqyyOPBzAIZTBgwuowg3UoA4EBDzDK7w5ynlx3p2PWWvOyWb24Q+czx/PMI/B</latexit><latexit sha1_base64="WhqpjvzZkJRyO5oivUZFe7Ewafg=">AAAB7nicbVDLSgNBEOyNrxhfUY9eBoMQL2FXFD0GvXiMYB6QxDA76U2GzM4uM7NiWPIRXjwo4tXv8ebfOEn2oIkFDUVVN91dfiy4Nq777eRWVtfWN/Kbha3tnd294v5BQ0eJYlhnkYhUy6caBZdYN9wIbMUKaegLbPqjm6nffESleSTvzTjGbkgHkgecUWOl5tNDWhank16x5FbcGcgy8TJSggy1XvGr049YEqI0TFCt254bm25KleFM4KTQSTTGlI3oANuWShqi7qazcyfkxCp9EkTKljRkpv6eSGmo9Tj0bWdIzVAvelPxP6+dmOCqm3IZJwYlmy8KEkFMRKa/kz5XyIwYW0KZ4vZWwoZUUWZsQgUbgrf48jJpnFW8i4p7d16qXmdx5OEIjqEMHlxCFW6hBnVgMIJneIU3J3ZenHfnY96ac7KZQ/gD5/MH8hSPTw==</latexit>lT<latexit sha1_base64="Qmel+LIc7saRra7GZEa+jxY8Dg4=">AAACTnichVHLSgMxFM3UV62vqks3wSLUTZkRRZeiGxGEClaFTh0y6Z02mMkMyR2xDP1CN+LOz3DjQhFNH+ATvBA4OQ+SnISpFAZd99EpTExOTc8UZ0tz8wuLS+XllXOTZJpDgycy0ZchMyCFggYKlHCZamBxKOEivD4c6Bc3oI1I1Bn2UmjFrKNEJDhDSwVlOArUVV6Vm33qszTVyS31Q9GpHge5j11AFsg+jb5sBurm1Rk9+cz97w/KFbfmDof+Bt4YVMh46kH5wW8nPItBIZfMmKbnptjKmUbBJfRLfmYgZfyadaBpoWIxmFY+rKNPNyzTplGi7VJIh+zXRM5iY3pxaJ0xw675qQ3Iv7RmhtFeKxcqzRAUHx0UZZJiQgfd0rbQwFH2LGBcC3tXyrtMM472B0q2BO/nk3+D862at1NzT7cr+wfjOopkjayTKvHILtknR6ROGoSTO/JEXsirc+88O2/O+8hacMaZVfJtCsUP/cC06Q==</latexit>

(l)
n

✓l ✓l

(l)
n

✓l ✓l

<latexit sha1_base64="Iy8gGNomBDwO+VwPxBcTpXub/3k=">AAACFnicbVDLSgMxFM3UV62vUZdugkVoFy0zouhGKLoRoVChL2inQyZN29DMgyQjlmG+wo2/4saFIm7FnX9jpp1FbT0QOPece7m5xwkYFdIwfrTMyura+kZ2M7e1vbO7p+8fNIUfckwa2Gc+bztIEEY90pBUMtIOOEGuw0jLGd8kfuuBcEF9ry4nAbFcNPTogGIklWTrpWovKrCSWYyv7uzoMSmKcdyL6jGszgo4p9t63igbU8BlYqYkD1LUbP272/dx6BJPYoaE6JhGIK0IcUkxI3GuGwoSIDxGQ9JR1EMuEVY0PSuGJ0rpw4HP1fMknKrzExFyhZi4jup0kRyJRS8R//M6oRxcWhH1glASD88WDUIGpQ+TjGCfcoIlmyiCMKfqrxCPEEdYqiRzKgRz8eRl0jwtm+dl4/4sX7lO48iCI3AMCsAEF6ACbkENNAAGT+AFvIF37Vl71T60z1lrRktnDsEfaF+/BkaeDg==</latexit><latexit sha1_base64="0nAKAhbmQbxubm/i6x22zm3yG8g=">AAAB7nicbVDLSgNBEOyNrxhfUY9eBoMQL2FXFD0GvXgRIpgHJGuYnXSSIbOzy8ysEJZ8hBcPinj1e7z5N06SPWhiQUNR1U13VxALro3rfju5ldW19Y38ZmFre2d3r7h/0NBRohjWWSQi1QqoRsEl1g03AluxQhoGApvB6GbqN59QaR7JBzOO0Q/pQPI+Z9RYqXn3mJbF6aRbLLkVdwayTLyMlCBDrVv86vQiloQoDRNU67bnxsZPqTKcCZwUOonGmLIRHWDbUklD1H46O3dCTqzSI/1I2ZKGzNTfEykNtR6Hge0MqRnqRW8q/ue1E9O/8lMu48SgZPNF/UQQE5Hp76THFTIjxpZQpri9lbAhVZQZm1DBhuAtvrxMGmcV76Li3p+XqtdZHHk4gmMogweXUIVbqEEdGIzgGV7hzYmdF+fd+Zi35pxs5hD+wPn8Aa/mjyQ=</latexit><latexit sha1_base64="UvW+XnwzoMMH+JINEAhMSUz2ENk=">AAAB8HicbVBNSwMxEJ2tX7V+VT16CRahHiy7ouix6MWLUMF+SLuWbJptQ5PskmSFsvRXePGgiFd/jjf/jWm7B219MPB4b4aZeUHMmTau++3klpZXVtfy64WNza3tneLuXkNHiSK0TiIeqVaANeVM0rphhtNWrCgWAafNYHg98ZtPVGkWyXsziqkvcF+ykBFsrPRw+5iW+Yl3PO4WS27FnQItEi8jJchQ6xa/Or2IJIJKQzjWuu25sfFTrAwjnI4LnUTTGJMh7tO2pRILqv10evAYHVmlh8JI2ZIGTdXfEykWWo9EYDsFNgM9703E/7x2YsJLP2UyTgyVZLYoTDgyEZp8j3pMUWL4yBJMFLO3IjLAChNjMyrYELz5lxdJ47TinVfcu7NS9SqLIw8HcAhl8OACqnADNagDAQHP8ApvjnJenHfnY9aac7KZffgD5/MHjKyPlg==</latexit>

(l)

(l)

<latexit sha1_base64="6dnzjD/soRgoA21+kE0+suKmH3g=">AAACIXicbVDLSgMxFM34rPVVdekmWIS6KTOi2GXRTRGECn1BZzpk0rQNzTxI7ohlmF9x46+4caFId+LPmD4WtvVA4HDOudzc40WCKzDNb2NtfWNzazuzk93d2z84zB0dN1QYS8rqNBShbHlEMcEDVgcOgrUiyYjvCdb0hncTv/nEpOJhUINRxByf9APe45SAltxcqdJJCuIixTaJIhk+43s3sWHAgLiJSNNOUkvxwyyy6Li5vFk0p8CrxJqTPJqj6ubGdjeksc8CoIIo1bbMCJyESOBUsDRrx4pFhA5Jn7U1DYjPlJNML0zxuVa6uBdK/QLAU/XvREJ8pUa+p5M+gYFa9ibif147hl7JSXgQxcACOlvUiwWGEE/qwl0uGQUx0oRQyfVfMR0QSSjoUrO6BGv55FXSuCxa10Xz8Spfvp3XkUGn6AwVkIVuUBlVUBXVEUUv6A19oE/j1Xg3vozxLLpmzGdO0AKMn19SeaRG</latexit>lH_n^{(l)} \approx \big(J_{\theta_l} f_{\theta_l}\big)^T M_n^{(l)} \big(J_{\theta_l} f_{\theta_l}\big)\{M^{(l)}_n\}_{n=1}^Nx_n^{(l)} = f_{\theta_l}\big(x^{(l-1)}_n\big)M_n^{(l-1)} = \big(J_{x_n^{(l)}} f_{\theta_l}\big)^T M_n^{(l)} \big(J_{x_n^{(l)}} f_{\theta_l}\big)H^{(l)} \approx J_{\theta_{l}}^{T} M^{(l)}J_{\theta_{l}}5lAPPROXIMATIONSMEMORYTIME

Block diag.
KFAC
Exact diag.
Approx. diag. (ours)
Mixed diag. (ours)2
O(Rm+ Ws2 )
2
O(Rs + Ws )2O(Rm+ Ws )
O(Rm + Ws )O(Rs2 + Ws2 )O(Rs2 + Ws )O(Rs2 + Ws )O(Rs + Ws )<latexit sha1_base64="xsmfAL7DBjgKyvKNlvetK2OUgWw=">AAACFHicbVDLSgMxFL1TX7W+Rl26CRZBEMpMUXRZdOPCRQX7gE4tmTTThmYeJBmxDPMRbvwVNy4UcevCnX9jph1EWw8ETs65N7n3uBFnUlnWl1FYWFxaXimultbWNza3zO2dpgxjQWiDhDwUbRdLyllAG4opTtuRoNh3OW25o4vMb91RIVkY3KhxRLs+HgTMYwQrLfXMI8cTmCROhIVimN9W0x+O7vUNOT5WQ4J5cpX2zLJVsSZA88TOSRly1Hvmp9MPSezTQBGOpezYVqS6SfY84TQtObGkESYjPKAdTQPsU9lNJkul6EArfeSFQp9AoYn6uyPBvpRj39WV2Yhy1svE/7xOrLyzbsKCKFY0INOPvJgjFaIsIdRnghLFx5pgIpieFZEh1ikpnWNJh2DPrjxPmtWKfVKxro/LtfM8jiLswT4cgg2nUINLqEMDCDzAE7zAq/FoPBtvxvu0tGDkPbvwB8bHN5ygnyk=</latexit>

O(Rm + Ws )

@2
L
@x2

O(Rs + Ws )<latexit sha1_base64="xsmfAL7DBjgKyvKNlvetK2OUgWw=">AAACFHicbVDLSgMxFL1TX7W+Rl26CRZBEMpMUXRZdOPCRQX7gE4tmTTThmYeJBmxDPMRbvwVNy4UcevCnX9jph1EWw8ETs65N7n3uBFnUlnWl1FYWFxaXimultbWNza3zO2dpgxjQWiDhDwUbRdLyllAG4opTtuRoNh3OW25o4vMb91RIVkY3KhxRLs+HgTMYwQrLfXMI8cTmCROhIVimN9W0x+O7vUNOT5WQ4J5cpX2zLJVsSZA88TOSRly1Hvmp9MPSezTQBGOpezYVqS6SfY84TQtObGkESYjPKAdTQPsU9lNJkul6EArfeSFQp9AoYn6uyPBvpRj39WV2Yhy1svE/7xOrLyzbsKCKFY0INOPvJgjFaIsIdRnghLFx5pgIpieFZEh1ikpnWNJh2DPrjxPmtWKfVKxro/LtfM8jiLswT4cgg2nUINLqEMDCDzAE7zAq/FoPBtvxvu0tGDkPbvwB8bHN5ygnyk=</latexit>

Table 1: Memory & time complexityof Hessian approximations. For an LlayerP
network, let Rm =max|x(l) |,
P l=0...L
(l)
2
(l) 2
Rs = Pl |x |, Rs = l |x | , andWs = l |θ l |. Only our approximation
scales linearly with both the output
resolution and parameters.

@2
L
@x2

<latexit sha1_base64="xsmfAL7DBjgKyvKNlvetK2OUgWw=">AAACFHicbVDLSgMxFL1TX7W+Rl26CRZBEMpMUXRZdOPCRQX7gE4tmTTThmYeJBmxDPMRbvwVNy4UcevCnX9jph1EWw8ETs65N7n3uBFnUlnWl1FYWFxaXimultbWNza3zO2dpgxjQWiDhDwUbRdLyllAG4opTtuRoNh3OW25o4vMb91RIVkY3KhxRLs+HgTMYwQrLfXMI8cTmCROhIVimN9W0x+O7vUNOT5WQ4J5cpX2zLJVsSZA88TOSRly1Hvmp9MPSezTQBGOpezYVqS6SfY84TQtObGkESYjPKAdTQPsU9lNJkul6EArfeSFQp9AoYn6uyPBvpRj39WV2Yhy1svE/7xOrLyzbsKCKFY0INOPvJgjFaIsIdRnghLFx5pgIpieFZEh1ikpnWNJh2DPrjxPmtWKfVKxro/LtfM8jiLswT4cgg2nUINLqEMDCDzAE7zAq/FoPBtvxvu0tGDkPbvwB8bHN5ygnyk=</latexit><latexit sha1_base64="xsmfAL7DBjgKyvKNlvetK2OUgWw=">AAACFHicbVDLSgMxFL1TX7W+Rl26CRZBEMpMUXRZdOPCRQX7gE4tmTTThmYeJBmxDPMRbvwVNy4UcevCnX9jph1EWw8ETs65N7n3uBFnUlnWl1FYWFxaXimultbWNza3zO2dpgxjQWiDhDwUbRdLyllAG4opTtuRoNh3OW25o4vMb91RIVkY3KhxRLs+HgTMYwQrLfXMI8cTmCROhIVimN9W0x+O7vUNOT5WQ4J5cpX2zLJVsSZA88TOSRly1Hvmp9MPSezTQBGOpezYVqS6SfY84TQtObGkESYjPKAdTQPsU9lNJkul6EArfeSFQp9AoYn6uyPBvpRj39WV2Yhy1svE/7xOrLyzbsKCKFY0INOPvJgjFaIsIdRnghLFx5pgIpieFZEh1ikpnWNJh2DPrjxPmtWKfVKxro/LtfM8jiLswT4cgg2nUINLqEMDCDzAE7zAq/FoPBtvxvu0tGDkPbvwB8bHN5ygnyk=</latexit>

@2
L
@x2

@2
L
@x2

M 6= 0<latexit sha1_base64="bv3CQJti9XknCUE+obCi6JHEkFc=">AAAB73icbVBNS8NAEJ34WetX1aOXxSJ4Kokoeix68SJUsB/QhrLZTtqlm026uxFK6J/w4kERr/4db/4bt20O2vpg4PHeDDPzgkRwbVz321lZXVvf2CxsFbd3dvf2SweHDR2nimGdxSJWrYBqFFxi3XAjsJUopFEgsBkMb6d+8wmV5rF8NOME/Yj2JQ85o8ZKrXvSkTgibrdUdivuDGSZeDkpQ45at/TV6cUsjVAaJqjWbc9NjJ9RZTgTOCl2Uo0JZUPax7alkkao/Wx274ScWqVHwljZkobM1N8TGY20HkeB7YyoGehFbyr+57VTE177GZdJalCy+aIwFcTEZPo86XGFzIixJZQpbm8lbEAVZcZGVLQheIsvL5PGecW7rLgPF+XqTR5HAY7hBM7Agyuowh3UoA4MBDzDK7w5I+fFeXc+5q0rTj5zBH/gfP4A0dGPKw==</latexit>

H 6= 0<latexit sha1_base64="5FCXk1UPvDoI2sll11rXm4lJxlw=">AAAB73icbVBNS8NAEJ3Ur1q/qh69LBbBU0lE0WPRS48VbC20oWy2k3bpZpPuboQS+ie8eFDEq3/Hm//GbZuDtj4YeLw3w8y8IBFcG9f9dgpr6xubW8Xt0s7u3v5B+fCopeNUMWyyWMSqHVCNgktsGm4EthOFNAoEPgaju5n/+IRK81g+mEmCfkQHkoecUWOldp10JY6J2ytX3Ko7B1klXk4qkKPRK391+zFLI5SGCap1x3MT42dUGc4ETkvdVGNC2YgOsGOppBFqP5vfOyVnVumTMFa2pCFz9fdERiOtJ1FgOyNqhnrZm4n/eZ3UhDd+xmWSGpRssShMBTExmT1P+lwhM2JiCWWK21sJG1JFmbERlWwI3vLLq6R1UfWuqu79ZaV2m8dRhBM4hXPw4BpqUIcGNIGBgGd4hTdn7Lw4787HorXg5DPH8AfO5w/KGo8m</latexit>

JxT M Jx<latexit sha1_base64="t4T06mox+BC37EaVt4laTYmaOeE=">AAAB+nicbZDLSsNAFIZP6q3WW6pLN4NFcFUSUXRZdCOCUKE3aGOYTCft0MmFmYlaYh7FjQtF3Pok7nwbp2kX2vrDwMd/zuGc+b2YM6ks69soLC2vrK4V10sbm1vbO2Z5tyWjRBDaJBGPRMfDknIW0qZiitNOLCgOPE7b3uhyUm/fUyFZFDbUOKZOgAch8xnBSluuWb5208fsLm1k6Abl7JoVq2rlQotgz6ACM9Vd86vXj0gS0FARjqXs2lasnBQLxQinWamXSBpjMsID2tUY4oBKJ81Pz9ChdvrIj4R+oUK5+3sixYGU48DTnQFWQzlfm5j/1bqJ8s+dlIVxomhIpov8hCMVoUkOqM8EJYqPNWAimL4VkSEWmCidVkmHYM9/eRFax1X7tGrdnlRqF7M4irAPB3AENpxBDa6gDk0g8ADP8ApvxpPxYrwbH9PWgjGb2YM/Mj5/AOQ6k8I=</latexit>

J✓T M J✓
<latexit sha1_base64="ZA7xr5Y6H1DtSiPBeSWhXg4LAcI=">AAACBHicbVDLSsNAFJ3UV62vqMtuBovgqiSi6LLoRgShQl/QxjCZTtqhkwczN0IJWbjxV9y4UMStH+HOv3HaBtTWAxcO59zLvfd4seAKLOvLKCwtr6yuFddLG5tb2zvm7l5LRYmkrEkjEcmORxQTPGRN4CBYJ5aMBJ5gbW90OfHb90wqHoUNGMfMCcgg5D6nBLTkmuVrN+3BkAHJ7tJGhm/wj+CaFatqTYEXiZ2TCspRd83PXj+iScBCoIIo1bWtGJyUSOBUsKzUSxSLCR2RAetqGpKAKSedPpHhQ630sR9JXSHgqfp7IiWBUuPA050BgaGa9ybif143Af/cSXkYJ8BCOlvkJwJDhCeJ4D6XjIIYa0Ko5PpWTIdEEgo6t5IOwZ5/eZG0jqv2adW6PanULvI4iqiMDtARstEZqqErVEdNRNEDekIv6NV4NJ6NN+N91low8pl99AfGxzeo5Jga</latexit><latexit sha1_base64="MCD+8sCYp3xlwD8YT9yxbh35sZs=">AAACAnicbVDLSgMxFM3UV62vUVfiJliEuikzouiyKILLCvYB7VAyadqGZjJDckdahuLGX3HjQhG3foU7/8ZMOwttPRA4nHMvN+f4keAaHOfbyi0tr6yu5dcLG5tb2zv27l5dh7GirEZDEaqmTzQTXLIacBCsGSlGAl+whj+8Tv3GA1Oah/IexhHzAtKXvMcpASN17IM2sBEkJf8E34wIBdzlpB9KIiYdu+iUnSnwInEzUkQZqh37q90NaRwwCVQQrVuuE4GXEAWcCjYptGPNIkKHpM9ahkoSMO0l0wgTfGyULu6FyjwJeKr+3khIoPU48M1kQGCg571U/M9rxdC79BIuoxiYpLNDvVhgCHHah8mrGAUxToNTxc1fMR0QZaowrRVMCe585EVSPy2752Xn7qxYucrqyKNDdIRKyEUXqIJuURXVEEWP6Bm9ojfryXqx3q2P2WjOynb20R9Ynz+UgJbp</latexit><latexit sha1_base64="4ncGXueUJ4wXkLTosyH6uA6yrv0=">AAACAnicbVDLSgMxFM34rPVVdSVugkWomzIjii5L3bisYB/QDiWTZtrQTDIkd8QyFDf+ihsXirj1K9z5N2baLrT1QOBwzr3cnBPEghtw3W9naXlldW09t5Hf3Nre2S3s7TeMSjRldaqE0q2AGCa4ZHXgIFgr1oxEgWDNYHid+c17pg1X8g5GMfMj0pc85JSAlbqFww6wB0hL5BRXhaJD3OOkryQR426h6JbdCfAi8WakiGaodQtfnZ6iScQkUEGMaXtuDH5KNHAq2DjfSQyLCR2SPmtbKknEjJ9OIozxiVV6OFTaPgl4ov7eSElkzCgK7GREYGDmvUz8z2snEF75KZdxAkzS6aEwERgUzvqweTWjIEZZcKq5/SumA6IJBdta3pbgzUdeJI2zsndRdm/Pi5XqrI4cOkLHqIQ8dIkq6AbVUB1R9Iie0St6c56cF+fd+ZiOLjmznQP0B87nD4NLlt4=</latexit>(a) Block diagonal(b) Exact diagonal<latexit sha1_base64="akSGlhteizpPasy2oZ3Q7mJmtWc=">AAACAnicbVC7SgNBFJ31GeMraiU2g0HQJuyKomXQxkaIYB6QLGF29m4cnJ1dZu5KwhJs/BUbC0Vs/Qo7/8bJo/B1YOBwzr3cOSdIpTDoup/OzOzc/MJiYam4vLK6tl7a2GyYJNMc6jyRiW4FzIAUCuooUEIr1cDiQEIzuD0f+c070EYk6hoHKfgx6ykRCc7QSt3Sdgehj/l+eEAvRR9CGgrWSxSTw26p7FbcMehf4k1JmUxR65Y+OmHCsxgUcsmMaXtuin7ONAouYVjsZAZSxm9ZD9qWKhaD8fNxhCHds0pIo0Tbp5CO1e8bOYuNGcSBnYwZ3pjf3kj8z2tnGJ36uVBphqD45FCUSYoJHfVh82rgKAej4FwL+1fKb5hmHG1rRVuC9zvyX9I4rHjHFffqqFw9m9ZRIDtkl+wTj5yQKrkgNVInnNyTR/JMXpwH58l5dd4mozPOdGeL/IDz/gWa6Zbt</latexit><latexit sha1_base64="vRx9GABdnUUeOpe2Wp8RfSbxudA=">AAACBHicbVDLTgIxFO3gC/GFumTTSExwQ2aMRpeoG5eYyCOBCemUAg2ddtLeMZAJCzf+ihsXGuPWj3Dn31hgFgqepMnJOffm9pwgEtyA6347mZXVtfWN7GZua3tndy+/f1A3KtaU1agSSjcDYpjgktWAg2DNSDMSBoI1guHN1G88MG24kvcwjpgfkr7kPU4JWKmTL7SBjSAp0RN8FUVajcq4y0lfSSImnXzRLbsz4GXipaSIUlQ7+a92V9E4ZBKoIMa0PDcCPyEaOBVskmvHhkWEDkmftSyVJGTGT2YhJvjYKl3cU9o+CXim/t5ISGjMOAzsZEhgYBa9qfif14qhd+knXEYxMEnnh3qxwKDwtBGbVzMKYjwNTjW3f8V0QDShYHvL2RK8xcjLpH5a9s7L7t1ZsXKd1pFFBXSESshDF6iCblEV1RBFj+gZvaI358l5cd6dj/loxkl3DtEfOJ8/AkaXsQ==</latexit>(c) Approx. diagonal(d) Mixed diagonal

Figure 6: Comparison of Hessian approximation methods. Common approximations (a–b) scalequadratically with the output resolution. Our proposed approximate and mixed diagonal Hessians(c–d) scale linearly with the resolution. This is essential for scaling the LAE to large images.
has linear complexity in both time and memory with respect to the output dimension (Tab. 1). This
makes it viable for our model.
We can further tailor this approximation to the autoencoder setting by leveraging the bottleneckarchitecture. We note that the quadratic scaling of the exact diagonal Hessian is less of an issue in thelayers near the bottleneck than in the layers closer to the output space. We can therefore dynamicallyswitch between our approximate diagonal and the exact one, depending on the feature dimension.
This lessens the approximation error while remaining tractable in practice. We provide more details
on the fast hessian computations in Appendix E.

4Related Work

Deep generative models, and particularly the family of variational auto-encoders (VAEs) (Kingma
and Welling, 2014; Rezende et al., 2014), address unsupervised representation learning from aprobabilistic viewpoint by approximating the posterior over the representation space. Despite theirwidespread adoption, model parameters are still deterministic and sensible to ill-suited local minima,e.g. over-fitted to training data (Zhang et al., 2021), which may cause poor generalization. The
Bayesian NNs favour inference over the NN weights for addressing such issues (MacKay, 1995;
Neal, 1996). This approach deduce distributions on data space by learning posterior distributions on
the parameter space. However, several shortcomings (Wenzel et al., 2020), e.g. expensive training,tuning and implementations, often limit their applicability to autoencoder-style models. Alternatively,
other methods such as deep ensembles (Lakshminarayanan et al., 2017), stochastic weight averaging(SWAG) (Maddox et al., 2019) or Monte-Carlo dropout (Gal and Ghahramani, 2016) also promise
Bayesian approximations to NN weight’s posterior, but at the cost of increased training time, poor
empirical performance or limited Bayesian interpretation.
As demonstrated by Daxberger et al. (2021) LA is a scalable and well-behaved alternative to Bayesian
NN s if used post-hoc to approximate the intractable posterior over the weights after maximum-aposteriori (MAP) training in classification and regression. The general utility of LA has also motivatedits use as an approximation to the marginal likelihood over NN weights. Recent methods, including
Daxberger et al. (2021), have explored this path to find model hyperparameters (Immer et al., 2021) or
learning invariances (Immer et al., 2022). However, its computational burden, for instance in Hessian
matrices, has prescribed diagonal or Kronecker factored approximations (Ritter et al., 2018; Martens
and Grosse, 2015a; Botev et al., 2017), which are now widely used for second-order optimization.
We provide more details on the connections to existing hessian based methods Daxberger et al.
6

8

50

7

Memory [Gb]Time [sec]

60

40
30
20
10

Approx
Exact
KFAC
Block

6
5
4
3
2
1

0
0

10

20

30

40

50

Input size

60

0

10

20

30

40

Input size

50

60

Figure 7: Memory & time usage of Hessian approximations. The exact and KFAC
scales poorly with the image
resolution. In contrast, ourproposed approximate diagonal Hessian scales linearly.

(2021); Zhang et al. (2017); Khan et al. (2017), Bayes by Backpropagation Blundell et al. (2015) andAdam Kingma and Ba (2015a) in Appendix C.
The full Bayesian perspective on VAE weights was first explored by Daxberger and Hernández-Lobato(2019) which we find similar in spirit to our work. In contrast to them (1) we follow a principled
Bayesian derivation. (2) Neither do we depend on Hamiltonian Monte-Carlo sampling, which is
generally hard to scale to efficient training.

5

Experiments

First we demonstrate the computational advantages of the proposed Hessian approximation, andthen that our sampling-based training leads to well calibrated uncertainties that can be used for OOD
detection, data imputation and semi-supervised learning. For all the downstream tasks we consider thefollowing baselines: AE (Hinton and Salakhutdinov, 2006) with constant and learned variance, VAE(Rezende et al., 2014; Kingma and Welling, 2014), Monte-Carlo dropout AE (Gal and Ghahramani,2016) and Ensembles of AE (Lakshminarayanan et al., 2017). We extend StochManDetlefsen et al.
(2021) with the Hessian backpropagation for the approximate and mixed diagonals. The training codeis implemented in PyTorch and available2 . Appendix A provides more details on the experimental
setup.
Efficient Hessian Approximation. For practical applications, training time and memory usage of the
Hessian approximation must be kept low. We here show that the proposed approximate diagonal Hessian is sufficient and even outperforms other approximations when combined with our online training.
Fig. 7 show the time and memory requirement for different approximation methods as a functionof input size for a 5 layer convolutional network that preserve channel and input dimension. Asbaselines we use efficient implementations of the exact and KFAC approximation (Daxberger et al.,2021; Dangel et al., 2020). The exact diagonal approximation run out of memory for an ∼ 36 × 36 × 3
image on a 11 Gb NVIDIA GeForce GTX 1080 Ti. In contrast, our approximate diagonal Hessian
scales linearly with the resolution, which is especially beneficial for convolutional layers.

µµ<latexit sha1_base64="8c3ZV3H1YNaQu9q/pgnifAdGrLc=">AAAB6nicbVDLSgMxFM34rPVVdalIsAiuhplidborunHZon1AO5RMmmlDk8yQZIQydOnSjQtF3PoR/Q53foM/YdqKqPXAhcM593LvPUHMqNKO824tLC4tr6xm1rLrG5tb27md3bqKEolJDUcsks0AKcKoIDVNNSPNWBLEA0YaweBy4jduiVQ0Ejd6GBOfo56gIcVIG+m6zZNOLu/YpdKpV/SgId6Z5xagaztTfJN8+WBc/bg7HFc6ubd2N8IJJ0JjhpRquU6s/RRJTTEjo2w7USRGeIB6pGWoQJwoP52eOoLHRunCMJKmhIZT9edEirhSQx6YTo50X/31JuJ/XivRoeenVMSJJgLPFoUJgzqCk79hl0qCNRsagrCk5laI+0girE06WRPC3MvzpF6w3aLtVE0aF2CGDNgHR+AEuOAclMEVqIAawKAH7sEjeLKY9WA9Wy+z1gXra2YP/IL1+gm+MpHP</latexit>µ<latexit sha1_base64="8c3ZV3H1YNaQu9q/pgnifAdGrLc=">AAAB6nicbVDLSgMxFM34rPVVdalIsAiuhplidborunHZon1AO5RMmmlDk8yQZIQydOnSjQtF3PoR/Q53foM/YdqKqPXAhcM593LvPUHMqNKO824tLC4tr6xm1rLrG5tb27md3bqKEolJDUcsks0AKcKoIDVNNSPNWBLEA0YaweBy4jduiVQ0Ejd6GBOfo56gIcVIG+m6zZNOLu/YpdKpV/SgId6Z5xagaztTfJN8+WBc/bg7HFc6ubd2N8IJJ0JjhpRquU6s/RRJTTEjo2w7USRGeIB6pGWoQJwoP52eOoLHRunCMJKmhIZT9edEirhSQx6YTo50X/31JuJ/XivRoeenVMSJJgLPFoUJgzqCk79hl0qCNRsagrCk5laI+0girE06WRPC3MvzpF6w3aLtVE0aF2CGDNgHR+AEuOAclMEVqIAawKAH7sEjeLKY9WA9Wy+z1gXra2YP/IL1+gm+MpHP</latexit>µ<latexit sha1_base64="8c3ZV3H1YNaQu9q/pgnifAdGrLc=">AAAB6nicbVDLSgMxFM34rPVVdalIsAiuhplidborunHZon1AO5RMmmlDk8yQZIQydOnSjQtF3PoR/Q53foM/YdqKqPXAhcM593LvPUHMqNKO824tLC4tr6xm1rLrG5tb27md3bqKEolJDUcsks0AKcKoIDVNNSPNWBLEA0YaweBy4jduiVQ0Ejd6GBOfo56gIcVIG+m6zZNOLu/YpdKpV/SgId6Z5xagaztTfJN8+WBc/bg7HFc6ubd2N8IJJ0JjhpRquU6s/RRJTTEjo2w7USRGeIB6pGWoQJwoP52eOoLHRunCMJKmhIZT9edEirhSQx6YTo50X/31JuJ/XivRoeenVMSJJgLPFoUJgzqCk79hl0qCNRsagrCk5laI+0girE06WRPC3MvzpF6w3aLtVE0aF2CGDNgHR+AEuOAclMEVqIAawKAH7sEjeLKY9WA9Wy+z1gXra2YP/IL1+gm+MpHP</latexit>µ<latexit sha1_base64="8c3ZV3H1YNaQu9q/pgnifAdGrLc=">AAAB6nicbVDLSgMxFM34rPVVdalIsAiuhplidborunHZon1AO5RMmmlDk8yQZIQydOnSjQtF3PoR/Q53foM/YdqKqPXAhcM593LvPUHMqNKO824tLC4tr6xm1rLrG5tb27md3bqKEolJDUcsks0AKcKoIDVNNSPNWBLEA0YaweBy4jduiVQ0Ejd6GBOfo56gIcVIG+m6zZNOLu/YpdKpV/SgId6Z5xagaztTfJN8+WBc/bg7HFc6ubd2N8IJJ0JjhpRquU6s/RRJTTEjo2w7USRGeIB6pGWoQJwoP52eOoLHRunCMJKmhIZT9edEirhSQx6YTo50X/31JuJ/XivRoeenVMSJJgLPFoUJgzqCk79hl0qCNRsagrCk5laI+0girE06WRPC3MvzpF6w3aLtVE0aF2CGDNgHR+AEuOAclMEVqIAawKAH7sEjeLKY9WA9Wy+z1gXra2YP/IL1+gm+MpHP</latexit><latexit sha1_base64="8c3ZV3H1YNaQu9q/pgnifAdGrLc=">AAAB6nicbVDLSgMxFM34rPVVdalIsAiuhplidborunHZon1AO5RMmmlDk8yQZIQydOnSjQtF3PoR/Q53foM/YdqKqPXAhcM593LvPUHMqNKO824tLC4tr6xm1rLrG5tb27md3bqKEolJDUcsks0AKcKoIDVNNSPNWBLEA0YaweBy4jduiVQ0Ejd6GBOfo56gIcVIG+m6zZNOLu/YpdKpV/SgId6Z5xagaztTfJN8+WBc/bg7HFc6ubd2N8IJJ0JjhpRquU6s/RRJTTEjo2w7USRGeIB6pGWoQJwoP52eOoLHRunCMJKmhIZT9edEirhSQx6YTo50X/31JuJ/XivRoeenVMSJJgLPFoUJgzqCk79hl0qCNRsagrCk5laI+0girE06WRPC3MvzpF6w3aLtVE0aF2CGDNgHR+AEuOAclMEVqIAawKAH7sEjeLKY9WA9Wy+z1gXra2YP/IL1+gm+MpHP</latexit>

HessianKFAC

2<latexit sha1_base64="e9x4FYT0OxfxQlbw8BOwWjo4fYM=">AAAB73icbZDLSgMxFIYz9VbHW9Wlm2ARXA0zxep0IRbduKxgL9COJZOmbWiSGZOMUIa+hBsXigiufBL3bsS3MW1F1PpD4OP/zyHnnDBmVGnX/bAyc/MLi0vZZXtldW19I7e5VVNRIjGp4ohFshEiRRgVpKqpZqQRS4J4yEg9HJyN8/oNkYpG4lIPYxJw1BO0SzHSxmq0FO1xdFVo5/KuUyod+EUfGvAPfa8APced6BvyJ6/2cfz8blfaubdWJ8IJJ0JjhpRqem6sgxRJTTEjI7uVKBIjPEA90jQoECcqSCfzjuCecTqwG0nzhIYT92dHirhSQx6aSo50X/3NxuZ/WTPRXT9IqYgTTQSeftRNGNQRHC8PO1QSrNnQAMKSmlkh7iOJsDYnss0RZlaehVrB8YqOe+Hmy6dgqizYAbtgH3jgCJTBOaiAKsCAgVtwDx6sa+vOerSepqUZ66tnG/yS9fIJmS+TVg==</latexit>

2<latexit sha1_base64="e9x4FYT0OxfxQlbw8BOwWjo4fYM=">AAAB73icbZDLSgMxFIYz9VbHW9Wlm2ARXA0zxep0IRbduKxgL9COJZOmbWiSGZOMUIa+hBsXigiufBL3bsS3MW1F1PpD4OP/zyHnnDBmVGnX/bAyc/MLi0vZZXtldW19I7e5VVNRIjGp4ohFshEiRRgVpKqpZqQRS4J4yEg9HJyN8/oNkYpG4lIPYxJw1BO0SzHSxmq0FO1xdFVo5/KuUyod+EUfGvAPfa8APced6BvyJ6/2cfz8blfaubdWJ8IJJ0JjhpRqem6sgxRJTTEjI7uVKBIjPEA90jQoECcqSCfzjuCecTqwG0nzhIYT92dHirhSQx6aSo50X/3NxuZ/WTPRXT9IqYgTTQSeftRNGNQRHC8PO1QSrNnQAMKSmlkh7iOJsDYnss0RZlaehVrB8YqOe+Hmy6dgqizYAbtgH3jgCJTBOaiAKsCAgVtwDx6sa+vOerSepqUZ66tnG/yS9fIJmS+TVg==</latexit>

Approx2<latexit sha1_base64="e9x4FYT0OxfxQlbw8BOwWjo4fYM=">AAAB73icbZDLSgMxFIYz9VbHW9Wlm2ARXA0zxep0IRbduKxgL9COJZOmbWiSGZOMUIa+hBsXigiufBL3bsS3MW1F1PpD4OP/zyHnnDBmVGnX/bAyc/MLi0vZZXtldW19I7e5VVNRIjGp4ohFshEiRRgVpKqpZqQRS4J4yEg9HJyN8/oNkYpG4lIPYxJw1BO0SzHSxmq0FO1xdFVo5/KuUyod+EUfGvAPfa8APced6BvyJ6/2cfz8blfaubdWJ8IJJ0JjhpRqem6sgxRJTTEjI7uVKBIjPEA90jQoECcqSCfzjuCecTqwG0nzhIYT92dHirhSQx6aSo50X/3NxuZ/WTPRXT9IqYgTTQSeftRNGNQRHC8PO1QSrNnQAMKSmlkh7iOJsDYnss0RZlaehVrB8YqOe+Hmy6dgqizYAbtgH3jgCJTBOaiAKsCAgVtwDx6sa+vOerSepqUZ66tnG/yS9fIJmS+TVg==</latexit><latexit sha1_base64="sdul4ZrV50cL6djVwP2wacNhUbo=">AAAB7XicbVDLSgMxFM3UV62vqks3wSK4GmaK1emu6sZlBfuAdiiZNNPGZpKQZMQy9B/cuFDErf/jzr8xfSC+Dlw4nHMv994TSUa18bwPJ7e0vLK6ll8vbGxube8Ud/eaWqQKkwYWTKh2hDRhlJOGoYaRtlQEJREjrWh0OfVbd0RpKviNGUsSJmjAaUwxMlZqnkupxH2vWPLcavUkqATQkuA08MvQd70ZvkgJLFDvFd+7fYHThHCDGdK643vShBlShmJGJoVuqolEeIQGpGMpRwnRYTa7dgKPrNKHsVC2uIEz9ftEhhKtx0lkOxNkhvq3NxX/8zqpiYMwo1ymhnA8XxSnDBoBp6/DPlUEGza2BGFF7a0QD5FC2NiACjaEPy//Jc2y61dc77pcql0s4siDA3AIjoEPzkANXIE6aAAMbsEDeALPjnAenRfndd6acxYz++AHnLdPIvGPgA==</latexit>

2<latexit sha1_base64="e9x4FYT0OxfxQlbw8BOwWjo4fYM=">AAAB73icbZDLSgMxFIYz9VbHW9Wlm2ARXA0zxep0IRbduKxgL9COJZOmbWiSGZOMUIa+hBsXigiufBL3bsS3MW1F1PpD4OP/zyHnnDBmVGnX/bAyc/MLi0vZZXtldW19I7e5VVNRIjGp4ohFshEiRRgVpKqpZqQRS4J4yEg9HJyN8/oNkYpG4lIPYxJw1BO0SzHSxmq0FO1xdFVo5/KuUyod+EUfGvAPfa8APced6BvyJ6/2cfz8blfaubdWJ8IJJ0JjhpRqem6sgxRJTTEjI7uVKBIjPEA90jQoECcqSCfzjuCecTqwG0nzhIYT92dHirhSQx6aSo50X/3NxuZ/WTPRXT9IqYgTTQSeftRNGNQRHC8PO1QSrNnQAMKSmlkh7iOJsDYnss0RZlaehVrB8YqOe+Hmy6dgqizYAbtgH3jgCJTBOaiAKsCAgVtwDx6sa+vOerSepqUZ66tnG/yS9fIJmS+TVg==</latexit>

Approx⇤<latexit sha1_base64="e2aHVoRIfdni2tXjeFYpZch0f84=">AAAB8XicbVDLSgMxFM34rPVVdekm2AriYpgpVqe7qhuXFewD27Fk0kwbmkmGJCOWoX/hxoUibv0bd/6N6QNR64ELh3Pu5d57gphRpR3n01pYXFpeWc2sZdc3Nre2czu7dSUSiUkNCyZkM0CKMMpJTVPNSDOWBEUBI41gcDn2G/dEKir4jR7GxI9Qj9OQYqSNdHsex1I8FO6OC51c3rHL5ROv5EFDvFPPLULXdib4JnkwQ7WT+2h3BU4iwjVmSKmW68TaT5HUFDMyyrYTRWKEB6hHWoZyFBHlp5OLR/DQKF0YCmmKazhRf06kKFJqGAWmM0K6r/56Y/E/r5Xo0PNTyuNEE46ni8KEQS3g+H3YpZJgzYaGICypuRXiPpIIaxNS1oQw9/I8qRdtt2Q718V85WIWRwbsgwNwBFxwBirgClRBDWDAwSN4Bi+Wsp6sV+tt2rpgzWb2wC9Y71/9PpB4</latexit>

Exact<latexit sha1_base64="aIYcy8blz/H7qro6+83FyfoKIE0=">AAAB7HicbVBNS8NAEJ3Ur1q/qh69LBbBU0iK1fRWFMFjBdMW2lA22027dLMJuxuxlP4GLx4U8eoP8ua/cfuBqPXBwOO9GWbmhSlnSjvOp5VbWV1b38hvFra2d3b3ivsHDZVkklCfJDyRrRArypmgvmaa01YqKY5DTpvh8GrqN++pVCwRd3qU0iDGfcEiRrA2kn/9gInuFkuOXa2eeRUPGeKde24ZubYzwzcpwQL1bvGj00tIFlOhCcdKtV0n1cEYS80Ip5NCJ1M0xWSI+7RtqMAxVcF4duwEnRilh6JEmhIazdSfE2McKzWKQ9MZYz1Qf72p+J/XznTkBWMm0kxTQeaLoowjnaDp56jHJCWajwzBRDJzKyIDLE0AJp+CCWHp5WXSKNtuxXZuy6Xa5SKOPBzBMZyCCxdQgxuogw8EGDzCM7xYwnqyXq23eWvOWswcwi9Y718yW47x</latexit><latexit sha1_base64="rcN6X1H8M30hyZs9xCuSNWF3O94=">AAAB83icbVBNS8NAEJ34WetX1aOXxSLUgyUpiB6LXjxWsB/QhrLZbtqlm2zYnQgl9G948aCIV/+MN/+N2zYHbX0w8Hhvhpl5QSKFQdf9dtbWNza3tgs7xd29/YPD0tFxy6hUM95kSirdCajhUsS8iQIl7ySa0yiQvB2M72Z++4lrI1T8iJOE+xEdxiIUjKKVehV6QRrK4OVIsX6p7FbdOcgq8XJShhyNfumrN1AsjXiMTFJjup6boJ9RjYJJPi32UsMTysZ0yLuWxjTixs/mN0/JuVUGJFTaVoxkrv6eyGhkzCQKbGdEcWSWvZn4n9dNMbzxMxEnKfKYLRaFqSSoyCwAMhCaM5QTSyjTwt5K2IhqytDGVLQheMsvr5JWrepdVd2HWrl+m8dRgFM4gwp4cA11uIcGNIFBAs/wCm9O6rw4787HonXNyWdO4A+czx+4PJDT</latexit>(a) Post-hoc

KFAC<latexit sha1_base64="tJFoW/sxXduvk2+Zt9mBLbEWjmE=">AAAB63icbVDLSgMxFM3UV62vqks3wSK4GmaK1emuWhDBTQX7gHYomTTThiaZIckIZegvuHGhiFt/yJ1/Y6YtotYDFw7n3Mu99wQxo0o7zqeVW1ldW9/Ibxa2tnd294r7By0VJRKTJo5YJDsBUoRRQZqaakY6sSSIB4y0g3E989sPRCoaiXs9iYnP0VDQkGKkM+n2+rLeL5Ycu1o98yoeNMQ799wydG1nhm9SAgs0+sWP3iDCCSdCY4aU6rpOrP0USU0xI9NCL1EkRniMhqRrqECcKD+d3TqFJ0YZwDCSpoSGM/XnRIq4UhMemE6O9Ej99TLxP6+b6NDzUyriRBOB54vChEEdwexxOKCSYM0mhiAsqbkV4hGSCGsTT8GEsPTyMmmVbbdiO3flUu1qEUceHIFjcApccAFq4AY0QBNgMAKP4Bm8WNx6sl6tt3lrzlrMHIJfsN6/ALfgjgc=</latexit>

2<latexit sha1_base64="e9x4FYT0OxfxQlbw8BOwWjo4fYM=">AAAB73icbZDLSgMxFIYz9VbHW9Wlm2ARXA0zxep0IRbduKxgL9COJZOmbWiSGZOMUIa+hBsXigiufBL3bsS3MW1F1PpD4OP/zyHnnDBmVGnX/bAyc/MLi0vZZXtldW19I7e5VVNRIjGp4ohFshEiRRgVpKqpZqQRS4J4yEg9HJyN8/oNkYpG4lIPYxJw1BO0SzHSxmq0FO1xdFVo5/KuUyod+EUfGvAPfa8APced6BvyJ6/2cfz8blfaubdWJ8IJJ0JjhpRqem6sgxRJTTEjI7uVKBIjPEA90jQoECcqSCfzjuCecTqwG0nzhIYT92dHirhSQx6aSo50X/3NxuZ/WTPRXT9IqYgTTQSeftRNGNQRHC8PO1QSrNnQAMKSmlkh7iOJsDYnss0RZlaehVrB8YqOe+Hmy6dgqizYAbtgH3jgCJTBOaiAKsCAgVtwDx6sa+vOerSepqUZ66tnG/yS9fIJmS+TVg==</latexit>

Exact⇤<latexit sha1_base64="UOgbZP9t1/OuMiJKnjxByknYAgM=">AAAB8HicbVDLSgMxFM34rPVVdekm2AriYpgpVqe7ogguK9iHtGPJpJk2NMkMSUYsQ7/CjQtF3Po57vwb0wei1gMXDufcy733BDGjSjvOp7WwuLS8sppZy65vbG5t53Z26ypKJCY1HLFINgOkCKOC1DTVjDRjSRAPGGkEg4ux37gnUtFI3OhhTHyOeoKGFCNtpNvLB4R14e640MnlHbtcPvFKHjTEO/XcInRtZ4JvkgczVDu5j3Y3wgknQmOGlGq5Tqz9FElNMSOjbDtRJEZ4gHqkZahAnCg/nRw8godG6cIwkqaEhhP150SKuFJDHphOjnRf/fXG4n9eK9Gh56dUxIkmAk8XhQmDOoLj72GXSoI1GxqCsKTmVoj7SJoQTEZZE8Lcy/OkXrTdku1cF/OV81kcGbAPDsARcMEZqIArUAU1gAEHj+AZvFjSerJerbdp64I1m9kDv2C9fwEKm4/p</latexit><latexit sha1_base64="pcvjhTHxVMgKf/ih+o01HcKy954=">AAAB8XicbVDLSgMxFL1TX7W+qi7dBItQN2WmILosunFnBfvAdiiZNNOGZpIhyQhl6F+4caGIW//GnX9jpp2Fth4IHM65h9x7gpgzbVz32ymsrW9sbhW3Szu7e/sH5cOjtpaJIrRFJJeqG2BNORO0ZZjhtBsriqOA004wucn8zhNVmknxYKYx9SM8EixkBBsrPVaDc3QnsvCgXHFr7hxolXg5qUCO5qD81R9KkkRUGMKx1j3PjY2fYmUY4XRW6ieaxphM8Ij2LBU4otpP5xvP0JlVhiiUyj5h0Fz9nUhxpPU0CuxkhM1YL3uZ+J/XS0x45adMxImhgiw+ChOOjETZ+WjIFCWGTy3BRDG7KyJjrDAxtqSSLcFbPnmVtOs176Lm3tcrjeu8jiKcwClUwYNLaMAtNKEFBAQ8wyu8Odp5cd6dj8Vowckzx/AHzucPaq+QGA==</latexit>(b) Online

Exact
ApproxExact∗
Approx∗log p(x) ↑-10065.02-299.95-242.24-26.10-26.11MSE ↓125.58
27.21
26.63
26.06
26.08

Table 2: Online training (indicatedby ∗ ) outperforms post-hoc LA. The
approximate diagonal always has
similar performance with the exact
diagonal.

Figure 8: Mean and variance estimates of 100 sampled NN.
Tab. 2 shows that the exact or approximate Hessian diagonal gives similar performance. Using
post-hoc LA results in good mean reconstructions (low MSE), but each sampled NN does not give2 https://github.com/FrederikWarburg/LaplaceAE

7

improves OoD detection. Furthermore, the online training improves the model calibration.

229

Missing Data imputation. Another application of stochastic representation learning is to provide
distributions over unobserved values (Rezende et al., 2014). In many application domains, sensor
readings go missing, which we may mimic by letting parts of an image be unobserved. Rezende et al.
(2014) show that we can then draw samples from the distribution of the entire image conditioned on

230
231
232
233

Methodlogp(x)Typicality

2
output2
latent

0.98
0.98
0.66
0.70
0.65
0.54
0.62
0.46
0.5
0.97
0.97
0.08
0.31
LAE (post-hoc)(post-hoc) 0.98
0.98
0.74
0.96
LAE
LAE (online)(online)
0.98
0.98
0.91
0.96
LAE
2Out of Distribution detection. In-distribution data FashionMnist and O O D data Mnist.
2
Table
3:
(b) output(a) latent(c) AUROC (") for multiple statistics.
Area Under Receiver Operator Curve (AUROC ") for log-likelihood, typicality score, latent space
spaceIn-distribution
Our proposed LAE
is ableto discriminate between
Figure 9: Out of Distribution
detection.
data FashionMnist
and ELBOO O D data
Mnist.
latent and output
output uncertainties.
andcurves
O O D using
the deduced
variancesin latent and(c)outputspace.
(a) and (b) shows the in
ROC
for latentand output
space uncertaines.
shows(AUROC ↑) for
log-likelihood, typicality score, latent space σlatent and output space σoutput uncertainties. Online LAE
is able to discriminate between in and O O D using the deduced variances in latent and output space.
AE MSEAEMSE
AE likelihood
AElikelihood
VAE
VAE
MC-AE
MC-AE
<latexit sha1_base64="GkorlYok46sJRyyLqOP3KqnXvj8=">AAAB9HicdVDLSgNBEJyN7/UV9ehlMAiellmDGA9iVAQvQkRjhGQJs5NZMzj7cKY3GJZ8hxcPSvDqb3j3Iv6Nk0RBRQsaiqpuurv8RAoNhLxbubHxicmp6Rl7dm5+YTG/tHyh41QxXmWxjNWlTzWXIuJVECD5ZaI4DX3Ja/714cCvdbjSIo7OoZtwL6RXkQgEo2AkrwH8FrL9I3xydtRr5gvEIUNg4mwVi5ukZIjrkh23hN1Pq7D3Yu8m/Te70sy/NloxS0MeAZNU67pLEvAyqkAwyXt2I9U8oeyaXvG6oRENufay4dE9vG6UFg5iZSoCPFS/T2Q01Lob+qYzpNDWv72B+JdXTyEoeZmIkhR4xEaLglRiiPEgAdwSijOQXUMoU8LcilmbKsrA5GSbEL4+xf+Ti03H3XLIKSmUD9AI02gVraEN5KJtVEbHqIKqiKEbdIce0KPVse6tvvU0as1ZnzMr6Aes5w+tq5UK</latexit><latexit sha1_base64="4PJccs66yHuJCvUE9sHfLVrvtCE=">AAAB/XicdVDLSgMxFM3U9/iqj52bYBFcDTMjVrsQXwguK9gqtKVk0rQNzUyG5I5Yh+KvCOJCEbd+gns34t+YtgoqeuDC4Zx7k3tPEAuuwXXfrczI6Nj4xOSUPT0zOzefXVgsa5koykpUCqnOA6KZ4BErAQfBzmPFSBgIdhZ0Dvv+2QVTmsvoFLoxq4WkFfEmpwSMVM8uV4FdQrp/hAXvmFfaUjZ69WzOdfwNr+D52HU28l6h4BmS9/MFfwt7jjtAbvfF3olv3+xiPftabUiahCwCKojWFc+NoZYSBZwK1rOriWYxoR3SYhVDIxIyXUsH2/fwmlEauCmVqQjwQP0+kZJQ624YmM6QQFv/9vriX14lgeZ2LeVRnACL6PCjZiIwSNyPAje4YhRE1xBCFTe7YtomilAwgdkmhK9L8f+k7DvepuOeuLm9AzTEJFpBq2gdeWgL7aFjVEQlRNEVukH36MG6tu6sR+tp2JqxPmeW0A9Yzx9395jm</latexit><latexit sha1_base64="KsWvzdLfdQEGADSVCTDjMIBQQ3w=">AAAB8XicdZDLSgMxFIYz3h1vVZdugkVwNWQGpe1CvCG4rGBbsS2SSdM2mMkMyRmxDH0LNy4U0aUP4t6N+DamrYKK/hD4+P9zyDknTKQwQMi7MzY+MTk1PTPrzs0vLC7llleqJk414xUWy1ifhdRwKRSvgADJzxLNaRRKXgsvDwd57YprI2J1Cr2ENyPaUaItGAVrnTeAX0NW3T/qX+TyxCv4fikoYuIFwRYplQawHWxZ8D0yVH73xd1Jnt7c8kXutdGKWRpxBUxSY+o+SaCZUQ2CSd53G6nhCWWXtMPrFhWNuGlmw4n7eMM6LdyOtX0K8ND93pHRyJheFNrKiELX/M4G5l9ZPYV2sZkJlaTAFRt91E4lhhgP1sctoTkD2bNAmRZ2Vsy6VFMG9kiuPcLXpvh/qAaev+2RE5LfO0AjzaA1tI42kY8KaA8dozKqIIYUukF36N4xzq3z4DyOSsecz55V9EPO8wdqYZRp</latexit><latexit sha1_base64="Z0snb615b7mF6omLKl9PGzT6wUM=">AAAB83icdVBNS0JBFJ1nX2pfVss2QxK06fH8etnOkqBNYJAaqci8cdTBeR/M3BfKw78RRIsi2kb/pV2/pkYtqKgDFw7n3Mu99ziB4Aos682Izc0vLC7FE8nlldW19dTGZk35oaSsSn3hy0uHKCa4x6rAQbDLQDLiOoLVnUF54tevmVTc9y5gFLCWS3oe73JKQEvNJrAhRGfl/aOTcTuVtsx8tli0bWyZucKhlZsQO58r2HmcMa0p0qVEcHv1MnyvtFOvzY5PQ5d5QAVRqpGxAmhFRAKngo2TzVCxgNAB6bGGph5xmWpF05vHeFcrHdz1pS4P8FT9PhERV6mR6+hOl0Bf/fYm4l9eI4RusRVxLwiBeXS2qBsKDD6eBIA7XDIKYqQJoZLrWzHtE0ko6JiSOoSvT/H/pJY1MwXTOtdpHKMZ4mgb7aA9lEEHqIROUQVVEUUBukH36MEIjTvj0XiatcaMz5kt9APG8wccSZV4</latexit><latexit sha1_base64="0QZGBkxY4a8H/OP4dR8d+p3maJw=">AAAB/nicdVDLSgNBEJz1GeMrKp68DIqgiMtuwibxFhXBg4cIRoUkhNnJJBkyu7PM9IphCfgR/oAXD4p49Tu85W+cJAoqWtBQVHXT3eVHgmtwnIE1MTk1PTObmkvPLywuLWdWVi+1jBVlFSqFVNc+0UzwkFWAg2DXkWIk8AW78rvHQ//qhinNZXgBvYjVA9IOeYtTAkZqZNZrwG4hOTs8wTuR1LDfkXS338hsOXYul/VcDzt23nOy+ZwhBa9YzB1g13ZG2Cpt1vbuB6VeuZF5rzUljQMWAhVE66rrRFBPiAJOBeuna7FmEaFd0mZVQ0MSMF1PRuf38bZRmrgllakQ8Ej9PpGQQOte4JvOgEBH//aG4l9eNYZWsZ7wMIqBhXS8qBULDBIPs8BNrhgF0TOEUMXNrZh2iCIUTGJpE8LXp/h/cpm1Xc92zk0aR2iMFNpAm2gHuaiASugUlVEFUZSgB/SEnq0769F6sV7HrRPW58wa+gHr7QPWB5hm</latexit><latexit sha1_base64="p9o4/EP9ReqGKmwDd+79M/HeFsI=">AAAB/HicbVDLSgMxFM34rPVV7dJNaBEqQpkRRJdVEVy4qGAf0BlKJk3b0ExmSO6Iw1C/wr0bF4q49UPc9W9MHwttPXDhcM69yb3HjwTXYNsja2l5ZXVtPbOR3dza3tnN7e3XdRgrymo0FKFq+kQzwSWrAQfBmpFiJPAFa/iDq7HfeGBK81DeQxIxLyA9ybucEjBSO5d3gT1CentxjUuhHL9yNGzninbZngAvEmdGipWCe/w8qiTVdu7b7YQ0DpgEKojWLceOwEuJAk4FG2bdWLOI0AHpsZahkgRMe+lk+SE+NEoHd0NlSgKeqL8nUhJonQS+6QwI9PW8Nxb/81oxdM+9lMsoBibp9KNuLDCEeJwE7nDFKIjEEEIVN7ti2ieKUDB5ZU0IzvzJi6R+UnZOy/adSeMSTZFBB6iASshBZ6iCblAV1RBFCXpBb+jderJerQ/rc9q6ZM1m8ugPrK8fCreXWQ==</latexit><latexit sha1_base64="jBZcmeivIzibP1+vDaDJXblS7Tg=">AAACB3icbVDLSsNAFJ34rPUVdSnIYCvUTUkKosuiG5cV7AOaGCbTSTt08mDmRiyhOzf+ihsXirj1F9z5N07bLLT1wMDhnHu5c46fCK7Asr6NpeWV1bX1wkZxc2t7Z9fc22+pOJWUNWksYtnxiWKCR6wJHATrJJKR0Bes7Q+vJn77nknF4+gWRglzQ9KPeMApAS155lHFP8VlR/F+SO5qXuYAe4AsTiFJYTwue2bJqlpT4EVi56SEcjQ888vpxTQNWQRUEKW6tpWAmxEJnAo2LjqpYgmhQ9JnXU0jEjLlZtMcY3yilR4OYqlfBHiq/t7ISKjUKPT1ZEhgoOa9ifif100huHAzHulQLKKzQ0EqMMR4UgrucckoiJEmhEqu/4rpgEhCQVdX1CXY85EXSatWtc+q1k2tVL/M6yigQ3SMKshG56iOrlEDNRFFj+gZvaI348l4Md6Nj9nokpHvHKA/MD5/AHKzmRI=</latexit><latexit sha1_base64="W8byag/ZGlhdEpIYtTa5YPTXdZo=">AAACB3icbVC7SgNBFJ31GeMrainIYCLEJuwGRMugjWUE84BkXe5OJsmQ2Qczd8WwpLPxV2wsFLH1F+z8GyePQhMPDBzOuYc79/ixFBpt+9taWl5ZXVvPbGQ3t7Z3dnN7+3UdJYrxGotkpJo+aC5FyGsoUPJmrDgEvuQNf3A19hv3XGkRhbc4jLkbQC8UXcEAjeTljopwSgttLXoB3JW9tI38AVMJyEMcjQpeLm+X7AnoInFmJE9mqHq5r3YnYklg4kyC1i3HjtFNQaFgko+y7UTzGNgAerxlaAgB1246uWNET4zSod1ImRcinai/EykEWg8D30wGgH09743F/7xWgt0LNxVhnJi72HRRN5EUIzouhXaE4gzl0BBgSpi/UtYHBQxNdVlTgjN/8iKpl0vOWcm+Kecrl7M6MuSQHJMiccg5qZBrUiU1wsgjeSav5M16sl6sd+tjOrpkzTIH5A+szx8x+Jjo</latexit><latexit sha1_base64="7vyzwEWx06uTkbuXL8gPMSPaleU=">AAACGHicbVC7TgJBFJ3FF+ILtbSZCCbQ4C6J0RKlsRONPBLYkNlhFifM7mxm7moI4TNs/BUbC42xpfNvHGALBU91cs65M/ceLxJcg21/W6mV1bX1jfRmZmt7Z3cvu3/Q0DJWlNWpFFK1PKKZ4CGrAwfBWpFiJPAEa3qD6tRvPjKluQzvYRgxNyD9kPucEjBSN3taoEV8Wb+7qeJCvhNHRCn5lC9iXyocxAJ4JBjWYNIaONWlbjZnl+wZ8DJxEpJDCWrd7KTTkzQOWAhUEK3bjh2BOyLKPCfYONOJNYsIHZA+axsakoBpdzQ7bIxPjNKb7eLLEPBM/T0xIoHWw8AzyYDAg170puJ/XjsG/8Id8TCKgYV0/pEfCwwST1vCPa4YBTE0hFDFza6YPhBFKJguM6YEZ/HkZdIol5yzkn1bzlWukjrS6AgdowJy0DmqoGtUQ3VE0TN6Re/ow3qx3qxP62seTVnJzCH6A2vyA+TgnoI=</latexit>

8
good reconstructions (low log p(x)). Using our online training procedure results in a much higher
log-likelihood. This indicate that every sampled NN predict good reconstructions.

ID<latexit sha1_base64="Xfsu/d6TbYJXzqSSrCh0QHxFAGE=">AAAB8nicbVBNS8NAEN34WetX1aOXxSJ4KklB9FjUg94q2A9oQ9lsJ+3STTbsTsQS+jO8eFDEq7/Gm//GbZuDtj4YeLw3w8y8IJHCoOt+Oyura+sbm4Wt4vbO7t5+6eCwaVSqOTS4kkq3A2ZAihgaKFBCO9HAokBCKxhdT/3WI2gjVPyA4wT8iA1iEQrO0EqdLsITGp7d3Ux6pbJbcWegy8TLSZnkqPdKX92+4mkEMXLJjOl4boJ+xjQKLmFS7KYGEsZHbAAdS2MWgfGz2ckTemqVPg2VthUjnam/JzIWGTOOAtsZMRyaRW8q/ud1Ugwv/UzESYoQ8/miMJUUFZ3+T/tCA0c5toRxLeytlA+ZZhxtSkUbgrf48jJpViveecW9r5ZrV3kcBXJMTsgZ8cgFqZFbUicNwokiz+SVvDnovDjvzse8dcXJZ47IHzifP4WbkWY=</latexit>

OoD<latexit sha1_base64="Wy9c8Vbhv0+b9BeO8NNHCHOUkfQ=">AAAB83icbVBNS8NAEJ34WetX1aOXYBE8laQgeizqwZsV7Ac0oWy223bpZjfsTsQS+je8eFDEq3/Gm//GbZuDtj4YeLw3w8y8KBHcoOd9Oyura+sbm4Wt4vbO7t5+6eCwaVSqKWtQJZRuR8QwwSVrIEfB2olmJI4Ea0Wj66nfemTacCUfcJywMCYDyfucErRSECB7QkOzO3Uz6ZbKXsWbwV0mfk7KkKPeLX0FPUXTmEmkghjT8b0Ew4xo5FSwSTFIDUsIHZEB61gqScxMmM1unrinVum5faVtSXRn6u+JjMTGjOPIdsYEh2bRm4r/eZ0U+5dhxmWSIpN0vqifCheVOw3A7XHNKIqxJYRqbm916ZBoQtHGVLQh+IsvL5NmteKfV7z7arl2lcdRgGM4gTPw4QJqcAt1aACFBJ7hFd6c1Hlx3p2PeeuKk88cwR84nz9goZHl</latexit>

Fig. 8 shows the latent representation, mean and variance of the reconstructions with the KFAC, exact
and approximate diagonal for both post-hoc and online setup. Note that the online training makes
the uncertainties more well fitted, both in latent and data space. These well-fitted uncertainties haveseveral practical downstream applications, which we demonstrate next.

Out-of-Distribution (O O D) Detection capabilities are critical for identifying distributional shifts,
outliers and irregular user inputs, which can hinder the propagation of erroneous decision in an
automated system. We evaluate O O D performance on the commonly used benchmarks (Nalisnick(a) AE likelihood
LAE (online)(c)(Xiao
VAE et al., 2017)(d)asLAE(post-hoc)and (e)et al., 2019b), where(b)weMC-AEuse FASHION M NIST
in-distribution
M NIST(Lecun et al., 1998) as O O D. Fig. 9 (c) shows that our online LAE outperforms existing models inboth log-likelihood and Typicality score (Nalisnick et al., 2019a). This stems from the calibrated
model uncertainties, which are exemplified in the models ability to detect O O D examples from theuncertainty deduced in latent and output space; see Fig. 9 (a,b) for ROC curves.
<latexit sha1_base64="qqrOroZOaCM6P480Vy9vddNAIIg=">AAACAXicbVDLSgMxFM3UV62vUTeCm9AiVIQyI4guqyK4rGAf0Cklk6ZtaGYyJHfEYagbf8FPcONCEbf+hbv+jeljodUDFw7n3Jvce/xIcA2OM7IyC4tLyyvZ1dza+sbmlr29U9MyVpRVqRRSNXyimeAhqwIHwRqRYiTwBav7g8uxX79jSnMZ3kISsVZAeiHvckrASG17zwN2D2mRHOLzKyz4wLzUl7IzbNsFp+RMgP8Sd0YK5bx39DQqJ5W2/eV1JI0DFgIVROum60TQSokCTgUb5rxYs4jQAemxpqEhCZhupZMLhvjAKB3clcpUCHii/pxISaB1EvimMyDQ1/PeWPzPa8bQPWulPIxiYCGdftSNBQaJx3HgDleMgkgMIVRxsyumfaIIBRNazoTgzp/8l9SOS+5JybkxaVygKbJoH+VREbnoFJXRNaqgKqLoAT2jV/RmPVov1rv1MW3NWLOZXfQL1uc3q6CZXw==</latexit><latexit sha1_base64="MPOyih6M+2RZtx/RB/FWg0JXvTU=">AAACAHicbVDLSgNBEJz1GeMr6sGDlyFBSBDCriB6jIrgwUMEkwhJCLOTjg6ZnV1mesWw5OI3+AdePCji1c/wlr9x8jhoYkFDUdU9011+JIVB1x04c/MLi0vLqZX06tr6xmZma7tqwlhzqPBQhvrWZwakUFBBgRJuIw0s8CXU/O750K89gDYiVDfYi6AZsDslOoIztFIrs9tAeMQkDwV6dXpB86EavlTotzI5t+iOQGeJNyG5UrZx8Dwo9cqtzHejHfI4AIVcMmPqnhthM2EaBZfQTzdiAxHjXXYHdUsVC8A0k9EBfbpvlTbthNqWQjpSf08kLDCmF/i2M2B4b6a9ofifV4+xc9JMhIpiBMXHH3ViSTGkwzRoW2jgKHuWMK6F3ZXye6YZR5tZ2obgTZ88S6qHRe+o6F7bNM7IGCmyR7IkTzxyTErkkpRJhXDSJy/kjbw7T86r8+F8jlvnnMnMDvkD5+sH+n+YVw==</latexit><latexit sha1_base64="JgkG9z/37DNvXKZK3idGHF7zlH4=">AAAB8HicbVBNSwMxEJ2tX7V+VT16CRahHiy7BdFjtQhehAr2Q9qlZNO0DU2yS5IVytJf4cWDIl79Od78N6btHrT1wcDjvRlm5gURZ9q47reTWVldW9/Ibua2tnd29/L7Bw0dxorQOgl5qFoB1pQzSeuGGU5bkaJYBJw2g1F16jefqNIslA9mHFFf4IFkfUawsdJjMThFd9Wzq5tuvuCW3BnQMvFSUoAUtW7+q9MLSSyoNIRjrdueGxk/wcowwukk14k1jTAZ4QFtWyqxoNpPZgdP0IlVeqgfKlvSoJn6eyLBQuuxCGynwGaoF72p+J/Xjk3/0k+YjGJDJZkv6sccmRBNv0c9pigxfGwJJorZWxEZYoWJsRnlbAje4svLpFEueecl975cqFyncWThCI6hCB5cQAVuoQZ1ICDgGV7hzVHOi/PufMxbM046cwh/4Hz+AIoTjuw=</latexit><latexit sha1_base64="pgV9ajDBpasXAFak2yi/+G6d3aE=">AAAB7nicbVDLSgNBEOyNrxhfUY9eBoMQL2E3IHqMiuAxgnlAsoTZyWwyZHZ2mekVQshHePGgiFe/x5t/4yTZgyYWNBRV3XR3BYkUBl3328mtrW9sbuW3Czu7e/sHxcOjpolTzXiDxTLW7YAaLoXiDRQoeTvRnEaB5K1gdDvzW09cGxGrRxwn3I/oQIlQMIpWapXZOWle3/WKJbfizkFWiZeREmSo94pf3X7M0ogrZJIa0/HcBP0J1SiY5NNCNzU8oWxEB7xjqaIRN/5kfu6UnFmlT8JY21JI5urviQmNjBlHge2MKA7NsjcT//M6KYZX/kSoJEWu2GJRmEqCMZn9TvpCc4ZybAllWthbCRtSTRnahAo2BG/55VXSrFa8i4r7UC3VbrI48nACp1AGDy6hBvdQhwYwGMEzvMKbkzgvzrvzsWjNOdnMMfyB8/kDoj6Ocg==</latexit><latexit sha1_base64="J83hmBtnn2EdKbOMuRgOM8u8X+4=">AAAB+3icbVBNS8NAEN34WetXrEcvi0VoD5akIHqsiuDBQwX7AW0om82mXbrJht2JWEr/ihcPinj1j3jz37htc9DWBwOP92aYmecngmtwnG9rZXVtfWMzt5Xf3tnd27cPCk0tU0VZg0ohVdsnmgkeswZwEKydKEYiX7CWP7ye+q1HpjSX8QOMEuZFpB/zkFMCRurZhVJQxneXN7iUSA2nA0nLPbvoVJwZ8DJxM1JEGeo9+6sbSJpGLAYqiNYd10nAGxMFnAo2yXdTzRJCh6TPOobGJGLaG89un+ATowQ4lMpUDHim/p4Yk0jrUeSbzojAQC96U/E/r5NCeOGNeZykwGI6XxSmAoPE0yBwwBWjIEaGEKq4uRXTAVGEgokrb0JwF19eJs1qxT2rOPfVYu0qiyOHjtAxKiEXnaMaukV11EAUPaFn9IrerIn1Yr1bH/PWFSubOUR/YH3+AELMkqY=</latexit>

Fig. 10 shows distribution of the output variances for in- and O O D data. This illustrates that using LA
improves OoD detection. Furthermore, the online training improves the model calibration.
Missing Data imputation. Another application of stochastic representation learning is to provide
distributions over unobserved values (Rezende et al., 2014). In many application domains, sensor
readings go missing, which we may mimic by letting parts of an image be unobserved. Rezende et al.
(2014) show that we can then draw samples from the distribution of the entire image conditioned on
the observed part, by imputing the missing pixels with noise and repeatedly encode and decode while
keeping observed pixels fixed. Fig. 11 show samples using this procedure from a VAE, a post-hoc
LAE and our online LAE , where we only observe the lower half of an MNIST image. This implies
ambiguity about the original digit, e.g. the lower half of a “5“ could be a “3“ and similarly a “7“ could
be a “9“. Our LAE captures this ambiguity, which is exemplified by the multi-modal reconstructions
from the sampled networks in Fig. 11. The baselines only capture unimodal reconstructions.
Capturing the ambiguity of partly missing data can improve downstream tasks such as the calibrationof an image classifier. In Fig. 11 (c) we demonstrate how averaging the predictions of a simple
classifier across reconstructions improves standard calibration metrics. This is because the classifier
inherits the uncertainty and ambiguity in the learned
representations. A deep ensemble of AEs perform
2(b) Output variance,(a) Latent variance, latent
similarly to ours, but comes at the cost of training
and storing multiple models.
2
output<latexit sha1_base64="RhkKxl4gQOhXYJmPIMG5D402ky8=">AAACGHicbVDLSgNBEJz1bXxFPXoZTAQFibsB0aPoxZsKJgaSGHonnTg4O7vM9IphyWd48Ve8eFDEqzf/xsnj4KtgoKiqpqcrTJS05Puf3sTk1PTM7Nx8bmFxaXklv7pWtXFqBFZErGJTC8GikhorJElhLTEIUajwKrw9GfhXd2isjPUl9RJsRtDVsiMFkJNa+b3tcIefpZSkxO/ASNACd3mxYWU3gutyK2sQ3lMWDxP9frGVL/glfwj+lwRjUmBjnLfyH412LNIINQkF1tYDP6FmBoakUNjPNVKLCYhb6GLdUQ0R2mY2PKzPt5zS5p3YuKeJD9XvExlE1vai0CUjoBv72xuI/3n1lDqHzUxqdxRqMVrUSRWnmA9a4m1pUJDqOQLCSPdXLm7AgCDXZc6VEPw++S+plkvBfsm/KBeOjsd1zLENtsm2WcAO2BE7ZeeswgR7YE/shb16j96z9+a9j6IT3nhmnf2A9/EFVtOgAg==</latexit><latexit sha1_base64="0KDebwMzIXU29QBkI5RMWdmysw8=">AAACGHicbVC7SgNBFJ2NrxhfUUubwUSIIHE3IFoGbSwsIpgHJOtydzKbDJl9MDMbDEs+w8ZfsbFQxDadf+Mk2UITDwwczrmXuee4EWdSmea3kVlZXVvfyG7mtrZ3dvfy+wcNGcaC0DoJeShaLkjKWUDriilOW5Gg4LucNt3BzdRvDqmQLAwe1Ciitg+9gHmMgNKSkz8vwSm+A0UDhYcgGASEnuFiR7KeD48VJ+ko+qQSPpsYj4tOvmCWzRnwMrFSUkApak5+0umGJPb1OuEgZdsyI2UnIBQjnI5znVjSCMgAerStaQA+lXYyCzbGJ1rpYi8U+ukDZ+rvjQR8KUe+qyd9UH256E3F/7x2rLwrO2FBFOtcZP6RF3OsQjxtCXeZoETxkSZABNO3YtIHAUTpLnO6BGsx8jJpVMrWRdm8rxSq12kdWXSEjlEJWegSVdEtqqE6IugZvaJ39GG8GG/Gp/E1H80Y6c4h+gNj8gPRa5+v</latexit>

ID<latexit sha1_base64="Xfsu/d6TbYJXzqSSrCh0QHxFAGE=">AAAB8nicbVBNS8NAEN34WetX1aOXxSJ4KklB9FjUg94q2A9oQ9lsJ+3STTbsTsQS+jO8eFDEq7/Gm//GbZuDtj4YeLw3w8y8IJHCoOt+Oyura+sbm4Wt4vbO7t5+6eCwaVSqOTS4kkq3A2ZAihgaKFBCO9HAokBCKxhdT/3WI2gjVPyA4wT8iA1iEQrO0EqdLsITGp7d3Ux6pbJbcWegy8TLSZnkqPdKX92+4mkEMXLJjOl4boJ+xjQKLmFS7KYGEsZHbAAdS2MWgfGz2ckTemqVPg2VthUjnam/JzIWGTOOAtsZMRyaRW8q/ud1Ugwv/UzESYoQ8/miMJUUFZ3+T/tCA0c5toRxLeytlA+ZZhxtSkUbgrf48jJpViveecW9r5ZrV3kcBXJMTsgZ8cgFqZFbUicNwokiz+SVvDnovDjvzse8dcXJZ47IHzifP4WbkWY=</latexit>

OoD<latexit sha1_base64="Wy9c8Vbhv0+b9BeO8NNHCHOUkfQ=">AAAB83icbVBNS8NAEJ34WetX1aOXYBE8laQgeizqwZsV7Ac0oWy223bpZjfsTsQS+je8eFDEq3/Gm//GbZuDtj4YeLw3w8y8KBHcoOd9Oyura+sbm4Wt4vbO7t5+6eCwaVSqKWtQJZRuR8QwwSVrIEfB2olmJI4Ea0Wj66nfemTacCUfcJywMCYDyfucErRSECB7QkOzO3Uz6ZbKXsWbwV0mfk7KkKPeLX0FPUXTmEmkghjT8b0Ew4xo5FSwSTFIDUsIHZEB61gqScxMmM1unrinVum5faVtSXRn6u+JjMTGjOPIdsYEh2bRm4r/eZ0U+5dhxmWSIpN0vqifCheVOw3A7XHNKIqxJYRqbm916ZBoQtHGVLQh+IsvL5NmteKfV7z7arl2lcdRgGM4gTPw4QJqcAt1aACFBJ7hFd6c1Hlx3p2PeeuKk88cwR84nz9goZHl</latexit><latexit sha1_base64="qqrOroZOaCM6P480Vy9vddNAIIg=">AAACAXicbVDLSgMxFM3UV62vUTeCm9AiVIQyI4guqyK4rGAf0Cklk6ZtaGYyJHfEYagbf8FPcONCEbf+hbv+jeljodUDFw7n3Jvce/xIcA2OM7IyC4tLyyvZ1dza+sbmlr29U9MyVpRVqRRSNXyimeAhqwIHwRqRYiTwBav7g8uxX79jSnMZ3kISsVZAeiHvckrASG17zwN2D2mRHOLzKyz4wLzUl7IzbNsFp+RMgP8Sd0YK5bx39DQqJ5W2/eV1JI0DFgIVROum60TQSokCTgUb5rxYs4jQAemxpqEhCZhupZMLhvjAKB3clcpUCHii/pxISaB1EvimMyDQ1/PeWPzPa8bQPWulPIxiYCGdftSNBQaJx3HgDleMgkgMIVRxsyumfaIIBRNazoTgzp/8l9SOS+5JybkxaVygKbJoH+VREbnoFJXRNaqgKqLoAT2jV/RmPVov1rv1MW3NWLOZXfQL1uc3q6CZXw==</latexit><latexit sha1_base64="MPOyih6M+2RZtx/RB/FWg0JXvTU=">AAACAHicbVDLSgNBEJz1GeMr6sGDlyFBSBDCriB6jIrgwUMEkwhJCLOTjg6ZnV1mesWw5OI3+AdePCji1c/wlr9x8jhoYkFDUdU9011+JIVB1x04c/MLi0vLqZX06tr6xmZma7tqwlhzqPBQhvrWZwakUFBBgRJuIw0s8CXU/O750K89gDYiVDfYi6AZsDslOoIztFIrs9tAeMQkDwV6dXpB86EavlTotzI5t+iOQGeJNyG5UrZx8Dwo9cqtzHejHfI4AIVcMmPqnhthM2EaBZfQTzdiAxHjXXYHdUsVC8A0k9EBfbpvlTbthNqWQjpSf08kLDCmF/i2M2B4b6a9ofifV4+xc9JMhIpiBMXHH3ViSTGkwzRoW2jgKHuWMK6F3ZXye6YZR5tZ2obgTZ88S6qHRe+o6F7bNM7IGCmyR7IkTzxyTErkkpRJhXDSJy/kjbw7T86r8+F8jlvnnMnMDvkD5+sH+n+YVw==</latexit><latexit sha1_base64="JgkG9z/37DNvXKZK3idGHF7zlH4=">AAAB8HicbVBNSwMxEJ2tX7V+VT16CRahHiy7BdFjtQhehAr2Q9qlZNO0DU2yS5IVytJf4cWDIl79Od78N6btHrT1wcDjvRlm5gURZ9q47reTWVldW9/Ibua2tnd29/L7Bw0dxorQOgl5qFoB1pQzSeuGGU5bkaJYBJw2g1F16jefqNIslA9mHFFf4IFkfUawsdJjMThFd9Wzq5tuvuCW3BnQMvFSUoAUtW7+q9MLSSyoNIRjrdueGxk/wcowwukk14k1jTAZ4QFtWyqxoNpPZgdP0IlVeqgfKlvSoJn6eyLBQuuxCGynwGaoF72p+J/Xjk3/0k+YjGJDJZkv6sccmRBNv0c9pigxfGwJJorZWxEZYoWJsRnlbAje4svLpFEueecl975cqFyncWThCI6hCB5cQAVuoQZ1ICDgGV7hzVHOi/PufMxbM046cwh/4Hz+AIoTjuw=</latexit>(a) AE likelihood(b) MC-AE<latexit sha1_base64="pgV9ajDBpasXAFak2yi/+G6d3aE=">AAAB7nicbVDLSgNBEOyNrxhfUY9eBoMQL2E3IHqMiuAxgnlAsoTZyWwyZHZ2mekVQshHePGgiFe/x5t/4yTZgyYWNBRV3XR3BYkUBl3328mtrW9sbuW3Czu7e/sHxcOjpolTzXiDxTLW7YAaLoXiDRQoeTvRnEaB5K1gdDvzW09cGxGrRxwn3I/oQIlQMIpWapXZOWle3/WKJbfizkFWiZeREmSo94pf3X7M0ogrZJIa0/HcBP0J1SiY5NNCNzU8oWxEB7xjqaIRN/5kfu6UnFmlT8JY21JI5urviQmNjBlHge2MKA7NsjcT//M6KYZX/kSoJEWu2GJRmEqCMZn9TvpCc4ZybAllWthbCRtSTRnahAo2BG/55VXSrFa8i4r7UC3VbrI48nACp1AGDy6hBvdQhwYwGMEzvMKbkzgvzrvzsWjNOdnMMfyB8/kDoj6Ocg==</latexit><latexit sha1_base64="J83hmBtnn2EdKbOMuRgOM8u8X+4=">AAAB+3icbVBNS8NAEN34WetXrEcvi0VoD5akIHqsiuDBQwX7AW0om82mXbrJht2JWEr/ihcPinj1j3jz37htc9DWBwOP92aYmecngmtwnG9rZXVtfWMzt5Xf3tnd27cPCk0tU0VZg0ohVdsnmgkeswZwEKydKEYiX7CWP7ye+q1HpjSX8QOMEuZFpB/zkFMCRurZhVJQxneXN7iUSA2nA0nLPbvoVJwZ8DJxM1JEGeo9+6sbSJpGLAYqiNYd10nAGxMFnAo2yXdTzRJCh6TPOobGJGLaG89un+ATowQ4lMpUDHim/p4Yk0jrUeSbzojAQC96U/E/r5NCeOGNeZykwGI6XxSmAoPE0yBwwBWjIEaGEKq4uRXTAVGEgokrb0JwF19eJs1qxT2rOPfVYu0qiyOHjtAxKiEXnaMaukV11EAUPaFn9IrerIn1Yr1bH/PWFSubOUR/YH3+AELMkqY=</latexit>(c) VAE(d) LAE (post-hoc)(e) LAE (online)

Figure 10: Histograms of variance for in- and O O D reconstructions in output space. Note that MCAE separates the distributions well, but the model assign higher variance to the in- than O O D data.
8

Figure 11: Missing data imputation & generative capabilities. Online training of LAE improves
representational robustness. This is exemplified by the multimodal behavior in the data imputation(top rows) that accurately model the ambiguity in the data. The bottom two rows, shows that the LAE
is able to generate crisp digits from random noise.

MethodMSE ↓Acc. ↑

log p(x) ↑

ECE ↓

MCE ↓

RMSCE ↓Classifier
0.53
0.16
0.25
0.18
VAE
104.73 -104.75
0.22
0.18
0.34
0.19
MC-AE
106.05 -106.05
0.45
0.28
0.38
0.29
AE Ensemble
96.23-100.94
0.53
0.12
0.2
0.13
LAE (post-hoc) 101.62 -107.25
0.51
0.16
0.27
0.18
LAE (online)
99.59-106.29
0.53
0.12
0.16
0.13
Table 3: Reconstruction quality measured by the MSE and log-likelihood for the data imputation.
Our well-calibrated uncertainties propagates to the MNIST classifier and improves the calibration
metrics ECE, MCE and RMSCE.
When the entire input image is missing, the imputation procedure can be seen as a sampling mechanism, such that our LAE can be viewed as a generative model. The bottom rows in Fig. 11 show that
the LAE indeed does generate sharp images from a multi-modal distribution.
Semi-supervised learning combines a small amount of label data with a large amount of unlabeled
data. The hope is that the structure in the unlabeled data can be used to infer properties of the data thatcannot be extracted from a few labeled points. Embedding the same labeled data point multiple times
using a stochastic representation scales up the amount of labeled data that is available during training.
Attribute

AE

VAE

MCAE

LAE
(online)Arched EyebrowsAttractive
Bald
Wearing Lipstick
Heavy Makeup

0.50
0.52
0.98
0.52
0.45

0.52
0.50
0.98
0.49
0.52

0.55
0.49
0.98
0.50
0.49

0.60
0.53
0.98
0.54
0.56

Overall

0.73

0.72

0.73

0.74

Figure 12: Accuracy as an function of the Table 4: Semi-supervised classification accuracy
on C ELEBA using only 10 labeled datapoints.
number of labels per class on MNIST.
Fig. 12 shows the accuracy of a K-nearest neighbour classifier trained on different amounts oflabeled data from the MNIST dataset. For all models with a stochastic encoder, we encode eachlabeled datapoint 100 times and repeat the experiment 5 times. When only few labels per classavailable (1-20) we clearly observe that our LAE model outperforms all other models, stochastic anddeterministic. Increasing the number of labels beyond 100 per class makes the AE and LAE equal intheir classification performance with the AE model eventually outperforming the LAE model.
In Tab. 4 we conduct a similar experiment on the C ELEBA (Liu et al., 2015) facial dataset, where thethe task is to predict 40 different binary labels per datapoint. When evaluating the overall accuracy of
predicting all 40 facial attributes, we see no significant difference in performance. However, when
we zoom in on specific facial attributes we gain a clear performance advantage over other models.
9

Limitations. Empirically, the LAE improvements are more significant for overparameterized networks. The additional capacity seems to help the optimizer find a local mode where a Gaussian fit is appropriate. It seems the regularization induced by marginalizing θ compensates for the added flexibility.

6

Conclusion

In this paper we have introduced a Bayesian autoencoder that is realized using Laplace approximations. Unlike current models, this Laplacian autoencoder produces well-behaved uncertainties
in both latent and data space. We have proposed a novel variational lower-bound of the autoencoder
evidence and a efficient way to compute its Hessian on high dimensional data that scales linearlywith data size. Empirically, we demonstrate that our proposed model predicts reliable stochastic
representations that are useful for a multitude of downstream tasks: out-of-distribution detection,
missing data imputation and semi-supervised classification. Our work opens the way for fully
Bayesian representation learning where we can marginalize the representation in downstream tasks.
We find this to consistently improve performance.

Acknowledgments and Disclosure of Funding
This work was supported by research grants (15334, 42062) from VILLUM FONDEN. This project
has also received funding from the European Research Council (ERC) under the European Union’s
Horizon 2020 research and innovation programme (grant agreement 757360). This work was fundedin part by the Novo Nordisk Foundation through the Center for Basic Machine Learning Research inLife Science (NNF20OC0062606).

10

— Supplementary Material —The supplementary material is organised as follows. First, we give more technical details on theexperiments. Second, we discuss the Laplace approximation and the “mean shift” issue encounteredin non-local maxima. Third, we elaborate more on the difference between related works and our
proposed method. Fourth, we present a much more thorough explanation of the proposed method
accompanied with relevant proofs. Fifth, we present more details on the hessian derivations.

A

Experimental details

All experiments was conducted on one of the following datasets: MNIST (Lecun et al., 1998),
FASHIONMNIST (Xiao et al., 2017) and CELEBA (Liu et al., 2015). For training and testing the default

splits were used. For validation, we sampled 5000 data points randomly from the training sets. All
images were normalized to be in the [0, 1] range and for the CELEBA dataset the images were resized
to 64 × 64 additionally.

If nothing else is stated, the models were trained with the following configuration: we used Adam
optimizer (Kingma and Ba, 2015b) with learning rate 0.001 and default Pytorch settings (Paszke
et al., 2019). The learning rate was adjusted using ReduceLROnPlateau learning rate scheduler withparameters factor=0.5 and patience=5, meaning that the learning rate was halved whenever 5 epochs
had passed where the validation loss had not decreased. The mean squared error loss was used as thereconstruction loss in all models. Models where trained until convergence, defined as whenever thevalidation loss had not decreased for 8 epochs. Models trained on MNIST and FASHIONMNIST used abatch size of 64 and on CELEBA a batch size of 128 was used.
Model specific details:• VAE: Models were trained with KL-scaling of 0.001. We use two encoders and two decoders,such that the model has twice the number of parameters compared to the other models.
• MC - AE: Models were trained with dropout between all trainable layers with probability
p = 0.2. We keep the same dropout rate during testing.
• EMSEMBLE - AE: Each ensemble consists of 5 models, each initialised with a different
random seed.
• LAE ( POSTHOC ): For experiments with linear layers we used the Laplace Redux (Daxberger
et al., 2021) implementation. For convolutions, we found it necessary to use our proposed
hessian approximation. We use a diagonal approximation of the hessian in all experiments. After fitting the hessian, we optimise for the prior precision using the marginal
likelihood (Daxberger et al., 2021). We use 100 MC sampling in all experiments.
• LAE ( ONLINE ): We use the exact diagonal in experiments with linear layers and the mixed
diagonal approximation in all experiments with convolutional layers. We use a hessian
memory factor of 0.0001 and sample only 1 network per iteration. We found that it was notnecessary to optimise for the prior precision when trained online.
A.1

Hessian approximation

We use a linear encoder-decoder with three layers in the encoder and decoder, TANH activation
functions, and latent size 2. We choose this architecture as Laplace Redux (Daxberger et al., 2021)
supports various hessian approximations for this simple network. We use Laplace Redux for all
post-hoc experiments except for the approximate diagonal hessian.
A.2

Out-of-distribution

We use a convolutional encoder-decoder architecture. The encoder consisted of a CONV 2 D, TANH,MAXPOOL 2 D , CONV 2 D , TANH, MAXPOOL 2 D , LINEAR , TANH , LINEAR , TANH, LINEAR , where thedecoder was mirrored but with nearest neighbour Upsampling rather than MAXPOOL 2 D. We used alatent size 2 in these experiments for all models.
11

A.3

Missing data imputation

To elaborate on the procedure, we reconstruct 5 samples from the half/fully masked image. For each ofthese reconstructions, we make 5 more reconstructions and take the average of these reconstructions.
The intuition is that the first stage explores the multi-modal behaviour of the reconstructions. In thesecond stage, the uncertainty of the reconstructed digit is reduced, and each sample will reconstruct
the same modality. By averaging over these modalities, we achieve a more crisp reconstruction. We
use the same architecture as in the hessian approximation experiment.
A.4

Semi-supervised learning

For the experiments on MNIST, we use a single convergence checkpoint for each model. We use
the same model architecture as in the hessian approximation. We did 5 repetitions for each model
where we first sampled n labels from each of the 10 classes from the validation set, then embedded
the 10 × n data points into latent space, trained a KNN classifier on the embedded points and finallyevaluated the accuracy of the classifier on the remaining validation set. This procedure was repeatedfor different values of n in the [1, 100] range. For the stochastic encoders (VAE, MC - AE, LAE), werepeated the embedding step 100 times with the goal that the uncertainty could help the downstream
classification. For the KNN classifier we use cross-validation (K = 2) to find the optimal number ofnearest neighbours.
For the experiments on CELEBA we repeated the exact same experiments but with a fixed value
if n = 10. Additionally, the classifier was changed to a multi-label version KNN-classifier toaccommodate the multiple binary features in the dataset. For CELEBA we use a convolutionalarchitecture. The encoder consists of 5 convolutional layers with TANH and MAXPOOL 2 D in between
each parametric layer. We use a latent size of 64. The decoder mirrors the encoder, but we replace
MAXPOOL 2 D with nearest neighbor upsampling.

B

Laplace Approximation

Laplace approximation is an operator that maps local properties (derivatives) to global properties.
The idea is to infer a density on every point based on the curvature in a single point. This is done
through a Taylor expansion.
Given a vectorial space Θ of size D, let P (θ) be an arbitrary distribution on Θ and let θ ∗ ∈ Θ be an
arbitrary point. Consider the second order Taylor expansion of the log density around θ ∗1
ln P (θ) = ln P (θ ∗ ) + ∇θ ln P (θ ∗ )(θ − θ ∗ ) + (θ − θ ∗ )> ∇2θ ln P (θ ∗ )(θ − θ ∗ ) + O(kθ − θ ∗ k3 ),2where[∇θ ln P (θ ∗ )]i =∂ln P (θ)
∂θ i
θ=θ ∗

[∇2θ ln P (θ ∗ )]ij =∂2.
ln P (θ)∂θ i ∂θ j
θ=θ ∗are the first and second order derivatives.
Note that if P (θ) is a Gaussian, then its log density is a second order polynomial and the second order
Taylor expansion is exact. This implies that, starting from a Gaussian, the Laplace approximation caninfer the full exact density just from the values of ∇θ ln P and ∇2θ ln P in a single point θ ∗ .
The main takeaway from the Laplace approximation is the strong, intrinsic, tie between the covariancematrix and the negative inverse of the hessian of the log probability: −(∇2θ ln P (θ ∗ ))−1 .
B.1

If θ ∗ is a local maxima

If ∇θ ln P (θ ∗ ) = 0 the Taylor expansion consists of only two terms and the Laplace derivation
is easier. We also present this in order to develop intuition, although this case is a subcase of the
non-local maxima case. In the next section we will consider the more general setting.
Define the Gaussian∗

LAPLACEmax (θ ; P ) := Nθ|µ = θ ∗ , σ 2 = −(∇2θ ln P (θ ∗ ))−1 ,
12

(13)

which has density
Q(θ) :=P (θ ∗ ) 1 (θ−θ∗ )> (∇2θ ln P (θ∗ ))−1 (θ−θ∗ )
e2
,
ZQ

p
where ZQ = P (θ ∗ ) (−2π)D det(∇2θ ln P (θ ∗ )) is the normalizing constant. Then, Q(θ) is a good
approximation of P (θ) in the sense that1ln Q(θ) + ln ZQ = ln P (θ ∗ ) + (θ − θ ∗ )> (∇2θ ln P (θ ∗ ))−1 (θ − θ ∗ ) ∼=2
1∼= ln P (θ ∗ ) + (θ − θ ∗ )> (∇2θ ln P (θ ∗ ))−1 (θ − θ ∗ ) + O(kθ − θ ∗ k3 ) =2= ln P (θ).

Notice that θ ∗ is in a local maxima, which ensures that the hessian ∇2θ ln P (θ ∗ ) is negative semidefinite. This in turn ensure that the normalizing constant ZQ exists and that LAPLACE(θ ∗ ; P ) is
well defined.
B.2

If θ ∗ is not a local maxima

In order to proceed to a similar derivation when ∇θ ln P (θ ∗ ) 6= 0, we first rearrange the terms in
the Taylor expansion. For a more compact notation, we write ∇ ln P instead of ∇θ ln P (θ ∗ ) and∇2 ln P instead of ∇2θ ln P (θ ∗ ).

1
ln P (θ) ∼= ln P (θ ∗ ) + ∇ ln P (θ − θ ∗ ) + (θ − θ ∗ )> ∇2 ln P (θ − θ ∗ ) =
2
1
∗
> 2
−1
= ln P (θ ) − ∇ ln P ∇ ln P ∇ ln P
2
1
+ (θ − θ ∗ + ∇2 ln P −1 ∇ ln P )> ∇2 ln P (θ − θ ∗ + ∇2 ln P −1 ∇ ln P ) =2
1
1= ln P (θ ∗ ) − ∇ ln P > ∇2 ln P −1 ∇ ln P + (θ − θ ∗1 )> ∇2 ln P (θ − θ ∗1 ),2
2
where we define the new point θ ∗1 as
Define the Gaussianθ ∗1 = θ ∗ − (∇2θ ln P (θ ∗ ))−1 ∇θ ln P (θ ∗ ).
∗which has density

LAPLACE(θ ; P ) := Nθ|µ = θ ∗1 , σ 2 = −(∇2θ ln P (θ ∗ )−1 ,

(14)
(15)

P (θ ∗ ) − 1 ∇ ln P > ∇2 ln P −1 ∇ ln P 1 (θ−θ∗1 )> (∇2θ ln P (θ∗ ))−1 (θ−θ∗1 )
e 2
e2
,
ZQ
p
> 2
−1
1
where ZQ = P (θ ∗ )e− 2 ∇ ln P ∇ ln P ∇ ln P (−2π)D det(∇2θ ln P (θ ∗ )) is the normalizing constant. Then, Q(θ) is a good approximation of P (θ) in the sense that
Q(θ) :=11
ln Q(θ) + ln ZQ = ln P (θ ∗ ) − ∇ ln P > ∇2 ln P −1 ∇ ln P + (θ − θ ∗1 )> ∇2 ln P (θ − θ ∗1 ) =
2
2
1
∗
∗
∗ > 2
= ln P (θ ) + ∇ ln P (θ − θ ) + (θ − θ ) ∇ ln P (θ − θ ∗ ) ∼=2
1∼= ln P (θ ∗ ) + ∇ ln P (θ − θ ∗ ) + (θ − θ ∗ )> ∇2 ln P (θ − θ ∗ ) + O(kθ − θ ∗ k3 ) =2= ln P (θ).

We highlight four points:(1) In order to ensure that the normalizing constant ZQ exists and consequently that LAPLACE(θ ∗ ; P )is well defined, the hessian ∇2θ ln P (θ ∗ ) must be negative semi-definite. Differently from the locallymaximal case, this is no longer guaranteed.
(2) The method is numerically unstable if ∇2θ ln P (θ ∗ ) has eigenvalues close to 0. This has been
empirically observed to commonly be the case (Sagun et al., 2016).
13

(3) We emphasize that the Taylor expansion is accurate around θ ∗ , but the Laplace Gaussian iscentered in θ ∗1 . This is often referred to as “mean shift” and implies that the sampled parameters fromthe normal distribution over weights are sampled far away from the actual mean.
(4) Now, assume ln P (θ) to be the composition of a function f (θ) and a loss l(f ). If, as in our
setting, f is linear and l is concave, then ln P (θ) is concave, its hessian is guaranteed to be negative
semi-definite, and the Laplace approximation is well-defined.

C

Extended related work

Our proposed method’s update rule resemble different methods in the literature. This is ratherinteresting since, despite arriving at a similar algorithm, they follow different derivation. In thissection, we seek to explain the nuances between existing methods and ours.
It is useful to recall our precision update rulet
Ht+1 = (1 − α)Ht + Jθ f·t (x)> Σ−1data Jθ f· (x)H0 = Σ−1prior .

(16)

To the best of our knowledge, a key conceptual difference with all related works in the literature
is Exact vs. Approximate Hessian: The second term on the RHS, J > Σ−1 J is identical in all the
update formulations, but comes from different derivation. In the formulation of Daxberger et al.
(2021) (and others) this is an approximation of the hessian, that happens to guarantee negative
definiteness. In our formulation, this term comes from the linearization. This term is the exact hessian,i.e. no approximation, such that negative definiteness is implied by the linearization. This results in amore intuitive understanding of the linearization error.
C.1

Differences with Laplace Redux

Daxberger et al. (2021) specifically develop a post-hoc method, which assume access to a MAP
parameter, nevertheless they also present an online training scheme. This appears to be grounded in
ideas from second order optimization, while ours is closer linked to the probabilistic model.
Neglecting their prior optimization procedure, Daxberger et al. (2021) considerst> −1t
Ht+1 = Σ−1prior + Jθ f· (x) Σdata Jθ f· (x).

(17)

We highlight two main differences.
Recursively Updating the Hessian vs. Using an Uninformed Prior: In our formulation the firstterm of the RHS is the previous precision, which we “discount” with a forgetting term α thatcomes from the error inherited by moving away from the previous linearization. In the formulationof Daxberger et al. (2021) this is a prior that is assumed to be in the form γ 2 I, where the scalarγ is optimized at every step through maximization of the evidence (Eq. 6 in their paper). Thisevidence maximization comes with some drawbacks: (1) It tends to make the Laplace approximation
overconfident to outliers. They partially address this issue by adding to the evidence an auxiliary termthat depends on an OOD dataset, penalizing it (Eq. 12 in their Appendix). (2) Besides inducing the
avoidable need for an OOD dataset, this technique have the same pitfall as VAEs, namely uncertainty
should be a derived quantity, not a learned one.
Computing the hessian at every iteration vs. every epoch: We update the hessian estimate every
iteration. In practice, Daxberger et al. (2021) updates the hessian every epoch. In principle, they
could update the hessian more regularly, but to the best of our knowledge, they do not explore this
path.
C.2

Differences with Variational Adaptive Newton Method

Khan et al. (2017) propose the following precision update rulet
Ht+1 = Ht + αJθ f·t (x)> Σ−1data Jθ f· (x).

(18)

The updates are very similar to ours, where the only difference is that the scaling is made on the
Jacobian product instead of on the previous precision Ht . Both their and our updates are “additive”
14

in the sense that the magnitude is increasing, for them strictly monotonically, for us on averagedepending on the magnitude of α. Thus, in both cases, the variance will approach 0 in the limit,
modelling epistemic uncertainty disappearing for infinitely long training.
Their update rule comes from the Variational Optimization setting, that can be viewed as an instanceof Variational Inference neglecting the KL term. They highlight strong similarities with Newton’s
method.
C.3

Differences with Noisy Natural GradientZhang et al. (2017) propose the precision update rulet
Ht+1 = (1 − α)Ht + αJθ f·t (x)> Σ−1(19)data Jθ f· (x).
Again, the update rule is very similar to ours, besides the scaling being applied to the Jacobian product.
They use a convex sum, which makes the update rule “norm preserving”. Thus, these updates havevery different asymptotic behaviour than ours.

Their update rule comes from applying, in the context of Variational Inference, natural gradient tothe variational posterior distribution, instead of directly on the parameters space. Natural gradient isdeeply connected with LAE and natural parameters highlight the importance of updating the precision
matrix instead of the covariance.
They also extend their derivation to the KFAC approximation of the hessian (in place of the exact
diagonal), this is made through the use of matrix variate Gaussian. Despite not considered in thiswork, a similar derivation is in principle feasible in our setting.
C.4

Connection with Adam

Both Zhang et al. (2017) and Khan et al. (2017) highlight strong similarities with the Adam
method (Kingma and Ba, 2015a). We share these similarities and we highlight them too. The“connection point” is a noisy version of Adam. Zhang et al. (2017) describe this method and call it“Noisy Adam” (Algorithm 1 in their paper). The difference to the vanilla version is that at each step,
instead of using the current parameter, they use a noisy version of it. The noise magnitude is thepseudo-second order term (vt in the original Adam paper (Kingma and Ba, 2015b)).
We can then interpret Noisy Adam as an instance of our variational setting, specifically where theexpectation estimate Eq(θ) [·] is made through Monte Carlo estimation with N = 1 samples. Having
the methods in the same setting, we can compare them and highlight the two main differences.
First, we emphasize the difference is in the second order term. Similarly to Zhang et al. (2017) andKhan et al. (2017), we use the diagonal of the hessian (technically the diagonal of the GGN), while
Adam uses the pointwise square of the gradient, which they call second raw moment and is intended
as a cheaply computable approximation of the hessian.
Another difference is that Adam applies the square root. This is a minor point since, as pointed out
by Zhang et al. (2017), this change may affect optimization performance, but does not change thefixed points.
C.5

Connection with Bayes by Backpropagation

Bayes by Backprop (Blundell et al., 2015) can be viewed as a sample-based approximation of Laplace.
In order to show this relation, we recall two very powerful equations (Opper and Archambeau, 2009a).
Specifically, let µ, Σ be the parameter of a Gaussian distribution q(θ) ∼ N (µ, Σ), and let V (θ) be
an arbitrary L2 integrable function (the log-likelihood in our case). Then, we are interested in thederivatives of Eθ∼q [V (θ)]. By standard Fourier analysis and integration by part, we have∇µ Eθ∼q [V (θ)] = Eθ∼q [∇θ V (θ)],(20)
1Eθ∼q [∇2θ V (θ)].
(21)
2While the first equation is somehow trivial, the second highlight a very deep relationship. The LHS
can be rewritten as∇Σ Eθ∼q [V (θ)] = ∇Σ E∼N (0,1) [V (µ + Σ)] = E∼N (0,1) [∇Σ V (µ + Σ)].
(22)∇Σ Eθ∼q [V (θ)] =15

We can recognize that inside the expectation on the RHS is exactly the update rule for the variance inthe bayes by backprop method (Blundell et al., 2015). Thus, we can interpret bayes by backprop as a
one-sample Monte Carlo estimation of that expectation value. The equations shows that in the limitof infinite samples, the bayes by backprop update step ∇Σ V (µ + Σ) converges to the (averaged)
Laplace approximation step Eθ∼q [∇2θ V (θ)] up to a factor 2. This motivates both the power and theinstability of bayes by backprop.

D

Model

D.1

Overview

Let X = RD be the data space, let Y = RD be the reconstruction space and let Θ = RP be
the parameter space. Let F = (Θ → (X → Y)) be the space of operators from Θ to the space of
operators from X to Y. We will denote a function f ∈ F applied to a parameter θ ∈ Θ as fθ : X 7→ Y.
This will represent, for example, a NN f with a specific set of parameter (weights) θ ∈ Θ, that maps
some data x ∈ X to some reconstruction fθ (x) = y ∈ Y.
Let X × Y × Θ × F be a probability space. The only assumption we make on this space isp(y|x, θ, f ) ∼ N (y|µ = fθ (x), σ 2 = Σ)∀(x, θ, f ) ∈ X × Θ × F(23)

where Σ ∈ M(RD × RD ) is a fixed variance matrix. This is a common assumption for regression
tasks, and is sometimes referred to as the “data noise” or “reconstruction error”. In this paper, we fix
Σ = I, but the derivations hold for the general case. With only this assumption, the distribution is
undefined and multiple solutions can exist. Thus, we require more assumptions.
A dataset D = {xn } is a finite of infinite collection of data xn ∈ X that is assumed to follow acertain, fixed but unknown, distribution(24)

xn ∼ p(x).

Sacrificing slim notation for the sake of clarity, we introduce an operator I : X 7→ Y. This represent
the ideal reconstruction for a given x. In the standard supervised setting this would be the operator(defined on the dataset only) that maps each data input to its label. In our unsupervised setting, where
X and Y are the same space, the operator I is simply the identity (Indeed they are not the same space,
they are isomorphic spaces that we identify through the operator I). Since I is the identity, it is often
neglected in literature, which can lead to unclear and potentially ambiguous Bayesian derivations.
Thus, we choose to adopt this heavier, but more precise notation.
We assume access to a specific f N N ∈ F. Practically this will be our NN architecture, i.e. an operator
that, given a set of parameter θ ∈ Θ give rise to a function from X to Y. Having f N N fixed, one may
consider f not to be stochastic anymore, we choose to still explicitly condition on f in order to havea clearer notation in later stages. Note that, despite not being covered in this work, a proper stochastic
derivation also on the NN architecture should be feasible.
D.2

Objective

The NN’s parameter optimization process in this full Bayesian probabilistic framework can be viewed
as: given a fixed f N N ∈ F, namely the NN architecture, maximise the reconstruction probability ofI(xn ) over the dataset D
hiXExn ∼p(x) p(y|xn , f N N ) y=I(x ) =p(y|xn , f N N ) y=I(x ) ,
(25)
n

n

xn ∈D

where the untractable p(y|xn , f N N ) can be expanded in θ and thus related to our hypothesis (23) asp(y|xn , f N N ) = Eθ∼p(θ|xn ,f N N ) p(y|xn , θ, f N N ) .
(26)

Notice that the only unfixed quantity is the distribution on the parameters, which we will optimize
for. We are not interested in finding a datapoint-dependant distribution, but rather one that maximise
all reconstructions at the same time, i.e. p(θ|f N N ) = p(θ|xn , f N N ). We can then frame Bayesian
16

optimization as: find a distribution on parameters such thatihXq(θ) ∈ arg max
Eθ∼p(θ|f N N ) p(y|xn , θ, f N N ) y=I(xn )p(θ|f N N )xn ∈D= arg maxp(θ|f N N )

X

xn ∈D



Eθ∼p(θ|f N N ) p(I(xn ) N (fθN N (xn ), Σ)) .

(27)
(28)

Moreover, finding this optima in the space ∆(Θ) of all distributions on Θ is not tractable. So, as
commonly done, we restrict ourselves to the subset G(Θ) ⊂ ∆(Θ) of Gaussians over Θ. Then, asolution in our context isXq(θ) ∈ arg max
Eθ∼q(θ) p(I(xn ) N (fθN N (xn ), Σ)) .
(29)
q∈G(Θ)xn ∈DWe emphasize that this solution has no guarantees of being unique, but we are interested in finding
one of them.
D.3

Joint distribution for a fixed datapoint

Let us first get a better understanding of the joint distribution on Y×Θ conditional to a fixed datapointx ∈ X and a network architecture f ∈ F
p(y, θ|x, f ).

(30)p(y|x, f )(31)p(θ|x, f )(32)p(y|θ, x, f )(33)p(θ|y, x, f ).

(34)

This distribution has two marginals

and two conditionals

These four quantities must satisfy the system of two “recursive” equationsZp(y|x, f ) =p(y|θ, x, f )p(θ|x, f )dθ =

(35)

Θ

= Eθ∼p(θ|x,f ) [p(y|θ, x, f )]p(θ|x, f ) =Zp(θ|y, x, f )p(y|x, f )dy =

(36)

Y

= Ey∼p(y|x,f ) [p(θ|y, x, f )].
If these are satisfied then the joint is a well-defined distribution and we can apply Bayes rulep(y|θ, x, f )p(θ|x, f ) = p(y, θ|x, f ) = p(θ|y, x, f )p(y|x, f )(37)which in logarithmic form islog p(y|θ, x, f ) + log p(θ|x, f ) = log p(θ|y, x, f ) + log p(y|x, f ).

(38)

We can factor in the assumptions. The “data noise” assumption gives us one of the two conditionals:
p(y|θ, x, f ) ∼ N (y|µ = fθ (x), σ 2 = Σ).

(39)p(θ|x, f ) = q t (θ) ∼ N (θ|µ = θ t , σ 2 = H−1
t ).

(40)

p(y|x, f ) = Eθ∼qt (θ) [p(y|N (fθ (x), Σ))](41)

The “Gaussian parameter” assumption gives us one of the two marginals:

With these in place, the joint distribution is uniquely defined. The other marginal, by Eq. 35, is17and the other conditional, by Bayes rule, isp(θ|y, x, f ) =p(y|θ, x, f )p(θ|x, f ).
p(y|x, f )(42)Despite being uniquely defined, the integral in Eq. 41 is, with a general f , intractable, and so is thejoint distribution.
But why do we even care? The intractability of p(y|x, f ) in Eq. 41 may at first glance appear
irrelevant. This is the case, for instance, with bayes by backprop (Blundell et al., 2015) methods.
They simply need access to the gradient of this quantity. For this purpose a simple Monte Carloestimate of the expectation is enough.
On the other hand, we are interested in recovering a meaningful distribution on parameters. This
imply that we aim at using Eq. 41 to enforce that Eq. 36 holds. For this purpose we need access tothe the density p(y|x, f ), so a Monte Carlo estimate of Eq. 41 is not enough.
D.4

Linear f

Theorem 1. Given the data noise assumption from Eq. 39p(y|θ, x, f ) ∼ N (y|µ = fθ (x), σ 2 = Σ),(43)given the Gaussian parameter assumption from Eq. 40 for some θ t , Htp(θ|x, f ) = q t (θ) ∼ N (θ|µ = θ t , σ 2 = H−1
t ),

(44)

assume that f is linear in θ, i.e.
fθ (x) = f0 (x) + Jf· (x)θ∀θ ∈ Θ, ∀x ∈ X.

(45)

Then the joint distribution is Gaussian itselfp(θ, y|x, f ) ∼ N ((θ, y)|µt , Σt )
where
µt =





H−1
θt
t
, and Σt =fθt (x)Jf· (x)Σ

(46)


ΣJf· (x)>.
−1
Jf· (x)> Ht Jf· (x)+ΣProof. With the further assumption of linearity of f , we can explicitly carry out the integral in the
expectation in Eq. 41p(y|x, f ) = Eθ∼qt (θ) [p(y|N (fθ (x), Σ))] =Z=
p(y|N (fθ (x), Σ))q t (θ)dθ =ΘZ=p(y|N (fθ (x), Σ))p(θ|N (θ t , H−1t ))dθ =ΘZ=p(y|N (f0 (x) + Jf· (x)θ, Σ))p(θ|N (θ t , H−1t ))dθ =Θ= p(y|N (fθt (x), (Jf· (x)> Ht Jf· (x))−1 + Σ)).

We emphasize that as a consequence ∇2θ log p(y|x, f ) is not dependent on θ t .
Having Theorem 1 in place, we can go back to our original problem. We need to deal with a fixed
non-linear architecture f N N . We can exploit the Theorem 1 by defining f t : a linearization of f N N
with a first order Taylor expansion around θ tfθt (x) := TAYLOR(f N N , θ t )(x) =and it holds that= fθNt N (x) + Jθ fθNt N (x)(θ − θ t )fθN N (x) = fθt (x) + O(kθ − θ t k2 ).
18

(47)
(48)

Recalling Eq. 23 both for f N N and for f tp(y|x, θ, f N N ) ∼ N (y|µ = fθN N (x), σ 2 = Σ)p(y|x, θ, f t ) ∼ N (y|µ = fθt (x), σ 2 = Σ)that, together with Eq. 48, imply
N y|µ = fθN N (x), σ 2 = Σ ∼ N y|µ = fθt (x) + O(kθ − θ t k2 ), σ 2 = Σ ,

(49)
(50)
(51)

where we can interpret the unknown O(kθ − θ t k2 ) as θ-dependent noise. More specifically, calling
γ > 0 the scalar constant of the O-term, we assume that
O(kθ − θ t k2 ) ≈ (θ − θ t )2and thus, from Eq. 51, we havewhere  ∼ N (0, γI)
N y|µ = fθN N (x), σ 2 = Σ ∼ N (y|µ = fθt (x), σ 2 = Σ + γkθ − θ t k2 I).

(52)
(53)

Despite at this point integrals are not analytically tractable, and thus a proper proof is not feasible,
the intuition is that this increased variance reflects in an increased variance in p(θ|x, f N N )∇2θ log p(θ|x, f N N ) θ=θt+1 ≈ ∇2θ log p(θ|x, f t ) θ=θt+1 + γkθ t+1 − θ t k2 ,(54)

where we introduce an hyperparameter α > 0 to cope with this added variance. If we then assume
the Jacobian Jθ f N N to be a Lipschitz function, then the Lipschitz constant is an upper bound on γ,
as follows from the Taylor expansion of Eq. (47). If this Lipschitz constant is smaller than the inverseof the gradient step 1/kθ t+1 − θ t k (which is not an unreasonable assumption for the gradient ascent
to be stable) we haveγkθ t+1 − θ t k2 ≈ kθ t+1 − θ t k(55)
that gives us a plausible order of magnitude for choosing the hyperparameter α.
Motivation: During training, we produce a sequence of Gaussians q t (θ) ∼ N (θ t , H−1t ) that weassume to be the distribution p(θ|x, f t ), at every step t ≥ 0. This distribution q t is then used for (1)Gaussian derivation in the linear case, for (2) Monte Carlo sampling in the update rule θ t → θ t+1
and (3), as second order derivative, for update rule Ht → Ht+1 .

Moreover, given that this distribution q t is our "best guess so far", we assume it to be also the
distribution p(θ|x, f N N ). This, being f N N not linear, (1) cannot be used for Gaussian derivation,(2) can reasonably be used for sampling (and thus we derive the improved update rule Eq. (64)), and(3) can somehow be used as second order derivative (and thus we derive the improved update rule
Eq. (65)), but the latter requires some more care. That is why we introduce the parameter α.
D.5

Iterative learning

Our learning method produces a sequence of Gaussians
q t (θ) ∼ N (θ|µ = θ t , σ 2 = H−1
t )

(56)

f t = TAYLOR(f N N , θ t )(57)

and a sequence of linearized functions
for every t ≥ 0.

Initialization is trivially done using a Gaussian prior on the parameters.
θ 0 = θ prior

H0 = (Σprior )−1(58)

Iterative step is made in two steps. First, having access to θ t , we “generate” the linearization f t .
Practically this is equivalent to computing the two quantities fθNt N (x) and Jθ fθNt N (x), that, together
with the value θ t are actually equivalent to “generating” f t , as Eq. 47 shows.
Second, we compute the Gaussian parameters θ t+1 and Ht+1 .
19

Recalling our aim of maximizing the quantity Eq. (29), update on q(·) means, θ t → θ t+1 , is ideallymade through gradient ascent steps on p(y|x, f )|y=I(x) . As this is intractable, we instead do gradientsteps on the lower bound Ly of (the log of) Eq. 35
θ t+1 = θ t + λ∇θ Ly
where
and so(59)

y=I(x)

Ly = Eθ∼p(θ|x,f t ) [log p(y|θ, x, f t )] ≤ log Eθ∼p(θ|x,f t ) [p(y|θ, x, f t )]

(60)

h
i
∇θ Ly y=I(x) = Eθ∼p(θ|x,f t ) ∇θ log p(y|θ, x, f t ) y=I(x) == Eθ∼qt (θ) ∇θ log p(I(x) N (fθt (x), Σ)) .

(61)

Recalling the Laplace approximation Eq. (15), the negative precision, −Ht+1 , is ideally set to be thethe hessian of the log probability p(θ|x, f ), evaluated in θ t+1 . As this is intractable we instead set itto the hessian of the lower bound Lθ of (the log of) Eq. (36)
Ht+1 = −∇2θ Lθ
whereand so∇2θ Lθ

(62)

θ=θ t+1

Lθ = Ey∼p(y|x,f t ) [log p(θ|y, x, f t )] ≤ log Ey∼p(y|x,f t ) [p(θ|y, x, f t )]θ=θ t+1= Ey∼p(y|x,f t ) ∇2θ log p(θ|y, x, f t )θ=θ t+1

via Eq. (38)h
= Ey∼p(y|x,f t ) ∇2θ log p(y|θ, x, f t )− ∇2θ log p(y|x, f t )θ=θ t+1

via Theorem 1= Ey∼p(y|x,f t ) ∇2θ log p(y|θ, x, f t )θ=θ t+1+ ∇2θ log p(θ|x, f t )θ=θ t+1+ ∇2θ log p(θ|x, f t )

i

(63)

θ=θ t+1+

θ=θ t+1



via the chain rule Eq. (67)= Ey∼p(y|x,f t ) Jθ f·t (x)> ∇2y log p(y|θ t+1 , x, f t )Jθ f·t (x) + ∇2θ log p(θ|x, f t )via HP Eq. (23)= Ey∼p(y|x,f t )−Jθ f·t (x)> Σ−1 Jθ f·t (x) + ∇2θ log p(θ|x, f t )= −Jθ f·t (x)> Σ−1 Jθ f·t (x) + ∇2θ log p(θ|x, f t )= −Jθ f·t (x)> Σ−1 Jθ f·t (x) + ∇2θ log q t (θ)θ=θ t+1

θ=θ t+1θ=θ t+1

θ=θ t+1= −Jθ f·t (x)> Σ−1 Jθ f·t (x) − Ht .
D.5.1

Improved update rule

As said, the update θ t → θ t+1 is ideally made through gradient ascent steps on p(y|x, f )|y=I(x) ,but we instead use the tractable lower bound with f t . We can perform the same derivation using f N N
in place of f t . Assuming p(θ|x, f N N ) ∼ q t (θ) for sampling, leads to the improved update rule
θ t+1 = θ t + λEθ∼qt (θ) ∇θ log p(I(x) N (fθN N (x), Σ)) .
(64)

Similarly, the negative precision −Ht+1 is ideally set to be the hessian of the log probability p(θ|x, f ).
but instead we use the tractable lower bound with f t . Here we cannot perform the same derivation,
20

since Theorem 1 does not hold anymore. Instead, we can rely on the estimate Eq. (54) to improve theterm∇2θ log p(θ|x, f t )θ=θ t+1= −Ht∇2θ log p(θ|x, f N N )−→θ=θ t+1≈ −(1 − α)Ht ,and this leads to the improved update ruleHt+1 = (1 − α)Ht + Jθ f·t (x)> Σ−1 Jθ f·t (x).

E

(65)

Fast Hessian

We are interested in computing the hessian of a loss function. For this purpose, the Jacobian of the
NN w.r.t. parameters plays a crucial role. In this section we develop a better understanding of thisobject, we derive the backpropagation (also used by the BackPack library (Dangel et al., 2020)) andfinally we explain our approximated backpropagation that allows linear scaling.
E.1

Jacobian of a Neural Network

Let us first define some terminology that we will need for chain rule derivations. A NN is a compositionof l functions f := fLl ◦ fLl−1 ◦ . . . ◦ fL2 ◦ fL1 :x0 −−−−−→ x1 −→
fL1

...

−→ xi−1 −−−−−→ xi −→
fLi

...

−→ xl−1 −−−−−→ xl
fLl

where there are parametric and non-parametric function fLi . We here highlight two common
parametric functions and a common non-parametric function.
Parametric function such as a linear layer
xi = fLi (xi−1 ) = φLi xi−1
or convolution
xi = fLi (xi−1 ) = convfeat=φLi (xi−1 )where φLi ∈ M(|xi |, |xi−1 |)where φLi ∈ M(in channel, out channel, feat height, feat width)

Non-parametric such as activation functions Li = tanh, ReLU . . .
xi = fLi (xi−1 )where |xi | = |xi−1 |What is conventionally called layer is actually a composition of two function: a linear function andan activation function. For sake of clarity in our derivation we do not adopt this convention and we
use the word “layer” to indicate singular “function” component of the NN.
Let us now consider a NN with w parametric layers and l − w activation layers
xl = fφ (x0 ) = fLl ◦ · · · ◦ fL1 (x0 )where φ = (φ1 , . . . , φw ),and define the bijection
W:{1, . . . , w}p−→ {i| s.t. Li is parametric layer}
7−→
i s.t. Li has parameters φp⊆ {1, . . . , l}from the subset of parametric layers to the corresponding index in φ = (φ1 , . . . , φw ).
The Jacobian w.r.t. the parameters Jφ fφ (x0 ) ∈ M(|xl |, |φ1 | + · · · + |φw |) is a matrix with numberof-output |xl | rows and number-of-parameters |φ| columns. For reference, just storing this matrix can
exceed memory limits even with the smallest autoencoder working on MNIST.
There are two ways of looking at this matrix: (1) row by row, that is output by output or (2)block-of-columns by block-of-columns, that is layer by layer.
E.2

Jacobian per output

Each row of the Jacobian correspond to the gradient w.r.t. the parameters of an element of the output
∇φ [fφ (x0 )]1
..
Jφ fφ (x0 ) = .
∇φ [fφ (x0 )]|xl |21

This can be computed by defining |xl | loss functions
lossk (xl ) := [xl ]k

for k = 1, . . . , |xl |and backpropagating each of those to obtain one line at a time. The disadvantange with thisformulation is that we cannot reuse computation for one loss function to improve computation ofother loss functions (they are independent). Moreover, we need to store all these rows at the sametime in order to compute J > J, which is computationally inpractical.
E.3

Jacobian per layer

Each column of the Jacobian is the derivative of the output vector w.r.t. a single parameter. We canthen group the parameters (i.e. columns) layer by layer!
Jφ1 fφ (x0 )Jφ fφ (x0 ) =...

Jφw fφ (x0 )where Jφp fφ (x0 ) ∈ M(|xl |, |φp |). Let us focus on the computation of Jφp fφ (x0 ) for a fixed layerp = 1, . . . , w. First notice that the parameters φp in fφ = fLl ◦ . . . ◦ fL1 only appear in fLW(p) andso∂fLi(xi−1 ) = 0 if i 6= W(p).
∂φpChain rule (informal)∂xl∂fφ (x0 )
=
=
∂φp
∂φp
∂xl ∂xl−1==∂xl−1 ∂φp∂xl ∂xl−1=...
∂xl−1 ∂xl−2| {z } | {z }layer l layer l − 1
w.r.t. input w.r.t. input

Chain rule (formal)

Jφp fφ (x0 ) = 

k=lY

W(p)+1∂xW(p)+1∂xW(p)| {z }∂xW(p)∂φp| {z }

layer W(p) + 1 layer W(p)w.r.t. parameters
w.r.t. input

Jxk−1 fLk (xk−1 ) Jφp fLW(p) (xW(p)−1 )(66)

The intuition for the chain rule is that the Jacobian Jφp fφ (x0 ) is the composition of the Jacobians
w.r.t. input of subsequent layers times the Jacobian w.r.t. parameters of the specific layer. Thus, we
can reuse computation for one layer to improve computation of other layer, specifically the productof Jacobians w.r.t. input. Moreover, we can compute Jp> Jp layer by layer without ever storing thefull Jacobian.
E.3.1

Jacobian of a layer w.r.t. to input

The Jacobian of a standard linear layer w.r.t. to the input is
Jxp−1 fφp (xp−1 ) = φpand this remains the same also in the case with a bias. The Jacobian of a convolutional layer w.r.t.
to the input is
Jxp−1 convfeat=φp (xp−1 ) = M(convfeat=φp )The Jacobian of the activation function depends on the specific choice. Recall that, for each layer i,xi ∈ R|xi | is a vectorxi = [xi ]k k=1,...,|x |iwhere [xi ]k ∈ R is the value in position k of the vector xi .
22

If Li is tanh[xi ]k = [fLi (xi−1 )]k = tanh([xi−1 ]k ) =e[xi−1 ]k − e−[xi−1 ]ke[xi−1 ]k + e−[xi−1 ]k

for k = 1, . . . , |xi−1 |

then



[Jxi−1 fLi (xi−1 )]kj = δkj 1 − (tanh([xi−1 ]k ))2 = δkj 1 − [xi ]2k

for k, j = 1, . . . , |xi−1 |.

If Li is ReLU

for k = 1, . . . , |xi−1 |[xi ]k = [fLi (xi−1 )]k = ReLU([xi−1 ]k ) = max(0, [xi−1 ]k )
then
E.3.2

[Jxi−1 fLi (xi−1 )]kj = δkj (1 if [xi−1 ]k > 0 else 0)

for k, j = 1, . . . , |xi−1 |.

Jacobian of a layer w.r.t. to parameters

The Jacobian of a standard linear layer w.r.t. to the parameters isJφi fφi (xi−1 ) = I|xi | ⊗ xi−1

and in the case with bias bi ∈ R

xi

the Jacobian isJφi ,bi fφi ,bi (xi−1 ) = I|xi | ⊗ [xi−1 , 1]

∈ M(|xi |, |φi |)∈ M(|xi |, |φi | + |bi |)The Jacobian of a convolutional layer w.r.t. to the parameters isJφi convfeat=φi (xi−1 ) = Jφi convTfeat=rev(xi−1 ) (φi ) = M(convfeat=rev(xi−1 ) )Tand in the case with bias bi ∈ Roi the Jacobian isJφi convfeat=φi ,bias=bi (xi−1 ) = M(convfeat=rev(xi−1 ) )T I

E.4

Hessian of a Neural Network∈ M(|xi |, |φi |)∈ M(|xi |, |φi | + |bi |)Consider a function L : R|xl | → R from the output of the NN to scalar value. This later will beinterpreted as loss or likelihood, but for now let us stick to the general case.
We are interested in the hessian of this scalar value w.r.t. the parameters of the NN
!
w
w


X
X
∇2φ L(fφ (x0 )) ∈ M|φi |,|φi | .
p=1

p=1

Similarly to the previous section, it is convenient to see this matrix as block matrices, separated layer
wise


∂2
∂2
∇2φ1 L(fφ (x0 ))...
∂φ1 ∂φ2 L(fφ (x0 ))∂φ1 ∂φw L(fφ (x0 ))2 ∂φ∂∂φ L(fφ (x0 ))
∇2φ2 L(fφ (x0 ))
2
1




2


∇φ L(fφ (x0 )) = 
..

..


.
.




2
∂
2
∇φw L(fφ (x0 )∂φw ∂φ1 L(fφ (x0 ))First common assumption is to consider layers to be independent to each other, i.e.
∂2
L(fφ (x0 )) = 0∂φi ∂φj∀i 6= jLet us now fix a layer p = 1, . . . , w and focus on a single diagonal block.
∇2φp L(fφ (x0 )) ∈ M(|φp |, |φp |).
23

According to the chain rule|xl |X
∇2φp L(fφ (x0 )) = Jφp fφ (x0 )T · ∇2xl L(xl ) · Jφp fφ (x0 ) +[∇xl L(xl )]i · ∇2φp [fφ (x0 )]i
|
{z
} i=1

(67)

=:G(φ)The second term of the RHS is equal to 0 if the model perfectly fits the dataset, ∇xl L(xl ) = 0, OR if
f is linear in the parameters, Hφp [fφ (x0 )]i = 0.
The first term of the RHS of Eq. 67, G(φ), is in literature referred to as Generalized Gauss Newton(GGN) matrix. It can be computed efficiently thanks to the view of the Jacobian as layer by layer.
Using equation (66), the expression for the approximated hessian w.r.t. to φp is thenG(φ) = Jφp fφ (x0 )T · Hxl L(xl ) · Jφp fφ (x0 ) =




l
k=l
Y
Y
= Jφp fLTW(p) 
Jxk fLTk  Hxl L(xl ) 
Jxk fLk  Jφp fLW(p)W(p)+1k=W(p)+1

and from this we can build an efficient backpropagation algorithm.
Algorithm 1 Algorithm for Jφ f T · ∇2 L · Jφ f

M = ∇2xl L(xl )for k = l, l − 1, . . . , 1 do
if Lk is parametric with φp (i.e. k = W(p)): thenHp = Jφp fL>k · M · Jφp fLk
end if
M = Jxk fL>k · M · Jxk fLk
end forreturn (H1 , . . . , Hw )As it is written, each Hp is a matrix |φp |×|φp | so we technically obtain the hessian in a block-diagonal
form


H1
0


..


.
0
Hw

if we are interested in the diagonal only, we can construct that by concatenation of the diagonals foreach Hp .
E.4.1

Fast approximated version of the Algorithm

The idea is to backpropagate only the diagonal of the matrix M , neglecting all the non-diagonal
elements. In a single backpropagation step we have
M 0 = Jxp fL>p · M · Jxp fLp
In order to backpropagate the diagonal only, we need to use the operator
diag(M ) 7→ diag(M 0 )For linear layers and activation function this operator is trivial. For convolutional layer it turns outthat this operator is itself a convolution
diag(M 0 ) = convfeat=φ(2) (diag(M ))p(2)where the kernel tensor φp is the pointwise square of the kernel tensor φp .
24

E.5

Hessian of a Reconstruction Loss

In the previous Algorithm, the backpropagated quantity is initialized as the hessian of the loss w.r.t.
the output of the NN ∇2xl L(xl ), or, in an equivalent but more compact notation ∇2f L(f ). The valueof this hessian clearly depends on the specific choice of loss function L.

The most common choice of likelihood for regression is the Gaussian distribution, while for classification it is the Bernoulli distribution. The Gaussian log likelihood is√
1
L(f ) := log p(x|µ = fθ (x), σ 2 = σd2 ) = − 2 kx − fθ (x)k2 − log( 2πσd )(68)
2σdand its hessian is identity scaled with σd∇2f log p(x|µ = fθ (x), σ 2 = σd2 ) = −(σd )−2 I(69)

The bernoulli log likelihood is
L(f ) := log p(c|fθ (x)) = log[softmax(fθ (x))]c = [fθ (x)]c − logXie[fθ (x)]i

!

(70)

and its hessian can be written in terms of the vector π = softmax(fθ (x)) of predicted probabilities!
X∇2f log p(c|fθ (x)) = −∇2f loge[fθ (x)]i = diag(π) − ππ T
(71)i

We highlight that both hessians are independent on the label, and thus the GGN is equal to the Fisher
matrix, this is true every time p(y|fθ (x)) is an exponential family distribution with natural parametersfθ (x). In this paper, we focus on the Gaussian log likelihood, but we emphasize that the method is
not limited to this distribution.

25

References
C. Blundell, J. Cornebise, K. Kavukcuoglu, and D. Wierstra. Weight uncertainty in neural networks. In
Proceedings of the 32nd International Conference on International Conference on Machine Learning Volume 37, ICML’15, page 1613–1622. JMLR.org, 2015.
A. Botev, H. Ritter, and D. Barber. Practical Gauss-Newton optimisation for deep learning. In International
Conference on Machine Learning (ICML), pages 557–565. PMLR, 2017.
O. Cappé. Online sequential monte carlo em algorithm. In 2009 IEEE/SP 15th Workshop on Statistical Signal
Processing, pages 37–40. IEEE, 2009.
F. Dangel, F. Kunstner, and P. Hennig. BackPACK: Packing more into backprop. In International Conference on
Learning Representations (ICLR), 2020.
E. Daxberger and J. M. Hernández-Lobato. Bayesian variational autoencoders for unsupervised out-ofdistribution detection. arXiv preprint arXiv:1912.05651, 2019.
E. Daxberger, A. Kristiadi, A. Immer, R. Eschenhagen, M. Bauer, and P. Hennig. Laplace Redux - Effortless
Bayesian deep learning. Advances in Neural Information Processing Systems (NeurIPS), 34, 2021.
N.

S. Detlefsen, A. Pouplin, C. W. Feldager, C. Geng,
M. G. Duque, F. Warburg, M. Miani, and S. Hauberg.
https://github.com/MachineLearningLifeScience/stochman/, 2021.

D. Kalatzis,
Stochman.

H. Hauschultz,
GitHub. Note:

Y. Gal and Z. Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deeplearning. In International Conference on Machine Learning (ICML), volume 48, pages 1050–1059, New
York, New York, USA, 20–22 Jun 2016. PMLR.
A. Gelb et al. Applied optimal estimation. MIT press, 1974.
G. E. Hinton and R. R. Salakhutdinov. Reducing the dimensionality of data with neural networks. Science, 313(5786):504–507, 2006.
A. Immer, M. Bauer, V. Fortuin, G. Rätsch, and K. M. Emtiyaz. Scalable marginal likelihood estimation
for model selection in deep learning. In International Conference on Machine Learning (ICML), pages
4563–4573. PMLR, 2021.
A. Immer, T. F. van der Ouderaa, V. Fortuin, G. Rätsch, and M. van der Wilk. Invariance learning in deep neuralnetworks with differentiable Laplace approximations. arXiv preprint arXiv:2202.10638, 2022.
M. I. Jordan, Z. Ghahramani, T. S. Jaakkola, and L. K. Saul. An introduction to variational methods for graphicalmodels. Machine Learning, 37(2):183–233, 1999.
L. V. Jospin, W. Buntine, F. Boussaid, H. Laga, and M. Bennamoun. Hands-on bayesian neural networks – atutorial for deep learning users, 2020. URL https://arxiv.org/abs/2007.06823.
M. E. Khan, W. Lin, V. Tangkaratt, Z. Liu, and D. Nielsen. Variational adaptive-newton method for explorative
learning, 2017. URL https://arxiv.org/abs/1711.05560.
D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In Y. Bengio and Y. LeCun, editors, 3rdInternational Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015,
Conference Track Proceedings, 2015a. URL http://arxiv.org/abs/1412.6980.
D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In International Conference on Learning
Representations (ICLR), 2015b.
D. P. Kingma and M. Welling. Auto-encoding variational Bayes. In International Conference on Learning
Representations (ICLR), 2014.
B. Lakshminarayanan, A. Pritzel, and C. Blundell. Simple and scalable predictive uncertainty estimation using
deep ensembles. In Advances in Neural Information Processing Systems (NIPS), volume 30, 2017.
P. S. Laplace. Mémoire sur la probabilité des causes par les événements. Mémoire de l’Académie Royale desSciences, 1774.
Y. Lecun, L. Bottou, Y. Bengio, and P. Haffner. Gradient-based learning applied to document recognition.
Proceedings of the IEEE, 86(11):2278–2324, 1998. doi: 10.1109/5.726791.

26

Z. Liu, P. Luo, X. Wang, and X. Tang. Deep learning face attributes in the wild. In Proceedings of International
Conference on Computer Vision (ICCV), December 2015.
D. J. C. MacKay. Bayesian interpolation. Neural Computation, 4(3):415–447, 05 1992.
D. J. C. MacKay. Probable networks and plausible predictions: A review of practical Bayesian methods forsupervised neural networks. Network: Computation in Neural Systems, 6(3):469, 1995.
W. J. Maddox, T. Garipov, P. Izmailov, D. P. Vetrov, and A. G. Wilson. A simple baseline for bayesian uncertainty
in deep learning. Neurips, abs/1902.02476, 2019.
J. Martens and R. Grosse. Optimizing neural networks with Kronecker-factored approximate curvature. In
International Conference on Machine Learning (ICML), pages 2408–2417. PMLR, 2015a.
J. Martens and R. B. Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In
International Conference on Machine Learning (ICML), 2015b.
E. Nalisnick, A. Matsukawa, Y. W. Teh, and B. Lakshminarayanan. Detecting out-of-distribution inputs to deep
generative models using typicality, 2019a. URL https://arxiv.org/abs/1906.02994.
E. T. Nalisnick, A. Matsukawa, Y. W. Teh, D. Görür, and B. Lakshminarayanan. Do deep generative models
know what they don’t know? In 7th International Conference on Learning Representations, ICLR 2019, New
Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019b. URL https://openreview.net/forum?id=
H1xwNhCcYm.
R. M. Neal. Bayesian leaning for neural networks. University of Toronto, 1996.
M. Opper and C. Archambeau. The Variational Gaussian Approximation Revisited. Neural Computation, 21(3):786–792, 03 2009a. ISSN 0899-7667. doi: 10.1162/neco.2008.08-07-592. URL https://doi.org/10.
1162/neco.2008.08-07-592.
M. Opper and C. Archambeau. The variational Gaussian approximation revisited. Neural Computation, 21(3):786–792, 2009b.
Y. Park, C. Kim, and G. Kim. Variational Laplace autoencoders. In International Conference on Machine
Learning (ICML), pages 5032–5041. PMLR, 2019.
A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga,
A. Desmaison, A. Kopf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang,
J. Bai, and S. Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in
Neural Information Processing Systems (NeurIPS), pages 8024–8035. 2019.
D. J. Rezende, S. Mohamed, and D. Wierstra. Stochastic backpropagation and approximate inference in deep
generative models. In International Conference on Machine Learning (ICML), volume 32 of Proceedings of
Machine Learning Research, pages 1278–1286. PMLR, 22–24 Jun 2014.
H. Ritter, A. Botev, and D. Barber. A scalable Laplace approximation for neural networks. In International
Conference on Learning Representations (ICLR), volume 6, 2018.
D. Rumelhart, G. Hinton, and R. Williams. Learning internal representations by error propagation. In Parallel
Distributed Processing: Explorations in the Microstructure of Cognition, vol. 1: Foundations, pages 318–362.
1986.
L. Sagun, L. Bottou, and Y. LeCun. Eigenvalues of the hessian in deep learning: Singularity and beyond. arXiv
preprint arXiv:1611.07476, 2016.
F. Wenzel, K. Roth, B. Veeling, J. Swiatkowski, L. Tran, S. Mandt, J. Snoek, T. Salimans, R. Jenatton, and
S. Nowozin. How good is the Bayes posterior in deep neural networks really? In International Conference on
Machine Learning (ICML), pages 10248–10259, 2020.
H. Xiao, K. Rasul, and R. Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learningalgorithms, 2017.
G. Zhang, S. Sun, D. Duvenaud, and R. B. Grosse. Noisy natural gradient as variational inference. CoRR, 2017.
URL http://arxiv.org/abs/1712.02390.
M. Zhang, P. Hayes, and D. Barber. Generalization gap in amortized inference. Workshop on Bayesian Deep
Learning @ NeurIPS, 2021.

27