Meta-Learning Deep Energy-Based Memory Models

1 INTRODUCTION . Associative memory has long been of interest to neuroscience and machine learning communities ( Willshaw et al. , 1969 ; Hopfield , 1982 ; Kanerva , 1988 ) . This interest has generated many proposals for associative memory models , both biological and synthetic . These models address the problem of storing a set of patterns in such a way that a stored pattern can be retrieved based on a partially known or distorted version . This kind of retrieval from memory is known as auto-association . Due to the generality of associative retrieval , successful implementations of associative memory models have the potential to impact many applications . Attractor networks provide one well-grounded foundation for associative memory models ( Amit & Amit , 1992 ) . Patterns are stored in such a way that they become attractors of the update dynamics defined by the network . Then , if a query pattern that preserves sufficient information for association lies in the basin of attraction for the original stored pattern , a trajectory initialized by the query will converge to the stored pattern . A variety of implementations of the general attractor principle have been proposed . The classical Hopfield network ( Hopfield , 1982 ) , for example , defines a simple quadratic energy function whose parameters serve as a memory . The update dynamics in Hopfield networks iteratively minimize the energy by changing elements of the pattern until it converges to a minimum , typically corresponding to one of the stored patterns . The goal of the writing process is to find parameter values such that the stored patterns become attractors for the optimization process and such that , ideally , no spurious attractors are created . Many different learning rules have been proposed for Hopfield energy models , and the simplicity of the model affords compelling closed-form analysis ( Storkey & Valabregue , 1999 ) . At the same time , Hopfield memory models have fundamental limitations : ( 1 ) It is not possible to add capacity for more stored patterns by increasing the number of parameters since the number of parameters in a Hopfield network is quadratic in the dimensionality of the patterns . ( 2 ) The model lacks a means of modelling the higher-order dependencies that exist in real-world data . In domains such as natural images , the potentially large dimensionality of an input makes it both ineffective and often unnecessary to model global dependencies among raw input measurements . In fact , many auto-correlations that exist in real-world perceptual data can be efficiently compressed without significant sacrifice of fidelity using either algorithmic ( Wallace , 1992 ; Candes & Tao , 2004 ) or machine learning tools ( Gregor et al. , 2016 ; Toderici et al. , 2017 ) . The success of existing deep learning techniques suggests a more efficient recipe for processing high-dimensional inputs by modelling a hierarchy of signals with restricted or local dependencies ( LeCun et al. , 1995 ) . In this paper we use a similar idea for building an associative memory : use a deep network ’ s weights to store and retrieve data . Fast writing rules A variety of energy-based memory models have been proposed since the original Hopfield network to mitigate its limitations ( Hinton et al. , 2006b ; Du & Mordatch , 2019 ) . Restricted Boltzmann Machines ( RBMs ) ( Hinton , 2012 ) add capacity to the model by introducing latent variables , and deep variants of RBMs ( Hinton et al. , 2006b ; Salakhutdinov & Larochelle , 2010 ) afford more expressive energy functions . Unfortunately , training Boltzmann machines remains challenging , and while recent probabilistic models such as variational auto-encoders ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) are easier to train , they nevertheless pay the price for expressivity in the form of slow writing . While Hopfield networks memorize patterns quickly using a simple Hebbian rule , deep probabilistic models are slow in that they rely on gradient training that requires many updates ( typically thousands or more ) to settle new inputs into the weights of a network . Hence , writing memories via parametric gradient based optimization is not straightforwardly applicable to memory problems where fast adaptation is a crucial requirement . In contrast , and by explicit design , our proposed method enjoys fast writing , requiring few parameter updates ( we employ just 5 steps ) to write new inputs into the weights of the net once meta-learning is complete . It also enjoys fast reading , requiring few gradient descent steps ( again just 5 in our experiments ) to retrieve a pattern . Furthermore , our writing rules are also fast in the sense that they use O ( N ) operations to store N patterns in the memory – this scaling is the best one can hope for without additional assumptions . We propose a novel approach that leverages meta-learning to enable fast storage of patterns into the weights of arbitrarily structured neural networks , as well as fast associative retrieval . Our networks output a single scalar value which we treat as an energy function whose parameters implement a distributed storage scheme . We use gradient-based reading dynamics and meta-learn a writing rule in the form of truncated gradient descent over the parameters defining the energy function . We show that the proposed approach enables compression via efficient utilization of network weights , as well as fast-converging attractor dynamics . 2 RETRIEVAL IN ENERGY-BASED MODELS . We focus on attractor networks as a basis for associative memory . Attractor networks define update dynamics for iterative evolution of the input pattern : x ( k+1 ) = f ( x ( k ) ) . For simplicity , we will assume that this process is discrete in time and deterministic , however there are examples of both continuous-time ( Yoon et al. , 2013 ) and stochastic dynamics ( Aarts & Korst , 1988 ) . A fixed-point attractor of deterministic dynamics can be defined as a point x for which it converges , i.e . x = f ( x ) . Learning the associative memory in the attractor network is then equivalent to learning the dynamics f such that its fixed-point attractors are the stored patterns and the corresponding basins of attraction are sufficiently wide for retrieval . An energy-based attractor network is defined by the energy function E ( x ) mapping an input object x ∈ X to a real scalar value . A particular model may then impose additional requirements on the energy function . For example if the model has a probabilistic interpretation , the energy function is usually a negative unnormalized logarithm of the object probability log p ( x ) = −E ( x ) + const , implying that the energy has to be well-behaved for the normalizing constant to exist . In our case no such constraints are put on the energy . The attractor dynamics in energy-based models is often implemented either by iterative energy optimization ( Hopfield , 1982 ) or sampling ( Aarts & Korst , 1988 ) . In the optimization case considered further in the paper , attractors are conveniently defined as local minimizers of the energy function . While a particular energy function may suggest a number of different optimization schemes for retrieval , convergence to a local minimum of an arbitrary function is NP-hard . Thus , we consider a class of energy functions that are differentiable on X ⊆ Rd , bounded from below and define the update dynamics over k = 1 , . . . , K steps via gradient descent : read ( x̃ ; θ ) = x ( K ) , x ( k+1 ) = x ( k ) − γ ( k ) ∇xE ( x ( k ) ) , x ( 0 ) = x̃ . ( 1 ) With appropriately set step sizes { γ ( k ) } Kk=0 this procedure asymptotically converges to a local minimum of energy E ( x ) ( Nesterov , 2013 ) . Since asymptotic convergence may be not enough for practical applications , we truncate the optimization procedure ( 1 ) at K steps and treat x ( K ) as a result of the retrieval . While vanilla gradient descent ( 1 ) is sufficient to implement retrieval , in our experiments we employ a number of extensions , such as the use of Nesterov momentum and projected gradients , which are thoroughly described in Appendix B. Relying on the generic optimization procedure allows us to translate the problem of designing update dynamics with desirable properties to constructing an appropriate energy function , which in general is equally difficult . In the next section we discuss how to tackle this difficulty . 3 META-LEARNING GRADIENT-BASED WRITING RULES . As discussed in previous sections , our ambition is to be able to use any scalar-output neural network as an energy function for associate retrieval . We assume a parametric model E ( x ; θ ) differentiable in both x and θ , and bounded from below as a function of x . These are mild assumptions that are often met in the existing neural architectures with an appropriate choice of activation functions , e.g . tanh . The writing rule then compresses input patterns X = { x1 , x2 , . . . , xN } into parameters θ such that each of the stored patterns becomes a local minimum of E ( x ; θ ) or , equivalently , creates a basin of attraction for gradient descent in the pattern space . This property can be practically quantified by the reconstruction error , e.g . mean squared error , between the stored pattern x and the pattern read ( x̃ ; θ ) retrieved from a distorted version of x : L ( X , θ ) = 1 N N∑ i=1 Ep ( x̃i|xi ) [ ||xi − read ( x̃i ; θ ) ||22 ] . ( 2 ) Here we assume a known , potentially stochastic distortion model p ( x̃|x ) such as randomly erasing certain number of dimensions , or salt-and-pepper noise . While one can consider loss ( 2 ) as a function of network parameters θ and call minimization of this loss with a conventional optimization method a writing rule — it will require many optimization steps to obtain a satisfactory solution and thus does not fall into our definition of fast writing rules ( Santoro et al. , 2016 ) . Hence , we explore a different approach to designing a fast writing rule inspired by recently proposed gradient-based meta-learning techniques ( Finn et al. , 2017 ) which we call meta-learning energy-based memory models ( EBMM ) . Namely we perform many write and read optimization procedures with a small number of iterations for several sets of write and read observations , and backpropagate into the initial parameters of θ — to learn a good starting location for fast optimization . As usual , we assume that we have access to the underlying data distribution pd ( X ) over data batches of interest X from which we can sample sufficiently many training datasets , even if the actual batch our memory model will be used to store ( at test time ) is not available at the training time ( Santoro et al. , 2016 ) . The straightforward application of gradient-based meta-learning to the loss ( 2 ) is problematic , because we generally can not evaluate or differentiate through the expectation over stochasticity of the distortion model in a way that is reliable enough for adaptation , because as the dimensionality of the pattern space grows the number of possible ( and representative ) distortions grows exponentially . Instead , we define a different writing loss W , minimizing which serves as a proxy for ensuring that input patterns are local minima for the energy E ( x ; θ ) , but does not require costly retrieval of exponential number of distorted queries . W ( x , θ ) = E ( x ; θ ) + α||∇xE ( x ; θ ) ||22 + β||θ − θ̄||22 . ( 3 ) As one can see , the writing loss ( 3 ) consists of three terms . The first term is simply the energy value which we would like to be small for stored patterns relative to non-stored patterns . The condition for x to be a local minimum of E ( x ; θ ) is two-fold : first , the gradient at x is zero , which is captured by the second term of the writing loss , and , second , the hessian is positive-definite . The latter condition is difficult to express in a form that admits efficient optimization and we found that meta-learning using just first two terms in the writing loss is sufficient . Finally , the third term limits deviation from initial or prior parameters θ̄ which we found helpful from optimization perspective ( see Appendix D for more details ) . We use truncated gradient descent on the writing loss ( 3 ) to implement the writing rule : write ( X ) = θ ( T ) , θ ( t+1 ) = θ ( t ) − η ( t ) 1 N N∑ i=1 ∇θW ( xi , θ ( t ) ) , θ ( 0 ) = θ̄ ( 4 ) To ensure that gradient updates ( 4 ) are useful for minimization of the reconstruction error ( 2 ) we train the combination of retrieval and writing rules end-to-end , meta-learning initial parameters θ̄ , learning rate schedules r = ( { γ ( k ) } Kk=1 , { η ( t ) } Tt=1 ) and meta-parameters τ = ( α , β ) to perform well on random sets of patterns from the batch distribution pd ( X ) : minimize EX∼pd ( X ) [ L ( X , write ( X ) ) ] for θ̄ , r , τ . ( 5 ) In our experiments , p ( X ) is simply a distribution over batches of certain size N sampled uniformly from the training ( or testing - during evaluation ) set . Parameters θ = write ( X ) are produced by storing X in the memory using the writing procedure ( 4 ) . Once stored , distorted versions of X can be retrieved and we can evaluate and optimize the reconstruction error ( 2 ) . Crucially , the proposed EBMM implements both read ( x , θ ) and write ( X ) operations via truncated gradient descent which can be itself differentiated through in order to set up a tractable meta-learning problem . While truncated gradient descent is not guaranteed to converge , reading and writing rules are trained jointly to minimize the reconstruction error ( 2 ) and thus ensure that they converge sufficiently fast . This property turns this potential drawback of the method to its advantage over provably convergent , but slow models . It also relaxes the necessity of stored patterns to create too well-behaved basins of attraction because if , for example , a stored pattern creates a nuisance attractor in the dangerous proximity of the main one , the gradient descent ( 1 ) might successfully pass it with appropriately learned step sizes γ .

Summary: The authors propose to tackle the associative memory problem by recasting read/write operations to read/write by optimizing the parameters/input of an energy based model. Writing is reformulated as training a parametric energy model (EBMM) to have local minima of energy w.r.t. the parameters at memorized data points. Reading is performed by performing (projected) gradient descent on the corrupted/incomplete input to minimize the energy. To ensure the operations are fast (read and write with minimal gradient steps), the authors propose to take inspiration from modern meta-learning literature and learn initialization parameters of the energy model (and other hyperparameters for GD during read/write) from which writing is fast while ensuring reading is also fast, since the models are trained to maximize read/write performance within a constrained number of gradient steps.  Experimentally, the authors show that EBMM reading performs similar to baseline methods (but better across many memory sizes) on the standard Omniglot task. On CIFAR-10 and downsized ImageNet, they show much better L2 reconstruction error of corrupted images. They also show that the learnt energy