Machine learning model distillation for protein design
A distilled machine learning model is produced via initializing a first model with initial weights. An input protein sequence is input into both the first model and a folding protein model, wherein the inputting to the first model generates logits, and wherein the inputting to the folding protein model generates one or more predictive metrics. The one or more predictive metrics are discretized into classes and a first cross-entropy loss is computed based on the logits and the classes. The first model is optimized based on the first cross-entropy loss so that the optimized first model is the distilled machine learning model and an additional machine learning model is trained, using the distilled machine learning model, to perform a downstream protein modeling task.
1 . A computer-implemented method comprising:
producing a distilled machine learning model via:
initializing a first model with initial weights;
inputting an input protein sequence into both the first model and a folding protein model, wherein the inputting to the first model generates logits, and wherein the inputting to the folding protein model generates one or more predictive metrics;
discretizing the one or more predictive metrics into classes;
computing a first cross-entropy loss based on the logits and the classes; and
optimizing the first model based on the first cross-entropy loss so that the optimized first model is the distilled machine learning model; and
training, using the distilled machine learning model, an additional machine learning model to perform a downstream protein modeling task.
2 . The method of claim 1 , further comprising:
producing a protein design using the trained additional machine learning model; and
synthesizing a protein based on the protein design.
3 . The method of claim 2 , wherein the training of the additional machine learning model comprises:
inputting, separately, protein structures into the additional machine learning model, wherein, in response, the additional machine learning model generates as output a respective predicted protein sequence corresponding to the input protein structures.
4 . The method of claim 1 , wherein the training of the additional machine learning model comprises:
feeding output from the additional machine learning model to the distilled machine learning model, wherein the distilled machine learning model, in response, generates a structure consistency score; and
optimizing the additional machine learning model based on a second cross entropy loss and on a loss of the structure consistency score, wherein the second cross entropy loss is based on a ground truth value of the output of the additional machine learning model.
5 . The method of claim 4 , wherein the output of the additional machine learning model comprises a protein design.
6 . The method of claim 4 , wherein the downstream protein modeling task comprises an inverse protein folding task and the structure consistency score regularizes the inverse protein folding task.
7 . The method of claim 4 , wherein the downstream protein modeling task comprises a protein infilling task and the structure consistency score regularizes the protein infilling task.
8 . The method of claim 1 , wherein the folding protein model:
infers protein structure based on a protein sequence that is input,
comprises a deep learning model and an attention network,
was trained from a public repository of protein sequences and structures, and
is larger than the distilled machine learning model.
9 . The method of claim 1 , wherein the producing a distilled machine learning model further comprises:
generating a set of one or more metrics via a ground truth three-dimensional structure; and
discretizing the set of the one or more metrics into the classes such that the cross-entropy loss is further based on the one or more metrics.
10 . A computer program product, comprising:
one or more tangible computer-readable storage media and program instructions stored on at least one of the one or more tangible computer-readable storage media, the program instructions executable by a processor to cause the processor to perform operations comprising:
producing a distilled machine learning model via:
initializing a first model with initial weights;
inputting an input protein sequence into both the first model and a folding protein model, wherein the inputting to the first model generates logits, and wherein the inputting to the folding protein model generates one or more predictive metrics;
discretizing the one or more predictive metrics into classes;
computing a first cross-entropy loss based on the logits and the classes; and
optimizing the first model based on the first cross-entropy loss so that the optimized first model is the distilled machine learning model; and
training, using the distilled machine learning model, an additional machine learning model to perform a downstream protein modeling task.
11 . The computer program product of claim 10 , the instructions further comprising:
producing a protein design using the trained additional machine learning model; and
synthesizing a protein based on the protein design.
12 . A system comprising:
a memory; and
at least one processor, coupled to said memory, and operative to perform operations comprising:
producing a distilled machine learning model via:
initializing a first model with initial weights;
inputting an input protein sequence into both the first model and a folding protein model, wherein the inputting to the first model generates logits, and wherein the inputting to the folding protein model generates one or more predictive metrics;
discretizing the one or more predictive metrics into classes;
computing a first cross-entropy loss based on the logits and the classes; and
optimizing the first model based on the first cross-entropy loss so that the optimized first model is the distilled machine learning model; and
training, using the distilled machine learning model, an additional machine learning model to perform a downstream protein modeling task.
13 . The system of claim 12 , the operations further comprising:
producing a protein design using the trained additional machine learning model; and
synthesizing a protein based on the protein design.
14 . The system of claim 13 , wherein the training of the additional machine learning model comprises:
inputting, separately, protein structures into the additional machine learning model, wherein, in response, the additional machine learning model generates as output a respective predicted protein sequence corresponding to the input protein structures.
15 . The system of claim 12 , wherein the training of the additional machine learning model comprises:
feeding output from the additional machine learning model to the distilled machine learning model, wherein the distilled machine learning model, in response, generates a structure consistency score; and
optimizing the additional machine learning model based on a second cross entropy loss and on a loss of the structure consistency score, wherein the second cross entropy loss is based on a ground truth value of the output of the additional machine learning model.
16 . The system of claim 15 , wherein the output of the additional machine learning model comprises a protein design.
17 . The system of claim 15 , wherein the downstream protein modeling task comprises an inverse protein folding task and the structure consistency score regularizes the inverse protein folding task.
18 . The system of claim 15 , wherein the downstream protein modeling task comprises a protein infilling task and the structure consistency score regularizes the protein infilling task.
19 . The system of claim 12 , wherein the folding protein model:
infers protein structure based on a protein sequence that is input,
comprises a deep learning model and an attention network,
was trained from a public repository of protein sequences and structures, and
is larger than the distilled machine learning model.
20 . The system of claim 12 , wherein the producing a distilled machine learning model further comprises:
generating a set of one or more metrics via a ground truth three-dimensional structure; and
discretizing the set of the one or more metrics into the classes such that the cross-entropy loss is further based on the one or more metrics.