Predicting protein structures using auxiliary folding networks

A system comprising:

• one or more computers;
• and
• one or more storage devices communicatively coupled to the one or more computers, wherein the one or more storage devices store instructions that, when executed by the one or more computers, cause the one or more computers to perform operations for training a structure prediction neural network, wherein the structure prediction neural network comprises (i) an embedding neural network having a plurality of embedding parameters and that is configured to receive a network input characterizing a protein and to process the network input in accordance with the embedding parameters to generate an embedding output for the network input and (ii) a main folding neural network having a plurality of main folding parameters and that is configured to receive the embedding output and to process the embedding output in accordance with the main folding parameters to generate a main structure prediction that defines a predicted structure of the protein, the operations comprising:
• obtaining a training network input characterizing a training protein and data specifying a target protein structure for the training protein;
• processing the training network input using the embedding neural network and in accordance with current values of the embedding parameters to generate a training embedding output for the training network input;
• processing the training embedding output using the main folding neural network and in accordance with current values of the main folding parameters to generate a main structure prediction that defines a main predicted structure of the training protein;
• for each auxiliary folding neural network in a set of one or more auxiliary folding neural networks that each have a respective plurality of auxiliary folding parameters, processing at least a corresponding intermediate output of the embedding neural network using the auxiliary folding neural network and in accordance with current values of the respective auxiliary folding parameters of the auxiliary folding neural network to generate an auxiliary structure prediction that defines an auxiliary predicted structure of the training protein;
• determining a gradient of an objective function that includes:
• a main structure loss term that characterizes a similarity between:
• (i) the main predicted structure defined by the main structure prediction, and (ii) the target protein structure for the training protein;
• and
• a respective auxiliary structure loss term for each of the auxiliary folding neural networks that characterizes a similarity between:
• (i) the auxiliary predicted structure defined by the auxiliary structure prediction generated by the auxiliary structure prediction neural network, and (ii) the target protein structure for the training protein;
• and
• updating the current values of the embedding network parameters, the main folding, parameters, and the respective auxiliary folding parameters of the one or more auxiliary folding neural networks based on the gradient.

One or more non-transitory computer storage media storing instructions that when executed by one or more computers cause the one or more computers to perform operations for training a structure prediction neural network, wherein the structure prediction neural network comprises (i) an embedding neural network having a plurality of embedding parameters and that is configured to receive a network input characterizing a protein and to process the network input in accordance with the embedding parameters to generate an embedding output for the network input and (ii) a main folding neural network having a plurality of main folding parameters and that is configured to receive the embedding output and to process the embedding output in accordance with the main folding parameters to generate a main structure prediction that defines a predicted structure of the protein, the operations comprising:

• obtaining a training network input characterizing a training protein and data specifying a target protein structure for the training protein;
• processing the training network input using the embedding neural network and in accordance with current values of the embedding parameters to generate a training embedding output for the training network input;
• processing the training embedding output using the main folding neural network and in accordance with current values of the main folding parameters to generate a main structure prediction that defines a main predicted structure of the training protein;
• for each auxiliary folding neural network in a set of one or more auxiliary folding neural networks that each have a respective plurality of auxiliary folding parameters, processing at least a corresponding intermediate output of the embedding neural network using the auxiliary folding neural network and in accordance with current values of the respective auxiliary folding parameters of the auxiliary folding neural network to generate an auxiliary structure prediction that defines an auxiliary predicted structure of the training protein;
• determining a gradient of an objective function that includes:
• a main structure loss term that characterizes a similarity between:
• (i) the main predicted structure defined by the main structure prediction, and (ii) the target protein structure for the training protein;
• and
• a respective auxiliary structure loss term for each of the auxiliary folding neural networks that characterizes a similarity between:
• (i) the auxiliary predicted structure defined by the auxiliary structure prediction generated by the auxiliary structure prediction neural network, and (ii) the target protein structure for the training protein;
• and
• updating the current values of the embedding network parameters, the main folding parameters, and the respective auxiliary folding parameters of the one or more auxiliary folding neural networks based on the gradient.