Original: C ONTEXT - ENRICHED MOLECULE
Revised: C ONTEXT - ENRICHED MOLECULE

Original: REPRESENTATIONS
Revised: REPRESENTATIONS

Original: IMPROVE FEW - SHOT DRUG DISCOVERY
Revised: IMPROVE FEW - SHOT DRUG DISCOVERY

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: A central task in computational drug discovery is to construct models from known active molecules to ﬁnd further promising molecules for subsequent screening.
Revised: A central task in computational drug discovery is to construct models from known active 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-shot learning methods have the potential to improve the effectiveness of this critical phase of the drug discovery process.
Revised: Therefore, few-shot learning methods have the potential to improve the effectiveness of this critical phase of the drug discovery process.

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

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

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

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

Original: Intuitively, this enrichment step is analogous to a human expert who would associate a given molecule with familiar molecules whose properties are known.
Revised: Intuitively, this enrichment step is analogous to a human expert who would associate a given molecule with familiar molecules whose properties are known.

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

Original: Our approach is compared with other few-shot methods for drug discovery on the FS-Mol benchmark dataset.
Revised: Our approach is compared with other few-shot methods for drug discovery on the FS-Mol benchmark dataset.

Original: On FS-Mol, our approach outperforms all compared methods and therefore sets a new state-of-the art for few-shot learning in drug discovery.
Revised: On FS-Mol, our approach outperforms all compared methods and therefore sets a new state-of-the art for few-shot learning in drug discovery.

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

Original: In a domain shift experiment, we further demonstrate the robustness of our method.
Revised: In a domain shift experiment, we further demonstrate the robustness of our method.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

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

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

Original: Deep learning methods have recently been shown to reduce time 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 costs of this process (Chen et al., 2018; Walters and Barzilay, 2021).

Original: They diminish the required number of both wet-lab measurements and molecules that must be synthesized (Merk et al., 2018; Schneider et al., 2020).
Revised: They diminish the required number of both wet-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 molecular information about the ligands after being trained on a large training set.
Revised: However, as of now, deep learning approaches use only the molecular information about the ligands after being trained on a large training set.

Original: At inference time, they yield highly 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; Seidl et al., 2022).
Revised: At inference time, they yield highly 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; Seidl et al., 2022).

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-performing predictive models for activity prediction tasks of ChEMBL have been trained with an average of 3, activity points per task – i.e., drug target – by Mayr et al.
Revised: For example, well-performing predictive models for activity prediction tasks of ChEMBL have been trained with an average of 3, activity points per task – i.e., drug target – by Mayr et al.

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

Original: Wu et al.
Revised: Wu et al.

Original: (2018) published a large scale benchmark for molecular machine learning, including prediction models for the
Revised: (2018) published a large scale benchmark for molecular machine learning, including prediction models for the

Original: SIDER dataset (Kuhn et al., 2016) with an average of 5,187 data points, Tox21 (Huang et al., 2016b;
Revised: SIDER dataset (Kuhn et al., 2016) with an average of 5,187 data points, Tox21 (Huang et al., 2016b;

Original: Mayr et al., 2016)
Revised: Mayr et al., 2016)

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

Original: However, for typical drug design projects, the amount of available measurements is very limited (Stanley et al., 2021;
Revised: However, for typical drug design projects, the amount of available measurements is very limited (Stanley et al., 2021;

Original: Waring et al., 2015; Hochreiter et al., 2018), since in-vitro experiments are expensive and time-consuming.
Revised: Waring et al., 2015; Hochreiter et al., 2018), since in-vitro experiments are expensive and time-consuming.

Original: Therefore, methods that need only few measurements to build precise prediction models are desirable.
Revised: Therefore, methods that need only few measurements to build precise prediction models are desirable.

Original: This problem – i.e., the challenge of learning from few data points – is the focus of machine learning areas like meta-learning (Schmidhuber, 1987; Bengio et al.,
Revised: This problem – i.e., the challenge of learning from few data points – is the focus of machine learning areas like meta-learning (Schmidhuber, 1987; Bengio et al.,

Original: Hochreiter et al., 2001) and few-shot learning (Miller et al., 2000; Bendre et al., 2020; Wang et al., 2020).
Revised: Hochreiter et al., 2001) and few-shot learning (Miller et al., 2000; Bendre et al., 2020; Wang et 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-shot learning methods have been predominantly developed and tested on image datasets (Bendre et al., 2020; Wang et al., 2020) and have recently been adapted to drug discovery problems (Altae-Tran et al., 2017;
Revised: Few-shot learning methods have been predominantly developed and tested on image datasets (Bendre et al., 2020; Wang et al., 2020) and have recently been adapted to drug discovery problems (Altae-Tran et al., 2017;

Original: Guo et al., 2021; Wang et al., 2021; Stanley et al., 2021; Chen et al., 2022).
Revised: Guo et al., 2021; Wang et al., 2021; Stanley et al., 2021; Chen et al., 2022).

Original: They are usually categorized into 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 main approach (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, more diverse data points (Chen et al., 2020; Zhao et al., 2019; Antoniou and Storkey, 2019).
Revised: a) Data-augmentation-based approaches augment the available samples and generate new, more diverse data points (Chen et al., 2020; Zhao et al., 2019; Antoniou and Storkey, 2019).

Original: b) Embedding-based and nearest neighbour approaches learn embedding space representations.
Revised: b) Embedding-based and nearest neighbour approaches learn embedding space representations.

Original: Predictive models can then be constructed from only few data points by comparing these embeddings.
Revised: Predictive models can then be constructed from only few data points by comparing these embeddings.

Original: For example, in Matching Networks (Vinyals et al., 2016) an attention mechanism that relies on embeddings is the basis for the predictions.
Revised: For example, in Matching Networks (Vinyals et al., 2016) an attention mechanism that relies on embeddings is the basis for the predictions.

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

Original: c) Optimization-based or ﬁne-tuning methods utilize a meta-optimizer that focuses on efﬁciently navigating the parameter space.
Revised: c) Optimization-based or ﬁne-tuning methods 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 be adapted to a novel task by few optimization steps (Finn et al., 2017).
Revised: For example, with MAML the meta-optimizer learns initial weights that can be adapted to a novel task by few optimization steps (Finn et al., 2017).

Original: Most of these approaches have already been applied to few-shot drug discovery (see Section 4).
Revised: Most of these approaches have already been applied to few-shot drug discovery (see Section 4).

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 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).

Original: We hypothesize that the underperformance of these methods stems from disregarding the context – both in terms of similar molecules and similar activities.
Revised: We hypothesize that the underperformance of these methods stems from disregarding the context – both in terms of similar molecules and similar activities.

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

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

Original: In data-scarce situations, humans extract co-occurrences and covariances by associating current perceptions with memories (Bonner and Epstein, 2021; Potter, 2012).
Revised: In data-scarce situations, humans extract co-occurrences and covariances by associating current perceptions with memories (Bonner and Epstein, 2021; Potter, 2012).

Original: When we show a small set of active molecules to a human expert in drug discovery, the expert associates them with known molecules to suggest further active molecules (Gomez, 2018;
Revised: When we show a small set of active molecules to a human expert in drug discovery, the expert associates them with known molecules 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 for few-shot learning uses associative memories to extract co-occurrences and the covariance structure of the original data and to amplify them in the representations (Fürst et al., 2022).
Revised: In an analogous manner, our novel concept for few-shot learning uses associative memories to extract co-occurrences and the covariance structure of the original data and to amplify them in the representations (Fürst et al., 2022).

Original: We use Modern
Revised: We use Modern

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

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

Original: Those retrieved representations have ampliﬁed co-occurrences and covariance structures, while peculiarities and spurious co-occurrences of the query and support set molecules are averaged out.
Revised: Those retrieved representations have ampliﬁed co-occurrences and covariance structures, while peculiarities and spurious co-occurrences of the query and support set molecules are averaged 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 FS-Mol.
Revised: • We achieve a new state-of-the-art on the benchmarking dataset FS-Mol.

Original: • We introduce a novel concept to enrich the molecule representations with context by associating them with a large set of context molecules.
Revised: • 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 all other 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 a domain shift experiment to further demonstrate the effectiveness of our new method.

Original:  P ROBLEM SETTING
Revised:  P ROBLEM 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 or activity ˆ y , given a representation m of an input molecule from a chemical space M .

Original: We consider machine learning models ˆ y =
Revised: We consider machine 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 training set.

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 molecule encoder can process different symbolic or low-level representations of molecules, such as molecular descriptors (Bender et al., 2004; Unterthiner et al., 2014; Mayr et al., 2016), SMILES (Weininger, 1988; Mayr et al., 2018;
Revised: The molecule encoder can process different symbolic or low-level representations of molecules, such as molecular descriptors (Bender et al., 2004; Unterthiner et al., 2014; Mayr et al., 2016), SMILES (Weininger, 1988; Mayr et al., 2018;

Original: Winter et al., 2019; Segler et al., 2018a), or molecular graphs (Merkwirth and Lengauer, 2005; Kearnes et al., 2016;
Revised: Winter et al., 2019; 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-trained on related property prediction tasks.
Revised: Yang et al., 2019; Jiang et al., 2021) and can be 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 of molecules { x 1 , . . .
Revised: For few-shot learning, the goal is to select a high-quality predictive model based on a small set of molecules { 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 are usually assumed to be binary y n ∈ {− 1 , 1 } , corresponding to the molecule being inactive or active.
Revised: The measurements are usually assumed to be binary y n ∈ {− 1 , 1 } , corresponding to the molecule being inactive or active.

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

Original: The goal is to construct a model that correctly predicts y for an x that is not in the support set – in other words, a model that generalizes well.
Revised: The goal is to construct a model that correctly predicts y for an x that is not in the support set – in other words, a model that generalizes well.

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

Original: These approaches learn the parameters w of the model g w from the support set in a supervised manner.
Revised: These approaches 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 generalize well to new data.
Revised: Therefore, few-shot learning methods are necessary to construct models from the support set that generalize well to new data.

Original:  MHN FS : H OPFIELD - BASED MOLECULAR CONTEXT ENRICHMENT FOR FEW - SHOT DRUG DISCOVERY
Revised:  MHN FS : H OPFIELD - BASED MOLECULAR CONTEXT ENRICHMENT FOR FEW - SHOT DRUG DISCOVERY

Original: We aim at increasing the generalization capabilities of few-shot learning methods in drug discovery by enriching the molecule representations with molecular context.
Revised: We aim at increasing the generalization capabilities of few-shot learning methods in drug discovery by 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 – includes information about a large chemical space.
Revised: In comparison to the support set, which encodes information about the task, the context set – i.e. a large set of molecules – includes information about a large chemical space.

Original: The query and the support set molecules perform a retrieval from the context set and thereby enrich their representations.
Revised: The query and the support set molecules perform a retrieval from the context set and thereby enrich their representations.

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

Original: 3. M ODEL ARCHITECTURE
Revised: 3. M ODEL ARCHITECTURE

Original: We propose an architecture which consists of three consecutive modules.
Revised: We propose an architecture which consists of three consecutive modules.

Original: The ﬁrst module – a) the context 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 a large set of molecules.

Original: 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 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 the support set molecules.

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

Original: [Equation]
Revised: [Equation]

Original: module
Revised: module

Original: [Equation]
Revised: [Equation]

Original: -attention
Revised: -attention

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 set molecules 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 set molecules 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 molecule embeddings, y are the support set labels, and ˆ y is the prediction for the query molecule.
Revised: C ∈ R d × M is a large set of stacked molecule embeddings, y are the support set labels, and ˆ y is the prediction for the query molecule.

Original: Square brackets indicate concatenation, for example [ m (cid:48) , X (cid:48) ] is a matrix with N + 1 columns.
Revised: Square brackets indicate concatenation, for example [ m (cid:48) , X (cid:48) ] is a matrix with N + 1 columns.

Original: The modules f CM , f CAM , and f SM are detailed in the paragraphs below.
Revised: The modules f CM , f CAM , and f SM are detailed in the paragraphs below.

Original: An overview of our architecture is given in Figure 1.
Revised: An overview of our architecture is given 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 ( . ) and layer normalization (Ba et al., 2016), which are not shown in Figure1.
Revised: and f CAM ( . ) and layer normalization (Ba et al., 2016), which are not shown in Figure1.

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 ) , the support set molecules x n = f ME ( x n ) , and the context molecules c m =
Revised: f ME ( m ) , the support set molecules x n = f ME ( x n ) , and the context molecules c m =

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

Original: There are many possible 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
Revised: There are many possible 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

Original: (Dahl et al., 2014; Mayr et al., 2016; 2018).
Revised: (Dahl et al., 2014; Mayr et al., 2016; 2018).

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

Original: [Equation]
Revised: [Equation]

Original: 3. C ONTEXT MODULE (CM)
Revised: 3. C ONTEXT MODULE (CM)

Original: The context module associates the query and support set molecules with a large set of context molecules, 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 context molecules, and represents them as weighted average of context molecule embeddings.

Original: The context module is realised by a continuous Modern Hopﬁeld Network (MHN) (Ramsauer et al., 2021).
Revised: The context module is realised by a continuous Modern Hopﬁeld Network (MHN) (Ramsauer et al., 2021).

Original: An
Revised: An

Original: MHN is a content-addressable associative memory which can be built into deep learning architectures.
Revised: MHN is a content-addressable associative memory which can be built into deep learning architectures.

Original: There exists an analogy between the energy update of MHNs and the attention mechanism of
Revised: There exists an analogy between the energy update of MHNs and the attention mechanism of

Original: Transformers (Vaswani et al., 2017; Ramsauer et al., 2021).
Revised: Transformers (Vaswani et al., 2017; Ramsauer et al., 2021).

Original: MHNs are capable of storing and retrieving patterns from a memory M ∈ R e
Revised: MHNs are capable of storing and retrieving patterns from a memory M ∈ R e

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

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

Original: [Equation]
Revised: [Equation]

Original: 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 (Seidl et al., 2022; Widrich et al., 2020).

Original: We use this mechanism in the form of a Hopﬁeld layer , which ﬁrst maps raw patterns to an associative space 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 associative space using linear transformations, and uses multiple simultaneous queries Ξ ∈ R d × N :

Original: [Equation]
Revised: [Equation]

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, softmax is applied column-wise, and β is a hyperparameter.
Revised: e × d are trainable parameters of the Hopﬁeld layer, softmax is applied column-wise, and β is a hyperparameter.

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

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

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

Original: The context module of our new architecture uses a Hopﬁeld layer, where the query patterns are the embeddings 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 the embeddings of the query molecule m and the support set molecules X .

Original: The memory is composed of embeddings 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 of embeddings 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 the retrieved embeddings, which are weighted averages of context molecule embeddings:
Revised: Then the original embeddings m and X are replaced by the retrieved embeddings, which are weighted averages of context molecule embeddings:

Original: [Equation]
Revised: [Equation]

Original: Hopﬁeld
Revised: Hopﬁeld

Original: [Equation]
Revised: [Equation]

Original: m
Revised: m

Original: [Equation]
Revised: [Equation]

Original: and X (cid:48)
Revised: and X (cid:48)

Original: [Equation]
Revised: [Equation]

Original: Hopﬁeld( X , C ) .
Revised: Hopﬁeld( X , C ) .

Original: [Equation]
Revised: [Equation]

Original: This retrieval step reinforces the covariance structure of the retrieved representations (see Appendix A.6), which usually enhances robustness of the models (Fürst et al., 2022) by removing noise.
Revised: This retrieval step reinforces the covariance structure of the retrieved representations (see Appendix A.8), which usually enhances robustness of the models (Fürst et al., 2022) by removing noise.

Original: Note that the embeddings of the query and the support set molecules have not yet inﬂuenced each other.
Revised: Note that the embeddings of the query and the support set molecules have not yet inﬂuenced each 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: Exemplary retrievals from the context module are included in Appendix A.5.
Revised: Exemplary retrievals from the context module are included in Appendix A.7.

Original: 3. C ROSS - ATTENTION MODULE (CAM)
Revised: 3. C ROSS - 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 information exchange between query and support set molecules.
Revised: (2017) showed that the representations of the molecules can be enriched, if the architecture allows information exchange 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 an iterative fashion being aware of each other.
Revised: (2017) uses an attentionenhanced LSTM variant, which updates the query and the support set molecule representations in an iterative fashion being aware of each other.

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

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

Original: [Equation]
Revised: [Equation]

Original: (9) where [ m (cid:48) , X (cid:48) ]
Revised: (9) 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, 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 representations m (cid:48)(cid:48) and X (cid:48)(cid:48) are passed to the similarity module.

Original: 3. S IMILARITY MODULE (SM)
Revised: 3. S IMILARITY 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 of a 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 of a 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 is predicted, building a weighted mean over the support set labels:
Revised: Based on these similarity values, the activity for the query molecule is predicted, building a weighted mean over the support set labels:

Original: [Equation]
Revised: [Equation]

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

Original: x
Revised: x

Original: [Equation]
Revised: [Equation]

Original: n )
Revised: n )

Original: = m (cid:48)(cid:48)
Revised: = 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 balancing strategy for the labels y (cid:48) n =
Revised: Note that we use a balancing strategy for the labels y (cid:48) n =

Original: [Equation]
Revised: [Equation]

Original: else , where N A is the number of actives and N
Revised: else , where N A is the number of actives and N

Original: I is the number of inactives of the support set.
Revised: I is the number of inactives of the support set.

Original: 3. A RCHITECTURE , HYPERPARAMETER SELECTION , AND TRAINING DETAILS
Revised: 3. A RCHITECTURE , HYPERPARAMETER SELECTION , AND TRAINING DETAILS

Original: Hyperparameters.
Revised: Hyperparameters.

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

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

Original: The molecule encoder consists of a single layer with output size d = 1024 and SELU activation (Klambauer et al., 2017).
Revised: The molecule encoder consists 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 set to 512 and β = 1 / √ e .
Revised: The dimension e of the association space is set to 512 and β = 1 / √ e .

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

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

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

Original: The SM uses τ = 32 .
Revised: The SM uses τ = 32 .

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

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

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

Original: We provide considered and selected hyperparameters in Appendix A.1.6.
Revised: We provide considered and selected hyperparameters in Appendix A.1.6.

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 use a learning rate scheduler which includes a warm up phase, followed by a section with a constant learning rate, which is 0 . 0001 , and a third phase in which the learning rate steadily decreases.
Revised: We use a learning rate scheduler which includes a warm up phase, followed by a section with a constant learning rate, which is 0 . 0001 , and a third phase in which the learning rate steadily decreases.

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

Original: The molecule encoder has a dropout with rate 0 .
Revised: The molecule encoder has 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 Appendix A.1.6).

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 hours were consumed for this work.
Revised: In total, roughly 15,000 GPU hours were consumed for this work.

Original:  R ELATED WORK
Revised:  R ELATED WORK

Original: Several approaches to few-shot learning in drug discovery have been suggested (Altae-Tran et al., 2017; Nguyen et al., 2020;
Revised: Several approaches to few-shot learning in drug discovery have been suggested (Altae-Tran et al., 2017; Nguyen et al., 2020;

Original: Guo et al., 2021; Wang et al., 2021).
Revised: Guo et al., 2021; Wang et al., 2021).

Original: Nguyen et al.
Revised: Nguyen et al.

Original: (2020) evaluated the applicability of MAML and its variants to graph neural networks (GNNs) and (Guo et al., 2021) also combine GNNs and meta-learning.
Revised: (2020) evaluated the applicability of MAML and its variants to graph neural networks (GNNs) and (Guo et al., 2021) also combine GNNs and meta-learning.

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

Original: (2017) suggested an approach called
Revised: (2017) suggested an approach called

Original: Iterative Reﬁnement Long Short-Term Memory, in which query and support set embeddings can share information and update their embeddings.
Revised: Iterative Reﬁnement Long Short-Term Memory, in which query and support set embeddings can share information and update their embeddings.

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

Original: Chen et al.
Revised: Chen et al.

Original: (2022) propose to adaptively learn kernels and apply their method to few-shot drug discovery with predictive performance for larger support set sizes.
Revised: (2022) propose to adaptively learn kernels and apply their method to few-shot drug discovery with predictive performance for larger support set sizes.

Original: Recently, Stanley et al.
Revised: Recently, Stanley et al.

Original: (2021) generated a benchmark dataset for few-shot learning methods in drug discovery and provided some baseline results.
Revised: (2021) generated a benchmark dataset for few-shot learning methods in drug discovery and provided some baseline results.

Original: Many successful deep neural network architectures use external memories, such as the neural Turing machine (Graves et al., 2014), memory networks (Weston et al., 2014), end-to-end memory networks (Sukhbaatar et al., 2015).
Revised: Many successful deep neural network architectures use external memories, such as the neural Turing machine (Graves et al., 2014), memory networks (Weston et al., 2014), end-to-end memory networks (Sukhbaatar et al., 2015).

Original: Recently, the connection between continuous modern Hopﬁeld networks (Ramsauer et al., 2021), which are content-addressable associative memories, and Transformer architectures (Vaswani et al., 2017) has been established.
Revised: Recently, the connection between continuous modern Hopﬁeld networks (Ramsauer et al., 2021), which are content-addressable associative memories, and Transformer architectures (Vaswani et al., 2017) has been established.

Original: We refer to Le (2021) for an extensive overview of memory-based architectures.
Revised: We refer to Le (2021) for an extensive overview of memory-based architectures.

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

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

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

Original:  E XPERIMENTS 5. B ENCHMARKING ON FS-M OL
Revised:  E XPERIMENTS 5. B ENCHMARKING ON FS-M OL

Original: Experimental setup.
Revised: Experimental setup.

Original: Recently, the dataset FS-Mol (Stanley et al., 2021) was proposed to benchmark few-shot learning methods in drug discovery.
Revised: Recently, the dataset FS-Mol (Stanley et al., 2021) was proposed to benchmark few-shot learning methods in drug discovery.

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

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

Original: The dataset is well balanced as the mean ratio of active and inactive molecules is close to 1.
Revised: The dataset is well balanced as the mean ratio of active and inactive molecules is close 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.
Revised: Stanley et al.

Original: (2021) precomputed extended connectivity ﬁngerprints (ECFP)
Revised: (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 of the input molecules, such as the molecular graph, we used a concatenation of these ECFPs and
Revised: While methods would be allowed to use other representations of the input molecules, such as the molecular graph, we used a concatenation of these ECFPs and

Original: RDKit-based descriptors.
Revised: RDKit-based descriptors.

Original: We use the main benchmark setting of FS-Mol with support set size 16,
Revised: We use the main benchmark setting of FS-Mol with support set size 16,

Original: [Table]
Revised: [Table]

Original: which is close to the 5- and 10-shot settings in computer vision, and stratiﬁed random split (Stanley et al., 2021, Table 2) for a fair method comparison (see also Section A.3).
Revised: which is close to the 5- and 10-shot settings in computer vision, and stratiﬁed random split (Stanley et al., 2021, Table 2) for a fair method comparison (see also Section A.5).

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 were compared against each other.

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

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

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

Original: Similarity search is a standard chemoinformatics technique, used in situations with single or few known actives.
Revised: Similarity search is a standard chemoinformatics technique, used in situations with single or few known actives.

Original: In the simplest case, the search ﬁnds similar molecules by computing a ﬁngerprint or descriptor-representation of the molecules and using a similarity measure k ( ., . )
Revised: In the simplest case, the search ﬁnds similar molecules by computing a ﬁngerprint or descriptor-representation of the molecules and using a similarity measure k ( ., . )

Original: – such as Tanimoto Similarity (Tanimoto, 1960).
Revised: – such as Tanimoto Similarity (Tanimoto, 1960).

Original: Thus, Similarity Search, as used in chemoinformatics, can be formally written as ˆ y = 1 /N (cid:80)
Revised: Thus, Similarity 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 descriptor calculation.

Original: A natural extension of Similarity Search with ﬁxed chemical descriptors is Neural Similarity Search or Siamese networks (Koch et al., 2015), which extend the classic similarity search by learning a molecule encoder:
Revised: A natural extension of Similarity Search with ﬁxed chemical descriptors is Neural Similarity Search or Siamese networks (Koch et al., 2015), which extend the classic similarity search by learning a molecule encoder:

Original: ˆ y = σ (cid:16) τ − 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
Revised: Furthermore, we re-implemented the

Original: IterRefLSTM (Altae-Tran et al., 2017) in PyTorch.
Revised: IterRefLSTM (Altae-Tran et al., 2017) in PyTorch.

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

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

Original: Second, an attention-enhanced LSTM variant, the actual IterRefLSTM , iteratively updates 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 actual IterRefLSTM , iteratively updates the query and support set molecules, enabling information sharing between the molecules: [ m (cid:48) , X (cid:48) ]

Original: =
Revised: =

Original: IterRefLSTM L ([ m , X ]) , where the hyperparameter L controls the number of iteration steps of the IterRefLSTM .
Revised: IterRefLSTM L ([ m , X ]) , where the hyperparameter L controls the number of 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 based on 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 =
Revised: These representations are then used for the ﬁnal prediction: ˆ y = (cid:80) Ni =

Original: 1 a
Revised: 1 a

Original: i y i .
Revised: i y i .

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

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

Original: The method GNN-ST is a graph neural network (Stanley et al., 2021; Gilmer et al., 2017) that is trained from scratch for each task.
Revised: The method GNN-ST is a graph neural network (Stanley et al., 2021; Gilmer et al., 2017) that is trained from scratch for each task.

Original: The GNN-MT uses a two step strategy:
Revised: The GNN-MT uses a two step strategy:

Original: First, the model is pretrained on a large dataset on related tasks; second, an output layer is constructed to the few-shot task via linear probing (Stanley et al., 2021; Alain and Bengio, 2016).
Revised: First, the model is pretrained on a large dataset on related tasks; second, an output layer is constructed to the few-shot task via linear probing (Stanley et al., 2021; Alain and Bengio, 2016).

Original: The Molecule Attention
Revised: The Molecule Attention

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

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

Original: Notably, GNN-MAML also can be seen as a proxy for Meta-MGNN (Guo et al., 2021), which enriches the gradient update step in the outer loop of the MAML-framework by an attention mechanism and uses an additional atom type prediction loss and a bond reconstruction
Revised: Notably, GNN-MAML also can be seen as a proxy for Meta-MGNN (Guo et al., 2021), which enriches the gradient update step in the outer loop of the MAML-framework by an attention mechanism and uses an additional atom type prediction loss and a bond reconstruction

Original: [Figure]
Revised: [Figure]

Original: This indicates that our two newly introduced modules, CMand CAM play a crucialrole loss.
Revised: This indicates that our two newly introduced modules, CM play a crucial loss.

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

Original: In this embedding space, prototypical representations of each class are built by taking the mean across all related support set molecules for each class (details in Appendix A.1.4).
Revised: In this embedding space, prototypical representations of each class are built by taking the mean across all related support set molecules for each class (details in Appendix A.1.4).

Original: The PAR model (Wang et al., 2021) includes a GNN which creates initial molecule embeddings.
Revised: The PAR model (Wang et al., 2021) includes a GNN which creates initial molecule embeddings.

Original: These molecule embeddings are then enriched by an attention mechanism.
Revised: These molecule embeddings are then enriched by an attention mechanism.

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

Original: The
Revised: The

Original: PAR model has shown good results for datasets which just include very few tasks such as Tox (Wang et al., 2021).
Revised: PAR model has shown good results for datasets which just include very few tasks such as Tox (Wang et al., 2021).

Original: Chen et al.
Revised: Chen et al.

Original: (2022) suggest a framework for learning deep kernels by interpolating between meta-learning and conventional deep kernels, which results in the ADKF-IFT model.
Revised: (2022) suggest a framework for learning deep kernels by interpolating between meta-learning and conventional deep kernels, which results in the ADKF-IFT model.

Original: The model has exhibited especially high performance for large support set sizes.
Revised: The model has exhibited especially high performance for large support set sizes.

Original: For all methods the most important hyperparameters were adjusted on the validation tasks of FS-Mol.
Revised: For all methods the most 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).
Revised: BSD license).

Original: We used PyTorch Lightning (Falcon et al., 2019, Apache 2.0 license) as a framework for training and test logic, hydra for conﬁg ﬁle handling (Yadan, 2019, Apache 2.0 license) and
Revised: We used PyTorch Lightning (Falcon et al., 2019, Apache 2.0 license) as a framework for training and test logic, hydra for conﬁg ﬁle handling (Yadan, 2019, Apache 2.0 license) and

Original: Weights & Biases (Biewald, 2020, MIT license) as an experiment tracking tool.
Revised: Weights & Biases (Biewald, 2020, MIT license) as an experiment tracking tool.

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

Original: Each model was evaluated ten times by drawing support sets with ten different seeds.
Revised: Each model was evaluated ten times by drawing support sets with ten 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 reported across tasks.
Revised: The standard error is reported across 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 and
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 and

Original: ADKF-IFT.
Revised: ADKF-IFT.

Original: MHNfs has outperformed all 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), and the PAR model ( p -value <1.0e-8, paired Wilcoxon test).
Revised: MHNfs has outperformed all 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), and the PAR model ( p -value <1.0e-8, paired Wilcoxon test).

Original: 5. A BLATION STUDY
Revised: 5. A BLATION STUDY

Original: MHNfs has two new main components compared to the most similar previous state-of-the-art method
Revised: MHNfs has two new main components compared to the most similar previous state-of-the-art method

Original: IterRefLSTM : i) the context module, and ii) the cross-attention module which replaces the LSTMlike module.
Revised: IterRefLSTM : i) the context module, and ii) the cross-attention module which replaces the LSTMlike module.

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

Original: Therefore, we compared MHNfs to a method that does not have the context module ("MHNfs -CM") and to a method that does not have the context module and uses an LSTM-like module instead of the CAM ("MHNfs -CM (cid:10) (CAM,IterRefLSTM)").
Revised: Therefore, we compared MHNfs to a method that does not have the context module ("MHNfs -CM") and to a method that does not have the context module and uses an LSTM-like module instead of the CAM ("MHNfs -CM (cid:10) (CAM,IterRefLSTM)").

Original: For the ablation study, we used all 5 training reruns and evaluatedtimes on the test set with different support sets.
Revised: For the ablation study, we used all 5 training reruns and evaluatedtimes on the test set with different support sets.

Original: The results of this ablation steps are presented in
Revised: The results of this ablation steps are presented in

Original: Figure 2.
Revised: Figure 2.

Original: Both removing the CM and exchanging the CAM with the IterRefLSTM module were detrimental for the performance of the method ( p -value 0.002 and 1 .
Revised: Both removing the CM and exchanging the CAM with the IterRefLSTM module were detrimental for the performance of the method ( p -value 0.002 and 1 .

Original: 72 e − 7 , respectively; paired
Revised: 72 e − 7 , respectively; paired

Original: Wilcoxon test).
Revised: Wilcoxon test).

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

Original: [Table]
Revised: [Table]

Original: Appendix A.2.2 contains a second ablation study that examines the overall effects of the context, the cross-attention, the similarity module, and the molecule encoder of MHNfs .
Revised: 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 encoder of MHNfs .

Original: 5. D OMAIN SHIFT EXPERIMENT
Revised: 5. D OMAIN SHIFT EXPERIMENT

Original: Experimental setup.
Revised: Experimental setup.

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

Original: This dataset consists of 12,707 chemical compounds, for which measurements for up to 12 different toxic effects are reported (Mayr et al., 2016; Huang et al., 2016a).
Revised: The Tox21 dataset consists of 12,707 chemical compounds, for which measurements for up to 12 different toxic effects are reported (Mayr et al., 2016; Huang et al., 2016a).

Original: It was published with a ﬁxed training, validation and test split.
Revised: It was published with a ﬁxed training, validation and test split.

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

Original: 871 (Klambauer et al., 2017; Duvenaud et al., 2015; Li et al., 2017; 2021; Zaslavskiy et al., 2019;
Revised: 871 (Klambauer et al., 2017; Duvenaud et al., 2015; Li et al., 2017; 2021; Zaslavskiy et al., 2019;

Original: Alperstein et al., 2019).
Revised: Alperstein et al., 2019).

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

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

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

Original: Note that there is a strong domain shift from drug-like molecules of FS-Mol to environmental chemicals, pesticides, food additives of Tox21.
Revised: Note that there is a 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 is present.
Revised: 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 is present.

Original: Methods compared.
Revised: 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).
Revised: 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.
Revised: Training and evaluation.

Original: We followed the procedure of Stanley et al.
Revised: 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.
Revised: (2021) for data-cleaning, preprocessing and extraction of the ﬁngerprints and descriptors used in FS-Mol.

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

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

Original: The test set molecules were used as query molecules.
Revised: The test set molecules were used as query molecules.

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

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

Original: Then, the performance of the models were evaluated for these support sets, using the area under precision-recall curve (AUC-PR), analogously to the FS-Mol benchmarking experiment reported as the difference to a random classiﬁer ( ∆ AUC-PR), and the area under receiver operating characteristic curve (AUC) metrics.
Revised: Then, the performance of the models were evaluated for these support sets, using the area under precision-recall curve (AUC-PR), analogously to the FS-Mol benchmarking experiment reported as the difference to a random classiﬁer ( ∆ AUC-PR), and the area under receiver operating characteristic curve (AUC) metrics.

Original: The performance values report the mean over all combinations regarding the training reruns and the support set sampling iterations.
Revised: The performance values report the mean over all combinations regarding the training reruns and the support set sampling iterations.

Original: Error bars indicate the standard deviation.
Revised: Error bars indicate the standard deviation.

Original: Results.
Revised: Results.

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

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

Original: 4 e − 5 , p AUC -value 2 .
Revised: 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-
Revised: e-

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

Original: 10 , paired Wilcoxon test) and therefore showed robust performance on the toxicity domain, see Table 2.
Revised: 10 , paired Wilcoxon test) and therefore showed robust performance on the toxicity domain, see Table 2.

Original: Notably, all models were trained on the FS-Mol dataset and then applied to the Tox21 dataset without adjusting any weight parameter.
Revised: Notably, all models were trained on the FS-Mol dataset and then applied to the Tox21 dataset without adjusting any weight parameter.

Original:  C ONCLUSION
Revised:  C ONCLUSION

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

Original: In a benchmarking experiment the architecture outperformed all other methods and in a domain shift study the robustness and transferability has been assessed.
Revised: In a benchmarking experiment the architecture outperformed all other methods and in a domain shift study the robustness and transferability has been assessed.

Original: We envision that the context module can be applied to many different areas, enriching learned representations analogously to our work.
Revised: We envision that the context module can be applied to many different areas, enriching learned representations analogously to our work.

Original: For discussion, see A.7.
Revised: For discussion, see A.9.

Original:  Appendix
Revised: A Appendix

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table] 
Revised: [Table] [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original:  context module
Revised: support set sizes A.6 Generalization to different context sets A.7 Details and insights on the context module

Original: A.6 Reinforcing the covariance structure in the data using modern Hopﬁeld networks . 
Revised: A.8 Reinforcing the covariance structure in the data using modern Hopﬁeld networks . A

Original: A.7 Discussion, limitations and broader inpact.
Revised: .9 Discussion, limitations and broader inpact

Original: A A PPENDIX A. D ETAILS ON METHODS
Revised: A A PPENDIX A. D ETAILS ON METHODS

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

Original: , ( x N , y N ) } 1 as additional input to predict a label ˆ y for a molecule m 
Revised: , ( x N , y N ) } 1 as additional input to predict a label ˆ y for a molecule m [Equation] A

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original:  where a ( . ) is a function that adapts w of g based on Z for example via gradient-descent.
Revised: 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 set molecules { x 1 , . . .
Revised: Embedding-based methods use a different approach and learn representations of the support set molecules { x 1 , . . .

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

Original: We describe the embedding-based methods Similarity Search in Section A.1.2, Neural Similarity Search in 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.
Revised: We describe the embedding-based methods Similarity Search in Section A.1.2, Neural Similarity Search in 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 described in Section A.1.1.
Revised: The "frequent hitters" baseline is described in Section A.1.1.

Original: A.1. F REQUENT HITTERS : DETAILS AND HYPERPARAMETERS
Revised: A.1. F REQUENT HITTERS : DETAILS AND HYPERPARAMETERS

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

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

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

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: Thus, during training for the same molecule m
Revised: Thus, during training for the same molecule m

Original: [Table]
Revised: [Table]

Original: the model might have to predict both y = 1 and y = −
Revised: the model might have to predict both y = 1 and y = −

Original: [Table]
Revised: [Table]

Original: since the molecule can be active in one task and inactive in another task.
Revised: since the molecule can be active in one task and inactive in another task.

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

Original: We chose an additive combination of the Morgan ﬁngerprints, RDKit ﬁngerprints
Revised: We chose an additive combination of the Morgan ﬁngerprints, RDKit ﬁngerprints

Original: [Table]
Revised: [Table]

Original: and MACCS keys for the input representation to the MLP.
Revised: and MACCS keys for the input representation to the MLP.

Original: Hyperparameter search.
Revised: Hyperparameter search.

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

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

Original: A.1. C LASSIC SIMILARITY SEARCH : DETAILS AND HYPERPARAMETERS
Revised: A.1. C LASSIC SIMILARITY SEARCH : DETAILS AND HYPERPARAMETERS

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

Original: In the simplest case, molecules that are similar to a 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).
Revised: In the simplest case, molecules that are similar to a 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:
Revised: Thus, the Similarity Search as used in chemoinformatics can be formally written as:

Original: [Equation]
Revised: [Equation]

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

Original: Then, an association function, consisting of a) the similarity measure k ( ., . ) and b) mean pooling across molecules weighted by their similarity and activity, is used to compute the predictions.
Revised: Then, an association function, consisting of a) the similarity measure k ( ., . ) and b) mean pooling across molecules weighted by their similarity and activity, is used to compute the predictions.

Original: Notably, there are many variants of Similarity Search (Cereto-Massagué et al., 2015; Wang et al., 2010; Eckert and Bajorath, 2007;
Revised: Notably, there are many variants of Similarity Search (Cereto-Massagué et al., 2015; Wang et al., 2010; Eckert and Bajorath, 2007;

Original: Geppert et al., 2008; Willett, 2014; Sheridan and Kearsley, 2002;
Revised: 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.
Revised: 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.
Revised: 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).
Revised: 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 of active or inactive molecules, which corresponds to different pooling strategies of the support set.
Revised: Riniker and Landrum (2013) are aware of different fusion strategies for sets of active or inactive molecules, which corresponds to different pooling strategies of the support set.

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

Original: Hyperparameter search.
Revised: 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 like shown in Section 3.4.
Revised: 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 like shown in Section 3.4.

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

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

Original: A.1. N EURAL S IMILARITY S EARCH OR S IAMESE NETWORKS : DETAILS AND HYPERPARAMETERS
Revised: A.1. N EURAL S IMILARITY S EARCH OR S IAMESE NETWORKS : DETAILS AND HYPERPARAMETERS

Original: If the ﬁxed encoder f desc of the Classic Similarity Search is replaced by learned encoders f ME w ,
Revised: If the ﬁxed encoder f desc of the Classic Similarity Search is replaced by learned encoders f ME w ,

Original: Neural Similarity Search variants naturally arise.
Revised: Neural Similarity Search variants naturally arise.

Original: A lot of related work already was done (Koch et al., 2015; Hertz et al., 2006; Ye and Guo, 2018;
Revised: 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).
Revised: Torres et al., 2020).

Original: We adapted these ideas, such that a fully-connected deep neural network followed by a Layer Normalization (Ba et al., 2016) operation, f ME w with adaptive parameters w , is used in a
Revised: We adapted these ideas, such that a fully-connected deep neural network followed by 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.
Revised: Siamese fashion to compute the embeddings for the input molecule and the support set molecules.

Original: Within an association function block, pairwise similarity values for the query molecule and each support set molecule are computed, associating both embeddings via the dot product.
Revised: Within an association function block, pairwise similarity values for the query molecule and each support set molecule are computed, associating both embeddings via the dot product.

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

Original: [Equation]
Revised: [Equation]

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

Original: Note that this method uses a balancing strategy for the labels y (cid:48) n =
Revised: Note that this method uses a balancing strategy for the labels y (cid:48) n =

Original: [Equation]
Revised: [Equation]

Original: else , where N A is the number of actives and N
Revised: else , where N A is the number of actives and N

Original: I is the number of inactives of the support set.
Revised: I is the number of inactives of the support set.

Original: Figure A1 provides an schematic overview of the Neural Similarity Search variant.
Revised: Figure A1 provides an schematic overview of the Neural Similarity Search variant.

Original: We trained the model using the Adam optimizer (Kingma and Ba, 2014) to minimize binary crossentropy loss.
Revised: We trained the model using the Adam optimizer (Kingma and Ba, 2014) to minimize binary crossentropy loss.

Original: Hyperparameter search.
Revised: Hyperparameter search.

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

Original: We report the explored hyperparameter space (Table A2).
Revised: We report the explored hyperparameter space (Table A2).

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

Original: A.1. P ROTO N ET : DETAILS AND HYPERPARAMETERS
Revised: A.1. P ROTO N ET : DETAILS AND HYPERPARAMETERS

Original: Prototypical Networks (ProtoNet) (Snell et al., 2017) learn a prototype r for each class.
Revised: 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
Revised: Concretely, the support set Z is class-wise separated into Z

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

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

Original: { ( x, y ) ∈
Revised: { ( x, y ) ∈

Original: Z | y = − 1 } .
Revised: Z | y = − 1 } .

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

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: cosine similarity, dot product , MinMax similarity and
Revised: cosine similarity, dot product , MinMax similarity and

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: where d is a distance metric.
Revised: where d is a distance metric.

Original: Hyperparameter search.
Revised: Hyperparameter search.

Original: Hyperparameter search has been done by Stanley et al.
Revised: Hyperparameter search has been done by Stanley et al.

Original: (2021), to which we refer here.
Revised: (2021), to which we refer here.

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

Original: Stanley et al.
Revised: Stanley et al.

Original: (2021) use the Mahalanobis distance to measure the similarity between a query molecule and the prototypical representations in the embedding space.
Revised: (2021) use the Mahalanobis distance to measure the similarity between a query molecule and the prototypical representations in the embedding space.

Original: The learning rate is 0 .
Revised: The learning rate is 0 .

Original: 001 and the batch size is 256.
Revised: 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 important hyperparameters are chosen here
Revised: The implementation can be found here https://github.com/microsoft/FS-Mol/blob/main/fs_mol/ protonet_train.py and important hyperparameters are chosen here

Original: https://github.
Revised: https://github.

Original: com/microsoft/FS-Mol/blob/main/fs_mol/utils/protonet_utils.py .
Revised: com/microsoft/FS-Mol/blob/main/fs_mol/utils/protonet_utils.py .

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

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

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

Original: Note that in our study, another difference is that ProtoNet uses a GNN for the encoder, whereas the encoder of the Siamese Networks is a descriptor-based fully-connected network.
Revised: Note that in our study, another difference is that ProtoNet uses a GNN for the encoder, whereas the encoder of the Siamese Networks is a descriptor-based fully-connected network.

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

Original: A.1. I TER R EF LSTM: DETAILS AND HYPERPARAMETERS
Revised: A.1. I TER R EF LSTM: DETAILS AND HYPERPARAMETERS

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

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

Original: The use of the
Revised: The use of the

Original: IterRefLSTM empowers the architecture to update not only the embeddings for the query molecule but also adjust the representations of the support set molecules.
Revised: IterRefLSTM empowers the architecture to update not only the embeddings for the query molecule but also adjust the representations of the support set molecules.

Original: For the IterRefLSTM model, query molecule embedding m =
Revised: For the IterRefLSTM model, query molecule embedding m =

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

Original: , x N .
Revised: , x N .

Original: The query and support set molecule
Revised: The query and support set molecule

Original: [Table]
Revised: [Table]

Original: • Number of hidden layers
Revised: • Number of hidden layers

Original: [Table]
Revised: [Table]

Original: • Number of units per hidden layer
Revised: • Number of units per hidden layer

Original: [Table]
Revised: [Table]

Original: • Output dimension
Revised: • Output dimension

Original: [Table]
Revised: [Table]

Original: • Activation function
Revised: • Activation function

Original: [Table]
Revised: [Table]

Original: • Input dropout
Revised: • Input dropout

Original: [Table]
Revised: [Table]

Original: • Dropout
Revised: • Dropout

Original: [Table]
Revised: [Table]

Original: IterRef embedding layer •
Revised: IterRef embedding layer •

Original: L
Revised: L

Original: [Table]
Revised: [Table]

Original: Similarity module: • Metric
Revised: Similarity module: • Metric

Original: [Table]
Revised: [Table]

Original: similarity • Similarity space dimension
Revised: similarity • Similarity space dimension

Original: [Table]
Revised: [Table]

Original: • Afﬁne
Revised: • Afﬁne

Original: [Table]
Revised: [Table]

Original: • Learning rate
Revised: • Learning rate

Original: [Table]
Revised: [Table]

Original: • Optimizer
Revised: • Optimizer

Original: [Table]
Revised: [Table]

Original: • Weight decay
Revised: • Weight decay

Original: [Table]
Revised: [Table]

Original: • Batch size
Revised: • Batch size

Original: [Table]
Revised: [Table]

Original: embeddings are then updated by an LSTM-like module – the actual IterRefLSTM:
Revised: embeddings are then updated by an LSTM-like module – the actual IterRefLSTM:

Original: [Equation]
Revised: [Equation]

Original: Here
Revised: Here

Original: [Table]
Revised: [Table]

Original: m (cid:48) and X (cid:48) contain the updated representations for the query molecule and the support set molecules.
Revised: m (cid:48) and X (cid:48) contain the updated representations for the query molecule and the support set molecules.

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

Original: The main 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) ]) =
Revised: The main 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 ])) .
Revised: IterRefLSTM L ( π ([ m , X ])) .

Original: Therefore, the full architecture is invariant to permutations of the support set elements.
Revised: Therefore, the full architecture is invariant to permutations of the support set elements.

Original: For details, we refer to Altae-Tran et al.
Revised: For details, we refer to Altae-Tran et al.

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

Original: The IterRefLSTM also includes a similarity module which computes the predictions based on the updated representations mentioned above
Revised: The IterRefLSTM also includes a similarity module which computes the predictions based on the updated representations mentioned above

Original: [Equation]
Revised: [Equation]

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

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

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

Original: Hyperparameter search.
Revised: Hyperparameter search.

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

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

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

Original: A.1. MHN FS : DETAILS AND HYPERPARAMETERS
Revised: A.1. MHN FS : DETAILS AND HYPERPARAMETERS

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

Original: The molecule encoder is a fully-connected Neural Network, consisting of one layer with 1024 units.
Revised: The molecule encoder is a fully-connected Neural Network, consisting of one layer 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.
Revised: 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 descriptors as the inputs for the MHNfs model.
Revised: We chose a concatenation of ECFPs and RDKit-based descriptors as the inputs for the MHNfs model.

Original: Notably, the RDKit-based descriptors were pre-processed in a way that instead of raw values quantils, which were computed by comparing a raw value with the
Revised: Notably, the RDKit-based descriptors were pre-processed in a way that instead of raw values quantils, which were computed by comparing a raw value with the

Original: [Table]
Revised: [Table]

Original: • Number of hidden layers
Revised: • Number of hidden layers

Original: [Table]
Revised: [Table]

Original: • Number of units per hidden layer
Revised: • Number of units per hidden layer

Original: [Table]
Revised: [Table]

Original: • Output dimension
Revised: • Output dimension

Original: [Table]
Revised: [Table]

Original: • Activation function
Revised: • Activation function

Original: [Table]
Revised: [Table]

Original: • Input dropout
Revised: • Input dropout

Original: [Table]
Revised: [Table]

Original: • Dropout
Revised: • Dropout

Original: [Table]
Revised: [Table]

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

Original: [Table]
Revised: [Table]

Original: • Association space dimension
Revised: • Association space dimension

Original: [Table]
Revised: [Table]

Original: • Dropout
Revised: • Dropout

Original: [Table]
Revised: [Table]

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

Original: [Table]
Revised: [Table]

Original: • Number units in the hidden feedforward layer
Revised: • Number units in the hidden feedforward layer

Original: [Table]
Revised: [Table]

Original: • Association space dimension
Revised: • Association space dimension

Original: [Table]
Revised: [Table]

Original: • Dropout
Revised: • Dropout

Original: [Table]
Revised: [Table]

Original: • Number of layers:
Revised: • Number of layers:

Original: [Table]
Revised: [Table]

Original: Similarity module: • Metric
Revised: Similarity module: • Metric

Original: [Table]
Revised: [Table]

Original: • Similarity space dimension
Revised: • Similarity space dimension

Original: [Table]
Revised: [Table]

Original: • τ
Revised: • τ

Original: [Table]
Revised: [Table]

Original: Layer-normalization
Revised: Layer-normalization

Original: [Table]
Revised: [Table]

Original: • Afﬁne
Revised: • Afﬁne

Original: [Table]
Revised: [Table]

Original: Training • Learning rate
Revised: • Learning rate

Original: [Table]
Revised: [Table]

Original: • Optimizer
Revised: • Optimizer

Original: [Table]
Revised: [Table]

Original: • Weight decay
Revised: • Weight decay

Original: [Table]
Revised: [Table]

Original: • Batch size
Revised: • Batch size

Original: [Table]
Revised: [Table]

Original: • Warm-up phase (epochs)
Revised: • Warm-up phase (epochs)

Original: [Table]
Revised: [Table]

Original: • Constant learning rate phase (epochs)
Revised: • Constant learning rate phase (epochs)

Original: [Table]
Revised: [Table]

Original: • Decay rate 0.
Revised: • Decay rate 0.

Original: • Max.
Revised: • Max.

Original: number of epochs
Revised: number of epochs

Original: [Table]
Revised: [Table]

Original: distributation of all FS-Mol training molecules, were used.
Revised: distributation of all FS-Mol training molecules, were used.

Original: All descriptors were normalized based on the FS-Mol training data.
Revised: All descriptors were normalized based on the FS-Mol training data.

Original: Hyperparameter search.
Revised: Hyperparameter search.

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

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

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

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

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

Original: Model training.
Revised: Model training.

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

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

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

Original: Performance improvements in comparison to a naive baseline.
Revised: Performance improvements in comparison to a naive baseline.

Original: Figure A3 shows a task-wise performance comparison between MHNfs and the Frequent Hitter model.
Revised: Figure A3 shows a task-wise performance comparison between MHNfs and the Frequent Hitter model.

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

Original: In 132 cases the MHNfs outperforms the frequent hitter model.
Revised: In 132 cases the MHNfs outperforms the frequent hitter model.

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

Original: [Figure]
Revised: [Figure]

Original: n t r
Revised: n t r

Original: [Figure]
Revised: [Figure]

Original: i
Revised: i

Original: [Figure]
Revised: [Figure]

Original: i n g
Revised: i n g

Original: [Figure]
Revised: [Figure]

Original: L o
Revised: L o

Original: [Figure]
Revised: [Figure]

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

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: • Inner learning

Original: [Table] [Table]
Revised: [Table] [Table]

Original: • Update steptest
Revised: • Update step

Original: [Table]
Revised: [Table]

Original: • Weight decay
Revised: • Weight decay

Original: [Table]
Revised: [Table]

Original: • Epochs
Revised: • Epochs

Original: [Table]
Revised: [Table]

Original: • Eval.
Revised: • Eval.

Original: steps
Revised: steps

Original: [Table]
Revised: [Table]

Original: • Map dimension
Revised: • Map dimension

Original: [Table]
Revised: [Table]

Original: • Map layer
Revised: • Map layer

Original: [Table]
Revised: [Table]

Original: • Pre fc layer
Revised: • Pre fc layer

Original: [Table]
Revised: [Table]

Original: • Map dropout
Revised: • Map dropout

Original: [Table]
Revised: [Table]

Original: • Context layer
Revised: • Context layer

Original: [Table]
Revised: [Table]

Original: • Hidden dimension
Revised: • Hidden dimension

Original: [Table]
Revised: [Table]

Original: • Number of layers
Revised: • Number of layers

Original: [Table]
Revised: [Table]

Original: • Number of layers for relation edge update
Revised: • Number of layers for relation edge update

Original: [Table]
Revised: [Table]

Original: • Batch norm
Revised: • Batch norm

Original: [Table]
Revised: [Table]

Original: • Relation dropout
Revised: • Relation dropout

Original: [Table]
Revised: [Table]

Original: • Relation dropout
Revised: • Relation dropout

Original: [Table]
Revised: [Table]

Original: A.1. PAR: DETAILS AND HYPERPARAMETERS
Revised: A.1. PAR: DETAILS AND HYPERPARAMETERS

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

Original: These embeddings are fed into an attention mechanism module which also uses activity information of the support set molecules to create enriched representations.
Revised: These embeddings are fed into an attention mechanism module which also uses activity information of the support set molecules to create enriched representations.

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

Original: Hyperparameter search.
Revised: Hyperparameter search.

Original: For details we refer to Wang et al.
Revised: For details we refer to Wang et al.

Original: (2021) and https://github.
Revised: (2021) and https://github.

Original: com/tata1661/PAR-NeurIPS21/blob/main/parser.py .
Revised: com/tata1661/PAR-NeurIPS21/blob/main/parser.py .

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

Original: The hyperparameter choice for Tox21 (Wang et al., 2021) was used as a starting point.
Revised: 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.
Revised: We report the explored hyperparameter space in Table A5.

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

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

Original: We used a training script provided by (Wang et al., 2021), which can be found here https://github.com/tata1661/PAR-NeurIPS21 .
Revised: We used a training script provided by (Wang et al., 2021), which can be found here https://github.com/tata1661/PAR-NeurIPS21 .

Original: A. D ETAILS ON THE ABLATION STUDY
Revised: A. D ETAILS ON THE FS-M OL BENCHMARKING EXPERIMENT

Revised: This section provides additional information for the FS-Mol benchmarking experiment (see Section 5).

Revised: Memory-based baselines.

Revised: The Classic Similarity Search can be considered as a method with associative memory, where the label is retrieved from the memory.

Revised: Notably, for this method, the associative memory is very limited since it is the support set.

Revised: Siamese Networks, analogously to the

Revised: Classic Similarity Search, retrieve the label from a memory, whereby the similarities are determined in a learned space.

Revised: Also, the IterRefLSTM-based method can be seen as having a memory, whereby the LSTM controls storing and removing information from the training data by the input and the forget gate.

Revised: In NLP, kNN-type memories are currently used.

Revised: Conceptually, they are very similar to the Modern Hopﬁeld Networks, setting the number of heads to one and choosing a suitable value for β .

Revised: [Table]

Revised: [Figure]

Revised: Results.

Revised: The reported performance metrics comprise three different sources of variation, namely variation across different tasks, variation across different support sets during test time, and variation across different training reruns.

Revised: While error bars in Table 1 report variation across tasks, error bars in Table A6 report variation across training reruns.

Revised: For ADKF-IFT, the authors provided error bars for every single test task.

Revised: Based on these error bars we sampled performance values to be able to compare ADKF-IFT with the MHNfs training reruns.

Revised: Figure A4 shows a task-wise model comparison between a) MHNfs and the IterRefLSTM-based method and b) MHNfs and ADKF-IFT.

Revised: For a) MHNfs performs better on 106 of 157 tasks and therefore signiﬁcantly outperforms the

Revised: IterRefLSTM-based method (binomial test p -value 6 . 8 · 10 − 6 ).

Revised: For b) MHNfs performs better ontasks and therefore signiﬁcantly outperforms ADKF-IFT (binomial test p -value 0 .

Revised: 001 ), too.

Revised: Notably,

Revised: ADKF-IFT performs better on non kinase-targets which can be seen in Table 1.

Revised: A. D ETAILS ON THE ABLATION STUDY

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

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

Original: A.2. A BLATION STUDY A: COMPARISON AGAINST I TER R EF LSTM
Revised: A.3. A BLATION STUDY A: COMPARISON AGAINST I TER R EF LSTM

Original: For a fair comparison between the cross-attention module and the IterRefLSTM we used a pruned
Revised: 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 IterRefLSTM model.
Revised: MHN version ("MHNfs -CM") which has no context module and compared it with the IterRefLSTM model.

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

Revised: [Table]

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

Original: [Table]
Revised: [Table]

Revised: can be seen in Table A7.

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

Original: Both p -values indicate high signiﬁcance.
Revised: Both p -values indicate high signiﬁcance.

Original: A.2. A BLATION STUDY B: ALL DESIGN ELEMENTS
Revised: A.3. A BLATION 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.
Revised: 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-PR performance values.
Revised: For this analysis, we start with the complete MHNfs which includes all modules and report AUC and ∆ AUC-PR performance values.

Original: Then, we iteratively omit the individual modules, measuring whether there is a signiﬁcant performance difference with and without the module.
Revised: Then, we iteratively omit the individual modules, measuring whether there is a signiﬁcant performance difference with and without the module.

Original: Table A6 shows the results [Table] where performance 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") is included.
Revised: Table A7 shows the results, where performance 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") is included.

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

Original: We evaluted the impact of the learned molecule encoder by replacing it with a ﬁxed encoder, which maps a molecule to its descriptors.
Revised: We evaluted the impact of the learned molecule encoder by replacing it with a ﬁxed encoder, which maps a molecule to its descriptors.

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

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

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

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

Original: We performed paired Wilcoxon rank sum tests on both metrics, comparing two methods in consecutive rows in the table.
Revised: We performed paired Wilcoxon rank sum tests on both metrics, comparing two methods in consecutive rows in the table.

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

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

Original: A.2. A BLATION STUDY C: U NDER DOMAIN SHIFT ON T OX
Revised: A.3. A BLATION STUDY C: U NDER DOMAIN SHIFT ON T OX

Original: Referring to Section A.2.2, the context module and the cross-attention module showed their importance for the global architecture.
Revised: 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 shift experiment on Tox21 as one can see in Table A7.
Revised: This importance gets even more pronounced for the domain shift experiment on Tox21 as one can see in Table A8.

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

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

Original: A. G ENERALIZATION TO DIFFERENT SUPPORT SET SIZES
Revised: A. D ETAILS ON THE DOMAIN SHIFT EXPERIMENTS

Revised: This section provides additional information for the Domain shift experminet on Tox21.

Revised: Results.

Revised: The reported performance metrics comprise three different sources of variation, namely variation across different tasks, variation across different support sets during test time, and variation across different training reruns.

Revised: While error bars in Table 2 report variation across both, drawn support sets and training reruns, error bars in Table A9 just report variation across training reruns.

Revised: Notably, for the Similarity Search, the performance values do not vary since the model does not include any trainable parameters.

Revised: [Table]

Revised: [Table]

Revised: The Similarity Search does not include any learned parameters.

Revised: Therefore, there is no variability across training

Revised: [Table]

Revised: A. G ENERALIZATION TO DIFFERENT SUPPORT SET SIZES

Original: In the following section, we test the ability of MHNfs to generalize to different support set sizes.
Revised: In the following section [Table] we test the ability of MHNfs to generalize to different support set sizes.

Original: During training in the FS-Mol benchmarking setting, the MHNfs model has access to support sets of size 16.
Revised: During training in the FS-Mol benchmarking setting, the MHNfs model has access to support sets of size 16.

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

Original: Figure A4 provides performance estimates of the support-set-size-16 MHNfs models on other support set sizes.
Revised: Figure A5 provides performance estimates of the support-set-size-16 MHNfs models on other support set sizes.

Original: Note that the estimates could be seen as approximate lower bounds of the predictive performance on settings with different support set sizes (y-axis labels).
Revised: Note that the estimates could be seen as approximate lower bounds of the predictive performance on settings with different support set sizes (y-axis labels).

Original: For a model used in production or in a real-world drug discovery setting, MHNfs should be trained with varying support set sizes that resemble the distribution of real drug discovery projects.[Table] [Figure] [Table]
Revised: For a model used in production or in a real-world drug discovery setting, MHNfs should be trained with varying support set sizes that resemble the distribution of real drug discovery projects.

Original: Triantaﬁllou et al.
Revised: Triantaﬁllou et al.

Original: (2019) analysed the performance of different few-shot models across different support set sizes.
Revised: (2019) analysed the performance of different few-shot models across different support set sizes.

Original: Their analysis showed that in very-low-data settings embedding-based methods, namely Prototypical Networks and fo-Proto-MAML, performed best.
Revised: Their analysis showed that in very-low-data settings embedding-based methods, namely Prototypical Networks and fo-Proto-MAML, performed best.

Original: In contrast to that, ﬁnetuning based method signiﬁcantly proﬁt from larger support set sizes (Triantaﬁllou et al., 2019).
Revised: In contrast to that, ﬁnetuningbased method signiﬁcantly proﬁt from larger support set sizes (Triantaﬁllou et al., 2019).

Original: MHNfs is an embedding-based method and performs – in accordance with the ﬁndings mentioned above (Triantaﬁllou et al., 2019) – well for small support set sizes (see Table 1).
Revised: MHNfs is an embedding-based method and performs – in accordance with the ﬁndings mentioned above (Triantaﬁllou et al., 2019) – well for small support set sizes (see Table 1).

Original: Following Triantaﬁllou et al.
Revised: Following Triantaﬁllou et al.

Original: (2019), it is exactly the settings related to these smaller support set sizes, e.g. a support set size of 16, which are suitable for MHNfs.
Revised: (2019), it is exactly the settings related to these smaller support set sizes, e.g. a support set size of 16, which are suitable for MHNfs.

Original: For large support set sizes, e.g. 64 or 128, we point to the work done by (Chen et al., 2022), in which the ﬁne-tuning method ADKF-IFT achives an ∆ AUC-PR-score > 0 .
Revised: For large support set sizes, e.g. 64 or 128, we point to the work done by Chen et al. (2022), in which the ﬁne-tuning method ADKF-IFT achives an ∆ AUC-PR-score > 0 .

Original: 3 for large support set sizes.
Revised: 3 for large support set sizes.

Original: A. G ENERALIZATION TO DIFFERENT CONTEXT SETS
Revised: A. G ENERALIZATION TO DIFFERENT CONTEXT SETS

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

Original: While the FS-Mol training split is used as a context during training, we assessed whether our model is robust to different context sets for inference.
Revised: While the FS-Mol training split is used as a context during training, we assessed whether our model is robust to different context 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.
Revised: 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.
Revised: From this set, we sample 10,000 molecules as context set for MHNfs.

Original: Because GEOM contains drug-like molecules, similar to FS-Mol the predictive performance remains stable (see Table A8).
Revised: Because GEOM contains drug-like molecules, similar to FS-Mol the predictive performance remains stable (see Table A10). [Figure] [Figure] [Figure]

Revised: [Figure]

Revised: [Figure]

Original: A. D ETAILS AND INSIGHTS ON THE CONTEXT MODULE
Revised: A U C - P R ( t e s t s e t ) [Table] Axelrod [Table]

Original:  The context module replaces the initial representations of query and support set molecules by a retrieval from the context set.
Revised: Bombarelli, [Table] A. D ETAILS AND INSIGHTS ON THE CONTEXT MODULE The context module replaces the initial representations of query and support set molecules by a retrieval from the context set.

Original: The context set is a large set of molecules and covers a large chemical space.
Revised: The context set is a large set of molecules and covers a large chemical space.

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

Original: Figure A5 shows the effect of the context module for the MHNfs model.
Revised: Figure A6 shows the effect of the context module for the MHNfs model.

Original: Extreme initial embeddings, such as the purple embedding on the right, are pulled more into the known chemical space, represented by the context molecules.
Revised: Extreme initial embeddings, such as the purple embedding on 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 embeddings contribute to the context-enriched representations due to skip-connections.
Revised: Notably, the replacement described above is a soft replacement, because also the initial embeddings contribute to the context-enriched representations due to skip-connections.

Original: A. R EINFORCING THE COVARIANCE STRUCTURE IN THE DATA USING MODERN H OPFIELD NETWORKS
Revised: A. R EINFORCING THE COVARIANCE STRUCTURE IN THE DATA USING MODERN H OPFIELD NETWORKS

Original: We follow the argumentation of (Fürst et al., 2022, Theorem A3) that retrieval from an associative memory of a MHN reinforces the covariance structure.
Revised: We follow the argumentation of (Fürst et al., 2022, Theorem A3) that retrieval from an associative memory 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 molecule embedding from the support set x ∈ R d and both have been enriched with the context module with memory C ∈ R d × M (ignoring linear mappings):
Revised: Let us assume that we have one molecule embedding from the query set m ∈ R d and one molecule embedding from the support set x ∈ R d and both have been enriched with the context module with memory C ∈ R d × M (ignoring linear mappings):

Original: [Equation]
Revised: [Equation]

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

Original: [Figure]
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: where c is the row mean of C and following the weighted covariances are used:
Revised: where c is the row mean of C and following the weighted covariances are used:

Original: [Equation]
Revised: [Equation]

Original: J
Revised: J

Original: [Equation]
Revised: [Equation]

Original: R
Revised: R

Original: [Equation]
Revised: [Equation]

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

Original: The Jacobian J of p = softmax( β a ) is J( β a )
Revised: The Jacobian J of p = softmax( β a ) is J( β a )

Original: = β (cid:0) diag( p ) − pp T (cid:1) .
Revised: = β (cid:0) diag( p ) − pp T (cid:1) .

Original: [Equation]
Revised: [Equation]

Original: this is the second moment minus the mean squared, which is the variance.
Revised: this is the second moment minus the mean squared, which is the variance.

Original: Therefore, b T J( β a ) b is β times the covariance of b if component i is drawn with probability p
Revised: Therefore, b T J( β a ) b is β times the covariance of b if component i is drawn with probability p

Original: i of the multinomial distribution p .
Revised: i of the multinomial distribution p .

Original: In our case the component i is context sample c
Revised: In our case the component i is context sample c

Original: i .
Revised: i .

Original: J m is the average of J( λ a ) over λ = 0 to λ = β .
Revised: J m is the average of J( λ a ) over λ = 0 to λ = β .

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

Original: [Equation]
Revised: [Equation]

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

Original: A. D ISCUSSION , LIMITATIONS AND BROADER INPACT
Revised: A. D ISCUSSION , LIMITATIONS AND BROADER INPACT

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

Original: In a domain shift study, the robustness and transferability of the learned models has been assessed and again MHNfs exhibited the best performance.
Revised: In a domain shift study, the robustness and transferability of the learned models has been assessed and again MHNfs exhibited the best performance.

Original: The resulting predictive models often reach an AUC larger than .
Revised: The resulting predictive models often reach an AUC larger than .

Original: 70 , which means that enrichment of active molecules is expected (Simm et al., 2018) when the models are used for virtual screening.
Revised: 70 , which means that enrichment of active molecules is expected (Simm et al., 2018) when the models are used for virtual screening.

Original: It has not escaped our notice that the speciﬁc context module we have proposed could immediately be used for few-shot learning tasks in computer vision, but might be hampered by computational constraints.
Revised: It has not escaped our notice that the speciﬁc context module we have proposed could immediately be used for few-shot learning tasks in computer vision, but might be hampered by computational constraints.

Original: Limitations.
Revised: Limitations.

Revised: In the FS-Mol benchmark experiment, the runner-up method ADKF-IFT (Chen et al., 2022) performed better on non kinase-tasks.

Revised: We hypothesize that we could improve the MHNfs performance for non kinase tasks by upsampling the other task sub-groups.

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

Original: The output domain of our method comprises biological effects, such that the prediction must be understood in that domain.
Revised: The output domain of our method comprises biological effects, such that the prediction must be understood in that domain.

Original: Our method demands higher computational costs and memory footprint as other embedding-based methods because of the calculations necessary for the context module.
Revised: Our method demands higher computational costs and memory footprint as other embedding-based methods because of the calculations necessary for the context module.

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

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

Original: Deep learning methods usually have a large number of hyperparameters, such as hidden dimensions, number of layers, learning rates, of which we were only able to explore the most important ones.
Revised: Deep learning methods usually have a large number of hyperparameters, such as hidden dimensions, number of layers, learning rates, of which we were only able to explore the most important ones.

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

Original: Broader impact.
Revised: Broader impact.

Original: Impact on machine learning and related 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 and material science datasets, (b) further improved biotechnologies, and (c) accounting for characteristics of individuals, the drug and materials discovery process will be made more efﬁcient.
Revised: We envision that with (a) the increasing availability of drug discovery and material science datasets, (b) further improved biotechnologies, and (c) accounting for characteristics of individuals, the drug and materials discovery process will be made more efﬁcient.

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

Original: Our approach also shows that such a system is more robust against domain shifts, which could be a step towards Broad AI (Chollet, 2019; Hochreiter, 2022).
Revised: Our approach also shows that such a system is more 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 costefﬁcient drug discovery process.
Revised: If the approach proves useful, it could lead to a faster and more costefﬁcient drug discovery process.

Original: Especially the COVID-19 pandemic has shown that it is crucial for humanity to speed up the drug discovery process to few years or even months.
Revised: Especially the COVID-19 pandemic has shown that it is crucial for 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 drugs developed faster.
Revised: We hope that this work contributes to this effort and eventually leads to safer drugs developed faster.

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

Original: As common with methods in machine learning, potential danger lies in the possibility that users rely too much on our new approach and use it without reﬂecting on the outcomes.
Revised: As common with methods in machine learning, potential danger lies in the possibility 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 unsuccessful wet lab validation and negative wet lab tests.
Revised: Failures of the proposed method would lead to unsuccessful wet lab validation and negative wet lab tests.

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

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

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

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

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