Knowledge-Augmented PreTrainingfor Reasoning
ACL 2023
Tutorial on Complex Reasoning in Natural Language
Michihiro Yasunaga
Stanford UniversityRecap: Knowledge
Knowledge is available in various forms
Text

Knowledge Graph (KG)

● Diverse & contextual knowledge● Structured background knowledge

2Knowledge helps complex reasoning
Reasoning often involves combining multiple pieces of knowledge

Lin+2019; Yasunaga+2021

3This section:
Knowledge-augmented Pre-Training

4

What is Pre-Training (and Pre-Finetuning)● Key: learn from diverse data (e.g. through self-supervised learning)

Wide range of tasks
Diverse data

Model

Question
Answering
Text
Classification

Pre-train(e.g. language modeling,self-supervised learning)

Pre-finetune(e.g. instruction tuning, RLHF)

Adapt(e.g. finetuning,prompting)Information
Retrieval...

5

Why Pre-Training (and Pre-Finetuning)?
● Help a broad range of downstream tasks● Make adaptation efficient (e.g. few-shot finetuning/prompting)Wide range of tasks
Diverse data

Model

Question
Answering
Text
Classification

Pre-train(e.g. language modeling,self-supervised learning)

Pre-finetune(e.g. instruction tuning, RLHF)

Adapt(e.g. finetuning,prompting)Information
Retrieval...

6Goal: Knowledge-augmented Pre-Training

Text

Knowledge-augmentedPre-training● Diverse & contextualknowledge

Knowledge Graph (KG)
● Structured backgroundknowledge

Knowledge- &Reasoning-intensive
Tasks

7

Outline of Knowledge-augmented Pre-training

Integrate textual knowledge

REALM [ICML 2020]
CDLM [EMNLP 2021]
LinkBERT [ACL 2022]

Knowledge graph as
training objective

WKLM [ICLR 2020]
KEPLER [TACL 2021]
JAKET [AAAI 2022]

Knowledge graph as
input context

ERNIE [ACL 2019]CoLAKE [COLING 2020]
DRAGON [NeurIPS 2022]

Integrate structured
knowledge

8Integrate Textual Knowledge

Integrate textual knowledge

REALM [ICML 2020]
CDLM [EMNLP 2021]
LinkBERT [ACL 2022]

Knowledge graph as
training objective

WKLM [ICLR 2020]
KEPLER [TACL 2021]
JAKET [AAAI 2022]

Knowledge graph as
input context

ERNIE [ACL 2019]
CoLAKE [COLING 2020]
DRAGON [NeurIPS 2022]

Integrate structured
knowledge

9Integrate Textual Knowledge● We will focus on text retrieval, which helps to make reasoning process moreexplicit

10

REALM: Retrieval-Augmented Language Model Pre-Training
Method● When doing masked token prediction,
retrieve relevant documents from aknowledge corpus as reference
Key● DPR was retrieval-augmented finetuning
for QA. REALM is a self-supervised
pre-training version.

Input doc &retrieved doc

Guu, et al. "Realm: Retrieval-augmented language model pre-training." ICML 2020.

11REALM: Retrieval-Augmented Language Model Pre-Training
Result● Improve knowledge-intensive NLP (e.g. open-domain QA)

15x larger
than REALM

Needs gold
retrieval docs

Figure from https://web.stanford.edu/class/archive/cs/cs224n/cs224n.1224/slides/cs224n-2022-lecture15-guu.pdf

12

CDLM: Cross-Document Language Modeling
Method● Retrieve related docs and pre-train LM on concatenated context● Internalize knowledge during pretraining. Retrieval is optional during inference.
Set of related docs

LM pre-training

Multiple retrieved docs
Caciularu, et al. "CDLM: Cross-Document Language Modeling." EMNLP 2021.

13

CDLM: Cross-Document Language ModelingResult● Improve cross-document NLP (e.g. citation recommendation, coreference resolution)
Takeaway● Retrieval-augmented pre-training helps cross-document reasoning(Citation recommendation illustration: https://arxiv.org/pdf/2108.07571.pdf)

Caciularu, et al. "CDLM: Cross-Document Language Modeling." EMNLP 2021.

14

LinkBERT: Pretraining Language Models with Document Links
Method● Retrieve related docs and pre-train LM on concatenated context● Include various doc relations (e.g. hyperlink, citation, dense retrieval)● Internalize knowledge during pretraining. Retrieval is optional during inference.
Document relations

LM pre-trainingMultiple retrieved docs
Yasunaga, et al. "LinkBERT: Pretraining Language Models with Document Links". ACL 2022.

15

LinkBERT: Pretraining Language Models with Document Links
Result● Improve knowledge-intensive NLP● Improve multi-hop & multi-document reasoning

Yasunaga, et al. "LinkBERT: Pretraining Language Models with Document Links". ACL 2022.

16

LinkBERT: Pretraining Language Models with Document Links
Takeaway● Retrieval-augmented pre-training (⇒ multi-document in context) helps learn
multi-hop reasoning

Yasunaga, et al. "LinkBERT: Pretraining Language Models with Document Links". ACL 2022.

17

Summary so far

Integrate textual knowledge

Integrate
structured
knowledge

REALM - retrieve relevant docs
CDLM - learn relevant docs
LinkBERT - learn relevant docs & doc relations

Knowledge graph as
training objective

WKLM - KG entity objective
KEPLER - KG link objective
JAKET - both entity and link objectives

Knowledge graph as
input context

ERNIE - contextualize entity emb
CoLAKE - contextualize KG triplet
DRAGON - contextualize KG subgraph
18Integrate Structured Knowledge

Integrate textual knowledge

REALM [ICML 2020]
CDLM [EMNLP 2021]
LinkBERT [ACL 2022]

Knowledge graph as
training objective

WKLM [ICLR 2020]
KEPLER [TACL 2021]
JAKET [AAAI 2022]

Knowledge graph asinput context

ERNIE [ACL 2019]
CoLAKE [COLING 2020]
DRAGON [NeurIPS 2022]

Integrate structuredknowledge

19Integrate Knowledge Graph (KG)
KG as objective (output)

KG as input● Convenient – KG not needed at test time

● Expressive model

KGTextModelModelText
Text

KG20

Integrate Knowledge Graph as Training ObjectiveIntegrate textual knowledge

REALM [ICML 2020]
CDLM [EMNLP 2021]
LinkBERT [ACL 2022]

Knowledge graph as
training objective

WKLM [ICLR 2020]
KEPLER [TACL 2021]
JAKET [AAAI 2022]

Knowledge graph as
input context

ERNIE [ACL 2019]
CoLAKE [COLING 2020]
DRAGON [NeurIPS 2022]

Integrate structured
knowledge21

Information in Knowledge Graph● Entity information (e.g., category, definition)● Link information (e.g., reasoning about entity relations)Figure from http://web.stanford.edu/class/cs224w/slides/10-kg.pdf

22

WKLM: Weakly Supervised Knowledge-Pretrained Language Model
Idea● Add entity prediction objective

KG-based objective(entity prediction)

Method● Replace entity mentions in text
by false entities in the samecategory● Predict true/false entities

Xiong, et al. "Weakly Supervised Knowledge-Pretrained Language Model." ICLR 2020.

23

WKLM: Weakly Supervised Knowledge-Pretrained Language Model
Result● Improve knowledge-intensive NLP (e.g. QA, entity typing)
Takeaway● Seminal work in using KG for LM pre-training objective

Xiong, et al. "Weakly Supervised Knowledge-Pretrained Language Model." ICLR 2020.

24

KEPLER: A Unified Model for Knowledge Embedding and Pre-trained
Language Representation
Idea● Add KG link prediction objective

KG-based objective(link prediction)

Method● Predict whether(head, relation, tail)forms a link or not● Use TransE head: || h + r − t ||Wang, et al. "KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation". TACL 2021.

25

KEPLER: A Unified Model for Knowledge Embedding and Pre-trained
Language Representation
Result● Improve knowledge-intensive NLP and KG link prediction
Takeaway● Besides entities, KG links (structure) can augment LM pre-training objective

Wang, et al. "KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation". TACL 2021.

26

JAKET: Joint Pre-training of Knowledge Graph and Language Understanding

Idea● Add both KG entity and link prediction
objectives

MetaQA example (2-hop): “Who acted in themovies directed by Erik Poppe?”

Result● Further improvement on multi-hop
reasoning (e.g., MetaQA)

Yu, Zhu, et al. "JAKET: Joint Pre-training of Knowledge Graph and Language Understanding". AAAI 2022.

27

Summary so far

Integrate textual knowledge

Integratestructuredknowledge

REALM - retrieve relevant docs
CDLM - learn relevant docs
LinkBERT - learn relevant docs & doc relations

Knowledge graph as
training objective

WKLM - KG entity objective
KEPLER - KG link objective
JAKET - both entity and link objectives

Knowledge graph as
input context

ERNIE - contextualize entity emb
CoLAKE - contextualize KG triplet
DRAGON - contextualize KG subgraph
28

Integrate Knowledge Graph as Input Context

Integrate textual knowledge

REALM [ICML 2020]
CDLM [EMNLP 2021]
LinkBERT [ACL 2022]

Knowledge graph astraining objective

WKLM [ICLR 2020]
KEPLER [TACL 2021]
JAKET [AAAI 2022]

Knowledge graph asinput context

ERNIE [ACL 2019]
CoLAKE [COLING 2020]
DRAGON [NeurIPS 2022]

Integrate structured
knowledge

29

ERNIE: Enhanced Language Representation with Informative Entities
Method● Add KG entity embeddings in LMcontext● Entity embs are concatenated tocorresponding word embs

Contextualize

ModelTakeaway● Seminal work in using KG info as input
in LM pre-training
TextZhang, Zhengyan, et al. "ERNIE: Enhanced Language Representation with Informative Entities." ACL 2019.

KG entity
embeddings
30

CoLAKE: Contextualized Language and Knowledge Embedding
Method● Add KG triplets in LM context(⇒ bring neighbor entities andrelations in context, besides direct
entity mentions in text)

Contextualize

Model● Masked token prediction for bothtext and KG sides

KG triplets
Text
Harry Potter points his wand at Lord Voldemort
Sun, et al. "CoLAKE: Contextualized Language and Knowledge Embedding." COLING 2020

(Harry Potter, mother, Lily Potter),(Harry Potter, spouse, Ginny Weasley),(Harry Potter, enemy of, Lord Voldemort)
31

CoLAKE: Contextualized Language and Knowledge Embedding
Result● Improve knowledge-intensive NLP
and KG link prediction
Takeaway● KG triplets provide backgroundknowledge and help reason about
related entities

Sun, et al. "CoLAKE: Contextualized Language and Knowledge Embedding." COLING 202032

DRAGON: Deep Bidirectional Language-Knowledge Graph Pretraining
Method● Add KG subgraph in input context● Text is contextualized by LM and
KG is contextualized by GNN. The
two are then contextualized
bidirectionally

Contextualize

Self-supervised
Objective

Masked LM

KG link prediction

art supplies(round brush, at, hair)

LM HeadLinkPred Head

INT

Model

If

it······

Int

MASK···LM Layer

Cross-modal
Encoder···

hair

Text

AtLocation

x.....

Yasunaga, et al. "DRAGON: Deep Bidirectional Language-Knowledge
Graph Pretraining" NeurIPS 2022

INTGNN Layer
Int···

Interaction
nodeLM Layer[INT] If it is not used forhair, a round brush is an
example of [MASK] [MASK].

···

MIntFusion LayerxNInt

INTxM

To
edlatReroundbrush
UsedForpaintinghairbrushartsupplyFor
Used

KG subgraph33

DRAGON: Deep Bidirectional Language-Knowledge Graph Pretraining
Result● Improve broad reasoning tasks (QA, commonsense, link prediction)● Improve complex reasoning (multi-hop, logical)

Yasunaga, et al. "DRAGON: Deep Bidirectional Language-Knowledge Graph Pretraining" NeurIPS 2022

34

DRAGON: Deep Bidirectional Language-Knowledge Graph Pretraining
Takeaway● KG graph structure provides LM with a scaffold to perform complex reasoning
about entities

After several layers of fusion,
attention weight from text over hair decreases,
but attention weight over round brush andpainting increases,
adjusting for the negation in text

Prediction
Yasunaga, et al. "DRAGON: Deep Bidirectional Language-Knowledge Graph Pretraining" NeurIPS 202235

DRAGON: Deep Bidirectional Language-Knowledge Graph Pretraining
Takeaway● KG graph structure provides LM with a scaffold to perform complex reasoning
about entities

Yasunaga, et al. "DRAGON: Deep Bidirectional Language-Knowledge Graph Pretraining" NeurIPS 2022

36

Summary so far

Integrate textual knowledge

Integrate
structuredknowledge

REALM - retrieve relevant docs
CDLM - learn relevant docs
LinkBERT - learn relevant docs & doc relations

Knowledge graph as
training objective

WKLM - KG entity objective
KEPLER - KG link objective
JAKET - both entity and link objectives

Knowledge graph as
input context

ERNIE - contextualize entity emb
CoLAKE - contextualize KG triplet
DRAGON - contextualize KG subgraph
37SummaryIntegrate textual knowledge

Integratestructuredknowledge

REALM - retrieve relevant docs
CDLM - learn relevant docs
LinkBERT - learn relevant docs & doc relations

Knowledge graph as
training objective

WKLM - KG entity objective
KEPLER - KG link objective
JAKET - both entity and link objectives

Knowledge graph as
input context

ERNIE - contextualize entity emb
CoLAKE - contextualize KG triplet
DRAGON - contextualize KG subgraph38

Conclusion
Takeaways● Knowledge can be integrated into LM in self-supervised ways● Help a wide range of reasoning tasks

Open questions● Can we integrate knowledge in pre-finetuning (e.g. instruction tuning, RLHF)?
● How can we build a unified model with all various knowledge sources?
● How can we ensure the models use and reason about knowledge faithfully?

39

Thak you!
https://cs.stanford.edu/~myasu/@michiyasunaga

40