Protein structure prediction from amino acid sequences using self-attention neural networks

1. A method performed by one or more computers for determining a predicted three-dimensional structure of a protein that is specified by an amino acid sequence, the method comprising:

obtaining a multiple sequence alignment for the protein, wherein multiple sequence alignment defines a sequence alignment of the amino acid sequence of the protein with a respective amino acid sequence of each of a plurality of homologous proteins;

determining, from the multiple sequence alignment and for each pair of amino acids in the amino acid sequence of the protein, a respective initial embedding of the pair of amino acids;

processing the initial embeddings of the pairs of amino acids using a pair embedding neural network comprising a plurality of self-attention neural network layers to generate a final embedding of each pair of amino acids; and

determining the predicted three-dimensional structure of the protein based on the final embedding of each pair of amino acids.

2. The method of claim 1 , wherein each self-attention neural network layer of the pair embedding neural network is configured to perform operations comprising:

receiving a current embedding of each pair of amino acids; and

updating the current embedding of each pair of amino acids using attention over the current embeddings of the pairs of amino acids.

3. The method of claim 2 , wherein for each pair of amino acids, updating the current embedding of the pair of amino acids using attention over the current embeddings of the pairs of amino acids comprises:

updating the current embedding of the pair of amino acids using attention over only a proper subset of the current embeddings of the pairs of amino acids.

4. The method of claim 3 , wherein:

one or more of the self-attention neural network layers are row-wise self-attention neural network layers; and

for each row-wise self-attention neural network layer, updating the current embedding of the pair of amino acids using attention over only a proper subset of the current embeddings of the pairs of amino acids comprises:

updating the current embedding of the pair of amino acids using attention over only current embeddings of pairs of amino acids that are located in a same row as the current embedding of the pair of amino acids in an arrangement of the current embeddings of the pairs of amino acids into a two-dimensional array.

5. The method of claim 4 , wherein:

one or more of the self-attention neural network layers are column-wise self-attention neural network layers; and

for each column-wise self-attention neural network layer, updating the current embedding of the pair of amino acids using attention over only a proper subset of the current embeddings of the pairs of amino acids comprises:

updating the current embedding of the pair of amino acids using attention over only current embeddings of pairs of amino acids that are located in a same column as the current embedding of the pair of amino acids in the arrangement of the current embeddings of the pairs of amino acids into the two-dimensional array.

6. The method of claim 5 , wherein the plurality of self-attention neural network layers of the pair embedding neural network comprise an alternating sequence of row-wise self-attention neural network layers and column-wise self-attention neural network layers.

7. The method of claim 1 , wherein determining the predicted structure of the protein based on the final embedding of each pair of amino acids comprises:

determining a respective initial embedding of each amino acid in the amino acid sequence of the protein based on the final embeddings of the pairs of amino acids; and

determining the predicted structure of the protein based on the initial embedding of each amino acid in the amino acid sequence.

8. The method of claim 1 , wherein determining, from the multiple sequence alignment and for each pair of amino acids in the amino acid sequence of the protein, a respective initial embedding of the pair of amino acids comprises:

partitioning the multiple sequence alignment into: (i) a set of cluster amino acid sequences, and (ii) a set of extra amino acid sequences;

generating: (i) an embedding of the set of cluster amino acid sequences, and (ii) an embedding of the set of extra amino acid sequences;

processing a network input comprising: (i) the embedding of the cluster amino acid sequences, and (ii) the embedding of the extra amino acid sequences, using a cross-attention neural network to update the embedding of the cluster amino acid sequences; and

determining the initial embeddings of the pairs of amino acids based on the updated embedding of the cluster amino acid sequences.

9. The method of claim 8 , wherein processing a network input comprising: (i) the embedding of the cluster amino acid sequences, and (ii) the embedding of the extra amino acid sequences, using a cross-attention neural network to update the embedding of the cluster amino acid sequences comprises repeatedly performing operations including:

updating the embedding of the cluster amino acid sequences using attention over the embedding of the extra amino acid sequences; and

updating the embedding of the extra amino acid sequences using attention over the embedding of the cluster amino acid sequences.

10. 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 determining a predicted three-dimensional structure of a protein that is specified by an amino acid sequence, the operations comprising:

obtaining a multiple sequence alignment for the protein, wherein multiple sequence alignment defines a sequence alignment of the amino acid sequence of the protein with a respective amino acid sequence of each of a plurality of homologous proteins;

determining, from the multiple sequence alignment and for each pair of amino acids in the amino acid sequence of the protein, a respective initial embedding of the pair of amino acids;

processing the initial embeddings of the pairs of amino acids using a pair embedding neural network comprising a plurality of self-attention neural network layers to generate a final embedding of each pair of amino acids; and

determining the predicted three-dimensional structure of the protein based on the final embedding of each pair of amino acids.

11. The system of claim 10 , wherein each self-attention neural network layer of the pair embedding neural network is configured to perform operations comprising:

receiving a current embedding of each pair of amino acids; and

updating the current embedding of each pair of amino acids using attention over the current embeddings of the pairs of amino acids.

12. The system of claim 11 , wherein for each pair of amino acids, updating the current embedding of the pair of amino acids using attention over the current embeddings of the pairs of amino acids comprises:

updating the current embedding of the pair of amino acids using attention over only a proper subset of the current embeddings of the pairs of amino acids.

13. The system of claim 12 , wherein:

one or more of the self-attention neural network layers are row-wise self-attention neural network layers; and

for each row-wise self-attention neural network layer, updating the current embedding of the pair of amino acids using attention over only a proper subset of the current embeddings of the pairs of amino acids comprises:

updating the current embedding of the pair of amino acids using attention over only current embeddings of pairs of amino acids that are located in a same row as the current embedding of the pair of amino acids in an arrangement of the current embeddings of the pairs of amino acids into a two-dimensional array.

14. The system of claim 13 , wherein:

one or more of the self-attention neural network layers are column-wise self-attention neural network layers; and

for each column-wise self-attention neural network layer, updating the current embedding of the pair of amino acids using attention over only a proper subset of the current embeddings of the pairs of amino acids comprises:

updating the current embedding of the pair of amino acids using attention over only current embeddings of pairs of amino acids that are located in a same column as the current embedding of the pair of amino acids in the arrangement of the current embeddings of the pairs of amino acids into the two-dimensional array.

15. The system of claim 14 , wherein the plurality of self-attention neural network layers of the pair embedding neural network comprise an alternating sequence of row-wise self-attention neural network layers and column-wise self-attention neural network layers.

16. 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 determining a predicted three-dimensional structure of a protein that is specified by an amino acid sequence, the operations comprising:

obtaining a multiple sequence alignment for the protein, wherein multiple sequence alignment defines a sequence alignment of the amino acid sequence of the protein with a respective amino acid sequence of each of a plurality of homologous proteins;

determining, from the multiple sequence alignment and for each pair of amino acids in the amino acid sequence of the protein, a respective initial embedding of the pair of amino acids;

processing the initial embeddings of the pairs of amino acids using a pair embedding neural network comprising a plurality of self-attention neural network layers to generate a final embedding of each pair of amino acids; and

determining the predicted three-dimensional structure of the protein based on the final embedding of each pair of amino acids.

17. The non-transitory computer storage media of claim 16 , wherein each self-attention neural network layer of the pair embedding neural network is configured to perform operations comprising:

receiving a current embedding of each pair of amino acids; and

updating the current embedding of each pair of amino acids using attention over the current embeddings of the pairs of amino acids.

18. The non-transitory computer storage media of claim 17 , wherein for each pair of amino acids, updating the current embedding of the pair of amino acids using attention over the current embeddings of the pairs of amino acids comprises:

updating the current embedding of the pair of amino acids using attention over only a proper subset of the current embeddings of the pairs of amino acids.

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

one or more of the self-attention neural network layers are row-wise self-attention neural network layers; and

for each row-wise self-attention neural network layer, updating the current embedding of the pair of amino acids using attention over only a proper subset of the current embeddings of the pairs of amino acids comprises:

updating the current embedding of the pair of amino acids using attention over only current embeddings of pairs of amino acids that are located in a same row as the current embedding of the pair of amino acids in an arrangement of the current embeddings of the pairs of amino acids into a two-dimensional array.

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

one or more of the self-attention neural network layers are column-wise self-attention neural network layers; and

for each column-wise self-attention neural network layer, updating the current embedding of the pair of amino acids using attention over only a proper subset of the current embeddings of the pairs of amino acids comprises:

updating the current embedding of the pair of amino acids using attention over only current embeddings of pairs of amino acids that are located in a same column as the current embedding of the pair of amino acids in the arrangement of the current embeddings of the pairs of amino acids into the two-dimensional array.