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

Summary

The paper tackles fine-grained flat named entity recognition for Arabic in subtask 1 of the WojoodNER 2024 shared task. Arabic NER is hard because the script has no capitalization, the morphology is rich, spellings vary, and labeled data is scarce. The authors fine-tune AraBERTv02 on WojoodFine with two jointly trained heads. One head is a softmax over main entity tags in BIO format. The other is a multi-label sigmoid head over 31 subtypes. After training, they cache the contextual vector and gold main label of every training token in a datastore. At inference, each test token retrieves its nearest cached neighbors by cosine similarity. The neighbors' labels form a second probability distribution, which is linearly mixed with the model's own prediction. This step adds no training. It lifts test micro-F1 from 90 to 91 and puts the system first on the leaderboard. The work shows that a cheap, retrieval-based correction layer can improve an already fine-tuned tagger.

Key points

• Data: WojoodFine extends Wojood, which has about 550K manually annotated tokens and 21 entity types, mostly MSA with some Palestinian and Lebanese dialect.
  • It adds 31 subtypes attached to GPE, ORG, LOC and FAC.
  • A single entity can carry up to two levels of subtypes.
• Model heads:
  • The main-tag head has 43 outputs (O plus B-/I- for each of the 21 tags), with softmax and cross-entropy.
  • The subtype head has 62 sigmoid outputs, trained with binary cross-entropy and thresholded at 0.5.
  • KNN was applied only to main tags.
• Results (Table 1):
  • Test F1 rose from 90.00 to 91.00.
  • Dev F1 rose from 91.87 to 92.15. The conclusion reports different dev figures (91.62 → 92.90), an inconsistency within the paper.
  • On the leaderboard, mucAI scored 91, muNERa and Addax 90, and the baseline 89.
• KNN settings: the neighbor count was tuned over powers of two from 8 to 512, and λ from 0 to 1 in steps of 0.1. Temperature was fixed at 1.
• KNN-only failure mode: using the retrieved distribution alone is competitive at small neighbor counts. It falls to 88.5 at 512 neighbors, because the abundant O label floods the retrieved set.

Interesting ideas

• I1. Caching every training token's contextual vector with its gold label turns the training set into a lookup memory, which can correct a trained model at inference without retraining.
• I2. Mixing model and retrieval distributions through a single interpolation weight, tuned on dev data, gives a simple, inspectable knob for how far to trust memory over the parametric model.
• I3. Retrieval votes degrade as the neighbor count grows, because frequent "background" classes (here, the O label) swamp rare ones, so the neighbor count must be tuned or the votes class-balanced.
• I4. Main types use a softmax head while subtypes use independent sigmoids, so one token can carry zero or several fine labels without forcing a single choice.

For Mova

How these ideas could serve Mova, our small open language model — its files, stages and tests:

• I1 (tree): gparse could keep a datastore of token embeddings with gold UD relations and blend a neighbor-vote distribution into its arc-label scores, and each retrieved neighbor doubles as a human-readable justification for the chosen label.
• I2 (tokens): the en tagger could interpolate its POS/morphology scores with a nearest-neighbor vote from treebank tokens, and report which source dominated each decision.
• I3 (learn): when witness parsers or a KNN vote produce silver labels, Mova should cap or reweight frequent default labels so the silver data does not drift toward the majority class.
• I4 (world): entity typing against WordNet/Wikidata classes could use a single main type plus a multi-label set of finer ontology tags, matching how real entities belong to several categories at once.