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: This paper seeks to remedy the low-resource drug discovery problem by transferring the knowledge from previous assays, namely in-vivo experiments, by different laboratories and against various target proteins. The authors propose a functional rationalized meta-learning algorithm FRML for such knowledge transfer. The compositionality of the model improves the capacity of generalization to various and even out-of-distribution tasks. Empirical results on both virtual screening and ADMET prediction validate the superiority of FRML over state-of-the-art baselines powered with interpretability in assay relationships.