Context-enriched molecule representations improve few-shot drug discovery

Abstract: A central task in computational drug discovery is to construct models from known active molecules to find further promising molecules for subsequent screening. However, typically only very few active molecules are known. Therefore, few-shot learning methods have the potential to improve the effectiveness of this critical phase of the drug discovery process. We introduce a new method for few-shot drug discovery. Its main idea is to enrich a molecule representation by knowledge about known context or reference molecules. Our novel concept for molecule representation enrichment is to associate molecules from both the support set and the query set with a large set of reference (context) molecules through a modern Hopfield network. Intuitively, this enrichment step is analogous to a human expert who would associate a given molecule with familiar molecules whose properties are known. The enrichment step reinforces and amplifies the covariance structure of the data and simultaneously removes spurious correlations arising from the decoration of molecules. We analyze our novel method on FS-Mol, which is the only established few-shot learning benchmark dataset for drug discovery. An ablation study shows that the enrichment step of our method is key to improving the predictive quality. In a domain shift experiment, our new method is more robust than other methods. On FS-Mol, our new method achieves a new state-of-the-art and outperforms all other few-shot methods.

Reviews and author responses:

 MHNfs is able to outperform all other methods included in Table 1 on the FS-Mol test set (*). The FS-Mol test set consists of 157 tasks and the performances on the test set can be found in Table 1, column “ALL”. 

[a3] really performs well for two sub-categories (oxidoreductases, hydrolases).  But this is not at all in any contradiction to what we wrote here (*). This primarily means that the MHNfs and [a3] focused on different things during training. Up-sampling oxidoreductases- or hydrolases-related datapoints during training could improve the performance for these categories.

If there happened a mistake during the rebuttal phase regarding [a1], we would like to apologize. Thanks for flagging that. We will look into that and correct it.

ProtoNet, GNN-MAML, PAR, Siamese Networks, IterRefLSTM, and ADKF-IFT are included in Table 1. We would not consider any of these methods a naive baseline.


 [Larger support set sizes]

We are currently preparing experiments. We will add the results to the paper.


[Pretrained and frozen GNN]

Yes, good point! Actually, we have already thought about this, but just prioritized it lower compared to other experiments. We agree; further exploration here might be interesting.


[Figure A5]

Indeed, the provided plot might suggest the aforementioned.. But using a t-sne embedding, it looks like there could be different attraction points. We will additionally add the t-sne embeddings. As this still is underexplored, we don’t see ourselves in an informed position yet from which we can give a certain answer, but will investigate this further.


 [1] Initial loss and difference between losses

Thank you for this question. We want to point out that for few-shot drug discovery the scale of the loss and the characteristics of the loss curves are quite different from those of standard many-shot computer vision tasks.

Within our training pipeline we start reporting the validation loss from epoch 0 on (\~1.1), but the training loss from epoch 1 on (\~0.22).  We will update the plot to make this more consistent.

The difference between training and validation loss decreases until the model starts to overfit on the training data. During training, the loss is computed with dropout while during evaluation dropout is turned off. This is  a reason for the gap between training and validation curve. Furthermore, in few-shot settings this gap usually tends to be higher compared with standard multi-task settings, in which training and validation split include the same tasks.

Due to usually high noise in the labels, which is a strong difference to labels of computer vision tasks, the sigmoid output often tends to stay in an area around 0.5. Nevertheless, models can learn to rank molecules correctly according to their activity, which is exactly the model’s property which is needed for screening. Therefore, metrics which take a correct ranking into account - like the AUC or the ΔAUC-PR - are usefully meaningful and often preferred.

[2] 

We are glad the reviewer likes table A9 and has no concerns about the validity of our method.
The reviewer is right,  it is possible to explain our modules with queries, keys and values. Again, the reviewer is right explaining that - for our context module - the associative memory is the basis for the keys and values. Nevertheless we think - as the idea of an associative memory is crucial for our approach - describing our approach by using Modern Hopfield Networks is most intuitive.

[Originality] Our main contribution is to provide context as additional information to the model. Indeed we do this by using existing techniques (MHN, attention) in a new way, but we envision that follow-up works will find different ways to bring in context.Our method strongly differs from previous few-shot methods: In contrast to previous methods, which have only the query and support set as inputs, our method has three inputs: the context set, the query and the support set.

[Significance] Drug discovery projects usually start around just a few measurements. This is why few-shot methods are needed for early-stages. Any progress at few-shot drug discovery methods directly translates  to more efficient wet-lab experiments. Therefore, the real-world impact of improvements in few-shot drug discovery is large.

We wish to point out two more facts that underpin the significance of our work: a) we report the first successful domain shift experiment in few-shot drug discovery and b) multi-million dollar efforts of pharmaceutical companies have recently been aimed at improving representations of drug discovery tasks (see paragraphs below).

We point the reviewer to (Altae-Tran, 2017) and especially to Table 4. Already published in 2017, this paper can be seen as a milestone for few-shot drug discovery as the IterRefLSTM really showed impressive performance in few-shot settings. However, in a transfer experiment the authors showed that the IterRefLSTM does not achieve any predictive power on the transfer dataset.This domain shift problem is one of the next issues that few-shot learning methods in drug discovery have to solve. For the first time, we demonstrated (see Table A8) that MHNfs performed well under domain shift and that the associative memory was a crucial part to preserve predictive power.

Recently, there has been a multi-million dollar effort by 10 pharmaceutical companies to improve task representations  by federated learning (Humbeck, 2021). In our work, we improve task representations by our few-shot learning technique. Thus,we believe that our work is highly significant for both drug discovery and machine learning.

 [3] The reviewer is right that the context module might suppress insignificant features of a molecule. 
However, there must be a misunderstanding on what information is visualized because the figure neither shows the effect of multiple non-linear NN transformations nor the final embeddings, but  the effect of the retrieval step (Eq. (6)) and, thus, the context module. The figure depicts several initial patterns (light purple) $\boldsymbol \xi$ which only are used to determine the weights how the stored patterns (blue) in $\boldsymbol M$ are averaged to obtain the retrieved patterns (light purple). These patterns all live in the same space, which is why we visualize them together by their first two principal components in this figure. We apologize that we have not explained this well enough in the caption and will improve on the description. 

[4]
Thanks, we corrected it.


(Altae-Tran, 2017) Low Data Drug Discovery with One-Shot Learning, ACS Cent. Sci., 2017

(Humbeck, 2021) Don’t overweight weights: Evaluation of weighting strategies for multi-task bioactivity classification models, Molecules, 2021


 Hi, I have read all the reviewing comments and the authors' reply. 
Most of my concerns are cleared. For experimental analysis, I still think the current version is not satisfactory.  

Reviewer ZhVh wrote that `` I agree FS-Mol is newer and less established than some other datasets, although I wouldn't go as far as saying that the only available results there are naive baselines``. I also think FS-Mol is a valuable benchmark dataset for few-shot molecular property prediction. But I insist that only comparing with naive baselines provided by FS-Mol. As shown in the updated draft, 
the proposed MHNfs cannot beat [a3] in Table 1. 

Besides, please note that the claim``[a1] had already been included as 'GNN-MAML', but unfortunately we cited the original MAML paper instead of [a1]`` is wrong. 
In the original paper of FS-Mol, the authors wrote that ``Concretely, we consider Model-Agnostic Meta-Learning [18] procedure on top of a GNN operating on the molecular graph as a representative optimization-based approach and refer to it as GNN- MAML``. While Meta-MGNN [a1] further constructs attribute based self-supervised modules and incorporates self-attentive task weights during the episodic training of MAML. In short, Meta-MGNN is not equal to ``a GNN encoder+ MAML``. 


[a3] Meta-learning feature representations for adaptive Gaussian processes via implicit differentiation.




 Thank you for your response, and for providing lots of additional clarifications. Find a few follow-up comments below:

> Due to the computation time, we could not do this for the rebuttal version, but we provide a proxy in which we use the model trained on support set size 16 to extrapolate to other support set sizes (Appendix Figure A4).

Fair enough. However, I would encourage you to add the results for larger support set sizes eventually (either upon acceptance, or in the next version of the paper otherwise).

> End-to-end training with GNNs is feasible, but usually leads to inferior models.

Yes, this is understandable in a low-data setting. I wonder how well a pretrained (and frozen) GNN would do, but such investigation is orthogonal to the method developed here.

> Yes, we have some insights in what the context module is doing. The context module is reinforcing the correlation structure with similar molecules (see Appendix Section A7 and Figure A5).

Interesting, thanks for the pointer to the right Appendix. I also note that all examples in Figure A5 seem to be moving towards roughly the same "center", as if the retrieved embeddings (before combining with the initial embeddings) were all mostly the same? Or am I reading this incorrectly?

---

Finally, I also went through the other reviews, out of which the one from Reviewer LC2d is the most negative one (and hence farthest from mine); however, I find the authors' response to that review satisfactory. I agree FS-Mol is newer and less established than some other datasets, although I wouldn't go as far as saying that the only available results there are naive baselines. I would note that the setting in FS-Mol (i.e. a very large number of diverse tasks) is a bit closer to real-life property prediction than, say, experiments based on a small number of targets in Tox21. However, benchmarks mentioned by Reviewer LC2d are also worthwhile, just address a somewhat different setting. All in all, I'm so far not convinced I should lower my score, and so I maintain it as "Accept".

 [1] Minor question, in Figure A2, I was wondering why the initial validation loss (epoch0) is much greater than that in training (~1.1 vs ~0.22 CE loss)? And the loss difference remains large as training continues. The predicted likelihood associated with the validation set looks a bit concerning given this CE loss.

[2] Table A9 looks good. I don't have much concern about the validity of the proposed method, my main concerns are about the originality of the network architecture (again, everything can be explained using the key, value, query example, where the associative memories are key and value embeddings) and how significant the performance metric improvements are to the community.

[3] It's interesting to see in Figure A5 that the context-enriched repr. is like a scaled version of the initial embedding. My understanding is that some insignificant feature dimensions are ignored/compressed during learning. For PCA, the scale itself does not matter much (the eigen vectors can be arbitrarily scaled), thus, after all these non-linear NN transformations, the final embedding space is somewhat similar to the original one. Would you please help me understand the what is going on here?

[4] Very minor, in the new abstract, "state-of-the-art", missing hyphen.

 We thank the reviewers for their highly insightful, positive and constructive reviews, which helped a lot to improve the manuscript!  

All reviewers found the problem that we addressed, few-shot learning in drug discovery, relevant. Reviewer 1 and 2 found our proposed architecture and approach novel and appreciated our contribution. All reviewers found our approach clearly described. 

Reviewer 3 had four main points of criticism which we could address in this rebuttal version. 
First, missing comparisons against i) MAML-based GNN (GNN-MAML) and ii) property-aware relation (PAR) networks. However, we had already compared against i) which was not obvious because of the different naming of the method. We added comparisons with ii) in the new version of the manuscript.
Second, the missing related work section. We have added a “Related work” section, in which also the references suggested by the reviewer are added. 
Third, the motivation for using modern Hopfield networks. We now explain that modern Hopfield networks reinforce the covariance structure in the data via retrievals. 
Fourth, writing and typos: we improved writing and removed typos. Thanks for pointing this out!

Reviewer 1 and 2 commented on the connection between modern Hopfield networks and Transformer/attention. The reviewers are right: The update rule of MHNs is analogous to self-attention or cross-attention in Transformers. However, MHNs have two essential differences to Transformer attention: a) The Transformer attention mechanism combines activations of previous layers or layers of a second network based on the current input. In contrast, MHNs store embeddings that are not based on the current input but on an external reference set. MHNs can store external context, which is a mechanism that is not covered by Transformer attention. b) MHNs reinforce the data covariance structure which is a property of associative memories. In classical Hopfield networks the memory matrix constructed by Hebb learning is a covariance matrix, where retrieval is reinforcing the query covariance structure. Thus, there is a close connection between associative memories and the enhancement of the data covariance structure. 

Reviewer 1 and 2 mentioned our choice of molecule encoder, which is a descriptor-based fully-connected network. While we state that any reasonable molecule encoder can be used, we found those networks to perform better than GNN-based networks. There are several facts that support this finding: a) The same finding has been reported in FS-Mol (Stanley, 2021) where prototypical networks based on ECFP descriptors performed best. b) Descriptor-based fully-connected networks are usually better or on-par with GNNs for activity prediction (Jiang, 2021) and might be better suited for few-shot learning. GNN-based methods have to be combined with descriptors to perform well (Yang, 2019). c) the successful classic similarity search also relies on descriptors (Cereto-Massague, 2015). During development, we also tested GNN encoders, such as GIN, but their performance was always worse than with descriptor-based architectures. Nevertheless, we will add a comparison with an architecture based on a GNN molecule encoder in the final version. 

We appeal to the reviewers to re-consider their assessments and scores in light of our additional results and our improved version of the paper. 

### References
(Stanley, 2021)   Stanley, M., Bronskill, J. F., Maziarz, K., Misztela, H., Lanini, J., Segler, M., ... & Brockschmidt, M. (2021, August). Fs-mol: A few-shot learning dataset of molecules. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2).

(Jiang, 2021)   Jiang, D., Wu, Z., Hsieh, C. Y., Chen, G., Liao, B., Wang, Z., ... & Hou, T. (2021). Could graph neural networks learn better molecular representation for drug discovery? A comparison study of descriptor-based and graph-based models. Journal of cheminformatics, 13(1), 1-23.

(Yang, 2019)    Yang, K., Swanson, K., Jin, W., Coley, C., Eiden, P., Gao, H., ... & Barzilay, R. (2019). Analyzing learned molecular representations for property prediction. Journal of chemical information and modeling, 59(8), 3370-3388.

(Cereto-Massague, 2015)   Cereto-Massagué, A., Ojeda, M. J., Valls, C., Mulero, M., Garcia-Vallvé, S., and Pujadas, G. (2015).392 Molecular fingerprint similarity search in virtual screening. Methods, 71:58–63.

 # Strengths
Thanks for pointing out Figure 1 and the reproduction details given in the appendix. 



# Weaknesses
## Motivation for MHNs

We picked this particular type of external memories because of two reasons: a) other memories would not work because they do not reinforce the covariance structure (see Appendix Section A6) or they are not differentiable and b) the connection with transformer attention mechanisms, which are highly successful across all ML areas. See also answer to key question 1 (KQ1) below. 


## Adding related Works 

Indeed, we aimed to include all related work in the introduction section and we apologize that this has not met the requirements of the reviewer. We now have a separate Related Works section, in which we relate and cite the works [a1, a2]. Of note, [a1] had already been included in our Introduction section. 


## Experiments with other compared methods

We thank the reviewer for pointing this out! We follow the suggestions and compare [a1,a2] on FS-Mol. [a1] had already been included as “GNN-MAML”, but unfortunately we cited the original MAML paper instead of [a1]. We corrected this. We have also added results of the PAR [a2] approach. For PAR [a2] we performed hyperparameter search and could strongly improve performance over the hyperparameter settings suggested by the authors (see Appendix Section A1.7 for the selected hyperparameters). The results are added in Table 1. 


## Writing and typos.
We proofread the manuscript and corrected typos and grammar – thanks! In Table 1, we now mark the best method per column in bold and not just the overall best method. For the final version, we will have a professional editor proofread the manuscript.

 # Questions:
## [KQ1] Justification of MHNs
We considered MHNs because these associative memories are able to strengthen the covariance structure of the data with the retrieval step (see Appendix A6). Note that we can express the context-enriched representation of the molecules in terms of the weighted covariance matrix times the representation $\boldsymbol x' = (\phantom{d}\boldsymbol {c̅} + \mathrm{Cov}(\boldsymbol C, \boldsymbol x)^T \boldsymbol x\phantom{d})$ (see Appendix A7). The patterns that govern the covariance are weighted by their similarity to the given molecule embedding $\boldsymbol x$. The multiplication of the covariance matrix with $\boldsymbol x$ means that the directions of $\boldsymbol x$ that match with directions in the data are strengthened. 
Another reason is that the particular connection with the successful Transformer-like attention mechanisms (Ramsauer, 2021) make MHNs a prime candidate for use in a Deep Learning architecture. 
We also provide a perspective on architectures with external memories in Related Work including [b1] and [b2]. 
## [KQ2] Comparisons with additional methods
We now provide discussion and comparisons with [a1] and [a2]. 

## Limitations
We now mention that MHNfs increase the computational costs since the embeddings for the Hopfield memory must be precomputed. Fortunately, the memory can be reused for different mini-batches. We also mention that large context sets impose a large memory footprint which may exceed the capacity of GPUs. However, it is possible to subsample the context set to fit the GPU memory.




# References
[a1] Few-Shot Graph Learning for Molecular Property Prediction, WWW, 2021  
[a2] Property-Aware Relation Networks for Few-Shot Molecular Property Prediction, NeurIPS, 2021  
[a3] Low Data Drug Discovery with One-Shot Learning, ACS Cent. Sci., 2017  
[b1] Meta-learning with memory-augmented neural networks, ICML, 2017  
[b2] Few-shot visual learning with contextual memory and fine-grained calibration, IJCAI, 2020

 #  Strengths
Thanks, we appreciate the reviewers summary of the strengths S1, S2, and S3 of our method!
# Weaknesses
## W1: Other support set sizes
(W1) Indeed, we focus on the setting with a support set size of 16 for two reason: First, this was also the main setting of the original FS-Mol paper (Stanley, 2021, Table 2) and, second, this setting is close to the usual 5-shot and 10-shot (per class) setting of the computer vision papers. However, for the final version we will add a plot similar to Figure 2a in Stanley (2021), as suggested by the reviewer. Due to the computation time, we could not do this for the rebuttal version, but we provide a proxy in which we use the model trained on support set size 16 to extrapolate to other support set sizes (Appendix Figure A4). We note that these values are not comparable to the values in Stanley (2021), because the model has not been trained on these larger support sets and the hyperparameters were also adjusted towards the support-set-size-16 setting. As typical for embedding-based methods, we expect our method to work particularly well for smaller support set sizes up to 32 or 64 (Triantafillou, 2019, Figure 2). 

## (W2) Questions
Thanks for reading our paper in depth. We improved as suggested through the Questions (see below).


# Other comments
## O1: Missing recent result
Thanks, we included the comparisons with [2] in Table 1 and we will add performance metrics for larger support set sizes in the final version of the paper (see also W1). While [2] yields strong performances for larger support set sizes, embedding-based methods such as our MHNfs usually perform well for small support set sizes (Triantafillu, 2019, Figure 2b; Chen, 2019; Tian, 2020). 

## O2: Use of MHN terminology and relation to attention
The reviewer is right: MHNs are analogous to self-attention or cross-attention in Transformers. However, MHNs have two essential differences to Transformer attention:
a) The Transformer attention mechanism combines activations of previous layers or layers of a second network based on the current input. In contrast, MHNs store embeddings that are not based on the current input but on an external reference set. MHNs can store external context, which is a mechanism that is not covered by Transformer attention.
b) MHNs reinforce the data covariance structure which is a property of associative memories. In classical Hopfield networks the memory matrix constructed by Hebb learning is a covariance matrix, where retrieval is reinforcing the query covariance structure. Thus, there is a close connection between associative memories and the enhancement of the data covariance structure. 
We added some clarifications that make the synonymous cases to transformer attention mechanisms clear and derive the connection with the (weighted) data covariance in Appendix Section A7. 


# Nitpicks:
Thanks, corrected. 

 # Questions
## Choice of molecule encoder

During development of our method, we tried a more complex GNN-based encoder, but it always performed worse than descriptor-based fully-connected neural networks. This was also one of the findings in the FS-Mol paper (Stanley, 2021), where the descriptor-based methods (prototypical networks) outperformed the other methods. End-to-end training with GNNs is feasible, but usually leads to inferior models. We will add a comparison with MHNfs based on a recent GNN in the final version of the paper. 


## Q2: Balancing strategy
Thanks for your interest in this balancing strategy (see also reviewer 1). Indeed, we can provide an intuition. What we are suggesting here is a variance-conserving pooling strategy within each class (actives, inactives), and thus a trade-off between mean-pooling and sum-pooling. Assuming independence of the variables over which is summed, sum-pooling would increase the variance by a factor of $N_A$ (or $N_I$) and mean pooling would decrease the variance by $\frac{1}{N_A}$, whereas the factor $\frac{1}{\sqrt{N_A}}$ maintains variance. Very large or small values entering the sigmoid in Eq (10) would lead to problematic learning behavior, which is avoided by our balancing strategy.  


## Q3: input features

We apologize for the unclear statements about the input features. The molecule encoder of MHNfs uses a concatenation of ECFP and RDKit-based descriptors. These are exactly the descriptors which were provided in the FS-Mol paper (Stanley, 2021). We clarified the information about the input features in the main paper (line 246ff.) and added details in the appendix. 

## Q4: exploiting the frequent hitters model. 

During development, we have tried several ways to combine FrequentHitters (FH) models with task-specific models and different ways to combine the representations or prediction scores of the FH model, but found it non-trivial to increase predictive performance with such an approach. 
##Q5: intuition about context module
Yes, we have some insights in what the context module is doing. The context module is reinforcing the correlation structure with similar molecules (see Appendix Section A7 and Figure A5). The attention mechanism usually does not focus on a specific small subset of context molecules, but at a larger set of molecules over which an average is built. We have not investigated the popularity of active molecules or scaffolds, but we will add some visualized examples in the final version. 

 ## Limitations
We will investigate the relation between predictive performance and support set size in more detail in the final version. 



# References 

(Triantafillou, 2019)    Triantafillou, E., Zhu, T., Dumoulin, V., Lamblin, P., Evci, U., Xu, K., ... & Larochelle, H. (2019). Meta-dataset: A dataset of datasets for learning to learn from few examples. arXiv preprint arXiv:1903.03096.

(Tian, 2020)  Tian, Y., Wang, Y., Krishnan, D., Tenenbaum, J. B., & Isola, P. (2020, August). Rethinking few-shot image classification: a good embedding is all you need?. In European Conference on Computer Vision (pp. 266-282). Springer, Cham.

(Chen, 2019)  Chen, W. Y., Liu, Y. C., Kira, Z., Wang, Y. C. F., & Huang, J. B. (2019). A closer look at few-shot classification. arXiv preprint arXiv:1904.04232.

(Stanley, 2021)   Stanley, M., Bronskill, J. F., Maziarz, K., Misztela, H., Lanini, J., Segler, M., ... & Brockschmidt, M. (2021, August). Fs-mol: A few-shot learning dataset of molecules. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2).

 # Strengths
We thank the reviewer for pointing out the significance and novelty of our method and appreciate their insightful review!

# Weakness 

## [originality] Analogy with MHNs
The reviewer is right: MHNs are analogous to self-attention or cross-attention in Transformers. However, MHNs have two essential differences to Transformer attention:
 a) The Transformer attention mechanism combines activations of previous layers or layers of a second network based on the current input. In contrast, MHNs store embeddings that are not based on the current input but on an external reference set. MHNs can store external context, which is a mechanism that is not covered by Transformer attention.
b) MHNs reinforce the data covariance structure which is a property of associative memories. In classical Hopfield networks the memory matrix constructed by Hebb learning is a covariance matrix, where retrieval is reinforcing the query covariance structure. Thus, there is a close connection between associative memories and the enhancement of the data covariance structure. 
We, too, see our novel contribution in the context module, that separately handles molecules of the support and context set. We will put more emphasis on this, as the reviewer suggests.

## [quality] Insights into enrichment process
We agree with the reviewer that we should investigate the enrichment process more. We added an example in Appendix Section A5 and a theoretical consideration in Appendix Section A6. 

However, we disagree with the reviewer on the statement that for activity prediction there are better molecular encoders than descriptor-based fully-connected networks. While graph neural networks are state-of-the-art for physical modeling (e.g. the QM9 benchmark) and for encoding macromolecules, descriptor-based fully-connected networks are better or on par with GNNs for activity prediction of small molecules (Jiang, 2020; Mayr, 2018). The widely-used and proficient chemprop model (Yang, 2019) also strongly relies on chemical descriptors and only their combination  with graph-based representations yields high predictive performance. The initial FS-Mol study (Stanley, 2021) has also shown that descriptor-based deep neural networks are ahead of GNNs for few-shot learning. During development, we also used GNN encoders, but their performance was always worse than with descriptor-based architectures. Nevertheless, we will add a comparison with an architecture based on a GNN in the final version. 
## [significance]

We now provide a statement on the significance of such AUC values in the context of drug discovery (Section 6, line 336ff.). Briefly, for models with an AUC >0.7 the subsequent wet-lab screening is expected to bring an enrichment of active molecules, which is a strong improvement for early-stage drug discovery (Truchon, 2007; Simm, 2018). If a method can yield an AUC of 0.7 with only few measured molecules, this would be an improvement for real drug discovery problems. 

 # Questions 
## Context set in a real drug discovery problem
For a real drug discovery problem, we suggest choosing the context set exactly as we did during training: use ~10,000 randomly selected molecules from a large library, e.g. PubChem. This would ensure a similar predictive performance as the covariance can be similarly well estimated (see Appendix Section A7). We believe that there could be better ideas to choose the context set and we actively encourage the community to test those ideas.


## Geom drug as context set
In fact, this is what we are already doing: for the FS-Mol predictions, we use a general molecular database, i.e. ChEMBL, from which we pick the context set. The context set should cover a large part of the chemical space. Nevertheless, we will also provide experiments, in which we select the context set from GEOM-drugs (see Appendix Section A5). Because GEOM-drugs contains drug-like molecules, similar to FS-Mol, the predictive performance remains stable. 


## Amplification of co-occurences and covariances
We now provide a theoretical connection between the retrieved patterns and the covariance in terms of weighted covariance matrices (see Appendix Section A7). We also added a figure with a particular example that shows how molecules with special characteristics are put in relation to other molecule embeddings (Figure A5) 
## Inaccurate statement about standard supervised ML
The reviewer is right. We reformulated this statement.
## Number of training epochs and potential overfitting
There must be a misunderstanding on the number epochs to train the MHNfs models. The frequent hitters model is trained for a maximum of 20 epochs, but the MHNfs for 350 epochs and we use early-stopping based on validation metric We now write this clearer in the Appendix. An overfitting effect would lead to decreased performance on the test tasks, which we report in Table 1.  Additionally, we now also provide learning curves in the Appendix (Figure A2). 

# Minor concerns
## Typos and formulations in lines 79, 92, and 122
Thanks, corrected. 
## Balancing strategy 
Thanks for your interest in this balancing strategy (see also reviewer 2). What we are suggesting here is a variance-conserving pooling strategy within each class (actives, inactives), and thus a trade-off between mean-pooling and sum-pooling. Assuming independence of the variables, sum-pooling would increase the variance by a factor of $N_A$ or $N_I$ and mean pooling would decrease the variance by $\frac{1}{N_A}$, whereas the factor $\frac{1}{\sqrt{N_A}}$ maintains variance. Very large or small values entering the sigmoid in Eq. (10) would lead to problematic learning behavior, which is avoided by our balancing strategy.  
Indeed, what the reviewer suggests is basically used in our cross-attention module (CAM) which also contains the activity label. Using the kernel to guide an attention mechanism is, for example, used by IterRefLSTM (Eq.2, Altae-Tran, 2017). Therefore, we believe that the strategy that the reviewer suggests is already used in the best architectures. 


## Use of word “significantly”
Thanks! The “significantly” relates to the statistical test that we performed between the first ranked method and the other methods. The p-values of this Wilcoxon test are provided in brackets. 

 # Limitations
## Macromolecules
This is a very insightful comment by the reviewer that conceptually our approach can also be used for macromolecules, thanks! Indeed, the concept that we suggest can be used for larger molecules. Our comment related more to the limitation of our particular implementation.  We reformulated the statement.

## Different molecular representations
We emphasize that our model learns representations in a data-driven manner. From the many possible initial representations of small molecules (SMILES, descriptors, fingerprints, INCHI, molecular graph, etc) we picked fingerprints, but our architecture still learns a high-level representation of a molecule. We will add comparisons with a GNN based molecule encoder in the final version of the manuscript. 


#  References
(Stanley, 2021)   Stanley, M., Bronskill, J. F., Maziarz, K., Misztela, H., Lanini, J., Segler, M., ... & Brockschmidt, M. (2021, August). Fs-mol: A few-shot learning dataset of molecules. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2).

(Jiang, 2021)   Jiang, D., Wu, Z., Hsieh, C. Y., Chen, G., Liao, B., Wang, Z., ... & Hou, T. (2021). Could graph neural networks learn better molecular representation for drug discovery? A comparison study of descriptor-based and graph-based models. Journal of cheminformatics, 13(1), 1-23.

(Yang, 2019)    Yang, K., Swanson, K., Jin, W., Coley, C., Eiden, P., Gao, H., ... & Barzilay, R. (2019). Analyzing learned molecular representations for property prediction. Journal of chemical information and modeling, 59(8), 3370-3388.

(Mayr, 2018)  Mayr, A., Klambauer, G., Unterthiner, T., Steijaert, M., Wegner, J. K., Ceulemans, H., ... & Hochreiter, S. (2018). Large-scale comparison of machine learning methods for drug target prediction on ChEMBL. Chemical science, 9(24), 5441-5451.

(Simm, 2018)  Simm, J., Klambauer, G., Arany, A., Steijaert, M., Wegner, J. K., Gustin, E., ... & Ceulemans, H. (2018). Repurposing high-throughput image assays enables biological activity prediction for drug discovery. Cell chemical biology, 25(5), 611-618.


(Truchon, 2007) Truchon, J. F., & Bayly, C. I. (2007). Evaluating virtual screening methods: good and bad metrics for the “early recognition” problem. Journal of chemical information and modeling, 47(2), 488-508.

 The authors proposed a few-shot learning model to tackle the lack-of-data problem in drug discovery, where they demonstrated the proposed MHNfs model on benchmark dataset FS-Mol and the Tox21 dataset, reportedly outperforming existing methods. I find this paper clearly written with key concepts well explained. The central idea behind the proposed method, which is to enrich molecular representations through querying from a massive reference space (i.e., the context) is reasonable, since it is what a human expert would do in this scenario. Technical details are concisely discussed in the main text as well as the supplementary information, where non-domain experts should be able to understand and appreciate their work. **Strengths**

[*originality*] The proposed model architecture is novel for the few-shot learning task, compared to existing models listed in Table1.

[*quality*] The paper is well written with a clear story line and very few typos or grammar issues.

[*clarity*] The key idea is clearly explained and model architecture is accurately described.

[*significance*] The proposed method aims at solving a severe problem commonly encountered in drug discovery. The model seems to outperform all other existing models on the two prediction tasks demonstrated in the paper.

$\newline$

**Weakness**

[*originality*] The attention mechanism has been heavily exploited in the MHNfs model. Although the authors explained the difference between the proposed CM, CAM modules, they eventually both break down to the attention mechanism, which relies on computing the cosine similarity of linearly transformed key, query, and value embeddings. To that end, the novelty of the network architecture is not that significant in my opinion, regardless of the jargons and module names used to wrap these operations. That being said, the authors should probably spend more effort discussing and analyzing the significance of explicitly separating the support and context set, in comparison to other traditional methods. 

[*quality*] It would be ideal to have more discussions or insights on how the enrichment process benefits the predictive tasks (e.g., what did the model learn that made it better?), not just presenting tables with superior metrics. In addition, the authors used hand-crafted features (fingerprints and physical properties) for molecular embedding, I think it is widely accepted that there exists better molecular encoders, e.g., graph representations.

[*clarity*] No major concerns.

[*significance*] Regarding the quantitative model evaluation results, it would help the readers better appreciate the significance of this work by providing some sense of "a good model". Outperforming existing models is, of course, a significant contribution -- but readers who are not familiar with the specific predictive tasks may find it hard to judge whether the proposed method actually solved the proposed problem. For instance, in Table2, is AUC 0.679 good enough for the Tox21 task, or just a marginal performance gain compared to other methods? **Major concerns**

- For a real drug discovery problem, how should we choose the context dataset? The authors used part of the training set for the FS-Mol and Tox21 datasets, where the molecular data came from the same database, i.e., they are similar/related in some sense. However, when we only have access to, say, 16 known drug molecules for a particular disease, what would be the reasonable way to choose the context set? The authors should add more discussions about this problem. 

- A follow-up to the previous concern, is it possible to use a general molecular database, e.g., GEOM-drugs, as the context set, and use the same support and query set for FS-Mol and Tox21 predictions? If that works, then the proposed context-enrichment method really works.

- In the main text, line 82-84, "Those retrieved representations have amplified co-occurrences and covariance
83 structures, while peculiarities and spurious co-occurrences of the query and support set molecules are
84 averaged out." Is it possible to illustrate this with an example, or somehow support the claim in your case?

- In the main text, line 110-113, the authors discussed the limitations of standard supervised ML on data-scarce problems. I understand the point here, but I'm not sure if the statements here are accurate. Not all ML models will overfit the dataset, especially if somehow the model learns a proper similarity kernel function, which is exactly what the MHNfs model is doing.

- As far as I know attention-based models are data-hungry, and takes time and computation to extract meaningful hidden patterns. However, given the dimension of the hidden space (thousands) used in this work, the number of data points (thousands) in the context set, and the number of epochs (20), I was wondering if the model has been properly trained. Can you please show some learning curvs, if possible? I am also kind of curious to see if there is overfitting problem in this case.


**Minor concerns**
- In the main text, line 79, "as an associative memory".

- In the main text, line 92, "and a domain shift experiment".

- In the main text, line 122, "three consecutive modules that feed into each other". It is a bit confusing to me since the architecture is a simple feed-forward model, I don't see the modules feeding into *each other*.

- I am curious about the balancing strategy described in Equation 10. Naïvely, the kernel already provides a correlation between m and x, should we just attend to those with high correlation, and make some kind of vote from the highly correlated support set molecules, instead of using a weight based on the positive/negative ratio in the support set? Those molecules with small k(,) values may not contribute to the final voting. Any thoughts?

- In the main text, line 320, the authors claimed "significantly outperformed". From the numerical values, I get the "outperformed" part, but please justify the use of "significantly" in this case.



 - The authors wrote in line 336 that "the model cannot be used for molecules such as RNA, DNA, or proteins". I do not agree with this statement. Given the fact that the proposed method essentially captures the similarity through an attention-based kernel function by referencing from the context set, regardless of the underlying molecular representation, this method should be applicable to macromolecules as well. (E.g., I personally find the MSA blocks in AlphaFold2 similar to the CM module here) In fact, this few-shot learning method is potentially applicable to other data-scarce problems, as long as we could define a kernel function for similarity assessment.

- In conclusion, I think the key problem that this work tried to address is evaluating molecular similarity in some embedding space, and use that information to facilitate model predictions with few support data. I would love to see how different molecular representations perform in terms of similarity assessment, and whether hand-crafted features outperform representations learned in a data-driven manner.

 This paper proposes a context-enabled few shot learning method, which performs few shot classification by first looking at a group of context drug-like molecules, and then uses self-attention to distribute that information among the query and support set molecules. The authors show good empirical performance on FS-Mol, as well as ablations of the different parts of their model. Finally, the authors also present a FrequentHitters baseline, which performs very well despite ignoring the support set, yielding useful insight into why some of the FS-Mol baselines fail and how they could be improved.  === Strengths === 

- (S1): The paper considers an important (and still somewhat underexplored) problem of few-shot drug design, proposing a method based on an interesting mechanism of context lookup. The model design is clear and reasonable, and empirical performance is encouraging. 

- (S2): Apart from outlining their main method, the authors also discuss a FrequentHitters baseline, which ignores the support set and yet performs very well, which is intriguing, and could inform improvements to existing methods (see (Q4) for further thoughts on this). 

- (S3): The paper is well-written, with clear explanations and visualizations (in particular, Figure 1 is quite nice). The discussion is well-structured, most design decisions are clearly explained and ablated. The text flows smoothly, with very few minor typos or grammar errors. 

 

=== Weaknesses === 

- (W1): The paper only reports results with support set size 16, ignoring {32, 64, 128, 256} which are also commonly evaluated (e.g. in the original FS-Mol paper [1], and some other papers reporting results on that benchmark like [2]). It would be very interesting to see a plot akin to Figure 2a in [1]. As a start, it would be interesting to see the result for at least one larger support set size e.g. 128 (at 256 some FS-Mol-specific quirks kick in, and this leads to some artifacts in the performance of some of the methods). 

- (W2): Despite most aspects of the work being well-explained, some things are still not clear to me (see the "Questions" section). Further clarifying these would help polish the paper further. In particular, addressing (Q5) would be helpful for understanding why MHNfs works well.

 

=== Summary === 

Overall, this paper proposes a clear and reasonable approach, that (despite its simplicity) works very well on an important benchmark. The authors also propose an interesting baseline, which performs surprisingly well, leading to some useful insights. Finally, the paper is pleasant to read and the results seem reproducible. I think the paper is in good shape and a useful contribution to NeurIPS, hence I vote for acceptance. 

 

=== Other comments === 

Here I include various other comments, which had only a small impact on my score. 

- (O1): Table 1 misses a (very recent) result of [2], which (looking at their Figure 1a) gets a very similar context-size-16 result to MHNfs. I'm curious as to what the comparison is for larger support set sizes. 

- (O2): The authors say they use MHNs, as opposed to attention, attributing this to MHNs being more general. As far as I understand (but please correct me if I'm wrong), the main part where MHNs are practically more general than cross attention is that the formulation of MHNs considers the possibility of sequentially updating the "attention weights" several times before returning the output (this is different from having several sets of parallel attention weights, i.e. heads). However, it seems to me many papers using MHNs don't make use of this sequential update, which reduces to standard cross attention. In particular, while I agree with the authors that Equation (9) is a special case of Equation (7), the latter looks like standard cross attention to me (i.e. it's obviously not self-attention, but it looks like the kind of attention a decoder would do over the outputs of an encoder in a translation model). In general, I'm not too strongly opposed to calling cross-attention MHNs to make a connection to some theoretical results or leave the possibility of doing the iterative update of attention weights, but if the flavour of MHNs used is synonymous to attention (which I understand is the case for this work), I think this should be made more clear. 

 

=== Nitpicks === 

Here I include some final nitpicks, which did not affect my score; they are here just to help improve the paper.

- Line 92: "an domain shift experiment" - "an" -> "a" 

- Line 108: "predicts y for a x" - "a" -> "an"; "the model that generalizes" - I'd say "a model" 

- Missing space in lines 149 and 237

- Line 241: "wich" -> "which" 

- Line 254: "dataset on related task and then adapting" -> "dataset of related tasks and then adapt" 

- Line 262: "support support" 

- Line 324: the last sentence of the paragraph reads a bit off, might be useful to rephrase it (even though it may technically be grammatically correct) 

- Line 346: "material science dataset" -> "material science datasets" 

- Line 363: "This, might" - drop the comma 

- Line 647 (appendix): "fingerprints Morgan" -> "Morgan fingerprints" 

 

=== References === 

[1] Stanley et al, "FS-Mol: A Few-Shot Learning Dataset of Molecules" 

[2] Chen et al, "Meta-learning Adaptive Deep Kernel Gaussian Processes for Molecular Property Prediction"  - (Q1): It's interesting that MHNfs works so well with a very simple encoder (just a single linear layer on fingerprints/descriptors). Did the authors try using a more complex (e.g. GNN-based) encoder? If training this end-to-end is not feasible due to overfitting, one approach could be to replace the rdkit-based input to the linear encoder with a frozen pretrained feature extractor, such as a GNN or a graph transformer. 

- (Q2): Is there any intuition behind the label balancing strategy in Section 3.4? 

- (Q3): It is not clear to me from Section 4.1 what the input to the linear encoder is: whether it's a fingerprint, descriptors from rdkit, or a concatenation of both. Appendix A.1.1 suggests a concatenation (although calling it "an additive combination"?), but that part of the paper is only about the FrequentHitters baseline. Generally, I think the input to the encoder should be clear from reading the main paper. 

- (Q4): The fact that some of the methods developed for FS-Mol are outperformed by FrequentHitters seems to point to one issue with how some of these methods are designed: there doesn't seem to be a direct way for similarity-based models to learn a "frequent hitter bias" that only depends on the query molecule. I wonder if combining a frequent hitter branch (which would directly predict the output from the embedding of the query) with some of the methods would further improve them. In particular, given that MHNfs is similarity-based, could such an addition improve results further? It's not obvious that it would, as MHNfs already performs better than FrequentHitters, but maybe FH succeeds in some of the places where MHNfs fails. 

- (Q5): Do the authors have any intuition as to what the context module is doing? What is it attending to? What kind of information can be extracted from looking at a seemingly random collection of drug-like molecules? Could it be focusing on some implicit signal, e.g. that active molecules may be more "popular" i.e. their variations may be also present in the context set as they were tested for activity against other targets? The limitations of the method are discussed at length, and I have no concerns about that part of the paper. Perhaps one thing that would contribute to improving the understanding of the limitations is addressing (W1), as it would show whether MHNfs only works for very small support set sizes, or it can be considered a go-to method across a larger range. 

 Considering few-shot drug discovery, the paper proposes to enrich molecular representation by associating molecules from the support set and query set with additional context molecules. The proposed method consists of (i) context module designed upon Modern Hopfield  Network, (ii) cross attention module, and (iii) similarity module. 
The authors conduct experiments on FS-Mol dataset.

 Strength

+ Figure 1 is nicely plotted, which conveys the flow of the proposed MHNfs.
+ The authors provide concise reproduction details in appendix. 


Weakness

- Motivation: The authors did not provide convincing motivations or necessity for using Modern Hopfield Networks. 
- Related Works: The paper lacks a related work section. It seems that the authors plan to describe this part in the introduction section. However, the literature review is not enough. In the past, several few-shot drug discovery works appear [a1,a2]. But the authors neglect them all. 
- Experiments: Existing works [a1,a2,a3] use classic molecular property prediction datasets (i.e., Tox21, MUV, SIDER). While the authors only consider FS-Mol. FS-Mol is a new dataset published in dataset track of NeurIPS 2021, which naturally has only been tested by naive baselines. To show the proposed MHNfs indeed works the best on few-shot drug discovery, the authors can either (i) directly compare with [a1,a2] on FS-Mol, or (ii) test the proposed MHNfs on Tox21, MUV and SIDER, such that the performance can be compared with [a1,a2,a3]. 
- Writing: The paper should be carefully proofread to correct title case, typos and grammar issues. In Table 1, the author only mark one result in bold while neglecting the others. 

[a1] Few-Shot Graph Learning for Molecular Property Prediction, WWW, 2021
[a2] Property-Aware Relation Networks for Few-Shot Molecular Property Prediction, NeurIPS, 2021
[a3] Low Data Drug Discovery with One-Shot Learning, ACS Cent. Sci., 2017


 Please reply to my concerns in weakness. 

Key questions: 
- There exists a number of works using external memory, such as [b1,b2]. Why consider Modern Hopfield Networks instead? Any strength?
- Please provide discussion and empirical comparison w.r.t [a1,a2]. 

[b1] Meta-learning with memory-augmented neural networks, ICML, 2017
[b2] Few-shot visual learning with contextual memory and fine-grained calibration, IJCAI, 2020
 The authors provide discussion on limitations (including application domain and hyperparameter choices) and broader impact. The authors did not mention limitations of the proposed MHNfs from the perspective of methodology.  

Meta-review: This paper proposes a context-enriched molecular representation approach for few-shot drug discovery. Specifically, they show that the proposed MHNfs outperform existing methods on standard benchmarks including FS-Mol and the Tox21.
The proposed approach is analogous to the retrieval-based approach or generation models by editing prototypes in the deep learning community, which enrich molecular representations through querying from a massive reference space, namely first looks at a group of context drug-like molecules, and then using self-attention to distribute that information among the query and support set molecules. 
This is a borderline paper and the main concern of this paper is the limited insights in result analysis (e.g., what did the model learn from the context). I recommend the authors to address the weaknesses above and resubmit to another venue.

Decision: Reject