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: The manuscript describes GeneBERT, a self-supervised and multi-modal pre-training approach for genomic data. GeneBERT combines 1D genome sequence data with a 2D representation of regulatory elements in different cell-types in three different pre-training tasks. A large-scale single-cell ATAC dataset is used to identify pseudo-bulk ATAC profiles for 17 cell-types spanning > 1M genomic locations. The sequence of these genomic locations provides the 1D representation of the data which are used in self-supervised training with masked genome modeling and next genome segment prediction loss functions, both inspired by the BERT framework. A 2D representation is derived using accessibility per regulatory element and cell type, and is self-supervised using a infoNCE loss. The pre-training framework is then compared in a variety of biological tasks including promoter prediction, TFBS prediction, splice site prediction. The authors then present ablation studies highlighting the importance of the different components of the loss function.