Teaching Language Models What to Remember: Entity-Aware Pretraining and the Efficiency of Knowledge Acquisition

Pretrained language models have demonstrated remarkable capabilities in capturing linguistic patterns, yet their relationship with factual knowledge remains problematic. While transformers implicitly encode world knowledge through statistical patterns of co-occurrence, this indirect acquisition often yields inconsistent factual grounding, requiring costly fine-tuning or massive parameter scaling to remediate. The paper "Knowledge-Aware Language Model Pretraining" by Rosset et al. introduces a compelling alternative. Rather than increasing model capacity or adding complex architectural components, the authors demonstrate that simply alerting the model to entity boundaries during pretraining significantly improves factual retention and downstream performance. Their Knowledge-Aware Language Model (KALM) achieves these gains through minimal modifications to the input representation and training objective, suggesting that how we signal information to models may matter more than how many parameters we allocate to store it.

Encoding Knowledge Through Input Design

The technical innovation of KALM lies in its treatment of entities as first class citizens in the learning process. The authors implement an entity-extended tokenizer that operates alongside the standard subword tokenizer. Using a fuzzy frequency-based entity dictionary, the system annotates text spans with their corresponding entity identifiers, creating a parallel input channel where both word tokens and entity tokens are fed simultaneously into the transformer. This is not merely data augmentation; it fundamentally alters the learning signal by forcing the model to maintain coherence between surface forms and underlying referents.

The training objective extends beyond traditional next word prediction to include an entity prediction task. At each position, the model must predict not only the upcoming token but also the entity associated with the current context. This dual objective creates a pressure for the model to develop robust entity representations that persist across varied linguistic contexts. Crucially, this requires no modification to the transformer architecture itself, no external memory modules, and no additional inference-time computation. The knowledge becomes embedded in the standard parameter matrix through the modified training dynamics.

Empirical Evidence of Efficiency

The experimental results demonstrate that architectural minimalism yields substantial empirical gains. On the LAMA (LAnguage Model Analysis) knowledge probing benchmark, KALM outperforms standard GPT-2 by approximately 25% across diverse relation types, including commonsense knowledge, factual semantics, and relational predicates. Perhaps more surprisingly, the 24-layer KALM Large achieves performance comparable to the 17 billion parameter GPT-2 on certain knowledge intensive tasks, despite using orders of magnitude fewer parameters.

The advantages extend beyond static knowledge probes to functional capabilities. In zero-shot question answering, where models generate answers without task-specific fine-tuning, KALM demonstrates 20% to 100% improvements in exact match accuracy over equivalent GPT-2 baselines. The model achieves this using only a few example question-answer pairs as formatting templates, with no gradient updates on QA data. This suggests that the entity aware pretraining creates representations that are inherently more structured for knowledge retrieval and composition.

The authors employ edge probing analysis to dissect how these capabilities emerge during training. By examining hidden representations at various pretraining steps, they observe that standard transformers gradually acquire complex semantic understanding, while KALM exhibits accelerated learning of coreference resolution and entity type classification. The entity signals appear to act as scaffolding, allowing the model to construct higher level semantic abstractions more efficiently than through word co-occurrence alone.

Signal Design Versus Raw Capacity

This work presents a crucial insight for the current era of large language model development. The field has largely operated under the assumption that knowledge deficiencies are primarily capacity constraints, solvable through increased