Improving Generalizability of Protein Sequence Models via Data Augmentations

While protein sequence data is an emerging application domain for machine learning methods , small modifications to protein sequences can result in difficult-topredict changes to the protein ’ s function . Consequently , protein machine learning models typically do not use randomized data augmentation procedures analogous to those used in computer vision or natural language , e.g. , cropping or synonym substitution . In this paper , we empirically explore a set of simple string manipulations , which we use to augment protein sequence data when fine-tuning semisupervised protein models . We provide 276 different comparisons to the Tasks Assessing Protein Embeddings ( TAPE ) baseline models , with Transformer-based models and training datasets that vary from the baseline methods only in the data augmentations and representation learning procedure . For each TAPE validation task , we demonstrate improvements to the baseline scores when the learned protein representation is fixed between tasks . We also show that contrastive learning fine-tuning methods typically outperform masked-token prediction in these models , with increasing amounts of data augmentation generally improving performance for contrastive learning protein methods . We find the most consistent results across TAPE tasks when using domain-motivated transformations , such as amino acid replacement , as well as restricting the Transformer attention to randomly sampled sub-regions of the protein sequence . In rarer cases , we even find that information-destroying augmentations , such as randomly shuffling entire protein sequences , can improve downstream performance . 1 INTRODUCTION . Semi-supervised learning has proven to be an effective mechanism to promote generalizability for protein machine learning models , as task-specific labels are generally very sparse . However , with other common data types there are simple transformations that can be applied to the data in order to improve a model ’ s ability to generalize : for instance , vision models use cropping , rotations , or color distortion ; natural language models can employ synonym substitution ; and time series data models benefit from window restriction or noise injection . Scientific data , such as a corpus of protein sequences , have few obvious transformations that can be made to it that unambiguously preserve the meaningful information in the data . Often , an easily understood transformation to a protein sequence ( e.g. , replacing an amino acid with a chemically similar one ) will unpredictably produce either a very biologically similar or very biologically different mutant protein . In this paper , we take the uncertainty arising from the unknown effect of simple data augmentations in protein sequence modeling as an empirical challenge that deserves a robust assessment . To our knowledge , no study has been performed to find out whether simple data augmentation techniques improve a suite of protein tasks . We focus on fine-tuning previously published self-supervised models that are typically used for representation learning with protein sequences , viz . the transformerbased methods of Rao et al . ( 2019 ) which have shown the best ability to generalize on a set of biological tasks , which are referred to as Tasks Assessing Protein Embeddings ( TAPE ) . We test one or more of the following data augmentations : replacing an amino acid with a pre-defined alternative ; shuffling the input sequences either globally or locally ; reversing the sequence ; or subsampling the sequence to focus only on a local region ( see Fig . 1 ) . ( a ) Replacement ( Dictionary ) ( b ) Replacement ( Alanine ) ( c ) Global Random Shuffling We demonstrate that the protein sequence representations learned by fine-tuning the baseline models with data augmentations results in relative improvements between 1 % ( secondary structure accuracy ) and 41 % ( fluorescence ρ ) , as assessed with linear evaluation for all TAPE tasks we studied . When fine-tuning the same representations during supervised learning on each TAPE task , we show significant improvement as compared to baseline for 3 out of 4 TAPE tasks , with the fourth ( fluorescence ) within 1σ in performance . We also study the effect of increasingly aggressive data augmentations : when fine-tuning baseline models with contrastive learning ( Hadsell et al. , 2006 ; Chen et al. , 2020a ) we see a local maximum in downstream performance as a function of the quantity of data augmentation , with “ no augmentations ” generally under-performing modest amounts of data augmentations . Conversely , performing the same experiments but using masked-token prediction instead of contrastive learning , we detect a minor trend of decreasing performance on the TAPE tasks as we more frequently use data augmentations during fine-tuning . We interpret this as evidence that contrastive learning techniques , which require the use of data augmentation , are important methods that can be used to improve generalizibility of protein models . 2 RELATED WORKS . Self-supervised and semi-supervised methods have become the dominant paradigm in modeling protein sequences for use in downstream tasks . Rao et al . ( 2019 ) have studied next-token and maskedtoken prediction , inspired by the BERT natural language model ( Devlin et al. , 2018 ) . Riesselman et al . ( 2019 ) have extended this to autoregressive likelihoods ; and Rives et al . ( 2019 ) , Heinzinger et al . ( 2019 ) and Alley et al . ( 2019 ) have shown that unsupervised methods trained on unlabeled sequences are competitive with mutation effect predictors using evolutionary features.Of importance to this work are self-supervised learning algorithms employed for other data types that use or learn data augmentations . For example , Gidaris et al . ( 2018 ) learn image features through random rotations ; Dosovitskiy et al . ( 2014 ) and Noroozi & Favaro ( 2016 ) study image patches and their correlations to the original samples . van den Oord et al . ( 2018 ) uses contrastive methods to predicts future values of an input sequence . We consider sequence augmentations in natural language as the most relevant comparison for the data augmentations we study in this paper . Some commonly applied augmentations on strings include Lexical Substitution ( Zhang et al. , 2015 ) , Back Translation ( Xie et al. , 2019a ) , Text Surface TransformationPermalink ( Coulombe , 2018 ) , Random Noise Injection ( Xie et al. , 2019b ; Wei & Zou , 2019 ) , and Synonym Replacement , Random Swap , Random Deletion ( RD ) ( Wei & Zou , 2019 ) . However , sequence augmentations designed for natural languages often require the preservation of contextual meaning of the sentences , a factor that is less explicit for protein sequences . Contrastive Learning is a set of approaches that learn representations of data by distinguishing positive data pairs from negative pairs ( Hadsell et al. , 2006 ) . SimCLR ( v1 & v2 ) ( Chen et al. , 2020a ; b ) describes the current state-of-the-art contrastive learning technique . ; we use this approach liberally in this paper not only because it performs well , but because it requires data transformations to ex- ist . Since we focus on protein sequence transformations , the contrastive learning part described in both ( Chen et al. , 2020a ; b ) is our focus . Following Chen et al . ( 2020a ; b ) , we denote input data as x ∈ D , with D being our training set ; we then define an embedding function , fω : x 7→ h with h ∈ RN , and a mapping function gθ : h 7→ z , with z ∈ RM , where ω and θ are the learned model weights . For any x , we form two copies x1 = t1 ( x ) and x2 = t2 ( x ) given functions t1 , t2 ∼ T , where T denotes the distribution of the augmentation functions . Given D is of sizeN , the contrastive loss is written as : L = 1 2N N∑ k=1 [ l ( z ( 1 ) k , z ( 2 ) k ) + l ( z ( 2 ) k , z ( 1 ) k ) ] where l ( u , v ) ≡ − log esim ( u , v ) /τ∑ w 6=u e sim ( u , w ) /τ ( 1 ) Here , zi , k = gθ ( fω ( ti ( xk ) ) ) , sim ( · , · ) is cosine similarity , and τ ∈ ( 0 , ∞ ) is a scalar temperature ; we choose τ = 0.2 . By minimizing the contrastive loss , we obtain the learned h as the encoded feature for other downstream tasks . Note , the contrastive loss takes z ’ s as inputs , whereas the encoded feature is h , which is the variable after the function fω ( · ) and before gθ ( · ) . 3 METHOD . 3.1 EVALUATION PROCEDURE & APPROACH TO EXPERIMENT CONTROL . Our goal is to demonstrate that training self-supervised protein sequence models , with simple string manipulations as data augmentations , will lead to better performance on downstream tasks . To attempt to control external variables , we study the following restricted setting ; we provide the procedural diagram in Figure 2 and the corresponding explanations of the four major steps below ( See Appendix A for training setups in details . ) : Baseline.— A self-supervised model M0 is trained on non-augmented sequence data Dseq to do representation learning . To have a consistent baseline , we set M0 to the Transformer-based model trained and published in Rao et al . ( 2019 ) , without modification . This was trained with maskedtoken prediction on Pfam protein sequence data ( El-Gebali et al. , 2019 ) ; it has 12 self-attention layers , 8 heads per layer , and 512 hidden dimensions , yielding 38M parameters in total . Augmented training on validation set.— We fine-tune M0 on augmented subsets Dval ⊂ Dseq , given a set of pre-defined data transformations Taug . We define Maug as the final trained model derived from Taug ( Dseq ) with M0 as the initial conditions for the model parameters . We explore two different methods of fine-tuning on augmented data — a contrastive task ( as in Eq . 1 ) and a masked-token task ( exponentiated cross entropy loss ) — as well as different combinations of data augmentations . We use reduced subsets |Dval| |Dseq| to both reduce the computational cost of running bulk experiments , as well as to protect against overfitting . For consistency , we inherit the choice of Dval from the cross-validation split used to train M0 in Rao et al . ( 2019 ) . To adapt the same baseline model M0 to different self-supervised losses , we add a loss-specific randomlyinitialized layer to theM0 architecture : contrastive learning with a fully connected layer that outputs 256 dimensional vectors and masked-token uses fully connected layers with layer normalization to output one-hot vectors for each of the masked letters . We define our different choices of Taug in the next section . Linear evaluation on TAPE.— To assess the representations learned by Maug , we evaluate performance on four TAPE downstream training tasks ( Rao et al. , 2019 ) : stability , fluorescence , remote homology , and secondary structure . For consistency , we use the same training , validation , and testing sets . The first two tasks are evaluated by Spearman correlation ( ρ ) to the ground truth and the latter two by classification accuracy . However , we do not consider the fifth TAPE task , contact map prediction , as this relies only on the single CASP12 dataset , which has an incomplete test set due to data embargoes ( AlQuraishi , 2019 ) . Secondary structure prediction is a sequence-to-sequence task where each input amino acid is classified to a particular secondary structure type ( helix , beta sheet , or loop ) , which is evaluated on data from CASP12 , TS115 , and CB513 ( Berman et al. , 2000 ; Moult et al. , 2018 ; Klausen et al. , 2019 ) , with specifically “ 3-class ” classification accuracy being the metric in this paper . The remote homology task classifies sequences into one of 1195 classes , representing different possible protein folds , which are further grouped hierarchically into families , then superfamilies ; the datasets are derived from Fox et al . ( 2013 ) . The fluorescence task regresses a protein sequence to a real-valued log-fluorescence intensity measured in Sarkisyan et al . ( 2016 ) . The stability task regresses a sequence to a real-valued measure of the protein maintaining its fold above a concentration threshold ( Rocklin et al. , 2017 ) . We perform linear evaluation by training only a single linear layer for each downstream task for each contrastive-learning model Maug , but not changing the parameters of Maug , and its corresponding learned encodings , across all tasks . To compare the contrastive learning techniques to further fine-tuning with masked-token prediction , we identify the best-performing data augmentations per-task , then replace the MCLaug with the masked-token model with the same augmentation MMTaug , then also do linear evaluation on M MT aug . Full fine-tuning on TAPE.— For the best-performing augmented models in the linear evaluation task ( eitherMCLaug orM MT aug ) , we further study how the models improve when allowing the parameters of Maug to vary along with the linear model during the task-specific supervised model-tuning .

Summary: The paper explores the impact of different types of data augmentations for protein sequence data, and does a thorough benchmark analysis on them. The authors used a pre-trained transformer model, fine tuned the model on augmented data using two approaches, namely, contrastive learning and masked token prediction. This finetuned model was evaluated with an added linear layer on a range of tasks.