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 paper proposes Decoupled Kernel Neural Processes (DKNPs), which is a variant of neural processes first proposed a few years ago. The main contribution is to decouple the process of learning the mean and the covariance functions separately. Empirical evaluations are provided to demonstrate its usefulness.