Multi-constraint molecule generation method based on knowledge distillation and reinforcement learning

Abstract

The invention discloses a multi-constraint molecule generation method based on knowledge distillation and reinforcement learning, and belongs to the technical field of drug molecule generation models. The method comprises the steps that the maximum length of a generated molecule is set for a pre-trained agent model, model reasoning is conducted on the agent model, and the molecule meeting various attribute constraint conditions is output; wherein the training of the proxy model comprises the following steps: training based on a target attribute constraint condition and molecular SMILES expression data by adopting a conditional transformer to obtain a prior model; a distillation RNN model of a three-layer circulation unit is adopted, likelihood distribution of molecules sampled in the prior model is learned, a reinforcement learning method is used for fine adjustment of the distillation RNN model, and a proxy model is obtained; according to the method, the generation model is preprocessed without influencing the output diversity of the generation model in the multi-constraint task, and meanwhile, the efficiency of outputting required molecules by the molecule generation model is improved.

Description

Technical Field

[0001] The present invention relates to a multi-constraint molecule generation method based on knowledge distillation and reinforcement learning, and belongs to the technical field of drug molecule generation models. Background Art

[0002] Drug molecule design is essentially a multi-parameter optimization problem, the goal of which is to design and discover new molecules with certain ideal properties. Therefore, how to generate new molecular structures and optimize the key properties of molecules (such as bioactivity, drugability, safety, selectivity and synthesizability) are two key issues that affect the success or failure of drug design. It is estimated that the number of molecules that can be explored in the drug-like chemical space is as high as 10 23 ~10 60 How to intelligently generate molecular structures in such a huge chemical space and further achieve rapid structural evolution is a huge challenge facing drug molecule design.

[0003] With the rapid development of computer science and high-performance computing, artificial intelligence (AI) methods have achieved success in fields such as image processing, pattern recognition, and natural language processing. In recent years, machine learning, particularly deep learning, has been widely applied in drug discovery, for example, to predict the properties and activities of compounds and their interactions with protein targets. In recent years, a variety of deep molecular generative models based on AI techniques, such as recurrent neural networks (RNNs), variational autoencoders (VAEs), generative adversarial networks (GANs), and reinforcement learning (RL), have demonstrated great potential in exploring complex chemical spaces and are expected to become key technologies for de novo drug design. Summary of the Invention

[0004] The purpose of the present invention is to overcome the shortcomings of the existing technology and provide a multi-constrained molecule generation method based on knowledge distillation and reinforcement learning. By preprocessing the generative model without affecting its output diversity in multi-constrained tasks, the efficiency of the molecular generation model in outputting the desired molecules can be improved. In order to achieve the above objectives / solve the above technical problems, the present invention is implemented by adopting the following technical solutions: A multi-constraint molecule generation method based on knowledge distillation and reinforcement learning, the method comprising: Setting a maximum length for generated molecules for a pre-trained proxy model, and outputting molecules that meet multiple attribute constraints by performing model inference on the proxy model; The training of the proxy model includes: using a conditional transformer to train a priori model based on target attribute constraints and molecular SMILES expression data; using a distilled RNN model with three layers of recurrent units to learn the likelihood distribution of molecules sampled in the prior model, and using reinforcement learning methods to fine-tune the distilled RNN model to obtain a proxy model.

[0005] Optionally, the prior model is trained using a multi-attribute molecular dataset, and the multi-attribute molecular dataset includes a training set and a test set.

[0006] Optionally, the prior model target property constraints and the molecular SMILES expression data are jointly embedded, combined with position encoding, with the start token [BOS] + SMILES as input and the end token [EOS] + SMILES as the target output to generate molecules with target properties.

[0007] Optionally, the training method of the distilled RNN model includes: The distillation model is obtained by training an initial distillation model based on a set of molecules with target characteristics generated by a priori model.

[0008] Optionally, the target attribute constraint condition of the molecule of the target characteristic is consistent with the target attribute constraint condition of the prior model.

[0009] Optionally, the loss function for the proxy model training is defined as follows: The log-likelihood of the molecule SMILES string sampled from the surrogate model is recorded as , and then input the generated molecular SMILES string into the distillation model to estimate the log-likelihood median , the reward and punishment scores of each generated molecule , multiplied by a weighting coefficient, and added Get the augmented likelihood:

[0010] Among them, reward and punishment scores Obtained by the target attribute prediction model, is the weighted coefficient, and the final loss function is: .

[0011] Optionally, the conditional transformer removes the multi-head attention layer between the encoder and decoder, and adds an additional constraint code c to distinguish molecules with ideal properties in the training set.

[0012] Optionally, the constraint code c is an n-bit string, indicating whether a molecule satisfies each of the n constraint conditions.

[0013] Optionally, the reinforcement learning method adopts the reinforcement learning algorithm in REINVENT.

[0014] Optionally, the proxy model and the distilled RNN model are RNN models with consistent structures.

[0015] Compared with the prior art, the present invention has the following beneficial effects: This paper proposes a multi-constrained molecule generation method based on knowledge distillation and reinforcement learning. By pre-processing the generative model without affecting its output diversity in multi-constrained tasks, the efficiency of the molecule generation model in outputting the desired molecules is improved. This method uses conditional transformers to construct the generative model, adopts a knowledge distillation model to reduce the model's complexity and increase the diversity of the generated molecules, and then fine-tunes it through reinforcement learning to obtain the final molecule generation model under multi-constrained conditions. In the present invention, a three-stage model construction method is used to achieve a more efficient and diverse multi-constrained molecular generation model based on the conditional transformer c-Transformer model through knowledge distillation and reinforcement learning methods. BRIEF DESCRIPTION OF THE DRAWINGS

[0016] Figure 1 FIG2 is a schematic diagram of a model of a multi-constraint molecule generation method based on knowledge distillation and reinforcement learning according to the present invention; Figure 2 Shown is a flowchart of the prior model training of the multi-constraint molecule generation method based on knowledge distillation and reinforcement learning of the present invention. DETAILED DESCRIPTION

[0017] In order to make the technical means, creative features, objectives and effects achieved by the present invention easier to understand, the present invention is further described below in conjunction with specific implementation methods.

[0018] In the description of the present invention, it should be understood that the terms "center", "longitudinal", "lateral", "up", "down", "front", "back", "left", "right", "vertical", "horizontal", "top", "bottom", "inside", "outside" and the like indicate orientations or positional relationships based on the orientations or positional relationships shown in the accompanying drawings, and are only for the convenience of describing the present invention and simplifying the description, rather than indicating or implying that the device or element referred to must have a specific orientation, be constructed and operated in a specific orientation, and therefore cannot be understood as limiting the present invention. In addition, the terms "first", "second", etc. are only used for descriptive purposes and cannot be understood as indicating or implying relative importance or implicitly indicating the number of the indicated technical features. Therefore, features defined as "first", "second", etc. may explicitly or implicitly include one or more of the features. In the description of the present invention, unless otherwise specified, "multiple" means two or more.

[0019] In the description of the present invention, it should be noted that, unless otherwise expressly specified or limited, the terms "mounted," "connected," and "connected" should be understood in a broad sense. For example, they may refer to fixed connections, detachable connections, or integral connections; mechanical connections or electrical connections; direct connections or indirect connections through an intermediate medium; and internal communication between two components. Those skilled in the art will understand the specific meanings of the above terms in the present invention based on specific circumstances.

[0020] like Figure 1 As shown, the present invention discloses a multi-constraint molecule generation method based on knowledge distillation and reinforcement learning, the method comprising: Setting a maximum length for generated molecules for a pre-trained proxy model, and outputting molecules that meet multiple attribute constraints by performing model inference on the proxy model; The training of the proxy model includes: using a conditional transformer to train a priori model based on target attribute constraints and molecular SMILES expression data; using a distilled RNN model with three layers of recurrent units to learn the likelihood distribution of molecules sampled in the prior model, and using reinforcement learning methods to fine-tune the distilled RNN model to obtain a proxy model.

[0021] The present invention utilizes a three-stage model construction method to achieve the generation of multi-attribute constrained molecules, as specifically described below: This multi-constrained molecule generation method, based on knowledge distillation and reinforcement learning, consists of three main models: a prior model, a distillation model, and a proxy model. The prior model is trained using a conditional transformer (c-Transformer) based on target constraints and molecular SMILES expression data to generate a usable prior model.

[0022] The distillation model uses an RNN model with three layers of recurrent units to learn the likelihood distribution of molecules sampled in the prior model. Finally, the distilled RNN model is fine-tuned using reinforcement learning methods to obtain the proxy model.

[0023] In this embodiment, if Figure 2 As shown, the training method of the prior model is as follows: The dataset for the prior model can be selected from various public molecular datasets, such as zinc20, PubChem, ChEMBL, and various public literature datasets. In addition to SMILES expressions, the molecules in the dataset must also have target attribute feature indicators.

[0024] The prior model uses a conditional transformer (c-Transformer) as its model network architecture. Unlike traditional Transformers, this model achieves molecule generation by simplifying the decoder and removing the multi-head attention layer between the encoder and decoder. Furthermore, the c-Transformer adds an additional constraint code c to distinguish molecules with desirable properties in the training set. The constraint code c is an n-bit string indicating whether a molecule satisfies each of the n constraints (i.e., each property corresponds to a value of 0 or 1). During training, the prior model learns a joint embedding of the constraints and SMILES, incorporating positional encoding. It takes a start token (BOS) + SMILES as input and an end token (EOS) + SMILES as the target output, generating molecules with the desired properties. The generative model is constructed by performing multi-constraint analysis on attributes such as QED (quantitative estimate of drug-likeness), SA (synthetic accessibility score), and DRD2 (namely dopamine type 2 receptor) activity. Dataset Construction: Molecules containing properties such as QED, SA, and DRD2 activity were collected from Zinc20, ChEMBL, and published literature to construct a multi-attribute molecular dataset. The target attributes were set to: QED >= 0.6, SA <= 4, and DRD2 >= 0.5. A total of 983,135 data points were collected, split into training and test sets in an 8:2 ratio, resulting in a total of 786,508 training sets and 196,627 test sets.

[0025] The prior model training was performed on an NVIDIA V100 32G GPU server using the training set to obtain the prior model parameters that can generate molecules with target properties.

[0026] In this embodiment, the training method of the knowledge distillation model includes: the optimization efficiency of the reinforcement learning algorithm for the Transformer series model is low, and the chemical space reconstructed by the c-Transformer model used in the prior model is too concentrated, which may cause the generated results to fall into the dilemma of local optimal solution.

[0027] To address these issues, the present invention uses knowledge distillation to refine the prior model into a more compact structure that preserves only the conditional chemical space, enabling the model to generate molecules with greater diversity and novelty. This method employs a three-layer RNN model with a sample of molecules generated by the prior model as the training set. The knowledge distillation model is trained using the set of molecules generated by the prior model.

[0028] Based on the trained prior model, we generated a total of 1 million target molecules (QED >= 0.6, SA <= 4, DRD2 >= 0.5; consistent with the target settings of the prior model). We used these 1 million molecules generated by the prior model as the training set for the knowledge distillation RNN model and obtained the knowledge distillation RNN model parameters.

[0029] In this embodiment, the training of the proxy model is as follows: Reinforcement learning (RL) can be used to fine-tune the parameters of generative models to guide the free parameter space towards a set with the best objective function value (e.g., bioactivity). Although RL has been shown to be capable of finding a highly diverse set of bioactive molecules, it still requires a large number of optimization steps to learn these patterns through reward-based training. In previous studies, to alleviate this problem and improve the efficiency of RL, transfer learning (TL) was first used before RL operation to quickly focus on certain areas in the chemical space. However, a side effect of transfer learning (TL) is a significant contraction of the accessible chemical space, which makes the model prone to falling into local optima, meaning that the generated molecules are often highly similar to the compounds in the transfer learning (TL) training set. The surrogate model is a multi-objective molecule generation model that is finally usable by fine-tuning the knowledge distillation model through a reinforcement learning algorithm.

[0030] Specifically, the proxy model uses an RNN model with the same structure as the distillation model. The reinforcement learning algorithm regards the molecular SMILES expression generated by the model as a behavior and gives corresponding rewards and punishments for each behavior. The proxy model eventually learns the conditional probability p(A|S) of behavior A under state S.

[0031] The loss function of the agent model training guided by the reinforcement learning algorithm is defined as follows: the log likelihood of the molecular SMILES string sampled from the agent model is recorded as , and then input the generated molecular SMILES string into the distillation model to estimate the log-likelihood median , the reward and punishment scores of each generated molecule , multiplied by a weighting coefficient, and added Get the augmented likelihood:

[0032] Among them, reward and punishment scores Obtained by the target attribute prediction model, is the weighted coefficient, and the final loss function is: .

[0033] The resulting distilled RNN model was fine-tuned using the reinforcement learning algorithm in REINVENT to obtain the final performance-optimized multi-constrained molecular generation model. The QED and SA properties were calculated using the RDKit tool, while the DRD2 property was predicted using an open-source SVM support vector machine model.

[0034] The present invention is based on a trained proxy model and sets the maximum length of generated molecules, so as to directly perform model inference and generate target molecules under multiple constraints.

[0035] The above is only a preferred embodiment of the present invention. It should be pointed out that for ordinary technicians in this technical field, several improvements and modifications can be made without departing from the technical principles of the present invention. These improvements and modifications should also be regarded as the scope of protection of the present invention.

Claims

1. A multi-constrained molecule generation method based on knowledge distillation and reinforcement learning, characterized in that: The method comprises: Setting a maximum length of generated molecules for a pre-trained proxy model, and outputting molecules that meet multiple attribute constraints by performing model reasoning on the proxy model; Among them, the training of the proxy model includes: using a conditional transformer to train based on the target attribute constraints and the molecular SMILES expression data to obtain a prior model; using a distilled RNN model with three-layer recurrent units to learn the likelihood distribution of the sampled molecules in the prior model, and using a reinforcement learning method to fine-tune the distilled RNN model to obtain a proxy model.

2. The multi-constrained molecule generation method based on knowledge distillation and reinforcement learning according to claim 1, characterized in that: The prior model is trained using a multi-attribute molecular dataset, which includes a training set and a test set.

3. The multi-constrained molecule generation method based on knowledge distillation and reinforcement learning according to claim 1, characterized in that: The prior model target attribute constraints and the molecular SMILES expression data are jointly embedded, combined with position encoding, with the start token [BOS] + SMILES as input and the end token [EOS] + SMILES as the target output to generate molecules with target properties.

4. The multi-constrained molecule generation method based on knowledge distillation and reinforcement learning according to claim 3 is characterized in that: The training method of the distilled RNN model includes: Based on the molecular set of target characteristics generated by the prior model, the initial distillation model is trained to obtain the distillation model.

5. The multi-constrained molecule generation method based on knowledge distillation and reinforcement learning according to claim 4, characterized in that: The target attribute constraint condition of the molecule of the target characteristic is consistent with the target attribute constraint condition of the prior model.

6. The multi-constrained molecule generation method based on knowledge distillation and reinforcement learning according to claim 1, characterized in that: The loss function for the proxy model training is defined as follows: The log likelihood of the SMILES string of the molecule sampled from the proxy model is recorded as , and then input the resulting molecular SMILES string into the distillation model to estimate the log-likelihood median , the reward and penalty scores for each generated molecule , multiplied by a weighting factor, and then added Get the augmented likelihood: ; Among them, the reward and punishment scores Obtained by the target attribute prediction model, is the weighted coefficient, and the final loss function is: 。 7. The multi-constrained molecule generation method based on knowledge distillation and reinforcement learning according to claim 1, characterized in that: The conditional transformer removes the multi-head attention layer between the encoder and decoder, and adds an additional constraint code c to distinguish molecules with ideal properties in the training set.

8. The multi-constrained molecule generation method based on knowledge distillation and reinforcement learning according to claim 7, characterized in that: The constraint code c is an n-bit string, indicating whether a molecule satisfies each of the n constraint conditions.

9. The multi-constrained molecule generation method based on knowledge distillation and reinforcement learning according to claim 1, characterized in that: The reinforcement learning method adopts the reinforcement learning algorithm in REINVENT.

10. The multi-constrained molecule generation method based on knowledge distillation and reinforcement learning according to claim 1, characterized in that: The proxy model and the distillation RNN model are RNN models with consistent structures.

Images