Decoupled Kernel Neural Processes: Neural Network-Parameterized Stochastic Processes using Explicit Data-driven Kernel

Neural Processes ( NPs ) are a class of stochastic processes parametrized by neural networks . Unlike traditional stochastic processes ( e.g. , Gaussian processes ) , which require specifying explicit kernel functions , NPs implicitly learn kernel functions appropriate for a given task through observed data . While this data-driven learning of stochastic processes has been shown to model various types of data , the current NPs ’ implicit treatment of the mean and the covariance of the output variables limits its full potential when the underlying distribution of the given data is highly complex . To address this , we introduce a new class of neural stochastic processes , Decoupled Kernel Neural Processes ( DKNPs ) , which explicitly learn a separate mean and kernel function to directly model the covariance between output variables in a data-driven manner . By estimating kernel functions with cross-attentive neural networks , DKNPs demonstrate improved uncertainty estimation in terms of conditional likelihood and diversity in generated samples in 1-D and 2-D regression tasks , compared to other concurrent NP variants . Also , maintaining explicit kernel functions , a key component of stochastic processes , allows the model to reveal a deeper understanding of underlying distributions . 1 INTRODUCTION . Neural processes ( NPs ) ( Garnelo et al. , 2018a ; b ) are a class of stochastic processes parametrized by neural networks . By embracing statistical properties in stochastic processes , NPs can effectively estimate the uncertainty of underlying distributions of functions with a set of realizations and their data points . Different from traditional stochastic processes ( e.g. , Gaussian processes ( GP ) ( Rasmussen & Williams , 2006 ) ) , NPs learn data-driven stochastic processes without a need to specify or keep an explicit form of kernel functions . As a result of their simplicity and flexibility , there have been numerous efforts to further develop improved variants of NPs ( Kim et al. , 2019 ; Lee et al. , 2020 ; Gordon et al. , 2020 ) and apply them to various downstream tasks ( Singh et al. , 2019 ; Requeima et al. , 2019 ) . Though significant progress has been made in NPs , the current architectures of NPs either fails to capture output dependencies as in Conditional NPs ( Garnelo et al. , 2018a ; Gordon et al. , 2020 ) , or indirectly capture the full stochasticity present in the traditional stochastic processes . For instance , different from GPs , conventional NPs reserve stochasticity in a global latent variable and output variables separately . The output variables estimate point-wise uncertainty , which corresponds to the diagonal elements of a kernel matrix . Similarly , the global latent variable takes charge of the functional uncertainty and diversity , represented by the full covariance matrix in GPs . Due to this inductive bias of conventional NPs , the role of estimating functional stochasticity is mainly assigned to a fixed-length vector ( i.e . the global latent variable ) , and consequently , capturing the underlying distributions can be restricted in complex scenarios ( e.g. , variable relationships are periodic or abruptly changing at a certain point ) . Although several approaches attempt to alleviate the problem by introducing attention ( Kim et al. , 2019 ) and bootstrapping ( Lee et al. , 2020 ) on top of conventional NPs , the problem still exists as the architectural limitation ( i.e . implicit modeling of the mean and covariance ) has not been addressed directly . Besides this , as NPs implicitly learn the kernel functions inside the model , the interpretability of kernels such as in GPs ( Lloyd et al. , 2014 ) is diminished . To address this concern , we propose Decoupled Kernel Neural Processes ( DKNPs ) , a new class of neural stochastic processes that explicitly learn a separate mean and kernel function to directly model the covariance between output variables in a data-driven manner . Our experiments in 1-D and 2-D regression tasks reveal that the DKNP outperforms concurrent NP variants in terms of predictive likelihood , better global coherence of generated samples , and improved interpretability via explicitly learned kernels . 2 BACKGROUND . 2.1 NEURAL PROCESS . Given a stochastic process sample consisting of n points , let us denote the input and output as X = { xi } ni=1 and Y = { yi } ni=1 , respectively , where xi ∈ Rdx and yi ∈ Rdy . For a set of target input XT = { xi } i∈T ⊂ X , NPs model the conditional distribution of target outputs YT conditioned on the context set ( XC , YC ) = { ( xi , yi ) } i∈C using a factorized Gaussian distribution : log p ( YT |XT , XC , YC ) = ∑ i∈T log p ( yi|xi , XC , YC ) . ( 1 ) For obtaining the predictive distribution log p ( yi|xi , XC , YC ) , NPs use an encoder-decoder architecture that ensures the permutation invariance of the predictions of the target points given the context set ( XC , YC ) . Following Kim et al . ( 2019 ) and Lee et al . ( 2020 ) , we consider the NP encoder consisting of two separate paths , namely the deterministic path and the latent path . For the deterministic path , fθ represents each context points in { ( xi , yi ) } i∈C as ri ∈ Rdr , i.e. , ri = fθ ( xi , yi ) . Then , we aggregate the ri ’ s by averaging them across all context points , rC = 1nc ∑ i∈C ri where nc = |C| . This vector rC ∈ Rdr is the summarized representation of context points , and it is permutation invariant over the order of ( xi , yi ) ∈ ( XC , YC ) . The latent path of the NP encoder operates in a similar fashion to the deterministic path , i.e. , eC = 1 nc ∑ i∈C ei where ei = gϕ ( [ xi ; yi ] ) . Unlike the deterministic path , however , the la- tent path uses stochastic layers for obtaining a distribution of the latent variable z ∈ Rdz ; q ( z|eC ) = N ( z ; µz , σ2z ) where µz , σz are the output of the additional fully-connected layer applied to eC . Finally , by concatenating these aggregated vectors rC , z with the target inputs xi ∈ XT , the decoder hψ produces the predictive distribution p ( yt|xt , rC , z ) = N ( y ; µy , σ2y ) where ( µy , σy ) = hψ ( xt , rC , z ) . Note that the σy is used for capturing the point-wise uncertainty . During the training phase , the parameters are learned by maximizing the evidence lower bound of log p ( YT |XT , XC , YC ) via the reparametrization trick ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) as follows : log p ( YT |XT , XC , YC ) ≥ Eq ( z|XT , YT ) [ log p ( YT |XT , z ) ] −KL ( q ( z|eT ) ∥q ( z|eC ) ) ( 2 ) where q ( z|eT ) = q ( et|XT , YT ) is a distribution of latent variable by encoding the set of target points ( XT , YT ) = { ( xi , yi ) } i∈T . This objective function consists of two parts : 1 ) the reconstruction loss for the target points and 2 ) the KL divergence term , minimizing the divergence between two distributions q ( z|XT , YT ) and q ( z|XC , YC ) . Note that , in practice , we assume that XC ⊂ XT during the training phase . The KL divergence term encourages two distributions inferred by the context sets and target sets to be similar , which is reasonable because two sets are generated from the same function . Therefore , during the inference phase , the distribution of q ( z|XC , YC ) captures the functional stochasticity which is demonstrated with the coherent sample generation . It can be thought that z learns to capture the correlation of output variables of the stochastic processes . 2.2 ATTENTIVE NEURAL PROCESSES . Multihead Attention Given n key-value pairs of matrices K ∈ Rn×dmodel and V ∈ Rn×dmodel , and m queries Q ∈ Rm×dmodel , the scaled dot product attention is formulated as : Attention ( Q , K , V ) = Softmax ( QK⊤√ dmodel ) V ∈ Rm×dmodel , ( 3 ) where K , V , and Q are projected by learnable linear maps WKs , W V s , and W Q t from the source S and target T . The attention mechanism can be calculated from multiple subspaces , namely , multihead attention ( MHA ) ( Vaswani et al. , 2017 ) . Denoting a single attention head as Headi = Attention ( Qi , Ki , Vi ) , the aggregate attention from multiple subspaces can be expressed as : MHA ( Q , K , V ) = Concat ( Head1 , ... , HeadH ) WO ∈ Rm×dmodel , ( 4 ) where WO is the learnable linear map for aggretating the subspaces . Attentive NPs Attentive NPs ( ANPs ) leverage attention to resolve the underfitting issue in NPs . Unlike NPs that produce a single variable rC from the deterministic path , ANPs utilize a queryspecific variable r∗i by applying the attention score ai for each ri during the aggregation of the deterministic path , formulated as r∗i = ∑ i∈C ai · ri . The attention-based aggregation of ANPs resembles how the GPs utilize the correlation to estimate the predictive distribution of the context and the target points . 3 DECOUPLED KERNEL NEURAL PROCESSES . Decoupled Kernel Neural Processes use attention to explicitly learn a separate mean and kernel function so as to directly model covariances between output variables with related input variables in a data-driven fashion , which is contrary to conventional NPs that implicitly model the mean and kernel function through the latent variable z . As shown in Figure 2 , DKNPs estimate the predictive distribution as multivariate Gaussian N ( y ; µy , Σy ) by using two attention-based deterministic paths using Multihead Cross-Attention ( MCA ) to model the mean vector ( upper path ) and the full covariance matrix ( lower path ) . Here , attention modules are extensively utilized as Le et al . ( 2018 ) ; Kim et al . ( 2019 ) have demonstrated attention was helpful in achieving low predictive uncertainty near the context points . With the predictive distribution N ( y ; µy , Σy ) , DKNPs are trained and evaluated based on the function likelihood . The core design of DKNPs is motivated by the predictive posterior distribution of GPs , where XT , XC , YC are used for deriving the posterior mean,1 but only XT , XC for deriving the posterior covariance as follows : GP : p ( YT |XT , XC , YC ) = N ( YT ; Σ ⊤ XC , XT Σ −1 XC , XC YC , ΣXT , XT − Σ ⊤ XC , XT Σ −1 XC , XC ΣXC , XT ) , ( 5 ) DKNP : p ( YT |XT , XC , YC ) = N ( YT ; MLPψ ◦ MCA ( MLPθ ( XC , YC ) , MLPϕ ( XT , XC ) ) , MLPω ◦ MCA ( MLPϕ ( XT , XC ) ) ) , ( 6 ) where , in contrast to GPs , DKNPs ’ covariance is learned via attention.2 This decoupled process allows DKNPs to explicitly learn the prior of the given dataset , and thus act as a true generative process , ensuring the global consistency of all points in the given stochastic process samples . Unlike DKNPs , NPs pack all information ( XC , YC , XT ) into latent variables to derive both mean and covariance , thus inherently becoming a conditional process that requires a sufficient amount of context points , unable to explicitly learn a prior . Specifically , the DKNPs pass each context point { ( xi , yi ) } i∈C , the concatenation of xi and yi , to MLPθ and represent it as ei . Similarly , we produce the representation vector of xi , hi , using another MLPϕ for all inputs { xi } i∈C∪T . Then , ei and hi are passed to the MCA module to create the mean vector µy . All heads in the MCA module perform cross-attention , Q = { hi } i∈T , K = { hi } i∈C , and V = { ei } i∈C . We adopted the architecture of MCA used in image transformers ( Parmar et al. , 2018 ) , where the original query vectors are added to the output from the MCA through the residual path . This allows to do inference the output distributions without context points , which can be considered as prior distributions of DKNPs—the learned prior of DKNP . After the MCA , the last MLPψ generates the predictive mean vector µy for each data point . Intuitively , this can be interpreted as predicting the target mean based on the context and the correlation , which resembles the estimation of the predictive distribution in GPs and ANPs . Different from the NPs , DKNPs explicitly capture the correlation between the output variables using another multihead cross-attention ( MCA ) , where Q = { hi } i∈T , K = { hi } i∈C , and V = { ei } i∈C . Then MLPω produces the representation vector wi ∈ Rdw for each position , which are combined to generate the covariance matrix Σ = WW⊤ , where Wi , : = wi and Σij = kernel ( xi , xj ) = w⊤i wj . One might consider using self-attention to let the model learn the correlation between all data points . However , the self-attention module on only X as inputs receives no indication of context and target points and therefore fails to reduce the uncertainty near the points that have high confidence ( e.g. , context points ) . Also , the interaction between the target points through self-attention does not guarantee consistency under the marginal distribution of target points when the context points are given . Lastly , it is also important to note that the representation h is shared when modeling both the mean and the covariance . This motivation is drawn from Equation 5 that the calculation of the mean is also based on the kernel matrices , Σ⊤XC , XT and Σ −1 XC , XC . To train the DKNPs , the obtained mean vector µ and the covariance matrix Σ act as parameters of a predictive distribution N ( Y ; µY , ΣY ) . Instead of maximizing the lower bound of the log-likelihood as in most NP models , the training objective of DKNPs is to maximize the tractable log-likelihood of the Gaussian as follows : log p ( YT |XT , XC , YC ) = logN ( YT ; µY , ΣY ) where ΣY = WW⊤ . ( 7 ) 1We follow the typical GP formulation where the mean function is set to zero . 2Note that DKNPs assume C ⊂ T during training , which makes Eq . 5 and Eq . 6 technically different . This difference comes from DKNPs using all data points for better learning the kernels , unlike fixed-kernel GPs . DKNPs and GPs , however , are different methods and superior empirical performance led to the current design choice . Although the proposed objective function is equivalent to CNP ’ s ( Garnelo et al. , 2018a ) , modeling the correlation between output variables for capturing functional stochasticity shares the same motivation of NPs , thus DKNP being one of NP variants .

Summary: The authors proposed a novel extension of the neural process by explicitly and separately modeling the mean and the covariance of output variables via multi-head attention. The general NP has the problem that it doesn't well calibrate the output uncertainty when the data sample size is small. However, DKNP seems to mimic the uncertainty calibration behavior similar to GP which is encouraging. The experimental results look convincing with an extensive comparison over state-of-the-art NP variants.