Enhancing NER by Harnessing Multiple Datasets with Conditional Variational Autoencoders
Authors: Taku Oi, Makoto Miwa
Accepted at: acl25

Enhancing NER by Harnessing Multiple Datasets with Conditional Variational Autoencoders

Taku Oi and Makoto Miwa
Toyota Technological Institute
Correspondence: makoto-miwa@toyota-ti.ac.jp

Abstract

We propose a novel method to integrate a Conditional Variational Autoencoder (CVAE) into a span-based Named Entity Recognition (NER) model. This approach models shared and unshared information among labels in multiple datasets, thereby easing training on these datasets. Experimental results using multiple biomedical datasets demonstrate the effectiveness of the proposed method, showing improved performance on the BioRED dataset. Our source code for this implementation is publicly available at GitHub1.

1 Introduction

Named Entity Recognition (NER) is a fundamental Natural Language Processing (NLP) task, serving as a crucial first step in information extraction. Although Large Language Models (LLMs) with zero- or few-shot learning have been widely investigated, supervised learning or full fine-tuning remains essential for achieving high performance (Zhong and Chen, 2021; Luo et al., 2023; Wang et al., 2023; Munnangi et al., 2024; Zhou et al., 2024). The performance of such NER models depends on the quantity of labeled data, which is often costly to create manually.

One approach to increase labeled training data is to combine existing labeled datasets (Luo et al., 2022; Islamaj et al., 2021a). However, even datasets targeting the same types of named entities have different type definitions and annotation criteria, making it challenging to merge them into a unified training dataset. For instance, in the biomedical NER datasets NLMChem (Islamaj et al., 2021a) and BioRED (Luo et al., 2022), which both include a Chemical entity type, the term “hematoxylin” is labeled as Chemical in NLMChem, but not in BioRED. This discrepancy arises because BioRED’s annotation guidelines explicitly exclude staining reagents from the Chemical label.

A common strategy for utilizing multiple datasets is multi-task learning (MTL) (Wang et al., 2018; Zuo and Zhang, 2020; Rodriguez et al., 2022). MTL shares a base model among tasks while maintaining a separate classification layer for each dataset, which allows for learning without explicitly addressing dataset-specific differences. However, MTL cannot account for relationships among labels in different datasets. Luo et al. (2023) tackled this issue by manually editing additional datasets to align with the target dataset, which improved performance, but the manual editing process is costly.

Span-based approaches are widely used for NER (Sohrab and Miwa, 2018; Ouchi et al., 2020; Nguyen et al., 2023) due to their simplicity and effectiveness. Span representations are critical for these methods, as they directly influence the ability to accurately identify and classify named entities. Nguyen et al. (2023) explored integrating Variational Autoencoders (VAE) to enhance span representations. Inspired by this, we employ a span-based approach with Conditional VAE (CVAE) to model spans using label-specific conditions that are both shared and unshared across multiple datasets.

This study proposes a method that models the relationships between labels in different datasets as label-specific conditions using CVAE, thereby incorporating these conditions to alleviate differences across datasets. The contributions of this paper are as follows:

• We propose a span-based NER method to integrate multiple existing datasets, despite their differing type definitions and annotation criteria.

• We propose integrating CVAE to incorporate and model the relationships between labels from multiple datasets.

1https://github.com/tti-kde/CVAE-NER
Figure 1: Overview of the proposed method, illustrating the integration of CVAE into a span-based NER model. Within the CVAE component, the prior distribution vector provides information regarding shared and unshared labels to the model.

• We demonstrate the effectiveness of our method through experiments on multiple biomedical datasets, achieving improved performance on the BioRED (Luo et al., 2023) dataset with either of the two encoders used.

2 Related Work

2.1 Span-based NER model

Span-based models, which directly classify specific text spans within a document, have recently gained attention (Zhong and Chen, 2021; Sohrab and Miwa, 2018) due to their simplicity and ability to model entity spans directly, unlike sequence labeling models that process tokens to represent entities. Zhong and Chen (2021) represent each target span by concatenating the embeddings of the first and last tokens from the encoder with the span length embedding. The span representations are then passed through a fully-connected layer, followed by a Softmax layer for classification. The model demonstrated high performance in NER and achieved state-of-the-art performance in end-to-end entity and relation extraction.

Span representations are crucial in span-based NER models. Ouchi et al. (2020) enhanced interpretability by learning the similarity between span representations and assessing this similarity during prediction. Nguyen et al. (2023) proposed a span-based NER model that incorporates span reconstruction and synonym generation using VAE. Span reconstruction ensures that instance-specific information is retained in the span representations, thereby improving extraction performance.

2.2 NER using multiple datasets

MTL (Rodriguez et al., 2022; Zuo and Zhang, 2020) employs a shared encoder and a separate classification layer for each dataset during training. Recent approaches leverage the capabilities of LLMs for MTL (e.g., Zhou et al. (2024)). However, these approaches do not explicitly account for shared label information across diverse datasets.

Luo et al. (2023) trained NER on multiple datasets targeting six types of entities: Disease, Chemical, Gene, Species, Variant, and Cell line, to improve performance on the BioRED dataset. Additional datasets underwent manual editing to ensure each contained only a single label type through label integration, deletion, and adjustment of label spans to align annotation criteria. The method outperformed existing MTL methods but required costly manual annotations. Our approach uniquely addresses the label differences through direct modeling with CVAE, which has not been explored in existing methods.

3 Method

This study proposes integrating CVAE into an NER model to incorporate a prior distribution that provides the model with information about shared and unshared labels among the datasets as a label-specific condition. We anticipate that this condition will capture relationships between labels in multiple datasets, enabling the effective use of existing labeled datasets. An overview of the proposed model is illustrated in Figure 1.

We follow Zhong and Chen (2021) to model the
representation of a span \( x_1, \ldots, x_n \). The representation is constructed by concatenating the representations of the beginning token and end token from the encoder, along with the embeddings of the span length, as shown in Equation (1):

\[
h_{span} = \text{Linear}(\text{Concat}(x_1, x_n, \Phi(n))),
\]

where \( \text{Linear}(\cdot) \) is a fully-connected layer and \( \text{Concat}(\cdot) \) is vector concatenation, \( x_i \) is the output representation of a token \( x_i \) from the encoder, and \( \Phi(n) \) is the embedding of the span length \( n \). The dataset name is added as a special token to the input text to ensure the model recognizes inputs from different datasets. The span representation is classified using two fully-connected layers, followed by a Softmax layer. Cross-entropy loss (\( L_{CE} \)) is employed as the loss function.

In CVAE, the span representation is concatenated with a prior distribution vector, which functions as the conditional parameter during the reconstruction phase. The prior distribution vector is derived from the correct label and encodes the shared and unshared label information across datasets\footnote{The actual settings for the prior distribution vector will be shown in §4.2.}. After concatenating with the prior distribution vector, each span representation is processed through two fully-connected layers to estimate the mean (\( \mu \)) and variance (\( \sigma \)) parameters of the latent distribution. A vector \( z \), sampled from the estimated distribution, is further concatenated with the prior distribution vector and then processed through two fully-connected layers for reconstruction. The original loss function of CVAE \( L_{CVAE} \) (Kingma et al., 2014) is employed, which consists of the reconstruction error and the Kullback-Leibler (KL) divergence, as follows:

\[
L_{CVAE} = - \log p(h_{span}|z) \\
+ KL[q(z|h_{span})||p(z)].
\]

The prior distribution vector is prepared as trainable parameters. The vector is obtained by concatenating two one-hot vectors: one representing the labels across all datasets, with a “1” at the position corresponding to each label, and the other representing the corresponding (shared) label in the target dataset, including a negative label, with a “1” for each shared label.

As a simple example, consider the case where a target corpus tags Chemical\(_T\) and Protein\(_T\), and a source corpus tags Disease\(_S\) and Protein\(_S\), with Protein being a shared entity type across datasets. The first component of the prior distribution vector represents the labels across all datasets. For this example, this component would form a 6-dimensional space including negative labels: [O\(_T\), Chemical\(_T\), Protein\(_T\), O\(_S\), Disease\(_S\), Protein\(_S\)]. The second component represents the corresponding (shared) label in the target dataset including a negative label. For this example, this component would form a 3-dimensional space: [O, Chemical\(_T\), Protein\(_T\)]. Consequently, the entire prior distribution vector is expressed within a combined 9-dimensional space. For instance, for an entity labeled Protein\(_S\) from the source corpus, its prior distribution vector would be [0, 0, 0, 0, 0, 1, 0, 0, 1]. Here, the “1’s” are placed at the dimension corresponding to Proteins in the corpus-specific space and at the dimension corresponding to the shared Protein\(_T\) label in the shared target space.

Finally, the weighted sum of the two losses is used as the entire loss \( L \) for training:

\[
L = L_{CE} + \alpha L_{CVAE}.
\]

It is important to note that CVAE is not used during inference since it requires correct labels. Instead, the trained classifier alone is used to recognize named entities. Our method thus requires no additional computation during inference, making it suitable for practical applications.

4 Experimental Setup

We describe the datasets and models used in the experiments. Detailed settings and hyperparameters are provided in Appendix A.

4.1 Datasets

We employed 10 labeled biomedical datasets following Luo et al. (2023). The number of documents and the named entities labeled in these datasets are shown in Table 6 of Appendix B. We used the development set of BioRED for evaluation during training, while the development sets of other datasets were used as training data. Unlike Luo et al. (2023), which edited the spans and target labels of entities in additional datasets to follow the BioRED standards, we utilized the original datasets without any such edits.

4.2 Prior distribution vector

As explained in §3, the prior distribution vector provided as a condition to CVAE is constructed
<table>
  <tr>
    <th>Model</th>
    <th>Encoder</th>
    <th>All</th>
    <th>Disease</th>
    <th>Chemical</th>
    <th>Gene</th>
    <th>Variant</th>
    <th>Species</th>
    <th>Cell line</th>
  </tr>
  <tr>
    <td><b>Luo et al. (2023)<sup>3</sup></b></td>
    <td></td>
    <td>91.26</td>
    <td>88.07</td>
    <td>90.98</td>
    <td>92.40</td>
    <td>88.51</td>
    <td>97.50</td>
    <td>90.53</td>
  </tr>
  <tr>
    <td>Single</td>
    <td></td>
    <td>87.86</td>
    <td>82.33</td>
    <td>89.42</td>
    <td>89.63</td>
    <td>81.38</td>
    <td>96.53</td>
    <td>80.90</td>
  </tr>
  <tr>
    <td>Multi</td>
    <td>BERT</td>
    <td>93.22</td>
    <td>90.66</td>
    <td><b>93.41</b></td>
    <td>95.18</td>
    <td>86.84</td>
    <td>97.96</td>
    <td>79.55</td>
  </tr>
  <tr>
    <td>Single + CVAE</td>
    <td></td>
    <td>87.28</td>
    <td>81.93</td>
    <td>89.76</td>
    <td>89.64</td>
    <td>77.80</td>
    <td>94.74</td>
    <td>75.00</td>
  </tr>
  <tr>
    <td>Multi + CVAE</td>
    <td></td>
    <td>94.19</td>
    <td>91.88</td>
    <td>92.98</td>
    <td><b>96.47</b></td>
    <td>92.80</td>
    <td>97.48</td>
    <td>81.82</td>
  </tr>
  <tr>
    <td>Single</td>
    <td>T5</td>
    <td>89.89</td>
    <td>86.52</td>
    <td>87.99</td>
    <td>91.19</td>
    <td>90.56</td>
    <td>97.73</td>
    <td>86.67</td>
  </tr>
  <tr>
    <td>Multi</td>
    <td></td>
    <td>93.99</td>
    <td>92.82</td>
    <td>92.16</td>
    <td>95.89</td>
    <td>90.07</td>
    <td>98.00</td>
    <td>82.76</td>
  </tr>
  <tr>
    <td>Single + CVAE</td>
    <td></td>
    <td>89.87</td>
    <td>85.81</td>
    <td>89.45</td>
    <td>90.72</td>
    <td>89.18</td>
    <td>97.62</td>
    <td><b>91.84</b></td>
  </tr>
  <tr>
    <td>Multi + CVAE</td>
    <td></td>
    <td><b>94.36</b></td>
    <td><b>92.84</b></td>
    <td>92.89</td>
    <td>95.77</td>
    <td><b>94.29</b></td>
    <td><b>98.23</b></td>
    <td>78.57</td>
  </tr>
</table>

Table 1: F1 scores [%] on the BioRED test set, serving as the primary metric. The ‘Encoder’ column specifies T5-3B for ‘T5’ and PubMedBERT for ‘BERT’. The highest F1 score for each entity type is highlighted in bold.

by concatenating two one-hot vectors. The first vector is a one-hot vector with a dimension of 47, representing all labels listed in Table 6, including a negative label for each dataset. The second vector is a seven-dimensional one-hot vector (six types + negative) representing the shared labels based on the BioRED labels. We use the mapping of Luo et al. (2023), which mapped the labels from the additional datasets to one of the BioRED labels, including the negative label (e.g., treating the Gene and FamilyName labels from NLMGene as equivalent to the Gene label in BioRED). The details of this correspondence are shown in Appendix C.

4.3 Comparison settings

We employed the encoder part of T5 (Text-To-Text Transfer Transformer)-3B (Raffel et al., 2020) as the encoder of our model. We also employed PubMedBERT (Gu et al., 2021) as the encoder for a fair comparison with Luo et al. (2023). We employ the F1 score as our primary metric. We compared the following settings:

• Single
The span-based NER model trained on a single target dataset.

• Multi
The span-based NER model trained on multiple datasets in a pure MTL setting, where separate classification layers are prepared for each dataset while the base encoder layer is shared.

• Multi+CVAE
The span-based NER model with CVAE that incorporates the prior distribution vector described in §4.2.

• Multi+CVAE (Fixed)
The span-based NER model with CVAE that fixes the prior distribution vector after initialization.

• Multi+CVAE (Unshared)
The span-based NER model with CVAE that initializes the prior distribution vector using only the first one-hot vector (labels in each dataset) in §4.2.

• Multi+CVAE (Shared label)
The span-based NER model with CVAE that initializes the prior distribution vector using only the second one-hot vector (shared label information) in §4.2.

• Multi+CVAE (Random)
The span-based NER model with CVAE that initializes the prior distribution vector randomly.

• Multi (Shared label)
The span-based NER model trained on multiple datasets without CVAE, where shared labels across datasets are assigned to the same output layer.

5 Results

The overall performance is summarized in Table 1. Compared with the state-of-the-art (SOTA) model by Luo et al. (2023), our model demonstrates enhanced performance across all label types when using multiple datasets, with the exception of Cell line. This Cell line result can be attributed to the limited number of instances in the test set, as shown in Table 5 of Appendix B, which makes

3The score is taken from the original paper.
<table>
  <tr>
    <th>Model</th>
    <th>All</th>
  </tr>
  <tr>
    <td>Multi+CVAE</td>
    <td><b>94.36</b></td>
  </tr>
  <tr>
    <td>Multi+CVAE (Fixed)</td>
    <td>93.99</td>
  </tr>
  <tr>
    <td>Multi+CVAE (Unshared)</td>
    <td>94.12</td>
  </tr>
  <tr>
    <td>Multi+CVAE (Shared label)</td>
    <td>94.20</td>
  </tr>
  <tr>
    <td>Multi+CVAE (Random)</td>
    <td>94.07</td>
  </tr>
  <tr>
    <td>Multi (Shared label)</td>
    <td>92.52</td>
  </tr>
</table>

Table 2: Ablation study using the T5 encoder.

the evaluation sensitive to variations of individual cases. When trained on a single dataset with CVAE (Single+CVAE), the performance slightly degraded when the BERT encoder was employed. This may be because the representations of different labels are not sufficiently distant due to the interaction between labels caused by the CVAE loss. As for the performance of our models with multiple datasets, both the Multi and Multi+CVAE models achieved better F1 scores regardless of the encoder, compared to the corresponding Single model trained solely on BioRED. Furthermore, the CVAE model improved the F1 score for both encoders, suggesting its potential effectiveness in leveraging CVAE during training. This result indicates that the proposed model could capture some of the shared and unshared information among the datasets. This is further supported by the visualizations in Figure 4 of Appendix D, where instances with shared labels from different datasets have closer embeddings.

The ablation study with modified conditions for the prior distribution vector is shown in Table 2. All settings showed inferior performance compared to the proposed model, demonstrating the effectiveness of using the shared labels and tuning the prior distribution vector during training.

6 Conclusion

This study aimed to alleviate the differences in labeling criteria among different datasets to increase the amount of training data. To achieve this, we proposed a method that incorporates a CVAE-based loss function into a span-based NER model, which considers the differences in labeling criteria without the need for manual edits. Our experimental results showed that our CVAE-based approach can leverage multiple datasets while accommodating their labeling discrepancies for either of the two encoders used.

For future work, we will investigate a fully automated method by making the prior distribution vector independent from a specific dataset and eliminating the need for manual definition of the shared vectors, aiming to develop a model that can improve performance on additional datasets as well and be applied to other datasets. Additionally, to further demonstrate broader generalizability, we will evaluate CVAE with other backbone NER models.

Acknowledgments

The authors gratefully acknowledge the assistance provided by AI tools, specifically GitHub Copilot and Claude 3.5 Sonnet, in supporting our programming and writing processes.

Limitations

Our model is based on the approach proposed by Zhong and Chen (2021), where the maximum span length is set to 10. As a result, our model cannot extract named entities that exceed this length constraint. Additionally, as our model primarily utilizes an encoder-based architecture, it may not be straightforward to adapt it to LLMs. Although our method does not require manual annotation, it still requires manual effort to define the prior distribution vector. Due to this limitation, we have evaluated our method only on the BioRED dataset, and the application to other datasets has not been investigated.

Ethical Considerations

This paper utilizes the pre-trained LLM T5 as the base model, which may contain inherent biases due to its pre-training processes. However, we employ only the encoder part of T5 and fine-tune the model for NER, so such potential biases may be mitigated. Furthermore, we use public datasets only for scientific research purposes.

References

Cecilia Arighi, Lynette Hirschman, Thomas Lemberger, et al. 2017. Bio-id track overview. In BioCreative VI Workshop, pages 28–31, Bethesda, MD, USA. BioCreative.

Rezarta Islamaj Doğan, Robert Leaman, and Zhiyong Lu. 2014. Ncbi disease corpus: A resource for disease name recognition and concept normalization. Journal of Biomedical Informatics, 47:1–10.

Martin Gerner, Goran Nenadic, and Casey M Bergman. 2010. Linnaeus: a species name identification sys-
tem for biomedical literature. BMC bioinformatics, 11(1):1–17.

Yu Gu, Robert Tinn, Hao Cheng, Michael Lucas, Naoto Usuyama, Xiaodong Liu, Tristan Naumann, Jianfeng Gao, and Hoifung Poon. 2021. Domain-specific language model pretraining for biomedical natural language processing. ACM Trans. Comput. Healthcare, 3(1).

Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. 2022. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations.

Rezarta Islamaj, Robert Leaman, Sun Kim, Dongseop Kwon, Chih-Hsuan Wei, Donald C Comeau, Yifan Peng, David Cissel, Cathleen Coss, Carol Fisher, et al. 2021a. Nlm-chem, a new resource for chemical entity recognition in pubmed full text literature. Scientific data, 8(1):91.

Rezarta Islamaj, Chih-Hsuan Wei, David Cissel, Nicholas Miliaras, Olga Printseva, Oleg Rodionov, Keiko Sekiya, Janice Ward, and Zhiyong Lu. 2021b. Nlm-gene, a richly annotated gold standard dataset for gene entities that addresses ambiguity and multispecies gene recognition. Journal of biomedical informatics, 118:103779.

Durk P Kingma, Shakir Mohamed, Danilo Jimenez Rezende, and Max Welling. 2014. Semi-supervised learning with deep generative models. In Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc.

Jiao Li, Yueping Sun, Robin J. Johnson, Daniela Sciaky, Chih-Hsuan Wei, Robert Leaman, Allan Peter Davis, Carolyn J. Mattingly, Thomas C. Wiegers, and Zhiyong Lu. 2016. BioCreative V CDR task corpus: a resource for chemical disease relation extraction. Database, 2016:baw068.

Ling Luo, Po-Ting Lai, Chih-Hsuan Wei, Cecilia N Arighi, and Zhiyong Lu. 2022. Biored: a rich biomedical relation extraction dataset. Briefings in Bioinformatics, 23(5):bbac282.

Ling Luo, Chih-Hsuan Wei, Po-Ting Lai, Robert Leaman, Qingyu Chen, and Zhiyong Lu. 2023. AIONER: all-in-one scheme-based biomedical named entity recognition using deep learning. Bioinformatics, 39(5):btad310.

Monica Munnangi, Sergey Feldman, Byron Wallace, Silvio Amir, Tom Hope, and Aakanksha Naik. 2024. On-the-fly definition augmentation of LLMs for biomedical NER. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pages 3833–3854, Mexico City, Mexico. Association for Computational Linguistics.

Nhung T. H. Nguyen, Makoto Miwa, and Sophia Ananiadou. 2023. Span-based named entity recognition by generating and compressing information. In Proceedings of the 17th Conference of the European Chapter of the Association for Computational Linguistics, pages 1984–1996, Dubrovnik, Croatia. Association for Computational Linguistics.

Hiroki Ouchi, Jun Suzuki, Sosuke Kobayashi, Sho Yokoi, Tatsuki Kuribayashi, Ryuto Konno, and Kentaro Inui. 2020. Instance-based learning of span representations: A case study through named entity recognition. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 6452–6459, Online. Association for Computational Linguistics.

Evangelos Pafilis, Sune P Frankild, Lucia Fanini, Sarah Faulwetter, Christina Pavloudi, Aikaterini Vasileiadou, Christos Arvanitidis, and Lars Juhl Jensen. 2013. The species and organisms resources for fast and accurate identification of taxonomic names in text. PloS one, 8(6):e65390.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. 2020. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1–67.

Nicholas E. Rodriguez, Mai Nguyen, and Bridget T. McInnes. 2022. Effects of data and entity ablation on multitask learning models for biomedical entity recognition. Journal of Biomedical Informatics, 130:104062.

Mohammad Golam Sohrab and Makoto Miwa. 2018. Deep exhaustive model for nested named entity recognition. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pages 2843–2849, Brussels, Belgium. Association for Computational Linguistics.

Laurens van der Maaten and Geoffrey Hinton. 2008. Visualizing data using t-sne. Journal of Machine Learning Research, 9(86):2579–2605.

Xiao Wang, Weikang Zhou, Can Zu, Han Xia, Tianze Chen, Yuansen Zhang, Rui Zheng, Junjie Ye, Qi Zhang, Tao Gui, Jihua Kang, Jingsheng Yang, Siyuan Li, and Chuncai Du. 2023. Instructie: Multi-task instruction tuning for unified information extraction. Preprint, arXiv:2304.08085.

Xuan Wang, Yu Zhang, Xiang Ren, Yuhao Zhang, Marinka Zitnik, Jingbo Shang, Curtis Langlotz, and Jiawei Han. 2018. Cross-type biomedical named entity recognition with deep multi-task learning. Bioinformatics, 35(10):1745–1752.

Chih-Hsuan Wei, Alexis Allot, Kevin Riehle, Aleksandar Milosavljevic, and Zhiyong Lu. 2022. tmvar 3.0: an improved variant concept recognition and normalization tool. Bioinformatics, 38(18):4449–4451.
Chih-Hsuan Wei, Hung-Yu Kao, Zhiyong Lu, et al. 2015. Gnormplus: an integrative approach for tagging genes, gene families, and protein domains. BioMed research international, 2015.

Zexuan Zhong and Danqi Chen. 2021. A frustratingly easy approach for entity and relation extraction. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 50–61, Online. Association for Computational Linguistics.

Wenxuan Zhou, Sheng Zhang, Yu Gu, Muhao Chen, and Hoifung Poon. 2024. UniversalNER: Targeted distillation from large language models for open named entity recognition. In The Twelfth International Conference on Learning Representations.

Mei Zuo and Yang Zhang. 2020. Dataset-aware multi-task learning approaches for biomedical named entity recognition. Bioinformatics, 36(15):4331–4338.
A  Detailed training settings

We listed the number of parameters for the encoder models used in the experiments in Table 3. In experiments using the T5 encoder, we employed Low-Rank Adaptation (LoRA) (Hu et al., 2022) and model parameter quantization during training. In addition, the hyperparameters used in the experiments are summarized in Table 4. In the experiments, we took 16 examples from each dataset. We then performed training with a mini-batch size of 64, effectively including examples from multiple datasets by weight accumulation and combining losses from four datasets before updating the model. The weight of the CVAE loss (i.e., \( \alpha \)) was selected based on the results from Figure 2, which were obtained from experiments using the T5-base model. To ensure optimal training, the model was evaluated on the development set after each epoch, and training was terminated if no improvement in validation performance was observed for five consecutive epochs. All experiments were conducted on a computing infrastructure equipped with four NVIDIA V100 GPUs. All experimental results reported in this paper are based on a single run with a fixed random seed.

<table>
  <tr>
    <th>Model</th>
    <th>#Params</th>
  </tr>
  <tr>
    <td>T5-3B</td>
    <td>876,802,048</td>
  </tr>
  <tr>
    <td>T5-base</td>
    <td>86,043,264</td>
  </tr>
  <tr>
    <td>PubMedBERT</td>
    <td>109,482,240</td>
  </tr>
</table>

Table 3: Numbers of parameters (#Params) for the encoder models used in the experiments.

<table>
  <tr>
    <th>Parameter</th>
    <th>Value</th>
  </tr>
  <tr>
    <td>Learning rate</td>
    <td>7e-4</td>
  </tr>
  <tr>
    <td>\( \alpha \)</td>
    <td>1e-4</td>
  </tr>
  <tr>
    <td>Mini-batch size</td>
    <td>64</td>
  </tr>
  <tr>
    <td>CVAE hidden dim</td>
    <td>150</td>
  </tr>
  <tr>
    <td>Max span length</td>
    <td>10</td>
  </tr>
  <tr>
    <td>Span length embedding dim</td>
    <td>150</td>
  </tr>
  <tr>
    <td>LoRA rank</td>
    <td>32</td>
  </tr>
</table>

Table 4: Hyperparameter settings

<table>
  <tr>
    <th></th>
    <th>Train</th>
    <th>Dev</th>
    <th>Test</th>
  </tr>
  <tr>
    <th>Document</th>
    <td>400</td>
    <td>100</td>
    <td>100</td>
  </tr>
  <tr>
    <td>All</td>
    <td>13,351</td>
    <td>3,533</td>
    <td>3,535</td>
  </tr>
  <tr>
    <td><i>Disease</i></td>
    <td>3,646</td>
    <td>982</td>
    <td>917</td>
  </tr>
  <tr>
    <td><i>Chemical</i></td>
    <td>2,853</td>
    <td>822</td>
    <td>754</td>
  </tr>
  <tr>
    <td><i>Gene</i></td>
    <td>4,430</td>
    <td>1,087</td>
    <td>1,180</td>
  </tr>
  <tr>
    <td><i>Variant</i></td>
    <td>890</td>
    <td>250</td>
    <td>241</td>
  </tr>
  <tr>
    <td><i>Species</i></td>
    <td>1,429</td>
    <td>370</td>
    <td>393</td>
  </tr>
  <tr>
    <td><i>Cell line</i></td>
    <td>103</td>
    <td>22</td>
    <td>50</td>
  </tr>
</table>

Table 5: Detailed statistics of entity counts across train, development, and test sets of the BioRED dataset (Luo et al., 2022), categorized by the entity types.

![Line plot showing F1 scores (%) on the BioRED development set, plotted against the weight \( \alpha \) applied to the CVAE loss, using the T5-base encoder; \( \alpha = 0 \) indicates the F1 score obtained without the CVAE loss.](page_1012_1342_388_246.png)

Figure 2: F1 scores [%] on the BioRED development set, plotted against the weight \( \alpha \) applied to the CVAE loss, using the T5-base encoder; \( \alpha = 0 \) indicates the F1 score obtained without the CVAE loss.

B  Dataset statistics

Table 5 presents the statistics of entity counts in the train, development, and test sets of the BioRED dataset. The statistics of the datasets and the list of named entity labels are summarized in Table 6. Note that all datasets used in this study are in English.

C  Correspondence of labels

The correspondence of the labels of each dataset to the BioRED labels is summarized in Table 7. This correspondence is from Luo et al. (2023). For spans with no correspondence to the BioRED labels, the second one-hot vectors of their prior distribution vectors have no active values.

We visualize the prior distribution vectors after training using t-SNE in Figure 3. The visualization demonstrates that the labels in Table 7 remain clustered around their corresponding BioRED labels, and labels with no correspondence remain far from the BioRED labels. This suggests that the initialization information is mostly preserved throughout the training process.

D  Visualization of instance embeddings

Figure 4 visualizes the embeddings of instances in the development set of the BioRED dataset and the test sets of the nine additional datasets using
<table>
  <tr>
    <th>Datasets</th>
    <th>Size</th>
    <th>Labels</th>
  </tr>
  <tr>
    <td>BioRED (<a href="https://doi.org/10.1093/bioinformatics/btac487">Luo et al., 2022</a>)</td>
    <td>600</td>
    <td><i>Disease, Chemical, Gene, Variant, Species, Cell line</i></td>
  </tr>
  <tr>
    <td>BC5CDR (<a href="https://doi.org/10.1093/bioinformatics/btw370">Li et al., 2016</a>)</td>
    <td>1500</td>
    <td><i>Disease, Chemical</i></td>
  </tr>
  <tr>
    <td>BioID (<a href="https://doi.org/10.1093/bioinformatics/bty102">Arighi et al., 2017</a>)</td>
    <td>570 full<br>694</td>
    <td><i>Cell line</i><br><i>FamilyName, Gene, DomainMotif</i></td>
  </tr>
  <tr>
    <td>GNormPlus (<a href="https://doi.org/10.1093/bioinformatics/btv027">Wei et al., 2015</a>)</td>
    <td></td>
    <td><i>Species</i></td>
  </tr>
  <tr>
    <td>Linnaeus (<a href="https://doi.org/10.1093/bioinformatics/btp472">Gerner et al., 2010</a>)</td>
    <td>100 full</td>
    <td><i>Species</i></td>
  </tr>
  <tr>
    <td>NCBIdisease (<a href="https://doi.org/10.1093/bioinformatics/btu037">Doğan et al., 2014</a>)</td>
    <td>793</td>
    <td><i>DiseaseClass, SpecificDisease, CompositeMention, Modifier</i></td>
  </tr>
  <tr>
    <td>NLMChem (<a href="https://doi.org/10.1093/bioinformatics/btaa067">Islamaj et al., 2021a</a>)</td>
    <td>150 full</td>
    <td><i>Chemical, NonStandardRef, OTHER</i></td>
  </tr>
  <tr>
    <td>NLMGene (<a href="https://doi.org/10.1093/bioinformatics/btaa068">Islamaj et al., 2021b</a>)</td>
    <td>550</td>
    <td><i>Gene, FamilyName, Cell, DomainMotif, ChromosomeLocation</i></td>
  </tr>
  <tr>
    <td>SPECIES800 (<a href="https://doi.org/10.1093/bioinformatics/bts567">Pafilis et al., 2013</a>)</td>
    <td>800</td>
    <td><i>Species</i></td>
  </tr>
  <tr>
    <td>tmVar3 (<a href="https://doi.org/10.1093/bioinformatics/btaa069">Wei et al., 2022</a>)</td>
    <td>500</td>
    <td><i>Gene, Species, Disease, DNAMutation, ProteinMutation, OtherMutation, Cell line, AcidChange, SNP, DNAAllele, ProteinAllele</i></td>
  </tr>
</table>

Table 6: Summary of the datasets used. Datasets with “full” in the size column contain full-text data, while the others contain only abstracts.

![t-SNE visualization of prior distribution vectors](page_320_670_484_388.png)

Figure 3: Visualization of prior distribution vectors after training by t-SNE.

of the nine datasets. As discussed in the Conclusion, this limited improvement is attributed to using prior distribution vectors specifically designed for BioRED, highlighting an area that requires future investigation.

t-SNE (van der Maaten and Hinton, 2008). The instances are taken from 300 sentences for each dataset. Compared to the two figures on the left side by the Multi model, the two on the right side using the proposed Multi+CVAE model show a tendency for <i>Disease</i> and <i>Species</i> instances from different datasets to have closer embeddings. This tendency is particularly pronounced for <i>Species</i>, which may contribute to improved performance for <i>Species</i> in BioRED.

E  Evaluation on the additional datasets

We evaluated both Multi and Multi+CVAE models on nine additional datasets, with results shown in Appendix E. Our proposed method improved performance over the Multi baseline on only two
Figure 4: t-SNE Visualization of instance embeddings from the BioRED development set and additional datasets. Upper left: Disease Multi+CVAE, upper right: Disease Multi, lower left: Species Multi+CVAE, lower right: Species Multi
<table>
  <tr>
    <th>Dataset</th>
    <th>Label</th>
    <th>Corresponding label in BioRED</th>
  </tr>
  <tr>
    <td>BC5CDR (Li et al., 2016)</td>
    <td><i>Disease</i><br><i>Chemical</i></td>
    <td><i>Disease</i><br><i>Chemical</i></td>
  </tr>
  <tr>
    <td>BioID (Arighi et al., 2017)</td>
    <td><i>Cell line</i></td>
    <td><i>Cell line</i></td>
  </tr>
  <tr>
    <td>GNormPlus (Wei et al., 2015)</td>
    <td><i>Gene</i><br><i>FamilyName</i></td>
    <td><i>Gene</i><br><i>Gene</i></td>
  </tr>
  <tr>
    <td>Linnaeus (Gerner et al., 2010)</td>
    <td><i>Species</i></td>
    <td><i>Species</i></td>
  </tr>
  <tr>
    <td>NCBIdisease (Doğan et al., 2014)</td>
    <td><i>DiseaseClass</i><br><i>SpecificDisease</i><br><i>CompositeMention</i><br><i>Modifier</i></td>
    <td><i>Disease</i><br><i>Disease</i><br><i>Disease</i><br><i>Disease</i></td>
  </tr>
  <tr>
    <td>NLMchem (Islamaj et al., 2021a)</td>
    <td><i>Chemical</i></td>
    <td><i>Chemical</i></td>
  </tr>
  <tr>
    <td>NLMGene (Islamaj et al., 2021b)</td>
    <td><i>Gene</i><br><i>FamilyName</i></td>
    <td><i>Gene</i><br><i>Gene</i></td>
  </tr>
  <tr>
    <td>SPECIES800 (Pafilis et al., 2013)</td>
    <td><i>Species</i></td>
    <td><i>Species</i></td>
  </tr>
  <tr>
    <td>tmVar3 (Wei et al., 2022)</td>
    <td><i>DNAMutation</i><br><i>ProteinMutation</i><br><i>OtherMutation</i><br><i>AcidChange</i><br><i>SNP</i><br><i>DNAAllele</i><br><i>ProteinAllele</i><br><i>Disease</i><br><i>Chemical</i><br><i>Gene</i><br><i>Species</i><br><i>Cell line</i></td>
    <td><i>Variant</i><br><i>Variant</i><br><i>Variant</i><br><i>Variant</i><br><i>Variant</i><br><i>Variant</i><br><i>Variant</i><br><i>Disease</i><br><i>Chemical</i><br><i>Gene</i><br><i>Species</i><br><i>Cell line</i></td>
  </tr>
</table>

Table 7: Correspondence of labels of the additional datasets to BioRED labels. The labels <i>DomainMotif</i> in GNormplus, <i>NonStandardRef</i> and <i>OTHER</i> in NLMChem, and <i>Cell</i>, <i>DomainMotif</i>, and <i>ChromosomeLocation</i> in NLMGene have no correspondence to the BioRED labels.

<table>
  <tr>
    <th>Dataset</th>
    <th>Multi</th>
    <th>Multi+CVAE</th>
  </tr>
  <tr>
    <td>BC5CDR (Li et al., 2016)</td>
    <td>91.42</td>
    <td><b>91.59</b></td>
  </tr>
  <tr>
    <td>BioID (Arighi et al., 2017)</td>
    <td>90.71</td>
    <td>90.01</td>
  </tr>
  <tr>
    <td>GNormPlus (Wei et al., 2015)</td>
    <td>80.97</td>
    <td>80.68</td>
  </tr>
  <tr>
    <td>Linnaeus (Gerner et al., 2010)</td>
    <td>94.69</td>
    <td>93.84</td>
  </tr>
  <tr>
    <td>NCBIdisease (Doğan et al., 2014)</td>
    <td>85.77</td>
    <td>84.11</td>
  </tr>
  <tr>
    <td>NLMChem (Islamaj et al., 2021a)</td>
    <td>82.74</td>
    <td>82.12</td>
  </tr>
  <tr>
    <td>NLMGene (Islamaj et al., 2021b)</td>
    <td>85.12</td>
    <td>84.76</td>
  </tr>
  <tr>
    <td>SPECIES800 (Pafilis et al., 2013)</td>
    <td>81.12</td>
    <td>77.64</td>
  </tr>
  <tr>
    <td>tmVar3 (Wei et al., 2022)</td>
    <td>85.29</td>
    <td><b>91.87</b></td>
  </tr>
</table>

Table 8: F1 scores [%] on test data of the additional datasets.

Reviewer paper summaries: The paper describes a method for training on several different datasets in the named entity recognition method.
Since different datasets have different sets of labels, ingenuity is required when training models with several different data sets.
Conventional methods often learn as multi-task learning using different output layers for each dataset, 
but this does not effectively take the relationship of labels between data sets into account.
The authors therefore propose a named entity recognition method with an additional block 
that uses a conditional variational autoencoder (CVAE) with a prior distribution vector as reference data to learn the relationship between labels.
The prior distribution vector consists of two one-hot vectors, one representing the reference label and the other representing the label relationship between the data sets.

In their evaluation experiments, the authors confirmed that their method outperforms the conventional method of Liu et al.
In particular, the authors confirmed that CVAE with the proposed method's characteristic prior distribution vector 
improves performance when training with multiple data sets.

The authors propose a span-based NER method with Conditional Variational Autoencoders at training time only to allow for the incorporation of a prior distribution to inform the model of the shared and unshared labels among datasets under a multi-task learning setting. The prior distribution vector by concatenating two one-hot vectors, representing the available labels and the shared ones. The encoder of T5-3B or PubMedBERT is used for span representation. Experiments are conducted on 10 BioRED datasets, employing F1 as the primary metric. Overall, the proposed model demonstrates enhanced performance across five label types explored (Disease, Chemical, Gene, Variant, Species). Ablations are conducted in terms of training datasets and prior distribution vector conditions.
Summary of strengths: * The following points are considered novel and promising.
  * The authors proposed a method to incorporate the relationship of labels between datasets, 
    which is difficult to learn in multi-task learning using multiple datasets.
  * Mapping labels between data remains a manual process, but not so costly that it is easy to handle even when adding other datasets.
* Experiments confirmed the effectiveness in detail.
  * The authors confirmed performance exceeding that of conventional methods
  * Some careful analysis of the behavior of the method due to differences in the behavior of different 
    labels and the fact that CVAE is not effective when using a single training data set.
* The proposed method and the author's thinking behind it are easy to understand, 
  and the reasons for adding the new component parts are clearly written, making the methodology highly convincing.

- The proposed method effectively balances simplicity and sophistication by incorporating CVAE into a span-based NER model. The design adheres to a straightforward span classification approach, avoiding overly complex architectures while still addressing the intricate challenges of integrating multiple datasets with varying label definitions and annotation criteria.
- The paper presents thorough ablation studies and implementation details.
- Experimental results show that the proposed method achieves consistent improvements across different entity types.
Summary of weaknesses: * The prior distribution vector is one of the most important parts of the proposed method, 
  but the explanation is difficult to understand.
  * First mentioned in line 153-155 and explained in line 169-177, but I could not understand the content so far.
    I only understood it when I read this actual case described in section 4.2 (line203-217).
    I think a more detailed explanation is needed, such as giving specific examples around line 153-155 or line 169-177.
* The range of datasets to which the proposed method can be applied seems narrow.
  * This method would be difficult to use in cases where the number of different labels is small 
    and labels in other datasets cannot be successfully assigned, such as NLMGene, for example.
  * However, this is a short paper, and we consider the limited scope of application to be a minor point.

- Quality presentation must be revisited. The writing suffers from inconsistencies, such as redundant phrases (e.g., "multiple datasets" repeated excessively in a short span), acronyms that are introduced but inconsistently used or reintroduced (e.g., VAE at L57 and L104, LLMs at L15 and L314), and typographical errors (e.g., "encocoder" at L257). Additionally, the captions for tables and figures are either duplicated (Table 3 and Table 4) or insufficiently descriptive (Figure 1).
- The claim that supervised learning or fine-tuning is still essential for achieving high NER performance (L14–17) appears too hard and disregards recent advancements in state-of-the-art methods, including modern NER models leveraging auxiliary information for zero/few-shot learning settings. More contextualization and discussion are required for statement support.
- The Related Work section (particularly §2.2) lacks depth in discussing prior multi-dataset NER methods. More detailed comparisons to recent advancements, including motivations for choosing CVAE over other approaches, would strengthen the context.
- No examples are provided to illustrate the role of the prior distribution vector. Moreover, alternative methods for constructing the prior vector are not explored.
- The authors do not provide a strong justification for selecting the T5 encoder stack over other state-of-the-art architectures optimized for biomedical NER tasks, which weakens the rationale for architectural decisions.
- The experiments rely on a single baseline (Luo et al., 2023), limiting comparative insights. The reported results are based on single runs without specified random seeds, and there is no discussion of statistical significance for performance differences.
-  The absence of any commitment to public code release raises concerns about reproducibility, a critical aspect of modern NLP research.
- The evaluation is restricted to biomedical datasets, particularly BioRED, which limits generalizability to other domains.
Comments, suggestions and typos: * The key features of the proposed method is that it adds a CVAE block to the model 
  and uses the corresponding label as the reference data, 
  making it possible to learn the relationship between labels even between different datasets.
  In the experiments, the authors compared the presence or absence of the corresponding label when using CVAE, 
  but the authors did not conduct an experiment on the presence or absence of CVAE when using the corresponding label.
  For example, have there been any experiments using the corresponding labels as reference data for training the output layer 
  without using CVAE?
  * I think this kind of experimentation will make it clear that the addition of CVAE will allow for better learning of the corresponding label.
* Regarding the ablation study in Table 2, it is recommended to also experiment with Multi+CVAE (Only shared) 
  which uses only thecorresponding label for training CVAE.
* I think it is possible to train other datasets than BioRED using the corresponding label mapped to labels used in BioRED as the prior distribution vector.
  In particular, all of the corresponding labels in BioRED are used for tmVar3, 
  and I think that good results will be obtained when experiments are performed.
  * If this can improve performance, it would mean that the proposed method can be effectively utilized by mapping to 
    superordinate concept labels rather than to the labels of the dataset being evaluated, which would likely broaden the range of applications.
  * Of course, this is a short paper and I am not asking you to add this experiment to this paper.
    I am writing this because it is something that interests me as a reader.
* The authors mention the small number of training data as the reason for the low classification performance of the label of "Cell line", 
  but Table 5 only lists the entity distribution in BioRED.
  I think that the entity distribution in all datasets used for training should also be compared.

- Including additional baselines, such as ModernBERT, would strengthen the experimental validity.
- Given the declared hardware setup and the limited size of the T5-3B encoder stack, why the authors relied on adapter fine-tuning?
Overall assessment scores: 3.5

3
Metareview score: 4
Metareview: The paper introduces an approach to span-based Named Entity Recognition (NER) by integrating a Conditional Variational Autoencoder (CVAE) to facilitate training on multiple datasets with different label sets. The CVAE incorporates a prior distribution vector, composed of two one-hot vectors representing available and shared labels, to inform the model of label relationships across datasets. This vector is concatenated with the span representation to estimate the latent distribution, from which a sample vector is drawn to compute the CVAE loss using reconstruction error and Kullback-Leibler divergence. The proposed model employs T5-3B or PubMedBERT as the encoder and is evaluated on 10 BioRED datasets, with F1 as the primary metric. Experimental results show that the proposed method outperforms conventional multi-task learning approaches by effectively leveraging label relationships, leading to improved performance across five entity types: Disease, Chemical, Gene, Variant, and Species. Ablation studies further validate the contribution of the prior distribution vector in enhancing model performance when training with multiple datasets.
Summary of reasons to publish: - The method effectively balances simplicity and sophistication by integrating CVAE into a span-based NER model.
- The method efficiently integrates multiple datasets without adding computational overhead during inference, as the CVAE is only used during training.
- The reasoning behind the method and its components is clearly explained, making it easy to understand and highly convincing.
- Experimental results demonstrate that the proposed method consistently improves performance across different entity types.