[Paper Review] How Large Language Models Encode Context Knowledge? A Layer-Wise Probing Study

Tianjie Ju, Weiwei Sun+4 more|arXiv (Cornell University)|Feb 25, 2024

Topic Modeling4 citations

TL;DR

This paper proposes a layer-wise probing framework to study how large language models (LLMs) encode context knowledge, using ChatGPT-generated evidence and $Σ$-usable information as a metric. It finds that LLMs encode more context knowledge in upper layers, transfer knowledge from entity tokens to broader context over depth, and suffer from interference when exposed to irrelevant evidence, revealing key dynamics in their internal knowledge retention.

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.

Motivation & Objective

• To investigate how large language models encode context knowledge across their hidden layers, addressing a gap in understanding their internal mechanisms.
• To develop a reliable probing framework using diverse, coherent evidence generated by ChatGPT to assess layer-wise knowledge encoding.
• To measure and compare the capability of different LLM layers to encode conflicting or newly acquired knowledge using $Σ$-usable information as a discriminative metric.
• To examine the long-term memory capacity of LLMs by testing their retention of context knowledge under increasing irrelevant evidence.
• To uncover the role of self-attention and positional encoding in shaping knowledge representation across layers.

Proposed method

• Leverage ChatGPT to generate diverse, coherent evidence for facts and counterfactuals to construct probing datasets.
• Feed the evidence and questions into LLaMA 2 models and extract hidden layer representations for probing.
• Use $Σ$-usable information as the primary metric to assess how distinguishable context knowledge is in layer representations.
• Design probing tasks for conflicting knowledge (e.g., correcting parametric facts) and newly acquired knowledge (e.g., passwords) to evaluate encoding behavior.
• Systematically vary evidence relevance and positional encoding to test robustness and interference effects in intermediate layers.
• Visualize layer-wise probing results via heatmaps to identify patterns in knowledge encoding across layers and token types.

Experimental results

Research questions

• RQ1 How do LLMs distribute context knowledge encoding across their hidden layers?
• RQ2 Where and how is context knowledge encoded—primarily in entity tokens or across broader context?
• RQ3 How does the presence of irrelevant evidence affect the retention of newly acquired knowledge in intermediate layers?
• RQ4 To what extent does self-attention facilitate the transfer of knowledge from entity tokens to other tokens across layers?
• RQ5 Does positional encoding significantly influence the detectability of context knowledge in layer representations?

Key findings

• LLMs prefer to encode context knowledge in upper layers, with the highest $Σ$-usable information observed in deeper layers for conflicting knowledge scenarios.
• Knowledge-related entity tokens encode more context knowledge in lower layers, but this advantage diminishes and is overtaken by other tokens in upper layers.
• Self-attention mechanisms drive the progressive transfer of knowledge from entity tokens to surrounding context tokens as layers deepen.
• When exposed to irrelevant evidence, LLMs show significant performance degradation in intermediate layers, indicating interference and reduced retention of newly acquired knowledge.
• Positional encoding (e.g., adding newline tokens) has minimal impact on probing results, suggesting it acts as negligible noise in knowledge encoding detection.
• The model’s ability to encode context knowledge remains robust even for newly acquired facts like passwords, though performance degrades with increasing irrelevant context.