Linearity of Relation Decoding in Transformer Language Models

Authors: Evan Hernandez, Arnab Sen Sharma, Tal Haklay, Kevin Meng, Martin Wattenberg, Jacob Andreas, Yonatan Belinkov, David Bau

What

This paper investigates how transformer language models (LLMs) represent relational knowledge, finding that a subset of relations can be approximated by linear transformations applied to subject representations.

Why

This work sheds light on the internal mechanisms of LLMs, revealing that some aspects of their knowledge representation are surprisingly simple and interpretable. This finding contributes to our understanding of how LLMs store and process information, potentially enabling more transparent and controllable AI systems.

How

The authors manually curate a dataset of relations and corresponding subject-object pairs. They then estimate a linear relational embedding (LRE) for each relation by calculating the Jacobian of the model’s computation on a prompt designed to elicit the relation. They evaluate the faithfulness of the LRE by measuring how well it predicts the model’s output for new subjects, and its causality by using it to edit subject representations and induce the model to predict different objects.

Result

The research shows that LREs can faithfully approximate LLM relation decoding for a significant portion of the tested relations. They also demonstrate the causal influence of these LREs by successfully manipulating model predictions via representation editing. Interestingly, the study reveals that not all relations are linearly encoded, suggesting a more complex, non-linear processing mechanism for certain types of information.

LF

The paper acknowledges limitations in the dataset size, the reliance on first-token correctness as an evaluation metric, and the assumption of single correct objects for relations. Future work could address these limitations, exploring a wider range of relations, refining the evaluation scheme, and investigating how LREs could be used to understand and mitigate biases in LLMs.

Abstract

Much of the knowledge encoded in transformer language models (LMs) may be expressed in terms of relations: relations between words and their synonyms, entities and their attributes, etc. We show that, for a subset of relations, this computation is well-approximated by a single linear transformation on the subject representation. Linear relation representations may be obtained by constructing a first-order approximation to the LM from a single prompt, and they exist for a variety of factual, commonsense, and linguistic relations. However, we also identify many cases in which LM predictions capture relational knowledge accurately, but this knowledge is not linearly encoded in their representations. Our results thus reveal a simple, interpretable, but heterogeneously deployed knowledge representation strategy in transformer LMs.