Enhancing Phrase Representation by Information Bottleneck Guided Text Diffusion Process for Keyphrase Extraction

Automatic KeyPhrase Extraction (KPE) aims to extract a set of important and topical phrases from a given document, which can then be used in different tasks such as summarization Li et al. (2020), problem solving Huang et al. (2017, 2018b, 2018a), generation tasks Zhou and Huang (2019); Li et al. (2021), and so on.

Existing KPE technologies can be categorized as unsupervised and supervised methods. Unsupervised methods are mainly based on statistical information El-Beltagy and Rafea (2009); Florescu and Caragea (2017a); Campos et al. (2018), embedding features Mahata et al. (2018); Sun et al. (2020); Liang et al. (2021); Zhang et al. (2021); Ding and Luo (2021), and graph-based ranking algorithms Mihalcea and Tarau (2004); Florescu and Caragea (2017b); Boudin (2018). Supervised methods commonly formulate KPE as sequence tagging approaches Sahrawat et al. (2020); Alzaidy et al. (2019); Kulkarni et al. (2022), span-level classification Zhang et al. (2016); Xiong et al. (2019); Mu et al. (2020); Sun et al. (2021) or ranking Mu et al. (2020); Song et al. (2021); Sun et al. (2021); Song et al. (2022), or generative tasks Meng et al. (2017); Chen et al. (2018); Yuan et al. (2018); Kulkarni et al. (2022). Although supervised KPE methods require a lot of annotated data, their performance is significantly superior to unsupervised methods in many KPE benchmarks Sun et al. (2021); Meng et al. (2017).

Recently, some works have focused on constructing a zero-shot keyphrase extractor by prompting pre-trained large language models (LLMs). For example, Song et al. (2023) verified the performance of ChatGPT OpenAI (2022) and ChatGLM-6b Zeng et al. (2022) for the KPE task and found that they still have a lot of room for improvement in the KPE task compared to existing SOTA supervised models. Similar results can also be observed in Martínez-Cruz et al. (2023).

Diffusion models have been applied in many continuous domain generations like image, video, and audio Kong et al. (2020); Rombach et al. (2022); Ho et al. (2022); Yang et al. (2022). Recently, there are some works focused on applying the diffusion model to discrete text data. They usually generate continuous representations for the desired texts/words. For example, Diffusion-LM Li et al. (2022) first attempts to develop a continuous diffusion model to generate text by embedding rounding step. Following the work of Diffusion-LM, DiffuSeq Gong et al. (2022) and SeqDiffuSeq Yuan et al. (2022) designed a diffusion-based sequence-to-sequence model for the text generation task. To adapt the diffusion model to longer sequence generation, Zhang et al. (2023) proposed a sentence-level diffusion generation model for summary tasks, which directly generates sentence-level embeddings and matches from embeddings back to the original text.

Contrary to previous works, we apply the diffusion model to KPE, a phrase-level extraction task. In order to take the keyphrase information into consideration during extraction, we directly inject the keyphrase information generated by the diffusion model into each phrase representation.

Variational Information Bottleneck (VIB) is one of a group of Information Bottleneck (IB) methods. It aims to find compact representations of data that preserve the most relevant information while filtering out irrelevant or redundant information Tishby et al. (2000).

There are lots of studies that apply VIB to many NLP tasks. For example, Li and Eisner (2019) used VIB for parsing, and West et al. (2019) used it for text summarization. Recently, VIB was also used in Named Entity Recognition (NER) Wang et al. (2022); Nguyen et al. (2023), text classification Zhang et al. (2022), machine translation Ormazabal et al. (2022) and so on.

To enumerate and encode all the possible phrase representations, we first use pre-trained language model BERT Devlin et al. (2018) to encode document $D=\{w_{1},w_{2},...,w_{n}\}$, producing contextual word embeddings $\mathbf{E}=\{\mathbf{e}_{1},\mathbf{e}_{2},...,\mathbf{e}_{n}\}$. The word embeddings are then integrated into phrase representations using a set of Convolutional Neural Networks (CNNs):

$$\mathbf{s}_{i}^{k}=\mathbf{CNN}^{k}(\mathbf{e}_{i},\mathbf{e}_{i+1},...,\mathbf{e}_{i+n-1})$$

(1)

where $1\leq k\leq N$ and $N$ represents the pre-defined maximum length of phrase. The $i$th k-gram phrase representation $\mathbf{s}_{i}^{k}$ is calculated by its corresponding CNN^k.

In order to inject reference keyphrases information into each phrase representation, we use a continuous diffusion module to generate desired keyphrase embeddings.

After the diffusion generation process, the generated keyphrase embeddings $\tilde{\mathbf{H}}^{kp}$ are concatenated into each phrase representation $\mathbf{s}_{i}^{k}$. This aims to inject the information from keyphrases into each phrase, resulting in performance improvement of keyphrase ranking. Specifically, formulate the final phrase representation as:

$$\tilde{\mathbf{ss}}_{i}^{k}=\mathbf{s}_{i}^{k}||\mathbf{flat}(\mathbf{\tilde{H}}^{kp})$$

(7)

where $\mathbf{flat(\mathbf{x})}$ means that $\mathbf{x}$ is flattened to a vector. Equation 7 means that the final phrase representation not only contains the original phrase representation but also all the reconstructed keyphrase information.

For training the model to rank each phrase, we introduce a contrastive rank loss. Following the previous work Sun et al. (2021), we first take a feedforward layer to project the input representation $\tilde{\mathbf{ss}}_{i}^{k}$ to a scalar score:

$$r(\tilde{\mathbf{ss}}_{i}^{k})=\mathbf{FeedForward}(\tilde{\mathbf{ss}}_{i}^{k})$$

(8)

Then the margin rank loss is introduced to learn to rank keyphrase $\tilde{\mathbf{ss}}_{+}$ ahead of non-keyphrase $\tilde{\mathbf{ss}}_{-}$ for the given document $D$:

$$\mathcal{L}_{rank}=\sum_{\tilde{\mathbf{ss}}_{+},\tilde{\mathbf{ss}}_{-}\in D}\max(0,1-r(\tilde{\mathbf{ss}}_{+})+r(\tilde{\mathbf{ss}}_{-}))$$

(9)

Combining the keyphrase classification task during training can enhance the phraseness measurement of the phrase Sun et al. (2021); Song et al. (2021). Similar to previous work Xiong et al. (2019); Sun et al. (2021); Song et al. (2021), we introduce a classification loss for each final phrase representation for multi-task learning. We found that the use of supervised VIB substantially improves the ranking performance (See Ablation Study). Supervised VIB aims to preserve the information about the target classes in the latent space while filtering out irrelevant information from the input Voloshynovskiy et al. (2019). Given the final phrase representation $\tilde{\mathbf{ss}}_{i}^{k}$, the supervised VIB first compresses the input to a latent variable $z\sim q_{\phi_{1}}(z|\tilde{\mathbf{ss}}_{i}^{k})$. We apply two linear layers to construct the parameters $q$ using the following equations:

	$\displaystyle\mu$	$\displaystyle=\mathbf{W}_{\mu}\tilde{\mathbf{ss}}_{i}^{k}+\mathbf{b}_{\mu}$		(10)
	$\displaystyle\sigma^{2}$	$\displaystyle=\mathbf{W}_{\sigma}\tilde{\mathbf{ss}}_{i}^{k}+\mathbf{b}_{\sigma}$

where $\mu$ and $\sigma$ are the parameters of a multivariate Gaussian, representing the latent feature space of the phrase; $\mathbf{W}$ and $\mathbf{b}$ are weights and biases of the linear layer, respectively. The posterior distribution $z\sim q_{\phi_{1}}(z|\tilde{\mathbf{ss}}_{i}^{k})$ is approximated via reparameterisation trick Kingma and Welling (2013):

$$z=\mu+\sigma\epsilon,\text{where }\epsilon\sim\mathcal{N}(0,1)$$

(11)

Since the main objective of VIB is to preserve target class information while filtering out irrelevant information from the input, the objective loss function for the supervised VIB is based on classification loss and compression loss. Denoted by $y$ as the true label of the input phrase, the objective loss of the supervised VIB is defined as:

	$\displaystyle\mathcal{L}_{vib}(\phi)$	$\displaystyle=\mathbb{E}_{z}[-\log p_{\phi_{2}}(y|z)]$		(12)
		$\displaystyle+\alpha\mathbb{E}_{\tilde{\mathbf{ss}}_{i}^{k}}[D_{KL}(q_{\phi_{1}}(z|\tilde{\mathbf{ss}}_{i}^{k}),pr(z))]$

where $pr(z)$ is an estimate of the prior probability $q_{\phi_{1}}(z)$, $\alpha$ range in $[0,1]$, $\phi$ is the neural network parameters, and $D_{KL}$ is the Kullback-Leibler divergence. We use a multi-layer perceptron with one linear layer and softmax function to calculate $p_{\phi_{2}}(y|z)$. Note that Equation 12 can be approximated by the Monte Carlo sampling method with sample size $M$.

We jointly optimize the diffusion module, ranking network, and supervised VIB end-to-end. Specifically, the overall training objective loss can be represented as:

$$\mathcal{L}=\mathcal{L}_{dif}+\mathcal{L}_{vib}+\mathcal{L}_{rank}$$

(13)

For inference, the Transformer encoder first obtains the initial document embeddings $\mathbf{h}^{D}$, and then the one-step Markov $q(\mathbf{x}_{0}^{D}|\mathbf{h}^{D})$ is performed. To construct the noise keyphrase embedding $\mathbf{x}_{T}^{kp}$, we random sample $m$ Gaussian noise embeddings such that $\mathbf{x}_{T}^{kp}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$. Then the reverse process is applied to remove the Gaussian noise of $\mathbf{x}_{T}=\mathbf{x}_{0}^{D}||\mathbf{x}_{T}^{kp}$ iteratively and get the output keyphrase embeddings $\mathbf{\tilde{H}}^{kp}=[\tilde{\mathbf{h}}^{kp}_{1},\tilde{\mathbf{h}}^{kp}_{2},...,\tilde{\mathbf{h}}^{kp}_{m}]$. After that, each original phrase representation $\mathbf{s}_{i}^{k}$ is concatenated to the flattened keyphrase embeddings $\mathbf{\tilde{H}}^{kp}$ and input to the ranking network to obtain the final score for each phrase.

In this paper, we use seven KPE benchmark datasets in our experiments.

• •

  OpenKP Xiong et al. (2019) consists of around 150K web documents from the Bing search engine. We follow its official split of training (134K), development (6.6K), and testing (6.6K) sets. Each document in OpenKP was labeled with 1-3 keyphrases by expert annotators.

• •

  KP20K Meng et al. (2017) consists of a large amount of high-quality scientific metadata in the computer science domain from various online digital libraries Meng et al. (2017). We follow the original partition of training (528K), development (20K), and testing (20K) set.

• •

  SemEval-2010 Kim et al. (2013) contains 244 scientific documents. The official split of 100 testing documents is used for testing in our experiments.

• •

  SemEval-2017 Augenstein et al. (2017) contains 400 scientific documents. The official split of 100 testing documents isis used for testing in our experiments.

• •

  Nus Nguyen and Kan (2007) contains 211 scholarly documents. We treat all 211 documents as testing data.

• •

  Inspec Hulth (2003) contains 2000 paper abstracts. We use the original 500 testing papers and their corresponding controlled (extractive) keyphrases for testing.

• •

  Krapivin Krapivin et al. (2009) contains 2305 papers from scientific papers in ACM. We treat all 2305 papers as testing data.

Note that in order to verify the robustness of our model, we test the model trained with KP20K on the testing data of SemEval-2010, SemEval-2017, Nus, Inspec, and Krapivin. For all datasets, only the present keyphrases are used for training and testing. The statistics of the training set of OpenKP and KP20k are shown in Table 1.

To keep consistent with previous work Meng et al. (2017); Xiong et al. (2019); Mu et al. (2020); Sun et al. (2021), we compare our model with two categories of KPE methods: Traditional KPE baselines and Neural KPE baselines.

Traditional KPE baselines consist of two popular unsupervised KPE methods, statistical feature-based method TF-IDF Sparck Jones (1972) and graph-based method TextRank Mihalcea and Tarau (2004), and two feature-based KPE systems PROD Xiong et al. (2019) and Maui Medelyan et al. (2009).

Neural KPE baselines consist of a sequence-to-sequence generation-based model named CopyRNN Meng et al. (2017). Previous state-of-the-art method on OpenKP and KP20K, KIEMP Song et al. (2021) incorporating multiple perspectives estimation for phrase ranking. Another strong baseline, JointKPE Sun et al. (2021), including its two variants ChunkKPE and RankKPE are reproduced according to their open-source code¹¹1https://github.com/thunlp/BERT-KPE. HyperMatch Song et al. (2022), a new matching method for extracting keyphrase in the hyperbolic space. Two phrase-level classification-based models named SKE-Base-Cls Mu et al. (2020) and BLING-KPE Xiong et al. (2019). We also compare our model with BERT-based span extraction and sequence tagging methods, both of which come from the implementation of Sun et al. (2021). Note that since both of KeyBart and KBIR Kulkarni et al. (2022) are pre-trained with a well-defined pretraining strategy specifically on RoBERTa-large Liu et al. (2019), we do not compare our Diff-KPE with them for a fair comparison.

In addition, we also add the results of two large language models (LLMs) with zero-shot settings: ChatGLM2-6b Zeng et al. (2022) and ChatGPT²²2GPT-3.5-turbo, version: gpt-3.5-turbo-0125 OpenAI (2022). To restrict the output format of ChatGPT, we design the following prompt template:

[Instruction]

Please extract 1 to 15 keyphrases from the given document. Your extracted keyphrases should reasonably represent the topic of the document and must appear in the original text. You must give the keyphrases by strictly following this format: ‘‘[extracted keyphrases]", for example: ‘‘[machine learning, neural networks, NLP]"

[Document]

{document}

It should be noted that designing more complex prompts may improve the performance of LLMs, which is beyond the scope of this paper.

We use Recall (R), and F-measure (F1) of the top $K$ predicted keyphrases for evaluating the performance of the KPE models. Following the prior research Meng et al. (2017); Xiong et al. (2019), we utilize $K=\{1,3,5\}$ on OpenKP and $K=\{5,10\}$ on others. When determining the exact match of keyphrases, we first lowercase the candidate keyphrases and reference keyphrases, and then we apply Porter Stemmer Porter (1980) to both of them.

We truncate or zero-pad each document due to the input length limitations (512 tokens). We use the base version of BERT to generate initial word embeddings. We also use the base version of Sentence-BERT Reimers and Gurevych (2019) to generate initial fixed phrase embeddings for the diffusion module. The maximum length of k-gram is set to $N=5$ for all datasets. The maximum diffusion time steps $T$ is set to 100, $\alpha=2.8e-6$. The hidden size and number of layer in Transformer encoder in the diffusion module are set to 8 and 6 respectively. The latent dimension of the VIB model is set to 128. Sample size $M=5$. We optimize Diff-KPE using AdamW with 5e-5 learning rate, 0.1 warm-up proportion, and 32 batch size. The training used 8 NVIDIA Tesla V100 GPUs and took about 20 hours on 5 epochs. During training Diff-KPE, we also set a simple early stop strategy such that the model would stop training if the validation performance (F1@3 for OpenKP, F1@5 for KP20K) does not improve after 5 times consecutive evaluations (We evaluate the model every 200 optimization steps), and we select the model with the best validation performance. We run our model with 5 different random seeds and report their average score.

Table 2 shows the evaluation results of Diff-KPE and baselines. Based on the results, it is evident that the neural KPE methods outperform all the traditional KPE algorithms. Among the traditional methods, the unsupervised methods TF-IDF and TextRank show stable performance on both OpenKP and KP20k datasets, while the feature-based methods PROD and Maui outperform them on OpenKP and KP20k respectively. This is not surprising, as supervised methods benefit from large annotated data during training.

For neural KPE methods, CopyRNN performs the worst as it also focuses on generating abstractive keyphrases. HyperMatch, JointKPE and its variant RankKPE show powerful performance, outperforming other baselines such as the phrase classification-based models BLING-KPE, SKE-Base-Cls, BERT-Span, and the sequence tagging method BERT-SeqTag. It is worth noting that BERT-SeqTag and ChunkKPE exhibit competitive performance compared to RankKPE, indicating their robustness and strong performance.

Overall, Diff-KPE outperforms all baselines excluding KIEMP on both OpenKP and KP20K datasets. Compared to JointKPE, Diff-KPE shows slight improvements in F1@3 and F1@5 but a dramatic improvement in F1@1. Compared to the previous SOTA neural baseline method KIEMP, KIEMP outperforms our Diff-KPE in most F1@k scores on OpenKP and KP20k. However, Diff-KPE still exhibits performance improvements in F1@1, R@1, R@3, and R@5 on OpenKP. We hypothesize that the improvements in Recall benefit from our diffusion module, which is able to inject generated keyphrase embeddings into phrase representations, thereby enhancing the recall performance.

Moreover, to verify the robustness of Diff-KPE, we also evaluate our model trained with the KP20k dataset on five additional small scientific datasets, as shown in Table 3 ³³3We cannot evaluate KIEMP due to lack of open source code and models Song et al. (2021).. Diff-KPE demonstrates better or competitive results on all datasets compared to the best baseline JointKPE. We believe this phenomenon arises from the benefit of the diffusion module: during inference, the diffusion model can generate candidate keyphrase embeddings, providing keyphrase information for the ranking network to better rank each phrase.

To understand the effect of each component on our Diff-KPE model. We perform the ablation study on the OpenKP development set as following settings:

• •

  - w/o VIB: replace the VIB model with a single feedforward layer for keyphrase classification.

• •

  - w/o diffusion: the diffusion model is removed, and only use the phrase representations obtained from CNNs for ranking and classifying.

• •

  Diff-KPE: the original full joint model.

As shown in Table 4, the absence of the diffusion model or VIB model results in a dramatic drop in performance across all metrics, particularly in F1@1 (1.2 and 1.3 respectively). This performance decline indicates the crucial role of both the diffusion and VIB models in keyphrase ranking. The strong performance of Diff-KPE can be attributed to two main advantages. Firstly, the diffusion module directly incorporates the semantic information of keyphrases into the final phrase representations. Secondly, the supervised VIB module introduces an external classification loss during training, which indirectly enhances the diffusion module or CNNs to generate more informative n-gram embeddings. Therefore, it is evident that the addition of the diffusion module and supervised VIB greatly contributes to the overall performance improvement.