Laplacian Autoencoders for Learning Stochastic Representations
Abstract
Established methods for unsupervised representation learning such as variational autoencoders produce none or poorly calibrated uncertainty estimates making it difﬁcult 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.
This is maximized using Monte Carlo EM with a variational distribution that takes the shape of a Laplace approximation.
We develop a new Hessian approximation that 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 in improved performance across a multitude of downstream tasks.
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 which data is approximately reconstructed.
The nonlinearity of the model implies that sometimes small changes to data x give 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 quantiﬁcation of its uncertainty and we cannot easily answer when the representation is reliable and accurately reﬂects 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
But is the VAE ’s quantiﬁcation of reliability reliable?
To investigate, we ﬁt a VAE with a twodimensional latent representation to the MNIST dataset (Lecun et al., 1998), and illustrate the predicted uncertainty 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, but 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 governed by 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.
[Table]
+vRerFeh60p62tmGf2C9fYJwOGZXg==</latexit> (b) LAE (post-hoc) <latexit sha1_base64="I3lyBSIdfE5Wl35PdctE+++DNmc=">AAACAHicbVDLSgMxFM3Ud31VXbhwE1qEFqHMiI9xVxXBhQsF+4C2lEya1mAmMyR3xGHoxm/wD9y4UMStn+HOvzHTiqj1wIXDOfcm9x4vFFyDbX9YmYnJqemZ2bns/MLi0nJuZbWmg0hRVqWBCFTDI5oJLlkVOAjWCBUjvidY3bs+Tv36DVOaB/IS4pC1fdKXvMcpASN1custYLeQFGkJnx2e4GIg05dKg06uYJcPDnbcXRcb4u65zjZ2yvYQ36RQybe27j8q8Xkn997qBjTymQQqiNZNxw6hnRAFnAo2yLYizUJCr0mfNQ2VxGe6nQwPGOBNo3RxL1CmJOCh+nMiIb7Wse+ZTp/Alf7rpeJ/XjOCnttOuAwjYJKOPupFAkOA0zRwlytGQcSGEKq42RXTK6IIBZNZ1oQwdvI4qW2Xnd2yfWHSOEIjzKINlEdF5KB9VEGn6BxVEUUD9ICe0LN1Zz1aL9brqDVjfc2soV+w3j4BaeuYow==</latexit> (c) LAE (online)
From a Bayesian perspective the natural solution is to form an (approximate) posterior over the weights of the neural networks.
To investigate, we adapt a state-of-the-art implementation of a post-hoc Laplace approximation (Daxberger et al., 2021) of the weight posterior to the autoencoder domain.
This amounts to training a regular autoencoder, and thereafter approximating the weight uncertainty with the Hessian of the loss (Sec. 1.1).
Fig.
1b shows that uncertainty now grows, as intuitively expected, with the distance to the 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 a variational distribution on the neural network weights.
We ﬁrst propose a post-hoc LA for autoencoders; showcasing good out-of-distribution detection capabilities, but lack of properly calibrated uncertainties in-distribution.
To address this, we develop a fast and memory-efﬁcient Hessian approximation, which allows us to maximize a variational lower bound using Monte Carlo EM, such that model uncertainty is a key part of model training rather than estimated post-hoc.
Fig.
1c gives an example of the corresponding uncertainty, which exhibits a natural and semantically meaningful behavior.
1.1 Background
The VAE is a latent variable model that parametrize the data density p ( x ) =
[Equation]
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 renders the marginalization of z intractable, and a variational lower bound of p ( x ) is considered instead.
To arrive at this, one ﬁrst introduces a variational approximation q ( z | x ) ⇡ p ( z | x ) and write p ( x ) =
E q ( z | x ) [ p ( x | z ) p ( z ) / q
( z | x ) ] .
A lower bound on p ( x ) then follows by a direct application of Jensen’s inequality,
[Equation]
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 (Sec. 12).
Yet this is not the observed empirical behavior in Fig.
1a.
Since the  2 network is left untrained away from training data, its predictions depend on extrapolation.
In practice,  2 takes fairly small values near training data (assuming the mean µ provides a reasonable data ﬁt), and  2 extrapolates arbitrary even if this is suboptimal in terms of data likelihood.
Similar remarks hold for other likelihood 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 proposed solution.
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 ✓ ⇤
[Equation]
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 ﬁrst order term vanishes and the second order term is negative semi-deﬁnite.
We provide more details on the Laplace approximation in Sec. 8.
In practice, computing the full Hessian is too taxing both in terms of computation and memory, and various approximations are applied (Sec. 3).
Laplacian Autoencoders
We consider unsupervised representation learning from i.i.d.
data D = { x n } Nn =1 consisting of observations x n 2 R D .
We also deﬁne a continuous latent space such that representations z n 2 R K .
Similar to AE s (Hinton and Salakhutdinov, 2006), we consider two neural networks e : R D !
R K and d  : R K ! R D , 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 x rec =
f ✓ ( x ) given observation x , where the uncertainty comes from ✓ being unknown, p ( x rec | x , f
[Equation]
Notice that we explicitly condition on f , which is the operator ✓ 7! f ✓ , even if this is not stochastic; this conditioning will become important later on to distinguish between the distribution deduced by f and its linearization f ( t ) .
Mimicking the standard autoencoder reconstruction loss, we set p ( x rec | ✓ , x , f )=
N ( x rec | 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 to a Gaussian approximation, but rather than learning the variance variationally, we opt for LA .
This will allow us to sample NN s and deduce uncertainties in both latent and output space as illustrated in Fig.
Intractable joint distribution.
Any meaningful approximate posterior over ✓ should be similar to the marginal of the joint distribution p ( ✓ , x rec | x , f ) .
This marginal is
[Equation]
which can be bounded on a log-scale using Jensen’s inequality,
[Equation]
Our ﬁrst approximation is a LA of p ( ✓ | x , f ) ⇡ q t ( ✓ | x , f ) =
N ( ✓ | ✓ t , H  1 t ) , where we postpone the details on how to acquire ✓ t and H t .
These will eventually be iteratively computed from the 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 ( x rec | 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 r ✓ L = J ✓ f ✓ r f l ( f ) , where J ✓ f ✓ is the
[Table]
[Table]
[Table]
/4saFIm79C3f+jZM2C209MMzhnHu59x4/FVyBZX0bS8srq2vrlY3q5tb2zq65t99RSSYpa9NEJLLnE8UEj1kbOAjWSyUjkS9Y1x/dFH73gUnFk/gexilzIzKIecgpAS155qFzljt+IgI1jvSHHRgyIBMPPLNm1a0p8CKxS1JDJVqe+eUECc0iFgMVRKm+baXg5kQCp4JNqk6mWEroiAxYX9OYREy5+fSCCT7RSoDDROoXA56qvztyEqliQ10ZERiqea8Q
[Figure]
/P6GYRXbs7jNAMW09mgMBMYElzEgQMuGQUx1oRQyfWumA6JJBR0aFUdgj
[Figure]
Jacobian of f .
In a similar spirit, we linearize f in in order arrive tractable
[Equation]
where f ( t ) denote the associated ﬁrst-order approximation of f .
Under this approximation, the joint distribution p ( ✓ , x rec | x , f ) becomes Gaussian, p ( ✓ , x rec | x , f ( t ) )
⇡ N ( ✓ , x rec | µ t , ⌃ t ) , with µ t = ✓ ✓ t f ( t ) ✓ t ( x ) ◆ , and ⌃ t = ✓
[Equation]
This approximation is illustrated in Fig. 3(b).
We provide the proof in Sec. 10.4.
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 +  r ✓ L x rec , where
[Equation]
is a lower bound on Eq. 3, which we evaluate with a single Monte Carlo sample.
Following the LA , the covariance of q t +1 ( ✓ | x , f ) , 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)
[Equation]
The last term can be approximated since r 2 ✓ log q t ( ✓ | x , f [Equation] | ✓ = [-θ t +
-]✓ [Equation]
=
[Equation] + O ( k ✓ t +1  ✓ t [Equation]
To efﬁciently cope with the O -term, we introduce a parameter ↵ , such that the ﬁnal approximation is
[Equation]
where J ✓ f ( t ) ✓ is independent of which ✓ we evaluate in and r 2 x rec log p
x rec  ✓
[Equation]
, x , f ( t )  is trivial to compute for common losses, i.e. for MSE it is the identity.
The parameter ↵ can be viewed as a geometric running average that is useful for smoothing out results computed on a minibatch instead of on the full training set, similar to momentum-like training procedures.
It further allows for non-monotonically-increasing precision.
Note that we revisit data during training, and the precision
[Table]
[Figure]
matrix is updated for every revisit.
Thus, the forgetting induced by ↵ is a wanted behavior to avoid inﬁnite precision in the limit of inﬁnite training time.
In practice, we set ↵ equal to the learning rate, where the practical intuition is that when ↵ is small, the network uncertainty decreases faster.
The overall training procedure is summarized in Fig.
We initialize q 0 as a Gaussian with ✓ 0 =and H
= I .
We provide more details on the model, the linearization, and iterative learning in
Sec. 10.
Why not just. ..?
The proposed training procedure may at ﬁrst appear non-trivial, and it is reasonable 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 the supervised experiences of Jospin et al.
, we, unfortunately, found this approach too brittle to allow for practical model ﬁtting.
But then perhaps a post-hoc LA as proposed by Daxberger et al.
(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 .
Conceptually, we argue that our approach, while novel, is not entirely separate from existing methods.
Our reliance on lower bounds makes 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 by
Park et al.
Finally, we note that our linearization trick (6) has great similarities to classic extended Kalman ﬁltering (Gelb et al., 1974).
Scaling the Hessian to Large Images
The largest obstacle to apply LA in practice stems 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 of parameters.
The currently most efﬁcient Hessian implementations (Dangel et al., 2020;
Daxberger et al., 2021) builds on the generalized Gauss-Newton ( GGN ) approximation of the Hessian
[Equation]
for a single layer l , which neglects second order derivatives of f w.r.t.
the parameters.
Besides, the computational beneﬁts of this approximations, previous works on LA (Daxberger et al., 2021) relies on GGN to ensure that the Hessian is always semi-negative deﬁnite.
In contrast, the model presented in Sec. 2 implies that GGN is no longer a practical and unprincipled trick, but rather the exact Hessian (Immer et al., 2021b).
Albeit relying on ﬁrst order derivates, the layer-block-diagonal GGN , which assumes that layers are independent of each other, scales quadratically with the output dimension of the considered neural network f .
This lack of scaling is particularly detrimental for convolutional layers as these have low parameter counts, but potentially very high output dimensions.
Expanding J ✓ ( l ) f ✓ ( x ) with the chain rule, one realizes that the Jacobian can be computed as a function of the Jacobian of the next layer.
Fig.
6 illustrate that an intermediate quantity M , which is initialised as r 2 x rec L ( x rec ) , can be efﬁciently backpropagated through multiplication with the Jacobian
[Table] [Equation] [Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Table]
[Table]
[Table]
+ie8eFDEq3/Hm//GbZuDtj4YeLw3w8y8IBFcG9f9dgpr6xubW8Xt0s7u3v5B+fCopeNUMWyyWMSqHVCNgktsGm4EthOFNAoEPgaju5n/+IRK81g+mEmCfkQHkoecUWOldp10JY6J2ytX3Ko7B1klXk4qkKPRK391+zFLI5SGCap1x3MT42dUGc4ETkvdVGNC2YgOsGOppBFqP5vfOyVnVumTMFa2pCFz9fdERiOtJ1FgOyNqhnrZm4n/eZ3UhDd+xmWSGpRssShMBTExmT1P+lwhM2JiCWWK21sJG1JFmbERlWwI3vLLq6R1UfWuqu79ZaV2m8dRhBM4hXPw4BpqUIcGNIGBgGd4hTdn7Lw4787HorXg5DPH8AfO5w/KGo8m</latexit>
H  = w.r.t.
input of each layer.
This process leads to a block diagonal approximation of the Hessian as illustrated in Fig. 5(a).
However, diagonal blocks are generally too large to store and invert.
To combat this, each block can be further approximated by its exact diagonal (Daxberger et al., 2021) as depicted in Fig. 5(b).
This scales linearly w
[Figure]
r
[Figure]
t. parameters, but still scales quadratically w.r
[Figure]
t. the output resolution (Tab. 1).
[Figure]
[Table]
[Figure]
x for with extended backward pass in M is backpropagated to previous layers
Via the chain rule and M the Hessian of each layer can be computed efﬁciently
To scale our Laplacian autoencoder to high-dimensional 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. 5(c).
This assumes that features from the same layer are uncorrelated and consequently have 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 bottleneck architecture.
We note that the quadratic scaling of the exact diagonal Hessian is less of an issue in the layers near the bottleneck than in the layers closer to the output space.
We can therefore dynamically switch 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 Sec. 11.
Related Work
Deep generative models, and particularly the family of variational auto-encoders ( VAE s) (Kingma and Welling, 2014; Rezende et al., 2014), address unsupervised representation learning from a probabilistic viewpoint by approximating the posterior over the representation space.
Despite their widespread adoption, model parameters are still deterministic and sensitive to ill-suited local minima, e.g. over-ﬁtted to training data (Zhang et al., 2021), which may cause poor generalization.
The
Bayesian NN s favour inference over the NN weights for addressing such issues (MacKay, 1995;
Neal, 1996).
This approach deduces 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, 60 Input size T i m e [ s e c ]      60 Input size 
-]
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.
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 classiﬁcation and regression.
The general utility of LA has also motivated its use as an approximation to the marginal likelihood over NN weights.
Recent methods, including
Daxberger et al.
(2021), have explored this path to ﬁnd model hyperparameters (Immer et al., 2021a) 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.
(2021); Zhang et al.
; Khan et al.
(2017), Bayes by Backpropagation Blundell et al.
and
Adam Kingma and Ba (2015a) in Sec. 9.
The full Bayesian perspective on VAE weights was ﬁrst explored by Daxberger and Hernández-Lobato (2019) which we ﬁnd similar in spirit to our work.
In contrast to them (1) we follow a principled
Bayesian derivation.
Neither do we depend on Hamiltonian Monte-Carlo sampling, which is generally hard to scale to efﬁcient training.
Experiments
First, we demonstrate the computational advantages of the proposed Hessian approximation, and then 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 the following 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 code is implemented in PyTorch and available 2 .
Sec. 7 provides more details on the experimental setup.
Efﬁcient 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 sufﬁcient 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 function of input size for a 5-layer convolutional network that preserves channel and input dimension.
As baselines we use efﬁcient 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 ⇥image on a 11 Gb NVIDIA GeForce GTX 1080 Ti.
In contrast, our approximate diagonal Hessian scales linearly with the resolution, which is especially beneﬁcial for convolutional layers.
[Table] [Table]
/2cfz8blfaubdWJ8IJJ0JjhpRqem6sgxRJTTEjI7uVKBIjPEA90jQoECcqSCfzjuCecTqwG0nzhIYT92dHirhSQx6aSo50X/3NxuZ/WTPRXT9IqYgTTQSeftRNGNQRHC8PO1QSrNnQAMKSmlkh7iOJsDYnss0RZlaehVrB8YqOe+Hmy6dgqizYAbtgH3jgCJTBOaiAKsCAgVtwDx6sa+vOerSepqUZ66tnG/yS9fIJmS+TVg==</latexit
[Figure]
[Table]
Tab. 2 shows that the exact or approximate Hessian diagonal has similar performance for both post-hoc and online training.
Using post-hoc LA results in good mean reconstructions (low MSE), but each sampled NN does not give good reconstructions (low log p ( x ) ).
Using our online training procedure results in a much higher log-likelihood.
This indicates that every sampled NN predicts good reconstructions.
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 better ﬁtted, both in latent and data space.
These well-ﬁtted uncertainties have several 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 decisions in an automated system.
We evaluate O O D performance on the commonly used benchmarks (Nalisnick et al., 2019b), where we use F ASHION M NIST (Xiao et al., 2017) as in-distribution and MNIST
(Lecun et al., 1998) as O O D. Fig.
9 (c) shows that our online LAE outperforms existing models in both log-likelihood and Typicality score (Nalisnick et al., 2019a).
This stems from the calibrated model uncertainties, which are exempliﬁed in the models ability to detect O O D examples from the uncertainty deduced in latent and output space; see Fig.
9 (a,b) for ROC curves.
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. <
[Equation]
[Table]
[Equation]
[Table]
[Equation] [Table] [Equation]
[Table]
Q0DJWlNWpFFK1PKKZ4CGrAwfBWpFiJPAEa3qD6tRvPjKluQzvYRgxNyD9kPucEjBSN3taoEV8Wb+7qeJCvhNHRCn5lC9iXyocxAJ4JBjWYNIaONWlbjZnl+wZ8DJxEpJDCWrd7KTTkzQOWAhUEK3bjh2BOyLKPCfYONOJNYsIHZA+axsakoBpdzQ7bIxPjNKb7eLLEPBM/T0xIoHWw8AzyYDAg170puJ/XjsG/8Id8TCKgYV0/pEfCwwST1vCPa4YBTE0hFDFza6YPhBFKJguM6YEZ/HkZdIol5yzkn1bzlWukjrS6AgdowJy0DmqoGtUQ3VE0TN6Re/ow3qx3qxP62seTVnJzCH6A2vyA+TgnoI=<
[Figure]
the uncertainties more well ﬁtted, both in latent and data space.
These well-ﬁtted 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 anautomated system.
We evaluate O O D performance on the commonly used benchmarks (Nalisnicket al., 2019b), where we use F ASHION M NIST (Xiao et al., 2017) as in-distribution and M NIST(Lecun et al., 1998) as O O D. Tab.
3 shows that our online LAE outperforms existing models inboth log-likelihood and Typicality score (Nalisnick et al., 2019a).
This stems from the calibratedmodel uncertainties, which are exempliﬁed in the models ability to detect O O D examples from theuncertainty deduced in latent and output space; see Fig. 9 for ROC curves.
Fig.
10 shows distribution of the output variances for in- and O O D data.
This illustrates that using LAimproves OoD detection.
Furthermore, the online training improves the model calibration.
Missing Data imputation.
Another application of stochastic representation learning is to providedistributions over unobserved values (Rezende et al., 2014).
In many application domains, sensorreadings 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
[Table]
[Equation]
[Table]
(online)
[Table]
AE MSE
[Table]
[Figure]
<
sha
=" /d6TbYJXzqSSrCh0QHxFAGE=">AAAB8nicbVBNS8NAEN34WetX1aOXxSJ4KklB9FjUg94q2A9oQ9lsJ+3STTbsTsQS+jO8eFDEq7/Gm//GbZuDtj4YeLw3w8y8IJHCoOt+Oyura+sbm4Wt4vbO7t5+6eCwaVSqOTS4kkq3A2ZAihgaKFBCO9HAokBCKxhdT/3WI2gjVPyA4wT8iA1iEQrO0EqdLsITGp7d3Ux6pbJbcWegy8TLSZnkqPdKX92+4mkEMXLJjOl4boJ+xjQKLmFS7KYGEsZHbAAdS2MWgfGz2ckTemqVPg2VthUjnam/JzIWGTOOAtsZMRyaRW8q/ud1Ugwv/UzESYoQ8/miMJUUFZ3+T/tCA0c5toRxLeytlA+ZZhxtSkUbgrf48jJpViveecW9r5ZrV3kcBXJMTsgZ8cgFqZFbUicNwokiz+SVvDnovDjvzse8dcXJZ47IHzifP4WbkWY=</latexit> ID <latexit sha1_base64="Wy9c8Vbhv
[Table]
[Figure]
[Table]
[Figure]
[Figure]
[Table]
Our well-calibrated uncertainties propagates to the MNIST classiﬁer and improves the calibration metrics ECE, MCE and
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 ﬁxed.
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 exempliﬁed by the multi-modal reconstructions from the sampled networks in Fig.
The baselines only capture unimodal reconstructions.
Capturing the ambiguity of partly missing data can improve downstream tasks such as the calibration of an image classiﬁer.
In Fig.
11 (c) we demonstrate how averaging the predictions of a simple classiﬁer across reconstructions improves standard calibration metrics.
This is because the classiﬁer inherits the uncertainty and ambiguity in the learned representations.
A deep ensemble of AE s performs similarly to ours, but comes at the cost of training and storing multiple models.
When the entire input image is missing [Table] 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.
[Table]
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 that cannot 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.
Fig.
12 shows the accuracy of a K -nearest neighbor classiﬁer trained on different amounts of labeled data from the MNIST dataset.
For all models with a stochastic encoder, we encode each labeled datapoint 100 times and repeat the experiment 5 times.
When only a few labels per class are available ( 1 - 20 ) we clearly observe that our LAE model outperforms all other models, stochastic and deterministic.
Increasing the number of labels beyond 100 per class makes the AE and LAE equal in their classiﬁcation performance with the AE model eventually outperforming the LAE model.
In Tab. 4 we conduct a similar experiment on the C ELEB A
(Liu et al., 2015) facial dataset, where the the task is to predict 40 different binary labels per data point.
When evaluating the overall accuracy of predicting all 40 facial attributes, we see no signiﬁcant difference in performance.
However, when we zoom in on speciﬁc facial attributes we gain a clear performance advantage over other models.
Fig.shows the mean and variance of ﬁve reconstructed images.
The online LAE produces well-calibrated uncertainties in the output space and scales to large images.
Limitations.
Empirically, the LAE improvements are more signiﬁcant for overparameterized networks.
The additional capacity seems to help the optimizer ﬁnd a local mode where a Gaussian ﬁt is appropriate.
It seems the regularization induced by marginalizing ✓ compensates for the added ﬂexibility.
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 an efﬁcient way to compute its Hessian on high dimensional data that scales linearly with 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 classiﬁcation.
Our work opens the way for fully
Bayesian representation learning where we can marginalize the representation in downstream tasks.
We ﬁnd 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 funded in part by the Novo Nordisk Foundation through the Center for Basic Machine Learning Research in
Life Science (NNF20OC0062606) and by the Pioneer Centre for AI, DNRF grant number P1.