UNER: A Unified Prediction Head for Named Entity Recognition in Visually-rich Documents

Assuming that we need to extract $C$ types of entities from the document. when using a sequence-labeling-based (SL) head with the BIO-tagging scheme, the entity extraction task is treated as a $(2C+1)$-way classification problem for each token embedding $\mathbf{e}_{i}$:

The function $\mathcal{F}_{sl}$ maps the token embedding $\mathbf{e}_{i}$ to a BIO-tagging label. However, this approach has a limitation in that it can only handle a fixed number of pre-defined entity types.

To address this issue, we propose the Query-aware Token Classification (QTC) module. This module first uses a query encoder to encode the entity names or descriptions as classification queries. It then leverages these queries as semantic clues to perform binary classification for all tokens.

To handle $C$ entity types, we start by encoding each entity name using a query encoder. This encoder transforms each entity name into a $h$-dimensional embedding, resulting in a query embedding matrix $\mathbf{Q}\in\mathcal{R}^{C\times h}$. In this matrix, each row corresponds to an entity type, and the $j$-th entry of $\mathbf{Q}$, denoted as $\mathbf{q}_{j}\in\mathcal{R}^{h}$, represents the embedding of the $j$-th query.

Then we combine the token embedding matrix and query embedding matrix to generate the query-aware token embedding tensor: $\mathbf{E}\oplus\mathbf{Q}\to\mathbf{S}\in\mathcal{R}^{L\times C\times h}$, where $\oplus$ stands for the broadcastable addition operation. We use $\mathbf{s}_{i,j}$ to denote $(i,j)$-th entry of $\mathbf{S}$, which is a query-aware token embedding for the $i$-th token and $j$-th query and satisfies: $\mathbf{s}_{i,j}=\mathbf{e}_{i}+\mathbf{q}_{j}$. For each $\mathbf{s}_{i,j}$, we will conduct a binary classification to determine if the $i$-th token belongs to the $j$-th entity type:

We will conduct $L\times C$ binary classifications for all token-query pairs and the total classification loss is calculated as :

Here $\mathrm{CE}(\cdot,\cdot)$ is the cross-entropy loss function and $\hat{s}_{i,j}\in\{0,1\}$ is the ground-truth label.

In QTC, the shape of the query-aware token embedding tensor is dynamically determined by the number of input queries. This flexibility allows us to input any number and type of queries into the module, thereby avoiding the limitation of pre-defined entity types in the SL head.

Moreover, as we use independent binary classifications, each token can be possibly classified into multiple categories. Such a design makes the UNER head capable of extracting nested and overlapped entities. While these entities are not fully represented in current VrD-NER benchmarks, they are common in real-world applications, as users may have different extraction goals within the same document and the same token may belong to multiple entity types. For instance, in Figure-1 case (1), a UNER-based model can simultaneously extract the full address (denoted in green), the city name (“Batang Kal”), and the state name (“Selangor”) by using “address”, “city”, and “state” as input queries. In this case, these entities are overlapped and cannot be extracted with an exclusive $C$-way classification head.

Furthermore, the QTC module utilizes entity names as semantic clues for classification, which enables it to extract unseen entity types through semantic knowledge transfer.

After token classification, it is necessary to arrange the tokens in order to form correct entity spans. In the SL head, the BIO-tagging scheme labels the “Beginning” and “Inside” information, which helps in token arrangement. However, it fails to handle discontinuous entities or tokens that are incorrectly ordered in VrDs. To overcome this limitation, we introduce a Token Order Prediction (TOP) module that predicts the token orders in VrDs.

The TOP module treats the token order prediction task as an edge prediction task on a token graph. For a document with $L$ tokens, we represent its token order using a directed graph with $L$ vertices. This graph is represented by a binary matrix $\mathbf{G}\in\{0,1\}^{L\times L}$. If the $k$-th token is the successor of the $i$-th token, the $(i,k)$-th entry of $\mathbf{G}$, which represents the directed edge from $i$ to $k$, is set to 1. Otherwise, it is 0. By formulating the task in this way, the token order prediction becomes equivalent to predicting the edges on the graph.

During training, TOP performs binary classification for all directed edges in the graph:

The objective of this task is to optimize the following loss function:

where $\hat{g}_{i,k}\in\{0,1\}$ represents the ground-truth label for each edge in the graph.

It is important to note that $\hat{g}_{i,k}$ is derived from the entity annotations in the VrD-NER dataset, so it only covers the tokens that are part of entities. The order of the non-entity tokens is unknown during training, and therefore, the corresponding losses in Equation-5 should be ignored. Fortunately, after training with annotated documents, the knowledge of reading order learned from the entity-related tokens can be transferred to non-entity tokens and tokens in other types of documents. Such knowledge transfer is discussed in detail in Section-4.4.

The total training loss of the UNER head can be represented as follows:

where $\lambda$ is a hyper-parameter that controls the balance between the two losses. As shown in Figure-2, the ground truth labels in $\mathcal{L}_{qtc}$ and $\mathcal{L}_{qtc}$ are predominantly zeros, with only a small portion being ones. This label imbalance issue poses a challenge for loss optimization, so we employ the ZLPR loss (Su et al., 2022) to enhance the training of the two losses and improve optimization efficiency.

During inference, we utilize the predictions of the QTC and TOP modules to construct full entity spans using the following steps:

1. (1)

   For each entity type (e.g., the $j$-th type), we identify all tokens that meet the condition $\mathcal{F}{qtc}(\mathbf{s}{i,j})>0$, along with their valid token orders satisfying ${\mathcal{F}_{top}(\mathbf{e}_{i},\mathbf{e}_{k})>0}$.

2. (2)

   Based on the token orders, we identify the tokens without any predecessor tokens and use them as the starting tokens to create multiple token series.

3. (3)

   For each token series, we select the best subsequent token for the last token by choosing the token with the highest token order classification score and then adding it to the current token series.

4. (4)

   Repeat step (3) until the last token does not have any valid subsequent tokens, and use the resulting token series to form entity spans.

It should be noted that the UNER head enables multiple entity names to be queried for each document during both training and inference. When provided with a document and $C$ entity names, we only need to compute the document transformer and TOP module once and reuse their outputs to collaborate with the QTC module for $C$ iterations. This feature significantly alleviates the computational burden for recomputing the backbones when handling dense entity extraction tasks and enables a more flexible model deployment.

In our experiments, we utilize 7 commonly used VrD-NER benchmarks to compare the performance of different methods. To evaluate the effectiveness of supervised pre-training in VrD-NER, we incorporate two additional datasets for pre-training, which encompass a wide range of document types and entity types. Detailed information of these datasets can be found in Table-1.

In our experiments, we compare different prediction heads using the base version of LayoutLMv3 and LayoutMask as transformer backbones, following the backbone setting in TPP. The query encoder in the QTC module is a 4-layer transformer with cross-attention to the transformer backbone. It is initialized with the first four layers of Q-Former in BLIP-2 (Li et al., 2023) with 38M parameters. We employ the tokenizer and the 1D-position embedding layer from the transformer backbone to process queries during query encoding. As each query can contain multiple tokens, we use the embedding of the first token as the query embedding. To determine the optimal value of the hyper-parameter $\lambda$, we conduct ablation studies. In our experiments, we set $\lambda=0.1$ and $\lambda=0.5$ for models with or without supervised pre-training, respectively. The batch size is set to be 16 for all fine-tuning datasets. For few-shot experiments, we repeat them five times and report the average scores. During the supervised pre-training stage, both the backbone and the UNER head are trained on the pre-training datasets for 20 epochs before fine-tuning.

It is important to highlight that the FUNSD-r and CORD-r datasets have discontinuous entities due to complex layouts and incorrectly ordered tokens. This significantly impacts the performance of the SL head, as it is unable to handle such entities. Additionally, the SL head is incapable of learning and conveying token order knowledge effectively, which limits its usefulness as a VrD-NER head.

While the TPP head is capable of handling discontinuous entities, it faces optimization issues that negatively affect its performance. This is because the TPP head attempts to jointly optimize token classification and token order prediction in a unified manner. For each entity, it predicts one token graph to represent entity-specific token paths, resulting in the need to predict a $L\times L\times C$ binary tensor and $L^{2}C$ classification sub-tasks. Such entity-specific token graphs in TPP also influence the transfer of token order knowledge among different entity types.

On the other hand, the UNER head has two separate optimization goals and reduces the number of binary classifications sub-tasks to $(L^{2}+LC)$. Additionally, UNER explicitly learns and represents token order knowledge in one unified token graph which is independent of the entity types, making it suited for transferring universal token order knowledge. This design makes the UNER head easier to optimize and more efficient with limited data.

To further highlight this advantage, we compared the TPP and UNER head on the CORD-r dataset in a few-shot setting, where the model is fine-tuned using only a subset of the training samples. The results of this comparison are summarized in Figure-3. The findings from these experiments highlight that UNER consistently outperforms TPP across all training percentages. It is noteworthy that the performance advantage of UNER becomes more pronounced as the percentage of training data decreases. For instance, when utilizing only 5% of the data, UNER achieves a notable improvement of +21.24% compared to TPP. This demonstrates the data efficiency of UNER in handling limited training data.

However, we also observed that the F1 scores of the UNER-based models in the multi-task setting were slightly lower compared to the scores in a single-task setting. We believe that this is due to the difficulty in balancing the number of training samples for different datasets. As our focus is not on searching for the best performance with a particular group of datasets, we simply combined these datasets for training without making careful adjustments in order to showcase the vanilla capability of UNER for unified VrD-NER.

The UNER head has a significant advantage as it can be utilized for supervised pre-training, which further enhances the self-supervised pre-trained backbones. To assess such potential, we leverage two annotated VrD-NER datasets, SVRD and DocILE, to incorporate supervised pre-training before fine-tuning. We use “UNER+LayoutMask” as the base model and compare its fine-tuning performance with and without supervised pre-training. We find that the inclusion of a supervised pre-training stage improves models with notable effectiveness in cross-linguistic and few-shot scenarios and brings zero-shot VrD-NER abilities.

Our observations indicate that supervised pre-training has significant benefits in few-shot scenarios. Specifically, the “UNER+SP” model achieves remarkably high performance even with a small portion of training samples, outperforming the “UNER” model by a large margin on all tasks and configurations. Notably, the f1 score of “UNER+SP” in “VrD-NER” is surprisingly high with extremely limited data (74.62% with only 5% training samples). Furthermore, “UNER+SP” with 5% training samples has a higher accuracy score in “TOP” than “UNER” with 25% samples (98.01% compared to 97.39%). These advantages highlight that the UNER-based model benefits significantly from the document knowledge learned during pre-training on large datasets. This knowledge facilitates the model’s adaptation to the other downstream VrD-NER dataset, particularly when the available data is limited.

Our analysis indicates that the zero-shot entity extraction ability is largely attributed to the inclusion of the SVRD dataset, which significantly improves the QTC for both SROIE (15.14%) and EPHOIE (57.38%). Furthermore, incorporating the DocILE dataset further enhances the VrD-NER abilities, particularly in SROIE where the F1 score increased from 0.07% to 21.90%.

In order to gain a better understanding of the entity extraction ability of the pre-trained model, we tested the best model (pre-trained on both SVRD and DocILE) with various documents and queries, and presented some prediction results in Figure-4. Our observations indicate that the model is effective in extracting entities with explicit trigger words. Using these trigger words as queries consistently yields satisfactory results, regardless of the layout alignment. However, the model exhibits limitations when handling entities without explicit trigger words, unnatural queries, or queries with varying meanings. For example, the entity names in the CORD-r dataset are specifically designed, such as “menu.sub.price” and “total.cashprice”. Similarly, the entity type “Header” is used in multiple datasets, but its meaning can differ across documents, leading to inconsistent predictions. These findings underscore the potential for using diverse pre-training datasets to develop a unified VrD-NER model with stronger zero-shot abilities, while also identifying challenges in transferring document understanding knowledge.