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 tackles the problem of suboptimality of beam search as it is not able to have a global perspective of the probability of the generated beams due to truncated beam size (less than vocabulary size) for space and time complexity purposes. This paper proposes a global-aware beam search which is calibrated using the global attention distribution to guide the local decisions at every time step during beam search. The global attention distribution prediction is formulated as a regression task, which is then fed into a global scoring function to guide the beam search-based generation. This approach solves the mismatch between the local optimality of beam search with the global optimal generation. This modified search strategy has been shown to boost the summarization of two SOTA models BART and PEGASUS.