You Are My Type! Type Embeddings for Pre-trained Language Models (2022)

Abstract: One reason for the positive impact of Pre-trained Language Models (PLMs) in NLP tasks is their ability to encode semantic types, such as ‘European City’ or ‘Woman’. While previous work has analyzed such information in the context of interpretability, it is not clear how to use types to steer the PLM output. For example, in a cloze statement, it is desirable to steer the model to generate a token that satisfies a user-specified type, e.g., predict a date rather than a location. In this work, we introduce Type Embeddings (TEs), an input embedding that promotes desired types in a PLM. Our proposal is to define a type by a small set of word examples. We empirically study the ability of TEs both in representing types and in steering masking predictions without changes to the prompt text in BERT. Finally, using the LAMA datasets, we show how TEs highly improve the precision in extracting facts from PLMs.

## introduction pre-trained language models (plms) based on transformers @xcite have achieved state-of-the-art results in several downstream nlp tasks @xcite . being trained in a self-supervised fashion, such models convey, to a certain extent, linguistic @xcite and factual knowledge @xcite . being able to faithfully extract the desired knowledge is a crucial aspect that has sparked lots of interest @xcite . however, querying the plm for information is not always reliable and requires more than a manually-written prompt as an input @xcite . this is opposed to a standard knowledge graph (kg), where users formulate a structured sparql query specifying exactly what to expect at the output. for example, the query "select ?x where wd:q76 wdt:p26 ?x" returns the spouse of barack obama, "michelle obama". in the plm setting, the sparql query could be replaced by a natural-language prompt, such as "the spouse of barack obama is [mask]". while the predictions of the prompt are reasonable (left-hand side of figure 1 ), they do not reflect the requirement of getting instances of a specific type (names of people) in the output. in fact, in bert's top-1 prediction on prompts where the desired output type is a musical instrument (e.g., "philip glass plays [mask]"), more than half of the predictions follow different types such as sport ("plays football") and character ("plays hamlet"), instead of the expected "plays piano". indeed, differently from the kg with typed entities, the type information is dismissed from the input prompt, thus bringing no guarantee about the expected type. while several works try to remedy this by engineering prompts to satisfy a desired type @xcite @xcite , or relying on external sources to enrich the prompt @xcite , these approaches do not fully exploit the latent concepts encoded in the plm @xcite . to fill up this gap, we introduce the notion of type embeddings (tes). similar to how positional embeddings in a plm encode information about the position of a token in an input @xcite , tes encode the expected type information of the output. the definition of a te requires neither supervised training nor external resources as it simply uses the existing plm token embeddings, e.g., people names, to obtain type information, e.g., for person. tes can be then naturally injected into the input embedding layer of a plm to embody the expected type in the output (right-hand side of figure 1 ). driving the model towards the expected type can help in applications exploiting plms, such as data integration @xcite , data cleaning @xcite , rule induction @xcite , and fact-checking @xcite . our contributions can be summarized as follows: • we introduce type embeddings (tes), which, similar to positional embeddings, can be added to the input of plms and effectively encode type information. we show how to compute these embeddings using only labeled tokens that adhere to the specific type; the main idea is to remove the first singular vector of the token embedding matrix (section 3). ## related work plms have been largely studied in the last years, with most analysis focusing on the attention mechanism @xcite and on the role of embeddings @xcite @xcite . however, none of those efforts study the notion of types that we introduce. one exception is the recent studies of how concepts are encoded in plms. one work analyzes bert by clustering contextual representations across layers, followed by a manual annotation to label clusters with meaningful concepts @xcite . another work starts from treating the feedforward network of a transformer as a key-value memory and studies how certain vectors encode concepts in the vocabulary space @xcite . our effort is different in two ways. first, we do not require the labeling of artifacts from the plm, but rather we rely on user-specified tokens to model their common type. second, we focus on type, which is one semantic concept, leaving the others, such as syntactic, morphological, and lexical to future work (section 6). our approach is related to the interpretation of a neural net's internal state in terms of a concept defined through a vector @xcite . the concept activation vector (cav) is derived from example images by finding the normal to the hyperplane separating examples without a concept and examples with a concept in the model's activation. cavs separate examples with and without the target concept in a model's activation at a certain layer. by testing with a cav (tcav), one can identify the importance of the color 'red' in fire-engine images for a neural network. we use cavs on textual input, rather than on images, to measure how sensitive the model is to a type after adding its te (section 4.3). however, while cav is a sensitivity measurement tool, tes steer the target type in the model's output. a work sharing the same spirit as ours uses a vector to steer output in a plm for style transfer between sentences @xcite . however, our method requires only 10 tokens per type as opposed to 100 labeled sentences for style transfer, and it works also with gpt. our work introduces a new kind of type embeddings to enrich the input to the plm, in analogy to positional embeddings @xcite . we show the benefit of our solution on the lama benchmark (petroni et al., 2019), which contains cloze statements to query the plm for a masked token. to enhance a plms' performance for such task, previous work improve prompts by mining or paraphrasing new prompts @xcite , by adding trigger tokens @xcite , by finding vectors for prompts in the embedding space without restriction to the plm's vocabulary @xcite , or by combining multiple prompts @xcite . as we simply add the type embedding to the input, our work is also different from approaches that pre-train an adapter to enhance plms' factual knowledge @xcite or rely on information retrieval to provide additional context for the prompt @xcite . finally, we steer the output while not changing the underlying model by triggering the neurons responsible for a prediction @xcite or by producing an alternative model with edited facts @xcite . ## type embedding in this section, we propose how to compute tes from plm token embeddings (section 3.1), and how to use them (section 3.2). following the work on latent concepts in bert @xcite , we focus on such model and report results on other plms in table a2 in the appendix. ## analysis of tes having obtained a te, we propose a series of analysis methods to assess its validity. we use the te as a simple type retriever (section 4.1), study the distribution of singular vectors (section 4.2), analyze the effect of the te w.r.t. the output and quantify the model's sensitivity w.r.t. typed tokens (section 4.3), perform layer-wise classification to identify the desired type (section 4.4), and measure tcav of a model equipped with a te (section 4.5). ## experiments the lama benchmark @xcite contains cloze statements to test plms' factual knowledge. first, we apply tes to bert and show increase in precision for most datasets (section 5.1). we then enforce a change in the output with tes (section 5.2). finally, we show the impact of the tokens that encode the te (section 5.3). ## conclusion and future work we have introduced tes as additional input for plms to better encode type information, proposed methods to analyze tes, and tested them on the lama dataset. while initial results are promising, we identify two directions of research. more precise type embeddings. further analysis on the examples can lead to better tes. one direction is to use also negative samples to compute the te. this implies learning a vector that separates between samples as cavs do. however, adding negative samples can bring more bias in the te. this could be alleviated by performing some statistical hypothesis testing, as with cavs @xcite . another way to improve the effectiveness of our proposal is to combine vectors. assuming a taxonomy of the types, different tes can be combined, for example by subtracting for the one at hand, say person, the ones that are not super or sub types, such as city and year, as discussed for dob in from types to concepts. while we focus on types and tes, our approach can be extended to include more generic concepts, as long as their tokens are in the plm's vocabulary. this could help alleviate the stereotypical and toxic content found in plms @xcite . to test our idea, we report an example for the task of natural language generation, where we "de-toxify" text generated by an autoregressive language model. we use a distilled gpt-2 model @xcite and the realtoxicityprompts dataset that contains 100k sentence-level prompts derived from a corpus of english text @xcite . we feed 10k samples to the model, thus producing the generated texts. we then measure the toxicity of such texts with the perspective api foot_0 . we consider a text toxic if the toxicity probability returned by the api @xmath0.5 and obtain 460 toxic prompts. we then compute a "toxicity concept embedding" using 6 manually picked tokens that convey toxicity. to de-toxify the generated text, we set the multiplier λ to negative values. instead of adding the embedding to the [mask] token only, we found better results when adding it to all tokens in the prompt. we believe that adding the te to all tokens helps to 'preserve' type information along the lengthy generation procedure, as opposed to mlm which decodes one token. we also test a sample of nontoxic prompts (same size as toxic prompts) to show the effect of our embedding. in addition to toxicity, we measure fluency (perplexity of generated continuations according to a larger plm) and diversity (mean number of distinct uni-/bi-/trigrams, normalized by the length of text for each prompt), as in other works for text generation @xcite . results in