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

Abstract: 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 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 self- and mixed 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.

Reviews and author responses:

 Thanks to the authors for their response. The lack of novelty remains the major issue of this work though I agree this work is well motivated.   The novelty is limited because based on my understanding, with a single target sample, the model reduces to the attentional neural processes without a probabilistic encoder path. In this sense, improved uncertainty seems to rely on a large amount of target points. That being said, if you have access to more target data (in a batch), you will naturally have a better uncertainty estimation. At least this has not been explicitly discussed in the paper. For regression tasks, it does not seem to make much sense if only uncertainty estimation is considered without a measure of prediction error, e.g., MSE. In addition, the planed addition of  the GT trajectory in figure 1 is not found in the appendix. Therefore, I would not increase our original rating.

 Dear authors,

I have read the rebuttal and increased my score accordingly. 
Thanks for fixing the consistency issue. The modification from mixed attention to cross attention makes sense to me. 
I'm still slightly preferring another round of review because of the lack of sufficient experimental evaluation in the current form of the paper. 


Neural processes are a recent proposed type of flexible stochastic processes that models distributions of functions. The distribution of functions implicitly captured by a global latent variable Z that are fed into neural networks. Unlike NPs, the function-space view of Gaussian processes model the function distribution directly as a Gaussian using mean and covariance functions. This work proposes to combine the strength of the two by using attention-based neural networks to model the conditional mean and covariance functions, getting rid of the global latent variable Z.  ## Summary

Neural processes are a recent proposed type of flexible stochastic processes that models distributions of functions. The distribution of functions implicitly captured by a global latent variable Z that are fed into neural networks. Unlike NPs, the function-space view of Gaussian processes model the function distribution directly as a Gaussian using mean and covariance functions. This work proposes to combine the strength of the two by using attention-based neural networks to model the conditional mean and covariance functions, getting rid of the global latent variable Z. 

## Strengths

* Fair amount of work seems to have been devoted to the design of the multi-headed attention architectures and they seem to work very well in the toy examples shown (better than architectures used in previous works including NP and Attentive NP).

## Weaknesses

* The main concern that I have after reading the text is the proposed model construction may not lead to a proper stochastic process. Specifically, the Kolmogorov extension theorem requires that all marginal distributions of the stochastic process to be consistent with each other. However, in the proposed model, I can't see how this could not be violated. How do you ensure that the marginal distributions will match for (f(x2), f(x3)) in the cases where X_T = (x1, x2, x3) and X_T = (x2, x3, x4)?

* The description of Neural Processes in the paper is very confusing and gave me the feeling that the nature of the model is mis-interpreted. The whole background section on NPs are focused on the architecture, or what the authors referred as "deterministic path" and "latent path", without mentioning why these architecture emerges (they are there because we want to do amortized inference). The change of latent variable name $z_T$ to $z_C$ is also problematic as there is only one shared latent variable in the definition of NPs. If you take a look at Garnelo et al. (2018b), there is only $z$. It also seems to me that the evidence lower bound given in (2) is not well-understood that there is an approximation made to the prior so that p(z|context) can be replaced by q(z|context). I strongly encourage the authors to revisit the NP paper and revise the whole submission so that it does not contain misleading information to readers.

* Besides, there seems to be confusion about NPs and conditional NPs. The proposed model are eventually trained in a way similar to conditional NPs (using the conditional distribution), which is different from other NPs trained by the conditional evidence lower bound with approximated priors.

* The figures are generally not very descriptive. It is unclear from Figure 2 how the attention module works (which is the main contribution of the work). In Figure 3, it is confusing to concat $h_c$ and $h_t$ horizontally. Shouldn't it be concatenated in the first axis?

* The empirical results are weak and are missing important baselines. I would expect ConvCNPs and GNPs to work equally well on the 1D regression examples. The proposed method also generates quite a few artifacts in the 2d image experiments. This might be related to the fact that the proposed model may not satisfy the consistency requirements of stochastic processes.

## Overall suggestions

I vote for rejection mainly due to the questionable validity of the proposed model as a stochastic processes, as the parameterization can violate conditions of the Kolmogorov extension theorem.

## Questions and minor points

* In section 5, you mentioned that the Gaussian Neural Processes share similar motivation to your work, but termed their seek of translation equivariance as weaknesses. However, is it also fair to say their construction ensures the model is a valid stochastic process but yours doesn't?
* You may want to mention $X_C \subset X_T$ at the beginning of section 2.1.
* Using \mathrm to wrap "softmax" and "Concat" makes it look better.
* It is better to remove claims like "we can predict DKNP will better adapt to even more change points composed of multiple kernel types than NP variants.  I vote for rejection mainly due to the questionable validity of the proposed model as a stochastic processes, as the parameterization can violate marginal consistency required by the Kolmogorov extension theorem.

 Q1. It would be helpful to add some computational analysis.

- As mentioned in [1], the computation costs for n context and m target points are below. DKNP has attention mechanisms as in ANP, so it costs O(n(n+m)) which is much computationally more efficient than GPs.
| GP | NP | ANP |
|:---:|:---:|:---:|
| O((n+m)^3) | O(n+m) | O(n(n+m)) |


[1] Kim et al., Attentive Neural Processes, ICLR, 2019


Q2. One confusing point is in figure 1, the mean is achieved only given the three black context points, right? If not seeing the target points, how does DKNP manage to achieve a periodic mean estimate? There is no corresponding explanation of figure 1 in the main text.

- Figure 1 shows the visualization of results after training models on samples generated from a periodic kernel with fixed hyperparameters. As a result, DKNP can estimate the periodic mean even with three context points. For clarification, we have added this explanation in the capture of the figure.


Q3. Is permutation invariance also guaranteed?
- DKNP guarantees permutation invariance. DKNP consists of attention (without positional encoding), position-wise fully-connected layers, and a multivariate normal distribution. Attention and position-wise fully-connected layers have the property of permutation equivariance, and a normal distribution has the property of permutation invariance as GP does; therefore, DKNP guarantees permutation invariance.

Q4. In the original ANP paper, they showed examples of the image completion task with a reasonable diversity. How comes the diversity in figure 9 here is so bad for ANP?
- Referring to Figure 4 in the original ANP paper, ANP fails to generate diverse images with 100 context points, as opposed to the result with much fewer context points (e.g., 10 and 30).On the other hand, Figure 9 in our paper shows that DKNP can generate diverse images with 100 context points compared to ANP's. This reveals that DKNP outperforms baseline models in terms of capturing the global stochasticity while generating coherent samples.



 
Q4. The empirical results are weak and are missing important baselines. I would expect ConvCNPs and GNPs to work equally well on the 1D regression examples.
- Thank you for pointing this out.  As we mentioned in related works, the focus of two papers, ConvCNP and GNP, is translation equivariance, which are orthogonal to our research direction. However, to best reflect your suggestion, we tried both ConvCNP and GNP in the 1D setup. 

- For each experiment, there exists some difference in the training settings of DKNP and ConvCNP: 
1) ConvCNP and GNP are trained using target log-likelihood, while DKNP is trained using both context and target log-likelihood
2) ConvCNP does not add noise in the standard deviation of the output variable, while DKNP adds 0.1 following the training conventions in NP. To note that GNP sets adding noise as a learnable parameter.
To compare models in more fair setups, we conducted experiments in all different setups. The results are as follows:

| DKNP setting | DKNP | ConvCNPXL |
|:---:|:---:|:---:|
| RBF | 1.109   (0.001) | 0.97   (0.013) |
| Periodic | 0.941   (0.006) | 0.67   (0.023) |
| Matern | 1.039 (0.004) | 0.887 (0.017) |

| ConvCNP setting | DKNP | ConvCNPXL |
|:---:|:---:|:---:|
| RBF | 1.335   (0.011) | 1.254   (0.045) |
| Periodic | 0.49 (0.098) | 0.802   (0.013) |
| Matern | 1.097 (0.118) | 1.097 (0.056) |

|  DKNP setting | DKNP  | GNP |
|:---:|:---:|:---:|
| RBF | 1.109   (0.001) | 1.187   (0.000) |
| Periodic | 0.941   (0.006) | 1.045 (0.003) |
| Matern | 1.039 (0.004) | 1.117 (0.001) |

| GNP setting | DKNP | GNP |
|:---:|:---:|:---:|
| RBF | 1.6   (0.110) | 2.112 (0.001) |
| Periodic | 1.422 (0.024) | 1.075   (0.006) |
| Matern | 1.358 (0.353) | 1.948 (0.010) |

- In ConvCNP's setup, DKNP outperforms ConvCNP. In GNP’s setup, GNP achieved better performance. However, as the original GNP paper states, GNP has a severe scalability issue, which limits the application of the model in high-dimensional scenarios (e.g., image completion). However, we believe that both ConvCNP and GNP have their strength in translation equivariance, and it would be an interesting future direction to incorporate their idea with our motivation.


Q5. The proposed method also generates quite a few artifacts in the 2d image experiments. This might be related to the fact that the proposed model may not satisfy the consistency requirements of stochastic processes.

- We would be grateful if you could elaborate more on the detailed description of the artifacts you mentioned. If we best guess that the artifact you mentioned is "kinks"—abrupt pixel change, those artifacts can also be observed in other NPs (which do satisfy the consistency requirement) that use attention as in Figure 8 and 9, while no artifact can be found in the vanilla NP. We hypothesize that this phenomenon is due to attention modules, rather than consistency requirements


Q6. The figures are generally not very descriptive. 

- As for the figure of the model architecture, we have reflected your suggestion.


Q7. In section 5, you mentioned that the Gaussian Neural Processes share similar motivation to your work, but termed their seek of translation equivariance as weaknesses. However, is it also fair to say their construction ensures the model is a valid stochastic process but yours doesn't?

- We meant that the focus of their paper is orthogonal to ours, not to claim that those are their weakness. Regarding the validity of the stochastic process, we have added further discussions in the paper. 


 Q1. The main concern that I have after reading the text is the proposed model construction may not lead to a proper stochastic process. Specifically, the Kolmogorov extension theorem requires that all marginal distributions of the stochastic process to be consistent with each other. How do you ensure that the marginal distributions will match for (f(x2), f(x3)) in the cases where X_T = (x1, x2, x3) and X_T = (x2, x3, x4)?

- As you pointed out, it is hard to prove that the current version of DKNP satisfies consistency. We admit that we have overlooked this property while focusing on other appealing properties of DKNP—capturing the relationship between the target variables and inferring the underlying prior distribution. 
If we replace mixed attention with cross-attention, the representation of target points, w, is calculated based only on the set of context points. Therefore, given the context points, the marginal distributions of target points are the same in the two cases you mentioned, satisfying consistency. 
To reflect your suggestion, we have replaced multihead mixed attention (MMA) in the covariance path with multihead cross-attention in the manuscript. By doing so, we empirically found that DKNP with cross-attention performs slightly better than the current version of DKNP in terms of log-likelihood. 
Additionally, due to the residual connection (using the same attention module in [1]), the revised model does not require the contexts (key, value) to infer the covariance matrix of the target points, allowing us to visualize the prior (i.e. the covariance matrix of the target points given no context points) while preserving the consistency requirement.
We deeply appreciate your insightful comment that enabled us to derive the cross-attention version of DKNP, which both satisfies consistency and adheres to the property and the initial motivation of ours. 

[1] Parmar et al., Image Transformer, ICML, 2018


Q2. The description of Neural Processes in the paper is very confusing and gave me the feeling that the nature of the model is mis-interpreted.
- Thank you for your detailed comment. 
We implemented NP architectures following [2] and [3], which have both deterministic and latent paths in NPs. We admit that we did not mention this modification to NPs in the paper.
As for the subscript in the latent variable, we attempted to make a distinction between the latent variables that encode context and target points, which may lead to confusion. 
We have updated the background section following your concerns.

[2] Kim et al., Attentive Neural Processes, ICLR, 2019 

[3] Lee et al., Bootstrapping Neural Processes, NeurIPS, 2020

Q3. There seems to be confusion about NPs and conditional NPs. The proposed model are eventually trained in a way similar to conditional NPs (using the conditional distribution), which is different from other NPs trained by the conditional evidence lower bound with approximated priors.
- As you well pointed out, the objective function of ours is close to CNP. However, we claim that modeling the correlation between output variables is closer to NP as the correlation allows to generate more coherent samples.



 Thank you for appreciating the value of our work. In particular, we are grateful to you for agreeing on one of the main contributions—decoupling the full covariance from the NP architecture. Also, we have fixed several grammatical errors you pointed out.

 Q1. The novelty is limited. In the architecture, the path for generating the mean vector is the deterministic path of ANPs. Learning kernel function by the attention mechanism is a common strategy in many NPs-based papers.
- As you well pointed out, the previous works have mainly focused on modeling the mean path. To the best of our knowledge, we first proposed modeling the covariance path using attention and demonstrated the behavior of the architecture compared to other NPs that utilize attention for modeling the mean path.  


Q2. What does the “decoupled” mean in the paper? The separate mean and variance? Actually, they shouldn’t be decoupled since they share kernel matrices, e.g., K(Xc, Xc) and K(Xc, Xt).
- As opposed to previous NP architectures, we intend to "decouple" an explicit covariance matrix from the NP architecture.  Regarding the calculation of the mean, you are correct on the fact that the kernel matrices (e.g., K(Xc, Xc) and K(Xc, Xt)) are used in the mean of GP posteriors. However, in DKNP, both the mean and variance paths are predicted based on the shared hidden representation, h. In other words, the exchange of information happens differently from GPs.


Q3. The comparison between the proposed method and baselines is not fair.
- The same context and target points are used to evaluate all models, including the order and the number of inputs. We would appreciate it if you could elaborate more on the part that seems unfair in our experiment.


Q4. Providing MSE in the experiments will make the results more convincing.
- Thank you for the suggestion. The main task we focus on is uncertainty estimation, where measuring the performance in log-likelihood is a metric convention in the literature.


Q5. In Figure 1 (left), 4 and 6, the ground truth (the blue line for mean and the shade for variance) should be added.
- Thank you for the suggestion. Aligning with our main contribution, we wanted to emphasize that the uncertainty bound predicted by DKNP well covers target points and therefore did not visualize the GT trajectory. To reflect your suggestion, we plan to add new figures that include the GT trajectory along with sample points and uncertainty bounds in Appendix.


Q6.What’s the prior of your model?
- The prior of DKNP is the output distribution that is predicted with target points without any context point. Unlike the previous NPs that require the representation of context points to make inference on target points, DKNP's prior is purely data-driven based on the training data and can generate samples without encoding any context point.


 We appreciate all reviewers for providing high-quality reviews and valuable feedback. We also apologize that it took us some time to conduct all experiments again and revise the paper, including figures and experimental results, to best reflect the reviewer’s feedback. For convenience, we denote Reviewer Xxs5, Reviewer Ngox, Reviewer 1GBU and Reviewer 65De as R1, R2, R3 and R4, respectively. 

We are excited to see that reviewers found our proposed method well-motivated (R2, R3); the superiority of DKNP  in terms of uncertainty estimation (R1, R2, R3, R4); interesting connections between GPs (R2, R4); well-designed architecture (R1, R3); the potential of decoupling the covariance function from the model architecture. 

The revised manuscript has been uploaded in response to reviewer feedback. We appreciate that the valuable suggestions by the reviewers have improved the quality of the manuscript. 

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.  Perhaps it is because of my lack of background on neural processes, but from reading this paper and from a brief scan of the literature, it seems that the major leap forward for neural processes was made by [Kim et al., 2019] with the use of attention. The authors make heavy use of attention in this paper too, and although experiments show that improvements are made upon ANP, I think only time will tell whether this decoupling of mean and covariance functions will prove to be a major advancement in the long term. As I believe is typical with deep learning papers, theoretical results are non-existent, but empirical results are sufficiently extensive for this not to be a drawback. 
However, from my very limited judgment, the proposed method is well-motivated, and has intuitively appealing properties, such as those outlined around equations (5) and (6), and I am leaning towards acceptance. 

As an aside, there are some English errors which I think should be corrected. For example, in the abstract, "we introduce a new neural stochastic processes" should be "we introduce a new class of neural stochastic processes", and just above the start of Section 2, it says "a neural stochastic processes that explicitly learn..." which should also be corrected similarly.  Despite my lack of knowledge of the relevant literature, I still lean towards acceptance, because I like the idea of decoupling the mean and covariance function and believe this idea may lead to more advancements in the future. 

This paper proposes Decoupled Kernel Neural Processes (DKNPs), a new neural stochastic process, which improves uncertainty estimation by learning a separate mean and kernel function to directly model the covariance between output variables in a data-driven manner. It is demonstrated that DKNPs can improve 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. Strengths:

(1) The motivation of this paper is clear and interesting, improving uncertainty estimation for the highly complex data distribution.

(2) The model architecture looks straightforward. 

Weaknesses:

(1) The novelty is limited. In the architecture, the path for generating the mean vector is the deterministic path of ANPs. Learning kernel function by the attention mechanism is a common strategy in many NPs-based papers.  

(2) What does the “decoupled” mean in the paper? The separate mean and variance? Actually, they shouldn’t be decoupled since they share kernel matrices, e.g., K(Xc, Xc) and K(Xc, Xt).

(3) The comparison between the proposed method and baselines is not fair. In the test stage, the proposed model takes all test samples as inputs simultaneously to predict the full covariance matrix. However, other models based on NPs task one batch of test samples as input. Could the author provide the results of the proposed method under the same test setting as other NPs-based methods? 

(4) Mean squared error (MSE) is also an important evaluation of regression tasks. Providing MSE in the experiments will make the results more convincing. 

(5) In Figure 1 (left), 4 and 6, the ground truth (the blue line for mean and the shade for variance) should be added. It will provide a straightforward comparison between each prediction and the ground truth.

(6) I am confused with this statement: “Note that without any context, …, allowing us to view the prior learned by DKNPs, unlike conventional NPs”. What’s the prior of your model? More details are needed to make the paper self-contained.

 The motivation is interesting, but the model somewhat lacks novelty.

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. 
 
I think the paper is well-written with sufficient background introduction for people to follow and understand the difference between DKNP and ANP. The comparison between GP and DKNP is helpful for understanding. I think it is an acceptable paper to the community with a significant contribution.

Some detailed comments:

1. It would be helpful to add some computational analysis. One drawback of GP is its scalability to large-scale datasets. But the GP's benefit is its good uncertainly calibration with small data. Where does DKNP locate in this spectrum, also compared with NPs?

2. One confusing point is in figure 1, the mean is achieved only given the three black context points, right? If not seeing the target points, how does DKNP manage to achieve a periodic mean estimate? There is no corresponding explanation of figure 1 in the main text. 

3. Is permutation invariance also guaranteed?

4. In the original ANP paper, they showed examples of the image completion task with a reasonable diversity. How comes the diversity in figure 9 here is so bad for ANP?



 I think the model is novel and has a significant contribution to the community. The connection between DKNP and GP is interesting. 

Meta-review: This paper proposes Decoupled Kernel Neural Processes (DKNPs), a new neural stochastic process, which learns a separate mean and kernel function to directly model the covariance between output variables. Numerical experiments on 1-D regression and 2-D image completion are provided.

There  was a concern that the original version of the proposed model was not a valid stochastic process, since the consistency condition of Kolmogorov Extension Theorem might not be satisfied. The authors fixed this by replacing multihead mixed attention (MMA) in the covariance path with multihead cross-attention.
Overall, reviewers find the work interesting, but there remain concerns that the novelty is limited and that  the current work lacks sufficient experimental evaluation.

Decision: Reject