Propose-and-Refine: A Two-Stage Set Prediction Network for Nested Named Entity Recognition

1 Introduction

Named Entity Recognition (NER) aims to detect the span and category of all entities in text, which is an essential task in natural language processing. Notably, named entities are often nested in other external entities. For instance, in the sentence "This indeed was one of Uday's homes", the entity "Uday" is nested in the entity "Uday's homes" while "Uday's homes" is also nested in another larger entity "one of Uday's homes". This is because natural language sentences are hierarchical. Smaller-scale entities might be nested in largerscale entities as sub-constituency trees.

∗Corresponding author

Various methods have been proposed to handle the nested NER task, such as optimized sequence-tagging methods [Ju et al., 2018; Strakov´a et al., 2019], hypergraph methods [Lu and Roth, 2015; Katiyar and Cardie, 2018], transition-based methods [Wang et al., 2018]. These methods, however, either require complex manual-designed tagging schemas or suffer from error propagation. Recently, span-based methods, which perform classification over features of candidate spans, have gained popularity and have achieved promising results in the NER task [Sohrab and Miwa, 2018; Tan et al., 2020; Shen et al., 2021; Wang et al., 2020]. Unlike previous methods, span-based prediction can naturally address the nested NER task without complex detecting schemas and does not suffer from error propagation. However, as shown in Figure 1, span-based methods still have the following two issues. First, the prediction of an entity may rely on other phrases in the sentence. But span representations are typically generated through features of tokens that constitute the span. Therefore the relationship between a span and other phrases or entities is not considered in span-based methods, making the spanbased methods error-prone. Second, the length of the enumerated span is always limited since exhaustive enumeration is computationally expensive. Therefore it is hard for spanbased methods to detect long entities.

This paper presents the Propose-and-Refine Network (PnRNet), a two-stage set prediction network for the nested NER. To address the two previously mentioned issues of the spanbased methods, we apply a two-stage decoding procedure to detect named entities, as shown in Figure 1. In the propose stage, we use a span-based predictor to generate a set of coarse entity predictions as proposals. In the refine stage, proposals are fed into the transformer decoder [Vaswani et al., 2017], where proposals interact with each other, and richer contextual information is aggregated into the proposal representations. Finally, the refined proposal representations are used to re-predict entity boundaries and classes. In this way, the prediction errors of the coarse proposals can be eliminated with enriched information, and the boundary prediction is not constrained by the enumeration length limitation of the spanbased predictor. The final predictions are considered as a set, and a permutation-invariant loss is applied to train the model.

Additionally, we build multi-scale sentence representations to provide richer contextual information in the decoder. As mentioned before, natural language sentences are hierarchical. Therefore, representing the input sentence as a hierarchical structure is natural and helps solve the nested NER problem. For that purpose, we collect the span features generated in the propose stage to form multi-scale sentence representations. In this way, proposals can directly interact with features of spans highly related to the predicted entity rather than token features in the refine stage, which can aggregate hierarchical contextual information more effectively.

Our main contributions are as follows:

• We present a novel propose-and-refine two-stage set prediction network for the nested NER task. With richer contextualized information aggregated in the refine stage, PnRNet can make more precise predictions than the span-based predictor. Moreover, PnRNet is not constrained by the span enumeration length because we re-predict entity boundaries and classes after proposal refinement.
• To model the hierarchical structure of natural language sentences and better detect nested named entities, we build multi-scale features for decoding to provide richer hierarchical contextual information.
• Experiments on ACE04, ACE05, GENIA, KBP17, and CoNLL03 show that our model outperforms all previous models. We also conduct a detailed ablation study to validate the effectiveness of these innovations.

2 Model

In this section, we will introduce PnRNet, a two-stage set prediction network for nested NER, as illustrated in Figure 2.

2.1 Stage I: Propose

Span Feature Generation. Given an input sentence X of length N, we concatenate the contextual embedding xplm i generated by a pre-trained language model, word embedding xword i , part-of-speech embedding xpos i , and character embedding xch i of each token, and then feed the concatenated embeddings into a BiLSTM [Hochreiter and Schmidhuber, 1997] for token-level representation xi:

We generate span features from token-level sentence representations in a bottom-up manner:

where hl,i denotes the feature of the span (l, i), which is the l-gram span starting from the i-th token. We limit the bottomup construction process to spans of length L since exhaustive span enumeration is computationally expensive, especially for long sentences.

Entity Proposal. A span-based predictor is used to classify the entity type of each span with the span features generated in the previous step. The classification scores of span (l, i) is computed as follows:

Then the likelihood of that span being an entity can be obtained by:

where pcls l,i (tyt ) indicates the probability of the span (l, i) to be an entity of type tyt . E represents all entities in the sentence and ∅ is a pseudo entity type which means this span is not an entity.

Span features are sorted by p(l,i)∈E in descending order, and top-K span features which are most likely to be entities will be picked as the entity proposals Q ∈ RK×d .

It is worth noting that in the nested NER task, the prediction of an entity may rely on other related phrases or entities. However, the span-based predictor does not model the relationship between a span and other phrases. Therefore, the span-based predictor is error-prone, and these entity proposals are just coarse predictions. We have to incorporate richer contextual information into the proposal representation in the refine stage to get more precise predictions.

2.2 Stage II: Refine

PnRNet uses a transformer decoder [Vaswani et al., 2017] to refine the coarse entity proposals. The transformer decoder is composed of a stack of M transformer decoder layers. We denote Um ∈ RK×d as the output of decoder layer m. The coarse entity proposals are fed into the transformer decoder as the input of the first decoder layer U0 = Q. The output of each decoder layer will be fed into the next layer, forming an iterative refining process.

Self-attention. Entities in a sentence are related to each other. Therefore, modeling the relationship between different entities is helpful for NER. In self-attention layer, entity proposals interact with each other through the multi-head attention mechanism:

Cross-attention with multi-scale features. In order to model the relationship between the proposal and other phrases in the input sentences, entity proposals interact with sentence representations through cross-attention so that richer contextual information can be aggregated into the representations of the entity proposals:

where H is sentence representation. Since natural language sentences are hierarchical, we use multi-scale sentence representations to provide hierarchical contextual information for the nested NER task. Therefore, we collect the span representations generated in the propose stage to form layered pyramid-like multi-scale sentence representations:

where Hl ∈ R(N-l+1)×d is the list of features of spans with length l, H ∈ Rc×d is the list of all span features, and c = (2N-L+1)L 2 is the number of the enumerated spans. Since H contains features of spans of different lengths, H can be viewed as the multi-scale representation of the input sentence. With multi-scale features, proposal representations can directly attend with features of related spans. Compared with token-level features, using multi-scale features as keys and values in cross-attention can aggregate hierarchical contextual information more effectively.

Feed-forward layer. The entity proposals processed by the self-attention layer and the cross-attention layer will be fed into a feed-forward layer to generate the refined proposals of the current decoder layer:

Re-Prediction. In order to eliminate the errors in coarse proposals with the information incorporated in the transformer decoder, we use the output of the last transformer decoder layer (UM) to re-predict entity classes and boundaries. For each refined proposal ui in UM, we compute the entity classification probability of ui:

For boundary detection, we first fuse refined entity proposal ui with 1-gram span features (token-level features):

And then we perform classification over the fused features to obtain the probability of each token to be the left and right boundary of the entity:

where MLP is multi-layer perceptron.

2.3 Training Objective

Proposal Loss. We first calculate the loss of the entity proposals generated in the propose stage. The span-based entity proposal generation is a type classification task, so we use cross-entropy to obtain the loss between ground truth entity type and span-based entity classification of all enumerated spans:

where cl,i is the ground truth entity type of span (l, i).

Refine Loss. The final entity predictions of our PnRNet are order-agnostic, so we consider them as a set ˆ y = {(pcls k , plk , prk ) | k = 1, . . . ,K}. Following Tan et al. [2021], we compute a permutation-invariant set loss between the final entity predictions and ground truths. We first define the match cost between gold entity yk = (ck, lk, rk) and prediction indexed by σ(k):

where 1 denotes the indicator function. Then, we find an optimal match between prediction set and gold entities:

This optimal assignment problem can be easily solved by the Hungarian algorithm [Kuhn, 1955]. The loss for the refine stage is defined as the sum of the classification loss and the boundary prediction loss of all K predictions:

where λcls , λb are loss weights. We train the model with auxiliary losses, i.e., using the output of each decoder layer to predict entities and sum losses of all layers up for fast convergence.

3 Experiments

3.1 Setting

Dataset. We conduct experiments on four wildly used nested NER datasets - ACE04 [Doddington et al., 2004], ACE05 [Walker et al., 2006], GENIA [Ji et al., 2017], and KBP17 [Ohta et al., 2002]. Following [Katiyar and Cardie, 2018], we split samples of ACE04 and ACE05 into train, dev, test set by 8:1:1, and split samples of GENIA into train/dev, test set by 9:1. For KBP17, we split all documents into 866/20/167 documents for train, dev, and test set, following [Lin et al., 2019]. We also conduct experiments on a flat NER dataset, CoNLL03 [Tjong Kim Sang and De Meulder, 2003].

Evaluation metric. Entity prediction is considered correct when both span and category are correctly predicted. We consider precision, recall, and F1 score as our evaluation metrics. We additionally report classification F1 score and localization F1 score in the ablation study for detailed analysis.

Implementation details. We use pre-trainedBERT [Devlin et al., 2019] as the contextual encoder. For a fair comparison, we use the BERT-base-cased model for the KBP17 dataset, BERT-large-cased model for ACE04, ACE05, and CoNLL03 datasets, and BioBERT-large-cased-v1.1 [Lee et al., 2020] for GENIA dataset. We use GloVe (100d) [Pennington et al., 2014] as our pre-trained word embedding in all experiments except GENIA and use BioWordVec [Chiu et al., 2016] for the GENIA dataset. We set the span enumeration length limit to L = 16, the number of layers of the transformer decoder to M = 3. We pick K = 60 proposals with the highest scores as entity proposals, a number significantly larger than the number of entities in most sentences.

3.2 Overall Performance

Table 1 demonstrates the overall performance of our PnRNet compared with various baselines. The experiments on nested NER datasets show that our PnRNet outperforms all previous methods by a large margin. Specifically, PnRNet achieves +0.71%, +0.58%, +1.31%, and +1.22% gain in F1score in ACE04, ACE05, GENIA, and KBP17. On the flat NER dataset CoNLL03, PnRNet also achieves SOTA performance. It shows that modeling interactions between entities and incorporating richer hierarchical contextual information into the entity proposals not only help in detecting nested entities but also improve the performance of the flat NER.

3.3 Ablation Study

We conduct the ablation study in the following three aspects, as shown in Table 2.

Span-based entity proposal. To validate the effectiveness of our proposal generation process, we replace these proposal features with a set of randomly initialized learnable embeddings. The F1-score drops by -4.7% and -0.41% in ACE04 and GENIA datasets without entity proposal. It shows that proposal representations generated in the first stage provide necessary information for entity recognition compared with randomly initialized vectors.

Proposal refinement. In the ablation experiment without proposal refinement, we directly evaluate the performance of the span-based predictor. The performance drops by -1.94% and -0.20% in ACE04 and GENIA compared with full PnRNet. This indicates aggregating richer contextual information and modeling the relationship with other phrases can benefit the performance of NER.

Multi-scale feature. In the ablation experiment without multi-scale features, we use the output of the sequence encoder (H′ = [x1, x2, . . . , xN]), which is the token-level sentence representation, to provide contextual information in proposal refinement. The performance drops by -1.13% and -0.89% in ACE04 and GENIA datasets. It shows that multi-scale sentence representations provide richer hierarchical contextual information, which is helpful in entity recognition.

3.4 Detailed Analysis of the Effect of the Proposal Refinement

For detailed analysis, we compare the performance between span-based proposals and the final predictions of PnRNet on entities of different lengths. As shown in Table 3, as the entity length grows, the performance of the span-based entity recognition (entity proposals) declines significantly. In contrast, the performance of the final prediction only gets a slight drop. Furthermore, the two-stage detection still has satisfactory performance on very long entities, even when it exceeds the prediction limits of the span-based predictor used in the propose stage. This indicates the refine stage of PnRNet, which performs interaction between proposals and incorporates multiscale contextual information into proposal features, helps a lot in recognizing nested named entities, especially for long entities.

3.5 Visualization of Multi-Scale Cross-Attention Weight

We visualize the cross-attention weight map of the last decoder layer of our PnRNet to confirm the effectiveness of the multi-scale features. As shown in Figure 3, four spans with the highest attention scores are "law professor" (the predicted entity of the proposal), "rick pildes" (the person name of the "law professor"), "law professor rick pildes" (an entity related to "law professor") and "you" (another entity mentioned in this sample). This indicates that through multi-scale features, the entity proposal can directly attend to features of spans that are highly related to the proposal in cross-attention.

" Health officials in [[Saudi [Arabia]GPE]GPE]GPE have asked pilgrims visiting its holy sites to wear masks in crowded places to stop the spread of the MERS coronavirus . "

On the contrary , the patent owner may continue to pursue [[[its]ORG]ORG]PER rights through the courts .

Table 4: Case study. We mark coarse proposals in red, the corresponding final predictions in blue, and ground truths in green. For simplicity, we omit other irrelevant proposals and predictions and only show one pair of proposal and its corresponding final prediction in each case.

With the power of multi-scale features, the transformer decoder can aggregate hierarchical information that is helpful to detect named entities, improving the performance of the nested NER.

3.6 Case Study

We illustrate some cases in Table 4 to show that our PnRNet can eliminate errors in coarse proposals through proposal refinement. In case 1, by aggregating richer contextual information, boundary errors in the coarse proposal ("Arabia") can be fixed in the final prediction ("Saudi Arabia"). In case 2, the pronoun entity "its" is misclassified by the span-based predictor as a PER. By interacting between the entity proposal and other proposals and contextual phrases, our PnRNet can correctly classify the entity class as an ORG.

4 RelatedWork

Various methods have been proposed to recognize nested named entities. Since traditional sequence tagging methods [Huang et al., 2015; Lample et al., 2016] which predict a label for each token cannot address nested named entities, some optimized tagging schemes are proposed to cope with the nested NER task [Ju et al., 2018; Strakov´a et al., 2019]. Hypergraph methods [Lu and Roth, 2015; Katiyar and Cardie, 2018] represent the parsing structure of the input sentence as a hypergraph and detect nested entities on the graph. Transition-based methods [Wang et al., 2018] generate a sequence of constituency parsing actions to detect nested entities.

Span-based methods predict entities with span representations. Sohrab and Miwa [2018] exhaustively enumerate spans and generate span representation with boundary token features and pooling features of span tokens. Tan et al. [2020] first predict boundary and then perform classification over span features. Wang et al. [2020] use a pyramid model to generate span representations layer by layer. Yu et al. [2020] use a bi-affine operation to compute span classification scores. Shen et al. [2021] perform boundary regression after span-based prediction. Span-based methods can naturally address the nested NER task without complex detecting schemas and have achieved promising performance. However, span representations does not model the relationship with other contextual phrases or entities. Besides, span-based methods have difficulty predicting long entities because the span enumeration length is limited to reduce computational complexity. Our PnRNet solves all these two issues through proposal refinement and re-prediction.

Other studies design new architectures or incorporate different paradigms for the nested NER task. Lin et al. [2019] first identify anchor words of entity mentions and then detect entity boundaries. Li et al. [2020] use a machine reading comprehension model for the nested NER. Yan et al. [2021] model the nested NER as a sequence generation task. Since the nested NER task is essentially an order-agnostic set prediction problem, Tan et al. [2021] use a sequenceto-set neural network to detect entities as a set and apply a permutation-invariant set loss for training. However, most of these methods only use token-level encodings as sentence representations, which have difficulty representing the hierarchical structure of natural language sentences. We mitigate this issue with multi-scale sentence representation.

5 Conclusion

This paper presents a novel two-stage set prediction network named Propose-and-Refine Network. Firstly, we use a span-based predictor to generate a set of coarse entity predictions as proposals. Then proposals are fed into a transformer decoder for further refinement and finally re-predict entity boundaries and entity classes. So prediction errors in coarse entity proposals can be eliminated, and the model can better detect long entities. Moreover, we generate multi-scale sentence representations to provide richer hierarchical contextual information of the input sentence. Finally, we apply a cross-entropy loss for the entity proposals and a permutationinvariant set loss for the final predictions. Experiments show that our model achieves state-of-the-art performance on flat and nested NER datasets.

Acknowledgments

This work is supported by the National Key Research and Development Project of China (No. 2018AAA0101900), the Key Research and Development Program of Zhejiang Province, China (No. 2021C01013), CKCEST, and MOE Engineering Research Center of Digital Library.