Score-Based Neural Processes

We would like to thank the reviewer for their review and very thorough questions.

We provide a response to all the questions in order in the response and take on board the constructive criticism regards to presentation of the paper. We first provide a response to the first 5 questions, with the rest in the next comment.

• Q1 Conditional Consistency is not achieved in MNP in practice [1], see the final paragraph in section 4. “In practice, the inference model $q_{x_{1:n}} (z^{(1:T )} | y_{1:n}; \phi)$ provides an approximate prior/posterior. Ideally, if it gives the exact posterior, conditional consistency would hold perfectly. However, when the inference model is approximate, the degree of conditional consistency depends on the discrepancy between the inference model and the true posterior." Moreover, any NP that uses a posterior $q(z| y_c, \phi)$, can by definition not be conditionally consistent. In ScoreNP, we form the posterior by a $p(z)p(y|z)$ which corresponds to an exact model posterior and therefore posteriors are generated from a single prior $p(z)$, satisfying conditional consistency. Hence, our model is the first NP variant to our knowledge to maintain conditional and marginal consistency.

• Q2 We understand the concerns with the empirical results, whilst the model doesn’t strictly outperform the existing NPs on GP regression it does have strong performance on 2 Non-GP datasets. We do only use an ELBO rather than IWAE or exact log-likelihoods, which may be a reason for the reduced performance, however, after finding an issue in our code we are unable to fully comment on the performance in the current form.

• Q3 We thank the reviewer for the feedback on the organization of the work. We will work to improve on the writing in a revised version of the paper and apologise if it caused further effort during reviewing.

• Q4 This is not a claim to say that in practice the direct operation in the function space is flawed or doesn’t perform well. It is more a comment on the disparity between practice and theory when considering infinite-dimensional spaces, as one is required to perform discretisation in order to perform the operations in practice which does sacrifice the theoretical definitions. The empirical performance comparisons between function space operations and NPs is an interesting topic in its own right.

• Q5 To acknowledge the first part of this question where we refer the reviewer to sec. 4.3 in MNP [1]. “The initial SP can be arbitrarily chosen, as long as we can evaluate its marginals $p_{x_{1:n}}(y_{1:n}^{(0)})$. In our experiments, we use a trivial SP where all the outputs are i.i.d. standard normal distributed.” This implies that the initial samples for the finite marginals are drawn from $N(0,I)$ as is replicated by terminating the diffusion process at $N(0,1) = p^{1}$. Moreover, we are aware that in a Hilbert space $N(0,I)$ does not exist and don't claim otherwise. Similarly to GeomNDP [2] the use of finite dimensional marginals allows us to utilise white noise as a Target distribution. Moreover, one can define a process that terminates at a valid stochastic process in $L^{2}$. In [2], they provide different GP kernels as terminal distributions and conclude that $N(0,I)$ provides the best empirical results. Under the lebesgue measure finite marginal distributions defined $N(0_{n},I_{n})$ satisfy the Kolomogorov consistency conditions resulting in the underlying process being a valid stochastic process, however correctly not in the $L^{2}$. In the sketch proof, we used the fact that operating the probability flow ODE is equivalent to discrete transition in MNPs. We don’t naively extend $d \to \infty$, our finite dimensional diffusion also operates pointwise at any point of the stochastic process, providing a valid transition step from one finite dimensional marginal distribution to another without violating Kolomogorov Consistency conditions. This aligns the transition under our diffusion model to be equivalently a Markov Transition Operator introduced in [1]. We note that our diffusion process operates for any number of points $d$. The use of Measure Theory generally goes beyond the literature of Neural Processes which abstract the need of measure theory by the application of Kolomogorov Extension Theorem and operating on the finite marginals instead of directly in a function space. Due to the majority of the proof relying on Markov Transition Operators, we considered a sketch to be sufficient, however we will work on providing a more rigorous proof in a revised version of this work.

References:

[1] Jin Xu, Emilien Dupont, Kaspar Martens, Tom Rainforth, and Yee Whye Teh. Deep Stochastic Processes ¨ via Functional Markov Transition Operators, May 2023

[2] Mathieu, E., Dutordoir, V., Hutchinson, M., De Bortoli, V., Teh, Y. W., & Turner, R. (2024). Geometric neural diffusion processes. Advances in Neural Information Processing Systems, 36.

We would like to thank the reviewer for their response and their extensive feedback on the presentation of this work. We will take on board your comments in a revised version of this work, particularly around the discussion of other similarly named works and the whole Neural Process literature.

We also provide answers to the listed questions:

• Q1 Encoder Collapse, is when the output of the encoder becomes degenerate, generally $Enc_{\phi}(\cdot) \approx 0$. This means the latent space is degenerate so the Latent diffusion becomes trivial but the data-diffusion is unable to learn anything and generates noise. This can when the loss function has too large a weighting on the latent diffusion component.

• Q2 In our revised work, we will explain more the concept of Conditional Consistency and its impact on Neural Processes.

• Q3 We can postulate that it is due to the much less smooth nature of the functions in comparison to the others in the dataset, more extensive tuning of the architecture can hopefully improve this but we will not definitively comment on this as we have found a bug in the code which may be a reason for the weak performance. However, we agree it is definitely something that should be addressed in a revised version of this work.

We thank the reviewers again for their feedback as it will be extremely useful in the construction of a revised version of this work and hope that the provided answers are adequate for the questions raised.

We would like to thank the reviewer for their time in providing a comprehensive review. We will take on the constructive criticism and work to improve the work in a revised version.

We provide answers to the questions raised by the reviewer:

• Q1 We have since found a bug which potentially eradicates the issue of difficulty in the hyperparameters, but have yet to test extensively without it so cannot comment definitively. However, we were primarily referring to hyperparameters in the loss function for balancing the loss between score matching (data and latent) and the encoder entropy. Naively choosing a minimal loss can lead to the latent score being perfectly learnt but the encoder collapsing.

• Q2 A couple of advantages are (i) the denoising score matching objective is generally an easier objective to optimise than Maximising Log likelihood (ii) having the latent space generated with a diffusion model itself means no matter what level of context is provided the sample will be generated from a well-supported region leading to strong results. This is akin to some of the improvements of Latent Diffusion Models over VAEs.

• Q3 No the method is not very computationally demanding, a couple of minutes was an overestimation of the sampling speed and was meant to imply that generation is rather quick yet will still be more expensive than MNPs or NPs which require much fewer sampling steps.

We hope this answers the questions adequately and thank the reviewer again.

We would like to thank the reviewer for their response and respond their provided questions. We note that we have found a bug which will potentially fix the issues in training and scalability but we are yet analyse this to the full extent so will not make any definitive comments on the empirical performance.

• Q1 There are two avenues by which this task scales Increasing the dimension of the domain, i.e going from $1d$ to $2d$ gives input parameters of $(x1, x2)$ over $(x1)$. In many real physical processes this will be limited to $3d$ for $(x,y,z)$ coordinates and potentially $4d$ if time is also considered. This scaling would likely require the use of a higher-dimensional latent vector to encode the functions, with the encoder being free from any permutation invariances, at each of these scales the data diffusion is simply a flattened vector as the model operates pointwise. So any scalability limitations will primarily come from learning an effective latent embedding and diffusion model The other way of scaling is for the range of the functions to increase in dimensionality, for example RGB images $(y1, y2, y3)$ compared to grayscale $(y1)$. Scaling of this fashion would require greater capacity of the latent variable in the form of channels. This form of task scaling shouldn’t cause any issues in the scalability. The scalability aligns with that of latent diffusion models, in a revised version we will analyse and discuss this more in depth.

• Q2 The hyperparameters that were mentioned as issues are the ones in the loss function, balancing the losses of the latent diffusion, data diffusion and the entropy of the encoder does not directly transfer from task to task and simply choosing the optimal parameters based off the lowest loss function can result in encoder collapse. We have since found a bug in the code which appears to stabilise this issue, but we have not had the chance to fully analyse this again. The model is not sensitive to other hyperparameters such as learning rates, other optimiser hyperparameters, architecture parameters etc.

We hope this answers the provided questions adequately and thank the reviewer again for their response.