Probing Language Models on Their Knowledge Source (2024)

Abstract: Large Language Models (LLMs) often encounter conflicts between their learned, internal (parametric knowledge, PK) and external knowledge provided during inference (contextual knowledge, CK). Understanding how LLMs models prioritize one knowledge source over the other remains a challenge. In this paper, we propose a novel probing framework to explore the mechanisms governing the selection between PK and CK in LLMs. Using controlled prompts designed to contradict the model’s PK, we demonstrate that specific model activations are indicative of the knowledge source employed. We evaluate this framework on various LLMs of different sizes and demonstrate that mid-layer activations, particularly those related to relations in the input, are crucial in predicting knowledge source selection, paving the way for more reliable models capable of handling knowledge conflicts effectively.

## introduction large language models (llms) have demonstrated remarkable proficiency in memorizing and retrieving massive amounts of information. despite these strengths, llms often struggle when exposed to novel information not seen during training @xcite or when there is a conflict between their parametric knowledge (pk) and the context knowledge (ck) provided at inference @xcite . such discrepancies can lead to erroneous outputs, a phenomenon that remains a significant challenge in llms applications @xcite . while several approaches, such as reinforcement learning and retrieval-augmented generation, have been proposed to mitigate these issues @xcite , the mechanisms by which llms select and prioritize knowledge sources are not well understood, suggesting a gap in current research methodologies. this paper explores the internal dynamics of llms, and more precisely decoder-only layers, focusing on their decision-making processes regarding the use of ck versus pk. by prompting the llm in a way that contradicts its pk, we probe the model's knowledge-sourcing behaviors. by training a linear classifier on model activations, our experiments reveal that certain activations correlate with determining whether context or parametric knowledge predominates in the generated outputs. in this paper, we make the following key findings and contributions: ## related work the understanding of the mechanisms and knowledge localization within transformers has progressed through various studies. on the one hand, some work investigated the pk-based outputs (factual setting) @xcite @xcite @xcite . these works hypothesized that llms store parametric information in the multi-layer perceptron (mlp), which acts as a key-value memory, subsequently accessed by the multi-head self-attention (mhsa) mechanisms. on the other hand, other studies focused on the ck-based outputs and concluded that processing ck, as opposed to pk, is not specifically localized in the llm's parameters @xcite . @xcite employed an attribution method @xcite to identify the most influential attention heads responsible for generating pk and ck outputs, and subsequently adjusted the weights of these heads to modify the source of knowledge. their work however focuses exclusively on knowledge specific to capital cities and relies on causal tracing, which is costly to compute. in contrast, our work utilizes a probing approach that uses a classifier on the llm activations to identify the source of knowledge, leveraging the insights from previous research on the respective roles of mlps and mhsas in the inference process. we extend the scope of @xcite by incorporating a dataset with a broader range of knowledge categories (pararel @xcite ), moving beyond just capital cities. ## methodology we aim to show that specific activations correlate with the used knowledge source, parametric or context knowledge. in order to probe llms, we construct prompts that are composed of inputs rep-resenting information about a subject s that contradicts what the model learned during training, followed by a query about the same subject (see figure 1 ). if the model answers according to the prompt, then it uses context knowledge. on the other hand, if the model answers according to what it learned, then it is using its parametric knowledge. in the following two sections, we define more formally pk and ck. ## probing set-up to build our probing dataset, we associate each tuple (s, r, o, ō) ∈ p k with a prompt relation group id relations #examples geographic-geopolitic-language is-headquarter, located-in, headquarters-in, locate, share-border, is-twincity-of, located, border-with, is-located, work-in-area, which-is-located, capital-city-of, part-of-continent, capital-of, headquarter, belong-tocontinent, based-in, is-citizen-of, that-originate-in, originate-in, is-in, found-in, share-common-border, is-native-to, is-originally-from, passaway-in, born-in, hold-citizenship-of, have-citizenship-of, citizen-of, start-in, formulate-in, legal-term, tie-diplomatic-relations, maintainsdiplomatic-relations, have-diplomatic-relations, native, mother-tongue, original-language-is, the-official-language, communicate 2815 corporate-products-employment product-manufacture-by, develop-by, owned-by, product-develope-by, product-release-by, create-by, product-of, produce-by, owner, is-productof, is-part-of, who-works-for, employed-by, who-employed-by, works-for, work-in-field, profession-is, found-employment prompt(s, r, ō) that corresponds to a natural language statement of (s, r, ō) followed by q(s, r) (see figure 1 ). each prompt is associated with a label among ck, pk, and nd, where and with nd (not defined) otherwise. in this work, we discard tuples associated with nd. we specifically probe the activations ō of the object, s q of the subject in the query, and r q the relation in the query. as each of these elements may have multiple tokens, we use their last tokens as their representative (e.g. for "washington" → ["wash", "inghton"], we consider the activations of the token "inghton"). the fact that this token representation summarizes the entity is intuitively true for decoder-only models and has been experimentally validated in @xcite . note that our first probe targets ō as this is where the knowledge conflict starts (e.g. bill gates is the founder of apple(ō). bill gates(s q ) is the founder of(r q ) _). ## results and discussion in figure 5 of the relation token, the success rate increases significantly, reaching about 80%. this finding is consistent with prior research, which indicates that llms primarily store knowledge in the mlps @xcite . moreover, it supports geva et al. ( 2023 )'s insights on the information extraction process, where the relation token retrieves attributes from s q (a process referred to as attribute extraction). additionally, it is noteworthy that the knowledge source can be detected directly starting from the perturbing object ō. this shows that detecting a potentially harmful conflict knowledge statement is possible early in the llm inference process. mhsa activations are less connected to the used knowledge source than mlp and mlp-l1 activations. the results of the control experiments conducted on the first token of the input indicate that the learned patterns in the object, subject, and relation are not arbitrary. the success rates of most llms for the first token appear to be random (about 0.5), with the exception of pythia-1.4b, where the first token provides a slight indication of the knowledge source, although no significant fluctuations are observed. finally, compared to @xcite , we show in this work that it is possible to predict the knowledge source based on the sole activations of an llm, and, even more importantly, that we predict this for multiple relations rather than being limited to a single relation. ## subject frequency vs. knowledge source to understand what makes an llm select the ck object over the pk object, we observed the subject ## probing dataset examples ## conclusion in this study, we introduced a novel probing framework to investigate if we can detect when llms switch from pk to ck. our findings reveal that specific model activations are significantly correlated with the used knowledge source. this opens the door for future work investigating the mechanism at play when such a switch occurs, and finally to building models that can better control this behavior.