Original: Context-enriched molecule representations improve few-shot drug discovery
Revised: Context-enriched molecule representations improve few-shot drug discovery

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AfﬁliationAddress email
Revised: AfﬁliationAddress email

Original: Abstract
Revised: Abstract

Original: A central task in computational drug discovery is to construct models from knownactive molecules to ﬁnd further promising molecules for subsequent screening.
Revised: A central task in computational drug discovery is to construct models from knownactive molecules to ﬁnd further promising molecules for subsequent screening.

Original: However, typically only very few active molecules are known.
Revised: However, typically only very few active molecules are known.

Original: Therefore, few-shotlearning methods have the potential to improve the effectiveness of this criticalphase of the drug discovery process.
Revised: Therefore, few-shotlearning methods have the potential to improve the effectiveness of this criticalphase of the drug discovery process.

Original: We introduce a new method for few-shotdrug discovery.
Revised: We introduce a new method for few-shotdrug discovery.

Original: Its main idea is to enrich a molecule representation by knowledgeabout known context or reference molecules.
Revised: Its main idea is to enrich a molecule representation by knowledgeabout known context or reference molecules.

Original: Our novel concept for moleculerepresentation enrichment is to associate molecules from both the support set andthe query set with a large set of reference (context) molecules through a modern
Revised: Our novel concept for moleculerepresentation enrichment is to associate molecules from both the support set andthe query set with a large set of reference (context) molecules through a modern

Original: Hopﬁeld network.
Revised: Hopﬁeld network.

Original: Intuitively, this enrichment step is analogous to a human expertwho would associate a given molecule with familiar molecules whose propertiesare known.
Revised: Intuitively, this enrichment step is analogous to a human expertwho would associate a given molecule with familiar molecules whose propertiesare known.

Original: The enrichment step reinforces and ampliﬁes the covariance structureof the data, while simultaneously removing spurious correlations arising from thedecoration of molecules.
Revised: The enrichment step reinforces and ampliﬁes the covariance structureof the data and simultaneously removes spurious correlations arising from thedecoration of molecules.

Original: Our approach is compared with other few-shot methodsfor drug discovery onthe FS-Mol benchmark dataset. On FS-Mol, our approachoutperforms all compared methods and therefore sets a new state-of-the art forfew-shot learning in drug discovery.
Revised: We analyze our novel method on FS-Mol, which isthe only established few-shot learning benchmark dataset for drug discovery.

Original: An ablation study shows that the enrichmentstep of our method isthe key to improve the predictive quality.
Revised: Anablation study shows that the enrichment step of our method is key to improvingthe predictive quality.

Original: In a domain shiftexperiment, we further demonstrate the robustness of our method.
Revised: In a domain shift experiment, our new method is more robustthan other methods.

Revised: On FS-Mol, our new method achieves a new state-of-the-artand outperforms all other few-shot methods.

Original:  Introduction
Revised:  Introduction

Original: To improve human health, combat diseases, and tackle pandemics there is a steady need of discoveringnew drugs in a fast and efﬁcient way.
Revised: To improve human health, combat diseases and tackle pandemics, there is a steady need of discoveringnew drugs in a fast and efﬁcient way.

Original: However, the drug discovery process is time-consuming andcost-intensive (Arrowsmith, 2011).
Revised: The drug discovery process is time-consuming and cost-intensive(Arrowsmith, 2011).

Original: Deep learning methods have recently been shown to reducetime and costs of this process (Chen et al., 2018; Walters and Barzilay, 2021).
Revised: Deep Learning methods have recently been shown to reduce time and monetaryresources (Chen et al., 2018; Walters and Barzilay, 2021).

Original: They diminish therequired number of both wet-lab measurements and molecules that must be synthesized (Merk et al. Schneider et al., 2020).
Revised: They diminish the required number of bothwet-lab measurements and molecules that must be synthesized (Merk et al., 2018; Schneider et al.,2020).

Original: However, as of now, deep learning approaches use only the molecularinformation about the ligands after being trained on a large training set.
Revised: However, they use only the molecular information about the ligands after being trained on alarge training set.

Original: At inference time, they yieldhighly accurate property and activity prediction (Mayr et al., 2018; Yang et al., 2019), generative(Segler et al., 2018a; Gómez-Bombarelli et al., 2018), or synthesis models (Segler et al., 2018b; Seidlet al., 2022).
Revised: At inference time, they yield highly accurate property and activity prediction (Mayret al., 2018; Yang et al., 2019), generative (Segler et al., 2018a; Gómez-Bombarelli et al., 2018), orsynthesis models (Segler et al., 2018b; Seidl et al., 2021).

Original: Deep learning methods in drug discovery usually require large amounts of biological measurements.
Revised: Deep Learning methods in drug discovery usually require large amounts of biological measurements.

Original: To train deep learning-based activity and property prediction models with high predictive performance, hundreds or thousands of data points per task are required.
Revised: To train Deep Learning-based activity and property prediction models with high predictive performance, hundreds or thousands of data points per task are required.

Original: For example, well-performingpredictive models for activity prediction tasks of ChEMBL have been trained with an average of 3,621activity points per task, i.e., drug target, by (Mayr et al., 2018).
Revised: For example, well-performingpredictive models for activity prediction tasks of ChEMBL have been trained with an average of 3,621activity points per task, i.e. drug target, by Mayr et al.

Original: The ExCAPE-DB dataset provides onaverage 42,501 measurements per task (Sun et al., 2017; Sturm et al., 2020).
Revised: The ExCAPE-DB dataset provideson average 42,501 measurements per task (Sun et al., 2017; Sturm et al., 2020).

Original: (Wu et al., 2018) pub- 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., 2016b;
Revised: Wu et al. (2018)published a large scale benchmark for molecular machine learning, including among others prediction models for the SIDER dataset with an average of 5,187 data points, Tox21 (Huang et al., 2016b;

Original: Mayr et al., 2016) with on average 9,031, and ClinTox (Wu et al., 2018) with 1,491 measurementsper task.
Revised: Mayr et al., 2016) with on average 9,031, and ClinTox (Wu et al., 2018) with 1,491 measurementsper task.

Original: However, for typical drug design projects, the amount of available measurements is verylimited (Stanley et al., 2021; Waring et al., 2015; Hochreiter et al., 2018), since in-vitro experimentsare expensive and time-consuming.
Revised: However, for typical drug design projects, the amount of available measurements is verylimited (Stanley et al., 2021; Waring et al., 2015; Hochreiter et al., 2018), since in-vitro experimentsare expensive and time-consuming.

Original: Therefore, methods that need only few measurements to buildprecise  prediction models are desirable.
Revised: Therefore, methods that need only few measurements to buildprecise property and activity prediction models are desirable.

Original: This problem— i.e., the challenge of learning from fewdata points — is the focus of machine learning areas like meta-learning (Schmidhuber, 1987; Bengioet al., 1991; Hochreiter et al., 2001) and few-shot learning (Miller et al., 2000; Bendre et al., 2020;
Revised: This problem of learning from few datapoints is the focus of machine learning areas like meta-learning (Schmidhuber, 1987; Bengio et al. Hochreiter et al., 2001) and few-shot learning (Miller et al., 2000; Bendre et al., 2020;

Original: Wang et al., 2020).
Revised: Wanget al., 2020).

Original: Few-shot learning tackles the low-data problem that is ubiquitous in drug discovery.
Revised: Few-shot learning tackles the low-data problem that is ubiquitous in drug discovery.

Original: Few-shotlearning methods have been predominantly developed and tested on image datasets (Bendre et al.
Revised: Few-shotlearning methods have been predominantly developed and tested on image datasets (Bendre et al.

Original: ,2020; Wang et al., 2020), and have recently been adapted to drug discovery problems (Chen et al.,2022; Wang et al., 2021; Stanley et al., 2021;
Revised: ,2020; Wang et al., 2020), and have recently been adapted to drug discovery problems (Stanley et al.

Original:  Altae-Tran et al., 2017).
Revised: ,2021; Altae-Tran et al., 2017).

Original: They are usually categorizedinto three groups according to their main approach (Bendre et al., 2020; Wang et al., 2020; Adler et al.,2020).
Revised: They are usually categorized into three groups according to their mainapproach (Bendre et al., 2020; Wang et al., 2020; Adler et al., 2020).

Original: a) Data-augmentation-based approaches augment the available samples and generate new, morediverse data points (Chen et al., 2020; Zhao et al., 2019; Antoniou and Storkey, 2019).
Revised: a) Data-augmentation-basedapproaches counter the problem of too few data by augmenting the available samples and generatingnew, more diverse data points (Chen et al., 2020; Zhao et al., 2019; Antoniou and Storkey, 2019).

Original: b) Embeddingbased and nearest neighbour approaches learn embedding space representations.
Revised: b) Embedding-based and nearest neighbour approaches learn representations in a meaningful embeddingspace.

Original: Predictive modelscan then  be constructed from only few net data points by comparing these embeddings.
Revised: If the embeddings have a proper representation, then models can be constructed from onlyfew new data points by comparing these embeddings.

Original: Forexamplein Matching Networks (Vinyals et al., 2016) an attention mechanism that relies on embeddings is thebasis for the predictions.
Revised: For Matching Networks (Vinyals et al., 2016),e.g., an attention mechanism, which relies on these embeddings, is the basis for the predictions.

Original: Prototypical Networks (Snell et al., 2017) create prototype representationsfor each class using the above mentioned representations in the embedding space.
Revised: Prototypical Networks (Snell et al., 2017) create prototype representations for each class, usingthe above mentioned representations in the embedding space.

Original: c) Optimizationbased or ﬁne-tuning methods utilize a meta-optimizer that focuses on efﬁciently navigating theparameter space.
Revised: c) Optimization-based or ﬁne-tuningmethods utilize a meta-optimizer that focuses on efﬁciently navigating the parameter space.

Original: For example, with MAML the meta-optimizer learns initial weights that can beadapted to a novel task by few optimization steps (Finn et al., 2017). 
Revised: Forexample, with MAML the meta-optimizer learns initial weights that can be adapted to a noveltask by few optimization steps (Finn et al., 2017). Nguyen et al. (2020) evaluated the applicabilityof MAML and its variants to graph neural networks (GNNs) and Guo et al. (2021) also combine

Revised: GNNs and meta-learning.

Revised: Altae-Tran et al.

Original: Most of these approaches have already been applied to few-shot drug discovery (see Sec. 4).
Revised: (2017) suggested an approach called Iterative Reﬁnement Long Short-Term Memory, in which query and support set embeddings can share information andupdate their embeddings. Recently, Stanley et al. (2021) generated a benchmark dataset for few-shotlearning methods in drug discovery and provided some baseline results.

Original: Surprisingly, almost all these few-shot learning methods in drug discovery are worse than a naive baseline,which does not even use the support set (see Section 5).
Revised: Surprisingly, almost all thesefew-shot learning methods, are worse than a naive baseline, that does not even use the support set (see Section 4).

Original: We hypothesize that the under-performanceof these methods stems from disregarding the context — both in terms of similar molecules andsimilar activities.
Revised: We hypothesize that this under-performance of these methods stems from disregardingthe context, both in terms of similar molecules and similar activities.

Original: Therefore, we propose a method that informs the representations of the query andsupport set with a large number of context molecules covering the chemical space.
Revised: Therefore, we propose a methodthat informs the representations of the query and support set with a large number of context moleculescovering the chemical space.

Original: Enriching molecule representations with context using associative memories.
Revised: Enriching molecule representations with context using associative memories.

Original: In data-scarcesituations, humans extract co-occurrences and covariances by associating current perceptions withmemories (Bonner and Epstein, 2021; Potter, 2012).
Revised: Especially insituations with few data, humans extract co-occurrences and covariances by associating currentperceptions with memories (Bonner and Epstein, 2021; Potter, 2012).

Original: When we show a small set of active molecules toa human expert in drug discovery, the expert associates them with known molecules  to suggest furtheractive molecules (Gomez, 2018;
Revised: In drug discovery, a humanexpert, when shown a small set of active molecules, also associates them with known molecules andtheir properties to suggest further active molecules (Gomez, 2018;

Original: He et al., 2021).
Revised: He et al., 2021).

Original: In an analogous manner, our novel concept forfew-shot learning uses associative memories to extract co-occurrences and the covariance structureof the original data and to amplify them in the representations (Fürst et al., 2021).
Revised: In an analogousmanner, our novel concept for few-shot learning uses associative memories to extract co-occurrencesand the covariance structure of the original data and to amplify them in the representations (Fürstet al., 2021).

Original: We use Modern Hopﬁeld Networks (MHNs) as an associative memory, since they can store a large set of contextmolecule representations (Ramsauer et al., 2021, Theorem 3).
Revised: We use Modern Hopﬁeld Networks (MHNs) as a associative memory, since theycan store a large set of context molecule representations (Ramsauer et al., 2020, Theorem 3).

Original: The representations that are retrievedfrom the MHNs replace the original representations of the query and support set molecules.
Revised: Therepresentations that are retrieved from the MHNs replace the original representations of the query andsupport set molecules.

Original: Thoseretrieved representations have ampliﬁed co-occurrences and covariance structures, while peculiaritiesand spurious co-occurrences of the query and support set molecules are averaged out.
Revised: Those retrieved representations have ampliﬁed co-occurrences and covariancestructures, while peculiarities and spurious co-occurrences of the query and support set molecules areaveraged out.

Original: In this work, our contributions are the following:
Revised: In this work, our contributions are the following:

Original: • We propose a new architecture MHNfs for few-shot learning in drug discovery.•
Revised: • We propose a new architecture MHNfs for few-shot learning in drug discovery.•

Original: We achieve a new state-of-the-art on the benchmarking dataset
Revised: We achieved a new state-of-the-art on the benchmarking dataset

Original: FS-Mol.• We introduce a novel concept to enrich the molecule representations with context by associating them with a large set of context molecules.
Revised: FS-Mol.• We introduce a novel concept to enrich the molecule representations with context by associating them with a large set of context molecules.•

Original: • We add a naive baseline to the FS-Mol benchmark that yields better results than almost all other published few-shot learning methods.
Revised: We add a naive baseline to the FS-Mol benchmark that yields better results than almost allother published few-shot learning methods.

Original: • We provide results of an ablation study and a domain shift experiment to further demonstrate the effectiveness of our new method.
Revised: • We provide results of an ablation study and an domain shift experiment to further demonstratethe effectiveness of our new method.

Original:  Problem setting
Revised:  Problem setting

Original: Drug discovery projects revolve around models g ( m ) that can predict a molecular property or activityˆ y, given a representation m of an input molecule from a chemical space M .
Revised: Drug discovery projects revolve around models g ( m ) that can predict a molecular property oractivity ˆ y given a representation m of an input molecule from a chemical space M .

Original: We consider machinelearning models ˆ y =
Revised: We considermachine-learning models ˆ y =

Original: g w ( m ) with parameters w that have been selected using a training set.
Revised: g w ( m ) with parameters w that have been selected using a trainingset.

Original: Typically,deep learning based property prediction uses a molecule encoder f ME : M → R d .
Revised: Typically, Deep Learning based property prediction uses a molecule encoder f ME : M → R d .

Original: The moleculeencoder can process different symbolic or low-level representations of molecules, such as moleculardescriptors (Bender et al., 2004;
Revised: The molecule encoder can process different symbolic or low-level representations of molecules, suchas molecular descriptors (Bender et al., 2004;

Original: Unterthiner et al., 2014; Mayr et al., 2016), SMILES (Weininger,1988; Mayr et al., 2018; 
Revised: Unterthiner et al., 2014; Mayr et al., 2016), SMILES(Weininger, 1988; Mayr et al., 2018; Winter et al., 2019;

Original: Winter et al., 2019; Segler et al., 2018a), or molecular graphs (Merkwirthand Lengauer, 2005; Kearnes et al., 2016;
Revised: Segler et al., 2018a), or molecular graphs(Merkwirth and Lengauer, 2005; Kearnes et al., 2016;

Original: Yang et al., 2019; Jiang et al., 2021) and can be pre-trainedon related property prediction tasks.
Revised: Yang et al., 2019; Jiang et al., 2021) and canbe pre-trained on related property prediction tasks.

Original: For few-shot learning, the goal is to select a high-quality predictive model based on a small set ofmolecules  { x 1 , . . .
Revised: For few-shot learning, the goal is to select a high-quality predictive model based on a small set ofmolecules X = { x 1 , . . .

Original: , x N } with associated measurements y = { y 1 , . .
Revised: , x N } with associated measurements y = { y 1 , . .

Original: . , y N } .
Revised: . , y N } .

Original: The measurements areusually assumed to be binary y n ∈ {− 1 , 1 } , corresponding to the molecule being inactive or active.
Revised: The measurementsare usually assumed to be binary y n ∈ {− 1 , 1 } , corresponding to the molecule being inactive oractive.

Original: The set { ( x n , yn ) }Nn =1 is called the support set that contains samples from a prediction task and N isthe support set size .
Revised: The set { X , y } is called the support set that contains samples from a prediction task.

Original: The goal is to construct a model that correctly predicts y for an x that is notinthe support set — in other words, a model that generalizes well.
Revised: Goal isto construct a model that correctly predicts y for a x that is not in the support set, which means themodel that generalizes well.

Original: Standard supervised machine learning approaches typically just show limited predictive power atthis task (Stanley et al., 2021) since they tend to overﬁt on the support set due to a small number oftraining samples.
Revised: Standard supervised machine learning approaches fail at this task (Stanley et al., 2021).

Original: These approaches learn the parameters w of the model g w from the support set ina supervised manner.
Revised: Theseapproaches learn the parameters w of the model g w from the support set in a supervised manner.

Original: However, they heavily overﬁt to the support set when N is small.
Revised: However, they heavily overﬁt to the support set when N is small.

Original: Therefore,few-shot learning methods are necessary to construct models from the support set that generalizewell to new data.
Revised: Therefore, few-shot learningmethods are necessary to construct models from the support set that generalize well to new data.

Original:  MHNfs: Hopﬁeld-basedmolecular context enrichment for few-shot drugdiscovery
Revised:  MHNfs: Hopﬁeld-based

Revised: molecular context enrichment for few-shot drug

Original: We aim at increasing the generalization capabilities of few-shot learning methods in drug discoveryby enriching the molecule representations with molecular context.
Revised: We aim at increasing the generalization capabilities of few-shot learning methods in drug discoveryby enriching the molecule representations with molecular context.

Original: In comparison to the support set,which encodes information about the task, the context set – i.e. a large set of molecules – includesinformation about a large chemical space. The query andthe support set molecules perform a retrievalfrom the context set and thereby enrich their representations.
Revised: This context, i.e. a large set ofmolecules that covers the chemical space, is provided to the model within a context module, in whichthe query and support set molecules perform a retrieval from the context set and thereby enrich theirrepresentations.

Original: We detail this in the following.
Revised: We detail this in the following.

Original: 3.1 Model architecture
Revised: 3.1 Model architecture

Original: We propose an architecture which consists of three consecutive modules.
Revised: We propose to include the molecular context via three consecutive modules that feed into each other.

Original: The ﬁrst module, a) thecontext module f CM , enriches molecule representations by retrieving from a large set of molecules.
Revised: The ﬁrst module, a) the context module f CM , enriches molecule representations by retrieving from alarge set of molecules.

Original: The second module, b) the cross-attention module f CAM (Hou et al., 2019; Chen et al., 2021), enablesthe effective exchange of information between the query molecule and the support set molecules.
Revised: The second module, b) the cross-attention module f CAM (Hou et al., 2019; Chen et al., 2021), enables the effective exchange of information between the query molecule and thesupport set molecules.

Original: Finally the prediction for the query molecule is computed by using the usual c) similarity modulef SM (Koch et al., 2015; Altae-Tran et al., 2017):
Revised: Finally the prediction for the query molecule is computed by using the usualc) similarity module f SM (Koch et al., 2015; Altae-Tran et al. , 2017):context module: [Equation] cross-attention module

Original: [Equation]
Revised: [Equation]

Revised: similarity module:

Original: [Equation]
Revised: [Equation]

Original: where m ∈ R
Revised: where m ∈ R

Original: d is a molecule embedding from a trainable or ﬁxed molecule encoder, and m (cid:48) and m (cid:48)(cid:48)are enriched versions of it.
Revised: d is a molecule embedding from a trainable or ﬁxed molecule encoder, and m (cid:48) and m (cid:48)(cid:48)are enriched versions of it.

Original: Similarly, X ∈ R d × N contains the stacked embeddings of the support setmolecules and X (cid:48) and X (cid:48)(cid:48) are their enriched versions.
Revised: Similarly, X ∈ R d × N contains the stacked embeddings of the support setmolecules and X (cid:48) and X (cid:48)(cid:48) are their enriched versions.

Original: C ∈ R d × M is a large set of stacked moleculeembeddings,y are the support set labels, and ˆ y is the prediction for the query molecule. Squarebrackets indicate concatenation, for example [ m (cid:48) , X (cid:48) ] is a matrix with N + 1 columns. The modulesf CM , f CAM , and f SM are detailed inthe paragraphs below. An overview of our architecture is givenin Figure 1.
Revised: C ∈ R d × M is a large set of stacked moleculeembeddings, in Figure 1.

Original: The architecture also includes skip connections bypassing f CM ( ., . )
Revised: The architecture also includes skip connections bypassing f CM ( ., . )

Original: and f CAM ( . ) andlayer normalization (Ba et al., 2016), which are not shown in Figure1.
Revised: and f CAM ( . ) , whichare not shown in the ﬁgure.

Original: A shared molecule encoder f
Revised: A shared molecule encoder f

Original: ME creates embeddings for the query molecule m =
Revised: ME creates embeddings for the query molecule m =

Original: f ME ( m ) , thesupport set molecules x n = f ME ( x n ) , and the context molecules
Revised: f ME ( m ) , thesupport set molecules x n = f ME ( x n ) , and the context molecules

Original: c m =
Revised: c m =

Original: f ME ( c m ) .
Revised: f ME ( c m ) .

Original: There are manypossible choices for ﬁxed or adaptive molecule encoders (see Section 2), of which we use descriptorbased fully-connected networks because of their computational efﬁciency and good accuracy (Dahlet al., 2014; Mayr et al., 2016, 2018).
Revised: There are manypossible choices for ﬁxed or adaptive molecule encoders (see Section 2), of which we use descriptorbased fully-connected networks because of their computational efﬁciency and good accuracy (Dahlet al., 2014; Mayr et al., 2016, 2018).

Original: For notational clarity we denote the course of the representationsthrough the architecture:
Revised: For notational clarity we denote the course of the representationsthrough the architecture:

Original: [Equation]
Revised: [Equation]

Original: 3.2 Context module (CM)
Revised: 3.2 Context module (CM)

Original: The context module associates the query and support set molecules with a large set of contextmolecules, and represents them as weighted average of context molecule embeddings.
Revised: The context module associates the query and support set molecules with a large set of contextmolecules, and represents them as weighted average of context molecule embeddings.

Original: The contextmodule is realised by acontinuous Modern Hopﬁeld Network (MHN) (Ramsauer et al., 2021).
Revised: The contextmodule is realised by a Modern Hopﬁeld Network (MHN) (Ramsauer et al., 2020).

Original: An MHN is a content-addressable associative memory which can be built into deep learning architectures. There exists an analogy between the energy update of MHNs and the attention mechanism of
Revised: An MHN is anassociative memory which is capable of storing and retrieving patterns from a memory M ∈ R

Original: Transformers (Vaswani et al., 2017; Ramsauer et al., 2021). MHNs are capable of storing andretrieving patterns from a memory M ∈ R e
Revised: e

Original: × M given a state pattern ξ ∈ R e that represents the query.
Revised: × Mgiven a state pattern ξ ∈ R e that represents the query.

Original: The retrieved pattern ξ new∈ R e is obtained by
Revised: The retrieved pattern ξ new is obtained by

Original: [Equation]
Revised: [Equation]

Original: 6) where p is called the vector of associations and β is a scaling factor or inverse temperature.
Revised: where p is called the vector of associations and β is a scaling factor or inverse temperature.

Original: Modern
Revised: Modern

Original: Hopﬁeld Networks have been successfully applied to chemistry and computational immunology(Seidl et al., 2022; Widrich et al., 2020).
Revised: Hopﬁeld Networks have been successfully applied to chemistry and computational immunology(Seidlet al., 2021; Widrich et al., 2020).

Original: We use this mechanism in the form of a Hopﬁeld layer , which ﬁrst maps raw patterns to an associativespace using linear transformations, and uses multiple simultaneous queries Ξ ∈ R d × N :
Revised: We use this mechanism in the form of a Hopﬁeld layer , which ﬁrst maps raw patterns to an associativespace using linear transformations, and uses multiple simultaneous queries Ξ ∈ R d × N :

Original: [Equation]
Revised: [Equation]

Original: where W E ∈ R d × d and W C , W Ξ ∈ R
Revised: where W E ∈ R d × d and W C , W Ξ ∈ R

Original: e × d are trainable parameters of the Hopﬁeld layer, softmaxis applied column-wise, and β is a hyperparameter.
Revised: e × d are trainable parameters of the Hopﬁeld layer, softmaxis applied column-wise, and β is a hyperparameter.

Original: Note that in principle the Ξ and C could have adifferent second dimension as long as the linear transformations map to the same dimension e .
Revised: Note that in principle the Ξ and C could have adifferent second dimension as long as the linear transformations map to the same dimension e .

Original: Notethat all embeddings that enter this module are ﬁrst layer normalized (Ba et al., 2016).
Revised: Notethat all embeddings that enter this module are ﬁrst layer normalized (Ba et al., 2016).

Original: Several ofthese Hopﬁeld layers can run in parallel and we refer to them as "heads" in analogy to Transformers(Vaswani et al., 2017). 
Revised: Several ofthese Hopﬁeld layers can run in parallel and we refer to them as "heads" in analogy to Transformers(Vaswani et al., 2017). From Eq.

Revised: 7, Transformers arise as special cases of these layers (see below).

Original: The context module of our new architecture uses a Hopﬁeld layer, where the query patterns are theembeddings of the query molecule m and the support set molecules X .
Revised: The context module of our new architecture uses a Hopﬁeld layer, where the query patterns are theembeddings of the query molecule m and the support set molecules X .

Original: The memory is composed ofembeddings of a large set of M molecules from a chemical space, for example reference molecules,here called context molecules C .
Revised: The memory is composed ofembeddings of a large set of M molecules from a chemical space, for example reference molecules,here called context molecules C .

Original: Then the original embeddings m and X are replaced by theretrieved embeddings, which are weighted averages of context molecule embeddings:
Revised: Then the original embeddings m and X are replaced by theretrieved embeddings, which are weighted averages of context molecule embeddings:

Original: [Equation]
Revised: [Equation]

Original: This retrieval step reinforces the covariance structure of the retrieved representations (see Appendix A.7). Note that the embeddings of the query  and the support set molecules have not yetinﬂuenced each other.
Revised: Note that the embeddings of the query molecule and the support set molecules have not yet inﬂuencedeach other.

Original: These updated representations m (cid:48) , X (cid:48) are passed to the cross-attention module.
Revised: These updated representations m (cid:48) , X (cid:48) are passed to the cross-attention module.

Original: 3.3 Cross-attention module (CAM)
Revised: 3.3 Cross-attention module (CAM)

Original: For embedding-based few-shot learning methods in the ﬁeld of drug discovery, Altae-Tran et al.
Revised: For embedding-based few-shot learning methods in the ﬁeld of drug discovery, Altae-Tran et al.

Original: (2017)showed that the representations of the molecules can be enriched, if the architecture allows informationexchange between query and support set molecules.
Revised: (2017)showed that the representations of the molecules can be enriched, if the architecture allows informationexchange between query and support set molecules.

Original: Altae-Tran et al.
Revised: Altae-Tran et al.

Original: (2017) uses an attentionenhanced LSTM variant which updates the query and the support set molecule representations in aniterative fashion, being aware of each other.
Revised: (2017) uses an attentionenhanced LSTM variant which updates the query and the support set molecule representations in aniterative fashion, being aware of each other.

Original: We further develop this idea and combine it with the ideaof using a transformer encoder layer (Vaswani et al., 2017) as a cross-attention module (Hou et al.
Revised: We further develop this idea and combine it with the ideaof using a transformer encoder layer (Vaswani et al., 2017) as a cross-attention module (Hou et al.

Original: ,2019; Chen et al., 2021).
Revised: ,2019; Chen et al., 2021).

Original: The cross-attention module updates the query molecule representation m (cid:48) and the support setmolecule representations X (cid:48) by mutually exchanging information, using the usual Transformermechanism:
Revised: The cross-attention module updates the query molecule representation m (cid:48) and the support setmolecule representations X (cid:48) by mutually exchanging information, using the usual Transformermechanism:

Original: [Equation]
Revised: [Equation]

Original: ) where [ m (cid:48) , X (cid:48) ]
Revised: where [ m (cid:48) , X (cid:48) ]

Original: ∈ R d × ( N +1) is the concatenation of the representations of the query molecule m (cid:48)with the support set molecules X (cid:48) and we exploited that the Transformer is a special case of the
Revised: ∈ R d × ( N +1) is the concatenation of the representations of the query molecule m (cid:48)with the support set molecules X (cid:48) and we exploited that the Transformer is a special case of the

Original: Hopﬁeld layer.
Revised: Hopﬁeld layer.

Original: Again, normalization is applied (Ba et al., 2016) and multiple Hopﬁeld layers, i.e.,heads, can run in parallel,  be stacked, and equipped with skip-connections.
Revised: Again, normalization is applied (Ba et al., 2016) and multiple Hopﬁeld layers, i.e.heads, can run in parallel, can be stacked and equipped with skip-connections.

Original: The representations m (cid:48)(cid:48)and X (cid:48)(cid:48) are passed to the similarity module.
Revised: The representationsm (cid:48)(cid:48) and X (cid:48)(cid:48) are passed to the similarity module.

Original: 3.4 Similarity module (SM)
Revised: 3.4 Similarity module (SM)

Original: In this module, pairwise similarity values k ( m (cid:48)(cid:48) , x (cid:48)(cid:48) n ) are computed between the representation ofa query molecule m (cid:48)(cid:48) and each molecule x (cid:48)(cid:48) n in the support set as done recently (Koch et al., 2015;
Revised: In this module, pairwise similarity values k ( m (cid:48)(cid:48) , x (cid:48)(cid:48) n ) are computed between the representation ofthe query molecule m (cid:48)(cid:48) and each molecule x (cid:48)(cid:48) n in the support set as done recently (Koch et al., 2015;

Original: Altae-Tran et al., 2017).
Revised: Altae-Tran et al., 2017).

Original: Based on these similarity values, the activity for the query molecule ispredicted, building a weighted mean over the support set labels:
Revised: Based on these similarity values, the activity for the query molecule ispredicted, building the weighted mean over the support set labels:

Original: [Equation]
Revised: [Equation]

Original: (10) where our architecture employs dot product similarity of normalized representations k ( m (cid:48)(cid:48) , x (cid:48)(cid:48) n ) =m (cid:48)(cid:48)
Revised: (10) where our architecture employs dot product similarity of normalized representations k ( m (cid:48)(cid:48) , x (cid:48)(cid:48) n ) =m (cid:48)(cid:48)

Original: T x (cid:48)(cid:48) n .
Revised: T x (cid:48)(cid:48) n .

Original: σ ( . ) is the sigmoid function and τ is a hyperparameter.
Revised: σ ( . ) is the sigmoid function and τ is a hyperparameter.

Original: Note that we use a balancingstrategy for the labels y (cid:48) n = (cid:26)

Original: [Equation]

Original:  if [Equation] [Equation] ) else , where N A is the number of actives and N Iis the number of inactives of the support set.
Revised: Note that we use a balancingstrategy for the labels y (cid:48) n = (cid:26) N/ (2 √ N A ) if y n =− N/ (2 √ N I ) else , where N A is the number of actives and N Iis the number of inactives of the support set.

Original: 3.5 Architecture, hyperparameter selection, and training details
Revised: 3.5 Architecture, hyperparameter selection, and training details

Original: Hyperparameters.
Revised: Hyperparameters.

Original: The main hyperparameters of our architecture are the number of heads, theembedding dimension, the dimension of the association space of the CAM and CM, the learning rate schedule, the scaling parameter β, and the molecule encoder.
Revised: The main hyperparameters of our architecture are the number of heads, theembedding dimension and the dimension of the association space of the CAM and CM, the learning rate schedule, the scaling parameter β and the molecule encoder.

Original: The following hyperparameterswere selected by manual hyperparameter selection on the validation tasks.
Revised: The following hyperparameterswere selected by manual hyperparameter selection on the validation tasks.

Original: The molecule encoderconsists of a single layer with output size d = 1024 and SELU activation (Klambauer et al., 2017).
Revised: The molecule encoderconsists of a single layer with output size d = 1024 and SELU activation (Klambauer et al., 2017).

Original: The CM consists of one Hopﬁeld layer with 8 heads.
Revised: The CM consists of one Hopﬁeld layer with 8 heads.

Original: The dimension e of the association space is setto 512 and β = 1 / √ e .
Revised: The dimension e of the association space is setto 512 and β = 1 / √ e .

Original: Since we use skip connections between all modules the output dimension ofthe CM and CAM matches the input dimension.
Revised: As we use skip connections between all modules, the output dimension ofthe CM and CAM matches the input dimension.

Original: The CAM comprises one layer with 8 heads and anassociation-space dimension of 1088.
Revised: The CAM comprises one layer with 8 heads andan association-space dimension of 1088.

Original: For the input to the CAM, an activity encoding was added tothe support set molecule representations to provide label information.
Revised: For the input to the CAM, an activity encoding was addedto the support set molecule representations to provide label information.

Original: The SM uses τ = 22 . 6 .
Revised: The SM uses τ = 22 . 6 .

Original: For thecontext set, we randomly sample 5% from a large set of molecules – i.e., the molecules in the FS-Moltraining split – for each batch.
Revised: For the context set, we randomly sample 5% from a large set of molecules, i.e. the molecules inthe FS-Mol training split, for each batch.

Original: For inference, we used a ﬁxed set of 5% of training set molecules asthe context set for each seed.
Revised: For inference, we used a ﬁxed set of 5% of training setmolecules as the context set for each seed.

Original: We hypothesize that these choices about the context could be furtherimproved (Section 6).
Revised: We hypothesize that these choices about the context couldbe further improved (see Section 5).

Original: We provide considered and selected hyperparameters in Appendix A.1.6.
Revised: For a table of considered and selected hyperparameters, see Supplementary.

Original: Loss function, regularization and optimization.
Revised: Loss function, regularization and optimization.

Original: We use the Adam optimizer (Kingma and Ba,2014) to minimize the cross-entropy loss between the predicted and known activity labels.
Revised: We use the Adam optimizer (Kingma and Ba,2014) to minimize the cross-entropy loss between the predicted and known activity labels.

Original: We usea learning rate scheduler which includes a warm up phase, followed by a section with a constantlearning rate, which is 0 . 0001 , and a third phase in which the learning rate steadily decreases.
Revised: We usea learning rate scheduler which includes a warm up phase, followed by a section with a constantlearning rate, which is 0 . 0001 , and a third phase in which the learning rate steadily decreases.

Original: As aregularization strategy, for the CM and the CAM a dropout rate of 0 . 5 is used.
Revised: As aregularization strategy, for the CM and the CAM a dropout rate of 0 . 5 is used.

Original: The molecule encoderhas a dropout with rate 0 .
Revised: The molecule encoderhas a dropout with rate 0 .

Original: 1 for the input and 0 .
Revised: 1 for the input and 0 .

Original: 5 elsewhere (see also Appendix A.1.6).
Revised: 5 elsewhere (see also Supplementary).

Original: Compute time and resources.
Revised: Compute time and resources.

Original: Training a single MHNfs model on the benchmarking dataset FSMol takes roughly 90 hours of wall-clock time on an A100 GPU.
Revised: Training a single MHNfs model on the benchmarking dataset FSMol takes roughly 90 hours of wall-clock time on an A100 GPU.

Original: In total, roughly 15,000 GPU hourswere consumed for this work.Related work [Figure]
Revised: In total, roughly 15,000 GPU hourswere consumed for this work.

Original:  Several approaches to few-shot learning in drug discovery have been suggested (Altae-Tran et al. ,2017; Nguyen et al., 2020; Guo et al., 2021; Wang et al., 2021). (Nguyen et al., 2020)
Revised:  Experiments4.1 Benchmarking

Original: evaluatedthe applicability of MAML and its variants to graph neural networks (GNNs) and (Guo et al. , 2021)also combine GNNs and meta-learning. (Altae-Tran et al., 2017) suggested an approach called Iterative Reﬁnement Long Short-Term Memory, in which query and support set embeddings canshare information and update their embeddings.
Revised: on FS-Mol

Original: Property-aware relation networks (PAR) (Wanget al., 2021) use an attention mechanism to enrich representations from cluster centers and then learna relation graph between molecules.

Original: (Chen et al., 2022) propose to adaptively learn kernels and applytheir method to 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 drugdiscovery and provided some baseline results. Many successful deep neural network architectures use external memories, such as the neural Turingmachine (Graves et al., 2014), memory networks (Weston et al., 2014), end-to-end memory networks(Sukhbaatar et al., 2015). Recently, the connection between continuous modern
Revised: Experimental setup.

Original: Hopﬁeld networks(Ramsauer et al., 2021), which are content-addressable associative memories, and Transformerarchitectures (Vaswani et al., 2017) has been established.

Original: We refer to (Le, 2021) for an extensiveoverview of memory-based architectures.

Original: Architectures with external memories have also been usedfor meta-learning (Vinyals et al., 2016;

Original: Santoro et al., 2016) and few-shot learning (Munkhdalai and

Original: Yu, 2017; Ramalho and Garnelo, 2018; Ma et al., 2021).

Original: 5.1 Benchmarking on FS-Mol Experimental setup. Recently, the dataset FS-Mol (Stanley et al., 2021) was proposed to benchmarkfew-shot learning methods in drug discovery.
Revised: Recently, the dataset FS-Mol (Stanley et al., 2021) was proposed to benchmarkfew-shot learning methods in drug discovery.

Original: It was extracted from ChEMBL27 and comprises intotal 489,133 measurements, 233,786 compounds and 5,120 tasks.
Revised: It was extracted from ChEMBL27 and comprises intotal 489,133 measurements, 233,786 compounds and 5,120 tasks.

Original: Per task, the mean number ofdata points is 94.
Revised: Per task, the mean number ofdata points is 94.

Original: The dataset is well balanced as the mean ratio of active and inactive molecules isclose to 1.
Revised: The dataset is well balanced as the mean ratio of active and inactive molecules isclose to 1.

Original: The FS-Mol benchmark dataset deﬁnes 4,938 training, 40 validation and 157 test tasks,guaranteeing disjoint task sets.
Revised: The FS-Mol benchmark dataset deﬁnes 4,938 training, 40 validation and 157 test tasks,guaranteeing disjoint task sets.

Original: (Stanley et al., 2021) precomputed extended connectivity ﬁngerprints(ECFP)
Revised: Stanley et al. (2021) precomputed extended connectivity ﬁngerprints(ECFP)

Original: (Rogers and Hahn, 2010) and key molecular physical descriptors, which were deﬁned by
Revised: (Rogers and Hahn, 2010) and key molecular physical descriptors, which were deﬁned by

Original: RDKit (Landrum et al., 2006).
Revised: RDKit (Landrum et al., 2006).

Original: While methods would be allowed to use other representations ofthe input molecules, such as the molecular graph, we used a concatenation of these ECFPs and [Table]
Revised: While methods would be allowed to use other representations of theinput molecules, such as the molecular graph, we use these provided descriptors as inputs for ourmethods.

Original: IterRefLSTM (Altae-Tran et al., 2017) [Table] RDKit-based descriptors. For the main benchmark, the support set size was ﬁxed to 16, using astratiﬁed random split.
Revised: For the main benchmark, the support set size was ﬁxed to 16, using a stratiﬁed randomsplit.

Original: We use all these settings of FS-Mol and therefore ensure a fair methodcomparison.
Revised: We use all these settings of FS-Mol and therefore ensure a fair method comparison.

Original: Methods compared.
Revised: Methods compared.

Original: Baselines for few-shot learning and our proposed method MHNfs were compared against each other.
Revised: Baselines for few-shot learning and our proposed method MHNfs werecompared against each other.

Original: The Frequent Hitters model is a naive baseline that ignores the providedsupport set and therefore has to learn to predict the average activity of a molecule.
Revised: The Frequent Hitters model is a naive baseline that ignores theprovided support set and therefore has to learn to predict the average activity of a molecule.

Original: This method canpotentially discriminate so-called frequent-hitter molecules (Stork et al., 2019) against moleculesthat are inactive across many tasks.
Revised: Thismethod can potentially discriminate so-called frequent-hitter molecules (Stork et al., 2019) againstmolecules that are inactive across many tasks.

Original: Wealso added Similarity Search (Cereto-Massagué et al., 2015)as a baseline.
Revised: We added Similarity Search (Cereto-Massagué et al.

Original: Similarity search is a standard chemoinformatics technique, used in situations withsingle or few known actives.
Revised: as a baseline, which is a standard chemoinformatics technique used in situations in whicha single or few actives are known.

Original: In the simplest case,the search ﬁnds similar molecules  by computinga ﬁngerprint or descriptor-representation of the molecules and using a similarity measure k ( ., . ) —such as Tanimoto Similarity (Tanimoto, 1960). 
Revised: In the simplest case, molecules that are similar to given activemolecules are searched by computing a ﬁngerprint or descriptor-representation of the molecules andusing a similarity measure k ( ., . ) , such as Tanimoto Similarity(Tanimoto, 1960). Thus, the Similarity

Original: Thus, Similarity Search, as used in chemoinformatics, can be formally written as ˆ y = 1 /N (cid:80)
Revised: Search as used in chemoinformatics can be formally written as ˆ y = 1 /N (cid:80)

Original: Nn =1 y n k
Revised: Nn =1 y n k

Original: ( m , x n ); where x 1 , . . 
Revised: ( m , x n ) ,where x 1 , . . .

Original: , x n come from aﬁxed molecule encoder, such as chemical ﬁngerprint or descriptor calculation.
Revised: , x n come from a ﬁxed molecule encoder, such as chemical ﬁngerprint or descriptorcalculation.

Original: A natural extension of Similarity Search with ﬁxed chemical descriptors is Neural Similarity Search or Siamesenetworks (Koch et al., 2015), which extend the classic similarity search by learning a moleculeencoder: ˆ y = σ (cid:16)
Revised: A natural extension of Similarity Search with ﬁxed chemical descriptors is Neural Similarity Search or Siamese networks (Koch et al., 2015) wich extend the classic similarity searchby learning a molecule encoder:

Original:  τ − 1 1 N (cid:80) Nn =1 y (cid:48)
Revised: ˆ y = σ (cid:16) τ − 1 1 N (cid:80) Nn =1 y (cid:48)

Original: n
Revised: n

Original: f ME w ( m ) T
Revised: f ME w ( m ) T

Original: f ME w ( x n ) (cid:17) .
Revised: f ME w ( x n ) (cid:17) .

Original: Furthermore, we re-implemented the IterRefLSTM (Altae-Tran et al., 2017) in Pytorch.
Revised: Furthermore, were-implemented the IterRefLSTM (Altae-Tran et al., 2017) in Pytorch.

Original: The IterRefLSTM model consists of threemodules.
Revised: The IterRefLSTM modelconsists of three modules.

Original: First, a molecule encoder maps the query and support set molecules to its representationsm and X .
Revised: First, a molecule encoder maps the query and support set moleculesto its representations m and X .

Original: Second, an attention-enhanced LSTM variant, the actual IterRefLSTM , iterativelyupdates the query and support set molecules, enabling information sharing between the molecules:[ m (cid:48) , X (cid:48) ]
Revised: Second, an attention-enhanced LSTM variant, the IterRefLSTM,iteratively updates the query and support set molecules, enabling information sharing between the [Table] metrics [Table]

Original: =
Revised: et [Table] [Table] molecules: [ m (cid:48) , X (cid:48) ]

Original:  IterRefLSTM L ([ m , X ]) , where the hyperparameter L controls the number of iterationsteps of the IterRefLSTM .
Revised: = IterRefLSTM L ([ m , X ]) , where the hyperparameter L controls the numberof iteration steps of the IterRefLSTM.

Original: Third, a similarity module computes attention weights based on the representations: a = softmax ( k ( m (cid:48) , X (cid:48) )) .
Revised: Third, a similarity module computes attention weights basedon the representations: a = softmax ( k ( m (cid:48) , X (cid:48) )) .

Original: These representations are then used for the ﬁnal prediction:ˆ y = (cid:80) Ni =1 a i y i .
Revised: These representations are then used for theﬁnal prediction: ˆ y = (cid:80) Ni = 1 a i y i .

Original: For further details, see Appendix A.1.5.
Revised: For further details, see Supplementary.

Original: The Random Forest baseline uses thechemical descriptors and is trained in standard supervised manner on the support set molecules foreach task.
Revised: The Random Forestbaseline uses the chemical descriptors and is trained in standard supervised manner on the supportset molecules for each task.

Original: The method GNN-ST is a graph neural network (Stanley et al., 2021;
Revised: The method GNN-ST is a graph neural network (Stanley et al., 2021;

Original: Gilmer et al. , 2017)that is trained from scratch for each task.
Revised: Gilmer et al., 2017) that is trained from scratch for each task.

Original: The GNN-MT uses a two step strategy: First, the modelispretrained on a large dataset on related tasks; second, an output layer is constructed to the few-shottask via linear probing (Stanley et al., 2021; Alain and Bengio, 2016).
Revised: The GNN-MT uses a two step strategy,which is ﬁrst to pre-train the model on a large dataset on related task and then adapting an outputlayer to the few-shot task via linear probing (Stanley et al., 2021; Alain and Bengio, 2016).

Original: TheMolecule Attention
Revised: The

Original:  Transformer (MAT) is pre-trained in a self-supervised fashion and ﬁne-tuning is performed for thefew-shot task (Maziarka et al., 2020).
Revised: Molecule Attention Transformer (MAT) is pre-trained in a self-supervised fashion and ﬁne-tuningis performed for the few-shot task (Maziarka et al., 2020).

Original: GNN-MAML is based on MAML (Finn et al., 2017), and usesa model-agnostic meta-learning strategy to ﬁnd a general core model from which one can easily adaptto single tasks.
Revised: GNN-MAML is based on MAML (Finnet al., 2017), and uses a model-agnostic meta-learning strategy to ﬁnd a general core model fromwhich one can easily adapt to single tasks.

Original: ProtoNet (Snell et al., 2017) includes a molecule encoder, which maps query andsupport set molecules to representations in an embedding space.In this embedding space, prototypical [Figure]
Revised: ProtoNet (Snell et al., 2017) includes a molecule encoder,which maps query and support set molecules to representations in an embedding space.

Original: This indicates that our two newly introduced modules, CM and CAM, play a crucial role in representations of each class are built by taking the mean across all related support set molecules foreach class (details in Appendix A.1.4). For all methods the most important hyperparameters wereadjusted on the validation tasks of FS-Mol. The PAR model (Wang et al., 2021) includes a GNNwhich creates initial molecule embeddings. These molecule embeddings are then enriched by anattention mechanism.
Revised: In thisembedding space, prototypical representations of each class are built by taking the mean across allrelated support support set molecules for each class (details in Supplementary).

Original: Finally, another GNN learns relations between support and query set molecules.

Original: The PAR model has shown good results for datasets which just include very few tasks such as Tox21(Wang et al., 2021). Chen et al. (2022) suggest a framework for learning deep kernels by interpolatingbetween meta-learning and conventional deep kernels, which results in the ADKF-IFT model. Themodel has exhibited especially high performance for large support set sizes.
Revised: For all methods themost important hyperparameters were adjusted on the validation tasks of FS-Mol.

Original: Training and evaluation.
Revised: Training and evaluation.

Original: For the model implementations, we used PyTorch (Paszke et al., 2019,
Revised: For the model implementations, we used PyTorch (Paszke et al., 2019,

Original: BSD license). We used PyTorch Lightning (Falcon et al., 2019, Apache 2.0 license) as a frameworkfor training and test logic, hydra for conﬁg ﬁle handling (Yadan, 2019, Apache 2.0 license) and Weights & Biases (Biewald, 2020, MIT license) as an experiment tracking tool.
Revised: BSD license), making use of PyTorch Lightning (Falcon et al., 2019, Apache 2.0 license) as aframework for training and test logic, hydra for conﬁg ﬁle handling (Yadan, 2019, Apache 2.0 license)and Weights & Biases (Biewald, 2020, MIT license) as an experiment tracking tool.

Original: We performed ﬁvetraining reruns with different seeds for all methods, except Classic Similarity Search as there isno variability across seeds.
Revised: We performedﬁve training reruns with different seeds for all methods, except Classic Similarity Search as there isno variability across seeds.

Original: Each  model was evaluated ten times by drawing support sets with tendifferent seeds.
Revised: Each trained model was evaluated ten times by drawing support sets withten different seeds.

Original: Results.
Revised: Results.

Original: The results in terms of area under precision-recall curve (AUC-PR) are presented in Table 1,where the difference to a random classiﬁer is reported ( ∆ AUC-PR).
Revised: The results in terms of area-under-precision-recall curve (AUC-PR) are presented in Table 1,where the difference to a random classiﬁer is reported ( ∆ AUC-PR).

Original: The standard error is reportedacross tasks.
Revised: The standard error is reportedacross tasks.

Original: Surprisingly, the naive baseline Frequent Hitters , that neglects the support set, has outperformed most of the few-shot learning methods, except for the embedding-based methods Siamese
Revised: Surprisingly, the naive baseline Frequent Hitters , that neglects the support set, has outperformed most of the few-shot learning methods, except for the embedding-based methods Siamese

Original: Networks , ProtoNet , IterRefLSTM , and MHNfs .
Revised: Networks , ProtoNet , IterRefLSTM , and MHNfs .

Original: IterRefLSTM , which has not been includedin the FS-Mol benchmark study, reaches the second best performance.
Revised: IterRefLSTM , which has not been includedin the FS-Mol benchmark study, reaches the second best performance.

Original: MHNfs has outperformedall other methods with respect to ∆ AUC-PR across all tasks, including the IterRefLSTM model( p -value 1.72e-7, paired Wilcoxon test), the ADKF-IFT model ( p -value <1.0e-8, Wilcoxon test), andthe PAR model ( p -value <1.0e-8, paired Wilcoxon test).
Revised: MHNfs has outperformedall other methods with respect to ∆ AUC-PR across all tasks, including the IterRefLSTM model( p -value 0.0014, paired Wilcoxon test).

Original: 5.2 Ablation study
Revised: 4.2 Ablation study

Original: MHNfs has two new main components compared to the previous state-of-the-art method IterRefLSTM : i) the context module, and ii) the cross-attention module which replaces the LSTM-likemodule.
Revised: MHNfs has two new main elements compared to the previous state-of-the-art method IterRefLSTM :i) the context module, and ii) the cross-attention module, which replaces the LSTM-like module.

Original: To assess the effects of these components, we performed an ablation study.
Revised: To assess the effects of these elements, we performed an ablation study.

Original: Therefore,we compared  MHNfs to a method that does not have the context module ("MHNfs -CM") and toa method that does not have the context module anduses an LSTM-like module instead of the 
Revised: Therefore, we compared [Figure] This indicates that our two newly introduced modules, CM and CAM, play a crucial role in MHNfs. MHNfs to a method that does not have the context module ("MHNfs -CM") and to a method thatdoes not have the context module and an LSTM-like module instead of the CAM ("MHNfs -CM 284 (cid:10) (CAM,IterRefLSTM)").

Original: CAM ("MHNfs -CM (cid:10) (CAM,IterRefLSTM)"). For the ablation study, we used all 5 training rerunsand evaluated each model 10 times on the test set with different support sets.
Revised: For the ablation study, we used all 5 training reruns and evaluatedeach model 10 times on the test set with different support sets.

Original: The results of thisablation steps are presented in Figure 2.
Revised: The results of this ablation steps arepresented in Figure 2.

Original: Both removing the CM and exchanging the CAM withthe IterRefLSTM module were detrimental for the performance of the method ( p -value 0.002 and1 . 
Revised: Both removing the CM and exchanging the CAM with the IterRefLSTM modulewere detrimental for the performance of the method ( p -value 0.002 and 1 . 3 e- 8 , respectively; paired

Original: 72 e − 7 , respectively; paired Wilcoxon test).
Revised: Wilcoxon test).

Original: The difference was even more pronounced underdomain shift (see Appendix A.3.3).
Revised: The difference was even more pronounced under domain shift (see Supplementary).

Original: Appendix A.3.2 contains a second ablation study that examines the overall effects of the context, the cross-attention, the similarity module, and the molecule encoderof MHNfs .
Revised: A second ablation study to examine the overall effects of the context module, the cross-attentionmodule, the similarity module, and the molecule encoder in the MHNfs model is reported in the Supplementary.

Original: 5.3 Domain shift experiment
Revised: 4.3 Domain shift experiment

Revised: Experimental setup.

Original: We performed an experiment in which we evaluate models, that were pretrained on FS-Mol, on the 
Revised: For the domain shift experiment, we used the Tox21 dataset. This datasetconsists of 12,707 chemical compounds, for which measurements for up to 12 different toxic effectsare reported (Mayr et al., 2016; Huang et al., 2016a). It was published with a ﬁxed training, validationand test split. State-of-the-art supervised learning methods that have access to the full training setreach AUC performance values between 0 . 845 and 0 .

Original: Tox21 (Mayr et al., 2016) dataset.
Revised: 871 (Klambauer et al., 2017; Duvenaud et al.

Original: There is a strong domain shift from the drug-like molecules of
Revised: ,2015; Li et al., 2017, 2021; Zaslavskiy et al., 2019; Alperstein et al., 2019). We re-cast Tox21 as afew-shot learning setting and draw small support sets from the 12 tasks.

Original:  FS-Molto the environmental chemicals, pesticides, and food additives of Tox21, such this datasetposes a challenging setting for few-shot learning methods.
Revised: The compared methods werepre-trained on FS-Mol and then obtain small support sets from Tox21.

Original:  The experiment is described in detailin Appendix A.2.
Revised: Based on those support sets,the methods had to predict the activities of the Tox21 test set. Note that there is a strong domainshift from drug-like molecules of FS-Mol to environmental chemicals, pesticides, food additivesof Tox21. The domain shift also concerns the outputs where a shift from kinases, hydrolases, andoxidoreductases of FS-Mol to nuclear receptors and stress responses of Tox21 is present. Methods compared.

Original: Our MHNfs approach has reached an AUC of .
Revised: The new method MHNfs , the runner-up method IterRefLSTM , and Similarity

Original: 679 ± .
Revised: Search , since it has been widely used for such purposes for decades (Cereto-Massagué et al., 2015),were compared.

Original:  and has signiﬁcantlyoutperformed the IterRefLSTM -based model
Revised: Training and evaluation. We followed the procedure of Stanley et al.

Original: ( p ∆AUC − PR -value 3 .
Revised: (2021) for data-cleaning,preprocessing and extraction of the ﬁngerprints and descriptors used in FS-Mol. After running thecleanup step, 8,423 molecules remained for the domain shift experiments. From the training set, 8active and 8 inactive molecules per task were randomly selected to build the support set. The test setmolecules were used as query molecules.

Original: 4 e − 5 , paired Wilcoxon test) andthe Classic Similarity Search ( p ∆AUC
Revised: The validation set molecules were not used at all. Duringtest-time, a support set was drawn ten times for each task. Then, the performance of the model wasevaluated for these support sets, using the area [Table]

Original: − PR -value 2 . 
 
Revised: The Hopﬁeld-based context retrieval method has signiﬁcantly outperformed the IterRefLSTM-based model ( p ∆AUPRC -value 0.0002, p AUC -value 9 . 4 e- 5 , paired Wilcoxon test)and the Classic Similarity Search ( p ∆AUPRC -value 2 . 
 e-

Original: e- 9  paired Wilcoxon test) and therefore showedrobust performance on the toxicity domain, see Table A6.
Revised: 9 , p AUC -value 7 . 6 e- 10 , paired Wilcoxontest) and therefore showed robust performance on the toxicity domain, see Table 2.

Revised: Notably, for noneof these methods, any weight parameter was adjusted to the new dataset.

Original:  Conclusion and discussion
Revised:  Conclusion and discussion

Original: We have introduced a new architecture for few-shot learning in drug discovery that is based onthe novel concept to enrich molecule representations with context.
Revised: We have introduced a new architecture for few shot learning in drug discovery that is based on thenovel concept to enrich molecule representations with context.

Revised: The architecture, called MHNfs ,comprises an adaptive molecule encoder, a context module, a cross-attention module and a similaritymodule.

Original: In a benchmarking experiment,the architecture was assessed for its ability to learn accurate predictive models from small sets oflabelled molecules and in this setting it outperformed all other methods.
Revised: In a benchmarking experiment, the architecture was assessed for its ability to learn accuratepredictive models from small sets of labelled molecules.

Revised: Indeed, MHNfs outperformed all previouslysuggested methods on FS-Mol.

Original: In a domain shift study, therobustness and transferability of the learned models has been assessed and again MHNfs exhibitedthe best performance.The resulting predictive models often reach an AUC larger than .
Revised: In a domain shift study, the robustness and transferability of thelearned models has been assessed and again MHNfs exhibited the best performance.

Original: 70 , whichmeans that enrichment of active molecules is expected (Simm et al., 2018) when the models are usedfor virtual screening. It has not escaped our notice that the speciﬁc context module we have proposedcould immediately be used for few-shot learning tasks in computer vision, but might be hamperedby computational constraints.
Revised: It has notescaped our notice that the speciﬁc context module we have proposed could immediately be used forfew-shot learning tasks in computer vision, but might be hampered by computational constraints.

Original: Limitations.
Revised: Limitations.

Original: While the implementation of our method is currentlylimited to small, organic drug-like molecules as inputs, our conceptual approach can also be usedfor macro-molecules such as RNA, DNA or proteins.
Revised: Our method is currently limited to small, organic drug-like molecules as inputs. Thus,the model cannot be used for molecules such as RNA, DNA or proteins.

Original: The output domainof our method comprisesbiological effects, such that the prediction must be understood in that domain.Our method demandshigher computational costs and memory footprint as other embedding-based methods because ofthe calculations necessary for the context module.
Revised: The output domain comprisesbiological effects focusing on effects on drug targets, such that the output of the trained models mustbe understood in that domain.

Original: While we hypothesize that our approach couldalso be successful for similar data in the materials science domain, this has not been assessed.
Revised: While we hypothesize that the approach could also be successful forsimilar data in the materials science domain, this has not been assessed.

Original: Ourstudy is also constrained by a limited amount of hyperparameter search for all methods.
Revised: Our study is also constrainedby a limited amount of hyperparameter search for all methods.

Original: Deeplearning 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.
Revised: Deep Learning methods usually havea large number of hyperparameters, such as hidden dimensions, number of layers, learning rates, ofwhich we were only able to explore the most important ones.

Original: Thecomposition and choice of the context set is also under-explored and might be improved by selectingreference molecules with an appropriate strategy.
Revised: The composition and choice of thecontext set is also underexplored and might be improved by selecting reference molecules with anappropriate strategy.

Revised: Further considerations can be found in the section below.

Original: Broader impact.
Revised: Broader impact.

Original: Impact on machine learning andrelated scientiﬁc ﬁelds.
Revised: Impact on machine learning and related scientiﬁc ﬁelds.

Original: We envision that with (a) the increasing availability of drug discovery andmaterial science datasets, (b) further improved biotechnologies, and (c) accounting for characteristicsof individuals, the drug and materials discovery process will be made more efﬁcient.
Revised: We envision that with(a) the increasing availability of drug discovery and material science dataset, (b) further improvedbiotechnologies, and (c) accounting for characteristics of individuals, the drug and materials discoveryprocess will be made more efﬁcient.

Original: For machinelearning and artiﬁcial intelligence, the novel way in which representations are enriched with contextmight strengthen the general research stream to include more context into deep learning systems.
Revised: For machine learning and artiﬁcial intelligence, the novel wayin which representations are enriched with context might strengthen the general research stream toinclude more context into Deep Learning systems.

Original: Ourapproach also shows that such a system is more robust against domain shifts, which could be a steptowards Broad AI (Chollet, 2019; Hochreiter, 2022).
Revised: Our approach also shows that such a system ismore robust against domain shifts, which could be a step towards Broad AI (Chollet, 2019; Hochreiter,2022).

Original: Impact on society.
Revised: Impact on society.

Original: If the approach proves useful,it could lead to a faster and more cost-efﬁcient drug discovery process.
Revised: If the approach proves itself useful, it could lead to a faster and morecost-efﬁcient drug discovery process.

Original: Especially the COVID-19pandemic has shown that it is crucial for humanity to speed up the drug discovery process to few yearsor even months.
Revised: Especially the COVID-19 pandemic has shown that it is crucialfor humanity to speed up the drug discovery process to few years or even months.

Original: We hope that this work contributes to this effort and eventually leads to safer drugsdeveloped faster.
Revised: We hope that thiswork contributes to this effort and eventually leads to safer drugs developed faster.

Original: Consequences of failures of the method.
Revised: Consequences offailures of the method.

Original: As common with methods in machinelearning, potential danger lies in the possibility that users rely too much on our new approach and useit without reﬂecting on the outcomes.
Revised: As common with methods in machine learning, potential danger lies in thepossibility that users rely too much on our new approach and use it without reﬂecting on the outcomes.

Original: Failures of the proposed method would lead to unsuccessfulwet lab validation and negative wet lab tests.
Revised: Failures of the proposed method would lead to unsuccessful wet lab validation and negative wet labtests.

Original: Since the proposed algorithm does not directly suggesttreatment or therapy, human beings are not directly at risk of being treated with a harmful therapy.
Revised: Since the proposed algorithm does not directly suggest treatment or therapy, human beings arenot directly at risk of being treated with a harmful therapy.

Original: Wet lab and in-vitro testing would indicate wrong decisions by the system.
Revised: Wet lab and in-vitro testing would indicatewrong decisions by the system.

Original: Leveraging of biases inthe data and potential discrimination.
Revised: Leveraging of biases in the data and potential discrimination.

Original: As for almost all machine learning methods, confoundingfactors, lab or batch effects, could be used for classiﬁcation.
Revised: Asfor almost all machine learning methods, confounding factors, lab or batch effects, could be usedfor classiﬁcation.

Original: This might lead to biases in predictionsor uneven predictive performance across different drug targets or bioassays.
Revised: This, might lead to biases in predictions or uneven predictive performance acrossdifferent drug targets or bioassays.

Original: Checklist
Revised: Checklist

Original: For all authors...
Revised: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’scontributions and scope?
Revised: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you describe the limitations of your work?
Revised: (b) Did you describe the limitations of your work?

Original: [Yes] See Section 6.(c)
Revised: [Yes] See Section 5.

Original:  Did you discuss any potential negative societal impacts of your work?
Revised: (c) Did you discuss any potential negative societal impacts of your work?

Original: [Yes] SeeSection 6.(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
Revised: [Yes] SeeSection 5. (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Original: [Yes]2.
Revised: [Yes]

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?
Revised: (a) Did you state the full set of assumptions of all theoretical results?

Original: [N/A](b)
Revised: [N/A]

Original:  Did you include complete proofs of all theoretical results?
Revised: (b) Did you include complete proofs of all theoretical results?

Original: [N/A]3.
Revised: [N/A]

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)?

Original: [Yes](b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
Revised: [Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?

Original: [Yes] We refer e.g. to Section 5 in which we provide this information.(c)
Revised: [Yes] We refer e.g. to Section 4 in which we provide this information.

Revised: All training details are provided in Supplementary Material.

Original:  Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
Revised: (c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)?

Original: [Yes] We report error bars for all performance metrics.
Revised: [Yes] We report error bars for all performance metrics.

Original: For allexperiments the variability across re-runs, different support sets and prediction taskswas assessed.(d)
Revised: For all experiments the variability across re-runs, different support sets and prediction taskswas assessed.

Original:  Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
Revised: (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?

Original: [Yes] See Section 3.5.4.
Revised: [Yes] See Section 3.5.

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?
Revised: (a) If your work uses existing assets, did you cite the creators?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you mention the license of the assets?
Revised: (b) Did you mention the license of the assets?

Original: [Yes](c)
Revised: [Yes]

Original:  Did you include any new assets either in the supplemental material or as a URL?
Revised: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [Yes]
Revised: [Yes]

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re
Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?
Revised: using/curating?

Original: [N/A]
Revised: [N/A]

Original: (e) Did you discuss whether the data you are using/curating contains personally identiﬁable
Revised: (e) Did you discuss whether the data you are using/curating contains personally identiﬁable

Original: information or offensive content?
Revised: information or offensive content?

Original: [Yes] See Section 6.
Revised: [Yes] See Section 5.

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: If you used crowdsourcing or conducted research with human subjects...

Original: (a) Did you include the full text of instructions given to participants and screenshots, if
Revised: (a) Did you include the full text of instructions given to participants and screenshots , if

Original: applicable?
Revised: applicable?

Original: [N/A]
Revised: [N/A]

Original: (b) Did you describe any potential participant risks, with links to Institutional ReviewBoard (IRB) approvals, if applicable? [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount
Revised: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: spent on participant compensation?

Original: [N/A]

Original: A

Original: [-Appendix 
-]
Revised: [-Appendix 
-]

Original: A.1 Details on methods . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Original: A.1.

Original: hyperparameters . . . . . . . . . . . . . . . .

Original: A.1.

Original: search:

Original: A.1.

Original: [-Search or and hyperparameters 
-]
Revised: [-Search or and hyperparameters 
-]

Original: A.1.4 ProtoNet: details and hyperparameters . . . . . . . . . . . . . . . . . . . .

Original: A.1.5 IterRefLSTM: details and hyperparameters . . . . . . . . . . . . . . . . .

Original: A.1.6 MHNfs: hyperparameters . . . . . . . . . . . . . . . . . . . . .

Original: A.1.7 PAR

Original: A.

Original: A

Original: A

Original: A

Original: A

Original: A

Original: A Appendix

Original: A.1 Details on methods

Original: Few-shot learning methods in drug discovery can be described as models with adaptive parameters wthat use a support set Z = { ( x 1 , y 1 ) , . . .

Original: , ( x N , y N ) } 1 as additional input to predict a label ˆ y for amolecule m

Original: [Equation]

Original: A

Original: [Equation]

Original: Optimization-based methods, such as MAML (Finn et al., 2017), use the support set to update theparameters w

Original: [Equation]

Original: A2) where a ( . ) is a function that adapts w of g based on Z for example via gradient-descent.

Original: Embedding-based methods use a different approach and learn representations of the support setmolecules { x 1 , . . .

Original: , x N } , sometimes written as stacked embeddings X ∈ R d × N , and the querymolecule m , and some function that associates these two types of information with each other.

Original: Wedescribe the embedding-based methods Similarity Search in Section A.1.2, Neural Similarity Searchin Section A.1.3, ProtoNet in Section A.1.4, IterRefLSTM in Section A.1.5, PAR in Section A.1.7,and MHNfs in the main paper and details in Section A.1.6.

Original: The "frequent hitters" baseline is describedin Section A.1.1.

Original: A.1.1 Frequent hitters: details and hyperparameters

Original: The "frequent hitters" model g FH is a baseline that we implemented and included in the methodcomparison.

Original: This method uses the usual training scheme of sampling a query molecule m with alabel y , having access to a support set Z .

Original: In contrast to the usual models of the type g w ( m , Z ) , thefrequent hitters model g FH neglects the support set:

Original: [Equation]

Original: Thus, during training for the same molecule m , the model might have to predict both y = 1 andy = − 1 , since the molecule can be active in one task and inactive in another task.

Original: Therefore, the

Original: [Table]

Original: model tends to predict average activity of a molecule to minimize the cross-entropy loss.

Original: We chosean additive combination of the Morgan ﬁngerprints, RDKit ﬁngerprints, and MACCS keys for theinput representation to the MLP.

Original: Hyperparameter search.

Original: We performed manual hyperparameter search on the validation set andreport the explored hyperparameter space (Table A1).

Original: We use early-stopping based on validationaverage-precision, a patience of 3 epochs and train for a maximum of 20 epochs with a linear warm-uplearning-rate schedule for the ﬁrst 3 epochs.

Original: A.1.2 Classic similarity search: details and hyperparameters

Original: Similarity Search (Cereto-Massagué et al., 2015) is a classic chemoinformatics technique used insituations in which a single or few actives are known.

Original: In the simplest case, molecules that are similar toa given active molecule are searched by computing a ﬁngerprint or descriptor-representation f desc ( m )of the molecules and using a similarity measure k ( ., . ) , such as Tanimoto Similarity(Tanimoto, 1960).

Original: Thus, the Similarity Search as used in chemoinformatics can be formally written as:

Original: [Equation]

Original: where the function f desc maps the molecule to its chemical descriptors or ﬁngerprints and takesthe role of both the molecule encoder and the support set encoder.

Original: The association function f assocconsists of a) the similarity measure k ( ., . )

Original: and then b) mean pooling across molecules weighted bytheir similarity and activity.

Original: Notably, there are many variants of Similarity Search (Cereto-Massagué et al., 2015; Wang et al.

Original: ,2010; Eckert and Bajorath, 2007;

Original: Geppert et al., 2008; Willett, 2014; Sheridan and Kearsley, 2002;

Original: Riniker and Landrum, 2013) of which some correspond to recent few-shot learning methods with aﬁxed molecule encoder.

Original: For example, (Geppert et al., 2008) suggest to use centroid molecules, i.e.,prototypes or averages of active molecules.

Original: This is equivalent to the idea of Prototypical Networks(Snell et al., 2017).

Original: Riniker and Landrum (2013) are aware of different fusion strategies for sets ofactive or inactive molecules, which corresponds to different pooling strategies of the support set.

Original: Overall, the variants of the classic Similarity Search are highly similar to embedding-based few-shotlearning methods except that they have a ﬁxed instead of a learned molecule encoder.

Original: Hyperparameter search.

Original: For the Similarity Search, there were two decisions to make which wasﬁrstly the similarity metric and secondly the question whether we should use a balancing strategy likeshown in Section 3.4.

Original: We decided for the dot-product as a similarity metric and using the balancingstrategy.

Original: These decisions were made by evaluating the models on the validation set.

Original: A.1.3 Neural Similarity Search or Siamese networks: details and

Original: A lot of related work already was done (Koch et al., 2015; Hertz et al., 2006; Ye and Guo, 2018;

Original: Torres et al., 2020).

Original: We adapted these ideas, such that a fully-connected deep neural network followedby a Layer Normalization (Ba et al., 2016) operation, f ME w with adaptive parameters w , is used in a

Original: Siamese fashion to compute the embeddings for the input molecule and the support set molecules.

Original: Within the association function block, pairwise similarity values for the input molecule and eachsupport set molecule are computed, associating both embeddings via the dot product.

Original: Based on thesesimilarity values, the activity for the input molecule is predicted, building the weighted mean overthe support set molecule labels:

Original: [Equation]

Original: where σ ( . ) is the sigmoid function and τ is a hyperparameter in the range of √ d .

Original: Note that thismethod uses a balancing strategy for the labels y (cid:48)

Original: n =

Original: [Equation]

Original: [Equation]

Original: 2 √ N I ) else , where N A is thenumber of actives and N I is the number of inactives of the support set.

Original: Figure A1 provides anschematic overview of the Neural Similarity Search variant.

Original: We trained the networks using the Adam optimizer (Kingma and Ba, 2014) to minimize binarycross-entropy loss.

Original: Hyperparameter search.

Original: We performed manual hyperparameter search on the validation set.

Original: Wereport the explored hyperparameter space (Table A2).

Original: Bold values indicate the selected hyperparameters for the ﬁnal model.

Original: A.1.4 ProtoNet: details and hyperparameters

Original: Prototypical Networks (ProtoNet) (Snell et al., 2017), learn a prototype r for each class.

Original: Concretely,the support set Z is class-wise separated into Z

Original: + := { ( x, y ) ∈

Original: Z | y = 1 } and Z − :=

Original: { ( x, y ) ∈

Original: Z |y = − 1 } .

Original: For the subsets Z + and Z − prototypical representations r + and r − can be computed by

Original: [Equation]

Original: and

Original: [Equation]

Original: The prototypical representations r + , r − ∈ R d and the query molecule embedding m ∈ R d are thenused to make the ﬁnal prediction:

Original: [Equation]

Original: A

Original: [Equation]

Original: [Table]

Original: where d is a distance metric.

Original: Hyperparameter search.

Original: Hyperparameter search has been done in Stanley et al.

Original: (2021), to whichwe refer here.

Original: ECFP ﬁngerprints and descriptors created by a GNN, which operates on the moleculargraph, are fed into a fully connected neural network, which maps the input into an embedding spacewith the dimension of 512.

Original: (Stanley et al., 2021) use the Mahalanobis distance to measure thesimilarity between a query molecule and the prototypical representations in the embedding space.

Original: The learning rate is 0 .

Original: 001 and the batch size is 256.

Original: The implementation can be found here https://github.com/microsoft/FS-Mol/blob/main/fs_mol/protonet_train.py and importanthyperparameters are chosen here https://github.com/microsoft/FS-Mol/blob/main/fs_mol/utils/protonet_utils.py .

Original: Connection to Siamese networks and contrastive learning with InfoNCE.

Original: If instead of the negative distance − d ( ., . )

Original: the dot product similarity measure with appropriate scaling is used, ProtoNetfor two classes becomes equivalent to Siamese Networks.

Original: Note that in our study, another difference is that ProtoNet uses a GNN as encoder, whereas Siamese Networks use a descriptor-basedfully-connected network as encoder.

Original: In case of dot product as similarity measure, the objective alsobecomes equivalent to contrastive learning with the InfoNCE objective (Oord et al., 2018).

Original: A.1.5 IterRefLSTM: details and hyperparameters

Original: (Altae-Tran et al., 2017)

Original: modiﬁed the idea of Matching Networks (Vinyals et al., 2016) by replacingthe LSTM with their Iterative Reﬁnement Long Short-Term Memory (IterRefLSTM).

Original: The use of the

Original: IterRefLSTM empowers the architecutre to update not only the embeddings for the input moleculebut also adjust the representations of the support set molecules.

Original: For IterRefLSTM, m =

Original: f ME θ 1 ( m ) and x n = f ME θ 2 ( x n ) are two potentially different molecule encoders for input molecule m and the support set molecules x 1 , . . .

Original: , x N .

Original: The next step in IterRefLSTMis:

Original: [Equation]

Original: Here, m (cid:48) and X (cid:48) contain the updated representations for the query molecule and the supportset molecules.

Original: The IterRefLSTM denotes the function which updates these representations.

Original: Themain property of the IterRefLSTM module is that it is permutation-equivariant, thus a permutation π ( . ) of the input elements results in the permutation of output elements: π ([ m (cid:48) , X (cid:48) ]) =

Original: IterRefLSTM L ( π ([ m , X ])) .

Original: The full architecture in invariant to permutations of the support setelements.

Original: For details, we refer to (Altae-Tran et al., 2017).

Original: The hyperparameter L ∈ N controls thenumber of iteration steps of the IterRefLSTM.

Original: [Table]

Original: • Number of hidden layers

Original: [Table]

Original: • Number of units per hidden layer

Original: [Table]

Original: • Output dimension

Original: [Table]

Original: • Activation function

Original: [Table]

Original: • Input dropout

Original: [Table]

Original: • Dropout

Original: [Table]

Original: IterRef embedding layer •

Original: L [Table]
Revised: Board

Original: Similarity module: • Metric

Original: [Table]

Original: • Similarity space dimension

Original: [Table]

Original: Layer-normalization

Original: [Table]

Original: • Afﬁne

Original: [Table]

Original: • Optimizer

Original: [Table]

Original: • Weight decay

Original: [Table] • Batch size [Table] As similarity module, the IterRefLSTM uses the following:
Revised: (IRB) approvals, if applicable?

Original: [Equation]

Original: where ˆ y is the prediction for the query molecule.

Original: For the computation of the attention values a , thesoftmax function is used.

Original: k is a similarity metric, such as the cosine similarity.

Original: Hyperparameter search.

Original: All hyperparameters were selected based on manual tuning on thevalidation set.

Original: We report the explored hyperparameter space in Table A3.

Original: Bold values indicate theselected hyperparameters for the ﬁnal model.

Original: A.1.6 MHNfs: details and hyperparameters

Original: The MHNfs consists of a molecule encoder, the context module, the cross-attention-module, and thesimilarity module.

Original: The molecule encoder is a fully-connected Neural Network, consisting of onelayer with 1024 units.

Original: For the context module, a Hopﬁeld layer with 8 heads is used and also the crossattention module include 8 heads.

Original: We chose a concatenation of ECFPs and RDKit-based descriptorsas the inputs for the MHNfs model.

Original: Notably, the RDKit-based descriptors were pre-processed in away that instead of raw values quantils, which were computed by comparing a raw value with thedistributation of all FS-Mol training molecules, were used.

Original: All descriptors were normalized based onthe FS-Mol training data.

Original: Hyperparameter search.

Original: All hyperparameters were selected based on manual tuning on thevalidation set.

Original: We report the explored hyperparameter space in Table A4.

Original: Bold values indicate theselected hyperparameters for the ﬁnal model.

Original: Early stopping points for the different re-runs arechosen based on the ∆ AUC-PR metric on the validation set.

Original: For the ﬁve re-runs the early-stoppingpoints, that were automatically chosen by their validation metrics, were the checkpoints at epoch 94,192, 253, 253 and 309.

Original: Model training.

Original: Figure A2 shows the learning curve of an exemplary training run of a MHNfsmodel on FS-Mol.

Original: The left plot shows the loss on the training set and the right plot shows the

Original: [Table]

Original: • Number of hidden layers

Original: [Table]

Original: • Number of units per hidden layer

Original: [Table]

Original: • Output dimension

Original: [Table]

Original: • Activation function

Original: [Table]

Original: • Input dropout

Original: [Table]

Original: • Dropout

Original: [Table]

Original: Context module (hopﬁeld layer) • Heads

Original: [Table]

Original: • Association space dimension

Original: [Table]

Original: • τ

Original: [Table]

Original: • Dropout

Original: [Table]

Original: Cross-attention module (transformer mechanism) • Heads

Original: [Table]

Original: • Number units in the hidden feedforward layer

Original: [Table]

Original: • Association space dimension

Original: [Table]

Original: • Dropout

Original: [Table]

Original: • Number of layers:

Original: [Table]

Original: Similarity module: • Metric

Original: [Table]

Original: • Similarity space dimension

Original: [Table]

Original: -normalization

Original: [Table]

Original: • Afﬁne

Original: [Table]

Original: Training•Learning rate

Original: [Table]

Original: • Optimizer

Original: [Table]

Original: • Weight decay

Original: [Table]

Original: • Batch size

Original: [Table]

Original: • Warm-up phase (epochs)

Original: [Table]

Original: • Constant learning rate phase (epochs)

Original: [Table]

Original: • Decay rate

Original: [Table]

Original: • Max.

Original: number of epochs

Original: [Table]

Original: validation loss.

Original: The dashed line indicates the checkpoint of the model which was saved and then usedfor inference on the test set, whereas the stopping point was evaluated maximizing the ∆ AUC-PRmetric on the validation set.

Original: Performance improvements in comparison to a naive baseline.

Original: Figure A3 shows a task-wiseperformance comparison between MHNfs and the frequent hitter model.

Original: Each point indicates a taskin the test set and is colored according to their super-class membership.

Original: In 132 cases the MHNfsoutperforms the frequent hitter model.

Original: In 25 cases the frequent hitter model yields better performance.

Original: A.1.7 PAR: details and hyperparameters

Original: The PAR model (Wang et al., 2021) includes a pre-trained GNN encoder, which creates initialembeddings for the query and support set molecules.

Original: These embeddings are fed into an attentionmechanism module which also uses activity information of the support set molecules to createenriched representations.

Original: Another GNN learns relations between query and support set molecules.

Original: Hyperparameter search.

Original: For details we refer to (Wang et al., 2021) and https://github.com/tata1661/PAR-NeurIPS21/blob/main/parser.py .

Original: All hyperparameters were selected basedon manual tuning on the validation set.

Original: The hyperparameter choice for Tox21 (Wang et al., 2021)was used as a starting point.

Original: We report the explored hyperparameter space in Table A5.

Original: Bold values

Original: [Figure]
Revised: [N/A]

Original: L o ss o n t r a i n i n g s e t       350 Epoch 0.

Original: L o ss o n v a li d a t i o n s e t

Original: [Figure]

Original: [Table]

Original: • Inner

Original: [Table]

Original: rate

Original: [Table]

Original: • Update step

Original: [Table]

Original: • Update step test

Original: [Table]

Original: • Weight decay

Original: [Table]

Original: • Epochs

Original: [Table]

Original: • Eval.

Original: steps

Original: [Table]

Original: • Map dimension

Original: [Table]

Original: • Map layer

Original: [Table]

Original: • Pre fc layer

Original: [Table]

Original: • Map dropout

Original: [Table]

Original: • Context layer

Original: [Table]

Original: Relation graph • Hidden dimension

Original: [Table]

Original: • Number of layers

Original: [Table]

Original: • Number of layers for relation edge update

Original: [Table]

Original: • Batch norm

Original: [Table]

Original: • Relation dropout

Original: [Table]

Original: • Relation dropout

Original: [Table]

Original: indicate the selected hyperparameters for the ﬁnal model.

Original: Notably, we just report hyperparameterchoices which were different from standard choices.

Original: We used a training script provided by (Wanget al., 2021)

Original: [Table]

Original: which can be found here https://github.com/tata1661/PAR-NeurIPS21 .

Original: A.2 Domain shift experiment

Original: Experimental setup.

Original: For the domain shift experiment, we used the Tox21 dataset.

Original: This datasetconsists of 12,707 chemical compounds, for which measurements for up to 12 different toxic effectsare reported (Mayr et al., 2016; Huang et al., 2016a).

Original: It was published with a ﬁxed training, validationand test split.

Original: State-of-the-art supervised learning methods that have access to the full training setreach AUC performance values between 0 . 845 and 0 .

Original: 871 (Klambauer et al., 2017; Duvenaud et al.

Original: ,2015; Li et al., 2017, 2021; Zaslavskiy et al., 2019;

Original: Alperstein et al., 2019).

Original: For our evaluation, were-cast Tox21 as a few-shot learning setting and draw small support sets from the 12 tasks.

Original: Thecompared methods were pre-trained on FS-Mol and obtain small support sets from Tox21.

Original: Basedon the support sets, the methods had to predict the activities of the Tox21 test set.

Original: Note that there isa strong domain shift from drug-like molecules of FS-Mol to environmental chemicals, pesticides,food additives of Tox21.

Original: The domain shift also concerns the outputs where a shift from kinases,hydrolases, and oxidoreductases of FS-Mol to nuclear receptors and stress responses of Tox21 ispresent.

Original: Methods compared.

Original: We compared the new method MHNfs , the runner-up method IterRefLSTM ,and Similarity Search — since it has been widely used for such purposes for decades (CeretoMassagué et al., 2015).

Original: Training and evaluation.

Original: We followed the procedure of Stanley et al.

Original: (2021) for data-cleaning,preprocessing and extraction of the ﬁngerprints and descriptors used in FS-Mol.

Original: After running thecleanup step, 8,423 molecules remained for the domain shift experiments.

Original: From the training set, 8active and 8 inactive molecules per task were randomly selected to build the support set.

Original: The test setmolecules were used as query molecules.

Original: The validation set molecules were not used at all.

Original: Duringtest-time, a support set was drawn ten times for each task.

Original: Then, the performance of the models were

Original: [Table]

Original: [Table]

Original: IterRefLSTM is compared to MHNfs -CM evaluated for these support sets, using the area under precision-recall curve (AUC-PR), analogously to

Original: [Table]

Original: the FS-Mol benchmarking experiment reported as the difference to a random classiﬁer ( ∆ AUC-PR),

Original: [Table]

Original: and the area under receiver operating characteristic curve (AUC) metrics.

Original: The performance valuesreport the

Original: Results.

Original: The Hopﬁeld-based context retrieval method has signiﬁcantly outperformed the

Original: IterRefLSTM-based model ( p ∆AUC − PR -value 3 .

Original: 4 e − 5 , p AUC -value 2 .

Original: [-5 e- 6 , paired Wilcoxon test)and the Classic Similarity Search ( p ∆AUC − PR -value 2 . 
-]
Revised: [-5 e- 6 , paired Wilcoxon test)and the Classic Similarity Search ( p ∆AUC − PR -value 2 . 
-]

Original: e-

Original: 9 , p AUC -value 7 . 6 e-

Original: 10 , paired Wilcoxontest) and therefore showed robust performance on the toxicity domain, see Table A6.

Original: Notably, allmodels were trained on the FS-Mol dataset and then applied to the Tox21 dataset without adjustingany weight parameter.

Original: A.3 Details on the ablation study

Original: The MHNfs has two new main elements compared to the previous state-of-the art method IterRefLSTM, which are the context module and the cross-attention-module.

Original: In this ablation studywe aim to investigate i) the importance of all design elements, which are the context module, thecross-attention module, and the similarity module, and ii) the superiority of the cross-attention modulecompared to the IterRefLSTM module.

Original: A.3.1 Ablation study A: comparison against IterRefLSTM

Original: For a fair comparison between the cross-attention module and the IterRefLSTM we used a pruned

Original: MHN version ("MHNfs -CM") which has no context module and compared it with the IterRefLSTMmodel.

Original: The evaluation includes ﬁve training re-runs each and ten different support set samplings.

Original: The results, reported as the mean across training re-runs and support sets, can be seen in Table A7.

Original: We performed a paired Wilcoxon rank sum test for both the AUC and the ∆ AUC-PR metric.

Original: Bothp -values indicate high signiﬁcance.

Original: A.3.2 Ablation study B: all design elements

Original: We evaluate the performance of all main elements within the MHNfs, which are the context module,the cross-attention module, the similarity module and the molecule encoder.

Original: For this analysis,we start with the complete MHNfs which includes all modules and report AUC and ∆ AUC-PRperformance values.

Original: Then, we iteratively omit the individual modules, measuring whether there is a

Original: [Table]

Original: signiﬁcant performance difference with and without the module.

Original: Table A7 shows the results, whereperformance values for the full MHNfs, a MHNfs model without the context module ("MHNfs -CM")and a MHNfs module without the context and the cross-attenion module ("MHNfs -CM -CAM") isincluded.

Original: Notably, the model without the context module and without the cross-attention modulejust consists of a learned molecule encoder and the similarity module.

Original: We evaluted the impact ofthe learned molecule encoder by replacing it with a ﬁxed encoder, which maps a molecule to itsdescriptors.

Original: The model with the ﬁxed encoder is a classic chemoinformatics method which is called

Original: Similarity Search (Cereto-Massagué et al., 2015).

Original: For the evaluation, we performed ﬁve training re-runs for every model and sampled ten differentsupport sets for every task.

Original: Table A7 shows the results in terms of AUC and ∆ AUC-PR.

Original: We performedpaired Wilcoxon rank sum tests on both metrics, comparing two methods in consecutive rows in thetable.

Original: The table shows that every module has a signiﬁcant impact as omitting a module results ina signiﬁcant performance drop.

Original: The comparison between the MHNfs version without the contextmodule and without the cross-attention module with the Similarity Search showed a signiﬁcantsuperiority of the learned molecule encoder in comparison to the ﬁxed encoder.

Original: A.3.3 Ablation study C: Under domain shift on Tox21

Original: Referring to Section A.3.2, the context module and the cross-attention module showed their importance for the global architecture.

Original: This importance gets even more pronounced for the domain shiftexperiment on Tox21 as one can see in Table A8.

Original: Again, ﬁve training re-runs and ten support set draws are used for evaluation.

Original: Including the contextmodule makes a clear and signiﬁcant difference for both metrics AUC and ∆ AUC-PR.

Original: A.4 Generalization to different support set sizes

Original: In this section, we test the ability of MHNfs to generalize to different support set sizes.

Original: Duringtraining in the FS-Mol benchmarking setting, the MHNfs model has access to support sets of size16.

Original: However, at inference, the support set size might be different.

Original: Figure A4 provides performanceestimates of the support-set-size-16 MHNfs models on other support set sizes.

Original: Note that the estimatescould be seen as approximate lower bounds of the predictive performance on settings with differentsupport set sizes (y-axis labels).

Original: For a model used in production or in a real-world drug discoverysetting, MHNfs should be trained with varying support set sizes that resemble the distribution of realdrug discovery projects.

Original: A.5 Generalization to different context sets

Original: In this section, we test the ability of MHNfs to generalize to different context sets.

Original: While the FS-Moltraining split is used as a context during training, we assessed whether our model is robust to differentcontext sets for inference.

Original: To this end we preprocessed the GEOM dataset (Axelrod and GomezBombarelli, 2022) from which we used 100,000 molecules that passed all pre-processing checks.

Original: From this set, we sample 10,000 molecules as context set for MHNfs.

Original: Because GEOM containsdrug-like molecules, similar to FS-Mol the predictive performance remains stable (see Table A9).

Original: [Figure]

Original: [Figure]

Original: A U C - P R ( t e s t s e t )

Original: [Table]

Original: A.6 Details and insights on the context module

Original: The context module replaces the initial representations of query and support set molecules by aretrieval from the context set.

Original: The context set is a large set of molecules and covers a large chemicalspace.

Original: The context module learns how to replace the initial molecule embeddings such that thecontext-enriched representations are put in relation to this large chemical space and still containsall necessary information for the similarity-based prediction part.

Original: Figure A5 shows the effect of thecontext module for the MHNfs model.

Original: Extreme initial embeddings, such as the purple embeddingon the right, are pulled more into the known chemical space, represented by the context molecules.

Original: Notably, the replacement described above is a soft replacement, because also the initial embeddingscontribute to the context-enriched representations due to skip-connections.

Original: A.7 Reinforcing the covariance structure in the data using modern

Original: Hopﬁeld

Original: networks

Original: We follow the argumentation of (Fürst et al., 2021, Theorem A3) that retrieval from an associativememory of a MHN reinforces the covariance structure.

Original: Let us assume that we have one molecule embedding from the query set m ∈ R d and one moleculeembedding from the support set x ∈ R d and both have been enriched with the context module withmemory C ∈ R d × M (ignoring linear mappings):

Original: [Equation]

Original: Then the similarity of the retrieved representations as measured by the dot product can be expressedin terms of covariances:

Original: [Figure] [Equation] where c is the row mean of C andfollowing the weighted covariances are used: [Equation]
Revised: (c) Did you include the estimated hourly wage paid to participants and the total amount

Original: J m : R M

Original: (cid:55)→ R M × M is a mean Jacobian function of the softmax (Fürst et al., 2021, Eq.(A172)).

Original: The Jacobian J of p = softmax( β a ) is J( β a ) = β (cid:0) diag( p ) − pp T (cid:1) .b T J( β a ) b = β b T (cid:0) diag( p ) − p p T (cid:1) b = β  (cid:88) i p i b
Revised: spent on participant compensation?

Original: [Equation]

Original: i − (cid:32) (cid:88)

Original: i p i

Original: b i (cid:33) 2 

Original:  , (A14) this is the second moment minus the mean squared, which is the variance.

Original: Therefore, b T J

Original: Note that we can express the enriched representations using these covariance functions:

Original: [Equation]
Revised: [N/A]

Original: which connects retrieval from MHNs with reinforcing the covariance structure of the data.