On the Relationship Between Variational Inference and Auto-Associative Memory
Anonymous Author(s)
AffiliationAddress email
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.
Taking thisinto consideration, Hopfield networks [11] have been proposed as an associative memory implementedby a Recurrent Neural Network (RNN).
Hopfield networks perform memory retrieval as an iterativeprocess where stored patterns constitute local attractors of the RNN dynamics.
In this work, we show how auto-associative memory models such as Hopfield networks can beformulated as variational inference methods using a memory dependent probabilistic model of theobserved data.
In variational inference algorithms such as Variational Auto-Encoders (VAE) [14, 23]and Predictive Coding (PC) [22, 8], the representation inferred from the observed data tries tominimize an energy function based on the probabilistic model p ( x ) generating the observed data.
Inthe Free Energy Principle (FEP) literature [6, 7], this function is called Variational Free-Energy (VFE),and is equivalent to the negative Evidence Lower Bound (ELBO) more often used in machine learning.
The probabilistic model p ( x ) can be decomposed according to a prior probability distribution overrepresentations p ( z ) , and a likelihood p ( x | z ) describing the probability of observed data based onthe representation.
We propose to make p ( z ) depend on the patterns stored in memory, which allowsus to derive a new expression of the VFE.
The obtained energy function has minima that should be close to one of the stored patterns while properly encoding the observed input.
As such, inferencebased on this energy function can be seen both as perceptual inference and memory retrieval.
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 auto-associative memory models minimizing this energyfunction.
In section 5, we evaluate the obtained algorithms on the task of memory retrieval on twoimage datasets, and compare their performance with other associative memory models.
The contributions brought by this work are the following:
• We propose a variational inference formulation of auto-associative memories. • We design four new auto-associative memory models based on this formulation.
• We draw a connection between PC and the modern continuous Hopfield network [21].
Related work
2.1 Auto-associative memory
Several approaches to the long-term storage of information in artificial neural networks have beenproposed, with different addressing mechanisms.
Memory networks [29], and End-to-End Memory Networks (MemN2N) [25] propose to store information in a memory matrix that can be addressedusing attention coefficients computed based on the content of different memory locations.
The Neural
Turing Machine (NTM) [9] and the differentiable neural computer [10] combine this content-basedaddressing using attention mechanisms with a location-based addressing allowing more computer-likememory accesses.
Standing out from these approaches that consider a feedforward reading mechanism, Hopfieldnetworks [11] and continuous Hopfield networks [12] 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 [17, 16, 5] propose several variants of the energy function using polynomial orexponential interactions.
Extending these models to the continuous case, [21] proposed the Modern
Continuous Hopfield Network (MCHN) with update rules implementing self attention, that theyrelate to the transformer model [27].
[19] introduce a general Hopfield network framework wherethe update rules are built using three components: a similarity function, a separation function, and aprojection function.
2.2 Predictive Coding
PC is a theory of brain function [22, 4] extending the idea that neural representations emerge as partof an inference process of the causes of sensory observations, as already suggested by Helmholtz in1867 [28].
The FEP [6, 7] provides a principled derivation of PC networks based on the minimizationof variational 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 [14, 23] 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 observed input 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 [21].
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, 26].
In this work, we also try to experiment with different combinations of variational inference approachesto minimize our memory dependent energy function.
Memory-based Variational Inference3.1 Framework definition
The proposed framework is based on a generative model p ( x , z ) that can be factored into a priorprobability p ( z ) and a likelihood p ( x | z ) .
The prior probability on z is defined as a memory dependentdistribution:
[Equation] [Equation]
where the vectors M k 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]
(3) 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 [3] 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 − 1 i =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[21].
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]
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[8, 2], 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 [20] 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[30, 18].
We provide in table 1 a simple classification of these approaches depending on whether they usea representation and memory component, and whether the inference mechanism is amortized oriterative.
[Table]
Methods
In this section, we present the auto-associative memory models that we have designed using thisframework.
All these models are based on a pre-trained VAE that is used to provide initial estimates˜ z and predictions of x based on ˜ z .
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]
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 eachlayer two populations of neurons, one encoding the layer’s representation, and one encoding the layerprediction error ϵ l = ˆ h l − f l θ ( ˆ h l +1 ) .
The forward pass through this model is given in algorithm 1.
Detailed derivations of this model are provided in appendix C. Since the RNN dynamics implement agradient descent, the initialization of this network is responsible for the local minimum to which itconverges.
We label VAE-PC-GMM the version of this algorithm where ˆ z is initialized using theresult of the amortized inference 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.
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, and Parameters: θ , α, σ, M Input: x Initialize (cid:8) ˆ z , ˆ h L − 1 , . . .
, ˆ h 1 (cid:9) for 0 ≤ t < T do /* Compute prediction errors */ for 0 ≤ l < L do ϵ l ← ˆ h l − f l θ ( ˆ h l +1 ) end / * Update hidden representations */ for 1 ≤ l < L do ˆ h l ← ˆ h l + α (cid:0) f l − 1 ′ θ ( ˆ h l ) · ϵ l − 1 − ϵ l (cid:1) end / * Update z */ ˆ z ← ˆ z
+ α (cid:0) f
L − 1 ′ θ (ˆ z ) · ϵ L − 1 (cid:1) + ασ 2 (cid:16) softmax (cid:0) − ∥
ˆ z − M ∥
22 2 σ(cid:1) · M ⊺ − ˆ
[Equation]
(cid:17) end that 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 is toretrieve scenes using as input a shifted image of the same scene, or an image of the same scene wherethe colors of the objects have been modified.
To properly accomplish these two tasks, the associativememory model needs to learn to give less importance to position information in the first case, and tocolor information 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 manually chosen.
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 [15](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 [13] (CC BY 4.0 License) consists of 64 × 64 RGB images of 3D scenescomposed of simple 3D objects.
We use a pretrained MONet [3] 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 1.
5.2 Benchmark models
We compare our associative memory models with Modern Continuous Hopfield Networks (MCHN),
Neural Turing Machines (NTMs) and End-to-End Memory Networks (MemN2N).
For fair comparison, in all benchmark models, the memory is initialized with the stored patterns (we do not use thewriting mechanisms of NTMs).
We also experiment with these models in the representation space,where the memory is initialized with representations of the stored images, and the networks receiveas input the representation 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 priordistribution, that we denote GMM in our experiments.
Equivalently, this is the version of ouralgorithm without a representation component.
5.3 Corrupted inputs
First [Table] we evaluate whether the auto-associative memory models are able to properly retrieve storedpatterns when given corrupted versions of these patterns.
We apply different level of noise to theimages ( σ ∈ { 0 . 2 , 0 . 4 , 0 . 6 , · · · , 2 . 0 } ), as well as different sizes of masks (16 and 24 pixels wideon the CIFAR10 dataset
[Table] and 32 and 48 pixels wide on the CLEVR dataset).
The results of theseexperiments are 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 associative memory model working on the representation level could perform well when we applytransformations on these specific object features.
We propose two such scenarios.
In the first scenario, we perform a RGB rotation on the images, which results in identical scenes with the exception of theobject colors.
In the second scenario, we take a shifted crop of the original CLEVR image, which cansimulate a shifted point of view of the same scene.
Examples images are displayed in figure 1.
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 associative memory models with high retrieval accuracy and robustness.
Combining representation and memory retrieval allows to retrieve patterns on the representationlevel. First, this is interesting with regard to the size of the memory store, which is proportional tothe dimensionality of the encoded patterns instead of the dimensionality of the sensory (here visual)space. Second, this allows to perform retrieval based on a similarity measured in the representationspace, where we can expect more meaningful features to appear.
For instance, we have shown thatthis allows to recognize visual scenes from a shifted point of view, while pixel-level associativememory models failed most of the time.
Our results also seem to demonstrate that distance basedattention as used in our models (see equation 11) always outperforms dot-product based attention asused in MCHNs.
An open question is whether this insight transfers to other deep learning dot-productbased architectures such as transformers.
Another work [24] has suggested implementing auto-associative memory with PC networks.
However,their method do not rely on a prior probability on representations z to form attractors, but insteadembed the stored patterns as minima of the VFE through training.
In contrast, our methods allowone-shot writing: a new pattern x can be inserted in memory simply by computing its representationand adding it to the memory matrix M .
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 mechanismsthat allow to only store relevant information in the memory.
Another limitation of this work is thatwe did not investigate the memory capacity and convergence properties of the proposed models.
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 credited in 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]