Search papers, labs, and topics across Lattice.
This paper investigates how relational knowledge is recalled by LLMs during text generation by training linear probes on various latent representations. They find that per-head attention contributions to the residual stream are the most effective features for linear relation classification. Further analysis reveals that probe accuracy correlates with relation specificity, entity connectedness, and the distribution of the signal across attention heads.
Forget scaling laws: the secret to extracting relational knowledge from LLMs lies in the specificity and connectedness of the relations themselves, and how their signals are distributed across attention heads.
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.
