C ONTRASTIVE LEARNING OF IMAGE - AND STRUCTUREBASED REPRESENTATIONS IN DRUG DISCOVERY
A BSTRACT
Contrastive learning for self-supervised representation learning has brought a strong improvement to many application areas, such as computer vision and natural language processing.
With the availability of large collections of unlabeled data in vision and language, contrastive learning of language and image representations has brought impressive results.
The contrastive learning methods CLIP and CLOOB have demonstrated that the learned representations are highly transferable to a large set of diverse tasks when trained on multi-modal data from two different domains.
In drug discovery, similar large, multi-modal datasets comprising both cell-based microscopy images and chemical structures of molecules are available.
However, contrastive learning has not been used for this type of multi-modal data in drug discovery, although transferable representations could be a remedy for the time-consuming and cost-expensive label acquisition in this domain.
In this work, we present a contrastive learning method for image-based and structure-based representations of small molecules for drug discovery.
Our method, Contrastive
Leave-One-Out boost for Molecule Encoders (CLOOME), is based on CLOOB and comprises an encoder for microscopy data, an encoder for chemical structures and a contrastive learning objective.
On the benchmark dataset ºCell Paintingº, we demonstrate the ability of our method to learn transferable representations by performing linear probing for activity prediction tasks.
Additionally, we show that the representations could also be useful for bioisosteric replacement tasks.
I NTRODUCTION
Contrastive learning has had a strong impact on computer vision and natural language processing.
Over the last decade, supervised deep learning methods have achieved outstanding results in the field of computer vision (Krizhevsky et al., 2012;
He et al., 2016).
These supervised methods require large amounts of labeled data, which may be very costly or unfeasible to obtain, and they have limited generalization abilities (Sun et al., 2017; Marcus, 2018).
This has led to the exploration of new methods that are able to learn robust representations of the data which can be transferred to different downstream tasks (Luo et al., 2017; Chen et al., 2020).
With contrastive learning methods (Gutmann and HyvÈarinen, 2010) and self-supervision these meaningful representations can be obtained without the need for large amounts of expensive manually-provided labels (He et al., 2020; Chen et al., 2020; Caron et al., 2020;
Grill et al., 2020).
While uni-modal methods typically use pre-text tasks (Chen et al., 2020), for multi-modal methods the self-supervision arises from the availability of two modalities of an instance, such as image and text (Radford et al., 2021; Devillers et al., 2021).
Both uni-modal and multi-modal contrastive learning methods have recently had a substantial impact in computer vision and natural language processing (Jaiswal et al., 2020).
CLIP for multi-modal data yields spectacular performance at zero-shot transfer learning and has recently been improved by CLOOB.
An outstanding multi-modal approach is Contrastive
Language-Image Pre-training (CLIP) (Radford et al., 2021), which learns both image- and textrepresentations simultaneously.
CLIP shows comparable performance to methods that are solely image-based and yields highly transferable representations, which is shown by its high performance at zero-shot transfer learning.
However, CLIP has recently been shown to suffer from the ºexplaining awayº effect (FÈurst et al., 2021; Pearl, 1988; Wellman and Henrion, 1993) (details in Section 2).
Considering these caveat, the ªContrastive Leave One Out Boostº (CLOOB) method has been proposed (FÈurst et al., 2021).
CLOOB uses a different objective, the ªInfoLOOBº (LOOB for ªLeave
One Out Boundº) objective (Poole et al., 2019), which does not include the positive pair in the denominator to avoid saturation effects (FÈurst et al., 2021).
Moreover, continuous modern Hopfield networks (Ramsauer et al., 2021) are used to reinforce the covariance structure of the data.
As a result, CLOOB has further improved zero-shot transfer learning.
The ability to learn transferable representation from multi-modal data makes CLOOB the prime candidate for learning representations of molecules in drug discovery.
Contrastive learning for molecule representations in drug discovery.
In drug discovery, the effect of the limited availability of data on molecules is even more severe, since the acquisition of a single bioactivity data point can cost several thousand dollars and take several weeks or months
(MacArron et al., 2011; Knight et al., 2006).
Therefore, methods that can learn transferable representations from unlabelled data are highly demanded.
Thus, a handful of contrastive learning approaches have been recently developed for different tasks in drug discovery.
MolCLR (Wang et al., 2022) uses contrastive molecule-to-molecule training by augmenting molecular graphs.
StÈark et al.
(2021) contrastively learn 3D and 2D molecule representations to inform the learned molecule encoder with 3D information.
Lee et al.
(2021) and Seidl et al.
(2022) use contrastive learning for molecules and chemical reactions, and Vall et al.
(2021) utilizes text representations of wet-lab procedures to enable zero-shot predictions.
However, none of these methods have exploited the wealth of information contained in microscopy images of molecule-perturbed cells (Bray et al., 2016) and demonstrated strong transferability of the learned molecule encoders.
Image-based profiling of small molecules has strongly improved the drug discovery process.
Characterizing a small molecule by the morphological changes it induces to a cell, is considered promising for accelerating drug discovery (Bray et al., 2016; Caicedo et al., 2017; Chandrasekaran et al., 2021; Simm et al., 2018).
The advantages of this biotechnology are that it is time- and cost- effective as compared to standard activity measurements.
Measuring the effects of a molecule on a biological system early in the drug discovery process might be useful to improve clinical success rates (Bender and CortÂes-Ciriano, 2021).
Particularly, microscopy image-based profiles of small molecules have been suggested to be effective together with deep learning methods (Chandrasekaran et al., 2021).
However, the current efforts are still in standard supervised learning settings based on extracted features (Simm et al., 2018) or deep architectures (Hofmarcher et al., 2019).
The amount of labeled images is in the range of few tens of thousands, although international efforts are currently building datasets which are magnitudes larger (JUMP-CP Consortium, 2022).
Instead of the currently used activity measurements as labels (Simm et al., 2018; Hofmarcher et al., 2019), we propose a self-supervised contrastive learning strategy of image- and structure-based molecule encoders:
Contrastive Leave-One-Out boost for Molecule Encoders (CLOOME).
CLOOME extends recent successful contrastive learning methods to the fields of biological imaging and drug discovery.
Our approach intends to overcome the limited transferability of current molecule encoders (Cai et al., 2020; Stanley et al., 2021).
Contributions
• We introduce a new contrastive learning approach for image- and structure-based representations of molecules, which is based on the contrastive learning method CLOOB.
• We show that the learned representations are highly transferable to relevant downstream tasks in drug discovery by linear probing on activity prediction tasks.
• We demonstrate that our approach learns meaningful representations of molecules which allow to retrieve potential bioisosteres.
CLOOME: C ONTRASTIVE L EAVE O NE O UT B OOST FOR M OLECULE E NCODERS
We propose contrastive learning of representations of molecules from pairs of microscopy images and chemical structures to obtain highly transferable molecule encoders (see Figure 1).
In contrast to previous approaches, in which molecule encoders learned representations using activity data (Hofmarcher et al., 2019; Simm et al., 2018) or used hand-crafted representations (Carpenter et al., 2006; Bray et al., 2017), CLOOME optimizes representations without activity data or human expertise.
The training dataset consists of N pairs of microscopy images of molecule-perturbed cells and chemical structures of molecules { ( x 1 , z 1 ) , . . .
, ( x N , z N ) } .
We assume that an adaptive imageencoder h x ( . ) and an adaptive structure-encoder h z ( . ) are available that map the images and chemical structures to their embeddings x n =
h x ( x n ) and z n = h z ( z n ) , respectively.
Note that the original image is denoted as x n , which is mapped to an image embedding x n by a neural network
h x ( . ) , e.g. a ResNet.
The stacked microscopy image embeddings are denoted as X =
( x 1 , . .
, x N ) and the stacked structure embeddings as Z = ( z 1 , . . .
, z N ) .
The embeddings are normalized to such that ∥ x n ∥ =
∥
z n ∥ = 1 ∀ n .
For notation, see also Table A1.
In a contrastive learning setting, methods aim to increase the similarity of matched pairs and decrease the similarity of unmatched pairs.
This task has often been approached by maximizing the mutual information of the embeddings using the InfoNCE loss (Radford et al., 2021; van den Oord et al., 2018;
Chen et al., 2020), which is also used in the CLIP approach (Radford et al., 2021).
The InfoNCE objective function has the following form:
[Equation]
where τ − 1 is the inverse temperature parameter, which is a hyperparameter of the method.
CLIP has the problem of ºexplaining awayº (Pearl, 1988; Wellman and Henrion, 1993; FÈurst et al., 2021).
Explaining away describes the effect in which few features are over-represented while others are neglected.
This effect can be present a) when learning focuses only on few features and/or b) when the covariance structure in the data is insufficiently extracted.
Explaining away can be caused by saturation of the InfoNCE objective (Yeh et al., 2021; Zhang et al., 2022; FÈurst et al., 2021).
To ameliorate these drawbacks, CLOOB (FÈurst et al., 2021) has introduced the InfoLOOB objective together with Hopfield networks as a promising method for contrastive learning.
Therefore, we base our method on CLOOB.
For our approach, first image- and structure-embeddings are retrieved from stored image embeddings U and structure embeddings V .
U x
i denotes an image-retrieved image embedding, U z
i a structureretrieved image embedding, V x i an image-retrieved structure embedding and V z i a structure-retrieved structure embedding.
In analogy to CLOOB, these retrievals from continuous modern Hopfield networks are computed as follows:
[Equation]
where β is a scaling parameter of the Hopfield network which is considered a hyperparameter.
These retrieved embeddings U x
i , U z
i , V x i ,
V z
i are also normalized to unit norm.
By default, we store the current minibatch in the modern Hopfield networks, that is, U = X and V
= Z .
Note that X contains the image embeddings ( Z the structure embeddings) and we use N ambiguously both as dataset size, but also as mini-batch size to keep the notation uncluttered.
The choice that U
=
X and V
= Z is mostly taken because of computational constraints, while U and V could hold the whole dataset or, alternatively, exemplars.
For further details on notation, see Table A1.
Then, the InfoLOOB objective (FÈurst et al., 2021; Poole et al., 2019) for the retrieved embeddings is used as objective function:
[Equation]
Microscopy image encoder.
Microscopy images differ from natural images in several aspects, for example the variable number of channels that depends on the staining procedure (Pepperkok and
Ellenberg, 2006; Hofmarcher et al., 2019).
Although standard image encoders, such as Residual
Networks (He et al., 2016) could be in principle used with minor adjustments, alternative architectures, such as multiple instance learning approaches, could be required for very high resolution datasets (Ilse et al., 2020).
In all our experiments, we use a ResNet-50 encoder with five input channels and downsized the microscopy images to 320x320.
Molecule structure encoder.
Since the advent of Deep Learning, a large number of architectures to encode molecules have been suggested (Lusci et al., 2013; Dahl et al., 2014; Unterthiner et al., 2014;
Kearnes et al., 2016; Jiang et al., 2021).
In contrast to computer vision and natural language processing, in which only few prominent architectures have emerged, there is yet no standard choice for molecule structure encoders.
Due to their computational efficiency and good predictive performance, CLOOME uses descriptor-based fully-connected networks (Mayr et al., 2016; 2018) with 4 hidden layers of 1024 units with ReLU activations and batch normalization (for further details see Sec. 3 and Sec. A.2).
However, also any graph (Merkwirth and Lengauer, 2005; Scarselli et al., 2008; Kearnes et al., 2016;
Xu et al., 2018), message-passing (Gilmer et al., 2017), or sequence-based (Alperstein et al., 2019) neural network with an appropriate pooling operation can be used as structure encoder.
E XPERIMENTS
Dataset and preprocessing.
CellPainting.
We use pairs of microscopy images and molecules from the CellPainting (Bray et al., 2016; 2017) dataset.
This dataset is a collection of high-throughput fluorescence microscopy images of U2OS cells treated with different small molecules (Bray et al., 2016).
The dataset consists of 919,265 five-channel images corresponding to 30,616 different molecules.
The experiment to obtain the microscopy images was conducted using 406 multi-well plates, and each one of the before mentioned individual images are views from a sample spanning the space in the corresponding well, so that six adjacent views belong to one single sample.
After
[Table]
Pre-processing .
We followed the pre-processing protocol of Hofmarcher et al.
(2019), which consisted of converting the original TIF images from 16-bit to 8-bit, simultaneously removing the 0.0028% of pixels with highest values.
For data augmentation and to allow large batch sizes, the images were cropped and re-scaled from the original 520x696 pixel resolution to 320x320 during training.
Moreover, the images were normalized using the mean and standard deviation calculated for the training split.
Concerning molecules, their corresponding SMILES strings were transformed to 1024-bit Morgan fingerprints with a radius of 3, taking chirality into account (Morgan, 1964; Rogers and Hahn, 2010).
Data splits.
We split our dataset into training, validation, and test set, using the splits of Hofmarcher et al.
Samples which have not been used in the previous study due to missing activity data, are assigned to the training split.
Note that all images belonging to the same molecular structure are moved into the same set.
Finally, training, validation and test set consist of 674,357, 28,632 and 56,793 image and molecule pairs, respectively.
Pre-training, architecture and hyperparameters.
We use the suggested hyperparameters of
OpenCLIP (Wortsman et al., 2021) and CLOOB (FÈurst et al., 2021) wherever applicable, and tuned a few critical hyperparameters, such as learning rate and the β parameter of the Hopfield layer on a validation set.
The architecture of the structure encoder was inspired by previous successful models (Mayr et al., 2016) and was not subject to substantial hyperparameter optimization.
Due to computational constraints, the search was limited to the hyperparameters shown in Table A2.
We used the Adam optimizer (Kingma et al., 2014) with decoupled weight decay regularization (Loshchilov and Hutter, 2019).
The value for weight decay was 0 .
For the learning rate scheduler, we used cosine annealing with a warm-up of 20,000 steps and hard restarts every 7 epochs (Loshchilov and Hutter, 2017).
We set the dimension of the embedding space to d = 512 , which determines the size of the output of both encoders.
We use a batch size of 256 as default due to computational constraints.
For activity prediction as downstream task, the inverse temperature parameter τ − 1 = 30 was used.
For the Hopfield layers, the scaling hyperparameter β = 22 was selected, and the model was trained for 63 epochs based on linear probing results in the corresponding validation set.
For the bioisosteric replacement task, higher validation performance was achieved by a CLIP-like architecture directly using the embeddings returned from the image and structure encoders and the InfoNCE loss.
In this case, the inverse temperature parameter τ − 1 was set to 14 . 3 , and the model was trained for 51 epochs based on the top-1 accuracy in validation.
Hence, different pre-training settings have been found to yield best results for bioactivity prediction and for the bioisosteric replacement task, respectively.
However, the large majority of hyperparameters were shared in both strategies.
Because of the limited exploration of the vast hyperparameter space, we expect potential improvements from further investigations.
For further details on the hyperparameter selection, see Sec.
A.2.
Activity prediction as downstream tasks.
In this experiment, we tested whether the representations learned by CLOOME are transferable by linear probing on 209 downstream activity prediction tasks.
The linear probing test (Alain and Bengio, 2016; Chen et al., 2020) on downstream tasks is often performed for contrastive learning approaches to check the transferability of learned features.
In such experiments, the representations of the pretrained encoders are used and only a single-layer network, such as logistic regression, is fit to the given labels for the supervised task.
If the linear probing test yields good predictive quality, usually below a fully supervised approach (Chen et al., 2020), the representations are considered transferable.
Linear probing evaluation.
The prediction tasks that we employed for linear probing evaluation is the same as used in Hofmarcher et al.
It is a subset of the Cell Painting dataset, consisting of 284,035 images for which the activity labels of the compound treatments were retrieved from
ChEMBL.
The retrieved labels correspond to 10,574 compounds across 209 activity prediction tasks, which are binary classification problems.
However, activity data points are not available for all compounds in all of the tasks, which results in a sparse label matrix.
The data was split into 70% training, 10% validation, and 20% test sets.
This split had been carried out by grouping views from samples treated with the same molecule.
We use image features taken from the penultimate layer of the image encoder, omitting the classification layer.
We train a logistic regression classifier, and report the corresponding metric for each task.
The L2 regularization strength λ was tuned individually for each one of the tasks, considering the values { 10 − 6 , 10 − 5 , . .
In order to evaluate model performance for this downstream task, we use the area under the ROC curve (AUC), which is one of the most prevalent metrics for drug discovery Simm et al.
Hofmarcher et al.
(2019), as it considers the order of the molecules regarding their activity.
We also show the number of tasks for which this metric is higher than the thresholds 0.9, 0.8 and 0.7, respectively.
These thresholds have been used in previous studies (Simm et al., 2018;
Hofmarcher et al., 2019) because models within those categories lead to certain levels of enrichment of hit rates in drug discovery projects.
Baselines.
As baselines we propose methods with the best performance in bioactivity prediction using microscopy images to date, reported in Hofmarcher et al.
These consist of different convolutional neural network architectures, used in a fully supervised setting, and a method (ºFNNº) that uses expert-designed cell features (Simm et al., 2018; Carpenter et al., 2006; Bray et al., 2017).
The compared methods were trained in a multi-task setting to predict activity labels for 209 tasks, extracted from ChEMBL.
[Table]
Results.
The predictive performance on the downstream activity prediction tasks is reported in Table 1.
CLOOME reached an average AUC of 0 .
714 , which indicates that the learned representations are indeed transferable.
CLOOME even outperformed fully supervised methods, such as M-CNN (Godinez et al., 2017) and SC-CNN (Hofmarcher et al., 2019), with respect to AUC.
Retrieval for bioisoteric replacement and scaffold hopping.
In this experiment, we assessed the ability of CLOOME to correctly retrieve the correct molecular structure given a microscopy image of cells treated with this molecule.
Notably, this is an extremely challenging task for human experts: given a microscopy image of cells, select the molecular structure with which they have been CLOOME can be used to retrieve molecules that could produce similar biological effects on treated cells, i e bioisosteres. treated from a set of thousand candidate structures.
Since cells often do not exhibit any or only subtle morphological changes, this task is highly ambitious.
This image-based retrieval task can also be understood a bioisosteric replacement task (Lipinski, 1986):
Bioisosteres are molecules with roughly the same biological properties or activities, which is highly relevant in drug discovery when a chemical scaffold should be replaced with another, but at the same time its biological activity should be kept.
With this experiment, we evaluate the ability of CLOOME to correctly rank the matched molecular structure given the image.
Other high-ranked structures could be potential bioisosteres, which makes this experiment a proxy for the bioisteric replacement problems (see Figure 1 b) ).
On hold-out data of 2,115 image and molecule pairs, CLOOME ranked the correct molecule in the first place for 3 .
215 % of the cases.
A random method would achieve a value of 1 / 2 , 115 ≈ 0 .
047 %, which indicates a ∼ 70 -fold improvement of CLOOB.
For this task, different hyperparameters and model were selected based on the appropriate validation metric (see Sec.
A.2).
The top-1, top-5, top-10 accuracy are given in Table 2.
Further, some examples are displayed in Figure 2.
This is, to our knowledge, the first system of cell-image-based retrieval of molecular structures.
D ISCUSSION AND CONCLUSION
We have introduced a contrastive learning method for learning representations of molecules based on microscopy images and chemical structures.
On the largest available dataset of this type, we demonstrate that CLOOME is able to learn transferable representations of molecules.
This opens the possibility to re-use the learned representations for activity or property prediction and for other tasks, such as finding bioisosteric replacements of molecules.
Limitations.
Our method currently has several limitations.
Our trained networks are restricted to a particular type of microscopy images, which are acquired with the Cell Painting protocol (Bray et al., 2016).
This protocol has been published and currently there are community efforts (JUMP-CP
Consortium, 2022) to increase the amount of available data.
Large and more diverse datasets of molecule-perturbed cells or internal pharmaceutical company datasets will likely improve the learned representations, both image and structure encoder (Sturm et al., 2020).
Due to the computational complexity, the hyperparameter and architecture space is currently under-explored such that we expect our method to further improve with better hyperparameters or encoder architectures.
Furthermore, it has not escaped our notice that the learned structure encoder can also be used for transfer learning on molecular activities and properties.
Also, it is worth noting that, although linear probing has been extensively used for the purpose of evaluating the quality of representations (Radford et al., 2021; FÈurst et al., 2021), if the latter are very high dimensional, this method presents the risk of overfitting (Alain and Bengio, 2016).
Having addressed these limitations, we nevertheless believe that the representations obtained with CLOOME could be highly useful for the community with respect to different drug discovery efforts.
Future work.
For the reasons mentioned above, we intend to further analyze and test the representations produced by our method for different datasets and in different settings in order to better assess its capabilities and pitfalls.
A CKNOWLEDGEMENTS
ASF’s research position was funded by the European Union’s Horizon 2020 research and innovation program under the Marie Skøodowska-Curie Innovative Training Network - European Industrial
Doctorate grant agreement No. 956832, ªAdvanced machine learning for Innovative Drug Discoveryº.
The ELLIS Unit Linz, the LIT AI Lab, the Institute for Machine Learning, are supported by the
Federal State Upper Austria.
IARAI is supported by Here Technologies.
We thank the projects AIMOTION
(LIT-2018-6-YOU-212), AI-SNN (LIT-2018-6-YOU-214), DeepFlood (LIT-2019-8-YOU213), Medical Cognitive Computing Center (MC3), INCONTROL-RL (FFG-881064), PRIMAL (FFG-873979), S3AI (FFG-872172), DL for GranularFlow (FFG-871302), AIRI FG 9-N (FWF36284, FWF-36235), ELISE (H2020-ICT-2019-
Healthcare KGaA, Verbund AG, Software Competence Center Hagenberg GmbH, T ÈUV Austria,
Frauscher Sensonic and the NVIDIA Corporation.
A A PPENDIX A. N OTATION OVERVIEW Definition
Symbol/Notation
Dimension molecule-perturbed microscopy image x image dimension, e.g. 320 × 320 ×chemical structure of molecule z symbolic, e.g. graph image embedding x d structure embedding z d stacked image embeddings
X d × N stacked structure embeddings
Z d × N stored image embeddings
U d × N stored structure embeddings
V d × N image-retrieved image embedding
U x i d structure-retrieved image embedding
U z i d image-retrieved structure embedding
V x i d structure-retrieved structure embedding
V z i d microscopy image encoder
h
x ( . )
R 320 × 320 × 5 → d molecule structure encoder h z ( . )
M → d
[Table]
parameter
[Table]
functions
[Table]
β
[Table]
A. H YPERPARAMETER SEARCH Hyperparameter
[Table]