Span-based Unified Named Entity Recognition Framework via Contrastive Learning

1 Introduction

Named Entity Recognition (NER) is a foundational task for Natural Language Processing (NLP), which aims to extract named entities in the given text and classify them into predefined entity types such as persons, organizations and locations. As a subtask of information extraction, NER serves as a crucial building block in many NLP applications, including entity linking [Ganea and Hofmann, 2017; Le and Titov, 2018], relation extraction [Zhong and Chen, 2021; Rathore et al., 2022] and knowledge graph construction [Sarhan and Spruit, 2021].

∗Corresponding author.

The current approaches for NER typically employ pretrained transformer architectures, like BERT [Devlin et al., 2019] or RoBERTa [Liu et al., 2019], as a base encoding component to perform in-domain supervised learning on individual datasets [Yu et al., 2020; Yan et al. Zhang et al., 2023], as shown in Figure 1(a). Although they demonstrate strong performance on each dataset, these dataset-specific methods struggle to generalize across new domains or transfer to unfamiliar entity types without costly full retraining. For instance, a NER model trained to identify

person and location entities in news text struggles to identify disease names in medical documents. To enable learned knowledge transfer across domains, recent prompt-based generative methods propose jointly training across diverse datasets with a unified model [Lu et al., 2022; Lu et al., 2023; Wang et al., 2023], as shown in Figure 1(b). To this end, they take prompt instructions and text as input, leveraging a sequence-to-sequence (seq2seq) generative framework to sequentially output entities in an autoregressive manner.

However, due to inherent limitations of autoregressive language models for entity extraction tasks, these prompt-based generative approaches have two major weaknesses. First, compared with span-based methods directly encoding entities at span-level, the sequential entity generation of approaches can only implicitly captures associations between spans and labels, as well as inter-span relationships, making it difficult to accurately predict span boundaries. Second, due to the token-by-token decoding in seq2seq frameworks, these models cannot generate multiple entity predictions in parallel and thus suffer from slower inference speed.

To address the above-mentioned challenges, we draw inspiration from CLIP [Radford et al., 2021] utilizing contrastive learning to enhance model generalizability, and propose a novel Span-based Unified NER approach via contrastive learning (SUNER), which can jointly handle multiple NER datasets, as shown in Figure 1(c). The key idea of our method is to directly align text spans and entity type representations within a shared semantic space, enabling parallel extraction across domains. In this manner, we break through limitations of entity categories by introducing textual descriptions of entity types. Furthermore, by utilizing pretrained language models that have learned general linguistic patterns from large unlabeled corpora, our model achieves enhanced semantic matching and superior generalization capabilities to new domains.

Specifically, the proposed SUNER comprises two modules: span detection and span classification. First, to adapt variations across different domains, detection module focuses solely on extracting mention spans of entities without considering entity types. This approach, separating the identification process from later domain-specific categorization, is more generalized compared to directly classifying entities across diverse datasets. Second, detection module introduces an ingenious entity marker framework that highlight candidate entity spans in the input sentence, effectively capturing entity relationships while retaining the original textual structure to achieve better span representations. Following this, leveraging a BERT-based contrastive learning approach, the module maps candidate spans and textual type descriptions into the same semantic space. This allows flexible extraction per specified category by relying on learned semantic similarities. In experiments, we jointly train SUNER on seven different datasets, supporting a total of 47 entity types. Extensive assessments under supervised conditions as well as zero/fewshot settings validate the effectiveness of our proposed model in achieving better performance and higher efficiency than previous state-of-the-art (SOTA) unified NER models.

In summary, our main contributions are as follows:

• Instead of relying on autoregressive generated models,
• we introduce a span-based unified NER approach via contrastive learning that can extract entities across domains based on semantic similarities.
• To produce enhanced span representations, we propose a well-designed entity marker that highlights candidate spans, capturing inter-entity relationships while retaining original textual structure.
• In experiments across various settings, our model achieves a 1.9% average improvement under supervised conditions and a 5.9%-6.4% absolute increase for zero/few-shot learning over previous SOTA unified models, while also maintaining higher efficiency.

2 RelatedWork

2.1 Named Entity Recognition

Dataset-Specific NER Transformer-based pretrained language models [Devlin et al., 2019; Liu et al., 2019; Lee et al., 2020] have made significant achievements in current NER tasks with their strong representational capabilities. For flat NER, classic sequence labeling methods [Souza et al., 2019; Li et al., 2020] combine BERT with Conditional Random Fields (CRF) [Ma and Hovy, 2016] to assign the BIO tag for each token. In a different approach, span-based methods directly model candidate spans and employ BERT to generate span representations for classification, achieving the best performance for both flat and nested NER task [Yu et al., 2020; Li et al., 2022; Zhu and Li, 2022; Yan et al., 2023; Zhang et al., 2023]. However, these methods typically perform indomain supervised learning on specific datasets and are limited to a set of fixed predefined entity categories. This leads to challenges in adapting to unfamiliar entity types without extensive retraining.

Unified NER To enable cross-domain knowledge transfer, recent prompt-based generative methods propose jointly training on diverse NER datasets and integrating label information into a unified model. Among these, T5-based approaches [Raffel et al., 2020] like PUnifiedNER [Lu et al., 2023] and UIE [Lu et al., 2022] employ a unified structural model using a seq2seq generative framework to sequentially output specified entity types. Meanwhile, InstructUIE [Wang et al., 2023] and UniNER [Zhou et al., 2023] based on Large Language Models (LLM) converts various NER datasets into instruction-following format, subsequently finetuning encoder-decoder generative models. However, these generative NER methods struggle with precise entity boundary detection and slow decoding, especially for LLM-based models, whose size and cost limit their use in resourcelimited settings. Although USM [Lou et al., 2023] proposes a token-based unified method, it cannot deal with overlapped datasets, and also suffer from inefficient decoding due to semantic linking for each word. In our setting, to improve generalization and time efficiency, we apply a contrastive learning framework for semantic matching between spans and textual type descriptions, enabling parallel entity extraction.

2.2 Contrastive Learning

Self-supervised contrastive learning has been widely utilized in diverse tasks to generate representations [Chuang et al.,

2020; Gao et al., 2021; Radford et al., 2021; Han et al., 2022; Tan et al., 2022]. The core concept of contrastive learning aims to pull positive pairs closer while pushing negative pairs apart. SimCSE [Gao et al., 2021] applies twice dropout in the forward process to refine the better sentence representation. The CLIP [Radford et al., 2021] employs a contrastive learning framework for pretraining on semantic matching between image and caption text. This approach aims to transcend the limitations of fixed object categories and improve model generalizability to new categories. In this paper, inspired by the CLIP model, we propose a span-based contrastive learning unified NER approach. Different from some contrastive NER [Das et al., 2022; Huang et al., 2022] methods that focus on token-level contrasts, our approach aligns spans with textual entity types in a shared semantic space. By learning semantic similarities between spans and entity types, our model can better generalize to unseen entity categories.

3 The ProposedMethod

The architecture of our proposed SUNER model is shown in Figure 2. It consists of two main modules. First, the span detection module focuses on identifying all mention spans irrespective of entity type. It leverages a BERT-based biaffine model to generate contextual span representations and then applies a sigmoid layer to extract candidate entity spans. Second, the span classification module inserts special marks to highlight these extracted spans in the input text . This marked text, along with entity type descriptions, is then fed into a BERT-based contrastive framework to obtain vector representations of the spans and type texts separately. By maximizing similarity between the span and type text vectors, our model classifies each span into an entity type. In the following, we will first provide a task definition, then present components of our method in detail.

3.1 Task Formulation

The input for our SUNER model includes a sentence S = [w1,w2, . . . ,wn] and entity types T (expressed in natural language), where n denotes the total number of tokens in the sentence. The goal of SUNER is to extract all entities E = {(li, ri, ti)}en i=0 based on semantic similarity, where en is the number of entities and li, ri, ti represent the left and right boundary indices and type of the i-th entity.

3.2 Span Detection

As discussed in the introduction, we aim to train the model on multiple datasets and learn shared features for domain generalization. To this end, we initially extract entity spans without specifying entity types, a more generalized approach than directly classifying entities across various domains.

Span Representation Encoder

To achieve contextualized span representation for detection, we first augment the input S with the dataset name D during supervised training to harmonize annotations discrepancies, while omitting this field during zero-shot evaluation. Then the augmented S′ = [wcls,D,wsep,w1,w2, . . . ,wn] is fed into BERT to get token-level representation H =

[h1, h2, . . . , hn], where the special tokens wcls, wsep correspond to the start and separator tokens in BERT.

Next, we follow [Yu et al., 2020] and employ a biaffine encoder to generate span representation. Given a sentence S with token embeddings H, we process H through two separate feed-forward networks (FFNs) to obtain the start-of-span representations Hs ∈ Rn×dh and end-of-span representations He ∈ Rn×dh , then the span representation matrix R is calculated as follows:

where U ∈ Rdh×dr×dh , W ∈ Rdr×2dh and b ∈ Rdr are learnable parameters, dh and dr are the hidden size, ⊕ denotes element-wise concatenation, and R ∈ Rn×n×dr . Each cell (i, j) in the matrix R represents the feature representation of span from the i-th token to the j-th token.

Utilizing the matrix parallel computation of the biaffine model, our approach can efficiently generate representations of all spans simultaneously.

Span Detection Objective

After getting spans representations from biaffine model, we obtain the span prediction logits P through a sigmoid layer as follows:

And then, we use the binary cross entropy to calculate span-

A common issue in pipeline systems is error propagation, where inaccurate mention boundaries will result in incorrect entity type classification. To enhance the precision of span detection in a more fine-grained level, we further introduce an auxiliary boundary detection task to generate high-quality candidate spans. Specifically, we feed the token representations hi into two separate multi-layer perceptron (MLP) classifiers, and apply a sigmoid function to obtain the probabilities Psi and Pei of token wi being the start or end of an entity span respectively:

Then, the boundary-based loss can be calculated as:

Lbdr = Lsbdr + Lebdr

Finally, we combine span-based loss Lspan with the auxiliary boundary-based loss Lbdr using a weight λ to produce the overall training objective for the span detection module:

3.3 Span Classification

To further improve generalization across domains, the span classification module employs an entity marker structure for representing entity spans and utilizes contrastive semantic matching with textual type descriptions for effective entity categorization.

Entity Mark Structure

In order to effectively leverage boundary information for span classification, we do not directly employ the biaffine encoder representations. This method simply concatenates the start and end tokens, only capturing contextual information around each individual entity. It fails to capture inter-dependencies between the spans, which are crucial for entity type confirmation, especially for nested entities.

Instead of processing each span at token level, we insert specialized entity markers at the input layer to highlight spans in the original text. This marker representation allows us to capture dependencies between entities at full span level. Specifically, given an input sentence S and candidate entity spans C = {c1, c2, . . . , cm}, we insert text markers [ei] and [/ei] before and after each span ci to construct a hierarchically marked sentence ˆ S:

where wstart(i) and wend(i) denote the start and end tokens of span ci respectively.

In addition, to preserve original textual structure of the input for the pretrained language encoder, we tie the positional embeddings of the markers with the start and end tokens of the corresponding span. For instance, the marker [ek] and the start token wstart(k), as well as the closing marker [/ek] and the end token wend(k) share the same position embedding. In

Contrastive Learning

this manner, we keep the position id of original text tokens unchanged. Furthermore, we impose a constrained attention mask to the attention layers. Text tokens only attend to text tokens and are isolated from entity marker tokens. Meanwhile, entity marker tokens can attend to all text tokens to aggregate information for their associated spans, as well as other marker tokens to capture inter-span dependencies.

Contrastive Learning Framework

Our contrastive learning framework comprises two modules: an entity type encoder and a span representation encoder. To conserve memory and enable more efficient training, both encoders utilize a shared BERT base model. In this task, we process two types of inputs: entity type descriptions and the modified sentence ˆ S with marked candidate entities.

Entity Type Encoder First, we leverage the entity type encoder to generate representations for each targeted entity. In this study, we describe each entity type with a sequence of natural language tokens. Specifically, for a given textual entity type Ek, we feed Ek into BERT and pass the resulting [CLS] representation through a MLP layer to obtain the projected type embedding ek:

By leveraging descriptive entity text rather than discrete labels, our model can fully utilize the robust representational power of pre-trained language models which are trained on vast unlabeled corpora. This approach significantly boosts the model's ability to generalize across new domains. Even for previously unseen entities, the use of explanatory type texts with BERT embeddings allows model to better understand and categorize these entities.

Span Representation Encoder Similar to the span representation encoder used in the span detection module, we feed the marked sentence ˆ S containing candidate entity spans into a BERT-based biaffine model to produce vector representations ˆri for each span ci, formulated as:

where ̂ start(i) and ̂ end(i) are the indices of [ei] and [/ei] in ˆ S. With this entity marker framework, we effectively capture entity relationships while preserving the original textual structure, thereby achieving better span representations.

Contrastive Object Based on the entity type embeddings and entity span embeddings discussed above, we calculate the supervised contrastive loss via the InfoNCE [Oord et al., 2018] formulation as follows:

where ti denotes the target type for span ci, 1[k̸=ti] indicates whether category k matches this label, cos calculates cosine similarity and τ refers to a tuned temperature parameter. Through this objective, we pull span representations closer to their corresponding type vector (positive pair) while pushing farther from unrelated categories (negative pairs) to distinguish entities across domains.

3.4 Training and Inference

Training For the span detection and classification modules, we fine-tune two separate pretrained BERT models, BERTd and BERTc, using the task-specific losses Ldection and Lclass respectively. When training the classification module, we only consider gold entity spans in each sentence as supervision signals rather than predicted entities with potential noise.

Inference During inference, for span detection, we first prune out non-entity spans by setting a threshold θ1, then greedily select the highest probability span proposals while ignoring conflicting proposals to produce candidate entity spans. For span classification, we categorize each span ci by selecting the type embedding ek with maximum cosine similarity cos (ˆri, ek) /τ to its span representation ˆri, and only retain entities with matching scores higher than threshold θ2.

4 Experiments

4.1 Experimental Setting

Datasets

We train and evaluate our model on seven existing public NER benchmarks including diverse domains such as news, biomedicine, movie, and restaurant, etc. The used datasets include three nested NER datasets: ACE 2004 1 , ACE 2005 2 and GENIA [Kim et al., 2003]; along with four flat NER datasets: CoNLL 2003 [Sang and De Meulder, 2003], OntoNotes 5 3 , MIT Restaurant and MIT Movie [Liu and Lane, 2017]. We use MIT Restaurant and MIT Movie datasets with standard train, dev, and test splits, while adopting the splits of Yu et al. [2020] for the remaining datasets. After unifying the labels with identical semantics across different datasets, we ultimately construct a unified corpus encompassing 47 entity categories.

1 https://catalog.ldc.upenn.edu/LDC2005T09

2 https://catalog.ldc.upenn.edu/LDC2006T06

Implementation Detail

To ensure a fair comparison with previous SOTA works [Zhu and Li, 2022; Yan et al., 2023], we utilize RoBERTa-base [Liu et al., 2019] as the pretrained transformer encoder for both the span detection and classification modules, resulting in a total parameter size of 220M for our SUNER model. In the span detection module, we set the auxiliary loss weight λ as 0.7, the biaffine encoder hidden size is 300. The filtering thresholds θ1 and θ2 are set to 0.5 and 0.4, respectively. During training, all parameters are optimized using Adam with a peak learning rate of 1.5e - 5, while Hyper-parameter tuning is performed based on validation set.

Evaluation

We use span-level precision, recall and F1 score to evaluate the performance. A predicted entity is confirmed correct only when its label and boundaries exactly match the ground truth. The reported results are averaged over three runs with different random seeds.

4.2 Results on Supervised Settings

Compared with Unified Methods

For supervised experiments, we first compare SUNER against three unified NER approaches:

• USM [Lou et al., 2023] decouples NER into two tokenlinking tasks using RoBERTa-Large encoder. However, it cannot jointly train on overlapped datasets.
• PUnifiedNER [Lu et al., 2023], utilizing a T5-based model, employs a prompt-based generative framework to sequentially output entities. We reproduce PUnifiedNER under the same setting.
• One-stage SUNER directly employs the span-based contrastive approach for classification without a separate span detection step.

The results are presented in Table 1, the following observations can be found: (1) SUNER significantly outperforms the generative PUnifiedNER method, achieving an average F1 score improvement of +1.90% across all datasets. Due to token-by-token decoding limitations, the generative approach struggles for direct span modeling and interaction. This limitation is particularly evident on datasets with intricate nested structures, such as Genia, where PUnifiedNER underperforms SUNER by a substantial margin of 4.14%. (2) When separating span detection from domain-specific entity categorization, SUNER on average surpasses the one-stage method by +0.56%. It illustrates that span detection is a more general task, enabling the model to capture shared features across domains and consequently enhance its performance.

3 https://catalog.ldc.upenn.edu/LDC2013T19

Compared with Dataset-Specific Methods

We further compare the performance of the proposed model against recent dataset-specific SOTA methods, as shown in Table 2. Different from other models that train separate models for each dataset, SUNER adopts a more challenging strategy by using a single model for joint training across diverse datasets. This unified strategy not only improves the model's ability to generalize but also optimizes storage efficiency. Additionally, our model achieves competitive performance to SOTA dataset-specific methods. Notably, due to benefit from handling of similar entity types in ACE, our model lags behind the SOTA by only 0.64% on ACE 2004 and 0.21% on ACE 2005. While there is a 2.78% performance gap on Genia, this can be attributed to their use of BioBERT [Lee et al., 2020], a specialized model for biomedical texts.

4.3 Results on Zero-shot/Few-shot Settings

Zero-shot Settings

For zero-shot experiments, we focus on in-domain CoNLL 2003 dataset and out-of-domain Restaurant and Movie datasets. In this setting, we first pretrain SUNER on six other corpora before applying model to the new target datasets.

As shown in Table 3, SUNER substantially exceeds the prompt-based PUnifiedNER across all datasets, achieving an impressive average F1 increase of +6.37%. Especially for out-of-domain generalization, SUNER demonstrates considerable gains of +8.13% and +7.09% F1 on Restaurant and Movie datasets respectively. Our approach even averages higher than the UniNER, which is built on a costly 7B parameter LLM that may be impractical in resource-limited settings. By aligning unseen entities to explanatory type texts under our span-based contrastive framework, SUNER better recognizes these unfamiliar categories.

Few-shot Settings

To further analyze model generalization with limited supervision, we conduct additional few-shot experiments based on the pretrained zero-shot model. Specifically, we sample just 20 training sentences per target dataset and repeat each experiment three times to avoid the effect of random sampling.

The results of few-shot learning are shown in Table 3. Again, our proposed SUNER model outperforms PUnifiedNER with +5.91% higher average F1, demonstrating the benefits of modeling entities at the span level. Although spanbased Biaffine approaches achieve high supervised performance in dataset-specific setting, directly classifying entities via MLP layer struggles to transfer across unfamiliar domains. In contrast, SUNER exhibits strong adaptation capabilities, improving average zero-shot scores by 14.26% using just 20 examples per dataset.

5 Analysis

5.1 Analysis of Inference Efficiency

In this section, we conduct experiments to compare the inference efficiency between SUNER and other baseline models. As illustrated in Table 4, compared to the generative PUnifiedNER, SUNER not only obtains higher accuracy under similar parameter settings but also achieves a notable 5.96× speed increase. Due to seq2seq decoding limitations, PUnifiedNER can only output entities step by step in an autoregressive manner. In contrast, our span-based contrastive framework enables parallel extraction, significantly boosting processing speed. When compared to the One-stage SUNER, our model exhibits the stronger generalization, improving performance on the GENIA dataset by 1.43%. Overall, SUNER achieves satisfactory accuracy at affordable computational cost, proving its effectiveness for unified NER task in practice applications.

5.2 Analysis of EntityMark Structure

To analysis the effect of entity mark structure, we compare our approach with three other variant methods. The results are shown in Table 5. All marker-augmented configurations outperform the raw text input setting, validating the importance of modeling span interactions in NER Task. Especially for nested task like GENIA, nested dependency relationships are crucial for entity type confirmation. In addition, removing either the constrained attention mask or position embedding tying leads to varying degrees of performance drops, indicating both are critical for effectively incorporating markers while preserving original textual structure.

We achieve this balance by isolating text tokens from entity marker tokens through the attention mask, and ensuring position embedding tying to preserve the position embeddings of the original tokens. As a result, the BERT context vectors of the text tokens remain unchanged. Yet, by integrating span vectors using entity markers, we facilitate interactions between spans, leading to superior performance.

5.3 Analysis of Generalization Capability

As previously highlighted, our proposed SUNER aims to improve model generalization to new domains through joint training across different datasets. To validate this point, we evaluate SUNER's performance against the baseline models with different training sizes, particularly focusing on out-ofdomain Restaurant and Movie datasets. As shown in Figure 3, even with just 50 training examples, SUNER achieves approximately 80% of the optimal performance, which demonstrates its quick adaptation ability to new domains. When increasing the dataset size to 200 examples, SUNER's performance continues to improve, lagging behind the SOTA by reasonable margins of 9.76% on Restaurant and 8.28% on Movie. By incorporating a semantic matching-based contrastive framework, SUNER can perform well with limited labeled data, highlighting its potential in practical applications where annotation scarcity hinders model development.

6 Conclusion

In this paper, we present a novel method SUNER that leverages span-based contrastive framework to extract entities in parallel for unified NER. To improve model generalization, we divide SUNER into two stages, separating span detection from domain-specific span classification. In the span classification stage, we utilize descriptive entity text instead of discrete labels to help model better understand and categorize entities. We conduct extensive experiments in various settings, including supervised and zero/few-shot conditions. The results demonstrate that our approach achieves better performance and higher efficiency compared to previous state-ofthe-art methods.

Acknowledgments

The work is supported by National Natural Science Foundation of China (No. 62172039, U21B2009 and 62276110) and MIIT Program(CEIEC-2022-ZM02-0247).