G ENERATIVE P SEUDO -I NVERSE M EMORY
A BSTRACT
We propose Generative Pseudo-Inverse Memory (GPM), a class of deep generative memory models that are fast to write in and read out.
Memory operations are recast as seeking robust solutions of linear systems, which naturally lead to the use of matrix pseudo-inverses.
The pseudo-inverses are iteratively approximated, with practical computation complexity of almost O (1) .
We prove theoretically and verify empirically that our model can retrieve exactly what have been written to the memory under mild conditions.
A key capability of GPM is iterative reading, during which the attractor dynamics towards ﬁxed points are enabled, allowing the model to iteratively improve sample quality in denoising and generating.
More impressively, GPM can store a large amount of data while maintaining key abilities of accurate retrieving of stored patterns, denoising of corrupted data and generating novel samples.
Empirically we demonstrate the efﬁciency and versatility of GPM on a comprehensive suite of experiments involving binarized MNIST, binarized
Omniglot, FashionMNIST, CIFAR10 & CIFAR100 and CelebA.
I NTRODUCTION
Memory is central to intelligence by facilitating information compression, reconstruction, manipulation, and generation.
The processing speed and storage capacity of the working memory are known to correlate with reasoning capacity (Jensen and Munro, 1979; Kyllonen and Christal, 1990).
Recent work in machine learning has explored slot-based external memory to augment neural networks (Graves et al., 2016;
Le et al., 2019) in which memory reads and writes proceed sequentially through attention mechanism, making training and pattern retrieval difﬁcult for very long sequences.
Global self-attention techniques (Ramsauer et al., 2021) create short-cuts in information paths, hence are easier to train, but require large memory and computation for long sequences.
Different from slot-based memories, Kanerva Machines (Wu et al., 2018a;b) are a class of generative memory inspired by Kanerva’s sparse distributed memory (Kanerva, 1984).
Here memory update/retrieval and addressing mechanism are treated as Bayesian inference where posteriors are updated when a new data episode arrives.
This iterative and dynamic inference handles noisy inputs better due to the convergence of the attractor dynamics to ﬁxed points (Wu et al., 2018b).
However, these models suffer from slow processing speed, which originates from the sequential writing mechanism.
This slowness prevents the model to adapt to large batch of data, thus makes it difﬁcult to conduct experiments on memory storage capacity.
On the other side, one cannot ensure theoretically how well those models retrieve information.
This creates an obstacle for theoretical research on the efﬁciency of memory models.
To overcome these challenges we propose a new model called Generative Pseudo-Inverse Memory (GPM), which is illustrated in Fig. 1.
GPM reformulates the Bayesian updates of memory and address as ﬁnding least-square solutions to linear systems.
Among these solutions, the smallest norm is found through applying the Moore-Penrose pseudo-inverse of matrices (Ben-Israel and Greville, 2001).
This helps us achieve rapid and accurate memory read/write.
Moreover, our model is able to not only store and retrieve information perfectly under ideal conditions but also generate new samples based
[Figure]
encoder
Z
W 0 M
M
W encoder
X q
Z q decoder Z r
X o iterative reading write read
[Figure]
on written data.
This coincides with recent inspiring results from functional neuroimaging (Addis et al., 2007; Okuda et al., 2003; Schacter et al., 2012), which showed that memory not only helps us memorize but also supports imagining.
We derive an energy function for GPM and show that its updates seek to locate local minima of the energy.
This dynamic enables GPM to perform iterative reading to improve sample quality during denoising and generating.
Finally, GPM can store a large amount of data (up to a dataset) while still maintaining key capabilities.
We validate these theoretical insights through a comprehensive suite of experiments on binarized MNIST (LeCun et al., 2010) , binarized Omniglot (Burda et al., 2016), FashionMNIST (Xiao et al., 2017), CIFAR10 & CIFAR (Krizhevsky, 2009) and CelebA (Liu et al., 2015), demonstrating superior results.
M ODEL
Consider a data episode of length T and let X denote the episode data matrix of T rows, where the i -th row of X is the i -th item of the episode.
Throughout this paper, we refer to the episode as its associated episode matrix.
We aim to compress the episode into a ﬁxed-size memory for later reconstruction and generation processes.
We maintain a memory matrix M , which is a deterministic real matrix of size K × C with trainable initialization M 0 .
For convention, M 0 may be considered as the common prior memory containing information of the whole dataset while M is the posterior memory dependent on the episode.
Remark.
In this paper, we reuse classical terms "prior memory" and "posterior memory", which are frequently used in prior works where the memory is modeled as a distribution.
In this paper, however, we only consider deterministic memories with an exception when we introduce the optimization problem corresponding to the Bayesian inference (Eq.
Hence, “prior memory” and “posterior memory” should be interpreted as the memory before and after data arrive, respectively.
In order to write to and read from the memory, we additionally deﬁne the writing weight matrix Wand reading weight matrix 1 W , both of size T × K .
We treat the reading weight W as a random variable, hence inducing the generative capability of the model.
For simplicity, the prior p (vec( W )) of W is the standard Gaussian N ( 0 , I T
× K ) while the posterior q (vec( W )) is N (cid:0) W ,
σ 2 W I T × K (cid:1) , where the mean W is adapted to each episode and σ W is a learnable parameter.
Read-outs from the memory are matrix products between the reading weight and the posterior memory.
Let e θ ( . ) be the (neural) encoder and d θ ( . ) be the decoder parameterized by θ .
In the following we drop the explicit subscript θ for clarity.
We further let Z denote the T × C encoding matrix of the episode X , i.e. Z = e
( X ) where e ( . ) is applied to each row of X .
During training, the model can only observe a noisy version of Z , which we denote by Z ξ .
Speciﬁcally, Z ξ = Z + ξ , where ξ ∼ N ( 0 , σ 2 ξ I ) is the observation noise.
We assume that the observed data Z ξ depends linearly on the weights and the memories: Z ξ = W 0 M 0 and Z ξ = WM , which we refer as linear assumptions over Z ξ .
2. O BJECTIVE FUNCTION
Given the memory M , we aim to maximize the conditional log-likelihood ln p ( X | M ) .
It is difﬁcult to maximize ln p ( X | M ) directly, so we instead maximize its evidence lower bound L , where
[Equation]
The ﬁrst term in L is usually called the negative reconstruction loss, while the second term is the
Kullback-Leibler divergence between the posterior q ( W ) and the prior p ( W ) of the reading weight.
Details of derivation will be given in Appendix I.
In our model, we use the encoding matrix Z as the data to be written to the memory instead of the original episode X .
Therefore, we would like Z to be a “good” representation of X .
We quantify this goodness by the usual auto-encoder loss L AE =
E X ∼ data ln p ( d ( e ( X ))) .
Thus during training, we will maximize the objective function:
[Equation]
2. G ENERATIVE P SEUDO -I NVERSE M EMORY
In our model, the posterior memory, the writing weight and the reading weight are all dynamic (i.e. depending on the data episode).
Hence it is necessary to understand the motivations and computation methods as introduced in the following.
Intuition
Only in this part, we will consider probabilistic memories to have uniﬁed settings with previous
Bayesian-based works such as the Kanerva Machines (Wu et al., 2018a;b; Marblestone et al., 2020).
When the observed data Z ξ arrive, the Kanerva Machines apply the Bayesian inference process to obtain the posterior memory from a given (trainable) prior distribution.
It is well-known that the posterior mean is the Bayes estimator with respect to the squared error risk (Jaynes, 2003).
Formally, given the observed data Z ξ and the writing weight W 0 , the posterior memory mean M is the solution of the optimization problem
[Equation]
where (cid:107) .
(cid:107) F is the Frobenius norm.
With the linear assumption over Z ξ and an additional linear assumption over M , Eq.
(3) has a unique and explicit solution (see Appendix H).
The solution M of Eq.
(3) tends to preserve information from M 0 .
However, in practice, what appears in the prior memory might be blurry information.
Forcing the short-term memory M to be close to M 0 may prevent M from storing sufﬁcient information, leading to incorrect read-out.
Memory inference
From now on, we will only consider deterministic memories.
We propose a weighted version of the optimization problem (3) to help improve read-out accuracy.
Our proposed optimization problem is
[Equation]
[Equation]
Intuitively, if M is a solution of Eq.
(4), then W 0 M is close to W 0 M 0 , which is equal to Z ξ due to the linear assumption over Z ξ .
In other words, the read-out after writing (i.e. W 0 M ) is similar to the data written to the memory (i.e. Z ξ ).
Substitute W 0 M 0 =
Z ξ , Eq.
(4) becomes
[Equation]
Any solution of Eq.
(5) is called the least-square solution of the linear system W 0 M = Z ξ .
However,
Eq.
(5) may admit more than one least-square solution, thus we are interested in the least-square solution that has the smallest norm, i.e. the minimum-norm least-square solution.
The minimumnorm condition may help to regularize in the sense that the posterior memory should contain as little redundant information as possible.
It can be proved that M = W +0 Z ξ is the minimum-norm least-square solution of the system W 0 M =
Z ξ (Ben-Israel and Greville, 2001), where W +0 is the pseudo-inverse of W 0 .
The pseudo-inverse of matrix is applicable to classical associative memory which involves ﬁnding a linear mapping from input to output (Stiles and Denq, 1985; Yen and Michel, 1991).
Readers are referred to Appendix D for more mathematical properties of the pseudo-inverse.
Computing writing and reading weights (dynamic addressing)
The discussion so far is about seeking the posterior memory M given the data Z and the writing weight W 0 .
In our settings, all weights are dynamic and they depend on either the prior memory and data (during writing) or the posterior memory and the query (during reading).
Formally written, we need to determine the writing weight W 0 (given prior memory M 0 and noisy data Z ξ ) and reading weight W (given posterior memory M and query Z ) 2 .
Ideally, W 0 and W should satisfy
W 0 M
=
Z ξ and WM =
Z ξ ≈ Z , according to the linear assumptions over Z ξ .
The minimumnorm least-square solutions of those linear systems are W 0 =
Z ξ M +0 and W =
ZM + , where
M +0 and M + are pseudo-inverses of M 0 and M , respectively.
While least-square condition is obvious since we would like to ﬁnd solutions W 0 and W that best ﬁt the systems, explanation for the minimum-norm condition will be given in Appendix J.
Computing pseudo-inverses
We do not always have an explicit formula for the pseudo-inverse, and even in such cases, the formula may consist of inverse matrices which are expensive to compute.
Ben-Israel and Cohen (1966) provide us with an iterative way to approximate the pseudo-inverse, as stated in the following theorem:
Theorem 2.1.
(Ben-Cohen algorithm)
Given a real matrix D and an initial matrix D 0 of the same size with the transpose D (cid:62) .
The sequence { D i } i ≥ 0 deﬁned recursively as D i +
= 2 D
i − D i DD i will converge to the pseudo-inverse of D .
With appropriate D 0 , the sequence will converge quadratically.
Ben-Cohen algorithm allows us to compute the matrix pseudo-inverse efﬁciently with nearly O (1) time complexity in practice (see Section 3.5 for experimental results).
For simplicity, we set the initial term D
= α D (cid:62) , where α is a hyper-parameter dependent on the dataset.
Deeper analysis of the initial condition for Ben-Cohen algorithm will be given in Appendix G.
Algorithm We are now ready to present the Generative Pseudo-Inverse Memory (GPM) in full.
Assume that the observation noise ξ is sampled from N ( 0 , σ 2 ξ I ) .
Algorithm 1 illustrates a single training step of GPM.
All pseudo-inverses are approximated by the Ben-Cohen algorithm.
2.2. R EMARK : R ELATION WITH EXISTING ALGORITHMS
Several properties of GPM are worth highlighting.
GPM can be viewed as 2-step EM algorithm (Dempster et al., 1977).
Details are given in Appendix J. Moreover, in noiseless setting, i.e., ξ → 0 ,
GPM and Dynamic Kanerva Machine (DKM) (Wu et al., 2018b) coincide.
In DKM, solving W is related to minimizing D KL ( q ( W ) ||
p
( W | X , M )) , which is equivalent to the optimization problem
[Equation]
This problem has a solution W σ ξ =
ZM (cid:62) ( MM (cid:62) + σ 2 ξ I ) − 1 .
According to Theorem D.5, W σ ξ converges to W ∗ =
ZM + as σ ξ → 0 .
Note that W ∗ is actually the reading weight of GPM.
2. E RROR BOUND
Using pseudo-inverses in Algorithm 1, we establish an error bound for the memory read-out, as stated in the following theorem: Writing
Compute episode embedding Z = e ( X ) .
Randomize noise ξ from N ( 0 , σ 2 ξ I ) .
// simulate observed data Compute weight W 0 = Z ξ M +0 . // dynamic addressing
Compute posterior memory M = W +0 Z ξ .
// complete writing phase
Reading
Compute episode embedding Z = e ( X ) .
// similar role with
Z
q in Figure 1b
Compute weight mean W
=
ZM + .
// dynamic addressing
Sample W ∼ N ( W , σ 2 W I ) .
// variational schema for generating samples
Compute read-out Z read-out = WM .
// similar role with Z r in Figure 1b
Compute reconstruction ˆ X =
d ( Z read-out ) .
// complete reading phase, no iterative reading
Updating model parameters
Compute the objective O =
L + L AE using obtained terms from previous steps.
Update parameters via gradient ascent to maximize O .
Theorem 2.2.
Let (cid:107) .
(cid:107) 2 denote the spectral norm.
Suppose there exist α, β ∈ [0 , 1) such that
[Equation]
Then WM =
Z ( I + E ) , where E is the error matrix satisfying
[Equation]
Corollary 2.3.
Assume that rank W 0 = rank Z ξ = T , i.e. rows of W 0 and Z are linearly independent.
Then WM =
Z .
Proofs are given in Appendix E. Theorem 2.2 and Corollary 2.3 establish that the error of memory retrieval can be bounded by other errors only dependent on the data Z and writing weight W 0 .
When rows of W 0 and Z ξ are linearly independent, the retrieval is perfect .
The trainable prior memory
M 0 helps the model balance approximation errors when ideal conditions are not met.
Corollary 2.3 also suggests that retrieval may be still accurate if rows of W 0 and Z are sufﬁciently linearly independent.
A well-known property on high dimensional space is that two arbitrary vectors are likely to be orthogonal (Gorban and Tyukin, 2018).
Thus we can expect that GPM will still work fairly well when T (cid:29) max { K, C } , i.e. the episode length is much greater than the number of memory slots and size of embedding vectors.
Empirical evidences will be given in Section 3.4.
2. A TTRACTOR DYNAMICS
We demonstrate that GPM inherits properties of energy-based models.
The key property of models induced from dynamical systems is iterative reading, i.e. to iteratively reproduce the correct pattern from a partially broken query.
With ﬁxed memory M , the energy function of GPM can be deﬁned as a function of two variables x and w :
[Equation]
Here we replaced z by e ( x ) and note that the residual r is constant after training.
Denote ( w i , x i ) the pair of weight and reading output of the i -th iterative reading step.
The output x i is iteratively fed back to the model until convergence.
With x i as the query for ( i + 1) -th step, solving for w i + minimizes (cid:107) e ( x ) − wM (cid:107) 2 F , which indicates E ( x i , w i +1 ) >
E ( x i , w i ) .
With w i +1 in hand, the way we train the model encourages x i +1 to maximize the likelihood p ( x | w i +1 , M ) .
If it is the case,
E ( x i +1 , w i +1 ) >
E ( x i , w i +1 ) and therefore E ( x i +1 , w i +1 ) >
E ( x i , w i ) .
In ideal cases, { x i } i ≥will converge to a local minimum associated with a stored pattern in the energy landscape.
[Table]
In practice [Table] the system might not converge to a stored pattern but a spurious state, which is a local minimum different from stored patterns.
Even with simple neural networks in low dimensional space, spurious local minima also exist (Safran and Shamir, 2018).
This phenomenon is not expected when retrieving, but it may help generate new images as demonstrated in Section 3.3.
E XPERIMENTS
We examine GPM on common generative memory benchmarking datasets including binarized MNIST (LeCun et al., 2010), binarized Omniglot (Burda et al., 2016), FashionMNIST (Xiao et al., 2017), CIFAR10, CIFAR100 (Krizhevsky, 2009) and CelebA (Liu et al., 2015).
To prove the effectiveness of our memory model, we only use simple encoders and decoders.
The details of image preprocessing, network architecture and training procedure are given in Appendix C. Codes are available at https://github.com/phamtienkha/generative-pseudoinverse-memory.
We compare our model to recent generative memory models such as Kanerva Machine (KM),
Dynamic Kanerva Machine (DKM) and Kanerva++.
In GPM, most of the parameters come from the encoder and the decoder since the trainable prior memory M 0 and the weight variance σ 2 W only account for only KC + 1 parameters.
Hence, the number of parameters of GPM is comparable to other baselines with similar architectures of encoder and decoder.
3. R ECONSTRUCTION
In this task, we validate the memory model’s ability to reconstruct image data.
For each dataset, we sample 32 images to form the episode as the input to the memory models during both encoding and decoding phases.
We then train all the models to optimize the objective in Eq.
(2) for 10,000-15, iterations, depending on the dataset.
Table 1 reports the negative evidence lower bound (Eq.
(1)) of
GPM and that of other baselines on 3 datasets, including the binarized MNIST, binarized Omniglot and CIFAR10.
Results for the FashionMNIST, CIFAR100 and CelebA datasets are reported in
Table 3 in Appendix A. For binarized datasets, we see a signiﬁcant improvement in the ELBO: 31.48 nats/image compared to current state-of-the-art 41.58 nats/images on MNIST dataset and 25.68 nats/image compared to current state-of-the-art 66.24 nats/image on Omniglot dataset.
On
RGB datasets, GPM shows competitive result.
We note that for complex image data, the role of the encoder and decoder is very important and GPM’s ones are simpler than those of Kanerva++.
Hence, we underperform Kanerva++ on CIFAR10 dataset.
However, compared to models with similar encoder-decoder like KM and DKM, GPM performs much better, which indicates the beneﬁt of our pseudo-inverse memory mechanisms.
To verify that our GPM supports exact memory retrieval, we measure the cosine similarity between the data written to the memory and memory read-out’s of GPM in Omniglot task (see Fig.
6 of
[Table]
Appendix).
As expected, GPM almost retrieves exactly what was written to the memory (average similarity ≈ 1 ), which is much better than that of Dynamic Kanerva Machine ( 0 . 82 ).
3. D ENOISING
The denoising capability of a memory model resembles human’s ability to retrieve old memories with blurry clues.
This process can be implemented using iterative reading, where in each step the model adds some more details to the image to reconstruct the original picture from a noisy query image.
Illustration for the denoising process of GPM is given in Appendix B.
To measure the denoising capability of GPM and DKM, we run denoising experiments with different episode lengths, types of noise and levels of noise on the binarized Omniglot dataset.
We use the
Hamming distance to measure the error between denoised images and correct ones.
In all experiments, we run iterative reading for 20 steps.
We report results for the case T = 16 in Fig.
2, while results for other cases can be found in Fig.
8 of Appendix.
Overall, GPM achieves a signiﬁcantly better performance than that of DKM.
Remarkably, GPM maintains a clear performance gap when the episode lengths or the level of noise vary.
To better quantify the difference, we reported the denoising success rate of GPM and DKM in Table 2.
Here we consider a denoising process to be successful if it can recover the original image with at most 1% error, which is equivalent to at most 7 pixels for the
Omniglot dataset.
The results show that GPM also outperforms DKM in all cases with a performance gap ranging from 15% to 58%.
salt & pepper
[Figure]
3. G ENERATION
We illustrate the ability to generate new images of GPM using Omniglot and CIFAR10 images in
Fig. 3a.
The generation process consists of 2 main steps: (1) sample images to form an episode and write it to the memory; (2) sample w ∼ p ( w ) ( w is a vector of size 1 × K ) and take z = wM to be the read-out vector.
The generated images will be improved via iterative reading.
We observe that GPM can generate images that do not appear in the written episode.
This phenomenon can be viewed in Fig.
3a-left, where we compare generated patterns with their most similar ones in terms of
Hamming distance.
Here the generated images with additional details are generally different from their nearest neighbors.
More interestingly, in Fig.
3a-right, we can see a generated aircraft image that inherits the background of the frog image (the 3rd row).
Theoretically speaking, these new images are formed from spurious local minima, as discussed in Section 2.4.
3. GPM AS A HUGE STORAGE
We conduct experiments to test the memory capacity of GPM.
The experiments still revolve around reconstructing, denoising and generating images, yet on larger episode length.
We use the whole
Omniglot dataset with 24 , 345 training images to train GPM for around 24 , 000 iterations for all [Figure] experiments.
For convention, we denote the T -GPM as the GPM trained on episodes of length T .
Results are shown in Fig.
In the reconstruction experiment, we keep the memory size ﬁxed and increase the episode length from minimum (equivalent to the number of memory slots) to maximum ( 24 , 345 , which is the whole set of training images).
We run the experiment with two memory conﬁgurations: 8 × 100 and 32 × 100 .
With the same episode length, the 32 × 100 memory always performs better, which is reasonable since bigger memory enables better storage capability.
We also observe a more interesting phenomenon in both conﬁgurations: from a certain point, the reconstruction loss seems to be stable even when we double the episode length.
This agrees with earlier discussion in Section 2.2, where we argue that
GPM can handle big episodes thanks to the orthogonality of random vectors in high-dimensional space.
We also test the denoising and generating capabilities of GPM when trained with a large episode size.
We are curious whether GPM can keep these abilities (to some extent) when storing lots of information.
We use the 24 , 345 -GPM with 32 × 100 memory in these experiments.
At test time, we ﬁrst write the whole test set of 8070 images to the memory and keep this posterior memory ﬁxed, then we either input noisy queries for denoising or sample from the memory for generating images.
Despite the huge amount of information compressed into the memory, the model can still denoise and generate images with moderate quality.
The average Hamming error when denoising is 97 .
08 , which is approximately 12 .
4% of the total pixels.
We note that other baselines such as KM and DKM are inapplicable to such a large-scale scenario due to their slow memory mechanisms.
In contrast, GPM can be trained efﬁciently even with episodes of 24 , 345 images.
More experiments regarding this point are conducted in the next section.
[Figure]
3. A BLATION : RUNNING TIME
GPM is designed to avoid the computation complexity
O ( K 3 ) of inverting matrices and O ( T ) of online Bayesian inference.
We conduct experiments to show the time efﬁciency of GPM compared to
DKM in practice.
All operations are computed on a single GPU.
We use the inverse function of
Pytorch 1.8.0 (Paszke et al., 2017) for batch matrix inverse.
Results are shown in Fig.
3b.
In the ﬁrst experiment, we run GPM with 7 approximation steps for the pseudo-inverses and compare the running time against DKM in different cases of episode and memory size.
As expected, the running time of DKM increases as the memory size K increases while GPM maintains almost the same performance.
Moreover, we can see a problem of Bayesian update on the running time: if we double the episode size and keep the same memory size, the running time also doubles.
In contrast, it only yields a modest increase in the running time of GPM.
In the second experiment, we would like to see how the number of pseudo-inverse approximation steps affects the running time of GPM in two cases: T = 64 , K = 64 and T = 128 , K = 64 .
The running time curves in the two cases share a similar trend: they depend linearly on the number of approximation steps.
Hence, it is important to choose an appropriate initial term for the approximation sequence to reduce the number of approximation steps.
More discussion is given in Appendix G.
R ELATED WORK
Memories are known to be critical to tasks that demand long-term dependencies (Hochreiter and
Schmidhuber, 1997; Graves et al., 2016; Santoro et al., 2016; Vaswani et al., 2017; Le et al., 2019;
Munkhdalai et al., 2019).
When all stored data are available to retrieve through attention mechanism (Sukhbaatar et al., 2015; Vaswani et al., 2017), we indeed perform retrieval over a content-addressable memory such as Hopﬁeld networks (Hopﬁeld, 1982; Ramsauer et al., 2021) and Kanerva’s
sparse distributed memory (SDM) (Kanerva, 1984;
Bricken and Pehlevan, 2021).
These works are primarily designed for storing and retrieval of patterns, mostly in a deterministic manner, and thus less concerned about noise and data generation (Le et al., 2018).
Kanerva machines (Wu et al., 2018a), inspired by the addressing mechanism of SDM, are generative memory models which ﬁll these gaps.
Here memory update/retrieval and addressing mechanism are treated as
Bayesian inference where posteriors are updated when a new data episode arrives.
This iterative and dynamic inference handles noisy inputs better due to the convergence of the attractor dynamics to ﬁxed points (Wu et al., 2018b).
The Kanerva machines thus belong to a class of latent variable models with a very expressive latent prior.
These generalize Variational AutoEncoders (VAE) (Kingma and
Welling, 2014), deep encoder-decoder neural networks that can generate high-quality data in high dimensional space.
VAEs are very powerful when equipped with a high-capacity decoder (Gregor et al., 2015; Bachman, 2016; Ma et al., 2019; Sadeghi et al., 2019) and a expressive prior (Chen et al., 2017; Graves et al., 2018; Tomczak and Welling, 2018).
Here memory-based priors in Kanerva machines offer a systematic approach to bring more representation power to VAEs.
However, Bayesian inference in the original Kanerva machines is expensive.
To address this drawback,
Marblestone et al.
(2020) introduced the Product Kanerva Machine to reduce the heavy cost of inverting matrices in the original model.
An enhanced variant known as Kanerva++ (Ramapuram et al., 2021) uses a more powerful encoder and regional memory writing-reading to help improve both the retrieval accuracy and training time.
Our work takes a radically alternative view in that we treat memory update as seeking least-square solutions to linear systems.
This permits the use of fast matrix pseudo-inverse for the memory operations (Ben-Israel and Greville, 2001).
Indeed the pseudo-inverse has been studied in the associative memory literature.
Examples include Kohohen-type linear associative memory, which can be referred as a pseudo-inverse neural network (Kohonen, 2012).
Personnaz et al.
(1986) proposed the projection rule (other than Hebb’s learning rule in standard Hopﬁeld Network) involving matrix pseudo-inverses to guarantee perfect storage and retrieval.
However, the equilibrium point produced by the projection rule may not be asymptotically stable in the sense that adding new vectors may affect the existing equilibria in the network.
To overcome this problem, Yen and Michel (1991) introduced a learning and forgetting algorithm with efﬁcient computation.
As far as we are aware, the pseudo-inverse is not commonly used in deep memory networks partly because of the difﬁculty to compute it exactly, although several attempts have been made to improve the computation efﬁciency (Greville, 1960; Ben-Israel and Cohen, 1966; Courrieu, 2008; Toutounian and Ataei, 2009).
D ISCUSSION
We have proposed Generative Pseudo-Inverse Memory (GPM) to overcome crucial limitations of previous generative memory models.
GPM is not only able to operate fast but also store a large amount of data at a time while maintaining abilities of denoising corrupted patterns and generating novel patterns.
We derived an associated energy function and showed that iterative updates in GPM induce attractor dynamics converging toward ﬁxed points.
We demonstrated these capacities on an extensive set of experiments, obtaining superior results.
These new capacities of GPM as a generative memory open up new exciting rooms of investigation.
A primary direction is positioning GPM as working memory whose size and processing speed are known to be critical for high-level reasoning (Kyllonen and Christal, 1990), i e., the capacity to deduce new knowledge through manipulating the old.
Here the query encoder will support arbitrary query format and the decoder will support answer generation.
As GPM is designed for episodic handling, streaming data can be periodically cached through local registers before sending to GPM, as suggested in slot-based memories (Le et al., 2019).
Finally, GPM may offer fast episodic learning and generative exploration in RL agents.
APPENDIX A R ECONSTRUCTIONS
We provide visualizations for the reconstructions of GPM on numerous datasets in Figure 5.
We also provide empirical evidence to show that the optimization problem associated with GPM (Eq.
(4)) is more direct than DKM (Eq.
(3)) in Figure (6): the cosine similarity between stored vectors and read-out’s of GPM is nearly 1, which is much better than DKM (around 0.82).
We report the ELBO of DKM and GPM over 3 datasets, including FashionMNIST, CIFAR100 and CelebA in Table 3.
B D ENOISING
We demonstrate denoising capibility of GPM in Fig.
The denoising process consists of two steps: we ﬁrst write the correct patterns into the memory, then read from the memory using noisy queries.
The image will be completed during iterative reading.
We also report the full results of the denoising experiment described in Section 3.2 in Figure 8.
[Figure]
C T RAINING DETAILS
The encoder consists of 4 layers, each of which is a convolution layer with 4 × 4 ﬁlter with stridefollowed by a Resnet block with bottleneck (He et al., 2016).
The decoder is simply a mirror of the encoder with transpose convolutional layer.
We use the swish activation function (Ramachandran et al., 2017) non-linear layers.
For binarized datasets, we use 16 ﬁlters for all convnets; for real-valued datasets, we use 256 ﬁlters to adapt with the complexity of the data.
The Omniglot is a dataset with great diversity, consists of 1623 different classes and 20 images in each class.
Following settings in (Burda et al., 2016), the dataset is splitted into 24,345 training and 8,070 test examples.
We use a 32 × 100 memory for this dataset.
The data noise when writing is sampled from N ( 0 , 0 .
5 2 I ) .
At each training step, we randomly sample 32 images to form an episode.
We run the Ben-Cohen algorithm for 7 steps to approximate the pseudo-inverses, with the initial term is 10 − 3 times the transpose of the matrix which we want to calculate the pseudo-inverse.
We use the
Bernoulli likelihood for all binarized datasets.
For real-valued datasets, we use a 128 × 512 memory with data noise sampled from N ( 0 , 0 .
5 2 I ) .
All other settings are the same with the Omniglot dataset, except for the initial terms of the Ben-Cohen algorithm: 10 − 4 times the transpose of the matrix when approximating pseudo-inverse of the weight matrix and 5 · 10 − 5 times the transpose of the matrix when approximating pseudo-inverse of the memory.
We use the discretized mixture of logistics (Salimans et al., 2017) for the output distribution.
In all experiments, we use the Adam optimizer with learning rate varying from 5 · 10 − 5 to 5 · 10 −depending on the dataset.
We use weight decay of 10 − 3 along with gradient clipping at threshold 10 .
For real-valued datasets, we additionally use a scheduler to reduce the learning rate by 0 .
99995 after each training epoch.
We train all models for 500-1000 epochs with exception 30 epochs for CelebA dataset.
D P ROPERTIES OF PSEUDO - INVERSE
In this section, we outline the deﬁnition and some important properties of the pseudo-inverse of a matrix.
Most of the proofs of following propositions and theorems can be found in (Ben-Israel and
Greville, 2001).
Deﬁnition D.1.
For every real matrix A (square or rectangular), there is a unique matrix Y satisfying all following conditions:
AYA =
A , YAY = Y , ( AY ) (cid:62) =
AY and ( YA ) (cid:62) =
YA , where (cid:62) denotes the transpose of a matrix.
The matrix Y is called the pseudo-inverse of A , denoted by A + .
If A is a square matrix and non-singular, then A + is the inverse A − 1 of A .
If A is rectangular of size m × n , then A + is of size n × m .
The following propositions shows some important relations between A and A + .
[Figure]
H a mm
i ng e rr o r
H a mm
i ng e rr o r
H a mm
i ng e rr o r
Proposition D.2.
Let A be a m × n real matrix.
Then
A + A =
I n if and only if rank A = n (i.e. columns of A are linearly independent), where I n is the n × n identity matrix; 2.
AA + = I m if and only if rank A
= m (i.e. rows of A are linearly independent), where I m is the m × m identity matrix.
Proposition D.3.
For every real matrix A :
( A + ) + =
A ; 2.
(cid:0) A (cid:62) (cid:1) + =
( A + ) (cid:62) ; 3.
A + = (cid:0) A (cid:62) A (cid:1) + A (cid:62)
=
A (cid:62) (cid:0) AA (cid:62) (cid:1) + .
If A and B are square invertible matrices such that AB is also invertible, then we have the wellknown property ( AB ) − 1 =
B − 1 A − 1 .
The same property does not always hold for pseudo-inverse.
However, with some sufﬁcient condition of independence, we will get the desired property.
Proposition D.4.
If A has linearly independent columns or B has linearly independent rows, then ( AB )
+ =
B + A
+ .
Theorem D.5.
For any real matrix A of size m × n , the limit lim σ (cid:38)
A (cid:62) ( AA (cid:62) + σ I m ) − 1 always exists.
Moreover, lim σ (cid:38)
A (cid:62) ( AA (cid:62) + σ I m ) − 1 =
A + .
The concept of pseudo-inverse arose when ones tried to solve the linear system
Ax = b , where b is some vector with appropriate size.
When the system has no solution, it is reasonable to ﬁnd some vector x that minimizes the Euclidean norm of the residual r =
b − Ax .
Any solution of the optimization problem min x (cid:107) b − Ax (cid:107) F (where (cid:107) .
(cid:107) F is the Frobenius norm) is called the leastsquare solution of the system
Ax = b .
It can be proven that x
= A
+ b is a least-square solution.
Moreover, among the least-square solutions of system Ax = b , x = A + b is the one of minimum
Euclidean norm, i.e. the minimum-norm least-square solution .
This result can be generalized as in the following proposition.
Proposition D.6.
Let A , B , D be real matrices with appropriate size.
The minimum-norm leastsquare solution of the system AXB =
D is X =
A + DB + .
E P ROOFS OF THEORETICAL RESULTS ON ERROR BOUNDS
We provide the proofs for Theorem 2.2 and Corollary 2.3.
Proof of Theorem 2.2.
For convention, let us denote E W 0 =
W 0 W +0 − I .
Following Algorithm 1, we get
[Equation]
Since (cid:107) I − Z ξ Z (cid:62) ξ (cid:107) 2 ≤ α , by (Rump, 2011, Lemma 2.2), it follows (cid:107) Z ξ (cid:107) 2 ≤ √ 1 + α and (cid:107) Z
+ ξ (cid:107) 2 ≤ 1 √ 1 − α .
Since (cid:107) .
(cid:107) 2 is a multiplicative norm, we ﬁnally have
[Equation]
+ α [Equation] , which is our desired result.
Proof of Corallary 2.3 : This is a special case of Theorem 2.2.
When rows of W 0 and Z are independent
, W 0 W +0 =
I and Z ξ Z + ξ =
I , which leads to E W 0 = 0 and β = 0 .
F I TERATIVE READING FROM THE VIEW OF FIXED POINT
In addition to the discussion in Section 2.4, we can alternatively think about the attractor as the ﬁxed point of the reading operation.
Assume that the current query is q t and the next query is determined by q t
+1 = f
( q t ) , where f represents for the reading operation.
If { q t } converges to some q ∗ , then q ∗ = f
( q ∗ ) , which means q ∗ is a ﬁxed point of f .
If the model is well-trained, stored patterns are ﬁxed points of f .
(Banach, 1922) provides a sufﬁcient condition for an iterative sequence to converge to a ﬁxed point.
Roughly speaking, if q ∗ is a ﬁxed point of f such that (cid:107) J ( q ∗ ) (cid:107) 2 < 1 (where J is the
Jacobian of f ) and the whole sequence { q t } t ≥ 0 locates in a small neighborhood of q ∗ , then q t → q ∗ .
This result demonstrates the ability to converge to a correct pattern given a good enough initial query.
Similar to spurious local minima in the energy landscape, there also exists many spurious ﬁxed points in the f landscape (see Fig. 9).
This phenomenon is not expected when retrieving, but it may be helpful in case we want to generate new images as demonstrated in Section 3.3.
G I NITIAL CONDITION FOR B EN -C OHEN ALGORITHM
Suppose A is a real matrix.
With 0 < α < 2 σ 21 , where σ 1 is the largest eigenvalue of AA (cid:62) , the sequence deﬁned as A 0 = α A (cid:62) ,
A i +1 =
2 A i − A
i AA i will converge to A
+ (Ben-Israel and
Cohen, 1966).
This choice of initial term may lead to slow convergence when the condition number κ ( A ) = (cid:107) A (cid:107) 2 (cid:107)
A + (cid:107) 2 is large (Söderström and Stewart, 1974).
If we start with an A 0 such that
A
A =
( A 0 A ) (cid:62) and A 0 is already close to A + , for example
A
= A (cid:62) ( AA (cid:62) + δ I ) − 1 where δ > 0 and small enough, then { A i } i ≥ 0 converges quadratically to A + .
However, this requires us to calculate the inverse of AA (cid:62) + δ I , which we tend to completely avoid in GPM.
Therefore we keep the original setting of Ben-Cohen for the initial term and let the model adapt with the approximation error.
H S OLUTION OF B AYESIAN UPDATE OPTIMIZATION PROBLEM
Let M 0 , M denote random variables of the prior and posterior memory, respectively, and suppose that both have matrix Gaussian distributions.
Speciﬁcally, M 0 ∼ N ( R 0 , U 0 , I ) and M ∼ N ( R , U , I ) .
Given the prior M 0 , the to-be-written data Z and the writing weight W 0 , the optimization problem
[Equation]
s.t.
M = AZ + B and Z = W 0 M 0 + ξ has a unique solution M ∼ N ( R , U , I ) , where
[Equation]
where Σ ξ is diagonal matrix whose diagonal elements are σ 2 ξ .
I D ERIVATION FOR EVIDENCE LOWER BOUND
The evidence lower bound of the conditional log-likelihood p ( X | M ) can be derived as
[Equation]
J R ELATION WITH EM ALGORITHM
We show that GPM can be viewed as 2-step EM algorithm (Dempster et al., 1977).
Let us remind ourselves the log-likelihood approximation derived in Section 2.1:
[Equation]
Keeping p (vec( W )) ∼ N ( 0 , I ) in mind, we can rewrite the conditional log-likelihood as
[Equation]
where R is residual term that is independent of W .
From the perspective of EM algorithm and
Eq.
((9)), Algorithm 1 can be thought as a 2-step process to maximize the log-likelihood: 1.
Keeping the likelihood unchanged and tightening the bound by minimizing (cid:107) Z − WM (cid:107) 2 F ; 2.
Maximizing E q ( W ) ln p ( X | W , M ) , i.e. the negative reconstruction loss.
Note that in training, we maximize the evidence lower bound L T =
E q ( W ) ln p ( X | W , M ) −
D KL ( q ( W ) || p
( W )) instead of E q ( W ) ln p
( X | W , M ) .
Therefore, for the maximum values of
L T and E q ( W ) ln p ( X | W , M ) to be as close as possible, we look for W that minimizes
D KL ( q ( W ) || p
( W )) .
Since p (vec( W )) ∼ N ( 0 , I ) , this is equivalent to seeking W with minimum norm.
This ﬁts with the step of ﬁnding the weight matrix W in the reading phase, where we determine W =
ZM + , which is the least-square minimum-norm solution of the linear system
WM =
Z .