Attention Temperature Matters in Abstractive Summarization Distillation (2022)

Abstract: Recent progress of abstractive text summarization largely relies on large pre-trained sequence-to-sequence Transformer models, which are computationally expensive. This paper aims to distill these large models into smaller ones for faster inference and with minimal performance loss. Pseudo-labeling based methods are popular in sequence-to-sequence model distillation. In this paper, we find simply manipulating attention temperatures in Transformers can make pseudo labels easier to learn for student models. Our experiments on three summarization datasets show our proposed method consistently improves vanilla pseudo-labeling based methods. Further empirical analysis shows that both pseudo labels and summaries produced by our students are shorter and more abstractive.

## introduction automatic document summarization is the task of rewriting a long document into its shorter form while still retaining its most important content. in the literature, there are mainly two kinds of methods for summarization: extractive summarization and abstractive summarization @xcite . in this work, we focus on abstractive summarization, which is viewed as a sequence-tosequence (seq2seq) learning problem, since recent abstractive models outperform their extractive counterparts and can produce more concise summaries @xcite @xcite . recent progress of abstractive summarization largely relies on large pre-trained transformer models @xcite @xcite @xcite . with these extremely large models, we can obtain state-of-theart summarization results, but they are slow for online inference, which makes them difficult to be used in the production environment even with cutting-edge hardware. this paper aims to distill these large transformer summarization models into smaller ones with minimal loss in performance. knowledge distillation is a class of methods that leverage the output of a (large) teacher model to guide the training of a (small) student model. in classification tasks, it is typically done by minimizing the distance between the teacher and student predictions @xcite . as to seq2seq models, an effective distillation method is called pseudo-labeling @xcite , where the teacher model generates pseudo summaries for all documents in the training set and the resulting document-pseudo-summary pairs are used to train the student model. in this paper, we argue that attention distributions of a seq2seq teacher model might be too sharp. as a result, pseudo labels generated from it are sub-optimal for student models. in the summarization task, we observe that 1) pseudo summaries generated from our teacher model copy more continuous text spans from original documents than reference summaries (56% 4-grams in pseudo summaries and 15% 4-grams in reference summaries are copied from their original documents on cnn/dailymail dataset); 2) pseudo summaries tend to summarize the leading part of a document (measured on cnn/dailymail, 74% of sentences in pseudo summaries and 64% of sentences in reference summaries are from the leading 40% sentences in original documents). we obtain the two numbers above by matching each sentence in a summary with the sentence in its original document that can produce maximum rouge @xcite score between them. we call the two biases above the copy bias and the leading bias. in order to have an intuitive feeling, we select a rep-resentative example foot_0 and visualize its cross attention weights foot_1 (see the left graph in figure 1 ). we observe that attention weights form three "lines", which indicates very time the decoder predicts the next word, its attention points to the next word in the input document. that may be the reason why multiple continuous spans of text are copied. another phenomenon we observe is that all high-value attention weights (in deeper color) concentrate on the first 200 words in the input document, which reflects the leading bias. in either case, the attention distribution is too sharp (i.e., attention weights of the next word position or the leading part is much larger than other positions), which means our teacher model is over-confident. based on the observations above, we propose a simple method called plate (as shorthand for pseudo-labeling with larger attention temperature) to smooth attention distributions of teacher models. specifically, we re-scale attention weights in all attention modules with a higher temperature, which leads to softer attention distributions. figure 1 intuitively shows the effect of using higher attention temperatures. compared with the left graph, the right graph with higher attention temperature has shorter lines (less copy bias) with high attention weights, and positions of high attention weights extend to the first 450 words (less leading bias). less copy bias in pseudo summaries encourages student models to be more abstractive, while less leading bias in pseudo summaries encourages student models to take advantage of longer context in documents. experiments on cnn/dailymail, xsum, and new york times datasets with student models of different sizes show plate consistently outperforms vanilla pseudo-labeling methods. further empirical analysis shows that, with plate, both pseudo summaries generated by teacher models and summaries generated by student models are shorter and more abstractive, which matches the goal of abstractive summarization. ## related work large pre-trained seq2seq transformer models largely improve results of generation tasks including text summarization @xcite @xcite ; token index in summary token index in document token index in summary attention temperature: 64 attention temperature: 96 zhang et al., 2020). these models are pre-trained using unsupervised text-to-text objectives. for example, t5 @xcite ) is pre-trained by predicting corrupted text spans. bart @xcite employs denoising auto-encoding objectives such as text infilling and sentence permutation during its pre-training. the pre-training objective of pegasus @xcite is tailored for the summarization task, which predicts the most "summary worthy" sentences in a document. our method aims to make these large models faster. in knowledge distillation, besides learning from gold labels in the training set, student models can learn from soft targets @xcite , intermediate hidden states @xcite , attentions @xcite , and target output derivatives @xcite of teacher models. recent work for distillation of pre-trained transformers (e.g., distilbert @xcite , tinybert @xcite , mobile-bert @xcite , bert-of-theseus @xcite , minilm @xcite ) focuses on natural language understanding tasks such as glue @xcite or squad @xcite benchmarks. most methods above are designed for classification models. in seq2seq learning tasks such as summarization, we can apply distillation methods above to each step of sequence model predictions. however, the sequence-level knowledge of teacher mod-els is not well utilized. therefore, @xcite introduce a sequence-level knowledge distillation method (i.e., pseudo-labeling), where a student model is trained with pseudo labels generated by the teacher model using beam search decoding. @xcite and later work @xcite @xcite show pseudo-labeling achieves competitive performance for seq2seq tasks such as machine translation. @xcite propose the shrink and fine-tune (sft) approach for pre-trained summarization distillation, which re-finetunes a teacher model with some layers removed, and they show sft outperforms pseudo-labeling and a modification of direct knowledge distillation @xcite ) on one of their datasets, but not others. our method, which builds on top of pseudo-labeling, is conceptually simple and improves pseudo-labeling across different summarization datasets. there is an interesting line of work called selfdistillation or self-training @xcite @xcite @xcite , where the size of the student model is identical to the size of the teacher model. our method can also be applied in selfdistillation and can potentially be combined with the self-distillation methods above. ## conclusions in this work, we propose a simple but effective extension of pseudo-labeling method plate for summarization distillation. experiments on three datasets demonstrate that our method can consistently outperform the vanilla pseudo-labeling method. further empirical analysis shows that by using our method, teacher models can generate more concise and abstractive summaries. as a result, summaries produced by student models also become more concise and abstractive. in the future, we would like to explore our method to other generation tasks as well as self-training with unlabeled data. we are also interested in combining our method with other distillation methods and extending our method for better teacher model training. model λ bleu transformer-big (teacher) -28.51 student transformer-base 1.0 27.79 transformer-base 1.5 27.90 transformer-base 2.0 27.85