Overview#
One of the reasons I am interested to improve medicine using deep learning is AlphaFold. It has pushed the frontiers of protein biology by beating the CASP and open-sourcing its structure predictions of millions of proteins. So, I wanted implement it myself and understand the novelties it introduced.
Feature Extraction#
The first step in any machine learning system is feature extraction. This is where domain knowledge is required to convert the raw data into tensors. AlphaFold2 uses the evolutionary information of amino acid sequences. For every amino acid sequence, it performs Multiple Sequence Alignment using JackHMMER and HHblits against protein databanks. Computing the MSAs, however, requires huge memory and we will be sticking with precomputed alignments instead.
The aligned sequences will contain insertions, deletions and substituitions. These are denoted as follows:
>Query
MKTAYVGDKLSP
>Seq1 Insertion
MKTAYVglyGDKLSP
>Seq2 Deletion
MK--YVGDKLSP
>Seq3 Substitution
MKTAYVGEKLSP
For consistency, it makes the aligned seqences to be of same length. To do so, it keeps track the number of insertions to the left of every amino acid and removes the insertions from the sequences, as shown below:
>Before
000000 300000000
MKTAYVglyGDKLSP---
>After
000000300000000
MKTAYVGDKLSP---
With the length being consistent, it proceeds with one-hot encoding the aligned sequences to vectors of length 22 (20 amino acids, 1 gap, 1 unknown), and the query sequence as a 21-class (20 amino acids, 1 unknown) vectors.
To reduce the computational costs, AlphaFold selects (N_{\text{clust}}) random sequences from the aligned as cluster centers. The remaining $N_{extra_seq}$ sequences are assigned to the clusters, after applying a mask (induce random mutations) to cluster centers for better robustness. The details of this procedure are at section 1.2.7 of the AlphaFold2 supplement paper.
After clustering, amino acid distributions and deletion statistics of the clusters are stacked together with one-hot representations of cluster centers to form msa_feat of shape [(N_{clust}), (N_{res}), 49]. The extra sequences form extra_msa_feat of shape [(N_{extra_seq}), (N_{res}), 25].
It also creates a one-hot representation for the query sequence, target_feat of shape [(N_{res}), 21] and a residue_index of shape [(N_{res})] for positional encodings. More details are at section 1.2.9 of the supplement paper.