SlowGuess/NAACL_2022_OCR huggingface dataset:
Shuyang Li
sh1008@ucsd.edu
Mukund Sridhar
harakere@amazon.com
Chandana Satya Prakash
chanprak@amazon.com
Jin Cao
jincao@amazon.com
Wael Hamza
waelhamz@amazon.com
Julian McAuley
jmcauley@ucsd.edu
Understanding human language often necessitates understanding entities and their place in a taxonomy of knowledge—their types. Previous methods to learn entity types rely on training classifiers on datasets with coarse, noisy, and incomplete labels. We introduce a method to instill fine-grained type knowledge in language models with text-to-text pre-training on type-centric questions leveraging knowledge base documents and knowledge graphs. We create the WikiWiki dataset: entities and passages from 10M Wikipedia articles linked to the Wikidata knowledge graph with 41K types. Models trained on WikiWiki achieve state-of-the-art performance in zero-shot dialog state tracking benchmarks, accurately infer entity types in Wikipedia articles, and can discover new types deemed useful by human judges.
Entities can be categorized by their types, which indicate where they belong in a taxonomy of knowledge. For example, Venus is a planet and thus also an astronomical body. Much like how knowledge acquisition in cognitive development progresses from recognizing concrete objects to gradually understanding their relations to one another (Lucariello et al., 1992), we aim to extend language models' existing rough understanding of entities (Heinzerling and Inui, 2021) to the types that govern how entities are related. Instilling type knowledge in multi-purpose models can improve performance in tasks like entity linking (Onoe and Durrett, 2020), question-answering (Févry et al., 2020a), and semantic parsing (Thirukovalluru et al., 2021).
While language models can memorize some facts (Petroni et al., 2019), they frequently hallucinate false information (Logan IV et al., 2019; Shuster et al., 2021). Current attempts to learn to infer types for entities are hampered by 1) the difficulty
Figure 1: Via the WikiWiki dataset, we train a model to answer questions about entities mentioned in Wikipedia articles (top) and WIkidata types that such entities are an instance of (P31) or subclass of (P279).
of collecting diverse, large-scale typing datasets; and 2) how existing corpora assume independence between types (Choi et al., 2018), while in reality types sit at levels of granularity that are useful in different settings: a pizza store may care whether a user likes Cheese Pizza; a restaurant recommender might care if the user wants Pizza; finally, a general dialog agent might only care if a user wants Food.
We address both issues by proposing a simple and effective approach for pre-training generative language models to answer questions about entities, types, and surface forms (mentions) in a large public knowledge graph (KG) consisting of Wikipedia articles and Wikidata nodes. We leverage high quality type labels in a large corpus of knowledge-rich text and an ordered, hierarchical type ontology.
To summarize our main contributions: 1) We create the new WikiWiki dataset comprising 10M Wikipedia articles linked to nodes from Wikidata; 2) We propose a pre-training scheme for generative language models using type-centric question-answering based on WikiWiki; 3) We achieve state-of-the-art (SOTA) performance in zero-shot domain adaptation for dialog state tracking using our type-instilled models, with average per-domain gains of  $14.9\%$  (49.4% relative) joint accuracy; and 4) We show that our models can precisely infer types for seen and unseen entities in new articles from WikiWiki, and propose novel types that hu
Findings of the Association for Computational Linguistics: NAACL 2022, pages 594 - 603
Training	Test	Test (New Ent)
Documents	10 M	5.0 K	5.0 K
Unique Entities	2.2 M	14.1 K	6.0 K
Unique Types	40.6 K	4.0 K	1.2 K
Num. of Mentions	38.7 M	19.3 K	6.4 K
Type References	43.8 M	21.5 K	6.5 K
Table 1: Unique documents/entities/types and number of mentions in each split of WikiWiki. Test (New Ent) comprises entities not seen in the training split.
mans judge to be accurate and appropriate.
Knowledge Grounding in Language Models Large pre-trained language models have been shown to memorize some facts (Petroni et al., 2019). One recent line of work aims to explicitly condition generation on knowledge bases by combining a retrieval module and a language model (Majumder et al., 2020; Guu et al., 2020; Lewis et al., 2020b; Mazaré et al., 2018). Peters et al. (2019) propose instead to align token representations from pre-trained language models with entity embeddings to reason over a limited set of entities. Yamada et al. (2020) explicitly denote entity tokens with a learned input embedding. Specific entity embeddings have also been learned jointly by using knowledge graphs as auxiliary inputs during language model pre-training (Sun et al., 2020a; Févry et al., 2020b; Zhang et al., 2021). Another line of work aims to model specific factual statements from knowledge bases (Wang et al., 2021) for reading comprehension (Lu et al., 2021) and trivia QA (Agarwal et al., 2021). We propose text-to-text pre-training on knowledge recovery tasks to instill type-awareness. Our models learn type knowledge that transfers to the type-adjacent downstream task of dialog state tracking and can infer unseen types.
Entity Representation Learning Many SOTA systems for knowledge retrieval and QA rely on learned dense embeddings of individual entities or types to perform multi-class classification (Ganea and Hofmann, 2017; Karpukhin et al., 2020; Wu et al., 2020a). Several recent frameworks aim to learn entity knowledge during language model pretraining via entity masking (Sun et al., 2020b) or contrastive learning (Qin et al., 2021). Systems for entity typing (Dai et al., 2021) and disambiguation (Yamada et al., 2019) also learn dense vector encodings that are later matched via dot-product scoring. Cao et al. (2021) aim to address some
Context: These included carbon dioxide by burning diamond, and mercuric oxide by heating mercury. This type of experiment contributed to the discovery of "dephlogisticated air" by Priestley, which became better known as oxygen, following Lavoisier's investigations.
Entity/Type Discovery  $(20\%)$ : List all concepts and types mentioned here.
Answer: Priestley (chemist), Lavoisier (chemist), mercuric oxide (chemical compound), mercury (chemical element), and dephlogisticated air (superseded scientific theory)
Entity Typing  $(30\%)$ : What is dephlogisticated air an example of?
Answer: superseded scientific theory
Entity Recognition (20%): What does Priestley refer to? Answer: Joseph Priestley (chemist)
Slot Filling  $(30\%)$ : Which chemists are mentioned here? Answer: Joseph Priestley and Antoine Lavoisier
Table 2: In pre-training, the model reads a Wikipedia article and answers questions from four tasks involving entities and types. It must generate answers containing terms not found verbatim in the text. Surface forms (mentions) in green, entities in red, and types in blue.
downsides of the above approaches—the linearly increasing space required to store learned representations and difficulties in negative sampling—by casting the task as generative language modeling: predict the name of an entity to be linked. We generalize this approach from entity names (which appear verbatim) to include types, which require a more nuanced understanding of a context.
3 Type-Centric Multitask Modeling
WikiWiki Corpus To train an entity- and type-aware language model, we build the WikiWiki dataset by combining Wikipedia articles with the Wikidata KG (Vrandecic, 2012). Wikipedia articles have been used to enrich corpora for dialog (Dinan et al., 2019), coreference resolution (Singh et al., 2012), and QA (Liu et al., 2020). KGs have been used for entity typing and relation extraction (Sakor et al., 2020). Yao et al. (2019) use Wikipedia pages as context for relation triples mined from Wikidata.
We link articles, entities, and types as in Figure 1: like Wu et al. (2020b), we take Wikipedia hyperlinks as links between entities (target page) and their mentions (link text); we link pages to Wikidata nodes via ID; and for each node we extract types  $T$  from Wikidata where  $t \in T$  is an instance/subclass of the node (discarding entities with no types). To address sparsity of hyperlinks,
<sup>1</sup>All humans on Wikidata are an instance of 'human'; we thus use the 'occupation' relation to determine their types.
we follow Yao et al. (2019) and use spaCy to identify additional entities. We sample 10M articles for training, with two disjoint 5K-article splits for evaluation, containing seen and unseen (New Ent) entities respectively (Table 1). The ontology of Wikidata types forms a directed acyclic graph with 41K type nodes applying to 2.2M entities. Previous entity typing datasets rely on annotations from small groups of crowd-workers and include a small type ontology in the hundreds (Ling and Weld, 2012) and/or sacrifice label accuracy (Choi et al., 2018). We instead rely on the cumulative, cross-checked annotations from tens of thousands of active Wikidata users.
Entities in Wikidata on average are assigned 1.28 types; for entities with multiple types, not all types are necessarily relevant to a context. For example, take the following passage: "Obama was elected to the Illinois Senate in 1996, succeeding Democratic State Senator Alice Palmer from Illinois's 13th District, which, at that time, spanned Chicago South Side neighborhoods from Hyde Park-Kenwood south to South Shore and west to Chicago Lawn."
While Wikidata entities may have  $5+$  types, many are not directly relevant to a context. For example, while Barack Obama has types including Politician, Jurist, Political Writer, Community Organizer, and Podcaster, the latter is not relevant to the context. To teach our models to infer types relevant to the context at hand, in pre-training data we take only types that are shared between Barack Obama and other entities in the document (e.g. Alice Palmer—Politician). We have made the WikiWiki dataset publicly available on Github.[2]
Pre-training Tasks To instill type-centric knowledge from WikiWiki, we train our models to answer four types of knowledge-based questions conditioned on a passage from Wikipedia (examples in Table 2). In entity/type discovery, the model is tasked to recover all surface forms (mentions) that reference an entity, along with their types—this is analogous to simultaneous entity recognition and typing. Entity typing consists of assigning types to an entity of interest. For entity recognition, we follow Cao et al. (2021) by training our model to respond with an entity's full name and type when queried with a surface form. In slot filling we ask our model to return all entities mentioned in the
User:	I'm looking for a place to stay during my upcoming trip to Cambridge.
System:	I can definitely help you with that! What area are you staying in, and what is the price range you are looking for?
User:	It should be located in the west and it should be cheap.
Belief State:	[hotel price range]: cheap; [hotel area]: west
Table 3: In Dialog State Tracking (DST), a model infers the belief state of a user given the dialog history thus far, comprising slots (red) and their values (blue). In Zero-shot DST, the model must infer the correct values for slots that it has not seen during training, requiring the agent to rely on general type knowledge.
passage belonging to a certain type. For multi-type entities, we use a subset of relevant types given other entities in the context (Appendix A). We treat QA as a universal format for diverse NLU tasks (McCann et al., 2018), and adopt the framework of Raffel et al. (2020) to treat each of our tasks as text-to-text generative modeling. We create 50M questions for pre-training.
Model Architecture We use an encoder-decoder (Sutskever et al., 2014) model initialized from BART—a Transformer (Vaswani et al., 2017) language model pre-trained via de-noising autoencoding (Lewis et al., 2020a). Our model generates an answer  $a$  as a text sequence given a document  $D$  of length  $t_d$  and question  $q$ . The document is encoded via the encoder—consisting of  $l$  Transformer layers of hidden dimensionality  $h$ , each applying 16-headed self-attention—to produce  $z := \operatorname{Enc}(D) \in \mathbb{R}^{t_d \times h}$ .
We train the model to perform QA via conditional language modeling. Instead of concatenating the question with the context in encoder input (Lin et al., 2021), the decoder generates a sequence consisting of the question and answer:  $x = [q; a]$ . We can thus cache the document encoding at inference to answer multiple questions. At training time we perform next-token prediction, calculating cross-entropy loss by maximizing the log likelihood of the question and answer conditioned on the document:  $P(q, a|D) = \prod_{t}^{T} P(x_{t}|x_{<t}, D)$ . We assess the impact of our pre-training on Base ( $l = 12$ ,  $h = 768$ ) and Large ( $l = 24$ ,  $h = 1024$ ) models.
We demonstrate the effectiveness of our pretraining on two tasks that require type understand
$^{2}$ https://github.com/amazon-research/wikiwiki-dataset/
# Params	R	H	A	T	X
TRADE	90M	12.6	14.2	20.1	22.4	59.2
MA-DST	90M	13.6	16.3	22.5	22.8	59.3
SUMBT	355M	16.5	19.8	22.6	22.5	59.5
GPT2-DST	355M	26.2	24.4	31.3	29.1	59.6
BART	139M	27.9	31.9	38.4	34.3	70.5
Ours (Base)	139M	40.4	36.5	39.8	36.1	70.9
Ours (Large)	406M	46.7	38.8	49.8	37.7	72.1
ing: zero-shot domain generalization in dialog state tracking (DST), and fine-grained entity typing.
Zero-Shot DST The goal of Dialog State Tracking (DST) is to infer user intent and goals from conversations by filling in belief slots (Lemon et al., 2006; Wang and Lemon, 2013). In many real-world settings, DST models must be able to predict new slot values (i.e. new entities that are not present in the training corpus) and new slot types (e.g. requirements for applications in new domains). This problem setting is known as zero-shot DST (Table 3). We follow the zero-shot setting in Campaigna et al. (2020): train a model on multi-domain DST data and evaluate on a held-out domain. We measure domain generalization performance via joint goal accuracy (JGA): the percent of turns in which a model successfully predicts values for all slots in the target domain. We use the MultiWOZ 2.1 benchmark (Eric et al., 2019), evaluating zero-shot JGA for the Restaurant, Hotel, Attraction, Train, and Taxi domains. At each turn, we ask the model a question about the preference for each slot. We compare against recent systems that can perform zero-shot DST: TRADE (Wu et al., 2019), MA-DST (Kumar et al., 2020), SUMBT (Lee et al., 2019), and GPT2-DST (Li et al., 2021). Our method is complementary to systems for creating synthetic in-domain dialogs (Kim et al., 2021).
As seen in Table 4, our type-centric pre-training allows a model to answer questions about unseen slots. BART-base itself achieves SOTA JGA across all domains, and our pre-training significantly increases the gain to  $10.6\%$  absolute /  $34.8\%$  relative JGA—despite only using one-third of the parameters. Our Large model achieves  $14.9\%$  absolute and  $49.4\%$  relative gain in JGA compared to previous SOTA. The most significant gains come in the Hotel and Restaurant domains, which contain the
Table 4: Zero-shot domain adaptation JGA (%) on MultiWOZ 2.1 test set on the (R)estaurant, (H)otel, (A)traction, (T)rain, and Ta(X)i domains. We achieve SOTA results on all domains by significant margins.
100%	50%	20%
Base (139M)	13.7	14.7	39.0
Large (406M)	0.9	1.6	4.8
Table 5: Relative gain  $(\%)$  in JGA for models trained on WikiWiki vs standard BART pre-training. Our method helps more in low-data regimes and for smaller models.
most categorical slots that resemble types (e.g. cuisine, hotel type). In Table 5 we compare our models against same-size BART models at different levels of training data availability to demonstrate the additive utility of our method. Our method is particularly helpful with less fine-tuning data (low-data regimes), with average gains of  $39\%$  for small models and  $4.8\%$  for large models at  $20\%$  data availability. Gains are magnified for smaller models, affirming that our method can effectively instill type knowledge in lightweight language models.
Ultra-Fine Entity Typing Our method improves generalization in type-adjacent tasks; we next aim to infer entity types in unseen documents. In preliminary experiments on the UltraFine dataset with 11K types (Choi et al., 2018), our models underperform SOTA (24.0 vs. 49.1 F1). Manual inspection of gold labels reveals two main causes for error: 1) inaccurate labels—e.g. “rare plants” as type “bird”; and 2) inconsistent usage of gold labels: different spellings (organization / organisation) or synonyms (car / automobile) are treated as distinct and often do not collocate. This suggests that label noise in UltraFine may make it unsuitable for assessing granular, hierarchical type knowledge.
We examine these annotation errors via human evaluation, presenting crowd-workers with 200 contexts from UltraFine (10% of the test set). Only 68% of gold type labels were judged accurate, and 21% inaccurate. We compare gold labels against zero-shot predictions from our model in a second trial with 200 pairs. Judges preferred our predictions 51% of the time compared to 29% for gold. We observed moderate inter-annotator agreement of  $\kappa = 0.4044$  (Fleiss, 1971). This suggests that our models can accurately infer types, but current benchmarks do not suitably measure typing quality.
Entity Typing on WikiWiki We turn to WikiWiki to evaluate fine-grained entity typing, leveraging type labels verified by active users of Wikidata. To verify the accuracy of ground-truth type labels in the WikiWiki test set, we asked human evalua
Entities	Model	Precision	Recall	F1
Seen	RoBERTa	62.35	59.38	60.82
Ours	78.13	72.39	75.15
Unseen	RoBERTa	48.88	47.96	48.41
Ours	66.65	63.71	65.14
Table 6: P/R/F1 of pred. vs. gold types on WikiWiki Test (seen) and Test New Ent (unseen entities) splits.
tors to judge the accuracy of 443 type labels from 200 randomly sampled contexts. We confirm that WikiWiki is a high-quality benchmark for entity typing, with  $85\%$  type precision assessed by human judges (compared to  $68\%$  for UltraFine).
We found that multi-label classifiers built on RoBERTa (Liu et al., 2019) that perform well on UltraFine require significant hyper-parameter tuning to output non-trivial predictions to classify our large and sparse (41K) type ontology. To perform entity typing with our model, we generate comma-delimited text sequences of types (Yang et al., 2018). This allows our models to infer and generate novel types while classifiers remain restricted to the training ontology. We confirm that our pre-training helps models better infer types for both seen (+14.3 F1) and unseen entities (+16.7 F1) in new contexts compared to classifiers (Table 6).
To investigate if our model can discover novel types, we perform another human evaluation over 557 such predictions from 300 contexts, with interannotator agreement of  $\kappa = 0.4086$ . Our model accurately extrapolates its type knowledge beyond the training ontology—we observe  $73.3\%$  precision when inferring new types (compared to  $74.5\%$  precision for seen types), demonstrating that our pre-training enables models to reason about types beyond simple memorization. Our model discovers complex and specific scientific types, correctly proposing that anorthosite (an aluminum silicate rock) is a metallurgical rock<sup>3</sup> and that speckled tortoises are monotrophs.<sup>4</sup> This reflects the robust taxonomy of types in scientific disciplines. Our model also proposes granular categories of events, and is judged to correctly type the 2015 Tour of Taiwan as an instance of the Tour de Taiwan cycling race. In the future, we seek methods to automatically assess the factual accuracy of new types.
In this paper, we 1) propose a text-to-text pretraining scheme to instill type knowledge in language models via QA and 2) release the WikiWiki dataset built from Wikipedia articles and the Wikidata KG. We show that WikiWiki is larger-scale and more accurate than existing fine-grained type recognition datasets. We demonstrate that our type-centric pre-training framework allows us to train language models that can better generalize to unseen domains, entities, and types—which in turn lead to improved model performance on downstream tasks like dialog state tracking (achieving SOTA results on zero-shot DST with average gains of  $14.9\%$  joint accuracy). Our models can extrapolate type knowledge and infer novel types that humans judge to be useful and precise. As the body of human knowledge grows, we see an opportunity to use life-long learning (Parisi et al., 2019) on news and publications to expand and model the taxonomy of knowledge.
We would like to thank Stephen Rawls, Ryan Gabbard, and anonymous reviewers for providing valuable feedback on this work. We also thank Nicolas Guénon des Mesnards and Victor Soto for their help setting up MTurk for human evaluations. Work was performed during first author's internship at Amazon Alexa AI. Findings and observations are of the authors only, and do not necessarily reflect the views of Amazon or UCSD.
Oshin Agarwal, Heming Ge, Siamak Shakeri, and Rami Al-Rfou. 2021. Knowledge graph based synthetic corpus generation for knowledge-enhanced language model pre-training. In *NAACL-HLT*, pages 3554-3565.
Giovanni Campagna, Agata Foryciarz, Mehrad Moradshahi, and Monica S. Lam. 2020. Zero-shot transfer learning with synthesized data for multi-domain dialogue state tracking. In ACL, pages 122-132.
Nicola De Cao, Gautier Izacard, Sebastian Riedel, and Fabio Petroni. 2021. Autoregressive entity retrieval. In ICLR.
Eunsol Choi, Omer Levy, Yejin Choi, and Luke Zettlemoyer. 2018. Ultra-fine entity typing. In ACL, pages 87-96.
Hongliang Dai, Yangqiu Song, and Haixun Wang. 2021. Ultra-fine entity typing with weak supervision
$^{3}$ rocks containing metallic compounds and properties  
 $^{4}$ has diet comprising one type of food (Herrera, 1976)