Iterative VAE as a predictive brain model for out-of-distribution generalization
Abstract
Our ability to generalize beyond training data to novel, out-of-distribution, image degradations is a hallmark of primate vision.
The predictive brain, exempliﬁed by predictive coding networks (PCNs), has become a prominent neuroscience theory of neural computation.
Motivated by the recent successes of variational autoencoders (VAEs) in machine learning, we rigorously derive a correspondence between PCNs and VAEs.
This motivates us to consider iterative extensions of VAEs (iVAEs) as plausible variational extensions of the PCNs.
We further demonstrate that iVAEs generalize to distributional shifts signiﬁcantly better than both PCNs and VAEs.
In addition, we propose a novel measure of recognizability for individual samples which can be tested against human psychophysical data.
Overall, we hope this work will spur interest in iVAEs as a promising new direction for modeling in neuroscience.
Introduction
Deep feedforward neural networks have profoundly impacted the development of computational neuroscience models of vision and have become de facto models of core object recognition.
Despite these successes, it is also becoming increasingly clear that current deep neural networks remain outmatched by the primate brain’s power and versatility [35].
The gap between human and machine vision is particularly obvious when artiﬁcial vision systems are required to generalize beyond training data to novel conditions including novel object transformations, occlusion, 3D viewpoints or other image degradations not seen during training [9, 12, 38, 1].
What brain mechanisms allow primate vision to generalize beyond training distributions to novel, out-of-distribution, image degradations?
An increasingly large body of cognitive neuroscience literature points to a critical role for cortical feedback as a key mechanism to help solve difﬁcult recognition problems [30, 43, 36, 17, 23].
However, the computational principles underlying feedback mechanisms are not well understood.
Starting with Helmholtz’s unconscious inference, predictive processing has now become one of the most prominent theories of neural computation (see [16, 39] for reviews).
The theory has taken multiple instantiations [32, 45, 11] but one of the core ideas is that the visual system learns generative models of the world – casting vision as an active inference process.
Predictive coding networks (PCNs) have now become popular computer vision algorithms because of their underlying biological basis [2] and ability to explain behavioral data [27, 41, 5].
In these models, feedback connections aim to reconstruct lower level (bottom-up) visual representations based on abstract (top-down) a priori knowledge derived from a brain’s internal model.
This mechanism is thought to bring out-of-distribution generalization by actively leveraging top-down knowledge to correct for distributional shifts that arise with novel image degradations [4, 6].
In parallel, progress in machine learning in the area of deep generative modeling has also been signiﬁcant.
In particular, variational autoencoders (VAEs) leverage amortized inference to model complex data [19] and have witnessed a widespread use in computer vision due to their high scalability [13, 40].
Recent reﬁnements include iterative extensions of VAEs (iVAEs) which leverage stochastic variational inference (SVI) and achieve greater performance in unsupervised learning [24, 18, 29].
While predictive coding and variational autoencoders appear to share some superﬁcial resemblance [28], a formal connection between the two has not been made explicit.
In this paper, we formally establish a connection between predictive coding and modern variational inference algorithms and consider the plausibility of iVAEs as computational neuroscience model extensions of PCNs.
Furthermore, we establish the superiority of iterative models, iVAEs and PCNs, over classic VAEs for out-of-distribution generalization.
We posit that the iterative and sample-speciﬁc inference of the iVAEs and PCNs is a key mechanism to gradually correct for the distributional shifts observed with out-of-distribution samples.
We experimentally validate this hypothesis and propose a new quantitative measure of image recognizability derived from the number of inference steps required for iVAEs to generate a recognizable reconstruction.
We hope to test these predictions in future psychophysics experiments and to motivate further discussions on the viability of iVAE as a biologically plausible vision model.
φ and θ are the set of variational parameters, amortized over the N training samples
From PCN to iVAE
Here, we formalize mathematically the link between PCN and iVAE – both in terms of their loss functions and their learning mechanisms.
We consider a Bayesian network with 2 random variables as shown in Fig.
The following theoretical derivation could be trivially extended to deeper hierarchical Bayesian networks.
Let us consider a dataset X
= { x ( i ) }
Ni =1 composed of N i.i.d samples of a random variable x .
We assume that x is generated by some random process involving an unobserved random variable z (see Eq. 1).
The latent variable z is sampled from a Gaussian distribution (see Eq. 2).
The mean of the likelihood is parametrized by µ θ (in which θ denotes the parameters) and its variance is considered constant.
[Equation]
The PCN loss is a particular case of the ELBO.
Variational inference algorithms such as VAE approximate the true posterior with a family of distributions q φ ( z | x ) parameterized by φ that match the latent variable density distribution over the entire training set [21].
Those algorithms maximize the log-likelihood of the marginalized input probability distribution by maximizing the evidence lower bound (ELBO).
The ELBO is commonly formalized as (see Demonstration S1):
[Equation]
As shown in Demonstration S2, the ELBO can be re-written as:
[Equation]
Below we show that PCN constitutes a speciﬁc choice of the posterior estimate such that q φ ( z | x ) is a delta distribution centered at the most likely latent variable z ∗ [10, 3]:
[Equation]
z ∗ = argmax z ( p ( z | x )) and δ z ∗ ( z ) = (cid:26) + ∞ if z = z ∗ else (5)
We then derive the PCN loss by replacing the approximate posterior by the delta distribution in Eq. 5:
[Equation]
where C 1 and C 2 are constants.
Note that the second integral on the right-hand side of Eq.
6 is constant as the output of a Dirac function is independent of its center, consequently it has no impact over the minimization process of the ELBO.
The PCN loss as deﬁned originally by Rao & Ballard [32] is shown in Eq.
[Equation]
Under the same hypothesis as the PCN, stating that the variance of both the likelihood and the prior are constant (and equal to unit variance), the maximization of Eq.
7 becomes equivalent to the minimization of Eq. 8.
The link between PCN and VAE becomes evident: the PCN loss is a special case of the VAE loss (i.e., the negative ELBO) with PCN using a point-wise estimate instead of an approximate posterior distribution.
Henceforth, we deﬁne the VAE loss using β as the disentanglement factor [14] in Eq.
9 ( β =1 by default).
[Equation]
Learning and inference in PCN, VAE and iVAE.
In PCN, the minimization of Eq. 8 is performed using the Expectation-Maximization (E-M) scheme [10].
The E-step, corresponding to inference, leads to an estimate of the most probable hypothesis z ∗ given the input x using K gradient descent steps w.r.t.
to the latent variable.
The M-step, which corresponds to learning, uses one step of gradient descent w.r.t.
θ to search for the model parameters that minimize the objective over the entire training set (see Alg. for more details).
VAE, as an amortized variational inference model, learns the parameters of the posterior and the likelihood that are common to the entire training set [21, 33].
These amortized parameters are known to lead to a looser lower bound of the input distribution.
This phenomenon is called the amortization gap [7].
Nevertheless, these methods provide highly scalable frameworks and fast inference (see
Alg. S3 for more details).
Recently, the iterative VAE (iVAE) was shown to mitigate the amortization gap [24, 18, 29].
In the iVAE inference scheme, the amortized posterior parametrization serves as an initialization for the stochastic variational inference (SVI).
Interestingly, the SVI is similar to the E-step used in the
PCN: both estimate instance-speciﬁc latent variables iteratively.
Said differently, the SVI could be considered as a variational version of the E-M algorithm [15].
Algorithm
Input: sample of dataset , model parameters θ , learning rates η z and η , inference steps K
Initialize: z 0 =for k = 1 ..
K do z k
[-+
 = z k − η z ∇ z L PC ( x , θ, z k ) θ = θ − η ∇ θ L PC ( x , θ, z K )-][Equation]
Algorithm
Input: sample x of dataset X , model parameters θ, φ , learning rates η ψ and η , inference steps K
Initialize: ψ 0 = (cid:0) µ φ ( x ) , σ 2 φ ( x ) (cid:1) for k = 1 ..
K do ψ
[Equation]
The link between PCN, SVI and iVAE becomes straightforward: they all share an iterative and sample-speciﬁc mechanism to infer an approximate posterior.
In contrast, VAE learns a posterior distribution over the entire training set.
The SVI can be seen as a variational extension of the PCN. In addition to this SVI process, the iVAE includes an additional amortized initialization step of the posterior.  Out-of-distribution generalization: deﬁnition and intuition In out-of-distribution generalization tasks, a model must, at test time, generalize to new data distributions that were not encountered during training.
Out-of-distribution generalization can be cast as an invariance problem: similar training and testing samples should both elicit the same representation even if they are drawn from different distributions [22]. In this paper, we simplify the out-of-distribution generalization task by only considering the robustness to relatively simple distributional changes such as additive noise, Gaussian blurring or salt & pepper degradation. We posit that the amortized parameters of the likelihood term, which reﬂect the training likelihood distribution, give to the ELBO optimization surface a high degree of invariance w.r.t the input perturbation. To test this hypothesis, we ﬁrst train VAE and iVAE toy models (with 2D latent variables models only) on the original (non-degraded)
MNIST dataset.
We then evaluate the optimization surface of both models with degraded inputs (see Fig. 2).
Importantly, during evaluation, none of the models are given access to the non-degraded input. Fig. 2 shows that the topology of the optimization surface (i.e. the ELBO) is relatively invariant to the input degradation. Intuitively, since the parameters of the VAE are speciﬁc to the training distribution, the generated posteriors for out-of-distribution samples won’t be aligned with the likelihood distribution as this is the case for training samples.
Therefore, in such a situation, the reconstruction generated by the VAE is likely to be shifted away from the training distribution (see red cross in the left map of Fig.away from green cross).
In contrast, the iVAE (SVI and PCN) have the ability to move in the ELBO landscape towards posteriors that are maximizing the likelihood probability.
We postulate that this hypothesis reﬁnement mechanism should lead to a better out-of-distribution generalization capability.
The following section is an empirical veriﬁcation of this hypothesis.
Experiments
Methods.
As a baseline, we also report the classiﬁcation accuracy for corrupted images (this baseline is denoted CL in Fig.
3-b). The SVI, PCN, VAE and iVAE were trained on clean images from the MNIST dataset [26] (see S4 for details regarding the choice of architectures and hyperparameters).
We subsequently evaluated the models with MNIST test images that we corrupted
[Figure]
with varying levels of noise.
Three types of noise degradations were used: Gaussian additive noise (white noise), salt & pepper noise (pixel intensities sampled from a Bernoulli distribution) and Gaussian blurring (convolution with a Gaussian kernel).
Note that during evaluation, none of the models are given access to the original (and non-degraded) training distribution.
Out-of-distribution generalization was evaluated as the accuracy of a convolutional neural network (CNN) fed with image reconstructions from the three generative models.
Importantly, this CNN classiﬁer was only trained on the original MNIST dataset so as to properly assess the generative models’ abilities to preserve the classiﬁer decision boundary.
Other metrics could be used to evaluate out-of-distribution generalization (e.g. the (cid:96) 2 -norm with the non-degraded input), we choose the classiﬁcation accuracy metric as it leads to a convenient parallel with behavioral data that can be collected from human participants during a psychophysics experiment (see subsection ‘ iVAE provides a testable prediction of the predictive brain theory) ’.
iVAE outperforms SVI, PCN and VAE in out-of-distribution generalization.
Fig.
3-a shows that the accuracy of SVI, PCN and iVAE increases over inference steps and exceeds that of VAE by a large margin.
The perceptual quality of the models’ reconstructions also improves with more inference steps (see Fig.
S6 for additional examples).
The accuracies of iVAE’s initial inference step and VAE’s single inference step are relatively similar due to the amortized initialization of the iVAE.
In contrast, the PCN’s and SVI’s initial accuracy are at the chance level because of the random initialization of the latent variable.
At the ﬁnal inference step, we observe that SVI constantly outperforms PCN in terms of accuracy.
This suggests that the variational estimation of the posterior in SVI is more efﬁcient than the point estimate used in PCN – at least in terms of out-of-distribution generalization.
Since iVAE combines an amortized initialization with a variational iterative inference process, it out-performs systematically PCN, SVI and VAE as shown in Fig.
3-b (see Fig.
S7 for the accuracy of the models on different types of noise).
Furthermore, the accuracies of all generative models largely exceed the classiﬁer baseline on noisy images.
This gap is attributed to the models’ abilities to push out-of-distribution latent variables back in the original distribution (see Fig.
2 for an illustration). We also observe that higher noise levels increase the performance gap between VAE and the iterative models (i.e., iVAE, SVI and PCN).
This suggests that such iterative inference processes are crucial for out-of-distribution generalization.
In Fig.
3-c, the disentanglement factor β which controls the prior constraint is varied under different noise levels.
As one might expect, the resulting accuracy is positively inﬂuenced by the prior when the noise level is high, and negatively when it is low.
All reported results are consistent across noise types (see Fig. S8) with the exception of the Gaussian blurring where increasing β has no signiﬁcant effect on the classiﬁcation accuracy. iVAE needs more inference iterations to classify atypical samples.
The ELBO measures the likelihood of an input image belonging to the training distribution.
Therefore, the ELBO should be able to capture the level of typicality of individual samples with more prototypical (resp.
atypical) samples associated to higher (resp.
lower) ELBO values.
Fig.
4-a qualitatively illustrates that higher
1 st 2 nd 3 rd 4 th
E L B O q u a r t il e 1 st 25 th 50 th 75 th 100 th ELBO centile
n b i t e r a t i o n f o r g
oo d c l a ss i f i c a t i o n
[Figure]
[Figure]
ELBO digits appear more prototypical while lower ELBO digits appear more atypical or ambiguous.
Interestingly, we also ﬁnd empirically that iVAE allows us to draw a link between the ELBO and the number of inference steps before the output reconstruction is correctly classiﬁed (see Fig. 4-b). iVAE provides a testable prediction of the predictive brain theory. Psychophysical data from prior studies have shown that observers’ response times increases exponentially as the visibility of a stimulus is being reduced either by reducing the luminance [31] and/or the contrast [8].
Our results suggest that the model response time measured by the number of iterations needed to achieve good classiﬁcation accuracy follows a similar trend when expressed as a function of the ELBO (see Fig.
4-b).
In addition, the ELBO appears qualitatively to also provide a good measure of a stimulus prototypicality (see Fig.
4-a).
While recognizability in psychophysics is often modeled by the distance to a classiﬁer decision boundary, exemplars or prototypes [34], we propose the ELBO as an alternative measure of recognizability.
This is consistent with a recent psychophysical study [37] which suggests that generative models account well for the perception of surface glossiness.
The backward masking protocol (i.e., reducing the visibility of a stimulus by presenting a noise mask shortly after the stimulus onset) is commonly used in psychophysics experiments to alter the effect of the feedback connections [25].
Masking has been shown to strongly impair recognition under challenging conditions including object occlusion and blurring [42] and it is widely believed that feedback mechanisms are necessary to help disambiguate degraded stimuli [30, 43, 36, 17, 23].
This is consistent with the results shown for the iVAE in Fig.
3-c.
Conclusion
We have shown that the iVAE constitutes a variational extension of the PCN.
In addition, we show that the iVAE signiﬁcantly outperforms SVI, PCN and VAE in out-of-distribution generalization.
Finally, we show that the ELBO of the iVAE could serve as a novel measure of recognizability for individual samples which can be tested against human psychophysical data.
Nevertheless, the link between the ELBO and standard cognitive science models based on examplars and prototypes is still an open question [34].
As a next step, we believe that a hierarchical iVAE would provide a better model of primate vision given the hierarchical organization of our own visual system; it is also likely to be necessary to extend this work to natural images.
Such an extension would allow to learn prior distributions directly from data and thus should provide more accurate predictions on the role of the feedback connection in the brain.
Last, but not least, we have suggested how speciﬁc measures of recognizability could be derived from the iVAE model resulting in testable predictions for psychophysics which we plan to test in future work.
Overall, we hope to spur interest from the community in considering iVAEs as a promising new direction for modeling vision beyond the feedforward sweep.
Broader impact statement
A key question for computational neuroscience is to understand the role of cortical feedback which is currently poorly understood.
By making explicit connections between cortical feedback and deep generative models, the present work may help further our understanding of brain mechanisms. Acknowledgment
This work was funded by ANITI (Artiﬁcial and Natural Intelligence Toulouse Institute, ANR-19PI3A-0004)
Supplementary Information S1 : Derivation of the ELBO
The objective of VAE is to estimate the true posterior p ( z | x ) with q φ ( z | x ) using the amortized parameters φ .
Using the Kullback-Leibler (KL) divergence between the two distributions, one can derive the evidence lower bound (ELBO) commonly used to optimize the variational parameters of a
VAE:
[Equation]
By deﬁnition
[Equation]
the KL being positive
[Equation]
one can derive a lower bound on the marginal input distribution:
[Equation]
The ELBO formulation in Eq.
4 is less commonly used compared to the one in Eq.
but the two are strictly equivalent:
[Equation]
[Equation]
PCN, VAE and iVAE leverage the standard generative model settings of VAE as described in the following equations:
[Equation]
All encoder (i.e. ψ ) and decoder (i.e. µ θ ) models use three fully connected layers with tanh activations, hidden dimensions h 1 = 512 , h 2 = 256 and latent dimension z = 15 .
For the VAE and iVAE, β was set by default to 1 and when speciﬁed we vary β with the following values : [0 . 0 , 1 . 0 , 2 . 0] .
SVI in the iVAE is run for 20 iterations with a inference update rate η ψ = 10 − 2 during training and 500 iterations with a inference update rate of η ψ = 10 − 3 during evaluation.
For the PCN, the total number of inner-loop iterations was set to 100 during training and 500 during evaluation with η z = 10 − 2 .
We used the MNIST dataset with normalized pixel values.
Fig.
4-b was produced with the 60 , 000 test samples of the QMNIST dataset [44].
All models were trained for 200 epochs with a batch size of 1,024, a learning rate η = 10 − 3 .
For all models, we use Adam [20] as an optimizer both for inference and learning.
The classiﬁer’s architecture and training parameters were the default used in the pytorch MNIST example.
S5 : Noise parameters.
Examples of an MNIST digit corrupted with different noise types and levels.
The level of noise applied to an image is controlled with the standard deviation (denoted σ ) of the
Gaussian distribution for the white noise and the Gaussian blurring, and the probability of pixel’s corruption (denoted p ) for the salt & pepper degradation.
[Table]
Gaussian [Table] σ = σ = σ = σ =
White Noise σ =0.
σ =0. σ =0.
σ =0.
Salt & pepper p =0.
p =0. p =0.
p =0.
input t= t= t= target iVAE
PCN
VAE iVAE
PCN
VAE iVAE
PCN
VAE
Salt
[Figure]
of the factor of
The in the are the same than in S5 and S7.
Note that increasing beta leads to better accuracy for types (i. Gaussian noise and salt & pepper), and has no signiﬁcant effect on local noise
[Figure]