[Paper Review] Knowledge-Aware Language Model Pretraining

Corby Rosset, Chenyan Xiong+4 more|arXiv (Cornell University)|Jun 29, 2020

Topic Modeling37 references44 citations

TL;DR

KALM adds entity signals to the input and an entity prediction objective during pretraining, achieving knowledge-rich representations and improved zero-shot and probing performance without changing Transformer architecture. It serves as a strong drop-in GPT-2 replacement with parameter-efficient gains.

ABSTRACT

How much knowledge do pretrained language models hold? Recent research observed that pretrained transformers are adept at modeling semantics but it is unclear to what degree they grasp human knowledge, or how to ensure they do so. In this paper we incorporate knowledge-awareness in language model pretraining without changing the transformer architecture, inserting explicit knowledge layers, or adding external storage of semantic information. Rather, we simply signal the existence of entities to the input of the transformer in pretraining, with an entity-extended tokenizer; and at the output, with an additional entity prediction task. Our experiments show that solely by adding these entity signals in pretraining, significantly more knowledge is packed into the transformer parameters: we observe improved language modeling accuracy, factual correctness in LAMA knowledge probing tasks, and semantics in the hidden representations through edge probing.We also show that our knowledge-aware language model (KALM) can serve as a drop-in replacement for GPT-2 models, significantly improving downstream tasks like zero-shot question-answering with no task-related training.

Motivation & Objective

• Motivate why standard pretraining lacks explicit real-world knowledge and explore a lightweight knowledge-aware alternative.
• Propose a knowledge-awareness mechanism that signals entities via an entity tokenizer and an auxiliary entity prediction task.
• Demonstrate that knowledge-aware pretraining improves knowledge probes, language modeling metrics, and zero-shot QA without increasing model size.
• Evaluate KALM against GPT-2 across LAMA probes, edge probing, perplexity, Lambada, and zero-shot QA tasks.

Proposed method

• Introduce an entity tokenizer that maps word-ngrams to world entities using a surface-form dictionary.
• Create a two-channel input: standard word tokens and aligned entity tokens with corresponding embeddings.
• Augment pretraining with a next-entity prediction task using a margin-based loss that contrasts the correct entity with negatives.
• Train the model end-to-end with a multi-task objective combining language modeling loss and entity prediction loss (l_KALM = l_W + alpha l_E).
• Keep Transformer architecture unchanged; only enlarge the token vocabulary with entity tokens and embeddings.
• Evaluate knowledge and language capabilities via LAMA probing, edge probing, and zero-shot QA.

Experimental results

Research questions

• RQ1 How does signaling entity knowledge during pretraining impact the knowledge encoded in transformer parameters?
• RQ2 Can a knowledge-aware pretraining approach improve factual correctness and commonsense/relational semantics without architectural changes?
• RQ3 Do knowledge-aware signals yield gains in zero-shot QA and probing tasks comparable to larger models?
• RQ4 What is the impact of KALM on standard language modeling metrics (perplexity, last-word accuracy) and representation quality (edge probing)?

Key findings

• KALM improves LAMA knowledge probing accuracy over GPT-2 baselines by about 40-80% across tasks, approaching GPT-2 scale performance on some relations.
• KALM yields better Lambada last-word accuracy and WikiText-103 perplexity compared to the baseline GPT-2, indicating improved language modeling.
• Zero-shot QA accuracy for KALM is 20%-100% higher than an equivalent GPT-2 model, demonstrating stronger zero-shot knowledge retrieval.
• Edge probing shows that KALM’s representations become more informative across tasks, especially for entity typing, relations, and semantic roles, and learning accelerates with knowledge-awareness.
• KALM Large matches GPT-2 17B performance on some hard relations (N-M) using only ~2% more parameters, highlighting parameter efficiency.
• KALM maintains GPT-2-like inference speed with only modest constant overhead from entity lookups and embeddings.