Generative Pseudo-Inverse Memory

1 INTRODUCTION . 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 difficult 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 fixed 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 difficult to conduct experiments on memory storage capacity . On the other side , one can not ensure theoretically how well those models retrieve information . This creates an obstacle for theoretical research on the efficiency 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 finding 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 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 & CIFAR100 ( Krizhevsky , 2009 ) and CelebA ( Liu et al. , 2015 ) , demonstrating superior results . 2 MODEL . 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 fixed-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 M0 . For convention , M0 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 . ( 3 ) ) . 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 define the writing weight matrix W0 and reading weight matrix1 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 , IT×K ) while the posterior q ( vec ( W ) ) is N ( W , σ2WIT×K ) , 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ξ . Specifically , Zξ=Z + ξ , where ξ∼ N ( 0 , σ2ξI ) is the observation noise . We assume that the observed data Zξ depends linearly on the 1We note that weights here are dynamic and unrelated to neural network parameters . The term “ dynamic ” means the writing/reading weights of an input not only depend on its own but also the episode it belongs to . weights and the memories : Zξ = W0M0 and Zξ = WM , which we refer as linear assumptions over Zξ . 2.1 OBJECTIVE FUNCTION . Given the memory M , we aim to maximize the conditional log-likelihood ln p ( X|M ) . It is difficult to maximize ln p ( X|M ) directly , so we instead maximize its evidence lower bound L , where L = Eq ( W ) ln p ( X|W , M ) −DKL ( q ( W ) ||p ( W ) ) . ( 1 ) The first 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 LAE = EX∼data ln p ( d ( e ( X ) ) ) . Thus during training , we will maximize the objective function : O = L+ LAE . ( 2 ) 2.2 GENERATIVE PSEUDO-INVERSE MEMORY . 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 unified 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 W0 , the posterior memory mean M is the solution of the optimization problem min M E ( ‖M−M0‖2F ∣∣∣∣Zξ , W0 ) , ( 3 ) where ‖.‖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 M0 . However , in practice , what appears in the prior memory might be blurry information . Forcing the short-term memory M to be close to M0 may prevent M from storing sufficient 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 min M ‖W0 ( M−M0 ) ‖2F . ( 4 ) Intuitively , if M is a solution of Eq . ( 4 ) , then W0M is close to W0M0 , which is equal to Zξ due to the linear assumption over Zξ . In other words , the read-out after writing ( i.e . W0M ) is similar to the data written to the memory ( i.e . Zξ ) . Substitute W0M0 = Zξ , Eq . ( 4 ) becomes min M ‖W0M− Zξ‖2F . ( 5 ) Any solution of Eq . ( 5 ) is called the least-square solution of the linear system W0M = 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 Algorithm 1 Single training step of Generative Pseudo-Inverse Memory Sample an episode X of length T . Writing 1 . Compute episode embedding Z = e ( X ) . 2 . Randomize noise ξ from N ( 0 , σ2ξI ) . // simulate observed data 3 . Compute weight W0 = ZξM+0 . // dynamic addressing 4 . Compute posterior memory M = W+0 Zξ . // complete writing phase Reading 1 . Compute episode embedding Z = e ( X ) . // similar role with Zq in Figure 1b 2 . Compute weight mean W = ZM+ . // dynamic addressing 3 . Sample W ∼ N ( W , σ2WI ) . // variational schema for generating samples4 . Compute read-out Zread-out = WM . // similar role with Zr in Figure 1b 5 . Compute reconstruction X̂ = d ( Zread-out ) . // complete reading phase , no iterative reading Updating model parameters 1 . Compute the objective O = L+ LAE using obtained terms from previous steps . 2 . Update parameters via gradient ascent to maximize O. little redundant information as possible . It can be proved that M = W+0 Zξ is the minimum-norm least-square solution of the system W0M = Zξ ( Ben-Israel and Greville , 2001 ) , where W+0 is the pseudo-inverse of W0 . The pseudo-inverse of matrix is applicable to classical associative memory which involves finding 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 W0 . 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 W0 ( given prior memory M0 and noisy data Zξ ) and reading weight W ( given posterior memory M and query Z ) 2 . Ideally , W0 and W should satisfy W0M0 = Zξ and WM = Zξ ≈ Z , according to the linear assumptions over Zξ . The minimumnorm least-square solutions of those linear systems are W0 = ZξM+0 and W = ZM + , where M+0 and M + are pseudo-inverses of M0 and M , respectively . While least-square condition is obvious since we would like to find solutions W0 and W that best fit 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 D0 of the same size with the transpose D > . The sequence { Di } i≥0 defined recursively as Di+1 = 2Di −DiDDi will converge to the pseudo-inverse of D. With appropriate D0 , the sequence will converge quadratically . Ben-Cohen algorithm allows us to compute the matrix pseudo-inverse efficiently with nearly O ( 1 ) time complexity in practice ( see Section 3.5 for experimental results ) . For simplicity , we set the initial term D0 = αD > , 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 . 2In general case ( e.g . when the query is a noisy version of Z ) , the same dynamic addressing holds .

Summary: The paper proposes a new memory model following the research line of Kanerva Machine, DKM, and Kanerva++. The proposed model reformulates the Bayesian updates of memory and address as finding least-square solution to linear systems. This requires matrix inversion operations and the authors proposed to approximate it iteratively by using Ben-Cohen algorithm for pseudo-inverse matrix. It results in a memory read/write system that is rapid/accurate and can store large batch of data. The evaluation is performed thoroughly on various datasets including binarized MNIST, binarized Omniglot, Fashion MNIST, CIFAR10/100, CelebA. It shows the superiority of the proposed method in negative ELBO of test likelihood, denoising success rate and hamming error, generation, memory capacity, and run-time per iteration.