Attentive Sequential Neural Processes

1 INTRODUCTION . Neural Processes ( NP ) ( Garnelo et al. , 2018b ) combine the strengths of neural networks and Gaussian processes ( GP ) such that , like Gaussian processes , it is flexible in learning a new stochastic process at test time and also provides fast O ( 1 ) prediction speed like neural networks . Learning from small datasets of multiple tasks ( i.e. , multiple stochastic processes ) , NP can be seen as a probabilistic latent variable framework for meta-learning . Sequential Neural Processes ( SNP ) ( Singh et al. , 2019 ) extend the power of NP to a sequence of stochastic processes thus introducing a new class of sequential latent generative models . It targets a large class of problems where the sequence of stochastic processes is governed by underlying transition dynamics and where learning to transfer information between stochastic processes is useful . Despite these strengths , the remained question is whether NP and SNP can scale to more complex stochastic processes when the individual observations about the true process are highly partial , which prevails in many situations . For example , an agent in a large complex landscape can only obtain very limited information from its immediate surroundings and it is desirable that it still maintains an accurate global scene representation that could also be changing with time . Studying NP from this perspective , Kim et al . ( 2019 ) demonstrates that the under-fitting problem significantly deteriorates the performance of NP . To address this , ANP ( Kim et al. , 2019 ) and the method in Rosenbaum et al . ( 2018 ) augment attention to NP and GQN ( Eslami et al. , 2018 ) , respectively . Given the severe performance degradation of NP due to the under-fitting , an important key question is whether SNP would also suffer from the under-fitting problem or not , and its follow-up questions such as if so , how severely and in which settings would it happen ? and how can we resolve the problem—would the same attention mechanism that worked for NP work for SNP as well ? In this paper , our goal is to answer these questions . We achieve this by not only showing that our proposed attention for SNP significantly improves the standard SNP under the under-fitting settings ( see Fig . 1 ) , but also by making a stronger claim i.e. , in a sequential setting like ours , it is not optimal to perform attention on a memory buffer that simply stores all the observed contexts . Instead , through our experiments , we claim that attention should be performed on a memory that is also updated ( learned ) at every time-step . Consequently , a ) this memory would learn to store an optimized representation of the past geared towards prediction and b ) it would require fewer storage locations as it does not need to naively store every observed context points . To implement such a memory , we introduce the concept of imaginary context that augments the set of real context points observed at any time-step . The imaginary context contains a fixed number of context points and is modeled as a random variable sampled at every time-step . Given a query , we then apply attention on the union of the real and the imaginary contexts to generate predictions . We call the proposed model as Attentive Sequential Neural Processes ( ASNP ) . In experiments on 1D regression tasks , we validate that SNP which simply attends on the memory buffer of all the past context points under-performs the proposed ASNP endowed with imaginary contexts . Then in a comprehensive set of experiments on 1D and 2D regression and rendering tasks , we demonstrate ASNP ’ s performance gains over NP , ANP , SNP , GQN , and TGQN in different context regimes . Our main contributions are as follows . a ) We show that SNP suffers from significant under-fitting by testing it in settings when the true stochastic process is complex ( e.g. , complex CelebA faces moving on a white canvas ) or when each observation is highly partial ( e.g. , each context point is just one pixel ) . b ) We introduce a new model , Attentive Sequential Neural Processes ( ASNP ) , that resolves the problem of under-fitting by augmenting SNP with an attention mechanism . To implement this , we introduce the novel concept of imaginary context . We also empirically demonstrate that SNP which simply performs attention on a lossless memory of all the past context points still underperforms the proposed model and therefore the use of imaginary context is a better design choice . c ) Using the attention paradigm introduced in ASNP , we extend TGQN to Attentive TGQN ( ATGQN ) for scene rendering . d ) We show that ASNP outperforms the baselines in various 1D and 2D tasks performing future prediction and meta-transfer learning i.e. , making simultaneous use of the past knowledge and the current contexts . 2 BACKGROUND . Neural Processes ( NP ) ( Garnelo et al. , 2018b ) learns to learn a stochastic process P that maps an input x ∈ Rdx to a random variable y ∈ Rdy . It does so using a set of context observations C = ( XC , YC ) = { xi , yi } i∈I ( C ) . Here , I ( C ) is the set of indices for the elements in set C. Specifically , to learn stochastic process from the context , NP uses an inference distribution P ( z|C ) and then uses the representation z ∼ P ( z|C ) of the inferred stochastic process to generate the prediction y given a query x . This is modeled by p ( y|x , z ) . The full generative process of NP conditioned on the context can be written as : P ( Y |X , C ) = ∫ P ( Y |X , z ) P ( z|C ) dz ( 1 ) where P ( Y |X , z ) = ∏ i∈I ( D ) P ( yi|xi , z ) and D contains the target observations D = ( X , Y ) = { xi , yi } i∈I ( D ) . The training data for NP is obtained by drawing a stochastic process and then sampling ( C , D ) from the sampled stochastic process . Treating each stochastic process as a task , NP can be seen as a probabilistic meta-learning framework . Generative Query Networks ( GQN ) ( Eslami et al. , 2018 ) uses modeling similar to NP for the problem of learning representation and rendering of 3D scenes from a set of partial 2D observations . In GQN , a query x becomes a camera viewpoint in a 3D environment and an output y corresponds to an image taken from the viewpoint x . The latent variable in the original GQN are query-dependent and may thus produce uncorrelated generations across queries . Consistent GQN ( CGQN ) ( Kumar et al. , 2018 ) resolves this problem by having a single stochastic scene representation . For simplicity , we use the term GQN to denote CGQN in the remainder of the paper . Attentive Neural Processes ( ANP ) ( Kim et al. , 2019 ) resolves the problem of under-fitting in NP . To this end , ANP uses two context-encoding paths . The first path is the same as the scene encoder in NP that summarizes all the context points agnostic to the target query x and produces a single vector of the global representation by average pooling . The other path computes representations for each target query using query-dependent attention . The set of representations returned r from the set of target queries X may be written as r = fattend ( X , XC , vC ) , where vC are the encoded representations of C via MLP or self-attention . The attention function fattend is a key-value based attention which computes a response for each query qj ∈ X by first computing its similarity with the key set XC as wji = fw ( ki , qj ) and then performing a weighted-sum on the value set i.e. , sj = ∑ i wjivi , . Here , the keys and values are the context queries xC and the representations vC , respectively . The generative process is as follows : P ( Y |X , C ) = ∫ P ( Y |X , z , fattend ( X , XC , vC ) ) P ( z|C ) dz . Kim et al . ( 2019 ) propose ANP as a general model class which may be realized using different attention mechanism choices such as Dot-Product , Laplace or Multi-Head ( Vaswani et al. , 2017 ) . ANP needs more computation than NP but Kim et al . ( 2019 ) mitigates this by showing faster training convergence than NP . Sequential Neural Processes ( SNP ) ( Singh et al. , 2019 ) extend the capabilities of NP to model stochastic processes which may be decomposed as a sequence of stochastic processesP1 , P2 , · · · PT such that Pt at time t depends on the past processes P < t through some underlying temporal dynamics . At any t , Pt is modeled as zt which depends on the current context Ct and the past processes z < t using the transition model p ( zt|z < t , Ct ) . This enables us to use the general past trend as well as the current context to meta-transfer learn the process Pt . Given zt and a query xt , the prediction yt is generated using the observation model p ( yt|xt , zt ) . In the SNP framework , at time-step t , we can not directly access the past contexts , but we can still directly access the current context Ct . Thus , to model Pt , we can use the stochastic encoding zt together with the direct encoding of Ct. Alternatively , this allows us to write the observation model as p ( yt|xt , zt , Ct ) . Combining the transition and the observation models , the complete generative process of SNP may be written as follows . P ( Y , Z|X , C ) = T∏ t=1 p ( yt|xt , zt , Ct ) p ( zt|z < t , Ct ) . ( 2 ) 3 ATTENTIVE SEQUENTIAL NEURAL PROCESSES . In this section , we describe the proposed Attentive Sequential Neural Processes ( ASNP ) that augments SNP with an attention mechanism . 3.1 GENERATIVE PROCESS We recall that ANP uses the observation model p ( y|x , z , C ) which , based on the query x , performs attention on the context points in C. In SNP , which uses the observation model p ( yt|xt , zt , Ct ) , an analogous approach would use xt-dependent attention onCt . However with this approach , attention is not effective because the current context is likely to be limited or even empty . Instead , we want to perform attention not only on the current but also on the past contexts . But to do that , we found that simply attending on a memory buffer of context history does not produce enough gains and leaves room for further improvement ( see Appendix D.3 ) . Instead , our hypothesis is that attention should be performed on a sequentially-updated memory so that the representations are optimized for prediction and are more size-efficient . To realize this , we propose augmenting the set of real context points Crt with a set of imaginary context points C i t resulting in the full context set Ct ≡ Crt ∪ Cit so that the observation model can now perform attention on this augmented Ct . Here , Cit = ( X i t , U i t ) consists of imaginary queries ( or keys ) X i t = { xit , k } Kk=1 and its corresponding values U it = { uit , k } Kk=1 where K is a hyperparameter similar to the number of slots in a memory . To generate Cit , we introduce a novel and attentive transition model P ( C i t |C < t , Crt ) treating Cit as a random variable . In this way , Cit is a noisy reconstruction of the past sampled using an attention mechanism on the past contexts ( real and imaginary ) and on the real contexts of the current time-step . Augmenting SNP with these , the generative model for ASNP may be written as follows . P ( Y , Z , Ci|X , Cr ) = T∏ t=1 P ( yt|xt , zt , Ct ) P ( zt|z < t , Ct ) P ( Cit |C < t , Crt ) . ( 3 ) Implementing the transition model P ( zt|z < t , Ct ) . With this model , we obtain the global representation zt which is a query-agnostic encoding of the current stochastic process . This module can be considered similar to its counterpart in SNP without attention . Note that it is a design choice whether to choose Ct or Crt for the conditioning context of this module . If we use C r t , zt becomes global context-encoding of real contexts . Otherwise , it becomes an encoding of the combined contexts . Implementing the context imagination P ( Cit |C < t , Crt ) . To implement this , we keep track of the K imaginary context points at all time-steps and their generative model is as follows : P ( Cit |C < t , Crt ) = K∏ k=1 P ( xit , k , u i t , k|C < t , Crt ) = K∏ k=1 P ( uit , k|xit , k , C < t , Crt ) P ( xit , k|C < t , Crt ) . ( 4 ) Algorithm 1 Attention with Imagination hXt = RNNX ( ( X i t−1 , v r t ) , h X t−1 ) Xit ∼ N ( fXµ ( hXt ) , fXσ ( hXt ) ) hut , k = RNNu ( ( x i t , k , u i t−1 , k ) , h u t−1 , k ) ait , k = fattend ( x i t , k , S i t ) uit , k ∼ N ( fuµ ( ait , k ) , fuσ ( ait , k ) ) rxt = fattend ( xt , Ct ) For each imagination slot k , we first sample the key xit , k . For this , we use an RNN that takes the past imaginary inputs Xit−1 and the encoding of the real contexts vrt = forder-inv ( C r t ) to produce a hidden state h X t from which we then sample the imaginary key set Xit . Given these keys , the corresponding imaginary values uit , k are generated next . For each imagination slot k , we maintain an imagination-tracker RNN that takes ( xit , k , u i t−1 , k ) as input and produces a hidden state hut , k . These hidden states act as intermediate imaginary values and we use them to construct a key-value set Sit = { ( xit , k , hut , k ) } ∪ { ( xrt , fy→u ( xrt , yrt ) ) } . On this set , we apply self-attention using only the imaginary xit , k as attention queries and obtain the corresponding attention encodings ait , k . Such encodings therefore also model the interactions among the context points . From these , we sample the final imagination values uit , k ∼ N ( fuµ ( ait , k ) , fuσ ( ait , k ) ) . Implementing the observation model P ( yt|xt , zt , Ct ) . Taking xt as the attention query , we attend onCt , the union of the realCrt and imaginedC i t context , as the key-value set and obtain the attention encoding rxt = fattend ( xt , Ct ) . This is done after each real context x ∈ Xrt and y ∈ Y rt is encoded using fy→u . The attention encoding rxt is then concatenated with the global encoding zt to perform the prediction yt . Algorithm 1 shows the described process of context imagination and encoding . Discussion . An interesting interpretation of the proposed modeling can be drawn in comparison to the memory retrieval process in the human brain ( Eichenbaum , 2017 ) . As is well-known in neuroscience , “ human memory is not a literal reproduction of the past , but instead relies on constructive processes that are sometimes prone to error and distortion ” ( Schacter , 2012 ) . Our model resembles this process in the sense that it ( i ) compresses the past observation and importantly what it has imagined in the past through an RNN ( similarly to memory consolidation ) , and then after observing current context , ( ii ) recalls from the compressed memory not a simple copy of the past but a constructive recreation of representations that is optimized to help prediction as illustrated in Fig . 2 .

Summary: Authors present a method to address the problem of underfitting found in sequential neural processes. They cover the literature appropriately in regards to neural processes and developments pertaining to tackling the underfitting problem by applying an attention mechanism. Although, this has successfully been achieved with Neural Processes, the case is different with sequential neural processes, as they cannot store the past context.