Functionally Regionalized Knowledge Transfer for Low-resource Drug Discovery

1 Introduction . Drug discovery brings new candidate medications to billions of people , helping them live longer , healthier and more productive lives . One crux step in drug discovery is virtual screening , which is a fast and cost-effective method that computationally predicts the activity value of a compound against the target protein of a disease . As shown in Figure 1 ( a ) , the hits screened out of large drug libraries of compounds by a virtual screening algorithm are further empirically validated against their in-vivo activities , resulting in leads . After optimizing the ADMET properties ( absorption , distribution , metabolism , excretion and toxicities ) of the leads , we obtain the drug candidates . There have been both traditional machine learning [ 33 ] and deep learning approaches [ 4 ] devoted to virtual screening , while the prediction performance ( i.e. , the hit rate ) is far from satisfactory . The crucial challenge lies in that the number of training compounds whose activities have been tested against the target protein of focus is severely limited . Though state-of-the-art deep learning algorithms typically rely on supervision in the form of thousands to millions of annotated data , it is highly expensive and almost impossible for in-vivo experiments to collect a sufficient set of drug ∗Part of the work was done when H.Y . was a student at Penn State University , correspondence to : Y.W . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . compounds with activity labels . In fact , virtual screening as a computational pre-screening method is desired precisely because of the prohibitive costs of an in-vivo experiment ( i.e. , an assay ) . Fortunately , previous assays conducted by different laboratories around the world towards a wide variety of diseases with different biological target proteins together provide a rich repository for learning the interactions between a protein and a compound . For example , as COVID-19 and SARS share high amino acid sequence identity , previous assays against SARS 3CLpro and PLpro proteases contribute a lot to learning a predictive model for COVID-19 [ 16 ] . We are highly motivated to transfer the knowledge of interactions from this repository to address the scarcity of labeled compounds in the assay against the target protein of our focus , which we name as the target assay for convenience . The challenges of such knowledge transfer are two-fold : ( 1 ) how to share the transferable knowledge but meanwhile accommodate the wide variance between assays , and ( 2 ) how to adequately identify the nearest neighbor assays to the target assay to reduce the risk of negative transfer . Since assays are from different institutions and against various target proteins , the compounds tested and the distribution of activity values vary a lot from assay to assay . As evidenced in Figure 1 ( b ) , there exists a large discrepancy between distributions of activity values for 10 randomly selected assays . The prevalent fine-tuning strategy in transfer learning [ 21 ] , trains a single model on previous assays and fine-tunes it to the target assay – it struggles in predicting accurately for each assay and confuses the most similar assays to the target with the others . Gradient-based meta-learning [ 7 ] has been a promising practice , which learns from previous assays an initialization for a shared predictive model and adapts the model from this initialization to each assay . while the initialization is learned so that the adapted model of each assay generalizes well on testing compounds , maintaining a shared initialization is still insufficient to handle wildly varying assays [ 37 ] and pinpoint the most similar assays . Recent efforts on heterogeneous meta-learning deal with this issue by modulating the shared initialization to different assays via task embedding [ 20 , 35 , 37 ] . Instead of only differentiating initializations , motivated by compositionality and brain functional specialization in neuroscience [ 5 , 26 ] , we aim to push ahead with distinguishing neural sub-networks , or so-called functional regions , each of which consists of a disparate set of parameters . This advancement brings at least the following two benefits . First , the similarity between assays is more accurately measured in a divide-and-conquer manner – only modulation for the initialization weights in those overlapping regions between two assays are considered for comparison . Second , the reduced parameter space prevents the predictive model from overfitting to a limited set of training compounds . We name the resulting meta-learning algorithm as FRML . The predictive model of the FRML is dissected into a sequence of hierarchically organized functional regions . Provided with an assay , the contrastive assay representation network forwards the learned assay embedding to a region localization network . The region localization network locates the most relevant functional regions for the assay in a recurrent manner , to be consistent with the hierarchical organization of functional regions . In the stage of meta-training on previous assays , both the region localization network and the weights for initializations of all functional regions are jointly learned . When it comes to meta-testing , FRML quickly adapts to the target assay via easy assembly of the located regions . We summarize our major contributions as follows . ( 1 ) We propose a novel meta-learning algorithm FRML , which pushes a step forward from differentiating initializations to differentiating neural sub-networks between tasks ; ( 2 ) We demonstrate the effectiveness of FRML on not only virtual screening but also the task of ADMET prediction . ( 3 ) FRML respects the key principle of machine learning models in healthcare – it is interpretable in the relationship between assays . 2 Notations and Problem Definition . In this section , we define some notations and discuss our problem . In drug prediction , we consider each task Ti as an assay which refers to an in-vivo experiment on a group of compounds , and all tasks are sampled from the distribution p ( T ) . Note that we use either task or assay alternatingly in the remainder of this paper . Assuming that we have N historical assays { Ti } Ni=1 as meta-training assays , we aim to generalize a meta-learner from these meta-training assays and quickly adapt it to unseen target assays { Tt } Ntt=1 even with limited amount of annotated data . Here , we define the process of learning well-generalized meta-knowledge from the meta-training assays as the meta-training phase and the adaption process on the target assays as the meta-testing phase . Concretely , for each task Ti , a support set of training samples Dsi = { Xsi , Ysi } = { ( xs , ys ) i , j } nsi j=1 and a query set of testing samples Dqi = { X q i , Y q i } = { ( x q , yq ) i , j } n q i j=1 are sampled from Ti , where nsi and nqi represent the number of support and query samples , respectively . Denote that the feature space is X and the label space is Y , a predictive model ( a.k.a. , base learner ) f : x 7→ ŷ is defined to map a sample x ∈ X to its predicted value ŷ ∈ Y . For each task Ti , the base learner f is updated from the initialization θ0 by minimizing the expected empirical loss L on Dsi , i.e. , minθ L ( θ ; Dsi ) , resulting in the optimal parameters θi . Specifically , the loss function L is defined as mean square error ( i.e. , ∑ ( x , y ) ∈Dsi ‖fθ ( x ) −y‖22 ) or cross-entropy loss ( i.e. , − ∑ ( x , y ) ∈Dsi log p ( y|x , fθ ) ) for regression and classification problems , respectively . In the meta-training phase , the query sets { Dqi } N i=1 of all meta-training assays are used to optimize the initialization of the base learner , so that the final initialization θ∗0 is well-generalized . θ∗0 can be further adapted to each meta-testing task Tt via the corresponding support set Dst . Formally , we define our problem as , Ŷqt = arg max Y q t p ( Yqt |X q i , D s t , fθ∗0 ) . ( 1 ) The well-generalized model initial weights θ∗0 encrypt the comprehensive knowledge learned from meta-training assays . We will detail how to learn θ∗0 in Section 3 . 3 Methodology . In this section , we introduce the proposed framework FRML whose overview is illustrated in Figure 2 . The goal of FRML is to improve the generalization ability for a wide range of and even out-ofdistribution target assays with limited training samples via discriminating functional regions between assays . To achieve this goal , we dissect the base learner into a sequence of functional regions . Given a new assay , we propose a region localization network taking the learned assay representation as input to locate and assemble the most relevant functional regions . Subsequently , FRML can be quickly adapted to the novel assay on the assembled functional region set . In the following subsections , we will first discuss the predictive models for virtual screening and ADMET classification as the base learner and our meta-learning pipeline . Then we elaborate the details of three key components ( i.e. , assay representation learning , localization strategy , and region localization network ) . 3.1 Predictive Models for Drug Discovery and Gradient-based Meta-Learning . We build predictive models for virtual screening and ADMET prediction , both of which are crucial for drug discovery . The input to the predictive models is a drug compound represented by 1024 dimensional Morgan fingerprints [ 28 ] , i.e. , x ∈ R1024 . For virtual screening , the output is the activity value of the compound against the target protein in this assay , i.e. , y ∈ R , while the output for ADMET prediction could be a discrete category or a real value . In our empirical study , we only consider those ADMET prediction tasks of classification , i.e. , y ∈ C , where C denotes the set of property categories . Building on these , we construct a neural network consisting of two fully connected layers as the predictive model , which also serves as the base learner f . We denote the weights for the base learner f to be θ . With the base learner f , we introduce gradient-based meta-learning as the backbone meta-learning framework , which regards the initialization θ0 for the base learner as the transferable knowledge . Apparently , it enjoys the advantage of being independent of problem types . Specifically , here we illustrate the gradient-based meta-learning by using model-agnostic meta-learning ( MAML ) [ 6 ] as an example . In the meta-training phase , MAML obtains the assay-specific model for each assay Ti by updating the parameters θ via the support set Dsi in a few gradient steps starting from θ0 , i.e. , θi = θ0 − α∇θL ( θ ; Dsi ) . ( 2 ) Here α denotes the learning rate for assay adaptation . Though only one gradient step is presented as exemplary in Eqn . ( 2 ) , it is easy to extend to several gradient steps . The crux is to evaluate the adapted assay-specific model θi on the query set Dqi and leverage the result as a feedback to meta-update the initializations θ0 as , θ0 ← θ0 − β 1 N ∑ Ti∈p ( T ) L ( θi ; Dqi ) , ( 3 ) where β is the learning rate for meta-updating . As a result of the meta-training phase , we get the well-generalized initialization θ∗0 for the base learner . In the meta-testing phase , the specific model θt for each target assay Tt with the support set Dst is achieved by a few gradient steps starting from the learned initialization θ∗0 , i.e. , θt = θ∗0 − α∇θL ( θ∗0 ; Dst ) . Finally , the performance is evaluated on the query set Dqt of the target assay Tt . Without loss of generality , we again take MAML as the backbone meta-learning framework of FRML and detail each component in the following .

Summary: In this paper, the authors propose a new meta-learning method for drug discovery. The proposed model transfers knowledge from source drug assays to target drug assays. To better leverage the similarity among drug assays, the proposed model disentangles the base learner to different functional regions and reassembles the selected functional regions for each drug assay. The proposed model is evaluated on virtual screening and ADMET prediction tasks and achieves better performance than other baseline methods.