Global-aware Beam Search for Neural Abstractive Summarization

1 Introduction . As the barriers exist from the auto-regressive design of neural probabilistic text generators to predicting the global optimum directly [ 30 ] , the heuristic algorithm beam search that factorizes global optimization to multiple local optimizations , has been popularly used for text decoding [ 24 ] . In the original beam search setting , the global optimum is a hypothesis g of the highest probability among all possible sentences , and consists of words in vocabulary V . Given the global optimum at step t denoted as g≤t , the local optimums l≤t refer to K candidate sequences with the highest probabilities at each step . While it is necessary to compromise on the beam size K V to ensure text quality [ 3 , 25 , 24 ] and search efficiency , beam search suffers from a major limitation due to its local property . Concretely , assuming that the global optimal hypothesis is within the K local optimal hypotheses of the highest probabilities , i.e . p ( g≤t ) ≥ p ( l≤t ) , for all t until the termination T , it operates solely with the local information available at each step . In practice , such assumption may however fail in the case that the probability of the global optimum at step τ < T is less than those of the local optimums , i.e . p ( g≤τ ) < p ( l≤τ ) , but is adjusted to a higher level in the later steps , p ( g > τ |g≤τ ) > p ( l > τ |l≤τ ) and p ( g≤τ ) p ( g > τ |g≤τ ) > p ( l≤τ ) p ( l > τ |l≤τ ) . This often leads beam search to get stuck in the local optimum from step τ onward in generating texts . ∗corresponding author 2https : //github.com/yema2018/global_aware 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . To cope with this limitation , this study proposes a calibrated beam-based algorithm with global awareness at all searching steps . Generally , our novel algorithm is implemented in two phases . Before the beam search ( Phase I ) , the global attention distribution is predicted in order to be included as a protocol to calibrate beam search at each step , encouraging the generated hypotheses to attend to the source in a more near-global optimal way . Specifically , the global attention distribution describes how all reference tokens should assign the attention to each source token ( illustrated in Figure 1.a ) , which could be predicted from the source by training an attention-prediction model . The training is fairly straightforward and resembles a sequence tagging task [ 12 ] , except that the predicted attention distribution from the source is a regression result . There are several advantages of using the attention distribution as the global protocol . 1 ) Attention distributions are sensitive to the decoder input , suggesting that any input to the decoder leads to a unique attention distribution with fixed model parameters ; 2 ) attention distributions are accessible for almost all text generation tasks thanks to the recent advances in attention models [ 27 , 4 , 31 ] ; 3 ) relying on the source only , the global attention distribution can be predicted before beam search , thus offering a rigorous mechanism to calibrate a global-aware beam search . During beam search ( Phase II ) , we develop a novel global scoring mechanism composed of attention scores and length rewards to guide beam search based on the predicted global attention distribution . As one main theoretical result , we show that the attention score can be considered as the probability that generated texts attend to sources in line with the predicted global attention distribution . Specifically , the generated tokens in each step update the local attention distribution to source tokens dynamically , where the attention values grow monotonically as the generation proceeds ( see Figure 1.b ) . Since the desired situation is that the local distribution reaches exactly the global distribution at the terminal step , we regulate the inference by discouraging local attention from exceeding their corresponding predicted global attention at all steps . With regards to the core target to investigate the possible paradigm that improves beam search with global awareness during decoding , contributions of this study are summarized as follows : • We argue that the limitation of beam search roots from its defect in finding the global optimal hypothesis . We improve the algorithm by proposing a global protocol to regulate beam search step-by-step . This paper is the first to predict and deploy the global attention distribution to calibrate the inference in a rigorous way , thus returning a hypothesis that attends to source tokens in a more near-global optimal manner . In contrast , previous works [ 37 , 7 , 20 , 29 , 21 ] try to use attention distributions to improve beam search , but ignore that the global attention distribution is predictable . • A novel global scoring mechanism is designed to evaluate the generated sequences at each step based on the desired situation described in Figure 1.c . As theoretically justified , its major component can be elegantly integrated into beam search in the form of a probability so that merely O ( K ) of the time complexity is increased in each step ( see Section § 3.1 for more details ) . • The proposed algorithm with global awareness manifests a robust and plug-and-play property in enhancing beam search for neural abstractive summarization . Without requiring any model or parameter modification , the global-aware inference shows excellent performance in generating meaningful texts , even if the attention distribution is corrupted or not of its own . Further , it is identified that summaries generated by global-aware inference are both higher-quality and different from beam search hypotheses ( see Global-aware in Table 1 ) . More interestingly , we find that the generation style of a dataset could be transferred by the designated global attention distribution . For instance , summaries of higher abstractness for CNN/DM could be generated by only replacing its global attention distribution with a highly abstractive distribution during inference , as presented in Global-aware† of Table 1 . • On the empirical side , we show that the proposed global-aware inference can stably and significantly boost two state-of-the-art summarization models BART [ 17 ] and PEGASUS [ 39 ] to produce higher quality summaries on nine datasets , even if only the empirical hyper-parameters are used . 2 Preliminary . The proposed decoding strategy is applied in BART [ 17 ] and PEGASUS [ 39 ] to perform summarization . BART is a pre-trained seq-to-seq model whose structure essentially follows a vanilla Transformer encoder-decoder [ 31 ] . PEGASUS has a similar structure but is pre-trained on a larger dataset differently . Notably , the fine-tuned parameters of both models are downloaded from HuggingFace Models3 and are fixed in all subsequent operations , where “ fine-tuned '' means the pre-trained model has been fine-tuned on a specific dataset . 2.1 Beam Search . Since the decoder of the seq-to-seq model is an auto-regressive model , the probability of a target sequence y = ( y0 , · · · , yt , · · · , yT ) can be factorized to the probabilities conditional on the source x = ( x1 , · · · , xi , · · · , xn ) and y < t = ( y0 , · · · , yt−1 ) , i.e. , p ( y|x ) = T∏ t=1 p ( yt|x , y < t ) , T ≥ 1 ( 1 ) When T = 0 , y = ( y0 ) and p ( y0|x ) = 1 because y0 is a fixed start token . Beam search [ 8 ] is a decoding strategy to predict a target sequence by maximizing this factorization . Given a vocabulary set V , at each inference step t , beam search selects a candidate beam set BKt = { bkt } Kk=1 ( where each beam bkt = ( b k 0 , · · · , bkt ) is a candidate sequence ) from an all-possible beam set Bt of size K × |V| , namely , Bt = { bkt−1 ◦ v | bkt−1 ∈ BKt−1 , v ∈ V } ( 2 ) 3https : //huggingface.co/models BKt = { bkt | bkt = argtopk ( log p ( bt|x ) ) , bt ∈ Bt } , t ≥ 1 ( 3 ) where argtopk ( · ) outputs K beams with the highest conditional probability , and ◦ is the concatenation operation . Besides , BK0 = { bk0 } Kk=1 where bk0 is the start token . By Eq . 1 , log p ( bt|x ) is an accumulated value . Its calculation can be simplified as : log p ( bt|x ) = { log p ( bkt−1|x ) + log p ( v|x , bkt−1 ) , t ≥ 2 log p ( v|x , bk0 ) , t = 1 ( 4 ) where the value of log p ( bkt−1|x ) is computed from the previous step . Therefore , at each step , we only need calculate the condition probability of each token in the vocabulary set . A beam is terminated after it generates an ending token , and the beam set of K terminated beams is defined as Y . The final hypothesis y∗ is chosen from Y based on the beam probability normalized by lengtha where a is a hyper-parameter of length [ 37 ] : y∗ = argmax yk∈Y log p ( yk | x ) ( |yk| − 1 ) a ( 5 ) where yk = ( yk0 , · · · , ykT ) . |yk| − 1 is used since the start token is not considered in calculating the length . 2.2 Attention Distribution . Attention distribution is a continuous vector whose element indicates the degree of attention paid to a source token . The element is formed by the accumulation of cross attention , i.e. , ∑ t αt , i , where αt , i refers to the cross attention that the tth token in the target sequence gives to the ith source token.4 Specially , cross attention is a scaled dot-product [ 31 ] of hidden states of the source x and the target sequence y . Notably , since Transformer-decoder is an auto-regressive model , the cross attention assigned by tth target token is actually calculated by y < t = ( y0 , · · · , yt−1 ) . Global Attention Distribution . The global attention distribution g = [ g1 , · · · , gi , · · · , gn ] ∈ Rn is the attention distribution given by the reference , where global attention gi refers to the total attention that the reference attends to the ith source token , and n is the source length . Optimal Length . The summation of g , namely ∑n i=1 gi , is equal to ∑n i=1 ∑T t=1 αt , i = T due to∑n i=1 αt , i = 1 in softmax operation , where T is the reference length , or equivalently the optimal length Z . Local Attention Distribution . The local attention distribution lkt = [ lkt,1 , · · · , lkt , i , · · · , lkt , n ] ∈ Rn is the attention distribution of the kth generated sequence and updated at each decoding step t. Thereinto , the local attention lkt , i denotes the total attention paid to the ith source token by the kth beam sequence ( bk1 , · · · , bkt ) and is dependent on the sequence generated before t , i.e. , bkt−1 = ( bk0 , · · · , bkt−1 ) .

Summary: This paper presents a method to improve beam search by learning and predicting the global attention. A global scoring function is developed to regulate beam search to generate summaries. In Phase I, the global attention distribution is predicted in order to be included as a protocol to calibrate beam search. In Phase II, a step-wise global scoring function is used to guide beam search based on the predicted global attention distribution.  Experiments on 9 datasets show that the global-aware inference improves state-of-the-art summarization models.