Sequence Embedding using embedding layer: how does the network architecture influence it?

Question

I want to obtain a dense vector representation of protein sequences so that I can meaningfully represent them in an embedding space. We can consider them as sequences of letters, in particular there are 21 unique symbols which are the amino acids (for example: MNTQILVFIACVLIEAKGDKICL).

My approach is to use a sequence embedding that can be learned as a part of a deep learning model (built with Python using Keras libraries), that is a classifier (supervised) neural network which I train to classify sequences according to the host species they belong to. The steps I follow are the following:

1. Tokenization. The only way I can tokenize these sequences of amino acids symbols is to consider single characters as tokens. I also found this example in Kaggle Deep Protein Sequence family Classification in Python that classifies different proteins and uses single amino acids as tokens.
2. Embedding. Stealing words from the answer to the question How the embedding layer is trained in Keras Embedding layer:

An embedding layer is a trainable layer that contains 1 embedding matrix, which is two dimensional, in one axis the number of unique values the categorical input can take (for example 26 in the case of lower case alphabet) and on the other axis the dimensionality of your embedding space. The role of the embedding layer is to map a category into a dense space in a way that is useful for the task at hand, at least in a supervised task. This usually means there is some semantic value in the embedding vectors and categories that are close in this space will be close in meaning for the task.

Moreover, useful words from the blog titled as Neural Network Embeddings Explained:

The main issue with one-hot encoding is that the transformation does not rely on any supervision. We can greatly improve embeddings by learning them using a neural network on a supervised task. The embeddings form the parameters — weights — of the network which are adjusted to minimize loss on the task. The resulting embedded vectors are representations of categories where similar categories — relative to the task — are closer to one another.

Thus putting all these pieces together in my case: I choose as tokens the single amino acids, so I have 21 unique symbols (i.e. number of amino acids) and I choose 10 as dimension of the embedding space, thus my embedding layer dimension is 21 x 10. This means that, once the neural network is trained, I can extract the weights of the embedding layer that are 21 vectors (one for each amino acid) and each vector is a 10 dimensional vector.

As the second extract explains, each element of these vectors are values that are adjusted to minimize the loss in the task. I could see each amino acid as if it was a letter in a word and I wanted to classify these words according to something (like positive or negative comment); or as if it was a word in a sentence and I chose as tokens words and I wanted to classify these sentences according to something (like positive or negative comment).

Since I want sequence representation, I have to find a way to put together the embeddings of single amino acids and the only way that I found feasible is to average all the 10 dimensional vectors so to obtain for each sequence a 10 dimensional vector that is the average of the embeddings of all the amino acids. This would for sure highlight if there are over-represented symbols in one sequence with respect to other. Furthermore, since each symbol is associated to vectors whose values are adjusted to minimize loss on the task, averaging should preserve each single amino acid meaningful embedding (meaningful embedding according to the task) and give "a global meaningful embedding" of the whole sequence that minimizes the loss on the task. This, in fact, seems to work for simple sentences embeddings. See the answer to a my question post: How to obtain vector representation of phrases using the embedding layer and do PCA with it in which each word of sentences were considered as tokens and I was looking for vectors embedding the whole sentences. Similarly, here I can see each single amino acid symbol as word and the whole sequence as sentence.

Hence this method should carry/embed two information: frequencies of letters of sequences and classes to which each sequence belong to (in this case host species).

But I would like it to consider also the global structure of letters positions (not only the relative frequencies), because with this method sequences like: MNTQILVFIACVLIEAKGDKICL (sequence 1) and AKGDKICLMNTQILVFIACVLIE (sequence 2) are represented by the same 10 dimensional vector.

Thus... from here the third point:

3. Choice of neural network architecture. Does the choice of neural network architecture influence the embedding of each single amino acid symbol and, consequently, the embedding of the entire sequence ? For example, if I use a LSTM neural network that should "memorize" the global structure of sequences, would it improve the dense vector representation (both in general and in this particular case) ? I would expect yes since, as reported also in the extracts, this embedding layer (so its weights) is trained as all the others so the backpropagation algorithm updates also these weights. In other words, if LSTM network has the power to recognize the importance of the position of each character according to the task (for example, it is able to learn that: if a letter is in position 1, it means it belongs to human, while instead if it is in position 2, it means it belongs to dog) then weights should be updated also according to this. Differently from a simple deep learning model that is not able to consider this information and to deal with sequences.

I understand that if I average the embeddings of the single amino acids, I would have anyway the problem that sequences like MNTQILVFIACVLIEAKGDKICL sequence 1 and AKGDKICLMNTQILVFIACVLIE sequence 2 would have the same dense vector representation. But does, also in this case, a better choice of the network architecture give a better result in some way ?

I apologize for the long question and I thank you in advance.

Answer 1

Premises: mine is not gonna be an exhaustive answer, also I'm more familiar with classic natural language processing than with embedding vectors applied to protein sequences. Said so, I think I can give a couple of suggestions.

3 Yes, the choice of architecture will produce vectors with different properties. This applies also to vectors generated with classic matrix factorization techniques, not only deep learning models. For example, transformers tend to perform better in regards with the position of a word in a sentence thanks to self attention and positional embeddings, LSTM suffer from long range dependencies issues (so the correlation between the first and last amino acids will be poorly captured).

2 Embedding is not just a trainable layer. In fact, many nlp applications leverage pretrained embeddings. So you could train your own embeddings prior to training your classifier using host species as target labels. There are a variety of approaches to do so, the classic one CBOW, Skip-Gram and GloVe. Of course to train good embeddings you need lot of documents (sequences in your case). This is why pretrianed embeddings are usually the way to go, cause they were trained on billions of documents (an operation that can take not hours but months). Still, even for small corpora, training embeddings with one of these approaches and then refine them on a classification task (so using them as a not randomly initialized but trainable layer) leads to better results.

2 (pt2) The order problem is not that hard to tackle as you might think. Averaging the vectors is indeed a good choice when it comes to combine embeddings to get a whole sentence representation. And indeed averaging doesn't provide any information about the order. I already mention that for example transformer are designed to account for order, but there are also simpler ways. The first that comes to my mind, specific for your case, is to simply increase the amount of embeddings to learn, and move from a 1-gram model (single characters) to 2/3-grams models, so training embeddings for A, B, AB, BA and so on. Computationally it won't be a problem since you're starting from 21, which is a ridiculously small amount of vectors compared to the millions of embeddings trained for classic NLP (and I would argue that also 10 as chosen final dimension is kinda small). Another idea could be to train "dummy" positional embeddings yourself, (dummy cause no transformers of self attention is involved). This would require to stick to a fixed maximum length though. Of course 21 becomes big if you think at all possible permutations of let's say 40 characters ($40^{21}$ is a monster number), but if you train separate embeddings for each combination of character+position, so A1, A2, B1, B2, and so on, then the total amount of vectors to train would be just 40*21=840, which again is nothing compared to normal embeddings dictionaries.