ToMMeR -- Efficient Entity Mention Detection from Large Language Models

AI summary

Overview

Research area: Natural Language Processing — information extraction, specifically untyped entity mention detection and named entity recognition (NER), combined with mechanistic interpretability of large language models.

Technical level: Intermediate. The paper is readable without deep mathematical background, but assumes familiarity with transformer architectures, tokenization, probing classifiers, recall/precision/F1, and NER terminology.

Scope: The paper introduces ToMMeR, a sub-300K-parameter probing model that recovers entity mention spans from the early hidden layers of a frozen LLM backbone in a single partial forward pass, and shows it transfers zero-shot across 13 NER benchmarks.

What This Paper Is About

Named entity recognition systems usually bundle two different things into one task: deciding where a mention begins and ends (mention detection), and deciding what type it is (entity typing). The authors argue this conflation hides a basic question — where and how do models learn to detect span boundaries? Their goal is to show that mention detection is a schema-invariant capability that already exists inside pretrained language models and can be pulled out cheaply, without prompting, fine-tuning the backbone, or specifying a schema at inference time.

Key Contributions

1. A lightweight probing architecture (ToMMeR). ToMMeR scores spans using pairwise "matching scores" between token representations plus token-level values, probed from the early layers of any frozen LLM. It adds under 300K parameters (274K in the main configuration), requires only one partial forward pass, and uses no prompting, schema specification, or text generation.

2. Evidence that mention boundaries are a convergent, emergent capability. Across LLMs from 14M to 15B parameters — including Pythia, LLaMA 3, Mistral-7B, and Phi families, plus encoder-only models like BERT, RoBERTa, and ModernBERT — auto-regressive backbones converge on similar mention boundaries (Sørensen–Dice scores above 0.75) without any shared supervision.

3. A precision evaluation protocol for untyped mention detection. Because typed benchmarks under-label legitimate mentions (nested, generic, pronominal), the authors estimate real precision with an LLM-as-a-judge protocol calibrated against human annotations and corrected using the Prediction-Powered Inference (PPI) framework.

4. A straightforward extension to full NER. Adding a span classification head (a two-layer MLP with hidden dimension 1024 producing span embeddings) yields competitive full NER performance of 80–87% F1 across standard benchmarks, with released models and code.

Main Findings

• High zero-shot recall at minimal cost: ToMMeR achieves 93% recall zero-shot across 13 NER benchmarks (aggregated recall of 92.6 in Table 1), reaching up to 98.6% recall on MultiNERD with threshold decoding. The overall aggregated mean recall is 88.2 under threshold decoding and 75.3 under greedy decoding.

• Estimated precision around 90%: Under the human-calibrated LLM-judge protocol, PPI-corrected precision is 90.8 ± 2% for MultiNERD and 87.1 ± 5% for GENIA (95% confidence), against raw LLM-judged precisions of 92.8 and 90.3 respectively. This contrasts sharply with the 23.2% aggregated precision measured against gold benchmark annotations, which the authors attribute to benchmarks labeling only schema-specific entities.

• Low precision on gold data is expected, not a failure: Because ToMMeR detects mentions beyond standard benchmark types, precision against gold NER labels is low (aggregated 23.2%). Greedy (flat, non-overlapping) decoding raises average precision from roughly 30% to about 40%.

• Mention signals appear in the earliest layers: Probing each of the 16 layers (0–15) of Llama-3.2-1B shows recall, precision, and F1 nearly reach their maximum from layer 0 onward, with a noticeable performance loss at the final layer, where the model likely discards these signals to prioritize next-token prediction.

• Different architectures share a notion of mention: Auto-regressive backbones (excluding small Pythia models below 160M parameters) reach Dice scores above 75% agreement on MultiNERD; encoder-only models differ by roughly 10%, which the authors hypothesize relates to bidirectional attention suppressing overlapping spans.

• Multilingual transfer works on Latin scripts: ToMMeR achieves 86% average recall on French, Spanish, and German in WikiANN, but drops to 41% recall on Chinese, attributed to tokenization and syntax encoding differences in the English-centric Llama backbone.

• ACE 2005 is the main outlier: Recall falls to 42% (versus 88.2% average), because ACE includes determiners (e.g., "the president") and pronouns that ToMMeR did not learn from Pile-NER annotations.

• Competitive full NER when extended with a typing head: Using 7.4M–7.9M trained parameters, ToMMeR plus span embedding reaches F1 scores such as 84.8 on CoNLL 2003 (LLaMA-3.2-1B, layer 6), 86.8 (LLaMA-3.2-3B, layer 5), 87.3 (RoBERTa-base, layer 5), and 92.2–93.3 on MultiNERD. On OntoNotes, the LLaMA-3.2-1B-based ToMMeR reaches 80.4% F1 despite only 73% zero-shot recall before fine-tuning. Reported baselines are GLiNER (209M trained parameters: 88.7 CoNLL, 78.9 GENIA, 93.8 MultiNERD, 89.0 OntoNotes, 87.8 ncbi) and EMBER (11.5M trained parameters: 85.1 CoNLL, 79.3 OntoNotes).

• Large speed advantage over prompting: On a sample passage with llama3.2-1B, a UniversalNER-style prompting baseline took 3.21 seconds (24GB GPU, HF transformers, 100 max generated tokens), while ToMMeR inference took 0.077 seconds — a 42× improvement. Cutting computation at a lower layer (e.g., layer 6 of 16 for Llama-3.2-1B) gives roughly a 2.6× speedup over a full forward pass.

• Human validation of the judge: Human–LLM agreement reaches 91.5% (κ = 0.449) on MultiNERD and 78.5% (κ = 0.239) on GENIA, with 1,200 and 600 annotated spans respectively, collected from 5 CS researchers. Aggregated agreement is 87.17% with κ = 0.359 and a bias correction of −3.06%. The LLM judge accepts gold-annotated entities in more than 99% of cases.

Methodology in Plain English

The authors treat mention detection as binary span classification. For a sentence, they enumerate contiguous spans up to 25 tokens long (a limit that covers 99.8% of mentions annotated in Pile-NER) and ask the model whether each span is a valid mention.

ToMMeR reads token representations from an early layer of a frozen LLM and computes two things. First, a matching score between every pair of tokens, produced by learned query and key projections into a small rank-r subspace, normalized with L2 instead of softmax so it behaves as a cosine similarity rather than a probability distribution. Second, token values, a learned linear probe on each token's representation. A logistic model with only five parameters combines the matching score of a span's endpoints, max and min pooling over the internal matching scores, and the values of the span's last token and its immediate context. This lets the model capture both which tokens bind together and where a span ends.

Training uses Pile-NER, a dataset of 45,889 samples from The Pile annotated with fine-grained entity types by GPT-3.5. Because non-entity spans vastly outnumber entity spans, the loss is Balanced Binary Cross-Entropy, which reweights the classes by a factor computed per batch as the number of negatives divided by the number of positives. To handle Pile-NER's incompleteness (for example, nested mentions are not labeled), training is two-stage: a first fit on available annotations, then self-distillation where the learned model augments the training data with mentions that were not annotated, reducing false negatives.

For evaluation, the authors run zero-shot mention detection on 13 NER corpora and also test greedy decoding, which iteratively picks the highest-scoring span that does not overlap previously chosen spans. Because gold benchmarks cannot distinguish true false positives from out-of-scope entities, they sample 10,000 spans and prompt gpt-4.1-mini to judge whether each predicted span qualifies as an entity under the Wikipedia definition, then calibrate those judgments with human annotations and correct residual bias using Prediction-Powered Inference. Finally, they freeze ToMMeR's detection stage and train a span classification head for full NER, which takes less than 20 minutes on a V100-32Gb for CoNLL with Llama-3.2-1B representations.

The architecture is grounded in the "binding ID" framework from mechanistic interpretability, which posits that transformers dynamically bind related tokens through learned signals. The authors hypothesize that LLMs implicitly group entity-mention tokens the same way, and that these latent signals are what ToMMeR reads out.

Why This Matters

For research, the paper reframes mention detection as a separable, largely solved-by-pretraining capability rather than a task requiring hundreds of millions of task-specific parameters. It provides evidence that structured entity representations exist in early transformer layers and are linearly recoverable, which connects to mechanistic interpretability work on binding, induction heads, and syntactic structure in hidden states. It also supplies a workable evaluation methodology for untyped mention detection, where no ground truth standard exists.

Real-world applications:

• Retrieval-augmented generation pipelines: ToMMeR can act as a first-stage mention filter, enabling entity-aware document chunking and supplying spans for downstream typing and linking, at a fraction of the cost of re-prompting an LLM.
• Streaming and low-latency extraction: The single partial forward pass allows entities to be identified as an LLM generates text with minimal latency penalty, useful for real-time systems.
• Multi-schema and multi-domain information extraction: Because boundaries are schema-invariant, one detection stage can feed different typing heads or knowledge bases without retraining the detector per schema.
• Cost-sensitive, high-volume processing: With 274K parameters and 0.077-second inference versus 3.21 seconds for a prompting baseline, the method suits large-scale or resource-constrained deployments.

Industry relevance centers on cost and adaptability: the detector is trained once on annotation boundaries, plugs into any frozen LLM backbone, and can be fine-tuned in minutes to match a new dataset-specific notion of mention, which lowers the barrier to building modular extraction pipelines instead of schema-specific NER models.

Future Directions

• Definitive ground truth for untyped mention detection. The authors state that comprehensive human annotation would be needed to definitively establish ToMMeR's precision, although their PPI analysis bounds true precision above 89%.

• Handling discontinuous and non-contiguous spans. The current formulation is restricted to contiguous spans of up to 25 tokens, excluding phenomena such as "New York" and "City" separated by other tokens.

• Reconciling architectural differences. Encoder-only models produce systematically different mention sets (roughly 10% Dice drop) and predict fewer nested entities; understanding why, and whether bidirectional context can be exploited within the probing framework, remains open.

• Reducing dependence on Pile-NER's biases. Since training annotations come from GPT-3.5 in 2023, the model inherits their biases and gaps, and the self-distillation strategy risks amplifying systematic bias present in the initial annotations.

• Improving weak-schema alignment. ACE 2005 recall of 42% shows the learned notion of mention does not match all annotation conventions; further study of how easily ToMMeR can be tuned to such schemes is suggested.

Target Audience

This paper is most useful to NLP researchers and graduate students working on information extraction, NER, and mechanistic interpretability of LLMs; to ML engineers building retrieval-augmented generation or entity-linking pipelines who want a cheap, schema-agnostic mention detector; and to practitioners who need high-recall candidate generation before a downstream typing or linking stage. Readers should be comfortable with standard NER evaluation metrics and basic transformer internals, but the core ideas are accessible without deep mathematical background.

Authors’ abstract

Identifying which text spans refer to entities - mention detection - is both foundational for information extraction and a known performance bottleneck. We introduce ToMMeR, a lightweight model (<300K parameters) probing mention detection capabilities from early LLM layers. Across 13 NER benchmarks, ToMMeR achieves 93% recall zero-shot, with an estimated 90% precision under a human-calibrated LLM-judge protocol, showing that ToMMeR rarely produces spurious predictions despite high recall. Cross-model analysis reveals that diverse architectures (14M-15B parameters) converge on similar mention boundaries (DICE >75%), confirming that mention detection emerges naturally from language modeling. When extended with span classification heads, ToMMeR achieves competitive NER performance (80-87% F1 on standard benchmarks). Our work provides evidence that structured entity representations exist in early transformer layers and can be efficiently recovered with minimal parameters.