Contrastive Learning Drug Response Models from Natural Language Supervision

1 Introduction

Phenotypic drug discovery (PDD) [Chen et al., 2023; Vincent et al., 2022a] demonstrates its superiority over target-based approaches [Lu et al., 2022; Li et al., 2021] by identifying drugs, targets, and mechanisms of action (MoA) [Maillard and Pascoe, 2023]. When molecular insights into a disease are limited [Vincent et al., 2022b; Eder et al., 2014], PDD research provides a framework for exploring the uncharted "dark biological matter" territory, which includes biological molecules and cellular processes linked to diseases through the utilization of proteomics and other molecular methods. This contributes to a more efficient drug discovery process [Moffat et al., 2017]. Moreover, the drug response prediction (DRP) task is an essential PDD step, which is gradually developing as a field and becoming a recognized discovery paradigm in academia and the pharmaceutical industry [Liu et al., 2023]. This sustained interest stems from notable successes over the past decade, such as treating schizophrenia with SEP-363856 [Begni et al., 2021], malaria with KAF156 [Ogutu et al., 2023], and atopic dermatitis with Crisaborole [Kim et al., 2023].

∗Corresponding author

DRP primarily focuses on regression tasks. Historically, representation learning has received less attention for regression than classification tasks [Zha et al., 2023]. This is because the DRP task inputs are one cell and drug type, and the output is the half maximal inhibitory concentration (IC50) [Bébéar and Robertson, 1996] of this process. Notably, deep learning-based DRP methods (e.g., [Chu et al., 2023; Liu et al., 2020; Liu et al., 2019b]) show excellent performance when both drug and cell line species are present in the training and test datasets (i.e., many/few-shot [Lee et al., 2022]). However, in practical applications, the IC50 of the tested compounds are unlabeled and not observed in the training datasets (i.e., zero-shot [Wang et al., 2023]). Different drugs exhibit diverse structures and properties within the same cell line, resulting in a significant performance decline under the zero-shot learning condition. This factor presents a serious challenge to the development of PDD. For instance, in the work of [Nguyen et al., 2022], the prediction performance of their models based on graph neural networks (GNNs) exceeded 90% in the Pearson correlation coefficient (PCC) under the many-shot condition. However, the PCC decreased to about 4-32% under the zero-shot condition.

Under the zero-shot condition, traditional methods represent sample features inadequately. This is mainly due to overlapping and disordered sample features. These challenges reduce the effectiveness of the DRP model. As illustrated in Figure 1, we compare our method's representations with that of traditional approaches using the genomics of drug sensitivity in cancer database (GDSC2) [Yang et al., 2013], visualized with the t-distributed stochastic neighbor embedding (tSNE) method. Applying the mean squared error (MSE) loss function as an example, traditional methods yield disordered, fragmented, aggregated, and overlapping mappings, leading to unsatisfactory results during various regression tasks. It is worth noting that when the model fails to capture the mapping pattern of numerical labels, it can only map unseen samples into the known space in a disordered manner, resulting in samples with different label values that cannot be distinguished. This has a detrimental impact on the DRP model's performance during practical applications. Therefore, there is an urgent need to enhance a model's ability to represent continuous values in an ordered and scattered manner to improve generalization performance under the zero-shot condition.

In recent years, contrastive language image pre-training (CLIP) [Zhang et al., 2022] has provided a powerful multimodal representation learning framework, enabling computers to better understand and process the semantic relationships between images and text, like GLIDE [Nichol et al., 2022]. Moreover, CLIP research indicates that state-of-theart (SOTA) image representations can be achieved with a simple pre-training task using a dataset of 400 million image-text pairs. Furthermore, contrastive learning with labels has been theoretically proven to enhance the performance of learned representations [Ji et al., 2023] in downstream tasks [Wang et al., 2022]. Consequently, in the realm of drug discovery, it becomes possible to establish a connection between drug response data and annotated text, learning representations from the text [Fang et al., 2022]. Subsequently, these representations may be used to enhance zero-shot learning performance during natural language supervision contrastive learning [Khosla et al., 2020; Gunel et al., 2020].

To address this challenge, we present the CLDR, a contrastive learning framework with natural language supervision for drug response prediction. The CLDR framework transforms numerical labels used to represent drug response into text using customized prompts. First, drugs and cell lines from the same samples are encoded using natural language and traditional fusion methods, respectively. Then, a contrastive learning strategy is employed to map the drug and cell line feature space and the text with labels using natural language into the shared representation space. This strategy maximizes the similarity between related samples while minimizing those between unrelated ones. In this study, we construct a common-sense numerical knowledge graph (CNKG), drawing inspiration from the ordinal number definition [Agustito et al., 2023]. The CN-KG constrains the text representation order, improving the fusion encoder representation for drugs and cell lines. Furthermore, through contrastive learning pre-training, ordered and scattered natural language representations are aligned and mapped to a unified highdimensional space with those of the fusion encoder for drugs and cells.

For the practical evaluation, we validate the method using a dataset comprising over 150,000 samples from the GDSC2 dataset. Consequently, all the methods display increases of 7.88%, 19.49%, 17.83%, 14.29%, 13.04%, and 31.46% after adopting our framework. The proposed CLDR method effectively establishes connections between drug response data and the corresponding labeled text, enhancing generalizability and improving the PDD success rate. Section 3.4 provides detailed theoretical proof of the CLDR method's validity. Our contributions are as follows:

• We propose the CLDR method, a novel contrastive learning framework with natural language supervision for the DRP task. The CLDR method effectively constructs links between drug response data and labeled text.
• We construct a CN-KG to capture the continuous nature of sample order to improve the fusion encoder representation for drugs and cell lines.
• The extensive experiments using the GDSC2 dataset demonstrate that employing the CLDR method leads to a notable improvement in the DRP methods, enhancing results by at least 7.8% and up to 31.4%.

2 RelatedWork

DRP exploration has become feasible due to drug response studies on a large number of cell lines, exemplified by GDSC2 [Yang et al., 2013] and cancer cell line encyclopedia (CCLE) [Barretina et al., 2012]. Various DRP techniques have been introduced, and categorized into 1DCNN, graph, and transformer methods. Generally, non-graph-based methods employ convolutional neural networks (CNNs) and multilayer perceptrons (MLPs) for information retrieval. These methods encode drug molecules in the simplified molecular input line entry specification (SMILES) string format and directly extract gene sequence features using 1DCNN, overlooking the pharmacological and structural attributes of the molecule. Conversely, graph-based methods [?; Liu et al., 2020] concentrate on converting drug molecules into graph structures, utilizing GNNs for representation instead of directly extracting features from the SMILES strings. Additionally, in contrast to GNNs, the transformer-based methods [Chu et al., 2023; Jiang et al., 2022] avoid introducing any structural inductive bias at intermediate layers [Chu et al., 2023], thereby mitigating the GNNs' expressivity limitations [Park et al., 2020].

(a) Contrastive pre-training

In summary, the conventional DRP method employs a unimodal end-to-end approach for IC50 training and prediction. However, its effective applicability is limited under the zeroshot condition, affecting the PDD's efficiency.

3 Methods

During the pre-training stage, the CLDR method principally comprised the fusion encoder for the drug and cell line features extraction and fusion, the caption encoder for the drug response process text, and the natural language supervision module incorporating the CN-KG.

Subsequently, during the fine-tuning phase, standard regression loss functions are utilized. The primary goal of the pre-training phase is to map the feature spaces of the drug, cell line, and the labeled text, along with the labeled text, into a shared representation space.

3.1 The Fusion Encoder

The DRP task's inputs were the drugs and cell lines, and the specific method for the fusion encoder is not discussed in detail in this section.

Moreover, the input types for the drug extractor were the SMILES sequence or molecular graph. We uniformly describe the feature representations process for the drug (denotes as di) and the cell line (denotes as ci) as follows:

where, N represents the total number of samples i ∈ [0,N], Φdrug and Φcell denotes the drug and cell line encoders of the initial DRP methods, while Pi and Qi are the representations of the drug di and the cell line ci. The fusion encoder can be expressed as follows:

where F ∈ RN×m is the feature after encoding using the fusion encoder, Φf denotes the fusion encoder, which can be replaced in different methods.

3.2 The Caption Encoder

In the caption encoder, the text-based description was the second coding modality for the drug response process. The inputs a, b and the output c in the DRP task are described uniformly with the prompt Fprompt(a, b) as follows:

The drug response value between [a] and [b] is [c]

where a denotes the drug's SMILES string representation, b represents the name of the cell line, and c is the quantitative output value of the reaction level between the drug and the cell lines (i.e., the IC50). To represent IC50 efficiently, we defined a transformation method Fnum2str(c) that converts it into English words in character order, where the decimal point is encoded as "point". For example, as shown in Figure 2(a), "0.513" can be converted to "zero point five one three."

To facilitate continuity constraints on the values, the caption encoder was designed to consist of two encoders, including Φtext, Φnumber, which encodes text Ttext = Fprompt(a, b) and value-based descriptions Tnumber = Fnum2str(c), respectively.

Then, the two feature vectors were merged, fusing the descriptive information of the drug reaction process using the multi-layer transformer Φcap as follows:

where T ∈ RN×m is the description text feature with the caption encoder encoding.

3.3 Supervising Natural Language Using a Knowledge Graph

To guarantee perceiving the number continuity when encoding numerical values in natural language, we constructed a

(b) Finetune

CN-KG according to [Liu et al., 2017; Duan et al., 2021], which uses linear structures to construct graphs that accurately and intuitively represent numbers and their relationships.

As shown in Figure 2(c), the CN-KG's entity sets are sequences of ordered numbers, denoted as E. These CN-KG entities are integers, the significant number of c multiplied by the specified precision, where the minimum and maximum numbers are customized for the specific task. The E are linked by a single relationship type called "is less than," denoted as L, which ensures the numbers' transfer properties are captured.

When dealing with numerical information, we must consider how to incorporate numerical features into the framework to represent their relationships more accurately. To enhance the number encoder's efficiency in the numerical relationships using the CN-KG, we proposed a margin-based loss function LKGE for the CN-KG embedding. We aim to minimize the differences in embedding vectors between the entities' set E and the single relationship L (i.e. is less than).

where, [x]+ denotes the positive part of x, γ > 0, and is a margin hyperparameter. The set S is composed of the triplets (h, l, t), with h, t ∈ E, l ∈ L. The embeddings h, l, t obtain values in Rk (k is a hyperparameter) and are denoted with the same letters in boldface characters. In addition, the L1 or L2-norm can be used for the similarity measure d.

Considering that the CN-KG can restrict number relationships, the set of entities E are encoded by the number encoder Φnumber (Tnumber(E)), and the relationship l is represented by a learnable embedding.

Pre-training For the i-th representations (di, ci) generated by the fusion encoder and the j-th captions (dj, cj, yj) produced by the caption encoder in a batch B, we normalized the feature vectors in a hyper-sphere using ui := Φf (di,ci) ∥Φf (di ,ci)∥ and vj := Φcap(dj,cj,yj) ∥Φcap(dj,cj,yj)∥ . The similarity between ui and vj was calculated as uTi vj. Finally, a supervised contrastive loss function was used to train the model:

where, N is the size of the batch B, and σ is the temperature for scaling the logits.

CLDR's goal during the pre-training phase is to jointly optimize the following contrast and CN-KG embedding loss functions:

where α represents the joint optimization weight adjustment factor for the two loss functions.

Fine-tuning During the fine-tuning stage, we employed the MSE loss function for supervised regression on the fusion encoder. A regression output layer Φmlp based on the MLP was designed after the fusion encoder as follows:

where yi is the normalized value of IC50 corresponding to {di, ci}.

3.4 Theoretical Analysis

In this section, we theoretically prove that jointly optimizing LNCE and LKGE enables the fusion encoder Φf to obtain continuous regression-aware representation.

Notations Let {xi, yi} be the inputs and outputs respectively, where yi is a sorted label with the ordering y1 ≤ y2 ≤ · · · ≤ yn. The δ ∈ (0, 1) denotes the minimum interval of normalized labels yi ∈ [0, 1].

First, based on the loss function LKGE and its expectation that t should be a nearest neighbor of h + l [Bordes et al., 2013] (see Section 3.3), we formulated:

where l is the only relationship embedding type in the CNKG and ϵ ∈ R is a model perturbation. If we denote Φnumber as N, then the following lemma can be inferred:

Lemma 1 (Equal interval representation of N). For any 0 < δ < 1, two perturbations ϵ0, ϵ1 exist to make:

Lemma 1 implies that Φnumber can learn continuous representations that capture the intrinsic sample order of the regression target. Next, based on the LNCE constraints on positive xi, yi and negative samples x-i , yi , we expect the following:

where C and F denote Φcap and Φf , respectively. Also, M := C(xi, yi) - F(xi) denotes the degree of alignment between two modalities. The subsequent lemma can be derived from Equation (9):

Lemma 2 (Upper bound of LNCE). For any i ∈ j, the formula ϵ2 > 0 such that ∥M∥ ≤ ϵ2 exists.

where ϵ2 is a value related to ϵ0 and ϵ1. Lemma 2 states that the upper bound on LNCE is equal to the upper bound on the distance between the distributions of features C and F in a uniform representation space.

Thus, under the LNCE constraint, the property of the equal interval representation of N will conditionally transfer to F:

Theorem 1 (Main theorem). For any i ∈ j, the formula θ ∈ (0, 1) such that if δ is close to 0, then ∥F(xi) -F(xi+1)∥ ≤ 2ϵ2 exists.

More generally, if Q := F(xi) - F(xi+1), then we can derive the following expression:

Theorem 1 suggests an upper bound on the learned continuous representation of F. The upper bound is jointly determined by that of M and δ, corresponding to LNCE and LKGE, respectively. In addition, we observed that reducing δ and LNCE enhances the representation of Φf .

3.5 AlgorithmComplexity

When a model encounters many unlabeled compounds, prediction time and accuracy become the most critical factors, while the training time and calculated cost may be deemed negligible. Our method exhibits a time complexity of O(2n2 + 3n) during pre-training. This complexity shows moderate variations when integrated with specific DRP models. During the fine-tuning and inference phases, the algorithm's complexity is solely determined by the specific model, and our method does not play a role in these processes.

4 Experiments

4.1 Experiment Settings

Validation strategy In the zero-shot learning context, simulating the PDD practical application scenario required clustering the response data by drug type. Then, the data was randomly divided into training, validation, and testing sets by a ratio of 8:1:1 and the drug type as the division standard. As a result, the DRP model encounters unknown compounds during the testing phase. This zero-shot learning condition presents a greater level of complexity compared to the random splitting of the entire dataset commonly employed in supervised learning [Bai et al., 2023; Wang et al., 2023]. During the ablation study, all the fusion encoder experiment branches adopted the TransEDRP structure [Li and Hu, 2022].

Metrics To comprehensively evaluate the CLDR's impact, we employed several evaluation metrics: root mean square error (RMSE) to gauge deviation, PCC [Cohen et al., 2009] to assess linear correlation, Spearman's rank correlation coefficient (SRC) [Sedgwick, 2014] to measure monotonicity, and margin ranking loss (Rank) [Liu et al., 2019a] to evaluate ranking performance.

4.2 Overall Experiment

This paper proposes the CLDR framework which leverages the caption encoder and CN-KG's powerful representation capabilities to align the fusion encoder for drug response features, improving the DRP model's generalizability for zeroshot learning.

Therefore, we broadly selected representative DRP methods based on deep learning models such as CNNs, GNNs, and transformers, to verify the CLDR's effectiveness and generalizability. In the overall experiment, we selected and tested six methods on the most widely employed GDSC2 dataset.

As shown in Table 1, the original models' prediction results are denoted by Original, while the combinedCLDR framework's outcomes are represented by +CLDR .

All the methods demonstrated increases of 7.88%, 19.49%, 17.83%, 14.29%, 13.04%, and 31.46% after our framework was adopted. The results show that the CLDR framework can be generally applied to various DRP methods to enhance performance under the zero-shot condition.

4.3 Ablation Study

CN-KG To verify that the CN-KG enhances the continuous numerical text representation, we designed the experiment without using the knowledge graph embedding. Specifically, we did not constrain the number encoder with LKGE. Also, we maintained pre-training and the fine-tuning approach remained unchanged. Based on the ablation experiments in Table 2, the model's performance increased by 9.6% with LKGE.

Furthermore, as shown in Figure 3, the numerical text representation vectors employing the CN-KG are compared. This is because contrastive learning enables the fusion encoder to learn the number encoder's continuous numerical representation capability. Although the fusion encoder without LKGE can still represent the continuous sample information in a regular manner (as shown in Figure 1), it fails to capture the continuous nature of sample orders, resulting in feature overlap and making distinguishing between different samples difficult.

Precision During pre-training, the constraint of continuous values with different precision has various effects on the performance of models based on Equation (10). Thus, we designed ablation experiments to investigate specific effects. With all other conditions consistent, the precision grids 0.1, 0.01, and 0.001 were created to pre-train the model. The results are shown in Table 2, where the performance of the model at 0.01 significantly exceeds that at 0.1 and 0.001. This is mainly because the model needs to balance the precision and perturbation. As shown in Table 1, the MSE loss of various DRP methods is around 0.005.

As a result, the model cannot accurately predict the thousandth part. Additionally, there is a possibility that the value of the thousandth part will be a perturbation factor during contrastive text learning.

Caption During the pre-training phase, the caption encoder's prompts contain the drug response inputs a, b and the IC50 results c. Since the inputs contain drugs and cell lines, to verify whether the inputs containing the drug reaction captions improve the fusion encoder's ability to discretely map various drug reaction pairs, we removed the input texts to compare the results. In the caption encoder, we used a and b as Text, and c as Number. As shown in Table 2, the model improves on all metrics with an average increase of 9.1% after adding the inputs to the text description. This is because the description of the drug response, a real biological phenomenon, requires at least three kinds of information, drug, cell type, and reaction result. Otherwise, the drug response process will collapse into a one-dimensional space, affecting the transformation of the feature space from the caption to the fusion encoder.

4.4 Computational Analysis

We performed a computational analysis to compare the time cost and parameter increases to the profits made during the pre-training phase.

The experiment utilized an Intel Xeon E5-2690 v3 processor with 12 cores (i.e., 24 threads) and a clock frequency of 2.60GHz. Additionally, the RTX 4090 GPU was utilized. During the training time experiments, each model underwent testing 10 times in the GPU environment.

As shown in Table 3, AVG w. CLDR represents an average metric enhancement facilitated by the CLDR method. Incorporating the CLDR framework significantly extends the training time and parameter count by an additional 4.21 seconds and 5,811M, respectively. Notably, the CLDR method demands fewer training epochs in the pre-training and finetuning phases, resulting in a shorter training duration than the original models.

In summary, integrating the CLDR framework while increasing the parameter count substantially enhances generalization performance. Our work represents a breakthrough when compared to traditional DRP methods.

4.5 Case Study Analysis

During the overall experiments, the CLDR method demonstrated high accuracy and generalizability, presenting an advanced DRP framework in drug discovery research. To validate that our framework improves numerical values and increases drug screening hit rates, we conducted a case study analysis as illustrated in Figure 4. We used unknown drugs as the screen molecules in our test dataset during preclinical drug screening.

Then, we compared the screening results of different methods with and without incorporating our framework. According to HR@n in recommendation systems [Lee et al., 2010], we defined a novel evaluation measure specifically designed for drug screening, denoted as Shot@x:

where T and P represent the true labels and predicted results, respectively. topx represents top number of drugs with the smallest drug response values, Ncell represents the number of c cell lines in the test dataset, and Hit(·, ·) is a counting function that is 1 when the two inputs are equal and 0 otherwise.

We present the capabilities of three representative methods, tCNNs, GraphDRP, and TransEDRP, under the original strategy (i.e., MSE) and the CLDR framework. As shown in Figure 4, We assessed the predicted IC50 rankings of multiple drugs using the model with the actual values in order. Then, the probability of scoring a hit with the first optimal drug when recommending x compounds.

Due to the application of the CLDR framework, the probability of the model scoring a hit with the first optimal drug was consistently above 32%. In other words, for a set of screened compounds, our model achieves a hit rate of about 32% on a single chance for about 320 cell lines, and it has a cumulative hit probability of between 46% and 51% on the second attempt. Notably, as the number of x attempts increases, the results between methods become more similar; however, the number of attempts is limited in real drug studies. Thus, the DRP task requires the model to efficiently recommend the most effective drug in as few attempts as possible. In this context, our method improves the success rate by about 10% compared to traditional methods when only one attempt is available. This is an astounding development.

5 Conclusion

In this paper, we propose the CLDR framework, a contrastive learning framework with natural language supervision for DRP. The CLDR framework converts regression labels into text, which is merged with the drug response captions as a second modality for the samples in contrast to traditional encoding modalities. To enhance the continuous representation capability of the numerical text, the CN-KG was proposed to constrain the caption encoder's ability to perceive continuous values. In addition, we provided detailed theoretical evidence of the CLDR method's validity and conducted validation experiments on the GDSC2 dataset. This demonstrated the CLDR framework's ability to establish a link between drug-response data and valuable labeled text, improving the PDD's generalizability and success rate.

Acknowledgments

The work of Wenbin Hu was supported by the National Key Research and Development Program of China (2023YFC2705700). This work was supported in part by the Natural Science Foundation of China (No. 82174230), Artificial Intelligence Innovation Project of Wuhan Science and Technology Bureau (No. 2022010702040070), Natural Science Foundation of Shenzhen City (No. JCYJ20230807090211021).