PaperlayerView exact revision on arXiv

Source-linked AI summary

Linearity of Relation Decoding in Transformer Language Models

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

arXiv:2308.09124v2cs.CL

TL;DR

The paper asks how transformer language models represent relational knowledge and tests whether relation decoding is linearly encoded in subject representations. It estimates affine mappings from local Jacobians and evaluates their predictive and causal behavior. Some relations fit this interpretable scheme, but many accurately predicted relations do not.

  • Problem

    How transformer LMs map enriched subject representations to language-based relational predictions has remained an open question.

  • Method

    The paper estimates affine linear relational embeddings from first-order Jacobian approximations of the LM’s subject-to-object computation and tests them through prediction matching and representation editing.

  • Results

    48% of 47 tested relations had robust LREs that faithfully recovered subject–object mappings for a majority of subjects, while some accurately predicted relations had no LRE.

  • Takeaways & Limitations

    Transformer LMs use a simple, interpretable relational representation for some knowledge, but encode and extract superficially similar facts in different ways.

  • Takeaways & Limitations

    Some relations are accurately predicted despite lacking linear decodability, and layer normalization can make affine estimates inaccurate in magnitude.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper shows that some transformer language-model relational knowledge is represented through simple affine mappings from subject representations to decodable object representations. This strategy is useful across diverse relations but is not deployed universally.

  • Core contribution: Transformer LMs can implicitly implement affine linear relational embeddings that map subject representations to object representations.The inputs are intermediate-layer subject states, while outputs are late-layer states decoded into next-token distributions.
  • Empirical results: For many relations, relational knowledge is decoded directly from subject representations through approximately affine transformations.These transformations can be computed from the LM Jacobian on a prompt expressing the relation.
  • Empirical results: 48% of the 47 tested relations had robust LREs that faithfully recovered subject–object mappings for a majority of subjects.The relations covered more than 10k facts involving famous entities, commonsense knowledge, and implicit biases.
  • Applications: LREs can be used to edit subject representations and control the language model’s predicted output.The editing procedure provides evidence that these representations are causally linked to model behavior.
  • Applications: The attribute lens visualizes object-token distributions across layers, revealing when relational knowledge is retrieved and when attributes remain internal.It can expose attribute knowledge that does not reach the model’s output.
  • Conclusion: Some relational knowledge is simple, interpretable, and structured, but superficially similar facts may use different encoding and extraction strategies.The paper therefore presents linear relational embeddings as a heterogeneous rather than universal representation system.

2 BACKGROUND: RELATIONS AND THEIR REPRESENTATIONS

Transformer language models enrich subject representations with factual information and must then decode relevant objects for prompted predictions. Linear relational embeddings provide one proposed connectionist representation for this relational readout, but the model’s mapping remains an open question.

  • Knowledge representations: Transformer MLP layers have been proposed as key–value stores that enrich entity representations with relevant factual knowledge.Across layers, these mechanisms populate a subject representation with facts associated with the entity.
  • Knowledge representations: After enrichment, the language model retrieves the fact relevant to the prompt and writes the object into its final representation before predicting the next token.This readout is required for prompts such as “Miles Davis plays the”.
  • Open question: How transformer LMs map enriched entity representations to language-based predictions has remained an open question.The paper investigates whether this learned readout is near-linear for a subset of relations.
  • Relational structures: Relational knowledge can be represented symbolically as subject–relation–object triplets, such as (Rome, is-capital-of, Italy).These triplets provide a structured representation of facts about entities and their relations.
  • Relational structures: Embedding relational structures in deep networks is difficult because vector geometry must accommodate the flexibility of symbolic triplets.Surveys identify more than 40 techniques for representing such structures.
  • Linear relational embeddings: Linear relational embedding represents a relation with an operator that maps a subject vector to an object vector, optionally adding a translation bias.The affine form learns both a relation-specific matrix W_r and translation b_r.

3 FINDING AND VALIDATING LINEAR RELATIONAL EMBEDDINGS

The paper estimates linear relational embeddings by locally approximating the transformer’s subject-to-object computation with Jacobians, then evaluates both predictive faithfulness and causal editing. Layer normalization limits magnitude accuracy, while practical edits use scaling and pseudoinverses.

  • Finding LREs: The transformer’s relation-decoding computation is modeled as o = F(s, c), mapping a subject representation to a directly decoded object representation in relation-evoking context.The subject is read at an intermediate layer and the object is the last-layer state used for the prediction.
  • Finding LREs: The method estimates an affine LRE from the Jacobian W = ∂F/∂s using a first-order Taylor approximation around a reference representation.This is the central local approximation used to obtain the relation operator and bias.
  • Finding LREs: The approximation requires F to behave nearly linearly across subject representations for the same relation.In practice, W and b are estimated as mean Jacobian and bias values across n examples under a zero-noise expectation assumption.
  • Limitations: Layer normalization prevents input-scale changes from being transmitted to output-scale changes, so the affine estimate may capture direction better than magnitude.The subject and object states are both affected by layer normalization in the decoding pipeline.
  • Finding LREs: A scalar β greater than 1 is used to make the estimated operator steeper because the calculated W underestimates the magnitude of change.This compensates for the empirical underestimation described in the method.
  • Evaluating LREs: Faithfulness measures whether LRE predictions from subject representations match the full transformer’s next-token predictions.The evaluation compares the LRE output with the transformer prediction using the LM decoder head.
  • Evaluating LREs: Causality is tested by inverting the LRE to compute an edit direction that should change the predicted object for a subject.The edited representation approximates another subject representation associated with the target object.
  • Evaluating LREs: Successful intervention requires the target object’s first token to become the LM’s top prediction after editing.The procedure uses a low-rank pseudoinverse rather than a potentially ill-conditioned matrix inverse.

4 EXPERIMENTS

Experiments evaluate LRE faithfulness, alternative linear approximations, causal influence, and layer dependence across diverse relations. LREs often track and causally influence LM behavior, but some accurately predicted relations resist linear decoding and vary substantially by layer.

  • Faithfulness: Over 60% faithfulness was achieved for almost half of the relations, indicating that those relations are linearly decodable from subject representations.Faithfulness measures how accurately LREs predict the transformer output for each relation.
  • Faithfulness: LRE captured LM behavior most faithfully across relation types, while low-performing subject-embedding, translation, and identity baselines show that enrichment, projection, and bias terms matter.The translation and identity baselines performed poorly, highlighting the necessity of both the projection and bias terms.
  • Faithfulness: For some relations, including Company CEO, no method exceeded 6% faithfulness despite GPT-J accurately predicting the CEOs of 69 companies.The result held across layers and random samples, suggesting more complicated, non-linear decoding for some relations.
  • Causality: LRE interventions provided causal evidence by approximating oracle performance, and faithfulness correlated strongly with causality at R = 0.84.For almost all relations, causality scores exceeded faithfulness scores, indicating that LREs could still enable successful edits when they did not fully capture computation.
  • Layer dependence: Faithfulness depended strongly on extraction layer: one relation remained linearly decodable from layers 7 through 17 before faithfulness plummeted, while removing relation-specific prompt context improved later-layer faithfulness.Qualitative post-edit generations remained fluent, and later layers may erase relational embeddings while supporting next-word prediction.

5 APPLICATION: THE ATTRIBUTE LENS

The attribute lens applies an LRE and decoder head to visualize relation-specific attribute predictions across hidden states and layers. It can expose latent knowledge even when distractions cause the model to output false statements.

  • Figure 7 compares LRE faithfulness across GPT-J layers for plays the sport of with and without relation-specific prompts.The caption states that faithfulness does not decrease in later layers when the object immediately follows the subject.
  • The attribute lens decodes relation-specific object-token distributions from hidden states using the LRE and transformer decoder head.It specializes the Logit Lens by visualizing attributes rather than only immediate next-token information.
  • Table 3 evaluates the attribute lens on 11,891 repetition-distracted and 11,891 instruction-distracted prompts designed to elicit false statements despite correct undistracted predictions.The table reports performance under these adversarial distraction conditions.
  • Figure 8 compares attribute-lens predictions for the relations plays sport and attended university on the same hidden states for Bill Bradley.Darker boxes indicate higher probabilities for the top predicted token.

6 RELATED WORK

The paper situates LREs among representation probes, knowledge-representation research, and methods for extracting knowledge from language models. Its approach differs from probing classifiers by extracting the relation map from the LM itself without training a new model.

  • Representation probes: Representation probes train classifiers to predict properties from neural representations, but such probes can overestimate knowledge when the classifier solves the task independently.Control tasks or constrained probe training can mitigate this concern.
  • Representation probes: This method differs from representation probing by extracting the LRE from the language model itself rather than introducing a separate training process.
  • Knowledge representation: Prior knowledge-representation work includes symbolic relational representations, knowledge-graph embeddings, and emergent knowledge representations in transformer language models.
  • Knowledge extraction: Prior knowledge-extraction methods query language models directly, while recent work investigates whether knowledge retrieval is localized within small model components.

7 CONCLUSION

The paper concludes that relation decoding can often be approximated by linear relational embeddings, while acknowledging that this approximation varies across relations. It also identifies potential uses for exposing internal biases and errors, alongside a limitation concerning superficial representations.

  • Relation decoding can often be well-approximated by linear relational embeddings, but some relations are better approximated than others.The authors relate this variation to relations having easier or harder random baselines.
  • LREs estimated from a small set of examples can produce faithful representations that are causally useful for studying model behavior.
  • LREs may expose internal model biases or errors before they are explicitly expressed in model output.Their causal effects could potentially support mitigation of undesired biases.
  • Representation-level interventions may be superficial unless internal model biases are corrected directly.The paper identifies exploring this issue as future work.
  • The released code and dataset support reproducing the experiments and inspecting the curated relation data.

A RELATIONS DATASET

The study uses a 47-relation dataset spanning factual, linguistic, bias, and commonsense knowledge, and estimates LREs from examples under an approximate-linearity assumption. Additional analyses examine estimator behavior and improve performance with a constant scaling factor.

  • A RELATIONS DATASET: The dataset contains 47 relations across factual, linguistic, bias, and commonsense groups.Six factual relations come from Wikidata, others from COUNTERFACT, and the remaining groups were newly curated.
  • A RELATIONS DATASET: Table 4 reports relation example counts and accurate predictions across language models using eight in-context examples, or five for LLaMA-13B.Results are means with standard deviations across 24 in-context-example selections.
  • LRE estimation: LRE parameter estimation assumes relation decoding is near-linear with a nonlinear error whose value and Jacobian have zero expectation over subject representations.
  • LRE estimation: The bias and projection terms are estimated by taking expectations over subject representations and using the mean Jacobian and bias.
  • Improving the estimate: βW with β > 1 provides a more accurate linear model of F than W alone, based on behavior between subject representation vectors.
  • Improving the estimate: A constant β for each language model is sufficient to achieve good performance across a range of relations.

D.1 DERIVATION OF EQN 7

The derivation motivates low-rank inversion for LRE-based edits and describes how rank, scaling, and hyperparameter choices affect faithfulness and causality.

  • D.1 DERIVATION OF EQN 7: The edit uses a low-rank inverse W † rather than a full inverse W −1 to avoid noise from small singular values.Singular values below a threshold can be excluded, preserving contributions from larger, more meaningful components.
  • D.1 DERIVATION OF EQN 7: Causality increases with rank until an optimum, then decreases as noisy singular values are introduced.The observed decline suggests that ranks beyond the optimum include less useful components.
  • D.1 DERIVATION OF EQN 7: The method requires a model-specific scalar β and relation-specific layer ℓr and inverse rank ρr.These parameters compensate for slope underestimation and determine where representations are extracted and how inversion is truncated.
  • D.1 DERIVATION OF EQN 7: A grid search selects ℓr and ρr for causal influence because faithfulness and causality agree strongly under this criterion.Selecting hyperparameters for faithfulness instead produces weaker agreement between the metrics.
  • D.1 DERIVATION OF EQN 7: For GPT-J, β = 2.25 gives the best correlation between faithfulness and causality, while causality is calculated with β = 1.0.The average causality score therefore remains constant across β values in the reported experiments.

E.2 LAYER-WISE LRE PERFORMANCE ON SELECTED RELATIONS (GPT-J)

Layer-wise and prompt-wise evaluations show that LRE performance varies across depth, examples, and interventions, while zero-shot transfer and causal behavior support the approximation for selected relations.

  • E.2 LAYER-WISE LRE PERFORMANCE ON SELECTED RELATIONS (GPT-J): Causality can decline before faithfulness reaches its limit because earlier subject states leak information during single-state activation patching.This leakage may reduce the efficacy of interventions at deeper layers and partly explains metric disagreement.
  • E.2 LAYER-WISE LRE PERFORMANCE ON SELECTED RELATIONS (GPT-J): LRE performance increases up to a layer and then drops, indicating a possible mode switch in later layers.The last row of Figure 11 contains relations for which LRE does not achieve satisfactory performance, consistent with non-linear decoding.
  • E.2 LAYER-WISE LRE PERFORMANCE ON SELECTED RELATIONS (GPT-J): For most relations, faithfulness and efficacy plateau after n = 5 examples, although the experiments use n = 8 for GPT-J.The reported hardware setup could fit at most eight examples on a single A6000.
  • E.2 LAYER-WISE LRE PERFORMANCE ON SELECTED RELATIONS (GPT-J): Different prompt templates produce minimal changes in faithfulness and causality scores.Performance scores are reported as means and standard deviations across 24 trials with different training-example sets.
  • E.2 LAYER-WISE LRE PERFORMANCE ON SELECTED RELATIONS (GPT-J): LRE trained with few-shot examples achieves similar results when applied to subject representations from zero-shot prompts.The zero-shot prompts contain only the subject and no context specifying the relation.
  • E.2 LAYER-WISE LRE PERFORMANCE ON SELECTED RELATIONS (GPT-J): LRE causality closely matches the oracle intervention across layers, while qualitative post-edit generations appear to preserve fluency.The oracle replaces the subject representation with the representation of another subject associated with the target object.

H LRE ON GPT2-XL AND LLAMA-13B

Evaluations on GPT2-xl and LLaMA-13B extend the analysis across models and relations, finding strong metric agreement and substantial cross-model similarity in relation-wise performance.

  • H LRE ON GPT2-XL AND LLAMA-13B: LRE performance is reported across factual, linguistic, bias, and commonsense relation categories for the three models.Figure 15 groups performance by relation category.
  • H LRE ON GPT2-XL AND LLAMA-13B: The experiments compare LRE with alternative linear decoding approaches on GPT2-xl and LLaMA-13B.The comparisons use the same baseline families described for the earlier GPT-J evaluation.
  • H LRE ON GPT2-XL AND LLAMA-13B: GPT2-xl and LLaMA-13B show high correlations between faithfulness and causality, with R = 0.85 and R = 0.83 respectively.Each plotted point represents LRE performance for a relation.
  • H LRE ON GPT2-XL AND LLAMA-13B: GPT-J relation-wise performance correlates strongly with GPT2-xl (R = 0.85) and LLaMA-13B (R = 0.71).GPT2-xl and LLaMA-13B are moderately correlated with each other at R = 0.58.

I LIMITATIONS

The paper identifies limitations in relation coverage, token-based correctness, and the treatment of relations with multiple valid objects.

  • I LIMITATIONS: The evaluation covers only 47 relations and omits numerical, physical, logical-inference, and multi-hop relations.The dataset spans several categories but does not represent all relation types.
  • I LIMITATIONS: First-token correctness can produce false positives when different objects share the same first token.The paper gives university names beginning with “University of” as an example and quantifies collision risk in Table 9.
  • I LIMITATIONS: The dataset records one canonical object even when a relation has multiple correct answers.Future evaluation could measure how well LRE estimates the LM’s distribution over candidate objects.
Loading 2308.09124v2…