How Large Language Models Encode Context Knowledge? A Layer-Wise Probing Study

TL;DR

This paper makes the first attempt to investigate the layer-wise capability of LLMs through probing tasks, leveraging the powerful generative capability of ChatGPT to construct probing datasets, providing diverse and coherent evidence corresponding to various facts.

Abstract

Previous work has showcased the intriguing capability of large language models (LLMs) in retrieving facts and processing context knowledge. However, only limited research exists on the layer-wise capability of LLMs to encode knowledge, which challenges our understanding of their internal mechanisms. In this paper, we devote the first attempt to investigate the layer-wise capability of LLMs through probing tasks. We leverage the powerful generative capability of ChatGPT to construct probing datasets, providing diverse and coherent evidence corresponding to various facts. We employ $\mathcal V$-usable information as the validation metric to better reflect the capability in encoding context knowledge across different layers. Our experiments on conflicting and newly acquired knowledge show that LLMs: (1) prefer to encode more context knowledge in the upper layers; (2) primarily encode context knowledge within knowledge-related entity tokens at lower layers while progressively expanding more knowledge within other tokens at upper layers; and (3) gradually forget the earlier context knowledge retained within the intermediate layers when provided with irrelevant evidence. Code is publicly available at https://github.com/Jometeorie/probing_llama.

Figures (9)

• Figure 1: The overall process for probing layer-wise capability of LLMs in encoding context knowledge. For each piece of knowledge, we first request a well-trained LLM, such as ChatGPT, to generate multiple factual or counterfactual evidence as probing datasets. Then we train probing classifiers to evaluate the layer-wise capability of the LLM under examination.
• Figure 2: The principal component analysis (PCA) visualization of how context knowledge is processed across different layers in LLaMA 2 Chat 13B, where instances denoted green and red signify factual and counterfactual evidence, respectively. $\mathcal{V}$-usable information (Vi) is found to be more effective in distinguishing between dataset difficulty than accuracy (Acc).
• Figure 3: The heatmap of probing results on LLaMA 2 Chat 13B. We select the question What is Mike Flanagan's occupation? as a case study and display the layer-wise $\mathcal{V}$-information for each token in the question.
• Figure 4: The average layer-wise $\mathcal{V}$-information (Vi) of the last token in the questions for each LLM.
• Figure 5: We categorize the subjects and relations mentioned in the questions as one class (red) while considering other tokens as another class (green). By comparing the differences of the average $\mathcal{V}$-information between these two classes, it is capable of detecting the LLM's level of attention to knowledge-related entity tokens.