Attentive Neural Processes

1 INTRODUCTION . Regression tasks are usually cast as modelling the distribution of a vector-valued output y given a vector-valued input x via a deterministic function , such as a neural network , taking x as an input . In this setting , the model is trained on a dataset of input-output pairs , and predictions of the outputs are independent of each other given the inputs . An alternative approach to regression involves using the training data to compute a distribution over functions that map inputs to outputs , and using draws from that distribution to make predictions on test inputs . This approach allows for reasoning about multiple functions consistent with the data , and can capture the co-variability in outputs given inputs . In the Bayesian machine learning literature , non-parametric models such as Gaussian Processes ( GPs ) are popular choices of this approach . Neural Processes ( NPs ) ( Garnelo et al. , 2018a ; b ) offer an efficient method to modelling a distribution over regression functions , with prediction complexity linear in the context set size . Once trained , they can predict the distribution of an arbitrary target output conditioned on a set of context inputoutput pairs of an arbitrary size . This flexibility of NPs enables them to model data that can be interpreted as being generated from a stochastic process . It is important to note however that NPs and GPs have different training regimes . NPs are trained on samples from multiple realisations of a stochastic process ( i.e . trained on many different functions ) , whereas GPs are usually trained on observations from one realisation of the stochastic process ( a single function ) . Hence a direct comparison between the two is usually not plausible . Despite their many appealing properties , one substantial weakness of NPs is that they tend to underfit the context set . This manifests in the 1D curve fitting example on the left half of Figure 1 as inaccurate predictive means and overestimated variances at the input locations of the context set . The right half of the figure shows this phenomenon when predicting the bottom half of a face image from its top half : although the prediction is globally coherent , the model ’ s reconstruction of the top-half is far from perfect . In an NP , the encoder aggregates the context set to a fixed-length latent summary via a permutation invariant function , and the decoder maps the latent and target input to the target output . We hypothesise that the underfitting behaviour is because the mean-aggregation step in the encoder acts as a bottleneck : since taking the mean across context representations gives the same weight to each context point , it is difficult for the decoder to learn which context points ∗Corresponding author : hyunjikk @ google.com provide relevant information for a given target prediction . In theory , increasing the dimensionality of the representation could address this issue , but we show in Section 4 that in practice , this is not sufficient . To address this issue , we draw inspiration from GPs , which also define a family of conditional distributions for regression . In GPs , the kernel can be interpreted as a measure of similarity among two points in the input domain , and shows which context points ( xi , yi ) are relevant for a given query x∗ . Hence when x∗ is close to some xi , its y-value prediction y∗ is necessarily close to yi ( assuming small likelihood noise ) , and there is no risk of underfitting . We implement a similar mechanism in NPs using differentiable attention that learns to attend to the contexts relevant to the given target , while preserving the permutation invariance in the contexts . We evaluate the resulting Attentive Neural Processes ( ANPs ) on 1D function regression and on 2D image regression . Our results show that ANPs greatly improve upon NPs in terms of reconstruction of contexts as well as speed of training , both against iterations and wall clock time . We also demonstrate that ANPs show enhanced expressiveness relative to the NP and is able to model a wider range of functions . 2 BACKGROUND . 2.1 NEURAL PROCESSES . The NP is a model for regression functions that map an input xi ∈ Rdx to an output yi ∈ Rdy . In particular , the NP defines a ( infinite ) family of conditional distributions , where one may condition on an arbitrary number of observed contexts ( xC , yC ) : = ( xi , yi ) i∈C to model an arbitrary number of targets ( xT , yT ) : = ( xi , yi ) i∈T in a way that is invariant to ordering of the contexts and ordering of the targets . The model is defined for arbitrary C and T but in practice we use C ⊂ T . The deterministic NP models these conditional distributions as : p ( yT |xT , xC , yC ) : = p ( yT |xT , rC ) ( 1 ) with rC : = r ( xC , yC ) ∈ Rd where r is a deterministic function that aggregates ( xC , yC ) into a finite dimensional representation with permutation invariance in C. In practice , each context ( x , y ) pair is passed through an MLP to form a representation of each pair , and these are aggregated by taking the mean to form rC . The likelihood p ( yT |xT , rC ) is modelled by a Gaussian factorised across the targets ( xi , yi ) i∈T with mean and variance given by passing xi and rC through an MLP . The unconditional distribution p ( yT |xT ) ( when C = ∅ ) is defined by letting r∅ be a fixed vector . The latent variable version of the NP model includes a global latent z to account for uncertainty in the predictions of yT for a given observed ( xC , yC ) . It is incorporated into the model via a latent path that complements the deterministic path described above . Here z is modelled by a factorised Gaussian parametrised by sC : = s ( xC , yC ) , with s being a function of the same properties as r p ( yT |xT , xC , yC ) : = ∫ p ( yT |xT , rC , z ) q ( z|sC ) dz ( 2 ) with q ( z|s∅ ) : = p ( z ) , the prior on z . The likelihood is referred to as the decoder , and q , r , s form the encoder . See Figure 2 for diagrams of these models . The motivation for having a global latent is to model different realisations of the data generating stochastic process — each sample of z would correspond to one realisation of the stochastic process . One can define the model using either just the deterministic path , just the latent path , or both . In this work we investigate the case of using both paths , which gives the most expressive model and also gives a sensible setup for incorporating attention , as we will show later in Section 3 . The parameters of the encoder and decoder are learned by maximising the following ELBO log p ( yT |xT , xC , yC ) ≥ Eq ( z|sT ) [ log p ( yT |xT , rC , z ) ] −DKL ( q ( z|sT ) ‖q ( z|sC ) ) ( 3 ) for a random subset of contexts C and targets T via the reparametrisation trick ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) . In other words , the NP learns to reconstruct targets , regularised by a KL term that encourages the summary of the contexts to be not too far from the summary of the targets . This is sensible since we are assuming that the contexts and targets come from the same realisation of the data-generating stochastic process , and especially so if targets contain contexts . At each training iteration , the number of contexts and targets are also chosen randomly ( as well as being randomly sampled from the training data ) , so that the NP can learn a wide family of conditional distributions . NPs have many desirable properties , namely ( i ) Scalability : computation scales linearly atO ( n+m ) for n contexts and m targets at train and prediction time . ( ii ) Flexibility : defines a very wide family of distributions , where one can condition on an arbitrary number of contexts to predict an arbitrary number of targets . ( iii ) Permutation invariance : the predictions of the targets are order invariant in the contexts . However these advantages come at the cost of not satisfying consistency in the contexts . For example , if y1 : m is generated given some context set , then its distribution need not match the distribution you would obtain if y1 : n is generated first , appended to the context set then yn+1 : m is generated . However maximum-likelihood learning can be interpreted as minimising the KL between the ( consistent ) conditional distributions of the data-generating stochastic process and the corresponding conditional distributions of the NP . Hence we could view the NP as approximating the conditionals of the consistent data-generating stochastic process . 2.2 ATTENTION . Given a set of key-value pairs ( ki , vi ) i∈I and a query q , an attention mechanism computes weights of each key with respect to the query , and aggregates the values with these weights to form the value corresponding to the query . In other words , the query attends to the key-value pairs . The queried values are invariant to the ordering of the key-value pairs ; this permutation invariance property of attention is key in its application to NPs . The idea of using a differentiable addressing mechanism that can be learned from the data has been applied successfully in various areas of Deep Learning , namely handwriting generation and recognition ( Graves , 2012 ) and neural machine translation ( Bahdanau et al. , 2015 ) . More recently , there has been work employing self-attention ( where keys and queries are identical ) to give expressive sequence-to-sequence mappings in natural language processing ( Vaswani et al. , 2017 ) and image modelling ( Parmar et al. , 2018 ) . We give some examples of attention mechanisms which are used in the paper . Suppose we have n key-value pairs arranged as matrices K ∈ Rn×dk , V ∈ Rn×dv , and m queries Q ∈ Rm×dk . Simple forms of attention based on locality ( weighting keys according to distance from query ) are given by various stationary kernels . For example , the ( normalised ) Laplace kernel gives the queried values as Laplace ( Q , K , V ) : =WV ∈ Rm×dv , Wi· : = softmax ( ( −||Qi· −Kj·||1 ) nj=1 ) ∈ Rn Similarly ( scaled ) dot-product attention uses the dot-product between the query and keys as a measure of similarity , and weights the keys according to the values DotProduct ( Q , K , V ) : = softmax ( QK > / √ dk ) V ∈ Rm×dv The use of dot-product attention allows the query values to be computed with two matrix multiplications and a softmax , allowing for use of highly optimised matrix multiplication code . multihead attention ( Vaswani et al. , 2017 ) is a parametrised extension where for each head , the keys , values and queries are linearly transformed , then dot-product attention is applied to give headspecific values . These values are concatenated and linearly transformed to produce the final values : MultiHead ( Q , K , V ) : = concat ( head1 , . . . , headH ) W ∈ Rm×dv where headh : = DotProduct ( QW Q h , KW K h , V W V h ) ∈ Rm×dv This multihead architecture allows the query to attend to different keys for each head and tends to give smoother query-values than dot-product attention ( c.f . Section 4 ) . 3 ATTENTIVE NEURAL PROCESSES Figure 2 describes how attention is incorporated into NP to give the Attentive NP ( ANP ) . In summary , self-attention is applied to the context points to compute representations of each ( x , y ) pair , and the target input attends to these context representations ( cross-attention ) to predict the target output . In detail , the representation of each context pair ( xi , yi ) i∈C before the mean-aggregation step is computed by a self-attention mechanism , in both the deterministic and latent path . The intuition for the self-attention is to model interactions between the context points . For example , if many context points overlap , then the query need not attend to all of these points , but only give high weight to one or a few . The self-attention will help obtain richer representations of the context points that encode these types of relations between the context points . We model higher order interactions by simply stacking the self-attention , as is done in Vaswani et al . ( 2017 ) . In the deterministic path , the mean-aggregation of the context representations that produces rC is replaced by a cross-attention mechanism , where each target query x∗ attends to the context xC : = ( xi ) i∈C to produce a query-specific representation r∗ : = r∗ ( xC , yC , x∗ ) . This is precisely where the model allows each query to attend more closely to the context points that it deems relevant for the prediction . The reason we do not have an analogous mechanism in the latent path is that we would like to preserve the global latent , that induces dependencies between the target predictions . The interpretation of the latent path is that z gives rise to correlations in the marginal distribution of the target predictions yT , modelling the global structure of the stochastic process realisation , whereas the deterministic path models the fine-grained local structure . The decoder remains the same , except we replace the shared context representation rC with the query-specific representation r∗ . Note that permutation invariance in the contexts is preserved with the attention mechanism . If we use uniform attention ( all contexts given the same weight ) throughout , we recover the NP . ANP is trained using the same loss ( 3 ) as the original NP , also using Gaussian likelihood p ( yi|xi , r∗ ( xC , yC , xi ) , z ) and diagonal Gaussian q ( z|sC ) . The added expressivity and resulting accuracy of the NP with attention comes at a cost . The computational complexity is raised from O ( n+m ) to O ( n ( n+m ) ) , since we apply self-attention across the contexts and for every target point we compute weights for all contexts . However most of the computation for the ( self- ) attention is done via matrix multiplication ( c.f . Section 2.2 ) , and so can be done in parallel across the contexts and across the targets . In practice , the training time for ANPs remains comparable to NPs , and in fact we show that ANPs learn significantly faster than NPs not only in terms of training iterations but also in wall-clock time , despite being slower at prediction time ( c.f . Section 4 ) .

Summary: The authors extend neural processes by incorporating two types of attention processes: self-attention for enriching the features of the context points and cross-attention for producing a query-specific representation. By replacing MLPs and mean pooling with these attention processes, the authors resolve the underfitting problem of NPs. The experimental results show that ANPs converge better and faster than NPs.