Multi-modal Self-supervised Pre-training for Regulatory Genome Across Cell Types

1 INTRODUCTION . In the gene biology research , many important regulatory problems including promoter ( Li et al. , 2006 ) classification and transaction factor binding sites ( Stewart et al. , 2012 ) prediction requires a well-designed approach for modeling regulatory genome sequences . As a multi-cellular organism , the human body is formed by different cell types , with each has its own gene expression patterns ( Schena et al. , 1995 ; Ross et al. , 2000 ) . Therefore , modeling regulatory genome , especially across different cell types , is crucial for both the understanding of this fundamental biological process and the development of RNA-level therapeutic intervention of a vast among of diseases . The main challenge is how to model the interaction between regulatory elements and its variability across different cell types . In recent years , some works ( Chen et al. , 2016 ; Kelley et al. , 2016 ; Romero et al. , 2017 ; Qiu et al. , 2018 ; Torada et al. , 2019 ; Chen et al. , 2020a ; Avsec et al. , 2021 ; Ji et al. , 2021 ) have been proposed to apply deep learning architectures on modeling the genome data . For example , Enformer ( Avsec et al. , 2021 ) combines dilated CNN and transformer architecture as well as multi-head output to predict gene expression and epigenomic marks , and gain some performance improvements against traditional methods . Most of these models are based on a supervised learning paradigm , which limits their abilities to learn the general interaction representations between different regulatory elements . As a result , each model is only useful for some specific downstream tasks . More recently , DNABERT ( Ji et al. , 2021 ) is introduced to formulate the whole DNA sequence as a sentence of nucleotide k-mers and utilize BERT to model the sequence generatively . In this way , the interactions between different regulatory elements could be well captured , like the language modeling approach in natural language processing ( NLP ) to capture the co-occur information between different words , and could be employed for modeling several related downstream tasks . However , DNABERT only performs well on the cell types in the training set , and generalizes poorly to unseen cell types . This is mainly because the whole DNA sequence is common for different cells , and pre-training on such data can not well reflect various interaction patterns across different cell types . Integration of genome data modalities across different cell types could help to build a more holistic model of gene expression regulation and benefit downstream applications such as mutation impact evaluation and disease risk prediction , as well as promote our understanding of the cell-type-specific regulatory programs and various development processes and disease etiology . Inspired by this idea , in this work , we present a simple yet effective method called GeneBERT , for pre-training large-scale genome data in a multi-modal and self-supervised manner . Specifically , we simultaneously take the 1D modality ( i.e . sequence ) and a 2D modality ( i.e . regulatory region ) of genome data as the input , where three pre-training tasks are proposed to improve the robustness and generalizability of our model . 1 ) masked genome modeling : we randomly mask some parts of the input k-mers with a special token ( i.e. , [ MASK ] ) , and the model is trained to predict the masked k-mers . 2 ) next genome segment prediction : we train the model using the embedding [ CLS ] to classify whether a pair of given sequences are two consecutive sequences in a cell . 3 ) sequence-region matching : a sequenceregion matching mechanism is proposed to capture the multi-modal alignment between sequence and regulatory region of genome data . We pre-train the GeneBERT on the ATAC-seq dataset ( Domcke et al. , 2020 ) with 17 million gene sequences . Furthermore , we conduct extensive experiments to evaluate our GeneBERT on various downstream tasks , including promoter prediction , transaction factor binding sites prediction , gene mutation localization , and personalized diseases prediction . Comprehensive ablation studies demonstrate the effectiveness of multi-modal self-supervised pre-training for large-scale genome data across different cell types . The main contributions of this work lie in the proposal of a simple yet effective method named GeneBERT , for large-scale genome data pre-training in a multi-modal and self-supervised manner . To the best of our knowledge , we are the first to incorporate different genome data modalities across various cell types into the pre-training for large-scale genome data , to tackle the regulatory genome modeling problem . Except for meaningful biological improvements , our model makes an important contribution to the machine learning community , by introducing a novel multi-modality construction . Different from existing multi-modal learning tasks , such as visual question answering , image caption , and image-text retrieval , the ‘ visual ’ modality in our work is constructed based on the regulatory property of the ‘ language ’ sequential units . This idea brings some inspiration to building a new ‘ visual ’ modality , based on text matching , part-of-speech tagging , and named entity recognition etc. , for the study of NLP . 2 RELATED WORK . Uni-modal language pre-training . Self-supervised pre-training models such as GPT ( Radford et al. , 2018 ) , BERT ( Devlin et al. , 2018 ) , RoBERTa ( Liu et al. , 2019 ) , and ERNIE ( Sun et al. , 2019 ) have led to dramatic improvement on a variety of natural language processing tasks in the past few years , significantly surpassing the traditional context-independent language model such as Word2Vec . Obviously , BERT ( Devlin et al. , 2018 ) uses the masked language model ( MLM ) and next sentence prediction ( NSP ) for pre-training , which represent the dynamic vector of words according to the context . RoBERTa ( Liu et al. , 2019 ) uses dynamic MLM and discards NSP to train the model with a long time . ERNIE ( Sun et al. , 2019 ) masks entities and phrases to learn more context relations . These self-supervised pre-training models are capable of performing natural language processing tasks more effectively by improving MLM . Multi-modal pre-training . Multi-modal pre-training has recently addressed researchers ’ attention for learning meaningful representations . Typically , Previous methods ( Radford et al. , 2021 ; Huo et al. , 2021 ) learn visual representations from text paired with images in unsupervised , selfsupervised , weakly supervised , and supervised ways . Since language and vision can share a similar semantic meaning , CLIP ( Radford et al. , 2021 ) is a commonly-used neural network trained on a variety of ( image , text ) pairs for learning transferable visual representations from natural language supervision . WenLan ( Huo et al. , 2021 ) applies a cross-modal contrastive learning framework called BriVL for image-text pre-training . However , in this work , we leverage the multi-modal self-supervised pre-training on the genome data to improve the robustness and generalizability of pre-trained models used for data-scarce scenarios . Genome data pre-training . Transformer models have been recently established to better understand the genotype-phenotype relationships ( Avsec et al. , 2021 ; Ji et al. , 2021 ) . DNABERT used the human genome to pre-train a BERT-based model , trying to decipher the regulatory code related to gene expression ( Ji et al. , 2021 ) . In order to adapt the DNA scenario , sequences were split into 5 to 510 base-pair long and tokenized to 3- to 6-mer representations . The model was then pre-trained using those k-mers by predicting around 15 % randomly masked regions of the sequence . After the pre-training , the model was fine-tuned on three downstream tasks related to gene regulation : prediction of promoters , transcription factor binding sites ( TFBSs ) , and splice sites . Furthermore , by analyzing the attention maps , DNABERT could visualize the important regions contributing to the model decision , which improved the interpretability of the model . 3 METHOD : GENEBERT . In this section , we propose a simple yet effective approach for pre-training genome data in a multimodal and self-supervised manner , as illustrated in Figure 1 . Specifically , we introduce two main objectives for sequence pre-training , including Masked Genome Modeling ( MGM ) and Next GenomeSegment Prediction ( NGSP ) . Next , we extract features of the region for region pre-training . Finally , we present the Sequence-Region Matching mechanism to explicitly align features of the sequence and region in the input genome data in a contrastive self-supervised manner . Therefore , our GeneBERT consists of three main components : sequence pre-training , region pre-training , and sequence-region matching . 3.1 INPUT DATA . We use potential regulatory elements mapped with ATAC-seq which measures the chromatin accessibility across the genome . The regions with open accessibility could be bound by transcription factors , which together with regulatory elements and RNA polymerase determines the expression level . In one cell type , the list of all regulatory elements can be viewed as a set of non-overlapped sub-regions across the genome . Each region or regulatory element can be viewed as a string like “ AATTCCT ... ” consist of 4 elements : [ A , T , C , G ] . The binding site of different transcription factors has different sequence patterns , some of which have been determined using experimental methods and have been modeled using probabilistic models like position weight matrix ( PWM ) . Using PWM , one can get the predicted affinity of a sequence segment to a given transcription factor . Each regulatory region can then be associated with a ‘ transcription factor binding vector ” which contains the in silico prediction of their affinity to a fixed number of transcription factors . By scanning through a regulatory element sequence with PWM of multiple transcription factors , we can get a transcription factor binding affinity vector for that regulatory element . When the collection of transcription factor binding models is comprehensive and precise enough , we can get a good representation of regulatory potential for a given sequence . In different cell types , the accessibility of regulatory elements is different . Some regulatory elements are shared across different cell types while some are more specific to one cell type . As a result , in a local region of the genome , the accessible transcription factor binding sites are determined by the combination of accessible regulatory regions . This combinatorial interaction between regulatory regions and transcription factors varies across different cell types such that precise control of expression in different cell types is possible . To capture the combinatorial interaction and the implicit cell-type specific information , we then group 10 consecutive regulatory elements into one train sample by concatenating their string ( 1D sequence ) and stacking their binding vector into a 2D region , as shown in Figure 1 . 3.2 SEQUENCE PRE-TRAINING . Similar to DNABERT , we take k-mer tokens as the sequence units . The k-mer refers to a sequence with length k , i.e. , for a sequence AGTCAG , the 3-mers are { AGT , GTC , TCA , CAG } , and the 4-mers are { AGTC , GTCA , TCAG } . Therefore for sequence pre-training , we input three types of embeddings : 1 ) a k-mer embedding xk for each k-mer in a sequence ; 2 ) a segment embedding xs indicating which part of the sequence the k-mer is from ; 3 ) a position embedding xp for the position of the k-mer in the sequence . Then we sum up all three embeddings to obtain a contextual representation en , n ∈ { 1 , 2 , ... , N } , where N denotes the number of k-mers in the sequence . After being fed into a BERT-based transformer fseq ( · ) , those contextual embeddings become Eseq , i.e . the obtained representation for [ CLS ] . Inspired by BERT ( Devlin et al. , 2018 ) , we adopt two objectives including Masked Genome Modeling ( MGM ) Lmgm and Next Genome Segment Prediction ( NGSP ) Lngsp . For Lmgm , we randomly mask some parts of the input k-mers with a special token ( i.e. , [ MASK ] ) , and the model is trained to predict the masked kmer . The objective is formulated as : Lmgm = max P ( Kmask|K ) ( 1 ) where K and Kmask denote the unmasked and masked k-mers , respectively . As for Lngsp , we train the model using the embedding [ CLS ] to classify whether a pair of given sequences are consecutive in a cell . The loss function of Lngsp is calcualted as : Lngsp = CrossEntropy ( P ( [ CLS ] ) , y ) ( 2 ) where P ( [ CLS ] ) and y represent the predicted probability and target for classification , respectively .

Summary: This work proposes an approach, called GeneBERT, for pre-training genome data in a multi-modal and self-supervised manner. They take the 1d sequence of genomic data and a 2d matrix of (transcription factors × regions) as the input and optimize three pre-training tasks to improve the robustness and generalizability the model. Specifically, they introduce two main objectives for sequence pre-training, Masked Genome Modeling (MGM) and Next Genome-Segment Prediction (NGSP). GeneBERT consists of three main components: sequence pre-training, region pre-training, and sequence-region matching. The main contribution of the paper is that they use transcription factor information in genomic regions which makes the model more generalizable to other cell types than the previous DNABERT that only uses the sequence information.