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: This paper proposes to improve attentive neural processes (ANP; Kim et al.) by replacing deterministic attention with a Bayesian attention module. Since the inference requires variational approximation, I think the proposed method is a variational counterpart of ANP, similar to the relation between VAEs and autoencoders.