SlowGuess/NAACL_2022_OCR huggingface dataset:
Mayank Kulkarni
Bloomberg, USA
mkulkarni24@bloomberg.net
Debanjan Mahata*
Moody's Analytics, USA
Debanjan.Mahata@moodys.com
Ravneet Arora
rarora62@bloomberg.net
Rajarshi Bhowmik
rbhowmik6@bloomberg.net
In this work, we explore how to train task-specific language models aimed towards learning rich representation of keyphrases from text documents. We experiment with different masking strategies for pre-training transformer language models (LMs) in discriminative as well as generative settings. In the discriminative setting, we introduce a new pre-training objective - Keyphrase Boundary Infilling with Replacement (KBIR), showing large gains in performance (upto 8.16 points in F1) over SOTA, when the LM pre-trained using KBIR is fine-tuned for the task of keyphrase extraction. In the generative setting, we introduce a new pre-training setup for BART - KeyBART, that reproduces the keyphrases related to the input text in the CatSeq format, instead of the denoised original input. This also led to gains in performance (upto 4.33 points in F1@M) over SOTA for keyphrase generation. Additionally, we also fine-tune the pre-trained language models on named entity recognition (NER), question answering (QA), relation extraction (RE), abstractive summarization and achieve comparable performance with that of the SOTA, showing that learning rich representation of keyphrases is indeed beneficial for many other fundamental NLP tasks.
1 Introduction and Background
Keyphrases capture the most salient topics of a document and facilitates extreme summarization. Identifying them in an automated way from a text document can be useful for several downstream tasks - classification (Hulth and Megyesi, 2006), clustering (Hammouda et al., 2005), summarization (Qazvinian et al., 2010; Zhang et al., 2004), reviewer and document recommendation (Augenstein et al., 2017), and many different information retrieval tasks such as enabling semantic and faceted search (Sanyal et al., 2019; Gutwin et al.,
1999), query expansion (Song et al., 2006), and interactive document retrieval (Jones and Staveley, 1999).
Keyphrases could either be extractive (part of the document) or abstractive (not part of the document). Prior works have referred to them as present and absent keyphrases, respectively. Automatically identifying them entails the process of detecting the extractive (Hasan and Ng, 2014) and generating the abstractive keyphrases (Cano and Bojar, 2019a) from a given document. While extractive approaches have mostly dominated over the generative ones with higher accuracies (Cano and Bojar, 2019b), the task is far from solved and the performances of the present systems are worse in comparison to many other NLP tasks (Liu et al., 2010). Some of the major challenges are the varied length of the documents to be processed, their structural inconsistency and developing strategies that can perform well in different domains.
Most of the prior work on identifying keyphrases using deep learning techniques have concentrated on developing new architectures and frameworks based on different training paradigms such as seq2seq (Meng et al., 2017; Yuan et al., 2018; Zhang et al., 2017a; Chen et al., 2018; Ye and Wang, 2018; Chen et al., 2019; Ye et al., 2021), sequence tagging (Alzaidy et al., 2019), reinforcement learning (Chan et al., 2019), adversarial training (Swaminathan et al., 2020) and game theory (Saxena et al., 2020). Although, there has been tremendous progress in learning better semantic and syntactic representation of language at different levels - characters, words, phrases, sentences and documents (Liu et al., 2020b), there hasn't been any effort in learning rich pre-trained representations of keyphrases, which is the major focus of this work.
Transformer language models when pre-trained on large corpora with different pre-training objectives (Qiu et al., 2020) have shown great success
*This work was done by the author Debanjan Mahata as an employee of Bloomberg
Findings of the Association for Computational Linguistics: NAACL 2022, pages 891 - 906
in various downstream tasks on fine-tuning, including the tasks of keyphrase extraction (Sahrawat et al., 2019; Martinc et al., 2020; Santosh et al., 2020) and generation (Liu et al., 2020a). However, pre-training objectives tailored towards learning better representation of keyphrases that can result in improving the performance of identifying and generating keyphrases from text have not yet been explored. This motivated us to look into this specific problem and make an attempt to answer the following questions:
Q1 - Can we formulate a pre-training objective for language models that can learn better representation of keyphrases?
Previous work explored training language models for learning better representation of text spans (Joshi et al., 2020), summary sentences (Zhang et al., 2020), and tokens for named entity recognition (Yamada et al., 2020). To effectively learn rich representation of keyphrases in a BERT like discriminative setup, we propose a new pretraining objective - Keyphrase Boundary Infilling with Replacement (KBIR) (Section 2) which utilizes a multi-task learning setup for optimizing a combined loss of random token Masked Language Modeling (MLM) (Devlin et al., 2018), Keyphrase Boundary Infilling (KBI) (Section 2.1) and Keyphrase Replacement Classification (KRC) (Section 2.2).
We also propose a new setup for pre-training BART (Lewis et al., 2019) - KeyBART (Section 3), focused towards learning better representation of keyphrases in a generative setting. Instead of reproducing the denoised input text as proposed in the original setup, we produce the keyphrases associated with the input document in the CatSeq (Meng et al., 2017) format from a corrupted input.
Q2 - Does learning rich representation of keyphrases in a language model lead to performance gains for the tasks of keyphrase extraction and generation?
One of the key contributions of this work is the introduction of KBIR, which is the combination of the KBI and KRC objectives with MLM that helps to learn good representation of keyphrases. This is validated by obtaining SOTA performance for the task of keyphrase extraction on three benchmark datasets (Section 4.2.1), surpassing the existing SOTA (Duan et al., 2021) by at most 8.16 F1 points on the SemEval 2017 corpus (Augenstein et al., 2017).
We also evaluated the KeyBART approach across five benchmark datasets for the task of keyphrase generation and obtained SOTA performances for both present and absent keyphrases (Section 4.2.2). Our best model surpassed the SOTA ONE2SEQ model (Ye et al., 2021) by 4.33 F1@M points and 0.72 F1@M points on Inspec (Hulth, 2003a) for present and absent keyphrases respectively.
Q3 - Do rich keyphrase representations aid other fundamental tasks in NLP such as NER, QA, RE and summarization?
It is to be noted, that although we trained our models on a large corpus of 23 million scientific articles, we find that it performs reasonably well when fine-tuned on datasets that do not belong to the scientific domain for different NLP tasks as shown in Section 4.2.3, 4.2.4, 4.2.5 and 4.2.6. This also suggests that identifying keyphrases in the context of an input text is a fundamental NLP task and a language model trained to learn optimal representation of keyphrases can aid many other tasks.
To summarize the main contributions that we make in this work are:
- We make the first attempt to train task-specific language models in discriminative as well as generative settings geared towards learning rich representation of keyphrases from text.
- We introduce a novel pre-training objective Keyphrase Boundary Infilling with Replacement (KBIR) and train a new language model that achieves SOTA performance for the task of keyphrase extraction.
- We propose a new setup - KeyBART for pretraining a generative language model for learning better representation of keyphrases and achieve SOTA performance on the task of keyphrase generation.
- We also empirically show how learning rich keyphrase representations from text is also useful for other NLP tasks like NER, RE, QA and summarization by achieving near SOTA performances in all of them using our language models trained using KBIR objective and KeyBART settings.
We have made our models<sup>1</sup> <sup>2</sup> publicly avail
https://huggingface.co/bloomberg/ KeyBART
$^{2}$ https://huggingface.co/bloomberg/KBIR
able. We also make our pre-training code available. Next, we give a detailed description of the methods that we propose in this work.
2 Keyphrase Boundary Infilling with Replacement (KBIR)
In the previous section, we mentioned various methods that aim at learning representations of text spans. Unlike LMs like SpanBERT (Joshi et al., 2020) and PEGASUS (Zhang et al., 2020) whose primary objective is to learn representations of random or heuristically chosen spans of text, the intuition behind learning good keyphrase representation is to provide the LM the ability to learn spans as well as to identify important phrases (in this case keyphrases) in the context of an input text. This motivated us to devise a framework that can optimize both of these objectives. Towards this effort, we propose a new pre-training objective Keyphrase Boundary Infilling with Replacement (KBIR) which is composed of two individual tasks - Keyphrase Boundary Infilling (KBI) and Keyphrase Replacement Classification (KRC) jointly learnt in a multi-task learning setup as shown in Figure 1. We build our framework on top of RoBERTa which implements random token Masked Language Modeling (MLM), therefore making our LM essentially optimizing MLM along with KBI and KRC objectives. In the following section, we describe the individual components of our framework.
2.1 Keyphrase Boundary Infilling (KBI)
To effectively learn span representations of keyphrases, we propose a new pre-training objective that builds upon the Span Boundary Objective (SBO) from SpanBERT (Joshi et al., 2020) and the Text Infilling setup from BART (Lewis et al., 2019). Similar to BART, we replace the entire span, in this case a keyphrase, with a single [MASK] token as shown in Figure 2 and predict the original tokens using positional embeddings in conjunction with boundary tokens. Text Infilling is a more challenging task than SpanBERT's objective of individual masked token predictions as the model must predict how many tokens correspond to a span (Lewis et al., 2019). Different from SpanBERT, which does not penalize incorrect predictions of a sequence of tokens within a masked span, we propose a cumulative loss (Equation 2) across
all tokens in the masked span to capture intra-span token relationships to learn better span representations. Text infilling, to the best of our knowledge, has not been explored in a discriminative setup as done in this work.
We denote the output of the transformer encoder for each token  $x_{l}$  in the sequence  $x_{1},\ldots ,x_{L}$  as  $\mathbf{x}_{1}$ . However, since the entire span of tokens  $(x_{s},\dots,x_{e})$  of a keyphrase  $y_{m}$  is masked with a mask token  $x_{m}$ , it is represented with a single vector  $\mathbf{x_m}$ , where  $(s,e)$  indicates its start and end positions and  $m$  represents the index of a masked keyphrase span. We set a maximum possible number of tokens corresponding to a keyphrase span,  $T_{max}$  such that  $i\in [1,T_{max}]$ . We then predict the sequence of tokens to replace  $x_{m}$  using the output encodings of the external boundary tokens  $x_{s - 1}$  and  $x_{e + 1}$ , as well as the position embedding  $\mathbf{p_i}$  of the target token as shown in Equation 1.
$$
\mathbf {y _ {i}} = f (\mathbf {x _ {s - 1}}, \mathbf {x _ {e + 1}}, \mathbf {p _ {i}}) \qquad (1)
$$
where positional embeddings use relative positions of the masked tokens with respect to the left boundary token  $x_{s-1}$ . We use Layer Normalization (Ba et al., 2016) and GeLU (Hendrycks and Gimpel, 2016) activation function to represent  $f(\Delta)$ . We then use the vector representation  $\mathbf{y_i}$  to predict the potential token  $x_i$  and compute the cumulative cross-entropy loss for each  $i$  present within the unmasked  $x_m$  as shown in Equation 2.
$$
\mathcal {L} _ {\text {I n f i l l}} (\boldsymbol {\theta}) = \sum_ {i = 1} ^ {T _ {\max }} \log p \left(x _ {i} \mid \mathbf {y} _ {\mathbf {i}}\right) \tag {2}
$$
In addition to predicting the actual tokens, we use a classification head to predict the expected number of tokens corresponding to the [MASK] in anticipation of providing a stronger learning signal. Each possible length of the [MASK] is represented as a class and therefore, the number of such classes is equal to the maximum number of possible tokens  $(T_{max})$ . The architecture used for classifying the number of tokens is a single linear layer which is trained with cross-entropy loss  $\mathcal{L}_{\mathrm{LP}}(x_m,z_m)$  along with the inflated masked token  $x_{m}$  and the corresponding actual length of the span class  $z_{m}$ .
The Keyphrase Boundary Infilling (KBI) objective is formally represented as:
$$
\mathcal {L} _ {\mathrm {K B I}} (\boldsymbol {\theta}) = \alpha \mathcal {L} _ {\mathrm {M L M}} (\boldsymbol {\theta}) + \gamma \mathcal {L} _ {\text {I n f i l l}} (\boldsymbol {\theta}) + \sigma \mathcal {L} _ {\mathrm {L P}} (\boldsymbol {\theta}) \tag {3}
$$
3https://github.com/bloomberg/kbir keybart
Figure 1: The KBIR model architecture for the training phase. The Random Token MASK is denoted in red, the Keyphrase MASK is denoted in blue and the Replaced Keyphrase in orange.
where  $\alpha, \gamma$  and  $\sigma$  are co-efficients applied to each loss and are primarily used to normalize the losses across the tasks.
We propose this pre-training objective to be used with keyphrases, however the objective is generic enough to be applied to any spans of text, these could be keyphrases, entities or even random spans.
2.2 Keyphrase Replacement Classification (KRC)
Apart from learning representations of keyphrase spans, we wanted our framework to have the ability to identify them within the context of a text input. Motivated by WKLM (Xiong et al., 2019) that explores pretraining a language model through weak supervision by replacing entities with random entities of the same type that belongs to a knowledge base, we adapt it to replace keyphrases by randomly choosing another keyphrase of variable length from the universe of keyphrases identified in a tagged corpus. The KRC task is then modeled as a binary classification task to determine whether a keyphrase is replaced or retained.
To implement this strategy, we construct a keyphrase universe by identifying the set of unique keyphrases tagged across the entire dataset. We then randomly shuffle this keyphrase universe and restrict it to 500,000 keyphrases for computational complexity. We use the concatenated representation of boundary tokens of a keyphrase  $x_{s-1}$  and  $x_{e+1}$  as input to a linear classifier as shown in Figure 1. Given the label  $y_k$  representing whether
a keyphrase was replaced or not, the objective here is to minimize the binary cross-entropy loss  $\mathcal{L}_{\mathrm{KRC}}((x_{s - 1} + x_{e + 1}),y_k)$
Finally, in order to train a LM with an objective of learning good keyphrase representations we use the KBIR pre-training strategy in which we jointly optimize the KBI loss and the KRC loss along with the already existing MLM loss  $(\mathcal{L}_{\mathrm{MLM}}(\pmb {\theta}))$  in RoBERTa. This is formally shown in Equation 4.
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {K B I R}} (\boldsymbol {\theta}) = \alpha \mathcal {L} _ {\mathrm {M L M}} (\boldsymbol {\theta}) + \gamma \mathcal {L} _ {\text {I n f i l l}} (\boldsymbol {\theta}) + \\ \sigma \mathcal {L} _ {\mathrm {L P}} (\boldsymbol {\theta}) + \delta \mathcal {L} _ {\mathrm {K R C}} (\boldsymbol {\theta}) \tag {4} \\ \end{array}
$$
3 KeyBART
We also explored learning a generative LM for the text generation tasks such as keyphrase generation and abstractive summarization. Our hypothesis behind the proposed setup is that masking and replacing task-specific spans, in this case keyphrases, that need to be re-generated should allow the generative model to learn a better representation of surrounding tokens and also the spans themselves.
BART (Lewis et al., 2019) generates sequences of different lengths from the input perturbed with [MASK] tokens along with token addition and deletion. On similar lines, we propose learning rich keyphrase representations by attempting to generate the Original Present Keyphrases in the Catseq format as proposed in (Meng et al., 2017) from an input perturbed with token masking, keyphrase masking, and keyphrase replacement as shown in
Figure 2. We call this method KeyBART. We maintain the order of occurrence of the keyphrases in the original document and remove duplicate occurrences. We also use the same method for finding keyphrase replacements as used in KRC (Section 2.2). We don't explicitly try to model the keyphrase replacement through a replacement classification head, but rely on learning this implicitly as part of the generation task. Similar to BART, we use a reconstruction loss objective during training which is a cross-entropy loss between the output and set of expected keyphrases.
Dataset - We use the OAGKX (Cano and Bojar, 2020) dataset which consists of 23 million scientific documents across multiple domains sampled from the Open Academic Graph with keyphrases tagged by the authors of the articles. The OAGKX contains keyphrases that appear in the abstract and also those which don't appear in the abstract, making it similar to the keyphrase generation setting with present and absent keyphrases. To the best of our knowledge, we are the first to explore OAGKX dataset for pre-training a large language model.
During LM pre-training we restricted the length of the input text for each sample to 512 tokens.
Note that we do not explicitly tag the keyphrases and use the readily available author tagged keyphrases associated with each document, which is a common practice in the scientific domain. This setup is analogous to how the Wikipedia corpus is used to perform entity specific pre-training in LUKE (Yamada et al., 2020) and WKLM (Xiong et al., 2019) among others.
We conducted a preliminary study of using TextRank, a baseline unsupervised keyphrase tagging techniques, to create a weakly supervised dataset but observed only marginal gains. More details are provided in Section 6 that would serve as a potential direction for future work. Additionally, we address Limitations and Ethical Concerns in Appendix 8.1.
Pre-training Strategies - We train LMs in different settings with different hyperparameters as shown in Table 1 (refer Appendix- 8.2 for more details). We pre-train in a discriminative setting with the KBIR method using RoBERTa (Liu et al., 2019) pre-trained weights. We also pre-train in a generative setting with the KeyBART method us
ing BART (Lewis et al., 2019) pre-trained weights. Ablations - Considering computational costs and environmental impact, we conduct a limited set of ablation studies to demonstrate the effectiveness of our proposed methods and also to demonstrate that the gains in performance are not due to additional data.4 We pre-train using basic random token masking strategy as RoBERTa-extended ablating both KBI and KRC from KBIR, using RoBERTa pretrained weights. We also pre-train using the KBI method, ablating KRC from KBIR using RoBERTa pre-trained weights. We pre-train BART's original denoising autoencoder strategy to recreate the original document as KeyBART-DOC by using BART pre-trained weights.
All our downstream evaluations are performed using HuggingFace Transformer's (Wolf et al., 2020) RoBERTa or BART architectures to facilitate reproducibility. We also specify all hyperparameters in Table 2. We add no additional parameters over RoBERTa or BART for the corresponding downstream evaluation architecture, demonstrating the effectiveness of our updated pre-trained weights.
4.2.1 Keyphrase Extraction
We report performance of our models for Keyphrase Extraction (KE) on Inspec (Hulth, 2003b), SemEval-2010 (SE10) (Kim et al., 2010), and SemEval-2017 (SE17) (Augenstein et al., 2017). (Sahrawat et al., 2019) explored KE as a sequence tagging task with contextual embeddings and demonstrate the effectiveness of a CRF. We compare our performance with RoBERTa+BiLSTM-CRF (Sahrawat et al., 2019), RoBERTa+TG-CRF (Chen et al., 2019), previous state-of-the-art model RoBERTa+Hypertnet-CRF, and SciBERT+Hypernet-CRF (Duan et al., 2021). However, different from these architectures, we do not use a LSTM/BiLSTM layer between the contextualized embeddings and the CRF. We fine-tune all the pre-trained language models on B-I-O tagged datasets for  $\mathrm{KE}^5$ . We use hyperparameters specified in (Sahrawat et al., 2019) and F1-score is used as the evaluation metric.
Table 3 shows our pre-trained LMs outperform SOTA by significant margins across all
<sup>4</sup>We attempted whole word masking keyphrases for both SBO and MLM for BASE model pre-training and observed no significant gains.
<sup>5</sup>https://github.com/midas-research/keyphrase-extraction-as-sequence-labeling-data
Model	Batch	Steps	Warmup	α	γ	σ	δ	MLM	KI	KR	MISL	MKR
RoBERTa-extended	4	130k	2.5k	1.0	0.0	0.0	0.0	0.15	0.0	0.0	-	-
KBI	4	130k	2.5k	1.0	0.33	1.0	0.0	0.15	0.2	0.0	10	-
KBIR	2	260k	5k	1.0	0.33	1.0	2.0	0.05	0.2	0.4	10	20
KeyBART	4	130k	2.5k	-	-	-	-	0.05	0.2	0.4	10	20
KeyBART-DOC	2	260k	5k	-	-	-	-	0.05	0.2	0.4	10	20
Table 1: Hyperparameters for our pre-training strategies. All models were trained using 8 Tesla V100 GPUs with the Adam (Kingma and Ba, 2015) optimizer and a learning rate of 1e-5. Difference in number of steps is to account for changes in batch size while seeing the same number of data points across training regimes. MLM, Keyphrase Infilling (KI) and Keyphrase Replacement (KR) show the probability of this perturbation occurring in the original text. MLM probability is reduced for KBIR in line with (Xiong et al., 2019). Maximum Infill Span Length (MISL) and Maximum Keyphrase Replacements (MKR), are based on averages from OAGKX and computational reasons. The coefficients for the loss are used to normalize the magnitude of loss across the different tasks.
Parameter	KE	NER	RE	QA	KG	SUM
Learning Rate	5e-5	1e-5	4e-5	3e-5	5e-5	5e-5
Batch	4	8	32	48	32	8
Epochs	100	5	10	2	300k	20k
GPUs	2	1	2	1	4	2
Table 2: Hyperparameters for our downstream task evaluations. KG and SUM specifies steps instead of epochs.
Model	Inspec	SE10	SE17
RoBERTa+BiLSTM-CRF	59.5	27.8	50.8
RoBERTa+TG-CRF	60.4	29.7	52.1
SciBERT+Hypernet-CRF	62.1	36.7	54.4
RoBERTa+Hypernet-CRF	62.3	34.8	53.3
RoBERTa-extended-CRF*	62.09	40.61	52.32
KBI-CRF*	62.61	40.81	59.7
KBIR-CRF*	62.72	40.15	62.56
Table 3: F1 scores for Keyphrase Extraction on Inspec, SE10 and SE17 datasets (* LMs trained by us).
three datasets despite having fewer parameters. While RoBERTa-extended, shows gains over RoBERTa+BiLSTM-CRF, this is expected since the domain of the continued pre-training data is more in line for KE evaluation. However, the models that explicitly learn keyphrase representations such as KBI and KBIR significantly outperform RoBERTa-extended. We believe the slight gain for SemEval-2010 is because of the small size of the dataset (130 - train, 100 - test).
4.2.2 Keyphrase Generation
We evaluate keyphrase generation (KG) performance on Inspec (Hulth, 2003b), NUS (Nguyen and Kan, 2007), Krapivin (Krapivin et al., 2009), SemEval (Kim et al., 2010) and KP20K (Meng et al., 2017). The task is to generate the CatSeq output of the present and absent keyphrases for a given concatenated title and abstract, as done in