Paper: Graph neural processes and their application to molecular functions (ICLR-2024)

Reviewer comment (Writing):
The description of the methodology itself (molecular graph attentive encoder) is not detailed enough and very superficial (in total 5 lines of the entire manuscript).

Author response:
We thank the reviewer for raising this point. We would like to explain that the architecture of the GNN used by our MG-NPs is based on a very popular architecture (Attentive FP [1]) whose paper we reference and can be checked by the readers. The adaptation procedure is a bit trickier to replicate than the architecture (especially the one using MAML), so we used more space on that section instead, given the length constraints. We explain the differences between Attentive FP and our GNN in Appendix C.3.2. In addition, if our article is accepted we will publish our code with the camera-ready version.

Reviewer comment (Novelty):
The fine-tuning and MAML approaches for parameter adaption described in the paper are of little novelty. In addition, the theoretical benefit and motivation of the MAML tuning is not clear to the reviewer since NPs can generally already be considered as meta learnerns. Empirically, the MAML tuning sometimes improves and sometimes worsens predictive performance (see, e.g., Table~1 FP-CNP with FP-CNP (MAML) or MG-LNP with MG-LNL (MAML)).

Author response:
As far as we know, we are the first to propose parameter adaptation for NPs during meta-testing. The reason why parameter adaptation could be helpful is that the meta-learning done by a NP, and the adaptation based on the contexts, are not necessarily perfect. Unlike a Gaussian process, which will make accurate predictions on the context points it has seen during its evaluation, a NP has no theoretical guarantees of accuracy on the context points. We exploit context predictions' errors to improve adaptation to meta-test functions during meta-testing, thus helping to solve the issue of the applicability of NPs when the meta-train and meta-test functions diverge.

Indeed, it may be surprising to combine MAML and NPs since both are meta-learning methods. However, as described above, NPs may not predict the context points accurately, which provides an opportunity for parameter adaptation using the loss of the context points. Such adaptation could be done in multiple ways. One is fine-tuning for multiple epochs, and another one is taking a single iteration of gradient descent after having trained on MAML. Here, we simple use MAML as a strategy to achieve fast adaptation. The meta-training regime of a NP trained with MAML is a bit tricky: it involves computing the loss of the NP's context predictions, taking one iteration of gradient descent on the context loss (the MAML inner update), computing the loss of the NP's target predictions, and taking one iteration of gradient descent through the inner update (the MAML outer update).

MAML always improves the performance of a single-task model in the low-data regime (Table 1, GNN with MAML vs GNN with just 20 context points), as expected. However, MAML doesn't always improve the performance of NPs. This may be because of training instabilities, an issue that is well documented for MAML [2]. We tried to improve robustness by implementing improvements from MAML++ but we still found high variance across different random seed initializations. Therefore, using MAML on NPs brings about a potential trade-off: MAML may improve the performance of NPs by improving adaptation to meta-test functions by taking a single step of gradient descent, but MAML may also worsen performance by hindering learning during meta-training due to training instabilities.

Reviewer comment (Presentation):
The reference section is incomplete and sometimes incorrect. For instance, the "Attention is all you need" paper is from 2017 and not from 2023 and misses the conference information.

Author response:
We thank the reviewer for catching this error! We have investigating the source and it appears the mistaken citations had been downloaded from arXiv. For example, the button "Export BibTeX citation" on the arXiv page of "Attention is all you need" ( https://arxiv.org/abs/1706.03762 ) yields a citation without the conference information and with the year 2023. We were unaware arXiv worked in this way. We apologize for this error, which we have now amended.

Reviewer comment (Novelty):
The authors fail to cite relevant literature on graph neural processes, e.g. [3].

Author response:
We thank the reviewer for pointing out this paper, which we were not aware of. Even though the names of the two models are similar, this model is actually very different from ours, since it is designed to learn node signals within a graph, and ours is designed to learn graph signals across many graphs. Therefore, their model is not readily applicable to our task. In addition, their application (weather prediction) is very different from ours (molecular property prediction). Because of these reasons it does not seem appropriate to cite this paper. However, if after reading our explanation the reviewer believes otherwise, we would be happy to reconsider.

Reviewer comment (Writing):
Some clarifications of the math of the encoder structure or an illustrative figure would in the reviewer's opinion improve the quality of the manuscript. While background on NPs is explained in sufficient detail (both in the main manuscript as well as the appendix), the actual method is not described at all.

Author response:
We thank the reviewer for raising this point. We would like to explain that the architecture of the GNN used by our MG-NPs is based on a very popular architecture (Attentive FP [1]) whose paper we reference and can be checked by the readers. The adaptation procedure is a bit trickier to replicate than the architecture (especially the one using MAML), so we used more space on that section instead, given the length constraints. We explain the differences between Attentive FP and our GNN in Appendix C.3.2. In addition, if our article is accepted we will publish our code with the camera-ready version.

Reviewer comment (Experiments):
As far as I can tell, the authors do not compare themselves against recent methods such as in [1]. Is this true and if so is there any reason for that?

Author response:
We thank the reviewer for raising this point. ADKF-IFT [3] is now included as a baseline.

Reviewer comment (Experiments):
LNPs are generally harder to train then CNPs. Is the poor performance of LNPs due to this fact or how can it be explained? Is it because the authors seemed to have trained only for a fixed number of iterations and not until converge (see Appendix C3.7)?

Author response:
Initially, we trained CNPs and LNPs for the same amount of iterations in order to promote an equalitarian allocation of resources and a fair comparison. However, we believed the LNP did not perform so well partly because its regularization term promoted underfitting and made reaching a good optimum challenging. Therefore, we attempted trainining during twice the number of iterations (6000 epochs). However, after 3000 epochs the results did not change substantially. In addition, in order to facilitate reaching a good optimum, we trained two versions of each LNP: one where the regularization term of the ELBO-like objective was computed analytically, and one where the regularization term was approximated with a single Monte Carlo sample (Appendix C.3.7). Then, we chose the version which performed best. At this point, we felt that further optimizing the LNP preferentially would have made the comparison unfair towards other models. All in all, we spent considerably more time and effort optimizing the LNP than the CNP.

Reviewer comment (Experiments):
The authors propose that existing datasets are highly homogeneous across tasks, while in reality the task diversity should be considered. However, there lacks evidence in this paper. No empirical results of existing popular datasets (e.g., fs-mol[2], moleculenet[3]), nor comparing them with real-world cases are provided.

Author response:
We would like to explain that when we referred to lack of heterogeneity in current benchmarks, we were referring to the toy datasets used by the broader meta-learning community, not to FS-Mol and MoleculeNet. Examples of the types of datasets we were referring to are the 1D sine function used in the original MAML paper [1], MNIST in the original CNP and LNP publications [2,3], or samples from a Gaussian process model with slightly different hyperparameters (the latter is often used in the NP literature, see e.g.[4] for a recent example). These simple datasets may be useful to develop novel meta-learning models but they do not represent the degree of heterogeneity of scientific datasets. Therefore, models that do well on such datasets may fail catastrophically when faced with divergent meta-train and meta-test functions. We showed this was the case for MAML and the CNP in Figure 1, where even a small change in the 1D sine function led to loss of meta-generalization.  

We thank the reviewer for raising this point, which is a fair one. We made a decision to focus on benchmarking docking scores in this publication and leave bioactivity endpoints for a future, more-biologically-focused publication. There are three compelling reasons to use docking scores when benchmarking drug-discovery algorithms. First, some widely-used bioactivity datasets such as MoleculeNet may present their own flaws, e.g. they present intra-function inconsistencies due to pooling from different publications [5,6]. Second, docking scores are highly useful for drug discovery, which makes them interesting to predict with ML by their own right. Some recent high-impact publications that had ML-prediction of docking scores at their core are [7,8] . Finally, docking depends on the simulation of a 3D interaction between a protein and a molecular ligand. Therefore, a model that is able to learn docking scores must be sophisticated enough to capture the structural relation between the protein target and the molecule. In fact, some docking scores are very challenging to predict, e.g. in the original publication of the dataset we are using (DOCKSTRING), the best model's prediction on ESR2 was a mediocre 0.627, even after training on 230k examples [9].  For this reason, docking scores have been proposed as promising benchmarks for drug discovery [9,10,11].  

We used two different datasets of docking scores: the original DOCKSTRING dataset, and an augmented dataset where scores were transformed either linearly or non-linearly (by taking the minimum or the maximum between each score and the median of the score distribution) and were combined with the quantitative estimate of drug-likelihood (QED).

Reviewer comment (Evaluation):
Could you explain the results presented in Figure F.1, where the $R^2$ did not decrease as the percentage of training points sampled increase?

Author response:
**Could you explain the results presented in Figure F.1, where the did not decrease as the percentage of training points sampled increase?**

This is indeed an interesting point. It appears that the NPs' mean prediction $\mu_\theta(x)$ was more robust to increasing the percentage of context points during meta-training than the variance prediction $\sigma_\theta^2(x)$. As a result, the NLPD deteriorated much more (Figure 4) than the $R^2$ (Figure F.1). This seems reasonable: if the model overfits to the training datapoints and memorizes their $y$ values, the mean predictions will fixate at the $y$ values and the variance predictions will become ever smaller (since smaller variances with the correct mean will lead to higher predictive likelihoods). Therefore, the model will learn to always predict inadequately small variances, which leads to catastrophic NLPD at meta-test times.

It is worth noting, however, that even if the NLPD deteriorated faster than the $R^2$, the $R^2$ also suffered if the fraction of context points at each iteration during meta-training was too high. This is shown in Figure F1, where the $R^2$ value of test datapoints was always lower at 100% than at 0.8, 2, 4 or 8%, for every meta-train and every meta-test function.

Reviewer comment (Experiments):
Have you considered studying the impact of context/target set randomization on calibration of uncertainty estimates?

Author response:
**Have you considered studying the impact of context/target set randomization on calibration of uncertainty estimates?**

We thank the reviewer for making this suggestion. The NLPD in Figure 4 was partly a measure of calibration since poorly calibrated models have larger NLPD. Indeed, when 100% of points of a function were taken as contexts and targets (hence having no randomization of contexts and targets during meta-training) the NLPD increased significantly both for ftrain, dtrain points and for all meta-test points (ftest, dtrain and ftest, dtest). In addition, we have now added to the manuscript Figure F.2, which shows the correlation of predicted variance and prediction error in a MG-CNP sampling a small fraction of training points as contexts and targets at each iteration, as usual, and a MG-CNP taking 100% of points as contexts and targets. The latter displayed low correlation, indicating poor calibration.