Towards Tracing Knowledge in Language Models Back to the Training Data (2022)

Abstract: Language models (LMs) have been shown to memorize a great deal of factual knowledge contained in their training data. But when an LM generates an assertion, it is often difficult to determine where it learned this information and whether it is true. In this paper, we propose the problem of fact tracing: identifying which training examples taught an LM to generate a particular factual assertion. Prior work on training data attribution (TDA) may offer effective tools for identifying such examples, known as “proponents”. We present the first quantitative benchmark to evaluate this. We compare two popular families of TDA methods — gradient-based and embedding-based — and find that much headroom remains. For example, both methods have lower proponent-retrieval precision than an information retrieval baseline (BM25) that does not have access to the LM at all. We identify key challenges that may be necessary for further improvement such as overcoming the problem of gradient saturation, and also show how several nuanced implementation details of existing neural TDA methods can significantly improve overall fact tracing performance.

## introduction research has shown that language models (lms) acquire significant amounts of world knowledge from the massive text corpora on which they are trained @xcite . this development has enabled exciting advances training data attribution ultimately, a language model's "knowledge" must derive from its training data. but there has been little research on attributing an lm's factual assertions back to specific training examples -a task we call fact tracing. training data attribution methods (tda) are the main literature concerned with linking predictions back to specific training examples (known as "proponents"). influence functions @xcite and tracin @xcite are among the first methods to do this for neural networks, by estimating the marginal effect of a training example on the loss of a testtime example. however, most work on tda has focused on classification and regression tasks that do not necessarily involve fine-grained factual information @xcite . several obstacles have limited research on fact tracing for large, pre-trained lms. first, since pretraining corpora are very large, it has not been clear how to obtain ground truth labels regarding which pre-training example was truly responsible for an lm's prediction. second, tda methods have traditionally been computationally prohibitive. in this paper, we present one of the first computationally tractable studies of fact tracing for lms. to do so, we construct: obtaining ground truth proponents to establish (1) ground truth data for fact tracing, we propose a new recipe, which we call "novel fact injection". first, suppose that we can identify a set of "facts" that the pre-trained lm does not know -we call these "novel facts". we can convert each novel fact into an lm training example, and then fine-tune the lm on these extra examples until it memorizes the novel facts (i.e. "injecting" them into the lm). with a few caveats, we now know that the lm must have learned these facts from our newly injected examples. we also know which examples are responsible for teaching each fact, since we constructed each example from a particular fact. hence, we now have ground-truth "proponents" for every novel fact, and can evaluate any tda method on its ability to identify these proponents -i.e. to retrieve the true proponents out of a large set of training examples. we implement this recipe using the trex dataset @xcite as our source of novel facts. trex is a large text corpus where each sentence has been comprehensively annotated with the facts that it expresses, in the form of relational knowledge tuples. to identify novel facts present in trex, we filter for knowledge tuples that the pretrained lm did not already know, as tested using masked lm prompting. the sentences in trex expressing these tuples are then "injected" via finetuning and labeled as proponents. we call this setup ftrace-trex. there are two caveats for the above setup. first, we must be careful about how we define what an lm "knows". for example, if an lm generates a particular assertion with 10% probability, does this count as "knowing" or not? second, some facts can be indirectly inferred from other facts. for example, suppose we want to know how an lm learned that barack obama was born in hawaii. it could learn this from a literal mention of the fact: "obama was born in hawaii", or indirectly infer it from "obama was born in honolulu". our trex setup only identifies literal proponents (the former), but not indirect proponents (the latter). to address these two issues, we introduce an additional, more controlled setup, ftrace-synth, featuring synthetically generated novel facts that could not have possibly been known by the pre-trained lm, and which also have no correlation with any existing facts -making indirect inferences impossible. ## retrieval methods we begin with a formal description of the different tda methods we study in this paper: gradientbased methods @xcite and embedding-based methods @xcite . to contextualize the performance of these two families of approaches, we also describe a widely used information retrieval baseline, bm25, which uses surface lexical similarity and thus tells us how effectively we can perform fact tracing without even having to access a model. ## fact tracing datasets we propose two datasets to measure fact tracing approaches: ftrace-trex, a natural language dataset with real facts derived from the trex dataset, and ftrace-synth, a synthetic dataset with novel facts using made-up entities and relations. for each dataset, we define an attribution set containing all lm training examples that might be considered proponents and a query set containing test examples, each annotated with their ground truth proponents from the attribution set. the examples in these sets are masked language modeling examples, each a (masked input, output, facts) tuple. ## experimental setup our experiments aim to answer the questions of (1) whether tda methods can be used as effective fact tracing tools (compared to simple ir baselines), (2) which configurations make them most effective (exploring many variations), and (3) analyzing the weaknesses of tracin, in particular its sensitivity to when the knowledge is learned (the aforementioned "saturation" hypothesis). ## conclusion we introduced a new dataset and benchmark for fact tracing: the task of tracing language models' assertions back to the training examples that provided evidence for those predictions. we evaluated gradient-based and embedding-based tda methods and found that they perform worse than a standard ir baseline (bm25) even in settings that favor tda methods. we investigated the effects of layer selection, model checkpoints and fine-tuning. our ablative analysis suggests that saturation is an important factor affecting the performance of current methods. much is needed to improve these methods before they can be reliably used for fact tracing. we hope that this benchmark will enable future research on fact tracing, by mitigating computational challenges and establishing a principled ground truth.