Context-enriched molecule representations improve few-shot drug discovery
Authors: Johannes Schimunek
Philipp Seidl
Lukas Friedrich
Daniel Kuhn
Friedrich Rippmann
Sepp Hochreiter
Günter Klambauer
Venue: NeurIPS 2022 Submitted (2022)

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.

Context-enriched molecule representations improve few-shot drug discovery
Anonymous Author(s)
Affiliation Address email
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 new method for few-shot drug discovery. Its main idea is to enrich a molecule representation by knowledge about known context reference molecules. Our novel concept for molecule representation enrichment to associate molecules from both the support set and the query set with a large set of reference (context) molecules through modern Hopfield network. Intuitively, this enrichment step is analogous to a human expert who would associate given molecule with familiar molecules whose properties are known. The enrichment step reinforces and amplifies the covariance structure of the data, while simultaneously removing spurious correlations arising from the decoration of molecules. Our approach is compared with other few-shot methods for drug discovery on the FS-Mol benchmark dataset. On FS-Mol, our approach outperforms all compared methods and therefore sets new state-of-the art for few-shot learning drug discovery. ablation study shows that the enrichment step of our method is the key to improve the predictive quality. In domain shift experiment, we further demonstrate the robustness of our method.
2
8
10
11
13
14
15
17
18
19
Introduction
20
21
To improve human health, combat diseases, and tackle pandemics there is a steady need of discovering new drugs in fast and efficient way. However, the drug discovery process is time-consuming and cost-intensive (Arrowsmith) 2011). Deep learning methods have recently been shown to reduce time and costs of this process (Chen et al. 2018 Walters and Barzilay 202T). They diminish the required number of both wet-lab measurements and molecules that must be synthesized (Merk et al., 2018: Schnnider al. 2020b. However, of now, deep learning approaches
23
24
25
only the molecular information about the ligands after being trained on large training set. At inference time, they yield
27
highly accurate property and activity prediction (Mayr al., 2018: Yang al. 2019). generative et al. 2018), or synthesis models (Segler et al. 2018b Seid!
29
Segler et al. 2018a
al. 2022
31
Deep learning methods in drug discovery usually require large amounts of biological measure- ments. To train deep learning-based activity and property prediction models with high predictive per- formance, hundreds or thousands of data points per task are required. For example, well-performing predictive models for activity prediction tasks of ChEMBL have been trained with average of 3,62 activity points per task. i.c., drug target, by (Mayr et al. 2018). The EXCAPE-DB dataset provides on average 42,501 measurements per task (Sun et al. 201 Sturm et al. 2020 (Wu et al. 2018) pub-
33
35
Submitted to 36th Conference on Neural Information Processing Systems (NeurIPS 2022). Do not distribute.

---

37
76 molecule representations (Ramsauer et al. 2021 Theorem 3). The representations that are retrieved from the MHNs replace the original representations of the query and support set molecules. Those lished a large scale benchmark for molecular machine learning, including prediction models for the SIDER dataset (Kuhn et al. 2016) with an average of 5,187 data points, Tox21 (Huang et al. 20T
38
39
Mayr et al. 2016) with on average 9,031, and Clin'Tox (Wu et al. 2018) with 1,491 measurements per task. However, for typical drug design projects, the amount of available measurements is very limited (Stanley et al. | Waring et al. 2015 Hochreiter et al.] 2018). since in-vitro experiments are expensive and cone-consu Therefore, methods that need only few measurements to build precise prediction models are desirable. This problem - i.e. the challenge of learning from few paaa points - is the focus of machine learning areas like meta-learning (Schmidhuberi 1987] Bengio
40
41
42
43
44
et al.,
1991 Hochreiter al. 2001) and few-shot learning (Miller et al. 2000 [endre al. 2020 Wang et al. 2020).
46
47
Few-shot learning tackles the low-data problem that is ubiquitous in drug discovery. Few-shot
48
learning methods have been predominantly developed and tested on image datasets (Bendre et al.
49
2020:
Wang et al. 2020). and have recently been adapted to drug discovery problems (Chen et al.
50
2022
Wang et al. 2021 Stanley et al., 2021 Altae- et al. 2017). They are usually categorized
51
into three groups according to their main approach (Bendre et al. 2020; Wang et al., 2020 Adler et al.,
52
2020). a) Data-augmennationn-oee
approaches augment the available samples and generate new, more diverse data points (Chen et al.) 2020; Zhao et al. 2019; Antoniou and Storkey 2019). b) Embedding- based and nearest neighbour approaches learn embedding space representations. Predictive models can then be constructed from only few net data points by comparing these embeddings. For example in Matching Networks (Vinyals et al. 2016) an attention mechanism that relies on embeddings is the basis for the predictions. Prototypical Networks (Snell et al., 2017) create prototype representations
53
54
56
58 for each class using the above mentioned representations in the embedding space. c) Optimization- based fine-tuning methods utilize meta-optimizer that focuses efficiently navigating the parameter space. For example, with MAML the meta-optimizer learns initial weights that can be adapted novel task by few optimization steps (Finn al., (017).
60
61
62 Most of these approaches have already been applied to few-shot drug discovery (see Sec. D. Surpris- ingly, almost all these few-shot learning methods in drug discovery are worse than naive baseline, which does not even use the support set (see Section 5). We hypothesize that the under-performet of these methods stems from disregarding the context both in terms of similar molecules and similar activities. Therefore, we propose a method that informs the representations of the query and support set with large number of context molecules covering the chemical space.
63
64
66
67
68
Enriching molecule representations with context using associative memories. In data-scarce situations, humans extract co-occurrences and covariances by associating current perceptions with memories (Bonner and Epstein 2021 Potter 2012). When we show a small set of active molecules to human expert in drug discovery, the expert associates them with known molecules to suggest further active molecules (Gomez 2018 He et al. 202T). In an analogous manner, our novel concept for
70
71
72
73 few-shot learning uses associative memories to extract ro-occurrences and the covariance structure
74 of the original data and to amplify them in the representations (Fürst et al. 2021). We use Modern Hopfield Networks (MHNs) as associative memory, since they can store large set of context
78 retrieved representations have amplified co-occurrences and covariance structures, while peculiarities
79 and spurious co-occurrences of the query and support set molecules are averaged out.
80 In this work, our contributions are the following:
81
We propose a new architecture MHNIs for few-shot learning in drug discovery.
82
We achieve a new state-of-the-art on the benchmarking dataset FS-Mol.
83
We introduce a novel concept to enrich the molecule representations with context by associ- ating them with large set of context molecules.
94
85
We add naive baseline to the FS-Mol benchmark that yields better results than almost all other published few-shot learning methods.
86
87
We provide results of an ablation study and a domain shift experiment to further demonstrate the effectiveness of our new method.
2

---

2
Problem setting
89
90
Drug discovery projects revolve around models g(m) that predict molecular property or activity D. given a representation m of an input molecule from a chemical space M.. We consider machine learning models 9 = Iw (m (m)) with parameters w that have been selected using training set. Typically, deep learning based property prediction uses molecule encoder fME .M - R The molecule encoder can process different symbolic or low-level representations of molecules. such as molecular Mayr et al. (016), SMILES (Weininger
91
92
94
descriptors (Bender et al. [004: Unterthiner et al.
2014:
96
1988 Mayr et al. 2018| Winter et al. 2019: Seg ler et al. 2018a), or molecular graphs (Merkwirth
97
and Lengauer 2005; Kearnes 200] Yang al. 2019 liang al. 2021 and can be pre-trained
98
on related property prediction tasks.
For few-shot learning, the goal is to select high-quality predictive model based small set of molecules (ZI,...,33, with associated measurements y = {y/,..... MN The measurements are usually assumed to be binary V/, E {-1,1}, corresponding to the molecule being inactive active. is called the support set that contains samples from a prediction task and N is the support set size. The goal is to construct a model that correctly predicts y for an that is not in the support set - in other words, a model that generalizes well.
100
The set {<.,. yn) )
N m=
102
104
105
Standard supervised machine learning approaches typically just show limited predictive power at this task (Stanley et al., (02T) since they tend to overfit on the support set due to a small number of training samples. These approaches learn the parameters w of the model Iw from the support set in a supervised manner. However, they heavily overfit to the support set when N is small. Therefore, few-shot learning methods are necessary to construct models from the support set that generalize well to new data.
106
107
108
110
MHNIs: Hopfield-based molecular context enrichment for few-shot drug discovery
111
113 We aim at increasing the generalization capabilities of few-shot learning methods in drug discovery
114 by enriching the molecule representations with molecular context. In comparison to the support set,
115 which encodes information about the task, the context set i i.e. a large set of molecules includes
116 information about large chemical space. The query and the support set molecules perform retrieval
117
from the context set and thereby enrich their representations. We detail this in the following.
118
3.1
Model architecture
119
We propose an architecture which consists of three consecutive modules. The first module, a) the context module fCM, enriches molecule representations by retrieving from large set of molecules. The second module, b) the cross-altention module fCAM (Hou et al. [2019] Chen et al. 2021). enables
120
121
122
the effective exchange of information between the query molecule and the support set molecules.
123
Finally the prediction for the query molecule is computed by using the usual c) similarity module
124
SSM
(Koch et al. 2015: Altae-Tran et al. 2017):
context module:
m' = -JOM(m, C)
X' = FCM (X,C) C),
(1)
cross-attention module:
Im" x = FCAM (Im', (Immm') ÿ = FSM xx" y):
(2)
similarity module:
(3)
where m E IRd is molecule embedding from trainable or fixed molecule encoder, and m' and m" are enriched versions of it. Similarly, X E R N contains the stacked embeddings of the support set molecules and X' and X" are their enriched versions. C E Rdd M is large set of stacked molecule embeddings, y are the support set labels, and Ø is the prediction for the query molecule. Square brackets indicate concatenation, for example [m' X') is matrix with N + I columns. The modules JCM PCAM, and JSM are detailed in the paragraphs below. An overview of our architecture is given in Figure 1 The architecture also includes skip connections bypassing f PCM (.,.) and (CAM( and
125
126
127
128
129
131
layer normalization (Ba et al. 2016), which are not shown in Figurq]
3

---

Prediction
Similarity Module
Final Representations
Retrieval
Attention Module
Multi-teed Attention
Context-enriched Representations Activity Encoding
Associations
Linear Projection
Context Module IHlopfield Layer)
Wx
Norm
Initial Embedding
Lincr Projection
Molecule Encoder
Structure
c a gog angg corni
Compounds
Label
Query
Support
Cross Attention Module
Figure 1: Schematic overview of architecture. Left: All molecules are fed through shared molecule encoder to obtain embeddings. Then, the context module (CM) enriches the representations by associating them with context molecules. The cross-attention module (CAM) enriches representa- tions by mutually associating the query and support set molecules. Finally, the similarity module computes the prediction for the query molecule. Right: Detailed depiction of the operations in the CM and the CAM.
shared molecule encoder fM creates embeddings for the query molecule m = fME fM))mm) the support set molecules n (x and the context molecules Cm. fM (cm). There many possible choices for fixed or adaptive molecule encoders (see Section 2 of which we use descriptor- based fully-connected networks because of their computational efficiency and good accuracy (Dahl
133
135
137 et al. [2014] Mayr et al. 2016 2018). For notational clarity we denote the course of the representations
through the architecture:
MR
CM
CAM
771 symbolic or low-level repr.
m malecule embedding
771 context repr.
m similarity repr.
(4)
MR
CM
CAA
Zy symbolic low-level repr.
It molecule embedding
T context repr.
x=" similarity repr.
(5)
139
3.2
Context module (CM)
140
The context module associates the query and support set molecules with large set context molecules, and represents them as weighted average of context molecule embeddings. The context module is realised by continuous Modern Hopfield Network (MHN) (Ramsauer al. 2021). An MHN is content-addressabll associative memory which can be built into deep learning architectures. attention mechanism of
142
144
There exists an analogy between the energy update of MHNs and Transformers (Vaswani et al., 2017 Ramsauer et al.] 2021). MHNs are capable of storing and Rex given a state pattern & E Re that represents query.
146
retrieving patterns from a memory M
The retrieved pattern gnow
IR is obtained by
nnew
M p M softmax (BMT&).
148
where p is called the vector of associations and 3 is scaling factor or inverse temperature. Modern Hopfield Networks have been successfully applied to chemistry and computational immunology
149
(Sedd et al. 2022 Widrich et al. 2020).
We this mechanism in the form Hopfield layer, which first maps raw patterns to associative space using linear transformations, and uses multiple simultaneous queries E E dx
152
Hopfield(E. (E, C) = (WgC) softmax (3 (WoC) (W=E)).

---

where WE E Iddd and Wc. W E Rexd are trainable parameters of the Hopfield layer, softmax applied column-wise, and is hyperparameter. Note that in principle the E and C could have different second dimension as long as the linear transformations map to the same dimension e. Note that all embeddings that enter this module are first layer normalized (Ba et al., (016). Several of these Hopfield layers can run in parallel and we refer to them as "heads" in analogy to Transformers Vasswan et al. 2017).
153
154
155
156
157
158
159
The context module of our new architecture uses a Hopfield layer, where the query patterns are the embeddings of the query molecule m and the support set molecules X. The memory is composed of embeddings of a large set of M molecules from chemical space, for example reference molecules, here called context molecules C. Then the original embeddings m and X are replaced by the retrieved embeddings, which are weighted averages of context molecule embeddings:
180
162
183
(8) This retrieval step reinforces the covariance structure of the retrieved representations (see Ap- pendix A.7). Note that the embeddings of the query and the support set molecules have not yet influenced each other. These updated representations m'. X'' are passed to the cross-attention module.
m' Hopfield C) and X' (opfield(X, ( C).
164
166
167
3.3
Cross-attention module (CAM)
168
For beaheedding-bsseee few-shot learning methods in the field of drug discovery, Altae- Tran et al. (2017) showed that the representations the molecules can be enriched, if the architecture allows information exchange between query and support set molecules. AAlaae Tran et al.) (017)) uses an attention-
169
170
171
enhanced LSTM variant which updates the query and the support set molecule representations in an
iterative fashion, being aware of each other. We further develop this idea and combine with the idea of using a transformer encoder layer (Vaswani et al., 2017) as a cross-attention module (Hou et al. Chen et al. (021).
173
2019
175
The cross-attention module updates the query molecule representation m' and the support set molecule representations X' by mutually exchanging information, using the usual Transformer mechanism:
178
177
im". x"I Hopfield((m', x'), [m', X')).
where Xm] E [N+(N+1) is the concatenation of the representations of the query molccule m'' with the support set molecules X' and exploited that the Transformer is special case of the Hopfield layer. Again, normalization is applied (Ba et al. 2016) and multiple Hopfield layers, i.c., heads,
178
180
in parallel, be stacked, and equipped with skip-connections. The representations m" and X" are passed to the similarity module.
182
183
3.4 Similarity module (SM)
184
In this module, pairwise similarity values k(m".a') .x") are computed between the representation of a query molecule m" and each molecule æç the support set as done recently (Koch al. 2015]
185
Altae-Tran et al. (017). Based on these similarity values, the activity for the query molecule is predicted, building weighted mean over the support set labels:
186
187
Σ Un k(mm R al))
(10)
N
188 where our architecture employs dot product similarity of normalized representations k(m", ((m",<)))
mn 'T x" (-) is the sigmoid function and is hyperparameter. Note that we use balancing
190 strategy for the labels v/a (/(2/N)) if yn =1 where NA is the number of actives and N₃ else
-N/(2/N)
is the number of inactives of the support set.
3.5 Architecture, hyperparameter selection, and training details
193 Hyperparameters. The main hyperparameters of our architecture are the number of heads, the embedding dimension, the dimension of the association space of the CAM and CM, the learning
5

---

195
rate schedule, the scaling parameter 3, and the molecule encoder. The following hyperparameters were selected by manual hyperparameter selection on the validation tasks. The molecule encoder consists of a single layer with output size d = 1024 and SELU activation (Klambauer et al. 2017).
196
197
196
The CM consists of one Hopfield layer with 8 heads. The dimension e of the association space is set
to 512 and B = V/.e. Since we use skip connections between all modules the output dimension of
199
200
the CM and CAM matches the input dimension. The CAM comprises one layer with 8 heads and an association- -pppac dimension of 1088. For the input to the CAM, an activity encoding was added to the support set molecule representations to provide label information. The SM uses T = 22.6. For the context set, randomly sample 5% from a large set of molecules i... the molecules in the FS-Mol training split for each batch. For inference, we used a fixed set of 5% of training set molecules as the context set for each seed. We hypothesize that these choices about the context could be further improved (Section 6). We provide considered and selected hyperparameters in Appendix A.1.6
202
204
206
207
Loss function, regularization and optimization. We use the Adam optimizer (Kingma and Ba 2014) to minimize the cross-entropy loss between the predicted and known activity labels. We use 208
learning rate scheduler which includes warm up phase, followed by section with constant learning rate, which is 0.0001, and a third phase in which the learning rate steadily decreases. As a regularization strategy, for the CM and the CAM dropout rate of 0.5 is used. The molecule encoder has a dropout with rate 0.1 for the input and 0.5 elsewhere (see also Appendix
210
212
213
Compute time and resources. Training a single MHNI's model on the benchmarking dataset FS- Mol takes roughly 90 hours of wall-clock time on an A100 GPU. In total, roughly 15,000 GPU hours were consumed for this work.
214
215
Related work
216
Several approaches to few-shot learning in drug discovery have been suggested (Altae- et al. 218
Nguyen et al. 2020 Guo et al. 2021 Wang et al., 202I). (Nuuyen et al., 2020) evaluated the applicability of MAML and its variants to graph neural networks (GNNs) and (Guo et al. 1
220
also combine GNNs and meta-learning. (Altae- Tran et al.] 2017) suggested an approach called Iterative Refinement Long Short-Term Memory, in which query and support set embeddings can share information and update their embeddings. Property-aware relation networks (PAR) (Wang al. )021) an attention mechanism to enrich representations from cluster centers and then learn a relation graph between molecules. (Chen et al.] 2022) propose to adaptively learn kernels and apply their method few-shot drug discovery with predictive performance for larger support set sizes. Recently, (Stanley et al., 2021) generated a benchmark dataset for few-shot learning methods in drug discovery and provided some baseline results.
222
228
224
225
226
227
228
Many successful deep neural network architectures use external memories, such as the neural Turing machine (Graves et al. 2014), memory networks (Weston et al. 2014), end-to-end memory networks SSkkhbataar et al.) 2015). Recently, the connection between continuous modern Hopfield networks
230
RRamsauer et al. 2021), ) which content- addressbbbd associative memories, and Transformer architectures (Vaswani et al.] (2017) has been established. We refer to (Le 2021) for an extensive overview of based architectures. Architectures with external memories have also been used for meta-learning (Vinyals el al. 2016 Santoro et al. (0166) and few-shot learning (Munkhdalai and
232
234
235
Yu
2017:
Ramalho and Garnelo, 2018 Ma et al. 2021
5
Experiments
236
237
5.1
Benchmarking on FS-Mol
238
Experimental setup.
Recently, the dataset FS-Mol (Stanlcy et al. 1021 was proposed to benchmark few-shot learning methods in drug discovery. was extracted from ChEMBL27 and comprises in total 489, measurements, 233, 333,786 compounds and 5,120 tasks. Per task, the mean number of data points 94. The dataset is well balanced as the mean ratio of active and inactive molecules is close to 1. The FS-Mol benchmark dataset defines 4,938 training, 40 validation and 157 test tasks, guaranteeing disjoint task sets. (Stanley et al. (021) precomputed extended connectivity fingerprints (ECFP) (Rogers and Hahn| 2010) and key molecular physical descriptors, which were defined by
239
240
241
242
243
245
RDKit (Landrum et 2006). While methods would be allowed to use other representations of the input molecules, such as the molecular graph, we used a concatenation of these ECFPs and

---

Table 1: Results on FS-MOL [AAUC-PR]. The best method is marked bold. Error bars represent standard errors across tasks according to Stanneyy et al. (202T). The metrics are also averaged across five
training re-runs and ten draws of support sets. In brackets the number of tasks per category is reported.
Method
All [157]
Kin. [125]
Hydrol. [20]
Oxid.[7]
GNN-ST (Stanley al. 2021
.029 .004
.027 + .004
.040 .018
.020 + .016
MAT (Maziarka et al. (2200)
.052 = .005
.043 = .005
.095 + .019
.062 + .024
Random Forest" (Breiman
2001
.092 .007
.081 = .009
.158 .028
.080 + .029
GNN-MT (Stanley et al.
202
.093 = .006
.093 + .006
.108 .025
.053 + .018
Similarity Search
.118 .008
.109 + .008
.166 + .029
.097 + .033
GNN-MAMI. (Guo et al.
2021
.159 = .009
.177 + .009
.105 + .024
.054 + .028
PAR(Wang et al. 2021)
.64 .008
.182 = .009
.109 .020
.039 = .008
Frequent hitters
.82 .010
.077 = .009
.098 .009
.041 0
ProtoNer' (Snell et al., 2017
.207 = .008
.215 = .009
.209 .030
.095 .029
Siamese Networks (Koch et al.
2015
223 .010
.211 = .010
.178 .026
.082 + .025
IterRefLSSTM (Altae Tran et al. 2017)
234 = .010
.211 = .010
.199 +99
.098 = .027
ADKF-IFT (Chen al. 2022)
234 .099
248 - .020
.217 0017
.106 + .008
MHNfs (ours)
.21 .009
.259
- 010
.199 - .027
.096 + .019
# metrics from Stanley et al. 2021
results from Chen et al. 2002
247
RDKit-based descriptors. For the main benchmark, the support set size was fixed to 16, using a stratified random split. We use all these settings of FS-Mol and therefore ensure fair method comparison.
249
Methods compared. Baselines for few-shot learning and our proposed method MIINfs were com- pared against each other. The Frequent Hitters model is a naive baseline that ignores the provided support set and therefore has to learn to predict the average activity of molecule. This method can potentially discriminate so-called frequent-hite hitter molecules (Stork et al. 2019) against molecules that are inactive across many tasks. We also added Similarity Search
251
253
et al. (015)) as a baseline. Similarity search is a standard chemoinformatics technique, used in situations with single few known actives. In the simplest case, the search finds similar molecules by computing a fingerprint or descriptoorrepressnnttion of the molecules and using a similarity measure k(.,-)
255
256
257
258
such as Tanimoto Similarity (Tanimoto 1960). Thus, Similarity Search, as used in chemoinfor- matics, can be formally written as 9 L/N Em Yn k(m, x..): where ZI..... Tn come from a fixed molecule encoder, such as chemical fingerprint or descriptor calculation. A natural exten- sion of Similarity Search with fixed chemical descriptors Neural Similarity Search or Siamese networks (Koch et al.
259
261
262
2015). which extend the classic similarity search by learning molecule
263
encoder:
T N Σm V/n (ME(m) PME (ax.)) Furthermore, we re-implemented the
264
RefSSSM (Altae- Tran et al.) (017)) in Pytorch. The IterReILSTM model consists of three modules. First, molecule encoder maps the query and support set molecules to its representations
285
286
LSTM variant, the actual ReILRSL iteratively updates the query and support set molecules, enabling information sharing between the molecules: RooRReLTSMM ((m, X)), where the hyperparameter controls the number of iteration steps of the IteRReILSTM. Third, a similarity module computes attention weights based on the rep- resentations: a = softmax (k (m', X')). These representations are then used for the final prediction:
m and X. Second, an invention-enhaneee
268
m' X
269
270
= ا i= 0:y/s- For further details, see Appendix A.1.5 The Random Forest baseline uses the chemical descriptors and is trained in standard supervised manner on the support set molecules for each task. The method GNN-ST is a graph neural network (Stanlcy et al., 2021 Gilmer et al.) 2017) that is trained from scratch for each task. The GNN-MT uses two step strategy: First, the model is pretrained on a large dataset on related tasks; second, an output layer is constructed to the few-shot task via linear probing (Stanley al. 2021] Alain and Bengio, 2016). The Molecule Attention Transformer (MAT) is pre-treeete trained in a self supervised fashion and fine-tuning is performed for the
271
272
273
274
275
276
277
278
few-shot task (Maziarka et al. 2020). GNN-MAML is based on MAML (Finn et al. 2017). and uses model-agnostic meta-learning strategy to find general core model from which one can easily adapt to single tasks. ProtoNet (Snell et al.) 2017) includes a molecule encoder, which maps query and
280
support set molecules to representations embedding space. In this embedding space, prototypical
7

---

0.250
0.745
0.245
dAUPRC
0.735
0.230
0.725
0.225
0.720
MHNI (CM+CAM+SMM SM)
MHNIS -CM
MHHS CCM (VCCAM, (CCCE.
MHNIS CAM+CAMMMM)
MHNNS -CM
MHNTS -CM (CCAM
Figure 2: Results of the ablation study. The boxes show the median, mean and the variability of the average predictive performance of the methods across training re-runs and draws of support sets. The performance significantly drops when the context module removed (light red bars), and when additionally the cross-attes coss-attention module is replaced with the IterReLLSTM module (light blue bars). This indicates that our two newly introduced modules, CM and CAM, play a crucial role in MHNIs.
representations of each class are built by taking the mean across all related support set molecules for 283
each class (details in Appendix A. 1.4. For all methods the most important hyperparameters were adjusted on the validation tasks of FS-Mol. The PAR model (Wang et al., (021)) includes GNN which creates initial molecule embeddings. These molecule embeddings are then enriched by an attention mechanism. Finally, another GNN learns relations between support and query set molecules.
285
287
The PAR model has shown good results for datasets which just include very few tasks such as Tox21 Wang al. 2021). CCen al, (2022) suggest framework for learning deep kernels by interpolating between meta letanlea and conventional deep kernels, which results in the ADKF-IFT model. The
289
290
model has exhibited especially high performance for large support sizes.
291
Training and evaluation. For the model implementations, we used PyTorch (Paszke et al. 2019 BSD license). We used PyTorch Lightning (Falcon et al. [019] Apache 2.0 license) as a framework for training and test logic, hydra for config file handling (Yadan [019 Apache 2.0 license) and Weights & Biases (Biewald [2000) MIT license) training reruns with different seeds for all methods, except Classic Similarity Search as there is variability across seeds. Each model was evaluated ten times by drawing support sets with ten different seeds.
293
experiment tracking tool. We performed five
295
297
The results in terms of area under precision-recall curve (AUC-PR) are presented in Table [ where the difference to a random classifier is reported (AAUC-PR). The standard error reported across tasks. Surprisingly, the naive baseline Frequent Hitters, that neglects the support set, has out- performed most of the few-shot learning methods, except for the embedding-based methods Siamese Networks, ProtoNet, IterReLLTMM, and MHNIs. IterRefLSTM, which has not been included in the FS-Mol benchmark study, reaches the second best performance. MHNIf has outperformed all other methods with respect AAUC-PR across all tasks, including the IterRefLSMM model va-velue 1.72e-7, paired Wilcoxon test), the ADKF-IFT model (p-value <1.0e-8, Wilcoxon test), and the PAR model (aluelue <1.0e-8, paired Wilcoxon test).
Results.
299
301
303
305
5.2
Ablation study
308
MHNfs has two new main components compared to the previous state-of-the-art method Iter- ReILSTM: i) the context module, and ii) the cross-attention module which replaces the LSTM-like
309
310
module. To assess the effects of these components, we performed an ablation study. Therefore, we compared MHNIS
method that does not have the context module ("MHNfs -CM") and to method that does not have the context module and uses an LSTM-like module instead of the For the ablation study, we used all 5 training reruns and evaluated each model 10 times on the test set with different support sets. The results of this ablation steps are presented Figure 2 Both removing the CM and exchanging the CAM with
312
313
CAM (" ("MHNfs CCM =(CAMM,teReffLSTMM
314
315
the IeerRefl.SS STM module were detrimental for the performance of the method (p-value 0.002 and 1.72e-7, 7, respectively; paired Wilcoxon test). The difference was even more pronounced under
317
domain shift (see Appendix A.3.3). Appendix A.3.2 contains second ablation study that examines

---

319
the overall effects of the context, the cross-attention the similarity module, and the molecule encoder of MHNI's
320
321
5.3 Domain shift experiment
We performed an experiment in which we evaluate models, that were pretrained on FS-Mol, on the Tox21 (Mayr et al.] 2016) dataset. There is a strong domain shift from the drug-like molecules of FS- --oo to the environmental chemicals, pesticides, and food additives of Tox21, such this dataset poses a challenging setting for few-shot learning methods. The experiment is described in detail in Appendix A.2 Our MHNf' approach has reached an AUC of .679 = .018 and has significantly outperformed the Iter SeeeLLST--asssd model (PAAUC-PR-Value 3.4e-5, 5, paired Wilcoxon test) and
323
324
325
328
327
328
the Classic Similarity Search (PAAUC-PR -vauue Pavalue 2.4e-9 paired Wilcoxon test) and therefore showed robust performance on the toxicity domain, see Table A6
329
6
Conclusion and discussion
330
331
We have introduced a new architecture for few-shot learning in drug discovery that is based on the novel concept to enrich molecule representations with context. In a benchmarking experiment, the architecture was assessed for its ability to learn accurate predictive models from small sets of labelled molecules and in this setting it outperformed all other methods. In domain shift study, the
332
333
334
335
robustness and transferability of the learned models has been assessed and again MHNIs exhibited the best performance. The resulting predictive models often reach AUC larger than .70, which means that enrichment of active molecules is expected (Simm et al. 2018) when the models are used for virtual screening. It has not escaped our notice that the specific context module we have proposed could immediately be used for few-shot learning tasks in computer vision, but might be hampered
336
337
338
339
340
by computational constraints. Limitations. While the implementation of our method is currently limited to small, organic drug-like molecules as inputs, our conceptual approach can also be used for mocro-molecules such as RNA, DNA or proteins. The output domain of our method comprises biological effects, such that the prediction must be understood in that domain. Our method demands higher computational costs and memory footprint as other embedding- based methods because of the calculations necessary for the context module. While we hypothesize that our approach could also be successful for similar data in the materials science domain, this has not been assessed. Our study is also constrained by a limited amount of hyperparameter search for all methods. Deep learning methods usually have a large number of hyperparameters, such as hidden dimensions, number of layers, learning rates, of which we were only able to explore the most important ones. The composition and choice of the context set is also under-explored and might be improved by selecting reference molecules with appropriate strategy. Broader impact. Impact on machine learning and related scientific fields. We envision that with (a) the increasing availability of drug discovery and material science datasets, (b) further improved biotechnologies, and (c) accounting for characteristics of individuals, the drug and materials discovery process will be made more efficient. For machine learning and artificial intelligence, the novel way in which representations are enriched with context might strengthen the general research stream to include more context into deep learning systems. Our approach also shows that such a system is more robust against domain shifts, which could be step towards Broad AI (Chollet 2019 Hochreiter 2022). Impact on society. If the approach proves useful, it could lead to faster and more cost-efficient drug discovery process. Especially the 1 pandemic has shown that it is crucial for humanity to speed up the drug discovery process to few years or even months. We hope that this work contributes to this effort and eventually leads to safer drugs developed faster. Consequences of failures of the method. As common with methods in machine learning, potential danger lies in the possibility that users rely too much on our new approach and use
341
342
343
344
345
346
348
350
352
354
356
357
358
359
360
361
382
383
364
it without reflecting on the outcomes. Failures of the proposed method would lead to unsuccessful wet lab validation and negative wet lab tests. Since the proposed algorithm does not directly suggest treatment or therapy, human beings are not directly at risk of being treated with a harmful therapy.
365
366
367
Wet lab and in-vitro testing would indicate wrong decisions by the system. Leveraging of biases in the data and potential discrinination. As for almost all machine learning methods, confounding factors, lab or batch effects, could be used for classification. This might lead to biases in predictions
369
or uneven predictive performance across different drug targets bioassays.

---

References
372
Adler, T., Brandstetter, J., Widrich, M., Mayr, A., Kreil, D., Kopp, M., Klambauer, G., and Hochreiter, S. (2020). Cross-domain few-shot learning by representation fusion. arXiv preprint arXiv:2010.06498.
374
Alain, G. and Bengio. Y. (2016). Understanding intermediate layers using linear classifier probes. arXiv preprint arXiv:1610.0160
376
377
A Alperstein, Z., Cherkasov, A., and Rolfe, J. T. (2019). All smiles variational autoencoder. arXiv preprint 1905 13343.
378
379
A Alran-Tran, H., Ramsundar, B., Pappu, A. S., and Pande, V. (2017). Low data drug discovery with one-shot learning. ACS central science, 3(4):283-293.
381
Antoniou, A. and Storkey, A. (2019). Assume, augment and learn: Unsupervised few-shot meta- learning via random labels and data augmentation. arXiv preprint arXiv:1902.0998 09884.
383
Arrowsmith, (2011). Phase ii failures: 2008-2010. Nature reviews drug discovery. 10(5).
384
Axelrod, S. and Gomez-Bombarelli, R. (2022). Geom, energy-annotated molecular conformations for property prediction and molecular generation. Scientific Data, 9(1):1-14.
385
386
Ba, L., Kiros, J. R., and Hinton, G. E. (2016). Layer normalization. arXiv preprint arXiv: 160%.01607
388
Bender, A., Mussa, H. Y., Glen, C., and Reiling, (2004). Similarity searching of chemical
databases using atom environment descriptors (molprint 2d): evaluation of performance. Journal of chemical information and computer sciences, 44(5):1708
390
381
Bendre, N., Marín, H. T., and Najafirad, P. (2020). Learning from few samples: A survey. arXiv preprint arXiv:2007.15488
392
393
Bengio, Y., Bengio, S., and Cloutier, J. (1991). Learning a synaptic learning rule. In Seattle
international joint conference neural networks.
385
Biewald, L. (2020). Experiment tracking with weights and biases. Software available from wandb.com.
397
Bonner. M. F. and Epstein, R. (2021). Object representations in the human brain reflect the
398
co-occurrence statistics of vision and language. Nature Communications, 12(4081).
Breiman, L. (2001). Random forests. Machine learning. 45(1):5-32.
400
Massague A., Ojeda, M. J., Valls, C., Mulero, M., Gareia- Gareia-Vallvé, S., and Pujadas, G. (2015). Molecular fingerprint similarity search in virtual screening. Methods, 71:58-63.
402
Chen, H., Engkvist, O., Wang, Y., Olivecrona, M., and Blaschke, T. (2018). The rise of deep learning in drug discovery. Drug discovery today, 23(6):1241-122
403
Chen, H., Li, H., Li, Y., and Chen, C. (2021). Sparse spatial transformers for few-shot learning. arXiv preprint Xrx::22109 129322
405
406
Chen, T., Kornblith, S., Norouzi, M., and Hinton, G. (2020). A simple framework for contrastive
407
learning of visual representations. In International conference on machine learning, pages 1,597 1607. PMLR.
409
Chen, W., Tripp. A., and Hernández-Lobato, J. M. (2022). Meta-leaming feature representations for adaptive gaussian processes via implicit differentiation. arXiv preprint arXiv:2205.02700
Chollet, F. (2019). On the measure of intelligence. arXiv preprint 1rXiv:1911.011
412 Dahl, E., Jaitly, N., and Salakhutdinov, R. (2014). Multi-task neural networks for qsar predictions. arXiv preprint arXiv: 1406.1241...1
10

---

414 Duvenaud, D., Maclaurin, D., eguulera-Iparraguirre J.,
R., Hirzel, T., Aspuru- Guzik, A., and Adams, P. (2015). Convolutional networks on graphs for learning molecular fingerprints. arXiv preprini arXiv: 1509.0992..
415
Eckert, H. and Bajorath, J. (2007). Molecular similarity analysis in virtual screening: foundations, limitations and novel approaches. Drug discovery today, 12(5-6):225-333.
418
Falcon, W. al. (2019). Pytorch lightning. hotpr//orhciiiiiiiioooooootteeee
GitHub.
Note: https://github.
420
3:6.
421
Finn, C., Abbeel, P., and Levine, S. (2017). Model-agnostic meta- meta-learning for fast adaptation of deep networks, International conference on machine learning, pages 1126-1135, PMLR.
422
423
Fürst, A., Rumetshofer, E., Tran, V., Ramsauer, H., Tang, F., Lehner, J., Kreil, D., Kopp, M..
424
Klambauer, G., Bitto-Nemling, A., al. (2021). Cloob: Modern hopfield networks with infoloob outperform clip, arXiv preprint arXiv:21 10.1:11
425
Geppert, H., Horváth, T., Gärtner, T., Wrobel, S., and Bajorath, J. (2008). Support-vecton pepport-veetor-mntte -ecchne- based ranking significantly improves the effectiveness of similarity searching using 2d fingerprints and multiple reference compounds. Journal of chemical information and modeling. 48(4))742-776.
427
429
Gilmer, J., Schoenholz, S. S., Riley, P. F. Vinyals, O., and Dahl, G. E. (2017). Neural message passing for quantum chemistry. In International conference machine learning, pages 1263-1272.
431
PMLR.
432
Gomez, L. (2018). Decision making in medicinal chemistry: The power of our intuition. ACS
433
Medicinal Chemistry Letters, 9(10):956-9) -588.
434
Bamee--ommaarlli, R., Wei, J. N., Duvenaud, D., Hernández-Lobato, J. M., Sánchez-Lengeling, B., Sheberla, D., Aguilera-Iparaguirre,, Hirzel, D., Adams, P., and Aspuru-Guzik, A. (2018). Automatic chemical design using data-driven continuous representation of molecules.
435
436
437
ACS central science, 4(2):268-276.
438
Graves, A., Wayne, G., and Danihelka, I. (2014). Neural turing machines. arXiv preprint 110:v:1410.540..
440
Guo, Z., Zhang, C., Yu, W., Herr, J., Wiest, O., Jiang, M., and Chawla, N. V. (2021). Few-shot graph learning for molecular property prediction. In Proceedings of the web conference 2021, pages 2559-2567.
442
443
He, J., You, H., Sandström, E., Nittinger, E., Bjerrum, E. J., Tyrchan, C., Czechtizky, W., and Engkvist, (2021). Molecular optimization by capturing chemist's s intuition using deep neural networks. Journal of cheminformatics, 13(1):1-17.
444
445
Hertz, T. Hillel, A. B., and Weinshall, D. (2006). Learning kernel function for classification with small training samples. In Proceedings of the 23rd international conference on machine learning, pages 401-408.
447
448
449
Hochreiter, S. (2022). Toward a broad ai. Communications of the ACM, 65(4))66557.
Hochreiter, S., Klambauer, G., and Rarey, M. (2018). Machine learning in drug discovery. Journal of Chemical Information and Modeling, 58(9):1723-)
451
452
Hochreiter, S., Younger, A. S., and Conwell, P. R. (2001). Learning to learn using gradient descent.
453
In International conference on artificial neural networks, pages 87-94. Springer.
454 Hou, R., Chang, H., Ma. B., Shan, S., and Chen, X. (2019). Cross attention network for few-shot classification. Advances in neural information processing systems 32.
456
Huang, R., Xia, M., Nguyen, D.-T., Zhao, T., Sakamuru, S., Zhao, J., Shahane, S. A., Rossoshek, A., and Simeonov, A. (2)16a). Tox2 challenge to build predictive models of nuclear receptor and 458 stress response pathways as mediated by exposure to environmental chemicals and drugs. Frontiers in Environmental Science, 3:85.

---

480
Huang, R., Xia, M., Sakamuru, S., Zhao, J., Shahane. S. A., Attene-Ramen M. Zhao, T., Austin, C. P., and Simeonov, A. (2016b). Modelling the tox21 10 chemical profiles for in vivo toxicity prediction and mechanism characterization. Nature communications, 7(1):1-10.
481
Jiang, D., Wu, Z,, Hsieh, C.-Y., Chen, G., Liao, B., Wang, Z.. Shen, C., Cao, D., Wu, J., and Hou, T. (2021). Could graph neural networks learn better molecular representation for drug discovery? comparison study of descriptor-based and graph-based models. Journal of cheminformalies, 13(1):1-23.
464
486
467
Kearnes, S., McCloskey, K., Berndl, M., Pande, V., and Riley, P. (2016). Molecular graph convolu- tions: moving beyond fingerprints. Journal of computer-aidee molecular design, 30(8):595-60) 608.
Kingma, D. P. and Ba, J. (2014), Adam: A method for stochastic optimization. arXiv preprint ArXiv: 1412,6990
470
471
Klambauer, G., Unterthiner, T., Mayr, A., and Hochreiter, S. (2017). Self-normalizing neural networks. In Advances in neural information processing systems 30, pages 972-981.
472
473
Koch, G., Zemel, R., Salakhutdinov, R., et al. (2015). Siamese neural networks for one-shot image recognition. In ICML deep learning workshop, volume 2. Lille.
474
475
Kuhn, M., Letunic, I., Jensen, L. J., and Bork, P. (2016). The sider database of drugs and side effects. Nucleic acids research, ((DD))DD07S-D107),
478
477
Landrum, G. et al. (2006). Rdkit: Open-souree cheminformatics.
478
Le, H. (2021). Memory and attention deep learning. arXiv preprint arXiv:2107.01390.
479
Li, J., Cai, D., and He, X. (2017). Learning graph-level representation for drug discovery. arXiv preprint
Li, P. Li, Y., Hsieh, C.-Y., Zhang, S., Liu, X., Liu, H., Song, S., and Yao, X. (2021). Trimnet: learning molecular representation from triplet messages for biomedicine. Briefings in Bioinformatics, 22(4):bbaa266.
482
Ma. Y., Liu, W., Bai, S., Zhang, Q. Liu, A., Chen, W., and Liu, X. (2021). Few-shot visual learning with contextual memory and ine-grained calibration. In Proceedings of the Twenty-Ninth International Conference on International Joint Conferences Artificial Intelligence, pages 811-817.
485
487
488
Mayr, A., Klambauer, G., Unterthiner, T., and Hochreiter, S. (2016). Deeptox: toxicity prediction using deep learning. Frontiers in environmental science, 3:80.
Mayr, A., Klambauer, G., Unterthiner, T., Steijaert, M., Wegner, J. K., Ceulemans, H.,, Clevert, D.-A., and Hochreiter, S. (2018). Large-scale comparison of machine learning methods for drug target prediction on chembl. Chemical science, :(44)5544135451
491
482
Maziarka, t.., Danel, T., Mucha, S., Rataj, K., Tabor, J., and Jastrzębski, S. (2020). Molecule attention transformer. arXiv preprint arXiv:2002.08264.
484
485
Merk, D., Friedrich, I.., Grisoni, F., and Schneider, G. (2018). De novo design of bioactive small
496
molecules by artificial intelligence. Molecular informatics, 37(1-2):170015;
497 Merkwirth, and Lengauer, (2005). Automatic generation of complementary descriptors with molecular graph networks. Journal of chemical information and modeling. (5(5):1159-1) 166-1168.
498
499
Miller, E. G., Matsakis, N. E., and Viola, P. A. (2000). Learning from one example through shared densities on transforms. In Proceedings ieee conference on computer vision and pattern recognition. cvpr 2000 (cat. no. PR00662), volume 1, pages 464-471.
500
501
502
Munkhdalai, T. and Yu, H. (2017). Meta networks. In International Conference on Machine Learning. PMLR.
503
pages
504
Nguyen, C. Q. Kreatsoulas, C., and Branson, K. M. (2020). Meta-learning gnn initializations for
low-resouree molecular property prediction. arXiv preprint arX1v:2003.05996
12

---

506
Oord, A. v. d.. Li, Y., and Vinyals, O. (2018). Representation learning with contrastive predictive coding. arXiv preprint arXiv:180700377
507
508
Paszke, A., Gross, S., Chintala, S., Chanan, G., Yang, E., DeVito, Z., Lin, Z., Desmaison, A., Antiga,
509
L.., and Lerer, A. (2019). Automatic differentiation in pytorch. In Conference on neural information
510
processing systems.
Potter, M. (2012). Conceptual short term memory in perception and thought. Frontiers in Psychology, 3:113.
512
Ramalho, T. and Garell,, M. (2018). Adaptive posterior learning: few-shot learning with surprise- based memory module. In International Conference on Learning Representations.
514
515
Ramsauer, H., Schäfl, B.,, Lehner. J., Seidl, P., Widrich, M., Gruber, L., Holzleitner, M., Adler, T., Kreil, D., Kopp, M. K., Klambauer, G., Brandstetter, J., and Hochreiter, S. (2021). Hopfield networks is all you need. In International conference learning representations.
516
518
Riniker, S. and Landrum, G. A. (2013). Open-source platform benchmark fingerprints for ligand- based virtual screening. Journal of cheminformatiis, 5(1):1-17.
519
520
Rogers, D. and Hahn, M. (2010). Extended-conneettvttyy fingerprints. Journal of chemical information and modeling, 50(5):742-754.
521
Santoro, A., Bartunov, S., Botvinick, M., Wierstra, D., and Lillicrap, T. (2016). Meta-learning with
523
memory-augmented neural networks. In International conference on machine learning, pages 1842-1850 PMLR.
525
Schmidhuber, J. (1987). Evolutionary principles in self-referential learning.
526
Schneider, P., Walters, W. P., Plowright, A. T., Sieroka, N., Listgarten, J., Goodnow, R. A., Fisher,
527
J., Jansen, J. M., Duca, J. S., Rush, T. S., et al. (2020). Rethinking drug design the artificial intelligence era. Nature reviews drug discovery, 19(5):353-28) 33.-64...
528
529
Segler, M. H., Kogej, T., Tyrchan, C., and Waller, M. (2018a). Generating focused molecule
530
libraries for drug discovery with recurrent neural networks. ACS central science, 4(1): ((1):220131)
Segler, M. H., Preuss, M., and Waller, M. P. (2018b). Planning chemical syntheses with deep neural networks and symbolic ai. Nature, 555(7698):60446610
532
Seidl, P., Renz, P., Dyubankova, N., Neves, P.. Verhoeven, J.. Wegner, J. K., Segler, M., Hochreiter, S., and Klambauer, G. (2022). Improving few-and zero-shot reaction template prediction using modern hopfield networks. Journal of chemical information and modeling, 62(9):2111 11-2120.
534
536
Sheridan, R. P. and Kearsley, S. K. (2002). Why do we need so many chemical similarity search methods? Drug discovery today, 7(17):903-911.
538
Simm, J., Klambauer, G., Arany, A., Steijaert, M., Wegner, J. K., Gustin, E., Chupakhin, V., Chong, Y. T., Vialard, J., Buijnsters, P., et al. (2018). Repurposed high-throughput image assays enables biological activity prediction for drug discovery. Cell Chemical Biology, page 108399.
539
540
541
Snell, J.. Swersky, K., and Zemel, R. S. (2017). Prototypical networks for few-shot learning, arXiv preprint arXiv:1703.05175
542
543 Stanley, M., Bronskill, J. F., Maziarz, K., Misztela, H., Lanini, J., Segler, M., Schneider, N., and Brockschmidt, M. (2021). Fs-mol: A few-shot learning dataset of molecules. In Conference on neural information processing systems workshop.
545
Stork, C.. Chen, Y. Sicho, M., and Kirchmair, J. (2019). Hit dexter 2.0: machine-learingg models for the prediction of frequent hitters. Journal of chemical information and modeling, 5933),0030-1043,
547
Sturm, N., Mayr, A., Le Van, T., Chupakhin, V., Ceulemans, H., Wegner, J., Golib-Dzib, J.-F., Jeliazkova, N., Vandriessche, Y., Böhm, S., et (2020). Industry-seale application and evaluation of deep learning for drug target prediction. Journal of Cheminformatics, 12(1):1-13.
549
13

---

551
Sukhbaatar, S., Weston, J., Fergus, R., et al. (2015). End-to-end memory networks. Advances in neural information processing systems, 28.
552
558
Sun, J., Jeliazkova, N., Chupakhin. V., Golib-Dzib, J.-F, Engkvist, O., Carlsson, L., Wegner, J., Ceulemans, H., Georgiev, I., Jeliazkov, V., al. (2017). Excape-db: integrated large scale dataset facilitating big data analysis in chemogenomics. Journal of cheminformatics, 9(1):199.
554
855
Tanimoto, T. (1960). Ibm type 704 medical diagnosis program. IRE transactions medical electronics, (4):280-283.
557
558
Torres, L., Monteiro, N., Oliveira, J., Arrais, J., and Ribeiro, B. (2020). Exploring a siamese neural network architecture for one-shot drug discovery. In 2020 ieee 20th international conference on bioinformatics and bioengineering (bibe), pages 168-175.
S59
Unterthiner, T., Mayr, A., Klambauer, G., Steijaert, M., Wegner, J. K., Ceulemans, H., and Hochreiter, S. (2014). Deep learning as an opportunity in virtual screening. In Advances in neural information processing systems workshop.
562
563
664
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.., Gomez, A. N., Kaiser, E., and
585
Polosukhin, I. (2017). Attention is all you need. In Advances in neural information processing
586
systems, pages 5998-6008.
Vinyals, O., Blundell, C., Lillicrap, T., Wierstra, D., et al. (2016). Matching networks for one shot learning. Advances in neural information processing systems, 29:3630-386 -3388
668
569
Walters, W. P. and Barzilay, R. (2021). Critical assessment of ai in drug discovery. Expert opinion on drug discovery, pages 1-11.
570
571
Wang, X., Huan, J., Smalter, A., and Lushington, G. H. (2010). Application of kernel functions for accurate similarity search in large chemical databases. In BMC bioinformatics, volume 11, pages
572
573
1-14. BioMed Central.
Wang, Y., Abuduweili, A., Yao, Q., and Dou, D. (2021). Property-aware relation networks for
575
few-shot molecular property prediction. Advances in Neural Information Processing Systems, 34.17441-17454.
577 Wang, Y., Yao, Q., Kwok, J. T., and Ni, L. M. (2020). Generalizing from a few examples: A survey on few-shot learning. ACM computing surveys (csur), 53(3):1-34.
579
Waring, M. J., Arrowsmith, J., Leach, A. R., Leeson, P. D., Mandrell, S., Owen, R. M., Pairaudeau, G., Pennie, W. D., Pickett, S. D., Wang, J., et al. (2015), An analysis of the attrition of drug candidates from four major pharmaceutical companies. Nature reviews drug discovery, 14(7):4754) 486.
580
581
582
Weininger, D. (1988). Smiles, a chemical language and information system. 1. introduction to methodology and encoding rules. Journal of chemical information and computer sciences, 28(1) 11:1: 36.
584
Weston, J.. Chopra, S., and Bordes, A. (2014). Memory networks. arXiv preprint arXiv:1410.3911
586
Widrich, M., Schäfl, B., Pavlović, M., Ramsauer, H., Gruber, L., Holzleitner, M., Brandstetter, J.,
587
Sandve, G. K., Greiff, V., Hochreiter, S., et al. (2020). Modern hopfield networks and attention for
588
immune repertoire classification. Advances in neural information processing systems 33.
589
Willett P. (2014). The calculation of molecular structural similarity: principles and practice. Molecu- lar informatics, 33(6-7)4403-41.
590
591
Winter, R., Montanari, F., Noé, F., and Clevert, D.-A. (2019). Learning continuous and data-driven
molecular descriptors by translating equivalent chemical representations. Chemical science, 1926-1162
583
Wu, Z., Ramsundar, B., Feinberg, E. N., Gomes, J., Geniesse, C., Pappu, A. S., Leswing, K., and Pande, (2018). Moleculenet: a benchmark for molecular machine learning. Chemical science, 9())5513-530.
585
14

---

597
Yadan, O. (2019). Hydra a framework for elegantly configuring complex upplications. Github. Visited 2022-04-25.
596
599
Yang, K., Swanson, K., Jin, W., Coley, C., Eiden, P., Gao, H., Guzman- A., Hopper, T., Kelley, B., Mathea, M., et al. (2019). Analyzing learned molecular representations for property prediction. Journal of chemical information and modeling, 59(8):3370-333 -388.
601
602
Ye, M. and Guo, Y. (2018). Deep triplet ranking networks for one-shot recognition. arXiv preprint arXiv: 1804.0777.
604
Zaslavskiy, M., Jégou, S., Tramel, E. W., and Wainrib, G. (2019). Toxicblend: Virtual screening of toxic compounds with ensemble predictors. Computational Toxicology, 10:81-88.
805
608
Zhao, A.. Balakrishnan, G., Durand, F., Guttag, J. V., and Dalca, A. V. (2019). Data augmentation
607
using learned transformations for one-shot medical image segmentation. Proceedings of the ieee conference on computer vision and pattern recognition, pages 8543 -5533.
15

---

Checklist
609
610
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
612
contributions and scope? [Yes]
613
(b) Did you descrihe the limitations of your work? [Yes] See Section 6
614
(c) Did you discuss any potential negative societal impacts of your work? [Yes] See
Section 6
615
616
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results... 618
619
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
620
(b) Did you include complete proofs of all theoretical results? [N/A]
621
3. If you ran experiments..
622
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URL)? [Yes]
(b) Did you specify all the training details (e.g.. data splits, hyperparameters, how they were chosen)? [Yes] We refer e.g. to Section 5 in which we provide this information. (c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes] We report error bars for all performance metrics. For all experiments the variability across re-runs, different support sets and prediction tasks was assessed.
625
626
627
629
630
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)? [Yes] See Section 3.5
4. If you using existing assets (e.g., code, data, models) or curating/releasing new assets...
633
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [Yes]
(c) Did you include any assets either in the supplemental material as URL? [Yes]
636
(d) Did you discuss whether and how consent was obtained from people whose data you're rou're using/curating? [N/A]
637
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [Yes] See Section
639
640
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, if applicable? [N/A]
642
643
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
644
645
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
16

---

Contents of the appendix
647
648
Appendix
17
A.1
Details
methods
17
650
A.1.1
Frequent hitters: details and hyperparameters
17
A.1.2
Classic similarity search: details and hyperparameter
18
652
A.1.3
Neural Similarity Search or Siamese networks: details and hyperparameters
19
653
A.1.4
ProtoNet: details and hyperparameters/
19
654
A.1.5
IteRefLSTMM details and hyperparameters
20
655
A.1.6
MHNís: details and hyperparameters)
21
A.1.7
PAR: details and hyperparametens)
22
657
A.2
Domain shift experiment]
24
A.3
Details on the ablation study
25
659
A.3.1 Ablation study A: comparison against IterRefl.ST STM
25
A.3.2
Ablation study B: all design elements
25
661
A.3.3 Ablation study C: Under domain shift on Tox21
26
A.4
Generalization to different support set sizes
26
663
A.5
Generalization to different context sets
26
A.6
Details and insights the context module
27
685
A.7
Reinforcing the covariance structure the data using modern Hoptield networks
27
686
A
Appendix
A.1
Details methods
668
Fho-shot learning methods drug discovery can be described as models with adaptive parameters wr that use support set Z { v1), (æ ]N)] [as additional input to predict a label for
670
molecule m
g = 9w (m, Z).
(A1)
671
Optimization-basee methods, such as MAML (Fin et al. 2017) use the support set to update the parameters
[a(w; Z) (m).
(A2)
673
where a(.) is a function that adapts w of g based on for example via gradient-descent.
674
Embedding- based methods use different approach and learn representations of the support set molecules {=1...... N sometimes written as stacked embeddings X E Rd and the query molecule m, and some function that associates these two types of information with each other. We describe the embedding-based methods Similarity Search Section A.1.2 Neural Similarity Search in Section A.1.3] ProtoNet in Section A.1.4] IterRefL in Section A.1. PAR in Section The "frequent hitters" baseline described
675
677
679
and MHNfs in the main paper and details Section 1.6
in Section A.I.I
G81
A.1.1
Frequent hitters: details and hyperparameters
The "frequent hitters" model g FH is baseline that we implemented and included in the method comparison. This method uses the usual training scheme of sampling query molecule m with label y. having access support set contrast to the usual models of the type Iwo (m, Z). the frequent hitters model FH neglects the support set:
682
684
② - g (m).
(A3)
Thus, during training for the same molecule m, the model might have predict both y 1 and 1, since the molecule can be active in one task and inactive in another task. Therefore, the
687
We use Z to denote the support set already embedded molecules to keep the notation uncluttered. More correctly, the methods have access to the raw support set Z = {(x1, (141,,)),. y.)..., (æN, yn)}, where x., is symbolic, such the molecular graph, low-level representation of the molecule.
17

---

Table A1: Hyperparameter space considered for the Frequent hitters model. The hyperparameters of the best configuration are marked bold.
Hyperparameter
Explored values
Number of hidden layers
1, 2, 4 1024, 2048, 4096
Number of units per hidden layer
Output dimension
512, 1024 ReLU
Activation function
Learning rate
0.0001, 0.001 Adam,A
Optimizer
Weight decay
0, 0.01
Batch size
32, 128, 512, 2048, 4096 0, 0.1
Input Dropout
Dropout
0.1, 0.2, 0.3, 0.4, 0.5 False, True False, True dot product
Layer-normalizztion
Affine
Similarity function
model tends to predict average activity of molecule to minimize the cross-entropy loss. We chose an additive combination of the Morgan fingerprints, RDKit fingerprints, and MACCS keys for the input representation to the MLP.
689
691
Hy perpaaraeete search. We performed manual hyperparameter search on the validation set and report the explored hyperparameter space (Table AT). We use early-stopping based validation seerage-precision a patience of 3 epochs and train for maximum 20 epochs with a linear warm-up learning-ral rate schedule for the first 3 epochs.
693
895
A.1.2
Classic similarity search: details and hy hyperparameters
696
Similarity Search (Cereto-Massagué Massagué et al., 2015) is a classic chemoinformatics technique used in
situations in which single or few actives are known. In the simplest case, molecules that are similar to
a given active molecule are searched by computing fingerprint or deserripoorreprrseenttiion prese (m
688
of the molecules and using similarity measure k(.,,.). such Tanimoto Cimilarity(Tanimmon 1960). Thus, the Similarity Search as used in chemoinformatics can be formally written as:
700
N
0 1/N
Σ Un fi(deesc pdess (I,)), n=1
(A4)
where the function pdess maps the molecule to its chemical descriptors fingerprints and takes the role of both the molecule encoder and the support set encoder. The association function fassoc consists of a) the similarity measure k(.,.) and then b) mean pooling across molecules weighted by their similarity and activity.
702
703
704
705
Notably, there are many variants of Similarity Search (Cereto-Massagué et al. 2015 Wang et al.
706
2010
|Eckrtt and Bajorath, 2007 Geppert et al. 2008 Willett 2014 Sheridan and Kearsley 2002
707
Riniker and Landrum (003)) 3) of which some correspond to recent sho- learning methods with
708
fixed molecule encoder. For example, (Geppert al., 2008) suggest use centroid molecules, i.c.
709
prototypes or averages of active molecules. This is equivalent to the idea of Prototypical Networks
710
(nell et al.) 2017). Riniker and Landrum (2013) are aware of different fusion strategies for sets of
active or inactive molecules, which corresponds to different pooling strategies of the support set. Overall, the variants the classic Similarity Search are highly similar to embedding- few-shot learning methods except that they have a fixed instead of learned molecule encoder.
712
714
Hyperparameter search. For the Similarity Search, there were two decisions to make which was firstly the similarity metric and secondly the question whether we should use balancing strategy like shown in Section 3.4. We decided for the dot-product as similarity metric and using the balancing strategy. These decisions made by evaluating the models on the validation set.
716
18

---

input molecule
FCDNN
LayerNorm
similarity function
molecule encoder
weights
embedding space
prediction
memory Goo are of agyiogl
FCDNN
LayerNom
target encoder
association function
Figure Al: Schematic overview of the implemented Neural Similarity Search variant
A.1.3 Neural Similarity Search Siamese networks: details and hyprpaarameters
lot of related work already done (Koch al., 2015 Hertz et al. 2006 Ye and Guo 2018
720
Torres et al. 2020). We adapted these ideas, such that a fully-connected deep neural network followed
by Layer Normalization (Ba et al. 2016) operation, fM with adaptive parameters w, is used in Siamese fashion to compute the embeddings for the input molecule and the support set molecules.
722
723
Within the association function block. pairwise similarity values for the input molecule and each support set molecule are computed, associating both embeddings via the dot product. Based on these similarity values, the activity for the input molecule is predicted, building the weighted mean over the support set molecule labels:
724
725
728
N Σ vn (M)(m) pME (M(() ==1
i -
(A5)
N
727
where (.) is the sigmoid function and T is hyperparameter in the range of vd. Note that this
method uses a balancing strategy for the labels v/= N/(2v/A) if Yn =1 -N/(2VN7) else where NA is the
number of actives and N, is the number of inactives of the support set. Figure Al
provides an
730
schematic overview of the Neural Similarity Search variant.
731
We trained the networks using the Adam optimizer (Kingma and Ba 2014) to minimize binary cross-entropy loss.
732
733
Hy pyerpaarreeter search. We performed manual hyperparameter search on the validation set. We
report the explored hyperparameter space (Table A2). Bold values indicate the selected hyperparame- ters for the final model.
735
736
A.1.4
ProtoNet: details and hyperparameters
737
Prototypical Networks (ProtoNet) (Snell et al. 2017). learn prototype r for each class. Concretely.
the support Z is class-wise separated into =F= {(z, 8) E y = 1} and Z ((x y) J.. For the subsets 7 and Z- prototypical representations t and r can be computed by
738
y
Σ
ME
(A6)
Z+
(*.y)
740 and
r
Σ
(x)
(x).
(A7)
7
(x.0)eZ
The prototypical representations r E Rd and the query molecule embedding m Id are then
742
used make the final prediction:
exp( (o(_d(m,r
(A8)
exp(-d(m,rt)) + exp( -d(mr--))
19

---

Table A2: Hyperparameter space considered for the Neural Search model selection. The hyperparameters of the best configuration are marked bold.
Hyperparameter
Explored values
Number of hidden layers
1. 2, 4
Number of units per hidden layer
1024, 4096
Output dimension
512, 1024
Activation function
ReLU, SELU
Learning rate
0.0001, 0.001, 0.01
Optimizer
Adam
Weight decay
0. 1 10-4
Batch size
4096
Input Dropout
0.1
Dropout
0.5
Layer-normalization
False, True
Affine
False
Similarity function
cosine similarity, dot product, MinMax similarity
743
where d is distance metric.
744
Hy perparameter search. Hyperparameter search has been done in Stanley al. (202I). which we refer here. ECFP fingerprints and descriptors created by GNN, which operates on the molecular graph, are fed into fully connected neural network, which maps the input into an embedding space with the dimension of 512. (Stanley et al. 2021) use the Mahalanobis distance to measure the similarity between query molecule and the prototypical representations in the embedding space. The learning rate is 0.001 and the batch size is 256. The implementation can be found here https
745
746
747
748
749
750
//gi www.tth..ooomcccaoo.c..o www..///w....ooooooooooooooooott /ttp///ooo////oo/ooo////////aaaa
wtt.mmol//oooo train. py and important http:///cccoooottttttoooo/.///.a
perparammetrs are chosen here
752
mo
753
Connection to Siamese networks and contrastive learning with InfoNCE. If instead of the neg- ative distance -d(...) the dot product similarity measure with appropriate scaling is used, ProtoNet for two classes becomes equivalent to Siamese Networks. Note that in our study, another differ- ence that ProtoNet uses GNN as encoder, whereas Siamese Networks use a descriptor-baser fully callecooneected network as encoder. In case of dot product similarity measure, the objective also becomes equivalent to contrastive learning with the InfoNCE objective (Oord et al. 2018).
754
796
758
A.1.5
IterRefLSTM: details and hyperparameters
780
Altae-Tran et al. 2017) modified the idea Matching Networks (Vinyals et al. 2016) by replacing
the LSTM with their Iterative Refinement Long Short-Term Memory (terRefLLTMM The use of the IterRefLSTM empowers the architecutre to update not only the embeddings for the input molecule
762
but also adjust the representations of the support set molecules.
For IterRefl. STM, m f (m) and æ, f ( are two potentially different molecule en-
764
785
coders for input molecule m and the support set molecules EI...., The next step in IterRffLSTM is:
786
m' X'I IterRefLSTMM X))-
787
Here, m' and X' contain the updated representations for the query molecule and the support set molecules. The IteRRefLSTM denotes the function which updates these representations. The main property of the IterRefLSTM module is that it is cermutation-qquiiariann, thus permu-
768
789
tation π (.) of the input elements results in the permutation of output elements: =(m'' X')) = (T XJ)) The full architecture in invariant to permutations of the support set elements. For details, we refer to FAltae-Tran et al. 2017). The hyperparameter T. E N controls the number of iteration steps of the IteRReLLSTM.
770
772
20

---

Table A3: Hyperparameter space considered for the IterRef model selection. The hyperparam- eters of the best configuration are marked bold.
Hyperparameter
Explored values
Molecule encoder
Number hidden layers
0, 1, 2, 4 1024, 4096
Number of units per hidden layer
Output dimension
512, 1024
Activation function
ReLU, SELU
Input dropout
0.1 0.5
Dropout
IterRef embedding layer
L
1,3
Similarity module:
Metric
cosine similarity, dot product, MinMax similarity
Similarity space dimension
512, 1024 False, True False, True
aayer-normalization
Affine
Training
Learning rate
0.0001, 0.001, 0.01 Adam, Adam 0. 0.0001 2048, 4096
Optimizer
Weight decay
Batch size
774
As similarity module, the IterRefI.STM uses the following:
a = softmax (k (m', X'))
ÿ On Un.
where ÿ is the prediction for the query molecule. For the computation of the attention values a, the softmax function is used. & is a similarity metric, such as the cosine similarity.
776
Hyperparameter search. All hyperparameters were selected based manual tuning the validation set. We report the explored hyperparameter space in Table A3 Bold values indicate selected hyperparameters for the final model.
778
780
A.1.6
MHNfs: details and hyperparameters
781
The MHNff consists of molecule encoder, the context module, the crossatttenioonnmddule and the similarity module. The molecule encoder is fully-connected Neural Network, consisting of one layer with 1024 units. For the context module, a Hopfield layer with 8 heads used and also the cross- attention module include 8 heads. We chose concatenation of ECFPs and RDKit-based descriptors as the inputs for the MHNff model. Notably, the RDKit-based descriptors were pre-processed in a way that instead of raw values quantils, which were computed by comparing distributation all PS-Mol training molecules, were used. All descriptors were normalized based on the Mol training data.
783
785
value with the
787
789
Hyperparameter search. All hyperparameters were selected based on manual tuning on the validation set. We report the explored hyperparameter space in Table A4] Bold values indicate the selected hyperparameters for the final model. Early stopping points for the different re-runs are chosen based on the AAUC-PR metric on the validation set. For the five re-runs the carly-stopping points, that were automatically chosen by their validation metries, were the checkpoints at epoch 94,
791
793
192, 253, 253 and 309.
795
Model training. Figure A2 shows the learning curve of an exemplary training run of MHNfs model on FS-Mol. The left plot shows the loss on the training set and the right plot shows the
21

---

Table A4: Hyperparameter space considered for the MHNff model selection. The hyperparameters of the best configuration are marked bold.
Hyperparameter
Explored values
Molecule encoder
Number hidden layers
0, 2,
Number of units per hidden layer
1024, 4096
Output dimension
512, 1024
Activation function
ReLU, SELU
Input dropout
0.1 0.5
Dropout
Context module (hopficld layer)
Heads
8, 16
Association space dimension
512 [512:20888
22.6 [15;40]
Dropout
0.1, 0.5
Cross-attention module (transformer mechanism)
Heads
1, 10, 16, 32, 64 567 [512: 4096]
Number units in the hidden feedforward layer
Association space dimension
1088 [512;2048] 0.1, 0.5, 0.6, 0.7 2,
Dropout
Number of layers:
Similarity module: Metric
cosine similarity, dot product, MinMax similarity
Similarity space dimension
512, 1024 False, True False, True
Layer- noyyerrormallzation
Affine
Training
Learning rate
0.0001, 0.001, 0.01 Adam, AdamW 0, 0.0001
Optimizer
Weight decay
Batch size
4096
Warm-up phase (epochs)
Constant learning rate phase (epochs)
25, 35
Decay rate
0.994
Max. number of epochs
350
validation loss. The dashed line indicates the checkpoint of the model which was saved and then used for inference on the test set, whereas the stopping point was evaluated maximizing the AAUC-PR metric on the validation set.
798
Performance improvements in comparison to naive baseline. Figure A3 shows task-wise performance comparison between MHNfs and the frequent hitter model. Fach point indicates a task in the test set and is colored according their super-class membership. In 132 cases the MHNfs outperforms the frequent hitter model. In 25 cases the frequent hitter model yields better performance.
801
803
804
A.1.7
PAR: details and hyperparameters
805
The PAR model (Wang et al. 2021) includes pre-trained GNN encoder, which creates initial embeddings for the query and support molecules. These embeddings are fed into an attention mechanism module which also uses activity information of the support set molecules to create enriched representations. Another GNN learns relations between query and support set molecules.
806
808
809
Hyperparameter search. For details we refer to (Wang al. 2021) and https://githhi:: com/ P htt.e//eeuusst////////aaaaaa/a/a
py All hyperparameters were selected based on manual tuning on the validation set. The hyperparameter choice for Tox21 (Wang et al. 202T) was used as starting point. We report the explored hyperparameter space in Table A5 Bold values
811
22

---

on training 3
1.1
Less validation
1.0
0.16
0.9
0.14
0.8
0.10
50
100
150
200
250
300
350
50
100
150 260
250
300
350
Epoch
Epoch
Figure A2: Exemplary MHNfs learning displayed and on the y-axis thee training loss (left) and the validation loss (right) The dashed line indicates the determined early-stopping point which is determined based AAUC-PR the validation set.
FS-Mol. On the x-axis the number of epochs
0.5
Oxidoreductase
Kinase
Lyase
Isoerase
0.4
Ligase
Transierase
0.3
MHNíS [AAUC-PR]
0.2
0.1
0.0
-0.1
0.11
0.0
0.1
0.2
0.3
0.4
0.5
Frequent hitter [AAUC-PR]
Performance comparison of MHNfs with the frequent hitter model. Each point refers to task the Dashed lines indicate variablility across training re-runs and different test support sets. The most points are located above the dashed line, which indicates that MHNfs performs better than den FH bascline this task.
Figure A3:
23

---

Table A5: Hyperparameter space considered for the PAR model selection. The hyperparameters of the best configuration are marked bold.
Hyperparameter
Explored values
Training
Meta learning rate
1.0 10-05 05
1.0 10 10-04,
1.0
10 03
1.0 10 02
Inner learning rate
0.01, 0.1
Update step
1,2 1,2
Update step test
Weight decay
5.0 10-10 05 1.0 10-03
Epochs
200000
Eval. steps
2000
Encoder
Use pre-trained GNN
yes, no
Attention-base module
Map dimension
128, 512
Map layer
2, 3
Pre fc layer
0,
Map dropout
0.1, 0.5
Context layer
2, 4
Relation graph
Hidden dimension
8. 122,
Number of layers
2, 2,
Number of layers for relation edge update
Batch norm
yes, no 0, 0.25, 0.5 0.2, 0.25, 0.5
Relation dropout 1
Relation dropout 2
indicate the selected hyperparameters for the final model. Notably, we just report hyperparameter choices which were different from standard choices. We used training script provided by (Wang al. 2021), which can be found here httpp:://githuu.ccmm///ttt wttp:thuu..cmm/tttaa66////Aaa. NerrA-
814
816
A.2
Domain shift experiment
817
Experimental setup. For the domain shift experiment, we used the Tox21 dataset. This dataset consists of 12, 707 chemical compounds, for which measurements for to 12 different toxic effects are reported (Mayr et al. 22066] [Hung et al. 20T6a). was published with a fixed training, validation and test split. ofaae--oftthe-ar supervised learning methods that have access to the full training set reach AUC performance values between 0.845 and 0.871 (Klambauer et al. 2017 Duvenaud et al.) Li al. 2017 2021 Zaslavskiy al. 2019: Alperstein et al. 2019). For evaluation, re-cast Tox21 as a few-shot learning setting and draw small support sets from the 12 tasks. The pre-trained FS-Mol and obtain small support sets from Tox21. Based on the support sets, the methods had to predict the activities of the Tox21 I test set. Note that there is strong domain shift from drug-liks like molecules of FS-Mol to environmental chemicals, pesticides, food additives of Tox21. The domain shift also concerns the outputs where shift from kinases, hydrolases, and oxidoreductases of FS-Mol to nuclear receptors and stress responses of Tox21 is present.
819
821
2015
823
compared methods
825
827
829
Methods compared. We compared the method MHNIs, the runner-up method IteRRefLSTM. and Similarity Search since it has been widely used for such purposes for decades (Cereto Massagué et al. 2015).
831
833
Training and evaluation. We followed the procedure of Stanley et al. 2021 for data-cleaning. preprocessing and extraction of the fingerprints and descriptors used in F--Mol. After running the cleanup step, 8,423 molecules remained for the domain shift experiments. From the training set, active and 8 inactive molecules per task were randomly selected to build the support set. The test set molecules were used as query molecules. The validation set molecules were used at all. During test-time, support set drawn ten times for each task. Then, the performance of the models were
835
837
24

---

Table A6: Results of the domain shift experiment Tox21 [AUC, AAUC-PRJ. The best method marked bold. Error bars represent standard deviation across training re-runs and draws of support sets
Method
AUC
AAUC-PR
Similarity Search (baseline)
.629 + .015
.061
.008
IterRefLSTM (Altae-Tran et al. 2017
.664 + .018
.067 = .008
MHNfs (ours)
.679 = .018
.073 .008
Table A7: Results of the ablation study on FS-Mol [AUC, AAUC-PR J. The error bars represent standard deviation across training re-runs and draws of support sets. The p-values indicate whether the difference between two models in consecutive rows significant.
Method
AUC
AAUC-PR
PAUC
PAAUC-PR
MHNís (CM+CAM+SM)
.739 + .005
.241 + .006
MHNfs- -CM
.737 + .004
.240
+ .005
0.030
0.002
MHNfs -CM -CAM
.719 + .006
.223
+ .006
1.0e-8
<1.0e-8
Similarity Search
.604 + .003
.113 = .004
<1.0c-8
1.0c-8
IterRefLSTM (Altae-Tran et al. 2017
.730 + .005
.234 + .005
<1.0e-8
8.73e-7
paired Wilcoxon rank sum test
b IteRRefLSTM is compared to MHNfs -CM
evaluated for these support sets, using the under precision-recall curve (AUC-PR), analogously the FS-Mol benchmarking experiment reported as the difference to random classifier (AAUC-PR). and the under receiver operating characteristic curve (AUC) metrics. The performance values report the mean over all combinations regarding the training reruns and the support set sampling iterations. Error bars indicate the standard deviation.
840
842
843
844
The Hopfield-basee context retrieval method has significantly outperformed the IterRefLSTM-basse model (PaAuC-PR-value 3. Pauc-value 2.5e-6, paired Wilcoxon test) and the Classic Similarity Search (PSAUC-PR-value 2.4e-9, PAuc-value 7.6e-10, paired Wilcoxon test) and therefore showed robust performance on the toxicity domain, see Table A6 Notably, all models were trained on the FS-Mol dataset and then applied to the Tox21 dataset without adjusting any weight parameter.
Results.
846
848
A.3
Details the ablation study
851
The MHNfs has two new main elements compared to the previous state-of-the art method Iter- RefLSTM, which are the context module and the coos-attention-moddee we aim to investigate i) the importance of all design elements, which are the context module, the cross-attention module, and the similarity module, and ii) the superiority of the cross-attention module compared to the HerRefl.STM module.
In this ablation study
853
855
856
A.3.1
Ablation study A: comparison against IterRefLSTM
For fair comparison between the cross-attention module and the IterRefLSTM we used a pruned MHN version ("MHNfs -CM") which has no context module and compared it with the IterReiLSTM model. The evaluation includes five training re-runs each and different support set samplings. The results, reported as the mean across training re-runs and support sets, can be seen in Table[A We performed paired Wilcoxon rank sum test for both the AUC and the AAUC-PR metrie. Both values indicate high significance.
858
859
860
861
862
863
A.3.2
Ablation study all design elements
evaluate the performance of all main elements within the MHNfs, which are context module, 864
the cross-attention module, the similarity module and the molecule encoder. For this analysis, we start with the complete MHNfs which includes all modules and report AUC and AUUC-PR performance values. Then, iteratively omit the individual modules, measuring whether there is
886
25

---

Table A8: Results of the ablation study Tox21 [AUC, AAUC-PR J. The error bars represent standard deviation across training re-runs and draws of support sets. The p-values indicate whether model is significantly different the MHNff in terms of the AUC and AAUC-PR metric.
Method
AUC
AAUC-PR
PAUC
PAAUC-PR
MHNfs (CM+CAM4SM)
.679 + .018
.073 + .008
MHNffs -CM
.622 = .028
.069 + .012
6.28e-8
0.002
MHNfs -CM -CAM
.640 + .018
.057 + .009
<1.0e-8
<1.0e-8
Similarity Search
.699 + .015
.061
.008
<1.0e-8
<1.0e-8
[terRefl.STT STM
.664 + .018
.067 + .008
2.53e-6
3.38e-5
paired Wilcoxon rank sum test
significant performance difference with and without the module. Table A7 shows the results, where performance values for the full MHNfs, a MHNfs model without the context module ("MHNfs -CM") and MHNff module without the context and the cross-attenion module ("MHNfs -CM -CAM") is included. Notably, the model without the context module and without the cross-attention module just consists of learned molecule encoder and the similarity module. We evaluted the impact of the learned molecule encoder by replacing it with a fixed encoder, which maps a molecule to its descriptors. The model with the fixed encoder is classic chemoinformatics method which is called Similarity Search (Cereto-Massagee et al. 2015).
869
871
873
875
For the evaluation, we performed five training re-runs for every model and sampled ten different support sets for every task. Table shows the results in terms of AUC and AAUC-PR. We performed paired Wilcoxon rank sum tests both metrics, comparing two methods in consecutive rows in the table. The table shows that every module has a significant impact as omitting a module results in significant performance drop. The comparison between the MHNfs version without the context module and without the cross-attention module with the Similarity Search showed a significant superiority of the learned molecule encoder in comparison the fixed encoder.
877
879
880
881
882
A.3.3
Ablation study C: Under domain shift on Tox21
Referring to Section A.3.2 the context module and the cross-ates rross-attention module showed their impor- lance for the global architecture. This importance gets even more pronounced for the domain shift experiment Tox21 as one can see in Table
885
886
887
Again, five training re-runs and ten support set draws are used for evaluation. Including the context module makes clear and significant difference for both metrics AUC and AAUC-PR.
889
A.4
Generalization to different support set sizes
890
In this section, we test the ability of MHNfs to generalize to different support set sizes. During training in the FS-Mol benchmarking setting, the MHNfs model has access to support sets of size 16. However, at inference, the support set size might be different. Figure A4 provides performance estimates of the support-set-size-16 MHNfs models on other support set sizes. Note that the estimates could be seen as approximate lower bounds of the predictive performance on settings with different support set sizes (y-axis labels). For model used in production or in a real-world drug discovery setting, MHNfs should be trained with varying support set sizes that resemble the distribution of real drug discovery projects.
881
882
883
884
895
897
898
A.5
Generalization to different context sets
899
In this section, we test the ability MHNfs to generalize to different context sets. While the FS-Mol training split is used as context during training, we assessed whether our model is robust to different context sets for inference. To this end we preprocessed the GEOM dataset (Axelrod and Gomez Bombarelli 2022) from which we used 100,000 molecules that passed all pre-processing checks. From this set, we sample 10,000 molecules as context set for MHNfs. Because GEOM contains drug- dike molecules, similar to FS-Mol the predictive performance remains stable (see Table A9).
900
901
903
26

---

0.28
0.24
MAUC-PR (tast se0)
0.22
0.20
0.18
0.14
12 14 20
32
64
128
(inference
Figure A4:
Performance of MHNís for different support set sizes during inference time, The MHNís models are trained with support sets of the size 16.
Table A9: MHNIs performance for different context sets [AAUC-PR
J. The error bars represent standard deviation across training re-runs and draws of support sets.
Dataset used as a context
AAUC-PR
FS-Mol (Stanley et al. 2021)
.2414 + .006 .2415 + .005
GEOM (Axelrod and Gomez Bombarell
2022
905
A.6
Details and insights on the context module
906
The context module replaces the initial representations of query and support set molecules by retrieval from the context set. The context
907
large of molecules and covers large chemical The context module learns how to replace the initial molecule embeddings such that the context-enriched representations are put in relation to this large chemical space and still contains all necessary information for the similarity-based prediction part. Figure A5 shows the effect of the context module for the MHNfs model. Extreme initial embeddings, such the purple embedding on the right, are pulled more into the known chemical space, represented by the context molecules. Notably, the replacement described above is soft replacement, because also the initial embeddings contribute the context-enriched representations due to skip-connections.
908
space.
909
911
913
915
A.7
Reinforcing the covariance structure in the data using modern Hopfield networks
We follow the argumentation of (Fürst al. 202T Theorem A3) that retrieval from associative memory of a MHN reinforces the covariance structure.
917
Let us assume that we have one molecule embedding from the query set E [ and one molecule embedding from the support set æ R&d and both have been enriched with the context module with memory C dx (ignoring linear mappings):
919
m' softmax(33* m)
(A9)
z' = C softmax(3C" x)
(A10)
921 Then the similarity of the retrieved representations as measured by the dot product can expressed in terms of covariances:
27

---

Context Initial Embedding
Context-enriched Representation
10
pcl
-10
-15
-20
-20
-10
20
pcO
Figure A5: PCA plot of molecule embeddings. Each dot the plot represents molecule embedding. of which the first two principal components are displayed on the x- and y-axis. Blue dots represent context molecules. Dark purple dots represent initial embeddings for some exemplary molecules, of which some exhibit extreme characteristics and are thus located away from the center. Arrows and light purple dots represent the enriched molecule embeddings after the retrieval step. Especially molecules from extreme positions are moved stronger to the center and thus are more similar to known molecules after retrieval.
n x' = softmax($C* m)" 23
(A11)
(c Cov(C, m)" m)" (c+ Cov(C,x)x.)
(A12)
923 where è the row mean of and following the weighted covariances are used:
Cov(C,m (3Cm)))
Cov(C, (BCx))))
(A13)
JmD RM MMX is mean Jacobian function of the softmax (Fürst et al.
2021 (A172))
The Jacobian J of p softmax(Ba) is J(ßa) 3 (diag(p) ppp").
6t bF b 6T (diag(p) pPd) b 3 Epe by
(Epu by
(A14)
this is the second moment minus the mean squared, which is the variance. Therefore, bt J=I(a 8 927 times the covariance of if component is drawn with probability ps the multinomial distribution
p. In case the component is context sample c, Jm is the average of (()a) over À to X 8.
28

---

929 Note that we can express the enriched representations using these covariance functions:
m' = (c + m))(C.m))mm x' (C + Cov(C, x)
(A15)
(A16)
930 which connects retrieval from MHNs with reinforcing the covariance structure of the data.
29


Reviews:
[{"review_id": "qd-dNrtV0or", "reviewer": "Reviewer_ZhVh", "summary": "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. ", "questions": "- (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. \n\n- (Q2): Is there any intuition behind the label balancing strategy in Section 3.4? \n\n- (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. \n\n- (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. \n\n- (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?", "limitations": "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. ", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "=== Strengths === \n\n- (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. \n\n- (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). \n\n- (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. \n\n \n\n=== Weaknesses === \n\n- (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). \n\n- (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.\n\n \n\n=== Summary === \n\nOverall, 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. \n\n \n\n=== Other comments === \n\nHere I include various other comments, which had only a small impact on my score. \n\n- (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. \n\n- (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. \n\n \n\n=== Nitpicks === \n\nHere I include some final nitpicks, which did not affect my score; they are here just to help improve the paper.\n\n- Line 92: \"an domain shift experiment\" - \"an\" -> \"a\" \n\n- Line 108: \"predicts y for a x\" - \"a\" -> \"an\"; \"the model that generalizes\" - I'd say \"a model\" \n\n- Missing space in lines 149 and 237\n\n- Line 241: \"wich\" -> \"which\" \n\n- Line 254: \"dataset on related task and then adapting\" -> \"dataset of related tasks and then adapt\" \n\n- Line 262: \"support support\" \n\n- 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) \n\n- Line 346: \"material science dataset\" -> \"material science datasets\" \n\n- Line 363: \"This, might\" - drop the comma \n\n- Line 647 (appendix): \"fingerprints Morgan\" -> \"Morgan fingerprints\" \n\n \n\n=== References === \n\n[1] Stanley et al, \"FS-Mol: A Few-Shot Learning Dataset of Molecules\" \n\n[2] Chen et al, \"Meta-learning Adaptive Deep Kernel Gaussian Processes for Molecular Property Prediction\" ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "EtAvz-y2OA", "reviewer": "Reviewer_LC2d", "summary": "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. \nThe authors conduct experiments on FS-Mol dataset.\n\n", "questions": "Please reply to my concerns in weakness. \n\nKey questions: \n- There exists a number of works using external memory, such as [b1,b2]. Why consider Modern Hopfield Networks instead? Any strength?\n- Please provide discussion and empirical comparison w.r.t [a1,a2]. \n\n[b1] Meta-learning with memory-augmented neural networks, ICML, 2017\n[b2] Few-shot visual learning with contextual memory and fine-grained calibration, IJCAI, 2020\n", "limitations": "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.  ", "rating": 4, "confidence": 4, "soundness": 1, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength\n\n+ Figure 1 is nicely plotted, which conveys the flow of the proposed MHNfs.\n+ The authors provide concise reproduction details in appendix. \n\n\nWeakness\n\n- Motivation: The authors did not provide convincing motivations or necessity for using Modern Hopfield Networks. \n- 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. \n- 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]. \n- 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. \n\n[a1] Few-Shot Graph Learning for Molecular Property Prediction, WWW, 2021\n[a2] Property-Aware Relation Networks for Few-Shot Molecular Property Prediction, NeurIPS, 2021\n[a3] Low Data Drug Discovery with One-Shot Learning, ACS Cent. Sci., 2017\n\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "EFplEJ9v71S", "reviewer": "Reviewer_2LtU", "summary": "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.", "questions": "**Major concerns**\n\n- 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. \n\n- 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.\n\n- In the main text, line 82-84, \"Those retrieved representations have amplified co-occurrences and covariance\n83 structures, while peculiarities and spurious co-occurrences of the query and support set molecules are\n84 averaged out.\" Is it possible to illustrate this with an example, or somehow support the claim in your case?\n\n- 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.\n\n- 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.\n\n\n**Minor concerns**\n- In the main text, line 79, \"as an associative memory\".\n\n- In the main text, line 92, \"and a domain shift experiment\".\n\n- 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*.\n\n- 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?\n\n- 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.\n\n\n\n", "limitations": "- 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.\n\n- 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.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n[*originality*] The proposed model architecture is novel for the few-shot learning task, compared to existing models listed in Table1.\n\n[*quality*] The paper is well written with a clear story line and very few typos or grammar issues.\n\n[*clarity*] The key idea is clearly explained and model architecture is accurately described.\n\n[*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.\n\n$\\newline$\n\n**Weakness**\n\n[*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. \n\n[*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.\n\n[*clarity*] No major concerns.\n\n[*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?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Reject