Using Deep Learning Sequence Models to Identify SARS-CoV-2 Divergence

The dataset for this problem is protein sequences that encode the spike protein of SARS-CoV-2, and with metadata describing the mutations, clade, and Pango lineage of each sequence, both obtained from GISAID [17, 18]. Using protein sequences provides the advantage of dismissing the insignificant effect of silent mutations at the DNA level, which our model would otherwise have to learn if we used raw genomic data. The encoded protein sequences are in FASTA format, and each entry has a title line with an EPI-ISL accession identifier. The encoded metadata is in TSV format, and each entry also has an EPI-ISL accession identifier used to match metadata entries with FASTA entries. Most protein sequences have about 1274 amino acids, with some minor differences in length due to insertion or deletion mutations. More details on this dataset, including example entries and the amino, clade, and Pango lineage mappings, are included in Appendix I.

Input format: FASTA (see appendix for more detailed example)

    >title line
    MXWXXX......BCG*
    (About 1274 amino acid sequence from spike protein,
    asterisk indicates ending of a sequence, X indicates unknown amino acid.)

We begin by reading the FASTA protein sequences and TSV metadata, parsing them into Python dictionaries indexed by EPI-ISL ID. Then, we iterate through the protein sequence and metadata dictionaries, building a list of EPI-ISL IDs that match between both files. Once the correspondence is made, we create a mapping from amino acids and clade labels to integers. We also use these mappings to convert each of our matched data points to a numerical representation. Finally, we set aside every 10th entry as validation data and keep the remaining entries as training data. Like clades, we perform a similar process for Pango lineage, whose label has 1271 classes compared to the clade’s 9 classes. The result is a set of training and validation instances whose pairs consist of tokenized-numeric versions of the input protein sequences that map to clade class labels.

One of the issues with our dataset above is that certain output classes dominate the dataset. Shown below is a histogram of clade and Pango lineage label frequencies. In both the training and validation data, clades 1, 2, 3, 4, and 0 appeared much more frequently than clades 5, 6, 7, and 8, with clade 1 appearing over 30x more often than clade 5 in both datasets. In addition, out of the 1271 Pango lineages in the dataset, the five most common lineages account for over 45% of the entire dataset, and many of the least common lineages occur fewer than ten times total across both datasets. For this reason, we use clade labels instead of Pango lineage labels as our classification targets for this task.

To further address this issue, we balanced the data by downsampling to ensure similar item counts for each output class. Specifically, we use RandomUnderSampler from the imbalanced-learn Python package¹¹1imbalanced-learn.org/stable/references/generated/imblearn.under_sampling.RandomUnderSampler.html. First, we determine how many times the least common output class occurred, then we select a subset of the data for all other output classes so that each class appears equally often.

This balancing step helped resolve two issues. First, it improved our training time by removing redundant data entries. Our original dataset was very large, taking over 2 hours to train per epoch on hardware provided by Google Colab, but training on the reduced dataset was over 14x faster. In addition, as some output classes dominated others in the original dataset, our model was just learning to output the most common classes, instead of learning a meaningful representation of the input data. Balancing the dataset prevents the model from achieving good performance simply by memorizing and outputting the most common class.

Given a set of inputs $\{\mathbf{x}_{i}\}_{i=1}^{N},\mathbf{x}_{i}\in\mathbb{R}^{M}$ and a set of corresponding labels $\{y_{i}\}_{i=1}^{N}$, where $M$ is the maximum length of a spike protein sequence and $N$ is the size of the dataset, the softmax layer outputs the predicted probability for the $j$-th class for the sample vector $\mathbf{x}_{i}$ and weight vector $\mathbf{w}$ as

$$\mathbb{P}(y_{i}=j\mid\mathbf{x}_{i})=\frac{e^{\mathbf{x}_{i}^{T}\mathbf{w}_{j}}}{\sum_{k=1}^{N}e^{\mathbf{x}_{i}^{T}\mathbf{w}_{k}}}$$

We output the class with the highest predicted probability. Denote $\hat{\mathbf{y}}$ as the output of the network, then the cross-entropy loss is

$$\ell(\mathbf{w})=-\sum_{i=1}^{N}\sum_{k=1}^{9}y_{k}^{(i)}\log\left(\hat{y}_{k}^{(i)}\right)$$

We evaluate model performance using cross-entropy loss and prediction accuracy. Given a set of predicted labels $\hat{\textbf{y}}=\{\hat{y_{i}}\}_{i=1}^{N}$ and ground truth labels $\textbf{y}=\{y_{i}\}_{i=1}^{N}$, the prediction accuracy is

$$Acc(\hat{\textbf{y}},\textbf{y})=\frac{\sum_{i=1}^{N}\textbf{1}_{\{y_{i}=\hat{y}_{i}\}}}{N}$$

All of the runs described below were implemented using PyTorch. We used the Adam optimizer with learning rate 0.002 and weight decay 0.000005, and the ReduceLROnPlateau learning rate scheduler with factor 0.5 and patience 2. Our batch size was set to be as large as possible without encountering out-of-memory errors, typically we used batch size 128.

This project is particularly challenging because there are no reference models to date in the literature that specifically model SARS-CoV-2 protein sequences. We postulate that the complexity of these models are no more complex than language modeling, whose hw3p2 performance motivates Specification A. There is a trade-off between having complex models that are difficult to iterate on versus smaller models that iterate quickly. In the coming sections, each of these specifications is expounded and their results are explained.

Inspired by the findings in [2], we use an LSTM-based architecture [7] as our baseline model for this problem, with convolutional layers as a feature extractor for sequential data. Convolutional layers are able to extract features by applying a non-linearity on an affine combination of inputs from each time window for each feature [19]. Furthermore, LSTMs are well-suited for processing sequential data, as LSTM cells are able to remember values across arbitrary intervals of the input sequence. In theory, this makes them well-suited for classifying protein sequences, as the presence of different combinations of mutations may have complex effects on how sequences get classified. For example, Lineage B.1.671 is identified by a combination of at least 6 signature mutations, 3 of which occur commonly in other globally circulating lineages [3]. A successful classifier would need to determine whether this combination of mutations is present to correctly classify these sequences.

Our baseline model consists of one convolutional layer with 128 output channels and a kernel size of 3, with BatchNorm and ReLU applied. Next, we have three bi-LSTM layers with 256 hidden states in each direction, and two fully connected layers using ReLU activation with 512 and 256 neurons respectively, as well as BatchNorm and Dropout=0.2. We use a log softmax layer as our final output. We selected this baseline based on findings in [2], as well as our own experiences in hw3p2.

==========================================================================
                      Kernel Shape       Output Shape     Params Mult-Adds
Layer
0_conv.Conv1d_0        [1, 128, 3]  [ 32,  128, 1274]      384.0  489.216k
1_conv.BatchNorm1d_1         [128]  [ 32,  128, 1274]      256.0     128.0
2_conv.ReLU_2                    -  [ 32,  128, 1274]          -         -
3_lstm                           -  [1274,  32,  512]  3.944448M  3.93216M
4_linear.Dropout_0               -  [1274,  32,  512]          -         -
5_linear.Linear_1       [512, 256]  [1274,  32,  256]   131.328k  131.072k
6_linear.Dropout_2               -  [1274,  32,  256]          -         -
7_linear.Linear_3         [256, 9]  [1274,  32,    9]     2.313k    2.304k
8_linear.LogSoftmax_4            -  [1274,  32,    9]          -         -
--------------------------------------------------------------------------
Total params          4.078729M            Non-trainable params        0.0
Trainable params      4.078729M            Mult-Adds              4.55488M
==========================================================================

Table 1: Model Metadata for Specification A

Although LSTM architectures avoid the issue of vanishing gradients that affects traditional RNN architectures [20], our baseline model remains very difficult to train as evidenced by our gradient flow analysis in Section 5.1. Inspired by [21], we propose a pyramidal LSTM architecture that halves the time resolution of the input at each layer, reducing the length of our data sequence and hopefully making it easier to train our network.

In a typical bi-LSTM network, the output at the $i$-th time step from the $j$-th layer is computed as:

$$h_{i}^{j}=BLSTM(h_{i-1}^{j},h_{i}^{j-1})$$

In a pyramidal bi-LSTM network, we concatenate the outputs at consecutive steps of each layer, as:

$$h_{i}^{j}=pBLSTM(h_{i-1}^{j},[h_{2i}^{j-1},h_{2i+1}^{j-1}])$$

In our model, following the approach of [21], we stack three pyramidal bi-LSTM layers on top of a regular bi-LSTM layer at the bottom. Our model description is as follows:

=============================================================================
                        Kernel Shape       Output Shape    Params   Mult-Adds
Layer
0_embedding                 [32, 25]  [1274,  32,   32]     800.0       800.0
1_conv.Conv1d_0         [32, 128, 3]  [  32, 128, 1274]   12.288k  15.654912M
2_conv.BatchNorm1d_1           [128]  [  32, 128, 1274]     256.0       128.0
3_conv.ReLU_2                      -  [  32, 128, 1274]         -           -
4_lstm                             -  [1274,  32,  256]  264.192k    262.144k
5_pblstm1.LSTM_blstm               -  [ 637,  32,  256]  657.408k     655.36k
6_pblstm2.LSTM_blstm               -  [ 318,  32,  256]  657.408k     655.36k
7_pblstm3.LSTM_blstm               -  [ 159,  32,  256]  657.408k     655.36k
8_linear.Linear_0         [256, 128]  [ 159,  32,  128]   32.896k     32.768k
9_linear.Dropout_1                 -  [ 159,  32,  128]         -           -
10_linear.Linear_2          [128, 9]  [ 159,  32,    9]    1.161k      1.152k
11_linear.LogSoftmax_3             -  [ 159,  32,    9]         -           -
-----------------------------------------------------------------------------
Total params           2.283817M             Non-trainable params         0.0
Trainable params       2.283817M             Mult-Adds             17.917984M
=============================================================================

Table 2: Model Metadata for Specification B

Our results are as follows:

We also plotted a histogram of model output frequencies vs. label frequencies, as well as a confusion matrix describing the classification error of the model:

 

This model seems to have much better performance, and has learned how to identify clades 2-5 with good accuracy, as well as partly how to identify clades 0-1 and 6-8. The model has also learned to group together clades 0-2 and clades 5-8, which reflects that these clades are relatively similar to each other according to the T-SNE embedding analysis discussed in Section 5.2. When the model is uncertain about how to label a particular input, it seems to default to either clade 2 or clade 5, which is consistent with the behavior of previous models. In total, this model achieves better performance and converges much more quickly than any of the standard LSTM-based models. We hypothesize that this is because it shortens the input sequences to a length that is reasonable for LSTM to handle.

Inspired by [22], we also experimented on Bert embedding schemes. Instead of simply encoding each amino acid as an index that has no specific biological meaning, we concatenate the Bert embedding pre-trained on protein sequences [23] with the amino acid hydrophilicity encoding. Hydrophilicity refers to how soluble an amino acid is in water, and it is an important property of amino acids with significant effect on protein function [24]. The resulting embeddings are then trained on an MLP classifier. We were able to see a converging behavior of our model.

===========================================================================
Layer                Kernel Shape  Output Shape  Param #       Mult-Adds
===========================================================================
Sequential: 1-1                 -  [32,    9]            -                -
Linear: 2-1          [2048, 1024]  [32, 1024]    2,098,176       67,108,864
ReLU: 2-2                       -  [32, 1024]            -                -
Linear: 2-3           [1024, 512]  [32,  512]      524,800       16,777,216
ReLU: 2-4                       -  [32,  512]            -                -
Linear: 2-5            [512, 256]  [32,  256]      131,328        4,194,304
ReLU: 2-6                       -  [32,  256]            -                -
Linear: 2-7              [256, 9]  [32,    9]        2,313           73,728
===========================================================================
Total params:          2.756617M            Non-trainable params        0.0
Trainable params:      2.756617M            Total mult-adds (M):     881.54
===========================================================================

Table 3: Model Metadata for Specification C

In addition to the models presented above, we performed various experiments on the baseline model to determine whether each component in our network was necessary. Our results are summarized in the following table, and more detailed descriptions and results are available in Appendix II. In summary, while adding a heuristical embedding layer seemed to improve our baseline performance, this improvement was reversed by the other changes we tried such as removing dropout, removing convolutional layers, and removing LSTM layers:

Architecture	Training Acc	Validation Acc
Specification A: Baseline Model	25.4%	24.9%
Specification D: Heuristical Embedding	35.9%	36.2%
Specification E: Dropout Removed	25.4%	25.3%
Specification F: CNN and Dropout Removed	25.4%	25.6%
Specification G: Single LSTM Layer	25.5%	25.7%

To quantify whether our model was learning, we performed gradient flow analysis on our models by plotting the average and maximum gradient for each set of parameters in the network. The gradient flows are averaged over the first five epochs of training on the balanced dataset. We could observe that the input layers in the baseline LSTM are not being updated, while the input layers in the PLSTMs are having small updates each batch.

The MLP model also suffers from vanishing gradients to some extent, since the average update on the weights in the first few linear layers is on the order of $10^{-4}$. This indicates that we may see a performance increase by adding more regularization techniques to the MLP model.

To visualize whether our raw numerical sequence data and Bert embedding data are distinguishable from each other, we used T-SNE²²2https://scikit-learn.org/stable/modules/generated/sklearn.manifold.TSNE.html from the scikit-learn package to plot the embeddings. More specifically, 10% of the data is sub-sampled from the balanced raw and embedding training datasets, with an equal amount of sequence per output clade label. Raw sequences are padded to the same length before performing T-SNE. The results are shown below:

While there are clades that form distinct clusters such as clades 3, 4 and 8 for both the raw and embedding data, most of the other clades’ data points are rather spread out such as clades 0 and 5. The clades that do form distinct clusters are also the labels that our neural network model described above prefer to output. This visualization also explains why many of our models had difficulty differentiating among clades 0-2 and and clades 6-8 as these clusters overlap on the T-SNE plots. Therefore, some of the clades are more similar to each other than others, and our naive numerical representation of amino acid sequence data might not be a good fit for this task. Using Bert embedding with hydrophillicity encoding does help to group data of the same clade together, but between group distances are reduced at the same time.

Because we were not able to find many existing papers that model our problem of classifying clades from protein sequencing data using neural networks, we lacked reference in how we should pre-process our data to account for all the nuances in biological data, as well as how to structure our networks. More specifically, our dataset is different from the frame-level phoneme classification task we were accustomed to because of the unique nature of protein folding - amino acids that are far downstream could have significant impacts on amino acids elsewhere (e.g. form disulfide bonds and other secondary or tertiary structures). Therefore, our LSTM model should not only learn to place an emphasis on nearby sequences, but also on amino acids much farther away which may have crucial impact on if a structural change in one site would lead to a clade change.

Because we had limited computational resources and time for this project, we were not able to refine our model with much more advanced LSTM models or dedicate time finding better encoding methods since we spent much time refining and expanding our dataset in the early phase and had to re-run all our models each time we updated our dataset. For example, when we originally tried to run our baseline BiLSTM model on the full, unbalanced dataset containing over 700k sequences, it took us over two hours per epoch using GPU provided by Google Colab Pro. We also experimented with adding convolutional and fully connected layers to our baseline model. The maximum hidden size without triggering an OOM error is 256 for a second convolutional layer and 512 for a fully connected layers, which did not show a significant improvement over our previous result. Therefore, monetary and time cost limited us from experimenting more model structures.