Neural Processes with Stochastic Attention: Paying more attention to the context dataset

1 INTRODUCTION . Neural processes ( NPs ) have been in the spotlight as they stochastically complete unseen target points considering a given context dataset without huge inference computation ( Garnelo et al. , 2018a ; b ; Kim et al. , 2019 ) . NPs leverage neural networks to derive an identifier suitable for a novel task using context representation , which contains information about given context data points . These methods enable us to handle considerable amounts of data points , such as in image-based applications , that the Gaussian process can not naively deal with . Many studies have revealed that the prediction performance relies on the way of context representation ( Kim et al. , 2019 ; Volpp et al. , 2021 ) . The variants of NPs have mainly investigated on context embedding approaches that generally design novel network architectures and aggregation functions that are permutation invariant . For example , Garnelo et al . ( 2018a ; b ) developed NPs using MLP layers to embed context set via mean aggregation . Bayesian aggregation satisfying permutation invariance enhanced the prediction accuracy by varying weights for individual context data points ( Volpp et al. , 2021 ) . From the prospective of network architectures to increase model capacity , Kim et al . ( 2019 ) suggested the way of local representation to consider relations between context dataset and target dataset using deterministic attention mechanism ( Vaswani et al. , 2017 ) . For robustness under noisy situations , bootstrapping method was proposed to orthogonally apply to variants of NPs ( Lee et al. , 2020 ) . Despite many appealing approaches , one significant drawback of previous NPs is that they still underfit when confronted with noisy situations like real-world problems . This manifests as inaccurate predictions at the locations of the context set as seen in Figure 1a . Additionally , in Figure 1b , the attentive neural process Kim et al . ( 2019 ) fails to capture contextual embeddings because the attention weights of all target points highlight on the lowest value or the maximum value in the context dataset . In the case of the Bootstrapping ANP Lee et al . ( 2020 ) and ANP with information dropout , the quality of heat-map is slightly improved , but it still falls short of ours . This indicates that the present NPs are unable to properly exploit context embeddings because the noisy situations impair the learning of the context embeddings during meta-training . To address this issue , we propose a newly designed neural process to fundamentally improve performance by paying more attention to the context dataset . The proposed method expedites a stochastic attention to adequately capture the dependency of the context and target datasets by adjusting stochasticity . It results in improving prediction performance due to maintained context information . We observe that our proposed algorithm works well by utilizing contextual information in an intended manner as seen in Figure 1 . This method outperforms current NPs and their regularization methods in all experiment settings . Thus , this paper clarifies the proposed method as regularized NPs in terms of information theory , explaining that the stochastic attention encourages the context embedding to be differentiated from the target dataset . This differentiated information induces NPs to become appropriate identifiers for target dataset by paying more attention to context representation . To summarize , we make the following contributions : • We propose the novel neural process that pay more attention to the context dataset . We claims for the first time , using the information theory framework , that critical conditions for contextual embeddings in NPs are independent of target features and close to contextual datasets . • Through comprehensive analyses , we illustrate how stochastic local embeddings are crucial for NPs to focus on capturing the dependencies of context and target datasets . Even when context dataset contains noise or is somewhat different from target dataset , as shown in Lee et al . ( 2020 ) , the proposed method is capable of adapting a novel task while preserving predictive performance . Particularly , this method significantly enhances performance without additional architectures and data augmentation compared to the attentive neural process ( Kim et al. , 2019 ) . • The experimental results show that the proposed model substantially outperforms conven- tional NPs in typical meta-regression problems . For instance , the proposed method achieves to obtain the state of the art score in the image completion task with the CelebA dataset . Especially , the proposed model maintains performance in the limited task distribution regimes such as the MovieLenz-10k dataset with a small number of users . This paper is organized as follows . We introduce background knowledge in section 2. section 3 presents a neural process with stochastic attention and related works illustrated in section 5 . Experimental results are shown in section 4 , concluding with section 6 . 2 BACKGROUND . 2.1 NEURAL PROCESSES . Suppose that we have an observation set X = { xi } ni=1 and a label set Y = { yi } n i=1 . NPs ( Garnelo et al. , 2018a ; b ; Kim et al. , 2019 ) are designed to obtain the probabilistic mapping from the observation set to the label set p ( Y |X , Xc , Yc ) given a small subset ( Xc , Yc ) = ( xj , yj ) mj=1 . Basically , it is built upon the neural network with an encoder-decoder architecture where the encoder fφ outputs a task representation by feed-forwarding ( Xc , Yc ) through permutation-invariant set encoding ( Zaheer et al. , 2017 ; Edwards & Storkey , 2017 ) and the decoder fθ models the distribution of Y ( e.g Gaussian case : estimating µ , σ ) using X along with the encoder outputs . Its objective is to maximize the log likelihood over the ( unknown ) task distribution p ( T ) . All tasks provided by the data generating process are considered as Monte Carlo samples , respectively : ∑ Tk∼p ( T ) log p ( Y |X , Xc , Yc ) = ∑ Tk∼p ( T ) n∑ i=1 logN ( yi|µi , σi ) where ( µi , σi ) = fθ ( xi , r ) , r = fφ ( { xj , yj } mj=1 ) ( 1 ) The set encoding architecture differentiates the type of NPs . In Conditional neural process ( CNP ) ( Garnelo et al. , 2018a ) , the representation r is a deterministic variable such that , by using mean aggregation , the encoder maps the context set into a single deterministic representation r = 1/m ∑m j=1 fφ ( xj , yj ) . Neural process ( NP ) ( Garnelo et al. , 2018b ) introduces a probabilistic latent variable z to model functional uncertainty as a stochastic process such that the parameters of output distribution may change according to the sampled value of z . Due to the intractable log-likelihood , a training objective is derived based on variational inference which can be decomposed into two terms , reconstruction term and regularization term : log p ( Y |X , Xc , Yc ) ≥ Eqφ ( z|X , Y ) [ log pθ ( Y |X , r , z ) pθ ( z|Xc , Yc ) qφ ( z|X , Y ) ] = Eqφ ( z|X , Y ) [ log pθ ( Y |X , r , z ) ] − KL ( qφ ( z|X , Y ) ||pθ ( z|Xc , Yc ) ) ( 2 ) For simplicity , the prior distribution pθ ( z|Xc , Yc ) is approximated by qφ ( z|Xc , Yc ) . However , as pointed out in Kim et al . ( 2019 ) , the mean aggregation over the context set is too restrictive to describe the dependencies of the set elements . To enhance the expressiveness of the task representation , Attentive neural process ( ANP ) accommodates an attention mechanism ( Vaswani et al. , 2017 ) into the encoder , which generates the local deterministic representation ri corresponding to a target data point xi , and addresses the underfitting issue in NP . Although additional set encoding methods such as the kernel method and the bayesian aggregation considering task information have been suggested ( Xu et al. , 2020 ; Volpp et al. , 2021 ) , the attentive neural process is mainly considered as the baseline in terms of model versatility . 2.2 BAYESIAN ATTENTION MODULE . Consider m key-value pairs , packed into a key matrix K ∈ Rm×dk and a value matrix V ∈ Rm×dv and n queries packed into Q ∈ Rn×dk , where the dimensions of the queries and keys are the same . Attention mechanisms aim to create the appropriate values O ∈ Rn×dv corresponding to Q based on the similarity metric to K , which are typically computed via an alignment score function g such that Φ = g ( Q , K ) . Then , a softmax function is applied to allow the attention weight W ∈ Rn×m to satisfy the simplex constraint so that the output features can be obtained by Oi , j = Wi · Vj : Wi , j = exp ( Φi , j ) ∑m j=1 exp ( Φi , j ) ( 3 ) Note that there are many options for the alignment score function g where a scaled dot-product or a neural network are widely used . Bayesian attention module ( Fan et al. , 2020 ) considers a stochastic attention weight W . Compared to other stochastic attention methods ( Shankar & Sarawagi , 2018 ; Lawson et al. , 2018 ; Bahuleyan et al. , 2018 ; Deng et al. , 2018 ) , it requires minimal modification to the deterministic attention mechanism described above so that it is compatible with the existing frameworks , which can be adopted in a straightforward manner . Specifically , un-normalized attention weights Ŵ such that Wi , j = Ŵi , j/ ∑m j=1 Ŵi , j are sampled via the variational distribution qφ ( Ŵ |Q , K ) , which can be trained via amortized variational inference : log p ( O|Q , K , V ) ≥ Eqφ ( Ŵ |Q , K ) [ log pθ ( O|V , Ŵ ) ] − KL ( qφ ( Ŵ |Q , K ) ||pθ ( Ŵ ) ) ( 4 ) Considering the random variable Ŵ to be non-negative such that it can satisfy the simplex constraint by normalization , the variational distribution qφ ( Ŵ |Q , K ) is set to Weibull ( k , λ ) and the prior distribution pθ ( Ŵ ) is set to Gamma ( α , β ) . This λ can be obtained by the standard attention mechanism , on the other hands , α can be either a learnable parameter or a hyper-parameter depending on that the model follows key-based contextual prior described in the following paragraph . The remaining variables k and β are regarded as user-defined hyper-parameters . By introducing Euler–Mascheroni constant γ ( Zhang et al. , 2018 ; Bauckhage , 2014 ) , the KL divergence in Equation 4 can be computed in an analytical expression ( shown in Appendix A ) . KL ( Weibull ( k , λ ) ||Gamma ( α , β ) ) = γα k −α log k+βλΓ ( 1+ 1 k ) −γ−1−α log β+ logΓ ( α ) ( 5 ) Samples from Weibull distribution can be obtained using a reparameterization trick exploiting an inverse CDF method : λ ( − log ( 1 − ) ) 1k , ∼ Uniform ( 0 , 1 ) . Note that mean and variance of Weibull ( k , λ ) are computed as λΓ ( 1 + 1/k ) and λ2 [ Γ ( 1 + 2/k ) − ( Γ ( 1 + 1/k ) ) 2 ] . It can be observed that the variance of obtained samples decreases as k increases . Key-based contextual prior To model the prior distribution of attention weights , key-based contextual prior was proposed . This method allows the neural network to calculate the shape parameter α of the gamma distribution as a prior distribution . This leads to the stabilization of KL divergence between standard attention weights and sampled attention weights , and it prevents overfitting of the attention weights ( Fan et al. , 2020 ) . In this paper , we explain the reason why the key-based context prior is important for capturing an appropriate representation focusing on a context dataset from the prosepective of information theory .

Summary: The paper proposes a novel method for the Attentive Neural Processes paradigm by adding stochasticity to the weights of the cross-attention module between the context and target representations. These weights are drawn from a proposal distribution which is a Weibull distribution with parameters determined by the context and target input points. A new KL regularization term is added to the total loss to enforce the proposal distribution be close to the gamma distribution determined only by the context points, in a similar fashion to the regularization term for the latent representation between the whole set of target and context points and that of the context points only. The authors include, from Fan’s paper, the closed-form KL divergence between the Weibull and Gamma distributions, making the whole loss be differentiable w.r.t. the network parameters thanks to the reparametrization trick, as in the original NP paper. The new loss is linked to the gain of the target distribution by means of information theory, by making the assumption that the new regularization implies not only that the mutual information between the target distribution and the target points is maximized, but also the mutual information between the representation for the context and target points and the latent representation of the context variables only is minimized, enforcing the latent variable to also consider the context points, and not otherwise as suggested by the original NP formulation.