Predicting protein structures using auxiliary folding networks

1 . A method of 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 method 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.

2 . The method of claim 1 , wherein each auxiliary folding neural network has a same neural network architecture as the main folding neural network.

3 . The method of claim 2 , wherein the set of one or more auxiliary folding neural networks comprises a plurality of auxiliary folding neural networks and wherein the objective function constrains the auxiliary folding neural networks to share parameter values.

4 . The method of claim 2 , wherein the objective function constrains the auxiliary folding neural networks and the main folding neural network to share parameter values.

5 . The method of claim 1 , wherein:

the network input comprises a respective initial pair embedding for each pair of amino acids in the protein,

the embedding neural network comprises a sequence of update blocks, wherein each update block performs operations comprising:

receiving a block input comprising a respective current pair embedding for each pair of amino acids in the protein; and

updating the respective current pair embeddings for each pair of amino acids in the protein to generate a respective updated pair embedding for each pair of amino acids in the protein, and

the embedding output comprises at least the updated pair embeddings generated by a last update block in the sequence.

6 . The method of claim 5 , wherein each auxiliary folding neural network corresponds to a different one of the update blocks in the sequence that is not the last update block in the sequence, and wherein each auxiliary folding neural network is configured to receive as input at least the updated pair embeddings generated by the corresponding update block.

7 . The method of claim 6 , wherein:

the network input further comprises an initial multiple sequence alignment (MSA) embedding that represents a respective multiple sequence alignment corresponding to each chain in the protein,

the block input to each of the update blocks further comprises a current MSA embedding; and

the operations performed by each update block further comprise:

updating the current MSA embedding to generate an updated MSA embedding.

8 . The method of claim 7 , wherein the input for each auxiliary folding neural network further comprises the updated MSA embedding generated by the corresponding update block.

9 . The method of claim 7 , wherein each auxiliary folding neural network is configured to:

generate, from the auxiliary structure prediction generated by the auxiliary folding neural network, a transformed structure prediction that has a same dimensionality as the updated pair embeddings;

combine the transformed structure prediction with the updated pair embeddings to generate further updated pair embeddings; and

provide the further updated pair embeddings as input to an update block that follows the current update block in the sequence.

10 . The method of claim 9 , wherein:

the auxiliary folding prediction comprises structure parameters that specify, for each amino acid, a predicted 3-D spatial location of a specified atom in the amino acid in the structure of the protein; and

generating the transformed structure prediction comprises:

generating, from the predicted 3-D spatial locations for the amino acids specified by the structure parameters, a distance map that characterizes, for each pair of amino acids in the protein, a respective estimated distance between the pair of amino acids in the structure of the protein; and

generating, from the distance map, a transformed distance map that has a same dimensionality as the updated pair embeddings.

11 . The method of claim 9 , wherein:

the auxiliary folding prediction comprise structure parameters that specify a distance map that characterizes, for each pair of amino acids in the protein, a respective estimated distance between the pair of amino acids in the structure of the protein; and

generating the transformed structure prediction comprises:

generating, from the distance map specified by the structure parameters, a transformed distance map that has a same dimensionality as the initial pair embeddings.

12 . The method of claim 1 , further comprising:

after the training, obtaining a new network input characterizing a new protein; and

processing the new network input using the trained structure prediction neural network to generate a new main structure prediction that defines a predicted structure of the new protein.

13 . 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.

14 . 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.

15 . The non-transitory computer storage media of claim 14 , wherein each auxiliary folding neural network has a same neural network architecture as the main folding neural network.

16 . The non-transitory computer storage media of claim 15 , wherein the set of one or more auxiliary folding neural networks comprises a plurality of auxiliary folding neural networks and wherein the objective function constrains the auxiliary folding neural networks to share parameter values.

17 . The non-transitory computer storage media of claim 15 , wherein the objective function constrains the auxiliary folding neural networks and the main folding neural network to share parameter values.

18 . The non-transitory computer storage media of claim 14 , wherein:

the network input comprises a respective initial pair embedding for each pair of amino acids in the protein,

the embedding neural network comprises a sequence of update blocks, wherein each update block performs operations comprising:

receiving a block input comprising a respective current pair embedding for each pair of amino acids in the protein; and

updating the respective current pair embeddings for each pair of amino acids in the protein to generate a respective updated pair embedding for each pair of amino acids in the protein, and

the embedding output comprises at least the updated pair embeddings generated by a last update block in the sequence.

19 . The non-transitory computer storage media of claim 18 , wherein each auxiliary folding neural network corresponds to a different one of the update blocks in the sequence that is not the last update block in the sequence, and wherein each auxiliary folding neural network is configured to receive as input at least the updated pair embeddings generated by the corresponding update block.

20 . The non-transitory computer storage media of claim 19 , wherein:

the network input further comprises an initial multiple sequence alignment (MSA) embedding that represents a respective multiple sequence alignment corresponding to each chain in the protein,

the block input to each of the update blocks further comprises a current MSA embedding; and

the operations performed by each update block further comprise:

updating the current MSA embedding to generate an updated MSA embedding.