Transformer Model for Genome Sequence Analysis
Anonymous Author(s)
AffiliationAddress email
Abstract
One major challenge of applying machine learning in genomics is the scarcityof labeled data, which often requires expensive and time-consuming physicalexperimentation under laboratory conditions to obtain.
However, the advent ofhigh throughput sequencing has made large quantities of unlabeled genome dataavailable.
This can be used to apply semi-supervised learning methods throughrepresentation learning.
In this paper, we investigate the impact of a popularand well-established language model, namely BERT [Devlin et al., 2018], forsequence genome analysis.
Specifically, we adapt DNABERT [Ji et al. , 2021]to GenomeNet-BERT in order to produce useful representations for downstreamtasks such as classification and semi-supervised learning.
We explore differentpretraining setups and compare their performance on a virus genome classificationtask to strictly supervised training and baselines on different training set size setups.
The conducted experiments show that this architecture provides an increase inperformance compared to existing methods at the cost of more resource-intensivetraining.
Introduction
Just as human beings use languages to communicate, nature created its own language: Genomes.
In order to understand this “language of life”, Natural Language Processing (NLP) methods havebeen used with the aim of decoding instructions and information contained within [Asgari and
Mofrad, 2015].
As to unravel the complex function and structures of cells hidden within theirgenomes semi-supervised learning can be applied to improve the capabilities to identify new genomestructures, impute missing nucleotides (NTs) and classify genomic data under sparse label conditions[BMBF, 2020].
Deep learning based methods have recently achieved breakthroughs in bioinformaticsby exceeding the performance of previous state-of-the-art approaches [Zhang et al., 2021].
Selfsupervised models have prevailed in NLP, as they take advantage of readily available amounts ofunlabeled data in the form of texts to pretrain model weights and representations in a self-supervisedmanner, thereby leading to higher performance on downstream language tasks.
[Devlin et al., 2018,
Radford et al., 2018, Peters et al., 2018].
One architecture among those models that have provenparticularly useful for representation learning of genomic data is the Bidirectional TransformerEncoder [Devlin et al., 2018, Rives et al., 2020, Ji et al., 2021, Le et al., 2021, Mo et al., 2021,
Avsec et al., 2021].
Using NLP methods for DNA data is an attractive idea, as both written humanlanguage and genome data coincide in their method of representing data as sequences of discreteinformation: Letters or words for language, and NTs in the case of DNA.
However, many NLPmethods, in particular BERT , rely on tokenization of text into discrete words or sub-word units [Wuet al., 2016].
While words as units of the information above the character-level are straightforward for humans to recognize and encode in the natural language domain, there is no readily apparent wayof tokenizing DNA sequences in general 1 .
A simple and yet effective approach to tokenization of
DNA data is to use k -mers:
Encoding “words” of DNA as units of k sequential NTs.
Ji et al.
[2021]introduce this method to decode human genome data with DNABERT and achieved promising results.
In this work, we analyze the potential of the Bidirectional Transformer-Encoder applied to virusgenome sequence data by implementing GenomeNet-BERT .
In a small test framework, optimalhyperparameter and tokenization settings are explored.
Furthermore, two additional strategies,differing in data preprocessing and pretraining setup to the original model implementation, arepursued.
The evaluation of model performance is conducted on the task of identifying bacteriophagesfrom short sequences of NTs over various label scarcity scenarios and sequence lengths, by comparingit to the performance of the same architecture trained in a fully supervised fashion, and the SelfGenomeNet [Gündüz et al., 2022] architecture.
Compared to fully supervised methods, this oneprovides reusability for a variety of downstream tasks and is well suited for further improvements ondata preprocessing and pretext task tuning, since its architecture is not dependent on these.
Giventhat the BERT architecture is highly explicable with its attention layers, its application helps tounderstand the importance of nucleotide snippets in terms of classification.
Further, the model usesraw nucleotide sequence data as input, which reduces the data preprocessing overhead [Buermansand den Dunnen, 2014].
Method
The BERT model is built of 12 stacked Transformer-Encoder blocks and relies solely on bidirectionalattention and fully connected layers to learn representations for each input token.
Two pretext tasksaid this purpose: Predicting randomly masked words and whether two input sequences are consecutivein the source they were extracted from.
DNABERT
[Ji et al., 2021] takes genome sequences as input.
Unlike BERT, next sentence prediction is not used as a pretext task.
To attain input sequences,genomes are split into non-overlapping sequences of sampled length and cut from randomly sampledlocations.
These are then tokenized using all permutations of the k -mer representation of size 768,which creates tokens with stride 1 (see Figure 1).
Because of this overlap of NTs per token, it ispossible to simply infer a masked token by its neighbors.
Therefore, instead of randomly sampling apercentage of m tokens to mask, k consecutive tokens per sampled masking location are masked.
Procedure We devised a three-step-method, starting with hyperparameter optimization within ascaled-down framework to find values for learning rate, masking percentage, weight decay, andothers, leading to the best performance for our tasks, as well as compare different strategies in regardto data preprocessing, tokenization and pretext task.
Subsequently, three architecture designs basedon DNABERT
[Ji et al., 2021] (referred to as GenomeNet-BERT in the following), differing mainly intokenization and pretext task, are pretrained full scale.
GenomeNet-BERT models were all pretrainedfor 100k steps due to loss plateauing and associated comparability reasons between models.
Finally,after subsequent supervised training on our bacteriophage classification task, the performance ofour models is compared to the same architecture fully supervised trained and self-genomenet , aself-supervised model proposed by Gündüz et al.
Fine-tuning is performed over differenttraining set sizes, representing various scenarios of label scarcity.
Additionally, two distinct inputsequence lengths (150 and 1000 nucleotide lengths, respectively) are examined.
The data used isdescribed in the Appendix (see Section A.2).
Architecture Designs After self-supervised pretraining, all setups are fine-tuned in a supervisedmanner on balanced, labeled subsets of the dataset, and macro-averaged recall Recall M (in %), aswell as F 1 -score, are then used to measure model performance on a separate prediction set.
Thefirst adaptation, GenomeNet-BERT , is a replica of the DNABERT setup adapted for our purposes.
Ittokenizes sequences of up to 510NTs to 6-mers and masks 6 consecutive tokens at 2.5% sampled token locations, leading to a token masking rate of 15%.
A learning rate of 4 × 10 − 4 , linearly warmedup over 5% of the total steps and AdamW-Optimizer [Loshchilov and Hutter, 2019] are applied.
In contrast to the other setups, the hyperparameters are based on the original settings proposed by
Ji et al.
The main reasoning behind altering pretraining strategies stems from how k-mertokenization and masking interact, to only hide 2.5% of distinct NTs for this first setup, which wasperceived as low.
Therefore, the second adaptation, GenomeNet-BERT-mask8 , is designed to maskmore NTs.
It masks 8 consecutive tokens at 2.875% sampled token locations, which leads to morethan double the amount of NTs hidden from the model, while the ratio of masked tokens remains thesame at 15% (see Figure 1).
The third adaptation, GenomeNet-BERT-stride3 , is intended to work withlonger input sequences and train faster and up to 1000NTs long sequences are tokenized to 6-mers ofstride 3 and 3 consecutive tokens at 5% sampled token locations are masked.
This leads to a sixfoldincrease in the number of NTs hidden during pretraining compared to the first GenomeNet-BERTimplementation.
Tokenizing sequences with up to 1000NTs to 6-mers with stride 3 leads to 332 inputtokens in total and the model is further hard-limited to input sequences up to 340 tokens compare tothe standard 512 tokens for BERT.
Experiments
Label availability scenarios are artificially created by limiting access to a specific subset of FASTAfiles during training.
As in the semi-supervised protocol of Henaff [2020], 1% and 10% labeled datais used.
In addition, a very sparse label setting of 0.1% is trained.
Since a k of 6 performs best in theexperiments by Ji et al.
[2021], 6-mers are used in all setups.
Hyperparameter Optimization
The impact of hyperparameter settings is evaluated using thebert-small configuration [Wolf et al., 2020], since it is closely related to the original emphBERTarchitecture, but at the same time allows for performing experiments in a less time-consuming fashion.
The detailed results of these experiments are listed in the Appendix in Figure A.3.
GenomeNet-BERTmask8 and GenomeNet-BERT-mstride3 are accordingly trained with an increased learning rate of1 × 10 − 3 and a longer linear warmup of 20k steps.
While no masking setup is consistently betteracross the range of learning rates trialed, setups that mask 8 consecutive tokens can perform equallywell or better than the standard setup at the same learning rate, producing the best model of all stride1 setups.
Fine-Tuning To fine-tune our models, network heads of the pretext task and all representationsexcept the first token are removed.
This starting token collects sequence-level information and is fedexclusively to a classifier via an additional projection layer to predict the class of an input sequence(see Figure A.2).
Sequences of length 1000NTs tokenized as k -mers with a stride of 1 surpass theinput token limitation of BERT -base.
They are split into two parts, traversed individually, and thenconcatenated again by a fully connected layer for all DNABERT -based models on the 1000NTs taskexcept for GenomeNet-BERT-stride3 , which can handle inputs up to 1024nt.
[Table]
Results and Discussion
Table 1 compares model performance for sequences of 150 and 1000NTs, respectively.
Throughoutall 6 scenarios, GenomeNet-BERT -based models show superior performance compared to selfgenomenet , the base GenomeNet-BERT appearing the best on average overall.
The mask8 variantperforms very similar to the base GenomeNet-BERT model, while the stride3 variant provides lesssuccessful class predictions for 150NTs input sequences.
However, it can be seen that the stride3variant exhibits similar or better accuracy than the other variants for 1000NTs input sequences.
Since the stride3 model variant has a shorter input length, it trains notably faster than the other
GenomeNet-BERT models.
Generally, the pretrained GenomeNet-BERT -model manifests an about20% increase in recall than the strictly supervised baseline in all scenarios and increasing in the labelscarcer setups of 1% and 0.1%.
The GenomeNet-BERT model also shows an impressive accuracy inthe low label scenario of 0.1%, outperforming self-genomenet by about 16% and 20% in recall for150 and 1000NTs, respectively.
We have shown that DNABERT , implemented for use with human genome sequences, is also capableof learning representations from virus genome sequences.
Our virus pretrained version, referredto as GenomeNet-BERT , outperforms the given baseline at all input length and label availabilitieson the task of identifying bacteriophages from read-level length genome sequence excerpts.
The
GenomeNet-BERT realization, which follows the original setup of DNABERT6 , also outperformsboth permutations ( mask8 and stride3 ) trialed in this task on average.
However, since the GenomeNetBERT realization was trained using the hyperparameters proposed by Ji et al.
[2021], and HPO wasperformed using the bert-small
[Wolf et al., 2020] configuration, it is possible that the method canbe further improved by HPO based on the full BERT model.
While GenomeNet-BERT-stride3 isless accurate on the shorter input length task, it provides the same level of accuracy at the 1000NTsinput length with much lower resource requirements for both pretraining and fine-tuning than thebase GenomeNet-BERT model.
In general, it must be acknowledged that the Transformer-Encodermodel is very resource and training time intensive, even compared to other self-supervised modelsfor genome sequence analysis.
An interesting observation in the experiments conducted is that all models trained in these experimentsoverpredicted the bacteriophage class in every setup.
It is possible that the model learns to classifymore noisy input sequences as phages, as these could be more diverse in short genome excerpts.
Fora more definitive evaluation of this model architecture, it is necessary to investigate its performanceon a higher number of more diverse downstream tasks.
A. The DNABERT Architecture
A. Virus Data
All self-supervised learning models are trained and evaluated on a collection of viral genomesequences.
On August 2nd, 2021, all available viral genome data was downloaded from GenBank[Sayers et al., 2019] and divided into two taxonomic classes: Bacteriophages, and Other Viruses.
This collection of about 40k FASTA files includes about 1 billion NTs for the bacteriophage class and0.5 billion for other viruses and poses the binary classification task of identifying whether a read-levellength nucleotide sequence is an excerpt of a bacteriophage genome.
All self-supervised models arepretrained using unlabeled nucleotide sequences generated from a training split of the data.
A. Trial Results
[Figure]
A. Computational Information
Pretraining was conducted for 190h on 8 nvidia-A100-40Gb GPUs for GenomeNet-BERT &
GenomeNet-BERT-mask8 and 141h on 5 of the same GPUs for GenomeNet-BERT-stride3 .