Tracing Relational Knowledge Recall in Large Language Models

2026-04-21

AI Summary

This study investigates how large language models encode relational knowledge during text generation, focusing on which internal representations are amenable to linear probing for relation classification and why certain relations are more linearly decodable than others. Through systematic evaluation of latent representations from attention heads and MLP outputs—combined with linear probes, feature attribution, attention decomposition, and residual stream tracing—we identify the contribution of attention heads to the residual stream as a key signal for relation classification. Our analysis reveals that probe performance strongly correlates with relation specificity, entity connectivity, and the distribution of relational signals across attention heads. The work pinpoints optimal representation sources, establishes an interpretable link between relational properties and probing efficacy, and introduces a fine-grained, token-level attribution method for detailed analysis.

Technology Category

• Reasoning under Uncertainty: Relational Probabilistic Models
• Natural Language Processing: Sentence-level Semantics, Textual Inference, etc.
• Machine Learning: Large Multimodal Models (LMMs)

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web query analysis, representation and understanding

Abstract

We study how large language models recall relational knowledge during text generation, with a focus on identifying latent representations suitable for relation classification via linear probes. Prior work shows how attention heads and MLPs interact to resolve subject, predicate, and object, but it remains unclear which representations support faithful linear relation classification and why some relation types are easier to capture linearly than others. We systematically evaluate different latent representations derived from attention head and MLP contributions, showing that per-head attention contributions to the residual stream are comparatively strong features for linear relation classification. Feature attribution analyses of the trained probes, as well as characteristics of the different relation types, reveal clear correlations between probe accuracy and relation specificity, entity connectedness, and how distributed the signal on which the probe relies is across attention heads. Finally, we show how token-level feature attribution of probe predictions can be used to reveal probe behavior in further detail.

Problem

• relational knowledge
• large language models
• linear probes
• relation classification
• latent representations

Innovation

• relational knowledge
• linear probing
• attention heads
• feature attribution
• large language models