[Paper Review] Linearity of Relation Decoding in Transformer Language Models

Evan Hernandez, Arnab Sen Sharma+6 more|arXiv (Cornell University)|Aug 17, 2023

Topic ModelingComputer Science3 citations

TL;DR

This paper investigates how transformer language models encode and decode relational knowledge, finding that for 48% of tested relations—such as 'plays the instrument' or 'has color'—this knowledge is well-approximated by a linear transformation (LRE) applied to subject representations. The method estimates these linear relations directly from the model's Jacobian on a single prompt, enabling interpretable knowledge editing and visualization via an 'attribute lens' that reveals hidden factual knowledge even when incorrect outputs are generated.

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.

Motivation & Objective

• To investigate whether relational knowledge in transformer language models is encoded via linear transformations on subject representations.
• To determine whether a single prompt can be used to estimate linear relational embeddings (LREs) that accurately predict object tokens for various relations.
• To evaluate the extent to which LREs can be used to edit model behavior and visualize hidden factual knowledge.
• To compare the performance of LREs against random baselines and assess their robustness across diverse relation types.
• To develop and validate an 'attribute lens' visualization tool that reveals when and where models retrieve specific relational facts.

Proposed method

• The authors estimate linear relational embeddings (LREs) by computing the Jacobian of the model’s output with respect to the subject representation from a single prompt, approximating the first-order behavior of the model.
• For each relation, they train a linear probe (W_r, b_r) to map subject representations to object representations, using a small set of subject-object pairs.
• They validate LREs by measuring how well the predicted object distribution matches the actual next-token distribution in the model.
• They introduce an 'attribute lens' visualization that projects hidden representations at each layer through the estimated LRE, showing object-token distributions for a given relation at each layer.
• They evaluate LREs on 47 relations across GPT and LLaMA models, including factual, commonsense, and biased relations, using both single-token and multi-token completions.
• They assess causal effects by editing subject representations using LREs and measuring changes in model output.

Experimental results

Research questions

• RQ1 To what extent is relational knowledge in transformer language models encoded as a linear transformation of subject representations?
• RQ2 Can a single prompt be used to reliably estimate a linear relational embedding (LRE) that captures a relation’s object prediction across diverse subjects?
• RQ3 How do LREs compare to random baselines in predicting correct object tokens for various relations?
• RQ4 In cases where LREs fail, what alternative mechanisms might the model use to encode and retrieve relational knowledge?
• RQ5 Can LREs be used to visualize and detect hidden factual knowledge in models, even when the model generates incorrect outputs?

Key findings

• For 48% of the 47 relations tested, linear relational embeddings (LREs) provide a robust and faithful approximation of the model’s relation decoding process.
• LREs estimated from a single prompt can accurately predict object tokens across diverse subjects, with performance significantly outperforming random baselines.
• The attribute lens visualization successfully reveals hidden factual knowledge—such as the correct capital of Norway—even when the model generates incorrect outputs like 'London'.
• LREs can be used to edit subject representations in a way that causally alters model outputs, demonstrating their interpretability and utility in model control.
• The study identifies a heterogeneous knowledge representation strategy: while some relations are linearly encoded, others rely on more complex, non-linear mechanisms, indicating that linearity is not universal.
• The method reveals that relations with easier or harder random baselines tend to be more or less amenable to LRE approximation, respectively.