1 Tel Aviv University2 Mila – Quebec AI Institute & McGill University
Cite LMEnt (TACL 2026)
{{bibtex}}
TL;DR
We present LMEnt, a suite including a knowledge-rich pretraining corpus, fully annotated with entity mentions based on Wikipedia; an entity-based retrieval method over pretraining data that outperforms existing tools by as much as 80.4%; and 12 pretrained LMs with up to 1B parameters and 4K intermediate checkpoints, with comparable performance to popular open-source models on knowledge tasks. Below, we cover our entity annotation method, the LMEnt models, and applications of LMEnt to mechanistic interpretability, learning dynamics, knowledge editing, plasticity, attribution, and hallucinations.
Introduction
Language models capture substantial amounts of knowledge and beliefs about the world from their training data. A fundamental, yet underexplored, question is how such knowledge representations are formed and shaped during pretraining. Namely, what is the interplay between data composition, training dynamics, and the internal knowledge mechanisms in LMs?
A key prerequisite for studying this question is the ability to locate exactly where specific knowledge appears in the pretraining data. However, since such metadata is rarely available for pretraining corpora, many works that study knowledge acquisition have used synthetic data, or post-hoc string-based search tools which lack robustness against variability in phrasing. For example, when retrieving information about the American football team the Buffalo Bills, these tools are less likely to retrieve documents that refer to the team as Buffalo, The Bills, or even the team.
In this work, we’re excited to introduce LMEnt (Language Models with annotated Entity mentions), a suite of pretrained models, a knowledge-rich pretraining corpus fully annotated with fine-grained entity mentions, and an entity-based retrieval method which allows full transparency into where a specific entity is mentioned in training.
Corpus
7.3M
unique entities annotated across English Wikipedia
Retrieval
10.5M
chunks in an Elasticsearch index, retrievable by an entity’s unique Wikidata identifier (QID)
Models
12
LMs with 170M, 600M, and 1B parameters, with 110 intermediate checkpoints per epoch and 4K checkpoints total
Annotating every entity mention
Entities can be mentioned explicitly by name, or implicitly through pronouns, aliases, and generic descriptors like the team. We link each mention to a Wikidata QID using three sources, each with its own score. Entity-based retrieval retrieves chunks from the pretraining data corpus that mention an entity matching a given entity QID and source score threshold, so the aliases of Buffalo the team and Buffalo the city are disambiguated.
Hyperlinks
Links that Wikipedia editors added by hand.
151M mentions
Entity linking
Finds and disambiguates additional named mentions, with a confidence score.
203M mentions
Coreference
Groups indirect mentions of the same entity, so he and the team inherit their cluster’s entity.
310M mentions
Interactive figure
Document
{{ch.label}}
{{ch.lead}}{{seg.text}}
{{lg.qid}}{{lg.name}}
filtered out
How to use this figure
1Hover over an entity mention (or click to pin it) to see its entity, its score from each source, and the rest of its coreference cluster.
2Toggle the source scores and move the thresholds in the panel on the right. Mentions that are filtered out from the entity-based retrieval get a dotted line, and the counter shows how many mentions still match.
Retrieval that understands entities
Because every mention maps to a Wikidata QID identifier, finding an entity in the training data becomes a QID lookup instead of exact matching on a string. Compared with string-based baselines, LMEnt retrieves more correct mentions for 66–80% of entities, and keeps precision above 97% even when 100K chunks are retrieved for the entity.
LMEnt wins
Tie
Baseline wins
{{wr.label}}
{{sg.t}}
Pairwise wins rate for entity-based retrieval (LMEnt) with multiple string-based retrieval methods and ablations. LMEnt’s entity-based retrieval outperforms string-based methods by 66.7%–80.4%. Ablations (bottom three rows) show that hyperlinks and entity linking are the most crucial sources of entity mentions.
k
1
5
10
100
1K
10K
25K
50K
100K
{{lbl.t}}
{{precChart.tip.title}}
{{precChart.tip.body}}
Precision at various retrieval depths (k), the top-k chunks retrieved per entity. As depth increases, LMEnt maintains > 97% precision, while the string-based retrieval variants CS-SS Canonical and CI-SS Expanded decrease to 84% and 27%.
Models you can trace…step by step, chunk by chunk, and entity by entity
We pretrain models at three sizes (170M, 600M and 1B parameters) for 1, 2, 4 and 6 epochs over the corpus: 12 models in total. A checkpoint is saved every 1,000 steps, resulting in 110 checkpoints per epoch and 4K checkpoints overall.
Each training sequence holds chunks from a single document, and chunks never split a mention. So for any training step we know exactly which chunks the model saw, and every entity they mention.
Pick an entity to mark every training step whose batch contains a chunk mentioning it. Select a step to read the matching chunk with the mention highlighted.
Example data · Wikipedia excerpts, illustrative steps
Training loss
{{plot.qid}}
{{plot.name}}
{{xt.label}}
Training steps
Step {{detail.step}}
chunk from the article {{detail.doc}}
{{detail.kind}}
{{detail.lead}}{{ds.text}}
“{{dm.text}}”{{dm.scores}}
How to use this figure
1Pick an entity with the buttons above the plot. Every training step whose batch contains a chunk mentioning it is marked on the training loss curve.
2Click a marker (or use the ‹ › arrows) to read the chunk from that step, with the entity’s mentions highlighted. Filled markers are explicit mentions; hollow markers are implicit mentions found only by coreference.
LMEnt models are valid resources for studying knowledge
LMEnt models only ever see Wikipedia, about 0.09% of the tokens OLMo-2 was trained on. Still, on PopQA questions whose subject and answer entities co-occur frequently, LMEnt models achieve comparable performance to other models with better compute efficiency.
PopQA (1K+ Subject Answer Chunks Shared)
LMEnt
OLMo
Pythia
SmolLM2
Baselines
{{lbl.t}}
{{paretoChart.tip.title}}
{{paretoChart.tip.body}}
Accuracy on PopQA questions, whose subject and answer entities co-occur frequently, as a function of compute budget.
What we learned using LMEnt so far
Beyond releasing this resource, we use LMEnt to explore how pretraining data shapes knowledge in LMs. We apply LMEnt to study knowledge acquisition and forgetting trends, and in-context editing.
1
Frequent facts are both learned and forgotten, more
LMs learn, and also forget, frequent facts more over the course of training, and these trends differ across relation types.
In-context editing success correlates with the frequency of entity co-occurrence in the data: LMs are more rigid to editing facts with strong associations.
Orig-Subj low · CF-Subj low
Orig-Subj low · CF-Subj high
Orig-Subj high · CF-Subj low
Orig-Subj high · CF-Subj high
{{lbl.t}}
{{editChart.tip.title}}
{{editChart.tip.body}}
In-context editing success across training checkpoints, split by the co-occurrence frequencies of Orig-Subj and CF-Subj pairs. LMs are more rigid to editing facts with strong associations with either pair (red < blue, dash < solid). 6E refers to the final LMEnt-1B model at epoch six.
2
Entity forms in training data matter
For frequent facts, a model is more likely to answer correctly when its training chunks name the entity explicitly rather than refer to it implicitly like “he” or “the team”. These analyses are infeasible with string search due to alias ambiguity.
Correct
Incorrect
{{lbl.t}}
{{boxChart.tip.title}}
{{boxChart.tip.body}}
LMEnt-1B is more accurate on facts for which entities co-occur in their explicit forms in the training data, implying that entity mention forms influence learning.
4
Brittle facts are easier to edit
In-context editing success increases with prediction brittleness: facts whose answers fluctuate more during training are easier to edit.
{{lbl.t}}
{{britChart.tip.title}}
{{britChart.tip.body}}
In-context editing success of the LMEnt-1B model increases with prediction brittleness.
What you can study with LMEnt
We release the corpus, annotations, entity-based retrieval index, models and checkpoints to promote research in the following areas, and more!
Knowledge plasticity
Identify steps during pretraining when models are more receptive to acquiring new knowledge.
Improving factuality
Experiment with different data ordering methods, like grouping together all the chunks mentioning the same entity, or augment the pretraining data by replacing implicit mentions with explicit entity names.
Effect of other data sources
Extend LMEnt with other sources to analyze if the addition of these tokens improves factuality. The annotation pipeline can also be applied to other knowledge-rich sources that lack hyperlinks.
Mechanistic interpretability
A controlled, yet natural, setup for studying the formation of latent knowledge representations and circuits in LMs.