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: The paper proposes a methodology to overcome the problem of processing long sequences with a pre-trained Transformer model, which suffers from high computational costs due to the complexity being quadratic in the length of the sequence. The authors also point out that BERT needs to be retrained from scratch if sequences longer than the specified maximum length (512) are to be processed. Their method (BERT-AL) chunks the input text into segments of maximum length. Each segment is propagated through several layers of a pre-trained Transformer layer followed by a multi-channel LSTM. Herein, the positional embeddings in each segment are the same. The model is applied to extractive summarization, and directly compared to the BERTSUM model, which can only process documents up to length 512. The comparison is made for 4 application scenarios, each corresponding to an artificial maximum length after which the BERTSUM model truncates the input: after 8, 16, 128, and 256 tokens. The experimental results suggest that BERT-AL outperforms BERTSUM in all 4 scenarios: Substantially for max length 8 and 16, and marginally for 128 and 256. BERTSUM without any truncation still performs best, however.