On the Relationship Between Variational Inference and Auto-Associative Memory
AffiliationAddress
Abstract
In this article, we propose a variational inference formulation of memory retrievalin auto-associative memories, allowing us to combine perceptual inference withrepresentation learning into the same mathematical framework.
In this formulation,the prior probability distribution onto representations is made memory dependent,thus pulling the inference process towards stored representations.
We then studyhow different neural network approaches to variational inference can be appliedin this framework.
We compare methods relying on amortized inference suchas Variational Auto Encoders and methods relying on iterative inference such as
Predictive Coding and suggest combining both approaches to design new autoassociative memory models.
We evaluate the obtained algorithms on the CIFAR10and CLEVR image datasets and compare them with other associative memorymodels such as Hopfield Networks, End-to-End Memory Networks and Neural
Turing Machines
Introduction
In the recent years, many methods have been proposed in order to augment deep learning modelswith long-term memories.
These models provide writing and reading mechanisms allowing to storeand retrieve certain patterns inside a memory.
Among these models, associative memories arememories that perform content-based addressing, meaning that they retrieve stored patterns based onan approximate or incomplete version of those.
While many models consider one-step reading mechanisms to retrieve patterns from the memory,biological memory retrieval is an iterative process that can take a variable amount of time.
Takingthis into consideration, Hopfield networks [14] have been proposed as an Associative Memory (AM)implemented by a Recurrent Neural Network (RNN).
Hopfield networks perform memory retrievalas an iterative process where stored patterns constitute local attractors of the RNN dynamics.
In this work, we show how AM models such as Hopfield networks can be formulated as variationalinference methods using a memory dependent probabilistic model of the observed data.
In variationalinference algorithms such as Variational Auto-Encoders (VAE) [18, 31] and Predictive Coding (PC)[30, 11], the representation inferred from the observed data tries to minimize an energy functionbased on the probabilistic model p ( x ) generating the observed data.
In the Free Energy Principle(FEP) literature [9, 10], this function is called Variational Free-Energy (VFE), and is equivalent tothe negative Evidence Lower Bound (ELBO) more often used in machine learning.
The probabilisticmodel p ( x ) can be decomposed according to a prior probability distribution over representationsp ( z ) , and a likelihood p ( x | z ) describing the probability of observed data based on the representation.
We propose to make p ( z ) depend on the patterns stored in memory, which allows us to derive anew expression of the VFE.
The obtained energy function has minima that should be close to oneof the stored patterns while properly encoding the observed input.
As such, inference based on this energy function can be seen both as perceptual inference and memory retrieval.
While usingmemory-dependent generative models or Gaussian Mixture Models (GMM) with VAEs [8, 2, 22] arenot novel ideas, the PC approach has never been applied to this problem.
The paper is organized as follows: in section 2, we provide a deeper presentation of the conceptsconnected to our work and review related approaches.
In section 3 we present an overview of the
FEP mathematical framework and derive an expression of the VFE depending on the patterns storedin memory.
In section 4 we design several AM models minimizing this energy function.
In section5, we evaluate the obtained algorithms on the task of memory retrieval on two image datasets, andcompare their performance with other AM models.
The contributions brought by this work are the following:
• We design four AM models based on a variational inference formulation of memory retrieval.
• We draw a connection between PC and the modern continuous Hopfield network [29].
Related work2.1 Auto-associative memory
Several approaches to the long-term storage of information in artificial neural networks have beenproposed.
Memory networks [37], and End-to-End Memory Networks (MemN2N) [33] propose tostore information in a memory matrix that can be addressed using attention coefficients computedbased on the content of different memory locations.
The Neural Turing Machine (NTM) [12]and the differentiable neural computer [13] combine this content-based addressing using attentionmechanisms with a location-based addressing allowing more computer-like memory accesses.
Standing out from these approaches that consider a feedforward reading mechanism, Hopfieldnetworks [14] and continuous Hopfield networks [15] instead store patterns as local attractors of an
RNN.
These memories can thus be addressed by initializing the RNN state with the input pattern, andretrieving the stored attractor after convergence.
In order to improve the memory capacity, modern
Hopfield networks [21, 20, 7] propose several variants of the energy function using polynomial orexponential interactions.
Extending these models to the continuous case, [29] proposed the Modern
Continuous Hopfield Network (MCHN) with update rules implementing self attention, that they relateto the transformer model [35].
In [25], the authors introduce a general Hopfield network frameworkwhere the update rules are built using three components: a similarity function, a separation function,and a projection function.
It has been shown that overparameterized auto-encoders also implement AM [27, 32].
These methodsembed the stored patterns as attractors through training, and retrieval is performed by iterating overthe auto-encoding loop.
In contrast, our methods allow one-shot writing: a new pattern x can beinserted in memory simply by computing its representation and adding it to the memory matrix M .
Another line of research takes inspiration from the Sparse Distributed Memory model [17], buildingconnections with attention mechanisms [3] and with the variational inference framework [39, 40, 28]. In particular, the Kanerva Machine [39, 40, 28] is similar to the models we build in many aspects:they use an iterative reading mechanism, a memory-dependent prior on the representation, all withinthe variational inference framework. Though, the PC toolbox provides methods that set the proposedmodels apart from these approaches, as detailed in the methods section. 2.2 Predictive Coding
PC is a theory of brain function [30, 6] extending the idea that neural representations emerge as part ofan inference process of the causes of sensory observations, as already suggested by Helmholtz in 1867[36].
The FEP [9, 10] provides a principled derivation of PC networks based on the minimization ofvariational free-energy (VFE), a quantity equivalent to the negative evidence lower bound (ELBO)used in variational Bayesian methods, and defined as:
[Equation]
where p ( x , z ) denotes the generative model and q ( z ) denotes the approximate posterior (also calledrecognition density) on z .
VAEs [18, 31] are a well-known method applying the idea of VFEminimization (equivalently ELBO maximization) to neural networks.
In VAEs, this quantity is onlyoptimized during the model training, and perceptual inference is performed as a simple forward passthrough the encoder.
In contrast, neural network models based on PC intertwine the prediction (decoder) and inference(encoder) mechanisms within a single hierarchical recurrent architecture comprising a populationof representation neurons and a population of prediction error neurons at each layer.
Perceptualinference is then an iterative mechanism supported by the dynamics of this RNN.
Given an observedinput x , the representation z is updated at each time step based on a bottom-up signal pushing ztowards values that minimize the reconstruction error, and on a top-down signal pushing z towardsvalues that maximize its prior probability p ( z ) .
While this prior probability on z is often ignored,we show that it can bring auto-associative capacities to the PC network, and we draw a connectionbetween the obtained mechanism and the modern continuous Hopfield network proposed in [29].
While both VAEs and PC networks output an estimation of the approximate posterior q ( z ) , theydiffer in the computational mechanisms used for inference: VAEs perform amortized inference via aforward pass through an encoder, while PC networks perform an iterative inference implementedby its recurrent dynamics.
Some works have suggested combining both approaches: the iterativeinference can be initialized using the estimation provided by the amortized inference method [1, 34].
Memory-based Variational Inference
3.1 Framework definition
Our framework is based on a generative model p ( x , z ) that can be factored into a prior probability p ( z )and a likelihood p ( x | z ) .
The prior probability on z is defined as a memory dependent distribution:
[Equation] [Equation]
where the vectors M k (the columns of M ) constitute a repertoire of representations.
The probability p ( x | z ) is defined as a hierarchical generative model featuring several intermediatevariables { h 1 , · · · , h L − 1 } , where L denotes the number of layers.
By extension we use the notationsh
= x and h L = z .
We assume that the generative model is a cascade of multivariate Gaussians:
[Equation]
where I is the covariance matrix of the Gaussians and is uniform across all layers, and f l θ are functions(typically neural network layers) parameterized by θ .
Note that this can be adapted to arbitrarycomputation graphs by replacing h l +1 by the set of parent nodes of h l for each node.
For simplicity,we assume in the following derivations that each node only has one parent.
In our experiments,the functions f l θ correspond to the different layers of a Convolutional Neural Network (CNN) onthe CIFAR10 dataset, and to a MONet [5] decoder on the CLEVR dataset.
We denote by f θ thecomposition of all layers f l θ .
We can derive the VFE corresponding to the described generative model.
The FEP formulation of PCuses different approximations that allow us to greatly simplify the expression of the VFE introducedin equation (1).
We refer to appendix A for detailed derivations and simply provide the simplifiedexpression:
[Equation]
where C is a quantity independent from { x , ˆ h 1 , . .
, ˆ h L − 1 , ˆ z } .
The vectors { ˆ h 1 , . . .
, ˆ h L − 1 , ˆ z }correspond to the means of the approximate posterior q ( h 1 , . . .
, h L − 1 , z )
=
q ( z ) (cid:81)
L − [Equation] =1 q ( h i ) .
The term corresponding to the prior probability is often omitted, which can be justified as beingequivalent to having no prior preferences over different values of z .
Key to our method is the ideathat this prior can pull the inference process towards values of z previously stored in the memory M ,turning the patterns M k into attractors of the PC network.
As such, a suitable distribution would beone that associates high probabilities for the patterns stored in memory.
3.2 Classification of related methods in this framework
Here, we show that we can formulate Modern Continuous Hopfield Networks (MCHN) as PCnetworks derived from this expression of the VFE under some conditions:
• There is no representation component: L = 0 .
• The PC network is initialized with ˆ z = x .
• The prior distribution p ( z ; M ) is defined as:
[Equation]
This prior distribution is a Gaussian Mixture Model (GMM) with mixture means corresponding tothe N stored patterns and mixing coefficients π k depending on the patterns’ Euclidean norm.
It isparameterized by a coefficient β > 0 .
Based on these assumptions, we can derive the followingexpression for the VFE:
[Equation]
which is, up to a constant C and a factor β , equivalent to the MCHN energy function proposed in[29].
The complete derivations of this expression are provided in appendix B.
We can note that VFEdoes not depend on the input x .
The input is not part of the energy function, but serves as an initialestimate of the approximate posterior mean ˆ z .
According to the FEP formulation of PC, applyinggradient descent on the VFE with regard to the approximate posterior mean ˆ z yields the update rulesof the PC network.
Using our expression of the VFE, we obtain:
[Equation]
[Equation] This equation is identical to the update rule of MCHN (detailed derivations are provided in appendix).
We have shown that our framework can be related to MCHNs, but variational inference methods forperceptual inference can also be retrieved in this framework by simply using a neutral prior.
For VAEs and variants, this prior probability distribution is exploited during learning, but theamortized inference cannot be dynamically adapted to this distribution.
Therefore, to take intoaccount a change in this prior, we would need to retrain the encoder.
In the derivations of PC networks[11, 4], the term p ( z ) is often ignored, which is equivalent to having a flat prior distribution.
However,theoretically the iterative inference could take into account a prior distribution, and dynamically adaptto changes in this distribution.
In the family of iterative inference algorithms, we can also mention the methods based on backpropagation (BP) to estimate the representation.
Instead of using neural computations to simulate thegradient descent on the energy function as done in PC, these methods directly use BP to optimizethe representation ˆ z .
For instance [26] minimize prediction error using BPTT to adjust a latentvariable in an RNN.
This method is in fact very similar to PC, since it has been shown that undersome conditions, PC approximates the update rules entailed by BP applied on the reconstruction error[38, 24].
We provide in table 1 a simple classification of these approaches depending on whether theyuse a representation and memory component, and whether the inference mechanism is amortized oriterative.
[Table]
Methods
In this section, we present the AM models that we have designed using this framework.
All thesemodels are based on a pre-trained VAE that is used to provide initial estimates ˜ z and predictions of xbased on ˜ z .
The proposed methods are represented in figure 1.
4.1 Probabilistic model
The prior probability over the representation z is crucial to obtain a model that infers representationsclose to stored patterns.
We have shown that MCHN can be obtained starting from a GMM priorbiased towards stored patterns of large Euclidean norm.
Since there is no intuitive justification forthis bias, we choose to instead start with a balanced GMM prior distribution:
[Equation]
(9) where Σ denotes the covariance matrix, uniform across the N mixtures.
In the following models, weconsider the simpler case where Σ = σ 2 I , except for the last model where P = Σ − 1 is trained.
4.2 PC based inference (VAE-PC-GMM)
Following the FEP formulation of PC, we can derive the gradient of F according to each layer’srepresentation ˆ h l .
As such, performing gradient descent on F yields a system of update rules oneach representation, that can be interpreted as RNN dynamics.
This PC network comprises at each layer two populations of neurons, one encoding the layer’s representation, and one encoding the layerprediction error ϵ l = ˆ h l − f l θ ( ˆ h l +1 ) .
Detailed derivations of this model are provided in appendix C, where the forward pass throughthe model is given in algorithm 1.
Since the RNN dynamics implement a gradient descent, theinitialization of this network is responsible for the local minimum to which it converges.
We label VAE-PC-GMM the version of this algorithm where ˆ z is initialized using the result of the amortizedinference via the encoder of the VAE.
4.3 BP based inference (VAE-BP-GMM)
In this second algorithm, we instead use BP as the iterative inference mechanism used to optimize ˆ z .
BP minimizes the following loss function:
[Equation] [Equation]
where γ is an hyperparameter weighting the influence of bottom-up and top-down mechanisms, andf θ denotes the decoder of the VAE.
This gradient descent is parameterized by a learning rate λ .
Onceagain, we can initialize the gradient descent using the estimate obtained with the encoder of the VAE.
Derivations of this model are provided in appendix C, where the forward pass through the model isgiven in algorithm 2.
Investigating the relationship between PC and BP based inference, we have found that PC networksneeded a larger number of iterations to convey information from the reconstruction error, andthat this number of iterations was exponential with regard to the depth of the decoder.
For thisreason, experiments with the VAE-PC-GMM model were prohibitively slow, and we only conductedexperiments with the remaining models.
4.4 Restricting the iterative inference to the memory component (VAE-GMM)
Backpropagating the reconstruction error might not be necessary when the estimate provided bythe VAE encoder already conveys enough information from the observed input x .
As such, we alsoexperiment with a simpler version of the previous model where iterative inference of ˆ z only considersthe top-down update rule coming from the memory.
In this simpler version, the update rule for ˆ zprovided in algorithm 1 becomes (proof in appendix C):
[Equation]
This update mechanism is computationally lighter that the BP-based inference model, and has theadvantage of being differentiable.
4.5 Training the precision matrix (VAE-GMM*)
For this last model, we suggest optimizing the precision coefficients P =
Σ − 1 of the GMM.
Thesecoefficients condition the shape of the Gaussian mixtures.
Some memory retrieval tasks might needto ignore partial information from the observed input x and as such could benefit from such anadaptation.
In our experiments, we design such tasks on the CLEVR dataset, where the objective isto retrieve scenes using as input a shifted image of the same scene, or an image of the same scenewhere the colors of the objects have been modified.
To properly accomplish these two tasks, the AMmodel needs to learn to give less importance to position information in the first case, and to colorinformation in the second case.
The model starts from the estimate ˆ z inferred by the VAE, and uses the following update rule:
[Equation]
During training, P is optimized using BP in order to reduce the mean squared error between theinferred representation ˆ z and the correct memory pattern M ∗ k .
During evaluation, the values of Pare fixed.
Experiments
In this section, we present the experiments performed to evaluate the proposed models.
In all experiments, we measure performance using the percentage of properly retrieved patterns from associativememories containing N = 100 patterns.
We consider that a pattern is properly retrieved if the distancebetween the inferred representation ˆ z and the correct memory pattern M k ∗ is lower than a thresholdvalue chosen manually.
All the implementation details (benchmark models, hyperparameters, training,dataset splits) are provided in appendix D.
5.1 Datasets
We evaluate the proposed models on two image datasets: CIFAR10 and CLEVR.
CIFAR10 [19](MIT License) consists of 32 × 32 RGB images of 10 classes.
We pretrain a VAE on the training setusing a convolutional architecture, with a latent space of dimension d = 16 .
The CLEVR dataset [16] (CC BY 4.0 License) consists of 64 × 64 RGB images of 3D scenescomposed of simple 3D objects.
We use a pretrained MONet [5] model as the VAE, with a latentspace of dimension d = 4 × 16 .
The MONet encoder infers a set representation given an input sceneimage.
In our experiments, we flatten this set representation, which can lead to wrong measures ofsimilarity between two representations.
Indeed, a permuted set representation still encodes the sameinformation, but once flattened the permuted and original representations might have a low similarityscore.
To limit this issue, we have capped the number of objects present in the scene images to 3, andcapped the size of the set representation to 4.
Examples of images from the CLEVR dataset alongwith possible transformations are presented in figure 2.
5.2 Benchmark models
We compare our AM models with Modern Continuous Hopfield Networks (MCHN), Neural Turing
Machines (NTM) and End-to-End Memory Networks (MemN2N).
For fair comparison, in allbenchmark models, the memory is initialized with the stored patterns (we do not use the writingmechanisms of NTMs).
We also experiment with these models in the representation space, where thememory is initialized with representations of the stored images, and the networks receive as input therepresentation of the corrupted image.
We indicate this modification with the VAE prefix.
The NTM and MemN2N models are trained using BP using as loss function the mean squared errorbetween the predicted output ˆ x (respectively ˆ z when using the VAE) and the correct memory pattern
M ∗ k .
For a fair comparison, this training is performed with clean inputs x = M ∗ k , to avoid givingany prior information about the type of transformation the inputs might be corrupted with duringevaluation.
Finally, we also experiment with a version of the MCHN using a balanced GMM as starting prior distribution, that we denote GMM in our experiments.
Equivalently, this is the versionof our algorithm without a representation component.
5.3 Corrupted inputs
First, we evaluate whether the associative memory models are able to properly retrieve stored patternswhen given corrupted versions of these patterns.
We apply different level of noise to the images( σ ∈ { 0 . 2 , 0 . 4 , 0 . 6 , · · · , 2 . 0 } ), as well as different sizes of masks (16 and 24 pixels wide on the
CIFAR10 dataset, and 32 and 48 pixels wide on the CLEVR dataset).
The results of these experimentsare reported in table 2 and 3.
More figures and examples are provided in appendix E.
[Table] [Table]
We can observe that the two models using dot-product attention, MCHN (Hopfield in the tables) and
MemN2N, perform poorly compared to other methods, not even reaching 100% of correctly retrievedpatterns when presented with clean inputs.
The NTM model, using cosine similarity based attention,works better and sometimes outperforms the proposed models.
Among the proposed models, we cansee that the VAE-BP-GMM does not perform significantly better than the simpler VAE-GMM model.
An important observation from these results is that in most scenarios, the best performing methodsdo not use a representation component.
In particular, the MONet model seems very sensitive to noise.
However, our intuition is that memory retrieval on the representation level should be more powerfulif the transformation applied on the observed inputs x has a limited effect on the representations.
5.4 Scene transformations
To verify this hypothesis, we perform two additional experiments on the CLEVR dataset.
If weassume that the representations provided by MONet encode object positions, shapes and colors, thenan AM model working on the representation level could perform well when we apply transformationson these specific object features.
We propose two such scenarios.
In the first scenario, we perform an RGB rotation on the images, which results in identical scenes with the exception of the object colors.
In the second scenario, we take a shifted crop of the original CLEVR image, which can simulate ashifted point of view of the same scene.
Examples images are displayed in figure 2.
We measure the success of memory retrieval with these transformed inputs using our models aswell the NTM and MemN2N models.
We also experiment with the VAE-GMM* model performingprecision coefficient training.
We expect the model to learn that some information conveyed by theencoder (for instance color in the first scenario) is irrelevant for memory retrieval.
For fair comparison,we also experiment with the VAE-NTM and VAE-MemN2N models trained in the two scenarios, thatwe denote with an asterisk.
The results are reported in table 4.
[Table]
Using a representation component improves the performance of the models in these scenarios.
Additionally, training on the retrieval task where the transformations are applied provides almostperfect retrieval scores for the VAE-GMM and VAE-NTM models.
Discussion
We have proposed a formulation of auto-associative memories based on the FEP and the variationalinference framework, where the stored patterns condition the prior probability on the representationz .
This framework has allowed us to draw a connection between MCHNs and the PC theory, as wellas to design several AM models with high retrieval accuracy and robustness.
Combining representation and memory allows to retrieve patterns on the representation level, whichdecreases the size of the memory store, and performs retrieval based on a similarity measured in therepresentation space, where we can expect more meaningful features to appear.
For instance, we haveshown that this allows to recognize visual scenes from a shifted point of view, while pixel-level AMmodels failed most of the time.
Our results also seem to demonstrate that distance based attention asused in our models (see equation 11) outperforms dot-product based attention as used in MCHNs,although this performance gap could be mitigated with normalization techniques, as hinted by theresults obtained with NTMs, that use cosine similarity instead of dot-product.
The proposed PC and BP based models take into account information from the query x at eachinference iteration, by directly minimizing an energy function that depends on x (eq. 4 and eq. 10).
In contrast, the Kanerva Machine [39, 40] or the VAE-GMM and VAE-GMM* models we proposedonly use this information to output an initial estimate ˆ z that is later optimized in order to minimizean energy function only depending on the memory. According to our experiments, this feature of the
PC and BP based models does not seem useful for memory retrieval. However, in other applicativesettings, these models could benefit from their ability to weight query pattern information and storedpatterns information, performing a form of memory-aided iterative perceptual inference.
The proposed formulation is currently limited to content-based reading mechanisms.
In contrast,models such as NTMs offer a larger variety of addressing schemes as well as writing mechanisms thatallow to only store relevant information in the memory.
Another limitation of this work is that we didnot investigate the memory capacity and convergence properties of the proposed models.
Future workshould focus on these analyses, and possible improvements of the PC-based inference, for instanceusing precision weighting, or the "learning to optimize" approach of [23].
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope? [Yes] We believe that the three claimed contributions in
Introduction are properly reflected in sections 3 and 4.
(b) Did you describe the limitations of your work?
[Yes] See the last paragraph of theDiscussion.(c)
Did you discuss any potential negative societal impacts of your work?
[No] We did notidentify any direct potential negative societal impact.
If an indirect potential impact is
brought to our attention, we would be happy to discuss it.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] The assump-
tions and definitions used in our derivations are stated (see sections 3.2 and 4, and
appendices A,B and C).
(b) Did you include complete proofs of all theoretical results?
[Yes] We provide detailed
derivations of our models and theoretical results in appendices A, B and C.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes] See appendixD.
(b) Did you specify all the training details (e
g., data splits, hyperparameters,
how they
were chosen)?
[Yes] See appendix D.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] See tables 2, 3 and 4, as well as figures in appendixE.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See appendix D.
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] Datasets were creditedin section 5, code assets are credited in appendix D.
(b) Did you mention the license of the assets?
[Yes] Licences for the datasets were
mentionned in section 5, the license for used code assets is mentionned in D.
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
Implementations of our models are linked in appendix D.
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[No] The data we used are commonly used ML
datasets.
CLEVR is an automatically generated dataset and to our knowledge, the
CIFAR10 dataset does not contain any personally identifiable of offensive content.
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]