LEARNING PROTEIN SEQUENCE EMBEDDINGS USING INFORMATION FROM STRUCTURE
Inferring the structural properties of a protein from its amino acid sequence is a challenging yet important problem in biology. Structures are not known for the vast majority of protein sequences, but structure is critical for understanding function. Existing approaches for detecting structural similarity between proteins from sequence are unable to recognize and exploit structural patterns when sequences have diverged too far, limiting our ability to transfer knowledge between structurally related proteins. We newly approach this problem through the lens of representation learning. We introduce a framework that maps any protein sequence to a sequence of vector embeddings -one per amino acid position -that encode structural information. We train bidirectional long short-term memory (LSTM) models on protein sequences with a two-part feedback mechanism that incorporates information from (i) global structural similarity between proteins and (ii) pairwise residue contact maps for individual proteins. To enable learning from structural similarity information, we define a novel similarity measure between arbitrarylength sequences of vector embeddings based on a soft symmetric alignment (SSA) between them. Our method is able to learn useful position-specific embeddings despite lacking direct observations of position-level correspondence between sequences. We show empirically that our multi-task framework outperforms other sequence-based methods and even a top-performing structure-based alignment method when predicting structural similarity, our goal. Finally, we demonstrate that our learned embeddings can be transferred to other protein sequence problems, improving the state-of-the-art in transmembrane domain prediction.
General comment
==============
The authors describe two loss functions for learning embeddings of protein amino acids based on i) predicting the global structural similarity of two proteins, and ii) predicting amino acid contacts within proteins. As far as I know, these loss functions are novel and the authors show clear improvements when using the learned embeddings in downstream tasks. The paper is well motivated and mostly clearly written. However, the evaluation must be strengthened and some aspects of it clarified. Provided that the authors address my comments below, I think it is a good ICLR application paper.
Major comments
=============
1. The authors should describe how they optimized hyperparameters such as the learning, lambda (loss section 3.3), or the smoothing factor (section 3.4). These should be optimized on an evaluation set, but the authors only mentioned that they split the dataset into training and holdout (test) set (section 4.1).
2. The way the authors present results in table 1 and table 2 is unclear. Both table 1 and table 2 contain results of the structural similarity tasks but with different baselines. ‘SSA w/ contact predictions’ is also undefined and can be interpreted as ‘with’ or ‘without’ contacts predictions. I therefore strongly recommend to show structural similarity results in table 1 and secondary structure results in table 2 and include in both tables i) ‘SSA full’, ‘SSA without contact predictions’, and ‘SSA without language model’.
3. The authors should compare SSA to the current state-of-the art in structure prediction in addition to baseline models.
4. The authors should evaluate how well their method predicts amino acid contact maps.
5. The authors should describe how they were dealing with variable-length protein sequences. Are sequences truncated and embedded to a fixed length? What is the mean and variance in protein sequence lengths in the considered datasets? The authors should point out that their method is limited to fixed length sequences.
6. The authors should briefly describe the training and inference time on a single GPU and CPU. How much memory is required for training with a certain sequence length, e.g. 400 amino acids per sequence? Does the model fit on a single GPU?
7. The authors should discuss limitations of their method, e.g. that it cannot handle variable length sequences and that the memory scales quadratically by the the sequence length.
8. CRF (SSA) (table 3) includes a biLSTM layer between SSA and the CRF. However, the biLSTM can learn a non-linear projection of embeddings learned by SSA such that it is unclear if improvements are due to the embeddings learned by SSA or the biLSTM+CRF architecture. The authors should therefore train a biLSTM+CRF model on one-hot encoded amino-acids and include it as baseline in table 3.
Minor comments
=============
9. The way the similarity score s’ is computed (section 3.2.1) should be motivated more clearly. Why do the authors compute the score s’ manually instead of predicting it, e.g. using a model that takes the embeddings z of both proteins as input and predicts a single scalar s’? 
10. How does ordinal regression (section 3.2.2) perform compared with a softmax layer? Why do the authors compute s’ and then train logistic regression classifiers on s’ to predict the similarity level, instead of predicting the similarity level directly based on the embeddings z?
11. Why do the authors use a distance threshold of 8A (section 3.3)? Is this common practice in the field?
12. Why do the authors use the not product and the absolute difference as features instead of the embeddings z directly? Which activation function is used to predict contact probabilities (sigmoid, softmax, …)?
13. The authors should reference and describe the results presented in table 1 more clearly.
14. Optional: the authors should analyze if learned embeddings are correlated with amino acid and structural properties such as their size, charge, or solvent accessibility. Do embeddings clusters by certain properties? This can be analyzed, e.g., using a tSNE plot. 
15. How does TMalign perform when using the maximum or geometric average instead of the arithmetic average of the two scores (section 4.1)