Machine learning model distillation for protein design

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.

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.