mucAI at WojoodNER 2024: Arabic Named Entity Recognition with Nearest Neighbor Search

TL;DR

This work tackles Arabic NER on the WojoodFine corpus, addressing morphology, lack of capitalization, and spelling variants by augmenting a fine-tuned AraBERTv02 model with a KNN-based inference component. The model uses two MLP heads to predict main entity types and subtypes, complemented by a KNN datastore built from training representations to produce a KNN-derived label distribution $P_{kNN}$. The final predictions interpolate the fine-tuned distribution $P_{main}$ with $P_{kNN}$ via $P_{final}=\lambda P_{main}+(1-\lambda)P_{kNN}$, with neighbors weighted by $\exp(sim(e_i,k)/\tau)$ where $sim(a,b)=\frac{a\cdot b}{|a||b|}$. On WojoodFine, this approach achieves about 91% micro-F1 on the test set, surpassing peers and ranking first, while requiring no further training beyond the initial fine-tuning and introducing additional but manageable inference-time KNN search. This demonstrates the practical value of caching training representations to boost NER performance for morphologically rich languages with fine-grained subtype annotations.

Abstract

Named Entity Recognition (NER) is a task in Natural Language Processing (NLP) that aims to identify and classify entities in text into predefined categories. However, when applied to Arabic data, NER encounters unique challenges stemming from the language's rich morphological inflections, absence of capitalization cues, and spelling variants, where a single word can comprise multiple morphemes. In this paper, we introduce Arabic KNN-NER, our submission to the Wojood NER Shared Task 2024 (ArabicNLP 2024). We have participated in the shared sub-task 1 Flat NER. In this shared sub-task, we tackle fine-grained flat-entity recognition for Arabic text, where we identify a single main entity and possibly zero or multiple sub-entities for each word. Arabic KNN-NER augments the probability distribution of a fine-tuned model with another label probability distribution derived from performing a KNN search over the cached training data. Our submission achieved 91% on the test set on the WojoodFine dataset, placing Arabic KNN-NER on top of the leaderboard for the shared task.

Figures (4)

• Figure 1: Examples from WojoodFine illustrating (a) same main entities with different subtypes. (b) main entity with zero and multiple subtypes.
• Figure 2: Distribution of NER tags in WojoodFine Subtask1 (i.e., FlatNER) across the training, development splits.
• Figure 3: The proposed Model workflow for flat NER with subtypes, jointly fine-tuned with two MLP heads for main entity and subtype prediction. KNN search is applied during inference to enhance prediction accuracy.
• Figure 4: Sensitivity of KNN search to number of neighbors ($N$), and interpolation factor $\lambda$.