Geometric attention for biological language reasoning
A sequence state including representations of neighboring amino acids included in a local physical structure for a specific amino acid is received, wherein the neighboring amino acids at least include a first neighboring amino acid and a second neighboring amino acid. A direction query vector and a direction key vector are determined including by applying a first directional rotation transformation to at least a portion of a representation of the first neighboring amino acid included in the representations and applying a second directional rotation transformation to at least a portion of a representation of the second neighboring amino acid included in the representations. A direction attention result is determined including by evaluating elements of the direction query vector and the direction key vector. At least the direction attention result is used to update the sequence state for an attention mechanism of a machine learning model.
1 . A method, comprising:
receiving, by one or more processors, a protein query including backbone coordinate data representing three-dimensional local structure of a protein, and generating a sequence state including representations of neighboring amino acids included in a local physical structure for a specific amino acid of the protein, wherein the neighboring amino acids at least include a first neighboring amino acid and a second neighboring amino acid, and the neighboring amino acids are specifically selected for the specific amino acid based on being physically closest in to the specific amino acid in a three-dimensional space based on the backbone coordinate data;
determining, using a geometric attention block of a biological protein language model trained to predict elements of protein sequence, structure, and/or function, comprising a plurality of transformer blocks, the geometric attention block comprising a direction attention block and a distance attention block:
determining, by the direction attention block, a direction query vector and a direction key vector by applying a first directional rotation transformation derived from a local backbone reference frame associated with the first neighboring amino acid to at least a portion of a representation of the first neighboring amino acid included in the representations and applying a second directional rotation transformation derived from a local backbone reference frame associated with the second neighboring amino acid to at least a portion of a representation of the second neighboring amino acid included in the representations;
determining, by the direction attention block, a direction attention result including by performing a dot product operation on the direction query vector and the direction key vector; and
determining, by the distance attention block, a distance attention result by performing a Euclidean norm operation based on representations of the first neighboring amino acid and the second neighboring amino acid that have been transformed into a common reference frame;
using at least the direction attention result and the distance attention result to update the sequence state within the geometric attention block of the biological protein language model to generate an updated sequence state; and
predicting, by the biological protein language model using the updated sequence state a structural property of the protein, wherein the structure property is represented as a structure token corresponding to the specific amino acid.
2 . The method of claim 1 , further comprising determining a weighted direction attention result including by applying a learned direction weight term to the determined direction attention result.
3 . The method of claim 1 , further comprising:
determining a distance query vector and a distance key vector including by applying a first distance rotation transformation and a first distance translation transformation to at least the portion of the representation of the first neighboring amino acid included in the representations and applying a second distance rotation transformation and a second distance translation transformation to at least the portion of the representation of the second neighboring amino acid included in the representations;
determining a distance attention result including by evaluating the elements of the distance query vector and the distance key vector; and
using at least the distance attention result to update the sequence state for an attention mechanism of the biological protein language model.
4 . The method of claim 3 , wherein the first direction rotation transformation and the first distance rotation transformation are the same; and wherein the second direction rotation transformation and the second distance rotation transformation are the same.
5 . The method of claim 3 , further comprising determining a weighted distance attention result including by applying a learned distance weight term to the determined direction attention result.
6 . The method of claim 3 , further comprising:
determining a weighted attention result including by subtracting a weighted distance attention result from a weighted direction attention result.
7 . The method of claim 6 , further comprising:
determining a value vector including by applying a first value rotation transformation to at least the portion of the representation of the first neighboring amino acid included in the representations and applying a second value rotation transformation to at least the portion of the representation of the second neighboring amino acid included in the representations;
determining a geometric attention result including by applying a matrix multiplication operation using the value vector and the weighted attention result as operands;
applying inverse rotation transformations to the determined geometric attention result; and
using at least the determined geometric attention result to update the sequence state for an attention mechanism of the biological protein language model.
8 . The method of claim 7 , wherein at least a subset of the applied inverse rotation transformations has an inverse relationship with the first value rotation transformation and the second value rotation transformation.
9 . A system, comprising:
one or more processors configured to:
receive a protein query including backbone coordinate data representing three-dimensional local structure of a protein, and generate a sequence state including representations of neighboring amino acids included in a local physical structure for a specific amino acid of the protein, wherein the neighboring amino acids at least include a first neighboring amino acid and a second neighboring amino acid, and the neighboring amino acids are specifically selected for the specific amino acid based on being physically closest in to the specific amino acid in a three-dimensional space based on the backbone coordinate data;
determine, using a geometric attention block of a biological protein language model trained to predict elements of protein sequence, structure, and/or function, comprising a plurality of transformer blocks, the geometric attention block comprising a direction attention block and a distance attention block:
determine, by the direction attention block, a direction query vector and a direction key vector including by being configured to apply a first directional rotation transformation derived from a local backbone reference frame associated with the first neighboring amino acid to at least a portion of a representation of the first neighboring amino acid included in the representations and applying a second directional rotation transformation derived from a local backbone reference frame associated with the second neighboring amino acid to at least a portion of a representation of the second neighboring amino acid included in the representations;
determine, by the direction attention block, a direction attention result including by performing a dot product operation on the direction query vector and the direction key vector; and
determine, by the distance attention block, a distance attention result by performing a Euclidean norm operation based on representations of the first neighboring amino acid and the second neighboring amino acid that have been transformed into a common reference frame;
use at least the direction attention result and the distance attention result to update the sequence state within the geometric attention block of the biological protein language model to generate an updated sequence state; and
predict, by the biological protein language model using the updated sequence state a structural property of the protein, wherein the structure property is represented as a structure token corresponding to the specific amino acid; and
a memory coupled to at least one of the one or more processors and configured to provide instructions.
10 . The system of claim 9 , wherein the one or more processors are configured to:
determine a weighted direction attention result including by applying a learned direction weight term to the determined direction attention result.
11 . The system of claim 9 , wherein the one or more processors are configured to:
determine a distance query vector and a distance key vector including by applying a first distance rotation transformation and a first distance translation transformation to at least the portion of the representation of the first neighboring amino acid included in the representations and applying a second distance rotation transformation and a second distance translation transformation to at least the portion of the representation of the second neighboring amino acid included in the representations;
determine a distance attention result including by evaluating the elements of the distance query vector and the distance key vector; and
update the sequence state for an attention mechanism of the biological protein language model using at least the distance attention result.
12 . The system of claim 11 , wherein the first direction rotation transformation and the first distance rotation transformation are the same; and wherein the second direction rotation transformation and the second distance rotation transformation are the same.
13 . The system of claim 11 , wherein the one or more processors are configured to:
determine a weighted distance attention result including by applying a learned distance weight term to the determined direction attention result.
14 . The system of claim 11 , wherein the one or more processors are configured to:
determine a weighted attention result including by being configured to subtract a weighted distance attention result from a weighted direction attention result.
15 . The system of claim 14 , wherein the one or more processors are configured to:
determine a value vector including by being configured to apply a first value rotation transformation to at least the portion of the representation of the first neighboring amino acid included in the representations and applying a second value rotation transformation to at least the portion of the representation of the second neighboring amino acid included in the representations;
determine a geometric attention result including by being configured to apply a matrix multiplication operation using the value vector and the weighted attention result as operands;
apply inverse rotation transformations to the determined geometric attention result; and
update the sequence state for an attention mechanism of the biological protein language model using at least the determined geometric attention result.
16 . A computer program product, the computer program product being embodied in a non-transitory computer readable storage medium and comprising computer instructions for:
receiving, by one or more processors, a protein query including backbone coordinate data representing three-dimensional local structure of a protein, and generating a sequence state including representations of neighboring amino acids included in a local physical structure for a specific amino acid of the protein, wherein the neighboring amino acids at least include a first neighboring amino acid and a second neighboring amino acid, and the neighboring amino acids are specifically selected for the specific amino acid based on being physically closest in to the specific amino acid in a three-dimensional space based on the backbone coordinate data;
determining, using a geometric attention block of a biological protein language model trained to predict elements of protein sequence, structure, and/or function, comprising a plurality of transformer blocks, the geometric attention block comprising a direction attention block and a distance attention block:
determining, by the direction attention block, a direction query vector and a direction key vector by applying a first directional rotation transformation derived from a local backbone reference frame associated with the first neighboring amino acid to at least a portion of a representation of the first neighboring amino acid included in the representations and applying a second directional rotation transformation derived from a local backbone reference frame associated with the second neighboring amino acid to at least a portion of a representation of the second neighboring amino acid included in the representations;
determining, by the direction attention block, a direction attention result including by performing a dot product operation on the direction query vector and the direction key vector; and
determining, by the distance attention block, a distance attention result by performing a Euclidean norm operation based on representations of the first neighboring amino acid and the second neighboring amino acid that have been transformed into a common reference frame;
using at least the direction attention result and the distance attention result to update the sequence state within the geometric attention block of the biological protein language model to generate an updated sequence state; and
predicting, by the biological protein language model using the updated sequence state a structural property of the protein, wherein the structure property is represented as a structure token corresponding to the specific amino acid.