Neural Bootstrapping Attention for Neural Processes

1 INTRODUCTION . Neural Processes ( NP ) ( Garnelo et al. , 2018b ) define distributions over functions given a set of observations , and are trained via a meta-learning framework so that it can adapt to new functions rapidly . NP learns to model a wide range of stochastic processes and can estimate the uncertainty over the predictions with less computational effort , compared to Gaussian Processes ( GP ) ( Rasmussen , 2003 ) . However , NP frequently suffers from a fundamental drawback of underfitting . As a remedy of the underfitting issue , Attentive Neural Processes ( ANP ) ( Kim et al. , 2018 ) applies the attention modules to the encoder network . Despite this modification , a single Gaussian latent variable of ( A ) NP has a limitation in inducing functional uncertainty ( Louizos et al. , 2019 ) , a global uncertainty that decides the distribution over the space of trajectories or functions . Appropriate modeling of functional uncertainty in stochastic processes improves the predictive performance and diversity in function realizations ( Le et al. , 2018 ) , thus provides a principled way to guide agents to find optimal candidates in sequential decision-making problems . In these tasks , a model needs to approximate a function and estimate uncertainty correctly to optimize a black-box function whose analytic information is not given . Although GP is widely used for these tasks , these are the promising area for the application of NP because GP is computationally expensive , and it can be hard to choose an appropriate prior . Recently , Bootstrapping Neural Processes ( B ( A ) NP ) ( Lee et al. , 2020 ) modify ( A ) NP to induce more robust uncertainty estimation by employing the residual bootstrapping . However , B ( A ) NP underperforms in capturing a functional uncertainty because the residual bootstrapping works in a homoscedastic way , removing the connection between the feature and the label in its bootstrapped samples . The bootstrap strategy used in B ( A ) NP demands a higher computational burden compared to ( A ) NP since it requires multiple computations of the encoder network and additional heuristics , including the adaptation layer and the lower bound on the variance1 . Furthermore , ANP and BANP tend to overfit for a simple regression task rather than learning the underlying heteroscedasticity . We observed that this problem is a by-product of the attention modules 1We colored the revised or added sentences in blue only for the rebuttal . This color will not appear in the final version of the manuscript . used in both models ( see Figure 1 and 5 ) . This finding suggests that effective regularization of attention modules is required to prevent overfitting . Additionally , as explained in Section 4 , ANP and BANP often tend to estimate homogeneous uncertainty regardless of whether the observation is given . To resolve the problems of previous NPs , we introduce a novel bootstrapping method for neural processes , Neural Bootstrapping Attentive Neural Processes ( NeuBANP ) . Motivated from the recent work on efficient bootstrapping of the neural network , Neural Bootstrapper ( NeuBoots ) ( Shin et al. , 2021 ) , we introduce bootstrapping of the attention in a computationally efficient way by simple modification of the input of attention modules and the loss function , instead of memory-inefficient resampling and contrived heuristics employed in BANP . The simplicity and computational efficiency of NeuBANP directly come from NeuBoots , but it does not guarantee the performance in modeling the functional uncertainty since NeuBANP operates on the meta-learning framework . Thus , we modify the method for training the model to learn the randomness present in the underlying function , allowing the model to estimate random functions generated by any stochastic process . This modification is simple but gives a strong consistency between the bootstrapped samples and the representations from the attention modules , unlike the residual bootstrapping used in BANP . Besides , our bootstrapping method , which implements the concatenation and multiplication of random bootstrap weights , operates as a regularizer on the attention networks , thus preventing overfitted predictions observed in previous attention-based NP methods . NeuBANP is trained to generate a valid predictive distribution by utilizing the uncertainty inherent in observations by bootstrapping instead of the uncertainty that depends on the prior assumption on the latent variable as in ( A ) NP . This leads to the success in capturing heteroscedasticity of the data and modeling functional uncertainty in stochastic processes . As a result , NeuBANP achieves the best performance in sequential decision-making problems such as Bayesian optimization and contextual multi-armed bandit . The experimental results demonstrate that our model provides promising capabilities as an efficient neural approximation of stochastic processes . Contributions We propose NeuBANP , an easy-to-implement and computationally efficient method for bootstrapping ANP . The proposed method has a novelty in learning a generator for bootstrapping stochastic processes under a meta-learning framework . NeuBANP resolves overfitting problem of attention modules and shows robust performance on heteroscedastic models without heuristics like the extra adaptation layer in BANP . NeuBANP estimates functional uncertainty better than BANP and achieves the best performance in stochastic optimization problems , including multi-dimensional Bayesian optimization and contextual multi-armed bandit , compared to previous NP methods . 2 PRELIMINARIES . 2.1 META-LEARNING FRAMEWORK OF NEURAL PROCESSES . Consider dataD = ( X , Y ) = { ( xi , yi ) } ni=1 ⊂ X ×Y , the pairs of inputs xi ∈ X and outputs yi ∈ Y . Let P be a probability distribution over functions f ∈ F ; yi = f ( xi ) + i where i∼N ( 0 , σ2i ) , hence P determines the distribution of D. For disjoint index sets C and T satisfying C ∪ T = [ n ] , define context DC = ( XC , YC ) = { ( xc , yc ) } c∈C and target DT = ( XT , YT ) = { ( xt , yt ) } t∈T , so that D = DC ∪ DT . The task is to learn the neural processes pθ that fits the stochastic processes f ∼ P given DC when C ∼ Pn is a randomly chosen subset of [ n ] , as follows : θ ? = argmin θ Ef∼P [ EC∼Pn [ − log pθ ( Y |X , DC ) ] ] . ( 1 ) This meta-learning framework allows pθ to learn diverse patterns in F via a prior distribution P ; hence NP can predict target points conditioned on contexts adaptively in the inference phase . 2.2 ( BOOTSTRAPPING ) ATTENTIVE NEURAL PROCESSES . ANP ANP is a variant of NP equipped with attention modules in the encoder part . See Appendix A for the detailed definition of attention operations . Let the context DC is given , and the model aims to infer y for a given feature x ∈ X . The encoder network of ANP maps ( x , DC ) into a pair of representation vectors r = ( z , h ) as follows : { sc } c∈C = SelfAttn ( DC ) , sC = mean ( { sc } c∈C ) , z ∼ N ( z|µz ( sC ) , σ2z ( sC ) ) ( 2 ) { hc } c∈C = SelfAttn ( DC ) , h = CrossAttn ( x , XC , { hc } c∈C ) ( 3 ) Here , µz and σz are single linear layers that map sC to mean and standard deviation of the latent variable z , respectively . In ANP , z and h refer to latent path and deterministic path , respectively . The deterministic path models the overall skeleton of the encoder network , while the latent path models the functional uncertainty using a stochastic global latent variable ( Garnelo et al. , 2018b ; Kim et al. , 2018 ) . Then the decoder network takes the representations r and target data x as inputs to predict µ ( x , r ) and σ ( x , r ) , the parameters of the conditional predictive distribution p ( y|x , DC ) = N ( y|µ , σ2 ) . BANP The global latent variable in ANP potentially limits the flexibility in expressing functional uncertainty . BANP proposes a method that utilizes paired bootstrapping and residual bootstrapping to model stochasticity in a data-driven way . First , they resample pairs of ( xc , yc ) with replacement to construct D̂ ( b ) C = ( X̂ ( b ) C , Ŷ ( b ) C ) . Here b = 1 , . . . , B implies the number of bootstrap samples . The resampled context is encoded into the representation vector r̂ ( b ) = ( ẑ ( b ) , ĥ ( b ) ) by replacing DC in ( 2 ) and ( 3 ) by D̂ ( b ) C . Then BANP conducts the residual bootstrapping to construct bootstrapped contexts again D̃ ( b ) C = ( X̃ ( b ) C , Ỹ ( b ) C ) . By using DC and D̃ ( b ) C in ( 2 ) and ( 3 ) separately , BANP gets the representations of contexts ( r , r̃ ( b ) ) . Finally , BANP uses an adaptation layer to merge ( r , r̃ ( b ) ) and obtain µ ( b ) , σ ( b ) through the decoder network ( see Figure 2 ) . p ( y|x , r , r̃ ( b ) ) = N ( y|µ ( b ) , ( σ ( b ) ) 2 ) , p ( y|x , DC ) ≈ 1 B B∑ b=1 p ( y|x , r , r̃ ( b ) ) . ( 4 ) Due to the residual bootstrapping , BANP conducts the encoder computation three times and the decoder computation twice for a single forward propagation . These additional calculations cause the computational bottleneck in the training and inference ( see Appendix B.6 ) . 2.3 NEURAL BOOTSTRAPPER . Repetitions of training restrain the practical use of bootstrap procedures in deep neural networks due to their high computational burden . To alleviate this , Shin et al . ( 2021 ) proposes Neural Bootstrapper ( NeuBoots ) . NeuBoots circumvents multiple training of networks by learning a bootstrap generator . Random Weight Bootstrapping Let ∑ c∈C ` ( f ( xc ) , yc ) be the loss function of interest for a neural network f . In standard bootstrap procedures , a bootstrapped neural network can be obtained by minimizing the loss function weighted by a random bootstrap weight wC : = { wc : c ∈ C } : L ( wC , f , DC ) = ∑ c∈C wc ` ( f ( xc ) , yc ) . ( 5 ) According to the choice of the distribution on wC , various bootstrap procedures can be represented under the form of ( 5 ) ; e.g. , the paired bootstrap by wC ∼ Multinomial ( n ; 1/n , . . . , 1/n ) and the Random Weight Bootstrapping ( RWB ) ( Præstgaard & Wellner , 1993 ; Newton & Raftery , 1994 ) by wC ∼ |C| × Dirichlet ( 1 , . . . , 1 ) . NeuBoots utilizes RWB to avoid the data discard problem which can occur in the standard bootstrapping . We then compute bootstrapped neural networks { f̂ ( b ) : b = 1 , . . . , B } via minimization of ( 5 ) for sampled w ( 1 ) C , . . . , w ( B ) C . Learning To Generate Bootstrap Distribution The main idea of NeuBoots is to construct a single generative network that models the bootstrapped neural networks with varying bootstrap weights in ( 5 ) . This formulation modifies the backbone network in a form of f ( x , wC ) that inputs both feature x and bootstrap weight wC . Shin et al . ( 2021 ) show that the minimizer of the following loss generates valid bootstrap evaluations that match the results of the standard bootstrap procedure : L ( f , DC ) = EwC∼|C|×Dirichlet ( 1 , ... ,1 ) [ L ( wC , f ( · , wC ) , DC ) ) ] , We call this weighted bootstrapping loss . Once this generator is trained via a single optimization procedure , we can efficiently generate bootstrapped predictions by plugging random bootstrap weights in the trained generator ; i.e. , for a feature of interest x , the trained generator inputs { w ( b ) C } Bb=1 and produces bootstrapped predictions ŷ ( b ) = f̂ ( x , w ( b ) C ) for b = 1 , . . . , B .

Summary: The authors introduce NeuBANP, an extension of B(A)NP that replaces the iterative prediction method of B(A)NP with a single prediction step following the idea behind the neural bootstrapper. This new model is computationally more efficient and able to produce more accurate uncertainty estimates as well as correctly model heteroscedastic uncertainties. The authors provide experimental results on nonparametric regression tasks, Bayesian optimisation, contextual multi-armed bandit tasks and image inpainting (in the appendix).