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 authors propose a novel deep generative memory model called Generative Pseudo-inverse Memory (GPM) that extends Dynamic Kanerva Machines (DKM, [1]). These models are both deep generative models that maintain a hidden state similar to recurrent neural networks called memory and have ways of writing into and reading from this memory. A key limitation of DKMs is that the time complexity of updating the memory scales cubically in the memory dimensionality and linearly in the episode length. The authors observe that by slightly altering the optimization objective that GPM has to solve, under mild assumptions the memory update can be performed much more efficiently, allowing them to write entire datasets into the memory of the GPM with little computational cost.  The authors perform some experiments on some toy datasets to validate their method and obtain good results against related methods.  [1] Y. Wu, G. Wayne, K. Gregor, and T. Lillicrap. Learning attractor dynamics for generative memory. NeurIPS 2018