BERT-AL: BERT for Arbitrarily Long Document Understanding

1 INTRODUCTION . In recent years , neural networks are proposed to solve various NLP tasks . Especially , pretrained language models ( Peters et al. , 2018 ; Radford et al. , 2018 ; Devlin et al. , 2019 ; Yang et al. , 2019 ) attract lots of attentions , which take advantage of the two-stages training process : pretraining on unlabeled corpus and then finetuning on specific tasks . The most famous model is BERT . BERT and its varieties are the state-of-the-art for many kinds of NLP tasks . The power of BERT does not only come from its architecture of networks , but also because it can be pretrained on a mass of text as a masked language model . BERT can be used to solve almost all NLP tasks , and especially it can perform best on datasets with short text , e.g. , GLUE ( Wang et al. , 2018 ) and Squad ( Rajpurkar et al. , 2016 ) . However , there are still many document-level tasks , e.g. , document-level text summarization ( Hermann et al. , 2015 ) , long-document machine reading comprehension ( Hewlett et al. , 2016 ) and long text classification ( Zhang et al. , 2015 ) . BERT can not be finetuned for such tasks with long text directly or perform good on these tasks , since it is limited by the fixed-length position embedding which was determined during pretraining . We employ document-level text summarization as an example , which usually has longer text than the maximum sequence length of BERT . Intuitively , there are two alternative solutions : 1 ) Truncating inputs by the maximum sequence length to fit the BERTs constraint . 2 ) Increasing the length of position embedding and re-pretraining the BERT from scratch . The first method will decrease performance , obviously , since some useful information placing behind the maximum sequence length is discarded by truncating . E.g. , for text summarization , if the key point sentence locates at the end of the document , it never be recalled even though the model is powerful . For the second method , re-pretraining the BERT from scratch will cost a mass of computing time and resources . Whats even worse is that the computational complexity will increase quadratically with the length , which will result in an unacceptable training time . For example , the XLNet-Large ( Yang et al. , 2019 ) costs 2.5 days on 512 TPU v3 chips and the RoBERTa ( Liu et al. , 2019 ) costs 1 day on 1024 V100 GPUs . To resolve these problems , we propose BERT-AL ( BERT for Arbitrarily Long Document Understanding ) model that extracts local features by applying parallel multi-layer Transformers into chunked input and employs multi-channel LSTMs to capture global information crossing Transformers . This fusion breaks the limitation of BERT by the ability of capture unlimited sequential information from LSTM , and makes it be able to process arbitrarily long text . On the other hand , it also skillfully avoids gradient vanishing and exploding problem ( Li et al. , 2018 ) of LSTM since only few steps are required by multi-channel LSTM . Therefore , BERT-AL can solve the problems of original BERT : 1 ) For document-level tasks , BERT-AL can directly take all text as the input without truncating . 2 ) When the input length is longer then maximum sequence length of BERT , BERT-AL still can load the pretrained BERT model , which avoids pretraining much longer BERT model from scratch . We demonstrate BERT-ALs effectiveness on text summarization by conducting experiments on the CNN/Daily Mail dataset ( Hermann et al. , 2015 ) with various maximum sequence lengths of pretrained BERT . The results prove that BERT-AL can consistently outperform BERTSUM ( Liu , 2019 ) which is the BERT-based state-of-the-art , when finetuning from the pretrained BERT model with the same maximum sequence length . Additionally , BERT-AL is a general NLP model which has no specific setting for text summarization , so it can be easily adapted to other tasks with long text , e.g. , document-level machine reading comprehension and long text classification . Furthermore , the method , applying multi-channel LSTM on hidden states from transformers , also can be used in other Transformer based pretrained models , e.g. , XLNet and RoBERTa . In summary , contributions of this paper are shown as follow . 1 ) We propose a new architecture that combines Transformer and LSTM , which resolve the problem Transformer can not be used in very long text , and skillfully avoid LSTMs drop backs . 2 ) We propose multi-channel LSTM only applied on the corresponding position across different Transformers , which can take fully advantage of LSTM without hurting Transformers representation too much . 3 ) We conduct experiments to prove BERT-AL can outperform other models with the BERT pretrained under the same maximum sequence length , and can perform very close to BERTSUM with at least twice the maximum length than ours . 2 BACKGROUND . 2.1 BERT . According to the original implementation described in Vaswani et al . ( 2017 ) , BERT is a multi-layer bidirectional Transformer encoder , and Multi-Head Self-Attention is the key structure of Transformer encoder . For the input H ∈ RL×D of each Transformer layer , H will be mapped to three different spaces , named as Q , K and V , respectively . Self-attention computes the dot products between Q and K to roughly get the weight matrix , and then multiply V to get the hidden representation . Multi-head mechanism promotes the power of Transformer because it allows the model to jointly attend to information from different representation subspaces at different positions . BERT employs two tasks for pretraining : Masked Language Model and Next Sentence Prediction . Masked Language Model : BERT proposes a bidirected language model , which replaces 15 % words in text by [ MASK ] label and then predicts which words they are . To mitigate the mismatch between pretraining and finetuning , i.e. , the [ MASK ] token does not appear during finetuning , they do not always replace masked words with the actual [ MASK ] token . Then , BERT will predict the original token with cross entropy loss . Next Sentence Prediction : BERT is also pretrained by the next sentence prediction task . Specifically , when choosing the sentences A and B for each pretraining example , 50 % of the time B is the actual next sentence that follows A ( labeled as IsNext ) , and 50 % of the time it is a random sentence from the corpus ( labeled as NotNext ) . Then , BERT will predict whether the sentence B is the actual next sentence that follows A or not with cross entropy loss . For most of all NLP tasks , BERT concatenates different parts of input into a sequence beginning with [ CLS ] token and inserts [ SEP ] token between two different parts . Before going through Transformer layers , BERT merges three different embeddings into one , i.e. , word embedding , position embedding and segment embedding . Thereinto , position embedding is the information about the relative or absolute position of the tokens in the sequence . Since the model contains no recurrence and no convolution , position embedding is added to make use of the order of the sequence . There are two implements of position embedding , i.e. , learned position embedding and sine and cosine functions of different frequencies , but both of them can not expand to longer without a performance regression ( Wang et al. , 2019a ) . 2.2 BERTSUM . BERTSUM is an extension of BERT on extractive text summarization task and it truncates only the first 512 tokens as input . To select sentences , BERTSUM adds [ CLS ] to the head of each sentence and [ SEP ] to the tail of each sentence indicating the end of that sentence . The following summarization layer will score each [ CLS ] which presents the importance of that sentence . Finally , sentences with top 3 highest scores compose the summary . There are three summarization layers proposed in Liu ( 2019 ) : 1 ) Simple Classifier : only adds a linear layer on each [ CLS ] and uses a sigmoid function to get the predicted score : Ŷi = σ ( WoTi + bo ) ( 1 ) where Ti is the logit of ith sentence , σ is the Sigmoid function . 2 ) Inter-Sentence Transformer : applies more Transformer layers into sentence-level representations as follows . h̃l = LN ( hl−1 +MHAtt ( hl−1 ) ) hl = LN ( h̃l + FFN ( h̃l ) ) Ŷi = σ ( Wohi + bo ) ( 2 ) where h0 = PosEmb ( T ) and T are the sentence vectors output by BERT , PosEmb is the function of adding position embeddings to T ; LN is the layer normalization operation ; FFN is a feedforward network ; MHAtt is the multi-head attention operation ; the superscript l indicates the depth of the stacked layer . The final output layer is still a sigmoid classifier and Ti is the logit of ith sentence . 3 ) Recurrent Neural Network : applies an LSTM layer over the BERT outputs to learn summarization-specific features . At time step i , the input to the LSTM layer is the BERT output Ti , and the output is calculated as : ( Fi Ii Oi Gi ) = LNh ( Whhi−1 + LNx ( WxTi ) ) Ci = σ ( Fi ⊙ Ci−1 ) + σ ( Ii ) ⊙ tanh ( Gi−1 ) hi = σ ( Ot ) ⊙ tanh ( LNc ( Ct ) ) Ŷi = σ ( Wohi + bo ) ( 3 ) where Fi , Ii , Oi are forget gates , input gates , output gates ; Gi is the hidden vector and Ci is the memory vector ; hi is the output vector ; LNh , LNx , LNc are there difference layer normalization operations ; The final output layer is also a sigmoid classifier and Ti is the logit of ith sentence . However , we argue that the way of truncation will result in the loss of information in the latter part of the document . If the key sentences locate at the end of the document , they will never be recalled even though the model is powerful . In contrast , if model can take all the tokens of the document as input , it will get these key sentences and produce better result . Therefore , we propose BERT-AL to solve this problem . For simplicity , we only employ the Simple Classifier as the summarization layer in the following sections . 3 BERT-AL . In this section , we will introduce the detail about BERT-AL , and as illustrated in Figure 1 , BERTAL has mainly two key components different with BERTSUM : one is the multi-channel LSTM and another is the positional encoding . 3.1 MULTI-CHANNEL LSTM . We first give the definition of our scenario : assume that the length of the document is ldoc , and we have only a pretrained BERT model with the maximum sequence length lBERT . We first split the document into nsegment segments and each segment has the length lBERT ( the length of nsegmentth segment is shorter than lBERT ) . We set lBERT−AL = nsegment ∗ lBERT as the maximum sequence length of BERT-AL , and then we have lBERT−AL ≥ ldoc , which means BERT-AL can take arbitrarily long document as input . As we know , LSTM has the capability of capturing information across arbitrarily long steps but is weak in capturing long-term dependencies . Therefore , within a segment , we use Transformer to capture long-term dependencies and extract the feature via self-attention among local positions , and then fully represent the segment at each layer . To this end , we propose a multi-channel LSTM to chain the representation of different segments . For the above definition , ldoc can be arbitrarily long and so is nsegment . Thus , we apply LSTM on nsegment dimension , and the steps of LSTM should also be nsegment . As Figure 1 shows , LSTM is following the segment-wise Transformers at each layer and takes the hidden states from Transformers as its input . To fit the next layer , we set LSTM ’ s hidden size also as dmodel , i.e. , 768 for BERT-Base and 1024 for BERT-Large . Thus , the computing in each layer is shown as follows . Hi−1 = Concat ( Hi−1,1 , Hi−1,2 , ... , Hi−1 , nsegment ) HTransi , j = Transformer ( Hi−1 , j ∈ R [ lBERT , dmodel ] ) Hi = LSTM ( H Trans i ∈ R [ lBERT , nsegment , dmodel ] ) ( 4 ) where Hi is the hidden states from ith layer , Hi , j is jth segment in Hi , and HTransi , j is the hidden states from jth Transformer . All segment-wise Transformers at the same layer share parameters . BERT-AL employs a single layer of unidirectional LSTM as the implement of multi-channel LSTM , and a channel is corresponding to a relative position within a segment . Per-gate layer normalization is applied in each LSTM cell ( Ba et al. , 2016 ) . We do not split the segments according to natural sentences , so [ CLS ] may be at the beginning , in the middle or at the end of a block . Similarly , the output at [ CLS ] positions also do not come from the same LSTM channel . The reason we apply multi-channel LSTM on nsegment rather than a 2D LSTM on the whole sequence is taking advantage of LSTM ’ s expandability on variable-length segments , but meanwhile , reducing the interference to self-attention among positions within a segment .

Summary: This paper proposed another variant of BERT, called BERT-AL, which can deal with arbitrarily long inputs. The authors constructed the proposed method by combining the segment-wise BERT with the multi-channel LSTM. The authors validated the proposed method on the text summarization task and achieved higher performance than existing works. Their proposed method can be combined with other transformer-based approaches, such as XLNet or RoBERTa.