CoLo: A Contrastive Learning based Re-ranking Framework for One-Stage Summarization

Two-stage paradigms Zhong et al. (2020); Liu and Liu (2021) improve summarization quality by re-ranking and selecting a candidate from a given set of candidates. MatchSum Zhong et al. (2020) forms a contrastive learning based re-ranking framework where they first generate a set of candidates summaries by a extractive summarization model and then feed them to a re-ranker. The re-ranker is trained to optimize a summary-level score and it can evaluate the candidate summaries holistically. SimCLS Liu and Liu (2021) is the abstractive version which replaces the extractive summarizer in Zhong et al. (2020) with a abstractive summarizer.

The training objective for summarization models is to estimate a conditional probability distribution $p(Y|X)$, where $X$ is the document and $Y$ is the output summary. Given a summarization model $\mathcal{M}$ that has already tuned under the conventional framework with loss function $\mathcal{L}_{sum}$ where $\mathcal{L}_{sum}$ could be binary cross entropy loss (BCELoss) or negative log likelihood loss (NLLLoss). The two-stage systems should first use a sampling algorithm e.g. beam search to sample a candidate set $\mathcal{C}=\{C_{1},C_{2},\ldots,C_{m}\}$ of size $m$ from the fixed model distribution $C_{i}\sim p_{\mathcal{M}}(Y|X)$. Candidates in $\mathcal{C}$ are sort by their ROUGE score in descending order. Then the they further train a separate re-ranker,e.g., BERT , with a contrastive-style ranking loss $\mathcal{L}_{rank}$ to select the the best candidate from $\mathcal{C}$ as the final output. The ranking loss used in the best re-ranking system for summarization is the triplet margin loss Kingma and Ba (2014). For a candidate pair $(C_{i},C_{j})$ where $i<j$, if $C_{i}$ has higher ROUGE score and it will be treated as the positive sample:

where $\mathbf{z}_{X},\,\mathbf{z}_{C_{i}},\,\mathbf{z}_{C_{j}}$ are the vector feature representation of $X,C_{i},C_{j}$ output by the re-ranker, and $\rho$ is the margin value. The final ranking loss is obtained by summing up all pairs: $\mathcal{L}_{rank}=\sum_{j}\sum_{i<j}\mathcal{L}_{i,j}$. The ranking loss ensures that candidates with higher ROUGE score is closer to the document in the embedding space.

Figure 1(a) illustrates the difference between the architecture of two-stage systems and CoLo. Although MatchSum and SimCLS significantly outperform all one-stage models, they mainly suffer from three drawbacks which strongly emphasize the necessity of designing an one-stage model:

(1) Training/inference inefficiency. Building the training set of the re-ranker and the second training stage consumes large amounts of GPU and CPU time (see details in Section 5.3). Moreover, the need of re-feeding generation results to another module also requires unaffordable computational resources.

(2) Coupling between the summarizer and re-ranker. Each improvement to one of these modules requires simultaneous updating or retraining of another module, which limits the use of such systems in the real world. For example, to try a larger candidate set or a different decoding method, we have to prepare the training set again for the second stage. In addition, how to tune the hyperparameters to be optimal in both modules at the same time is another tricky issue. Compared with two-stage systems, our one-stage system has a simple and clean implementation.

(3) Two-stage systems also face difficulties in long document summarization, because the input length of the re-ranker will drastically increase as the length of candidates increasing (see detailed analysis in Appendix A). Correspondingly, CoLo is not easily affected by length variance.

The goal of one-stage re-ranking systems is to enable both training and inference to score candidate summaries holistically without requiring a second stage of computation by a separate model. Ideally, an one-stage summarization model should both function as a summarizer and a re-ranker. A straightforward solution is multi-task learning. The naive training pipeline can be formulated as follows: (i) tuning $\mathcal{M}$ with $\mathcal{L}_{sum}$. (ii) Getting positive and negative samples from $p_{\mathcal{M}}(Y|X)$ via offline sampling for each datapoint $X$ in the training set. (iii) Building the ranking loss with these candidates and further tuning $\mathcal{M}$ with $\mathcal{L}_{rank}+\mathcal{L}_{sum}$. However, in practice, such training method is always suboptimal compared to the state-of-the-art two-stage models. We denote the model after multi-task learning as $\mathcal{M^{\prime}}$. There is a serious generalization error in the naive methods: via multi-task learning, $\mathcal{M}^{\prime}$ is only able to rank candidates drawn from the original model distribution $p_{\mathcal{M}}(Y|X)$ but not candidates from the new distribution $p_{\mathcal{M^{\prime}}}(Y|X)$. This error makes the naive approach unable to directly output a good summary in sequence-level generated by itself.

The first step of CoLo is also to train the summarization model with $\mathcal{L}_{sum}$ like the naive approach.

In CoLo, we discard using positive-negative samples that from a fixed model distribution, instead, we sample these candidates from a constantly shifting model distribution during multi-task learning. By doing so, we can mitigate the above mentioned generalization error as much as possible because candidates are dynamically changing with the parameters of the model distribution $p_{\mathcal{M}}(Y|X)$ updated by gradient descent. To implement this process, at each training step, we sample the newest candidates along with their feature presentations from the summarization model and calculate the ranking loss. We will give a detailed description about how we performing the online sampling process on mainstream extractive and abstractive summarization models in the following parts.

We conduct experiments on five mainstream datasets to evaluate the effectiveness of our approach.
CNN/DailyMail Hermann et al. (2015) is a classic benchmark which contains articles from the CNN/Daily Mail newspapers. We use the cased version from datasets³³3https://github.com/huggingface/datasets
XSum Narayan et al. (2018a) is a one-sentence summary dataset from BBC News. Gold summaries are professionally written by the authors of documents.
Reddit Kim et al. (2019) is collected from social media platform and we use the TIFU-long version.
PubMed Cohan et al. (2018) is a long document summarization dataset from scientific domain whose $avg$ summary length is about 4 times longer than CNN/DM.
SSN An et al. (2021a) consists of papers mainly from math, physics and computer science with the abstract section as gold reference.

For the simplity of experimental settings, both extractive model and abstractive mode are based on BART. We use the encoder of BART (170M) as the backbone and a 3-layer MLP as the classifier to implement the extractor. We add two special token ‘<cls>’ to generate the sentence representation and ‘<sep>’ as sentence separator. ‘<doc>’ token is used to generate the document feature representation. candidate size for each dataset can be found in 1 We use adam optimizer Kingma and Ba (2014) learning rate schedule follows the setting in transformer Vaswani et al. (2017). We train our model for 15000 steps with BCELoss and 32000 steps with BCELoss and RankingLoss where each step has a batch size of 36. The margin parameter $\gamma$ is set to 0.01. The size of generated candidates $|\mathcal{C}|$ is set to 20 for CNN/DM. We report the results. Other settings follow the default setting in Liu and Lapata (2019). Our model is trained on single GeForce RTX 3090 GPU for 8 hours. Both our abstractive model and extractive model are trained on 24G GeForce RTX 3090 GPUs and the inference process is on 12G GeForce GTX TITAN XP GPUs.

For abstractive model, we choose BART initialized with facebook/bart-large-cnn from transformers⁴⁴4https://github.com/huggingface/transformers as the basic summarizer. We further fintune this model by NLLLoss and RankingLoss for 15000 steps where each step with a batch size of 8. Other setting is the same with our extractive version. To encourage diversity, we use the diverse beam search Vijayakumar et al. (2016) to generate the candidates with beam size set to 16 and diversity penalty set to 1.0. Our model is trained on 8 GeForce RTX 3090 GPUs for about 18 hours.

We examine our approach with 4 metrics that measure the distance between generated summaries against the gold reference. ROUGE Lin (2004) where R-1 and R-2 measure informativeness based on n-gram overlapping and R-L represents fluency. JS-2 Divergence Louis and Nenkova (2013) measures Jensen-Shannon divergence between the bigram distributions of two input texts. BERTScore Zhang et al. (2019) measures soft overlap between BERT embeddings of two texts instead of using lexical matching methods. MoverScore Zhao et al. (2019) is also based on the neural model but applies a earth mover distance measure to contextualized BERT embeddings.

We compare our models with baseline models which has similar amount of parameters and decoding speed of our models in this section. Our extractive results on CNN/DM are shown in Table 2 We compare our model with previous strong extractive baseline built on pre-trained model Zhong et al. (2019); Bae et al. (2019); Liu and Lapata (2019) and strong multi-stage systems Zhong et al. (2020). From the third section of Table 2, we can see that our model CoLo_Ext beats the baseline model by 1.49 ROUGE-1 score and achieve the state-of-the-art among all end-to-end systems when input length set to 512 and the results can be further improved while extending the input length to 1024. Even compared with the BertSum-large (340M) Liu and Lapata (2019) which is built on large PTM, We still have an improvement of 0.42 with only the half number of parameters of theirs. Though RL-based methods hold the motivation of optimizing towards the evaluation metric, but it does not gain much improvement on performance in practice.

To verify whether our model is effective on datasets of various lengths, we also evaluate our model on datasets with short summaries (Reddit and XSum) and long document dataset PubMed and results are shown in Table 3. On reddit and XSum, we achieve the advantage of more than 1.0 point ROUGE-1 than baseline systems and close performance with the upper bound ORACLE. We also gain improvements when tested on the long document summarzation dataset PubMed. Detailed results on long document dataset can be found in Appendix A.

Early work also successfully applies reinforcement learning on abstractive summarization Paulus et al. (2017); Li et al. (2019). But we do not find related works that successfully combine reinforcement learning with strong pre-trained models. Therefore, most of our baselines are strong pertrained model finetuned with NLLLoss. Our results is shown in Table 4, due to the huge cost of using large pre-trained model with length set to 1024, we also report results with 512 input tokens and it is able to significantly outperform other baselines which has longer input length (1024). CoLo_Abs has an improvement of 2.17 R-1 socre on the very strong BART-large baseline without adding additional parameters or modules. Additionally, our method is able to outperform all one-stage baseline systems by a large margin. We also conduct experiments on long document summarzation datasets (see in Table 11 in Appendix).

Apart from the one-stage systems, we also compare our model with these powerful multi-stage systems: CTRLSum, multi-stage re-ranking models. CTRLSum needs other systems to previously produce a control signal.

In addition to ROUGE, we also select other metrics as the discriminator (shown in Table 7). ROUGE and JS-2 is based on lexical matching while BERTScore and MoverScore are based on the contextualized embedding from BERT. Our model generally obtains the best results on the metric used in training. Because these metrics are not actually separated, using one of these metrics as the discriminator can also gain significant improvements on other metrics. Overall, the neural evaluation metric BERTScore and MoverScore bring more improvements compared with metrics that based on the lexical matching. But incorporating neural model based metrics in training will obviously increase the training time.

We conduct a visualization experiment on our extractive model to get a close look on the distribution of candidates in semantic space. We randomly sample 100 documents with more than 10 sentences from the test set of CNN/DM. We first select the top 10 sentences based on the predicted score from the classifier. We set the possible number of sentences to $\{2,3\}$ resulting a candidate size of $C_{10}^{2}+C_{10}^{3}=165$ for each sample. We visualize the learned embedding of these candidates and the anchor in a two-dimensional space by applying the t-SNE algorithm. As shown in Figure 5, there is an obvious cluster of the top 50 candidates (colored in purple) and the candidates with higher score are closer to the anchor while the distribution of uninformative candidates (gray,cyan points) is relatively random.

We also conduct a human evaluation on our models to get more accurate results . We randomly select 30 articles from the test set of CNN/DM, and each articles have 5 candidate summaries 4 from automatic systems and 1 is the gold reference. We recruit 2 PhD students majoring in computer science and ask them to rank the candidate summries based on the fluency, informativeness. If two of these systems generate the same summary for the source document, this sample will be filtered out. As we can see from Table 8, the CoLo_Ext with the discriminator as BERTScore achieve the best result among all automatic systems. However, using BERTScore will bring much training time. We also suggest taking JS-2 divergence as the discriminator which also does a good job in human evaluation.