Embedded Named Entity Recognition using Probing Classifiers (2024)

Abstract: Streaming text generation, has become a common way of increasing the responsiveness of language model powered applications such as chat assistants. At the same time, extracting semantic information from generated text is a useful tool for applications such as automated fact checking or retrieval augmented generation. Currently, this requires either separate models during inference, which increases computational cost, or destructive fine-tuning of the language model. Instead, we propose an approach called EMBER which enables streaming named entity recognition in decoder-only language models without fine-tuning them and while incurring minimal additional computational cost at inference time. Specifically, our experiments show that EMBER maintains high token generation rates, with only a negligible decrease in speed of around 1% compared to a 43.64% slowdown measured for a baseline. We make our code and data available online, including a toolkit for training, testing, and deploying efficient token classification models optimized for streaming text generation.

## introduction combining pre-trained language models (lms) and external information at inference time is a widely used approach, for example as a means of improving the factual accuracy of generated texts in knowledge-intensive tasks @xcite @xcite . to effectively gather relevant information, there are generally two main strategies for extracting semantic data from the current context, each with its own drawbacks. the first strategy involves integrating the extraction process into text generation. this method, as seen in work by @xcite and @xcite , requires generating queries during the 1 https://github.com/nicpopovic/ember 2 https://github.com/nicpopovic/stoke figure 1 : ember enables simultaneous text generation and entity annotation by using a language model's internal representations as the feature space for classification. compared to using state-of-the-art ner models, this results in a substantially more efficient pipeline allowing for streaming named entity recognition. parameter and latency comparisons stated in this figure are based on the experiments conducted using gpt-2 xl , presented in section 6. inference phase. although this approach is direct, it has the downside of altering the lm through fine-tuning, which can lead to issues such as catastrophic forgetting @xcite ). the second strategy employs an external system for information extraction @xcite . while studies such as those by @xcite , @xcite , and @xcite show promising results, the required computational overhead is a major issue hindering adoption @xcite . for many applications (chat assistants, etc.), the delivery of generated text on a token-by-token basis as soon as they are available, known as streaming text generation, has become a common way of increasing responsiveness. an optimized solution for this setting is currently missing. meanwhile, research into the mechanistic interpretability of lms has shown that substantial semantic information can be recovered from individual internal representations. a common diagnostic tool are simple 3 classifiers, called probing classifiers @xcite , trained to perform specific tasks using a subset of the internal representations of a (frozen) lm as their feature space. while their validity as a means for understanding how and where information is stored in lms is debated @xcite , probing classifiers have been shown able to map internal representations to syntactic and semantic information @xcite @xcite @xcite . while the majority of this research has been conducted using encoder lms, studies have shown that similar information is recoverable from specific internal states of decoder-only lms @xcite @xcite . we therefore explore whether, rather than as a diagnostic tool, probing classifiers can be used for non-destructive, light-weight, and continuous ie in decoder-only lms at inference time. in this work, we develop an approach we call embedded named entity recognition (ember) for performing named entity recognition (ner), a central ie subtask consisting of mention detection and entity typing, using only a lm's internal representations as feature space without further finetuning thereof. as illustrated in figure 2 , the process involves two probing classifiers: the first performs tokenwise type classification based on the hidden state at a single transformer sublayer, while the second detects spans based on the lms attention weights between two tokens. finally, the outputs of both are fused into span-level entity predictions. we conduct a series of experiments using multiple lms, ner datasets, and task settings to evaluate the performance of ember and the factors which 3 typically small in the amount of trainable parameters and less complex in terms of architecture relative to the lm. influence it. in short, we find that while outperformed by finetuned state-of-the-art approaches in terms of raw benchmark scores (∼80 -85% f1 vs. >90% f1), our approach outperforms few-shot incontext learning approaches (∼50% f1) (section 5.2) and is significantly more efficient in the streaming text generation setting (approx. 80× faster than baseline (section 6)). in conclusion, we make the following contributions: we propose ember, the first nondestructive ner approach optimized specifically for use with decoder-only lms during streaming text generation. we show that our approach can achieve f1 scores of 80-85% while requiring minimal additional computational overhead. we provide insight into which architecture parameters of decoder-only lms determine how well our approach will work. lastly, we showcase efficient simultaneous text generation and ner, a novel usecase our approach is optimized for and provide a toolkit for training, testing, and deploying models. ## ember in this section we introduce our approach for building a ner system based on probing internal representations of pretrained, decoder-only lms. given a model m with l layers, a hidden state h l i is the output at a single transformer sublayer, where @xmath0, ..., l] is the index of the sublayer and i is the index of the input token. the attention weights foot_1 between two tokens are denoted as a j,i , where @xmath1. our approach entails two key steps, namely tokenwise entity type classification based on h l i (4.1) and span detection based on a j,i (4.2), the results of which are then combined to form a complete ner pipeline using a mechanism we call label propagation (4.3). ## experiments: non-streaming ner before examining the novel task setting of streaming named entity recognition (ner), we conduct experiments to determine how well ember performs in a variety of typical, non-streaming ner settings. we begin by evaluating which label propagation strategies work best (5.1). after identifying the best configuration, we evaluate its performance in the supervised learning setting (5.2). next, we analyse the measured results with respect to the effects of model scale and architecture parameters (5.3). finally, we evaluate our approach in heavily data-constrained settings (5.4) and show the efficient extraction of named entities during streaming text generation (6). ## conclusion we present ember, a lightweight approach for embedding ner capabilities into decoder-only lms without finetuning them. we find that, except in highly data-constrained settings (such as 1-shot or 5-shot), it surpasses in-context learning in classification accuracy while being significantly more efficient. our approach enables efficient simultaneous text generation and ner, with only a 1% reduction in token generation rate and less than 1% increase in model size, paving the way for novel applications such as a significantly more efficient integration of external structured knowledge into text generation. lastly, we include detailed observations about the factors which influence ember's performance and provide a toolkit for training, testing, and deploying models.