Paper deep dive

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

Year: 2023Venue: ICLR 2024Area: Representation AnalysisType: EmpiricalEmbeddings: 109

Models: GPT-2 XL, GPT-J, LLaMA-13B

Intelligence

Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 90%

Last extracted: 3/12/2026, 8:11:38 PM

Summary

The paper investigates the linearity of relation decoding in transformer language models (LMs). It demonstrates that for many factual, commonsense, and linguistic relations, the computation mapping subject representations to object predictions can be approximated by a linear relational embedding (LRE), defined as an affine transformation. The authors validate this through faithfulness and causality experiments, showing that these linear approximations can be used to edit model outputs, while also noting that this strategy is not universally deployed across all relational knowledge in LMs.

Entities (4)

Transformer Language Models · model-architecture · 100%Linear Relational Embedding · methodology · 98%Jacobian · mathematical-concept · 95%Attribute Lens · tool · 90%

Relation Signals (3)

Linear Relational Embedding → approximates → Relation Decoding

confidence 95% · a portion of transformer LMs’ (highly non-linear) computation can be well-approximated linearly in contexts requiring relation prediction.

Transformer Language Models → implements → Linear Relational Embedding

confidence 90% · we show that LMs employ a simple system for representing a portion of their relational knowledge: they implicitly implement (an affine version of) a linear relational embedding

Attribute Lens → visualizes → Relational Knowledge

confidence 85% · This lets us visualize where and when the LM finishes retrieving knowledge about a specific relation

Cypher Suggestions (0)

No Cypher suggestions yet.

Abstract

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.

PDF

Open source PDF →Open local PDF →

Trouble viewing inline? Open PDF directly →

Full Text

108,941 characters extracted from source content.

Expand or collapse full text

Linearity of Relation Decoding in Transformer Language Models Evan Hernandez1111 Massachusetts Institute of Technology, 22footnotemark: 2Northeastern University, 33footnotemark: 3Technion IIT, 44footnotemark: 4Harvard University. Martin Wattenberg4 &Arnab Sen Sharma2111 Massachusetts Institute of Technology, 22footnotemark: 2Northeastern University, 33footnotemark: 3Technion IIT, 44footnotemark: 4Harvard University. Jacob Andreas1 &Tal Haklay3 Yonatan Belinkov3 &Kevin Meng1 David Bau2 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 LMs. 11footnotetext: Equal contribution. Correspondence to: dez@mit.edu, sensharma.a@northeastern.edu. 1 Introduction How do neural language models (LMs) represent relations between entities? LMs store a wide variety of factual information in their weights, including facts about real world entities (e.g., John Adams was elected President of the United States in 1796) and common-sense knowledge about the world (e.g., doctors work in hospitals). Much of this knowledge can be represented in terms of relations between entities, properties, or lexical items. For example, the fact that Miles Davis is a trumpet player can be written as a relation (plays the instrument), connecting a subject entity (Miles Davis), with an object entity (trumpet). Categorically similar facts can be expressed in the same type of relation, as in e.g., (Carol Jantsch, plays the instrument, tuba). Prior studies of LMs (Li et al., 2021; Meng et al., 2022; Hernandez et al., 2023) have offered evidence that subject tokens act as keys for retrieving facts: after an input text mentions a subject, LMs construct enriched representations of subjects that encode information about those subjects. Recent studies of interventions (Hase et al., 2023) and attention mechanisms (Geva et al., 2023) suggest that the mechanism for retrieval of specific facts is complex, distributed across multiple layers and attention heads. Past work establishes where relational information is located: LMs extract relation and object information from subject representations. But these works have not yet described what computation LMs perform while resolving relations. In this paper, we show that LMs employ a simple system for representing a portion of their relational knowledge: they implicitly implement (an affine version of) a linear relational embedding (LRE) scheme (Paccanaro & Hinton, 2001). Given a relation r such as plays the instrument, a linear relational embedding is an affine function ⁢()=Wr⁢+brsubscriptsubscriptLRE(s)=W_rs+b_rLRE ( s ) = Witalic_r s + bitalic_r that maps any subject representation ss in the domain of the relation (e.g., Miles Davis, Carol Jantsch) to the corresponding object representation oo (e.g., trumpet, tuba). In LMs, the inputs to these implicit LREs are hidden representations of subjects at intermediate layers, and their outputs are hidden representations at late layers that can be decoded to distributions over next tokens. Thus, a portion of transformer LMs’ (highly non-linear) computation can be well-approximated linearly in contexts requiring relation prediction. More specifically, we find that for a variety of relations: (a) transformer LMs decode relational knowledge directly from subject entity representations (ss in Figure 1); (b) for each such relation, the decoding procedure is approximately affine (LRE); and (c) these affine transformations can be computed directly from the LM Jacobian on a prompt expressing the relation (i.e. ∂/∂ / ∂ o / ∂ s). However, this is not the only system that transformer LMs use to encode relational knowledge, and we also identify relations that are reliably predicted in LM outputs, but for which no LRE can be found. Figure 1: Within a transformer language model, (a) how it resolves many relations r, such as plays the instrument, can be well-approximated by (b) a linear function R that maps subject representations ss to object representations oo that can be directly decoded. In GPT and LLaMA models, we search for LREs encoding 47 different relations, covering more than 10k facts relating famous entities (The Space Needle, is located in, Seattle), commonsense knowledge (banana, has color, yellow), and implicit biases (doctor, has gender, man). In 48%percent4848\%48 % of the relations we tested, we find robust LREs that faithfully recover subject–object mappings for a majority of the subjects. Furthermore, we find that LREs can be used to edit subject representations (Hernandez et al., 2023) to control LM output. Finally, we use our dataset and LRE-estimating method to build a visualization tool we call an attribute lens. Instead of showing the next token distribution like Logit Lens (nostalgebraist, 2020) the attribute lens shows the object-token distribution at each layer for a given relation. This lets us visualize where and when the LM finishes retrieving knowledge about a specific relation, and can reveal the presence of knowledge about attributes even when that knowledge does not reach the output. Our results highlight two important facts about transformer LMs. First, some of their implicit knowledge is representated in a simple, interpretable, and structured format. Second, this representation system is not universally deployed, and superficially similar facts may be encoded and extracted in very different ways. 2 Background: Relations and Their Representations 2.1 Representations of knowledge in language models For LMs to generate factually correct statements, factual information must be represented somewhere in their weights. In transformer LMs, past work has suggested that most factual information is encoded in the multi-layer perceptron layers (Geva et al., 2020). These layers act as key–value stores, and work together across multiple layers to enrich the representation of an entity with relevant knowledge (Geva et al., 2022). For instance, in the example from Figure 1, the representation ss of Miles Davis goes through an enrichment process where LM populates ss with the fact that he plays the trumpet as well as other facts, like him being born in Alton, IL. By the halfway point of the LM’s computation, ss contains all the information needed to predict a fact about the subject entity when the LM is prompted to retrieve it. Once ss is populated with relevant facts, the LM must decode the fact most relevant to its current prediction task. Formally, a language model is a distribution pLM⁢(x)subscriptLMp_LM(x)pLM ( x ) over strings x, so this information must be retrieved when the LM is prompted to decode a specific fact, such as when it estimates pLM(⋅∣Miles Davis plays the)p_LM(· Davis plays the)pLM ( ⋅ ∣ Miles Davis plays the ). Internally, the object must be decoded and written into the final representation oo before the next word (trumpet) is predicted. Techniques like the logit lens (nostalgebraist, 2020) and linear shortcut approaches (Belrose et al., 2023; Din et al., 2023) reveal that the LM’s final prediction can be read off of oo well before the final layer, and recent work (Geva et al., 2023) suggests that this occurs because specific attention heads (before the final layer) specialize in reading specific relations. Meanwhile, prior work studying the structure of ss suggests that even though transformers are complex, non-linear neural networks, attributes of entities can be linearly decoded from their representations (Li et al., 2021; Hernandez et al., 2023). But how transformer LMs themselves map from enriched entity representations to language-based predictions has remained an open question. Here, we will show that for a subset of relations the transformer LMs implement the learned readout operation in a near-linear fashion. 2.2 Neural representations of relations Why might we expect a linear representation scheme for relational information in the first place? Separate from (and largely prior to) work on neural language models, a long line of artificial intelligence research has studied how to represent relational knowledge. A classic symbolic approach is to encode relational knowledge triplets of the form (subject, relation, object). For example, one might express the fact that Rome is the capital of Italy as (Rome, is-capital-of, Italy). This triplet format is extremely flexible, and has been used for a variety of tasks (Richens, 1956; Minsky, 1974; Lenat, 1995; Miller, 1995; Berners-Lee et al., 2001; Bollacker et al., 2008). While representing relational triplets symbolically is straightforward, it is far less clear how to embed these relational structures in deep networks or other connectionist systems. Surveys (Ji et al., 2021; Wang et al., 2017) list more than 40 techniques. These variations reflect the tension between the constraints of geometry and the flexibility of the triplet representation. In many approaches, subject and object entities s and o are represented as vectors ∈ℝm,∈ℝnformulae-sequencesuperscriptℝsuperscriptℝs ^m,o ^ns ∈ blackboard_Rm , o ∈ blackboard_Rn; for a given relation r, we define a relation function R:ℝm→ℝn:→superscriptℝsuperscriptℝR:R^m ^nR : blackboard_Rm → blackboard_Rn, with the property that when (s,r,o)(s,r,o)( s , r , o ) holds, we have ≈R⁢(). o≈ R( s).o ≈ R ( s ) . One way to implement R is to use linear transformations to represent relations. For instance, in linear relational embedding (Paccanaro & Hinton, 2001), the relation function has the form R⁢()=Wr⁢subscriptR( s)=W_r sR ( s ) = Witalic_r s where WrsubscriptW_rWitalic_r is a matrix depending on relation r. A modern example of this encoding can be seen in the positional encodings of many transformers (Vaswani et al., 2017). More generally, we can write R as an affine transformation, learning both a linear operator WrsubscriptW_rWitalic_r and a translation brsubscriptb_rbitalic_r (Lin et al., 2015; Yang et al., 2021). There are multiple variations on this idea, but the basic relation function is: R⁢()=Wr⁢+br.subscriptsubscriptR( s)=W_r s+b_r.R ( s ) = Witalic_r s + bitalic_r . (1) 3 Finding and Validating Linear Relational Embeddings 3.1 Finding LREs Consider a statement such as Miles Davis plays the trumpet, which expresses a fact (s,r,o)(s,r,o)( s , r , o ) connecting a subject s to an object o via relation r (see Figure 1). Within the transformer’s hidden states, let ss denote the representation555Following insights from Meng et al. (2022) and Geva et al. (2023), we read ss at the last token of the subject. of the subject s (Miles Davis) at layer ℓ ℓ, and let oo denote the last-layer hidden state that is directly decoded to get the prediction of the object’s first token o (trumpet). The transformer implements a calculation that obtains oo from ss within a textual context c that evokes the relation r, which we can write =F⁢(,c)o=F(s,c)o = F ( s , c ). Our main hypothesis is that F⁢(,c)F(s,c)F ( s , c ) can be well-approximated by a linear projection, which can be obtained from a local derivative of F. Denote the Jacobian of F as W=∂F/∂W= ∂ F W = / start_ARG ∂ F end_ARG start_ARG ∂ s end_ARG. Then a first-order Taylor approximation of F about 0subscript0s_0s0 is given by: F⁢(,c) F(s,c)F ( s , c ) ≈F⁢(0,c)+W⁢(−0)absentsubscript0subscript0 ≈ F(s_0,c)+W (s-s_0 )≈ F ( s0 , c ) + W ( s - s0 ) =W⁢+b,absent =Ws+b,= W s + b , (2) where ⁢bwhere bwhere b =F⁢(0,c)−W⁢0absentsubscript0subscript0 =F(s_0,c)-Ws_0= F ( s0 , c ) - W s0 This approximation would only be reasonable if F has near-linear behavior when decoding the relation from any ss. In practice, we estimate W and b as the mean Jacobian and bias at n examples i,cisubscriptsubscripts_i,c_isitalic_i , citalic_i within the same relation, which gives an unbiased estimate under the assumption that noise in F has zero value and zero Jacobian in expectation (see Appendix B). That is, we define: W=i,ci⁢[∂F∂|(i,ci)]subscriptsubscriptsubscriptdelimited-[]evaluated-atsubscriptsubscript W=E_s_i,c_i [ . ∂ F% |_(s_i,c_i) ]W = blackboard_Es start_POSTSUBSCRIPT i , citalic_i end_POSTSUBSCRIPT [ divide start_ARG ∂ F end_ARG start_ARG ∂ s end_ARG |( s start_POSTSUBSCRIPT i , citalic_i ) end_POSTSUBSCRIPT ] and ⁢b=i,ci⁢[F⁢(,c)−∂F∂⁢|(i,ci)] and subscriptsubscriptsubscriptdelimited-[]evaluated-atsubscriptsubscript \;\;\;and\;\;\;b=E_s_i,c_i [% .F(s,c)- ∂ F \;s |% _(s_i,c_i) ]and b = blackboard_Es start_POSTSUBSCRIPT i , citalic_i end_POSTSUBSCRIPT [ F ( s , c ) - divide start_ARG ∂ F end_ARG start_ARG ∂ s end_ARG s |( s start_POSTSUBSCRIPT i , citalic_i ) end_POSTSUBSCRIPT ] (3) This simple formulation has several limitations that arise due to the use of layer normalization (Ba et al., 2016) in the transformer: for example, ss is passed through layer normalization before contributing to the computation of oo, and oo is again passed through layer normalization before leading to token predictions, so in both cases, the transformer does not transmit changes in scale of inputs to changes in scale of outputs. That means that even if Equation 3.1 is a good estimate of the direction of change of F, it may not be an accurate estimate of the magnitude of change. In practice, we find that the magnitude of change in F⁢(,c)F(s,c)F ( s , c ) is underestimated in our calculated W (see Appendix C for empirical measurements). To remedy this underestimation we make W in Section 3.1 steeper by multiplying with a scalar constant β(>1)annotatedabsent1β\;(>1)β ( > 1 ). So, for a relation r we approximate the transformer calculation F⁢(,cr)subscriptF(s,c_r)F ( s , citalic_r ) as an affine transformation LRE on ss: F⁢(,cr)≈LRE⁢()=β⁢Wr⁢+brsubscriptLREsubscriptsubscript F(s,c_r)≈ LRE(s)=β\,W_r% s+b_rF ( s , citalic_r ) ≈ LRE ( s ) = β Witalic_r s + bitalic_r (4) 3.2 Evaluating LREs When a linear relation operator LRE is a good approximation of the transformer’s decoding algorithm, it should satisfy two properties: Faithfulness. When applied to new subjects s, the output of LRE⁢()LRE LRE(s)LRE ( s ) should make the same predictions as the transformer. Given the LM’s decoder head D, we define the transformer prediction o and LRE prediction o^ oover start_ARG o end_ARG as: o=argmax⁢D⁢(F⁢(,c))t⁢ and ⁢o^=argmax⁢D⁢(LRE⁢())targmaxsubscript and ^argmaxsubscriptLRE o= targmax\ D(F(s,c))_t\;% \;and\;\; o= targmax\ D( LRE(s))_to = undert start_ARG argmax end_ARG D ( F ( s , c ) )t and over start_ARG o end_ARG = undert start_ARG argmax end_ARG D ( LRE ( s ) )t And we define faithfulness as the success rate of o=?o^superscript?^o ?= oo start_RELOP SUPERSCRIPTOP start_ARG = end_ARG start_ARG ? end_ARG end_RELOP over start_ARG o end_ARG, i.e., the frequency with which predictions made by LRE from only ss match next-token predictions made by the full transformer: argmax⁢D⁢(F⁢(,c))t=?argmax⁢D⁢(LRE⁢())tsuperscript?argmaxsubscriptargmaxsubscriptLRE targmax\ D(F(s,c))_t ?% = targmax\ D( LRE(s))_tundert start_ARG argmax end_ARG D ( F ( s , c ) )t start_RELOP SUPERSCRIPTOP start_ARG = end_ARG start_ARG ? end_ARG end_RELOP undert start_ARG argmax end_ARG D ( LRE ( s ) )t (5) Causality. If a learned LRE is a good description of the LM’s decoding procedure, it should be able to model causal influence of the relational embedding on the LM’s predictions. If it does, then inverting LRE tells us how to perturb ss so that the LM decodes a different object o′. Formally, given a new object ′superscript′o o′, we use LRE to find an edit direction Δ⁢Δ Δ s that satisfies: LRE⁢(+Δ⁢)=′LREΔsuperscript′ LRE(s+ )=o LRE ( s + Δ s ) = o′ (6) With β=11β=1β = 1, we can edit ss as follows:666Full derivation in Appendix D.1 ~=+Δ⁢, where ⁢Δ⁢=Wr−1⁢(′−)formulae-sequence~Δ where Δsuperscriptsubscript1superscript′ s=s+ ,\;\;where\;% =W_r^-1(o -o)over~ start_ARG s end_ARG = s + Δ s , where Δ s = Witalic_r- 1 ( o′ - o ) (7) We obtain ′superscript′o o′ from a different subject s′ that is mapped by F to o′ under the relation r. ~~ sover~ start_ARG s end_ARG here is essentially an approximation of ′superscript′s s′. Figure 2 illustrates this procedure. Figure 2: Illustration of the representation editing used to measure causality. Under the relation r=absentr=\,r =plays the instrument, and given the subject s=absents=\,s =Miles Davis, LM will predict o=absento=\,o =trumpet (a); and given the subject s′ = Cat Stevens, the output is o′ = guitar (b). If the computation from ss to oo is well-approximated by LRE parameterized by WrsubscriptW_rWitalic_r and brsubscriptb_rbitalic_r (c), then Δ⁢Δ sΔ s (d) should tell us the direction of change from ss to ′superscript′s s′. Thus, ~=+Δ⁢~Δ s=s+ ~ start_ARG s end_ARG = s + Δ s would be an approximation of ′superscript′s s′ and patching ~~ sover~ start_ARG s end_ARG in place of ss should change the prediction to o′ = guitar (f). Note that Equation 7 requires inverting WrsubscriptW_rWitalic_r, but the inverted matrix might be ill-conditioned. To make edits more effective, we instead use a low-rank pseudoinverse Wr†superscriptsubscript†W_r Witalic_r†, which prevents the smaller singular values from washing out the contributions of the larger, more meaningful singular values. See Section D.2 for details. We call the intervention a success if o′ is the top prediction of the LM after the edit: o′=?argmax⁢D⁢(F⁢(,cr∣:=+Δ⁢))superscript?superscript′argmaxassignconditionalsubscriptΔo ?= targmax\;D(F(% s,c_r :=s+ ))o′ start_RELOP SUPERSCRIPTOP start_ARG = end_ARG start_ARG ? end_ARG end_RELOP undert start_ARG argmax end_ARG D ( F ( s , citalic_r ∣ s := s + Δ s ) ) (8) Note that for both faithfulness and causality we only consider the first token of the object when determining success. Limitations of this approach are discussed in Appendix I. 4 Experiments We now empirically evaluate how well LREs, estimated using the approach from Section 3, can approximate relation decoding in LMs for a variety of different relations. Models. In all of our experiments, we study autoregressive language models. Unless stated otherwise, reported results are for GPT-J (Wang & Komatsuzaki, 2021), and we include additional results for GPT-2-XL (Radford et al., 2019) and LLaMA-13B (Touvron et al., 2023) in Appendix H. Dataset. To support our evaluation, we manually curate a dataset of 47 relations spanning four categories: factual associations, commonsense knowledge, implicit biases, and linguistic knowledge. Each relation is associated with a number of example subject–object pairs (si,oi)subscriptsubscript(s_i,o_i)( sitalic_i , oitalic_i ), as well as a prompt template that leads the language model to predict o when s is filled in (e.g., [s] plays the). When evaluating each model, we filter the dataset to examples where the language model correctly predicts the object o given the prompt. Table 1 summarizes the dataset and filtering results. Further details on dataset construction are in Appendix A. Table 1: Information about the dataset of relations used to evaluate LM relation decoding in LMs. These relations are drawn from a variety of sources. Evaluation is always restricted to the subset of (s, r, o) triples for which the LM successfully decodes o when prompted with (s, r). Category # Rel. # Examples # GPT-J Corr. Factual 26 9696 4652 Commonsense 8 374 219 Linguistic 6 806 507 Bias 7 213 96 Implementation Details. We estimate LREs for each relation using the method discussed in Section 3 with n=88n=8n = 8. While calculating W and b for an individual example we prepend the remaining n−11n-1n - 1 training examples as few-shot examples so that the LM is more likely to generate the answer o given a s under the relation r over other plausible tokens. Then, an LRE is estimated with Equation 3 as the expectation of WWWs and b s calculated on n individual examples. We fix the scalar term β (from Equation 4) once per LM. We also have two hyperparameters specific to each relation r; ℓrsubscriptℓ _rℓitalic_r, the layer after which ss is to be extracted; and ρrsubscript _rρitalic_r, the rank of the inverse W†superscript†W W† (to check causality as in Equation 7). We select these hyperparameters with grid-search; see Appendix E for details. For each relation, we report average results over 24 trials with distinct sets of n examples randomly drawn from the dataset. LREs are evaluated according to faithfulness and causality metrics defined in Equations (5) and (8). 4.1 Are LREs faithful to relations? We first investigate whether LREs accurately predict the transformer output for different relations, that is, how faithful they are. Figure 3 shows faithfulness by relation. Our method achieves over 60%percent6060\%60 % faithfulness for almost half of the relations, indicating that those relations are linearly decodable from the subject representation. Figure 3: Relation-wise LRE faithfulness to LM computation F. Horizontal red lines per relation indicate accuracy of a random-guess baseline. LRE is consistenly better than random guess and is predictive of the behavior of the transformer on most relations. However, for some relations such as company CEO or task done by tool, the transformer LM deviates from LRE, suggesting non-linear model computation for those relations. We are also interested in whether relations are linearly decodable from ss by any other method. Figure 4 compares our method (from Section 3) to four other approaches for estimating linear relational functions. We first compare with Logit Lens (nostalgebraist, 2020), where ss is directly decoded with the LM decoder head D. This essentially tries to estimate F⁢(,c)F(s,c)F ( s , c ) as an Identity transformation on ss. Next, we try to approximate F as Translation⁢()=+bTranslationTranslation(s)=s+bTranslation ( s ) = s + b, where b is estimated as ⁢[−]delimited-[]E[o-s]blackboard_E [ o - s ] over n examples. This Translation baseline, inspired by Merullo et al. (2023) and traditional word embedding arithmetic (Mikolov et al., 2013), approximates F from the intermediate representation of the last token of s until oo is generated (Figure 1). Then we compare with a linear regression model trained with n examples to predict oo from ss. Finally, we apply LRE on the subject embedding ssubscripte_seitalic_s before initial layers of the LM get to enrich the representation. Figure 4 shows that our method LRE captures LM behavior most faithfully across all relation types. This effect is not explained by word identity, as evidenced by the low faithfulness of LRE⁢(s)LREsubscript LRE(e_s)LRE ( eitalic_s ). Also, low performance of the Translation and Identity baselines highlight that both the projection and bias terms of Equation 4 are necessary to approximate the decoding procedure as LRE. Figure 4: Faithfulness comparison of different linear approximations of LM decoding stratified across different relation types. (a) Our method, LRE, applied on ss extracted after ℓrsubscriptℓ _rℓitalic_r (b) LRE applied on the subject embedding ssubscripte_seitalic_s. The performance different between (a) and (b) shows the importance of the enrichment process ss goes through in the earlier layers. (c) shows the performance of a linear regression model trained with n=88n=8n = 8 examples, which is outperformed by a LRE calculated with similar number of examples, n. (d) is Translation⁢()=+bTranslation Translation(s)=s+bTranslation ( s ) = s + b, where b is estimated as ⁢[−]delimited-[]E[o-s]blackboard_E [ o - s ] over n samples. The performance drop in (d) compared to (a) emphasizes the necessity of the projection term W. In (e) ss is directly decoded with the decoder head D. However, it is also clear (from Figure 3) that some relations are not linearly decodable from intermediate representations of the subject, despite being accurately predicted by the LM. For example, no method reaches over 6%percent66\%6 % faithfulness on the Company CEO relation, despite GPT-J accurately predicting the CEOs of 69 companies when prompted. This is true across layers (Figure 11 of Section E.2) and random sampling of n examples for approximating LRE parameters. This indicates that some more complicated, non-linear decoding approach is employed by the model to make those predictions. Interestingly, the relations that exhibit this behavior the most are those where the range is the names of peoples or companies. One possible explanation is that these ranges are so large that the LM cannot reliably linearly encode them at a single layer, and relies on a more complicated encoding procedures possibly involving multiple layers. 4.2 Do LREs causally characterize model predictions? We now have evidence that some relations are linearly decodable from LM representations using a first-order approximation of the LM. However, it could be that these encodings are not used by the LM to predict the next word, and instead are correlative rather than causal. To show that LREs causally influence LM predictions, we follow the procedure described in Figure 2 to use the inverse of LRE to change the LM’s predicted object for a given subject. Figure 5: LRE causality compared with different baselines. (a) LRE causality on best performing hyperparameters (layer ℓrsubscriptℓ _rℓitalic_r and rank ρrsubscript _rρitalic_r) for each relation r. (b) is our oracle baseline, inserting the representation ′superscript′s s′ of target subject s′ in place of ss. (c) in place of ss inserting o′subscriptsuperscript′e_o eitalic_o′, the row in the decoder head matrix D corresponding to o′, and (d) inserting ′superscript′o o′, the output of F⁢(′,cr)superscript′subscriptF(s ,c_r)F ( s′ , citalic_r ). In Figure 5 we compare our causality intervention with 3 other approaches of replacing ss such that LM outputs o′. If the model computation F from ss to oo is well-approximated by the LRE, then our intervention should be equivalent to inserting ′superscript′s s′ in place of ss. This direct substitution procedure thus provides an oracle upper-bound. Besides the oracle approach we include 2 more baselines; in place of ss, (1) inserting ′superscript′o o′, the output of F⁢(′,cr)superscript′subscriptF(s ,c_r)F ( s′ , citalic_r ) (2) inserting o′subscriptsuperscript′e_o eitalic_o′, the row in the decoder head matrix D corresponding to o′ as the embedding of o′. These two additional baselines ensure that our approach is not trivially writing the answer ′superscript′o o′ on the position of ss. Figure 6: Faithfulness is strongly correlated with causality (R=0.840.84R=0.84R = 0.84) when hyperparameters are selected to achieve best causal influence (in GPT-J β=2.252.25β=2.25β = 2.25). Each dot represents LRE performance for one relation. Bold dots indicate relations for which LRE is evaluated on ≥30absent30≥ 30≥ 30 test examples. Figure 14 of Section G.2 compares our method with the baselines for selected relations and across layers. The graphs illustrate how our method matches the oracle’s performance and differs from the other two baselines. This provides causal evidence that LRE approximates these relations well. Figure 6 depicts a strong linear correlation between our metrics when the hyperparameters were selected to achieve best causal influence. This means that when an LRE causally influences the LM’s predictions, it is also faithful to the model.777However, when the hyperparameters are chosen to achieve best faithfulness we did not notice such strong agreement between faithfulness and causality. Discussion on Appendix E. Also from Figure 6, for almost all relations LRE causality score is higher than its faithfulness score. This suggests that, even in cases where LRE can not fully capture the LM’s computation of the relation, the linear approximation remains powerful enough to perform a successful edit. While our focus within this work is on the linearity of relation decoding and not on LM representation editing, a qualitative analysis of the post-edit generations reveals that the edits are nontrivial and preserve the LM’s fluency; see Table 8 of Section G.2. 4.3 Where in the network do representations exhibit LREs? In the previous experiments, for each relation we had fixed ℓrsubscriptℓ _rℓitalic_r (the layer after which ss is to be extracted) to achieve the best causal influence on the model. However, there are considerable differences in LRE faithfulness when estimated from different layers. Figure 7 highlights an example relation that appears to be linearly decodable from representations in layer 7 until layer 17, at which point faithfulness plummets. Figure 11 of Section E.2 shows similar plots for other relations. Why might this happen? One hypothesis is that a transformer’s hidden representations serve a dual purpose: they contain both information about the current word (its synonyms, physical attributes, etc.), and information necessary to predict the next token. At some point, the latter information structure must be preferred to the former in order for the LM to minimize its loss on the next-word prediction task. The steep drop in faithfulness might indicate that a mode switch is happening in the LM’s representations at later layers, where the LM decisively erases relational embeddings in support of predicting the next word. Table 2: Example of prompts with and without relation-specific context. With relation-specific context LeBron James plays the sport of basketball Roger Federer plays the sport of tennis Lionel Messi plays the sport of Without relation-specific context LeBron James basketball Roger Federer tennis Lionel Messi Figure 7: Later layers switching roles from enriching ss to predicting next token. LRE performance across different layers of GPT-J for the factual relation plays the sport of with and without relation-specific prompt (Table 2). Faithfulness does not decrease in later layers when o immediately follows the s in the prompt. We indeed see in Figure 7 that LRE faithfulness keeps improving in later layers when we remove relation specific texts from our prompt, meaning the o immediately follows s in the prompt (Table 2). 5 Application: The Attribute Lens We apply the LRE to create a novel probing method we call the attribute lens that provides a view into a LM’s knowledge of an attribute of a subject with respect to a relation r. Given a linear function LRE for the relation r, the attribute lens visualizes a hidden state hh by applying the LM decoder head D to decode D⁢(LRE⁢())LRED( LRE(h))D ( LRE ( h ) ) into language predictions. The attribute lens specializes the Logit Lens (nostalgebraist, 2020) (which visualizes next token information in a hidden state hh by decoding D⁢()D(h)D ( h )) and linear shortcut approaches (Belrose et al., 2023; Din et al., 2023) (where an affine probe AℓsubscriptℓA_ Aroman_ℓ is trained to skip computation after layer ℓ ℓ, directly decoding the next token as D⁢(Aℓ⁢(ℓ))subscriptℓsubscriptℓD(A_ (h_ ))D ( Aroman_ℓ ( hroman_ℓ ) ), where ℓsubscriptℓh_ hroman_ℓ is the hidden state after ℓ ℓ). However, unlike these approaches concerned with the immediate next token, the attribute lens is motivated by the observation that each high-dimensional hidden state hh may encode many pieces of information beyond predictions of the immediate next token. Traditional representation probes (Belinkov & Glass, 2019; Belinkov, 2022) also reveal specific facets of a representation, but unlike probing classifiers that divide the representation space into a small number of output classes, the attribute lens decodes a representation into an open-vocabulary distribution of output tokens. Figure 8 illustrates the use of one attribute lens to reveal knowledge representations that contain information about the sport played by a person, and another lens about university affiliation. Table 3: The performance of the attribute lens on repetition-distracted prompts and instruction-distracted prompts that (almost) never produce the correct statement of a fact. Each row tests 11,891 prompts on GPT-J. Condition R@1 R@2 R@3 Repetition-distracted prompt 0.02 0.33 0.41 Attribute lens on RD prompts 0.54 0.65 0.71 Instruction-distracted prompt 0.03 0.17 0.25 Attribute lens on ID prompts 0.63 0.73 0.78 This attribute lens can be applied to analyze LM falsehoods: in particular, it can identify cases where an LM outputs a falsehood that contradicts the LM’s own internal knowledge about a subject. To quantify the attribute lens’s ability to reveal such situations, we tested the attribute lens on a set of 11,891 “repetition distracted” (RD) and the same number of “instruction distracted” (ID) prompts where we deliberately bait the LM to output a wrong o, but the LM would have predicted the correct o without the distraction. For example, in order to bait an LM to predict that The capital city of England is… Oslo, a RD prompt states the falsehood The capital city of England is Oslo twice before asking the model to complete a third statement, and an ID prompt states the falsehood followed by the instruction Repeat exactly. Although in these cases, the LM will almost never output the true fact (it will predict Oslo instead of London), the attribute lens applied to the last mention of the subject (England) will typically reveal the true fact (e.g., London) within the top 3 predictions. In Table 3 we show performance of the attribute lens to reveal latent knowledge under this adversarial condition. Figure 8: The attribute lens applied to the hidden states of GPT-J processing Bill Bradley was a. First two grids visualize the same set of hidden states under the attribute lens for two different relations. The word in each rectangle is the most-likely token in the distribution D⁢(LRE⁢())LRED( LRE(h))D ( LRE ( h ) ), where D applies the transformer decoder head; darker boxes correspond to higher probabilities of the top prediction. In (a) the relation is plays sport, and in (b) attended university, and both these cases reveal high-scoring predictions for attributes on the subject. For comparison, (c) sets LRE=ILRE LRE=ILRE = I which produces the Logit Lens (nostalgebraist, 2020) visualization, in which the visualized relation can be thought of as next token. (Senator Bill Bradley was formerly a basketball player who went to school at Princeton.) 6 Related Work Representation probes. The structure of the information represented within a neural network is a foundational problem that has been studied from several perspectives. One approach is to identify properties encoded representations by training a probing classifier to predict properties from the representations (Ettinger et al., 2016; Shi et al., 2016; Hupkes et al., 2018; Conneau et al., 2018; Belinkov et al., 2017; Belinkov & Glass, 2019). However, such approaches can overestimate the knowledge contained in a network if the classifier learns to solve a task on its own (Belinkov, 2022); the problem can be mitigated by comparing to a control task (Hewitt & Liang, 2019) or by limiting the training of the probe (Voita & Titov, 2020). Our method differs from probing by avoiding the introduction of a training process entirely: we extract the LRE from the LM itself rather than training a new model. Knowledge representation. Ever since emergent neural representations of relations were first observed in the original backpropagation paper (Rumelhart et al., 1986), neural representations of knowledge and relations have been a central problem in artificial intelligence. Section 2 surveys work in this area including knowledge graph embedding (Wang et al., 2017; Yang et al., 2021) and emergent knowledge representations within a transformer language model (Li et al., 2021; Meng et al., 2022; Hase et al., 2023; Hernandez et al., 2023; Geva et al., 2023). This paper builds on past work in by showing that relational aspects of this knowledge are encoded linearly. Knowledge extraction. The most direct way to characterize knowledge in LMs is to prompt or query them directly (Petroni et al., 2019; Roberts et al., 2020; Jiang et al., 2020; Shin et al., 2020; Cohen et al., 2023). However, recent work has suggested that model knowledge and knowledge retrieval may be localized within small parts of a language model (Geva et al., 2020; Dai et al., 2021; Meng et al., 2022; Geva et al., 2023). In this paper we further investigate the localized retrieval of knowledge and ask whether knowledge about relations and objects can be separated, and whether relations are represented as a linear relational embedding. 7 Conclusion Reverse-engineering the full mechanism of an LLM is a daunting task. In this work, we have found that a certain kind of computation, relation decoding, can often be well-approximated by linear relational embeddings. We have also found that some relations are better-approximated as LREs than others; relations that have an easier or harder random baseline fall on either end of the spectrum. We have shown that LREs estimated from a small set of examples lead to faithful representations that are causally linked to the LM’s behavior. Furthermore, LRE can be used to provide specialized attribute lens on the LM’s intermediate computation, even revealing cases of LM falsehoods. Ethics Statement By revealing and decoding internal model relations before they are explicitly expressed in model output, LREs can potentially be used to provide information about internal biases or errors, and the causal effects could provide a way to mitigate undesired biases. However, such representation-level representation might be only superficial without correcting internal biases in the model; exploring such applications is a natural step for future work. Reproducibility Statement The code and dataset are available at lre.baulab.info. We include full details about dataset curation in Appendix A. In addition to the experiment details at the beginnings of Sections 4 and 5, we describe hyperparameter sweeps in Appendix E. We ran all experiments on workstations with 80GB NVIDIA A100 GPUs or 48GB A6000 GPUs using HuggingFace Transformers (Wolf et al., 2019) implemented in PyTorch (Paszke et al., 2019). Acknowledgements This research has been supported by an AI Alignment grant from Open Philanthropy, the Israel Science Foundation (grant No. 448/20), and an Azrieli Foundation Early Career Faculty Fellowship. We are also grateful to the Center for AI Safety (CAIS) for sharing their compute resources, which supported many of our experiments. References Ba et al. (2016) Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450, 2016. Belinkov (2022) Yonatan Belinkov. Probing classifiers: Promises, shortcomings, and advances. Computational Linguistics, 48(1):207–219, 2022. Belinkov & Glass (2019) Yonatan Belinkov and James Glass. Analysis methods in neural language processing: A survey. Transactions of the Association for Computational Linguistics, 7:49–72, 2019. Belinkov et al. (2017) Yonatan Belinkov, Nadir Durrani, Fahim Dalvi, Hassan Sajjad, and James Glass. What do neural machine translation models learn about morphology? arXiv preprint arXiv:1704.03471, 2017. Belrose et al. (2023) Nora Belrose, Zach Furman, Logan Smith, Danny Halawi, Igor Ostrovsky, Lev McKinney, Stella Biderman, and Jacob Steinhardt. Eliciting latent predictions from transformers with the tuned lens. arXiv preprint arXiv:2303.08112, 2023. Berners-Lee et al. (2001) Tim Berners-Lee, James Hendler, and Ora Lassila. The semantic web. Scientific american, 284(5):34–43, 2001. Bollacker et al. (2008) Kurt Bollacker, Colin Evans, Praveen Paritosh, Tim Sturge, and Jamie Taylor. Freebase: a collaboratively created graph database for structuring human knowledge. In Proceedings of the 2008 ACM SIGMOD international conference on Management of data, p. 1247–1250, 2008. Cohen et al. (2023) Roi Cohen, Mor Geva, Jonathan Berant, and Amir Globerson. Crawling the internal knowledge-base of language models. arXiv preprint arXiv:2301.12810, 2023. Conneau et al. (2018) Alexis Conneau, German Kruszewski, Guillaume Lample, Loïc Barrault, and Marco Baroni. What you can cram into a single vector: Probing sentence embeddings for linguistic properties. arXiv preprint arXiv:1805.01070, 2018. Dai et al. (2021) Damai Dai, Li Dong, Yaru Hao, Zhifang Sui, Baobao Chang, and Furu Wei. Knowledge neurons in pretrained transformers. arXiv preprint arXiv:2104.08696, 2021. Din et al. (2023) Alexander Yom Din, Taelin Karidi, Leshem Choshen, and Mor Geva. Jump to conclusions: Short-cutting transformers with linear transformations. arXiv preprint arXiv:2303.09435, 2023. Ettinger et al. (2016) Allyson Ettinger, Ahmed Elgohary, and Philip Resnik. Probing for semantic evidence of composition by means of simple classification tasks. In Proceedings of the 1st workshop on evaluating vector-space representations for nlp, p. 134–139, 2016. Geva et al. (2020) Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. arXiv preprint arXiv:2012.14913, 2020. Geva et al. (2022) Mor Geva, Avi Caciularu, Kevin Ro Wang, and Yoav Goldberg. Transformer feed-forward layers build predictions by promoting concepts in the vocabulary space. arXiv preprint arXiv:2203.14680, 2022. Geva et al. (2023) Mor Geva, Jasmijn Bastings, Katja Filippova, and Amir Globerson. Dissecting recall of factual associations in auto-regressive language models. arXiv preprint arXiv:2304.14767, 2023. Hase et al. (2023) Peter Hase, Mohit Bansal, Been Kim, and Asma Ghandeharioun. Does localization inform editing? surprising differences in causality-based localization vs. knowledge editing in language models. arXiv preprint arXiv:2301.04213, 2023. Hernandez et al. (2023) Evan Hernandez, Belinda Z Li, and Jacob Andreas. Inspecting and editing knowledge representations in language models. arXiv preprint arXiv:2304.00740, 2023. Hewitt & Liang (2019) John Hewitt and Percy Liang. Designing and interpreting probes with control tasks. arXiv preprint arXiv:1909.03368, 2019. Hupkes et al. (2018) Dieuwke Hupkes, Sara Veldhoen, and Willem Zuidema. Visualisation and’diagnostic classifiers’ reveal how recurrent and recursive neural networks process hierarchical structure. Journal of Artificial Intelligence Research, 61:907–926, 2018. Ji et al. (2021) Shaoxiong Ji, Shirui Pan, Erik Cambria, Pekka Marttinen, and S Yu Philip. A survey on knowledge graphs: Representation, acquisition, and applications. IEEE Transactions on Neural Networks and Learning Systems, 33(2):494–514, 2021. Jiang et al. (2020) Zhengbao Jiang, Frank F Xu, Jun Araki, and Graham Neubig. How can we know what language models know? Transactions of the Association for Computational Linguistics, 8:423–438, 2020. Lenat (1995) Douglas B Lenat. Cyc: A large-scale investment in knowledge infrastructure. Communications of the ACM, 38(11):33–38, 1995. Li et al. (2021) Belinda Z Li, Maxwell Nye, and Jacob Andreas. Implicit representations of meaning in neural language models. arXiv preprint arXiv:2106.00737, 2021. Lin et al. (2015) Yankai Lin, Zhiyuan Liu, Maosong Sun, Yang Liu, and Xuan Zhu. Learning entity and relation embeddings for knowledge graph completion. In Proceedings of the AAAI conference on artificial intelligence, volume 29, 2015. Meng et al. (2022) Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. Advances in Neural Information Processing Systems, 35:17359–17372, 2022. Merullo et al. (2023) Jack Merullo, Carsten Eickhoff, and Ellie Pavlick. Language models implement simple word2vec-style vector arithmetic. arXiv preprint arXiv:2305.16130, 2023. Mikolov et al. (2013) Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation of word representations in vector space. arXiv preprint arXiv:1301.3781, 2013. Miller (1995) George A Miller. Wordnet: a lexical database for english. Communications of the ACM, 38(11):39–41, 1995. Minsky (1974) Marvin Minsky. A framework for representing knowledge, 1974. nostalgebraist (2020) nostalgebraist. interpreting gpt: the logit lens, 2020. URL https://w.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens. Paccanaro & Hinton (2001) Alberto Paccanaro and Geoffrey E. Hinton. Learning distributed representations of concepts using linear relational embedding. IEEE Transactions on Knowledge and Data Engineering, 13(2):232–244, 2001. Paszke et al. (2019) Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019. Petroni et al. (2019) Fabio Petroni, Tim Rocktäschel, Sebastian Riedel, Patrick Lewis, Anton Bakhtin, Yuxiang Wu, and Alexander Miller. Language models as knowledge bases? In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), p. 2463–2473, 2019. Radford et al. (2019) Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019. Richens (1956) Richard H Richens. Preprogramming for mechanical translation. Mechanical Translation, 3(1):20–25, 1956. Roberts et al. (2020) Adam Roberts, Colin Raffel, and Noam Shazeer. How much knowledge can you pack into the parameters of a language model? In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), p. 5418–5426, 2020. Rumelhart et al. (1986) David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. Learning representations by back-propagating errors. nature, 323(6088):533–536, 1986. Shi et al. (2016) Xing Shi, Inkit Padhi, and Kevin Knight. Does string-based neural mt learn source syntax? In Proceedings of the 2016 conference on empirical methods in natural language processing, p. 1526–1534, 2016. Shin et al. (2020) Taylor Shin, Yasaman Razeghi, Robert L Logan IV, Eric Wallace, and Sameer Singh. Autoprompt: Eliciting knowledge from language models with automatically generated prompts. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), p. 4222–4235, 2020. Touvron et al. (2023) Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023. Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. Voita & Titov (2020) Elena Voita and Ivan Titov. Information-theoretic probing with minimum description length. arXiv preprint arXiv:2003.12298, 2020. Wang & Komatsuzaki (2021) Ben Wang and Aran Komatsuzaki. GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model. https://github.com/kingoflolz/mesh-transformer-jax, May 2021. Wang et al. (2017) Quan Wang, Zhendong Mao, Bin Wang, and Li Guo. Knowledge graph embedding: A survey of approaches and applications. IEEE Transactions on Knowledge and Data Engineering, 29(12):2724–2743, 2017. Wolf et al. (2019) Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Huggingface’s transformers: State-of-the-art natural language processing. arXiv preprint arXiv:1910.03771, 2019. Yang et al. (2021) Jinfa Yang, Yongjie Shi, Xin Tong, Robin Wang, Taiyan Chen, and Xianghua Ying. Improving knowledge graph embedding using affine transformations of entities corresponding to each relation. In Findings of the Association for Computational Linguistics: EMNLP 2021, p. 508–517, 2021. Appendix A Relations Dataset The dataset consists of 47 relations stratified across 4 groups; factual, linguistic, bias, and commonsense. Six of the factual relations were scraped from Wikidata while the rest were drawn from the CounterFact dataset by Meng et al. (2022). linguistic, bias, and commonsense relations were newly curated by the authors. See Table 4 for details. Table 4: Number of examples per relation and the count of accurate predictions by different LMs. Each of the examples were tested using n=88n=8n = 8 ICL examples (n=55n=5n = 5 for LLaMA-13B). Results presented as mean (±plus-or-minus± std) of the counts across 24 trials with different set of ICL examples. For cases where the count of examples accurately predicted by the LM is less then n, it was replaced with "—". LRE estimation was not calculated for such cases. **absent^**start_FLOATSUPERSCRIPT * * end_FLOATSUPERSCRIPTWe also do not calculate LRE for LLaMA-13B where o is a year (president birth year and president election year) as LLaMA tokenizer splits years by digits (see Table 9). Category Relation # # Correct GPT-J GPT2-xl LLaMA-13B factual person mother 994994994994 182.8±5.8plus-or-minus182.85.8182.8± 5.8182.8 ± 5.8 83.5±12.3plus-or-minus83.512.383.5± 12.383.5 ± 12.3 613.5±1.5plus-or-minus613.51.5613.5± 1.5613.5 ± 1.5 person father 991991991991 206.1±7.6plus-or-minus206.17.6206.1± 7.6206.1 ± 7.6 109.2±6.1plus-or-minus109.26.1109.2± 6.1109.2 ± 6.1 675.5±5.5plus-or-minus675.55.5675.5± 5.5675.5 ± 5.5 person sport position 952952952952 243.3±94.1plus-or-minus243.394.1243.3± 94.1243.3 ± 94.1 199.9±67.9plus-or-minus199.967.9199.9± 67.9199.9 ± 67.9 200.0±0.0plus-or-minus200.00.0200.0± 0.0200.0 ± 0.0 landmark on continent 947947947947 797.5±74.7plus-or-minus797.574.7797.5± 74.7797.5 ± 74.7 421.9±85.3plus-or-minus421.985.3421.9± 85.3421.9 ± 85.3 200.0±0.0plus-or-minus200.00.0200.0± 0.0200.0 ± 0.0 person native language 919919919919 720.2±31.3plus-or-minus720.231.3720.2± 31.3720.2 ± 31.3 697.3±16.1plus-or-minus697.316.1697.3± 16.1697.3 ± 16.1 200.0±0.0plus-or-minus200.00.0200.0± 0.0200.0 ± 0.0 landmark in country 836836836836 565.2±14.2plus-or-minus565.214.2565.2± 14.2565.2 ± 14.2 274.1±58.2plus-or-minus274.158.2274.1± 58.2274.1 ± 58.2 200.0±0.0plus-or-minus200.00.0200.0± 0.0200.0 ± 0.0 person occupation 821821821821 131.5±28.5plus-or-minus131.528.5131.5± 28.5131.5 ± 28.5 42.0±10.3plus-or-minus42.010.342.0± 10.342.0 ± 10.3 404.0±25.5plus-or-minus404.025.5404.0± 25.5404.0 ± 25.5 company hq 674674674674 316.0±10.4plus-or-minus316.010.4316.0± 10.4316.0 ± 10.4 148.8±36.6plus-or-minus148.836.6148.8± 36.6148.8 ± 36.6 200.0±0.0plus-or-minus200.00.0200.0± 0.0200.0 ± 0.0 product by company 522522522522 366.4±8.1plus-or-minus366.48.1366.4± 8.1366.4 ± 8.1 262.4±23.8plus-or-minus262.423.8262.4± 23.8262.4 ± 23.8 421.3±6.2plus-or-minus421.36.2421.3± 6.2421.3 ± 6.2 person plays instrument 513513513513 237.6±30.7plus-or-minus237.630.7237.6± 30.7237.6 ± 30.7 144.2±27.6plus-or-minus144.227.6144.2± 27.6144.2 ± 27.6 249.0±67.7plus-or-minus249.067.7249.0± 67.7249.0 ± 67.7 star constellation name 362362362362 276.3±3.4plus-or-minus276.33.4276.3± 3.4276.3 ± 3.4 210.3±16.7plus-or-minus210.316.7210.3± 16.7210.3 ± 16.7 200.0±0.0plus-or-minus200.00.0200.0± 0.0200.0 ± 0.0 plays pro sport 318318318318 244.2±10.8plus-or-minus244.210.8244.2± 10.8244.2 ± 10.8 195.1±13.1plus-or-minus195.113.1195.1± 13.1195.1 ± 13.1 294.8±1.2plus-or-minus294.81.2294.8± 1.2294.8 ± 1.2 company CEO 298298298298 90.2±7.2plus-or-minus90.27.290.2± 7.290.2 ± 7.2 15.4±6.1plus-or-minus15.46.115.4± 6.115.4 ± 6.1 155.5±4.7plus-or-minus155.54.7155.5± 4.7155.5 ± 4.7 superhero person 100100100100 49.8±2.0plus-or-minus49.82.049.8± 2.049.8 ± 2.0 36.6±2.4plus-or-minus36.62.436.6± 2.436.6 ± 2.4 71.9±2.3plus-or-minus71.92.371.9± 2.371.9 ± 2.3 superhero archnemesis 96969696 21.5±2.6plus-or-minus21.52.621.5± 2.621.5 ± 2.6 12.7±2.2plus-or-minus12.72.212.7± 2.212.7 ± 2.2 40.8±2.7plus-or-minus40.82.740.8± 2.740.8 ± 2.7 person university 91919191 41.0±1.6plus-or-minus41.01.641.0± 1.641.0 ± 1.6 35.4±1.9plus-or-minus35.41.935.4± 1.935.4 ± 1.9 44.8±2.2plus-or-minus44.82.244.8± 2.244.8 ± 2.2 pokemon evolution 44444444 28.7±3.9plus-or-minus28.73.928.7± 3.928.7 ± 3.9 30.3±1.1plus-or-minus30.31.130.3± 1.130.3 ± 1.1 35.7±0.5plus-or-minus35.70.535.7± 0.535.7 ± 0.5 country currency 30303030 19.4±0.9plus-or-minus19.40.919.4± 0.919.4 ± 0.9 20.2±0.8plus-or-minus20.20.820.2± 0.820.2 ± 0.8 22.0±0.0plus-or-minus22.00.022.0± 0.022.0 ± 0.0 food from country 30303030 16.4±1.1plus-or-minus16.41.116.4± 1.116.4 ± 1.1 11.2±1.4plus-or-minus11.21.411.2± 1.411.2 ± 1.4 18.8±1.4plus-or-minus18.81.418.8± 1.418.8 ± 1.4 city in country 27272727 18.0±0.9plus-or-minus18.00.918.0± 0.918.0 ± 0.9 18.0±1.1plus-or-minus18.01.118.0± 1.118.0 ± 1.1 18.1±0.7plus-or-minus18.10.718.1± 0.718.1 ± 0.7 country capital city 24242424 16.0±0.0plus-or-minus16.00.016.0± 0.016.0 ± 0.0 15.5±0.6plus-or-minus15.50.615.5± 0.615.5 ± 0.6 15.3±0.5plus-or-minus15.30.515.3± 0.515.3 ± 0.5 country language 24242424 15.4±0.6plus-or-minus15.40.615.4± 0.615.4 ± 0.6 14.9±0.6plus-or-minus14.90.614.9± 0.614.9 ± 0.6 15.8±0.4plus-or-minus15.80.415.8± 0.415.8 ± 0.4 country largest city 24242424 15.5±0.5plus-or-minus15.50.515.5± 0.515.5 ± 0.5 14.0±0.8plus-or-minus14.00.814.0± 0.814.0 ± 0.8 15.3±0.5plus-or-minus15.30.515.3± 0.515.3 ± 0.5 person lead singer of band 21212121 13.0±0.2plus-or-minus13.00.213.0± 0.213.0 ± 0.2 10.1±0.9plus-or-minus10.10.910.1± 0.910.1 ± 0.9 13.0±0.0plus-or-minus13.00.013.0± 0.013.0 ± 0.0 president birth year 19191919 11.0±0.0plus-or-minus11.00.011.0± 0.011.0 ± 0.0 — *⁣**** * president election year 19191919 9.5±0.5plus-or-minus9.50.59.5± 0.59.5 ± 0.5 9.9±0.6plus-or-minus9.90.69.9± 0.69.9 ± 0.6 *⁣**** * commonsense object superclass 76767676 52.7±1.5plus-or-minus52.71.552.7± 1.552.7 ± 1.5 51.5±2.2plus-or-minus51.52.251.5± 2.251.5 ± 2.2 54.4±1.7plus-or-minus54.41.754.4± 1.754.4 ± 1.7 word sentiment 60606060 47.2±4.1plus-or-minus47.24.147.2± 4.147.2 ± 4.1 42.8±5.1plus-or-minus42.85.142.8± 5.142.8 ± 5.1 50.5±2.0plus-or-minus50.52.050.5± 2.050.5 ± 2.0 task done by tool 52525252 30.6±1.5plus-or-minus30.61.530.6± 1.530.6 ± 1.5 25.8±1.9plus-or-minus25.81.925.8± 1.925.8 ± 1.9 33.7±1.8plus-or-minus33.71.833.7± 1.833.7 ± 1.8 substance phase of matter 50505050 30.4±6.8plus-or-minus30.46.830.4± 6.830.4 ± 6.8 34.2±3.3plus-or-minus34.23.334.2± 3.334.2 ± 3.3 40.5±1.9plus-or-minus40.51.940.5± 1.940.5 ± 1.9 work location 38383838 18.2±2.6plus-or-minus18.22.618.2± 2.618.2 ± 2.6 19.7±2.6plus-or-minus19.72.619.7± 2.619.7 ± 2.6 24.7±2.3plus-or-minus24.72.324.7± 2.324.7 ± 2.3 fruit inside color 36363636 10.2±1.0plus-or-minus10.21.010.2± 1.010.2 ± 1.0 9.0±0.0plus-or-minus9.00.09.0± 0.09.0 ± 0.0 17.3±1.6plus-or-minus17.31.617.3± 1.617.3 ± 1.6 task person type 32323232 18.0±1.2plus-or-minus18.01.218.0± 1.218.0 ± 1.2 16.1±1.5plus-or-minus16.11.516.1± 1.516.1 ± 1.5 19.2±2.1plus-or-minus19.22.119.2± 2.119.2 ± 2.1 fruit outside color 30303030 11.7±2.1plus-or-minus11.72.111.7± 2.111.7 ± 2.1 9.6±0.7plus-or-minus9.60.79.6± 0.79.6 ± 0.7 14.6±1.4plus-or-minus14.61.414.6± 1.414.6 ± 1.4 linguistic word first letter 241241241241 223.9±4.5plus-or-minus223.94.5223.9± 4.5223.9 ± 4.5 199.1±9.8plus-or-minus199.19.8199.1± 9.8199.1 ± 9.8 233.0±0.0plus-or-minus233.00.0233.0± 0.0233.0 ± 0.0 word last letter 241241241241 28.2±8.2plus-or-minus28.28.228.2± 8.228.2 ± 8.2 21.2±5.3plus-or-minus21.25.321.2± 5.321.2 ± 5.3 188.3±6.7plus-or-minus188.36.7188.3± 6.7188.3 ± 6.7 adjective antonym 100100100100 64.0±2.3plus-or-minus64.02.364.0± 2.364.0 ± 2.3 57.5±2.4plus-or-minus57.52.457.5± 2.457.5 ± 2.4 68.5±1.5plus-or-minus68.51.568.5± 1.568.5 ± 1.5 adjective superlative 80808080 70.5±0.9plus-or-minus70.50.970.5± 0.970.5 ± 0.9 64.4±3.3plus-or-minus64.43.364.4± 3.364.4 ± 3.3 70.5±0.7plus-or-minus70.50.770.5± 0.770.5 ± 0.7 verb past tense 76767676 61.0±4.6plus-or-minus61.04.661.0± 4.661.0 ± 4.6 54.0±3.0plus-or-minus54.03.054.0± 3.054.0 ± 3.0 65.8±3.9plus-or-minus65.83.965.8± 3.965.8 ± 3.9 adjective comparative 68686868 59.5±0.6plus-or-minus59.50.659.5± 0.659.5 ± 0.6 57.6±0.9plus-or-minus57.60.957.6± 0.957.6 ± 0.9 60.0±0.2plus-or-minus60.00.260.0± 0.260.0 ± 0.2 bias occupation age 45454545 25.7±2.6plus-or-minus25.72.625.7± 2.625.7 ± 2.6 22.9±3.8plus-or-minus22.93.822.9± 3.822.9 ± 3.8 32.8±2.6plus-or-minus32.82.632.8± 2.632.8 ± 2.6 univ degree gender 38383838 — 21.5±2.4plus-or-minus21.52.421.5± 2.421.5 ± 2.4 24.2±2.4plus-or-minus24.22.424.2± 2.424.2 ± 2.4 name birthplace 31313131 17.1±2.6plus-or-minus17.12.617.1± 2.617.1 ± 2.6 18.0±1.4plus-or-minus18.01.418.0± 1.418.0 ± 1.4 21.4±1.1plus-or-minus21.41.121.4± 1.121.4 ± 1.1 name religion 31313131 17.0±2.3plus-or-minus17.02.317.0± 2.317.0 ± 2.3 15.1±2.2plus-or-minus15.12.215.1± 2.215.1 ± 2.2 19.8±1.5plus-or-minus19.81.519.8± 1.519.8 ± 1.5 characteristic gender 30303030 15.9±2.7plus-or-minus15.92.715.9± 2.715.9 ± 2.7 15.8±2.2plus-or-minus15.82.215.8± 2.215.8 ± 2.2 19.7±1.2plus-or-minus19.71.219.7± 1.219.7 ± 1.2 name gender 19191919 11.0±0.0plus-or-minus11.00.011.0± 0.011.0 ± 0.0 10.7±0.6plus-or-minus10.70.610.7± 0.610.7 ± 0.6 10.8±0.4plus-or-minus10.80.410.8± 0.410.8 ± 0.4 occupation gender 19191919 9.6±0.8plus-or-minus9.60.89.6± 0.89.6 ± 0.8 9.8±0.7plus-or-minus9.80.79.8± 0.79.8 ± 0.7 10.8±0.4plus-or-minus10.80.410.8± 0.410.8 ± 0.4 Appendix B Assumptions underlying the LRE Approximation Our estimate of the LRE parameters is based on an assumption that the transformer LM implements relation decoding F⁢(,c)F(s,c)F ( s , c ) in a near-linear fashion that deviates from a linear model with a non-linear error term ε⁢() (s)ε ( s ) where both ε ε and ε′ are zero in expectation over ss, i.e., ⁢[ε⁢()]=0subscriptdelimited-[]0E_s[ (s)]=0blackboard_Es [ ε ( s ) ] = 0 and ⁢[ε′⁢()]=0subscriptdelimited-[]superscript′0E_s[ (s)]=0blackboard_Es [ ε′ ( s ) ] = 0. F⁢(,c) F(s,c)F ( s , c ) =b+W⁢+ε⁢()absent =b+Ws+ (s)= b + W s + ε ( s ) (9) Then passing to expectations over the distribution of ss we can estimate b and W: b b =F⁢(,c)−W⁢−ε⁢()absent =F(s,c)-Ws- (s)= F ( s , c ) - W s - ε ( s ) (10) b b =⁢[F⁢(,c)−W⁢]−⁢[ε⁢()]absentsubscriptdelimited-[]cancelsubscriptdelimited-[] =E_s[F(s,c)-Ws]- % E_s[ (s)]= blackboard_Es [ F ( s , c ) - W s ] - cancel blackboard_Es [ ε ( s ) ] (11) W W =F′⁢(,c)−ε′⁢()absentsuperscript′ =F (s,c)- (s)= F′ ( s , c ) - ε′ ( s ) (12) W W =⁢[F′⁢(,c)]−⁢[ε′⁢()]absentsubscriptdelimited-[]superscript′cancelsubscriptdelimited-[]superscript′ =E_s[F (s,c)]- % E_s[ (s)]= blackboard_Es [ F′ ( s , c ) ] - cancel blackboard_Es [ ε′ ( s ) ] (13) Equations 11 and 13 correspond to the bias term b and and projection term W of Equation 3 in the main paper. Appendix C Improving the estimate of F′⁢(,c)superscript′F (s,c)F′ ( s , c ) as β⁢Wβ Wβ W Empirically we have found that β⁢Wβ Wβ W (with β>11β>1β > 1) yields a more accurate linear model of F than W. In this section we measure the behavior of F′ in the region between subject representation vectors to provide some evidence on this. Take two subject representation vectors 1subscript1s_1s1 and 2subscript2s_2s2. The projection term of our LRE model W, based on the mean Jacobian of F calculated at subjects isubscripts_isitalic_i, yields this estimate of transformer’s behavior when traversing from one to the other F⁢(2)−F⁢(1)≈W⁢(2−1).subscript2subscript1subscript2subscript1 F(s_2)-F(s_1)≈ W(s_2-% s_1).F ( s2 ) - F ( s1 ) ≈ W ( s2 - s1 ) . (14) We can compare this to an exact calculation: the fundamental theorem of line integrals tells us that integrating the actual Jacobian along the path from 1subscript1s_1s1 to 2subscript2s_2s2 yields the actual change: F⁢(2)−F⁢(1)subscript2subscript1 F(s_2)-F(s_1)F ( s2 ) - F ( s1 ) =∫12F′⁢()⁢absentsuperscriptsubscriptsubscript1subscript2superscript′differential-d = _s_1^s_2F (s)d% s= ∫s start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPTs2 F′ ( s ) d s (15) ‖F⁢(2)−F⁢(1)‖normsubscript2subscript1 ||F(s_2)-F(s_1)||| | F ( s2 ) - F ( s1 ) | | =∫12T⁢F′⁢()⁢absentsuperscriptsubscriptsubscript1subscript2superscriptsuperscript′differential-d = _s_1^s_2u^TF (% s)ds= ∫s start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPTs2 uitalic_T F′ ( s ) d s (16) Here we reduce it to a one-dimensional problem, defining unit vectors ∝F⁢(2)−F⁢(1)proportional-tosubscript2subscript1u F(s_2)-F(s_1)u ∝ F ( s2 ) - F ( s1 ) and ∝2−1proportional-tosubscript2subscript1v _2-s_1v ∝ s2 - s1 in the row and column space respectively, so that ‖F⁢(2)−F⁢(1)‖normsubscript2subscript1 ||F(s_2)-F(s_1)||| | F ( s2 ) - F ( s1 ) | | ≈T⁢W⁢‖2−1‖absentsuperscriptnormsubscript2subscript1 ^TW\;v||s_2-s_1||≈ uitalic_T W v | | s2 - s1 | | (17) Table 5: Ratio between the right hand sides of Equation 16 and (17) for some of the relations. Relation Underestimation Ratio plays pro sport 2.517±1.043plus-or-minus2.5171.0432.517± 1.0432.517 ± 1.043 country capital city 4.198±0.954plus-or-minus4.1980.9544.198± 0.9544.198 ± 0.954 object superclass 3.058±0.457plus-or-minus3.0580.4573.058± 0.4573.058 ± 0.457 name birthplace 4.328±0.991plus-or-minus4.3280.9914.328± 0.9914.328 ± 0.991 By taking the ratio between the sides of (17) we can see how the actual rate of change from 1subscript1s_1s1 to 2subscript2s_2s2 is underestimated by W. Table 5 reports this value for some selected relations. In practice we find that setting β as a constant for an LM (instead of setting it per relation) is enough to attain good performance across a range of relations. Refer to Appendix E for further details. Appendix D Causality D.1 Derivation of Eqn 7 Under the same relation r = person plays instrument, consider two different (s,o)(s,o)( s , o ) pairs: (s = Miles Davis, o = trumpet) and (s′normal-′s s′ = Cat Stevens, o′normal-′o o′ = guitar). If an LRE defined by the projection term W and translation term b is an well-approximation of the model calculation F⁢(,c)F(s,c)F ( s , c ), then =β⁢W⁢()+b⁢ and ⁢′=β⁢W⁢(′)+b and superscript′ =β\,W(s)+b\;\;and\;\;o^% =β\,W(s )+bo = β W ( s ) + b and o′ = β W ( s′ ) + b Subtracting ′superscript′o o′ from oo ′−superscript′ -oo′ - o =β⁢W⁢′−β⁢W⁢absentsuperscript′ =β Ws -β Ws= β W s′ - β W s =β⁢W⁢(′−)⁢[since W is linear]absentsuperscript′[since W is linear] =β W(s -s)\;\;[since $W$ is % linear]= β W ( s′ - s ) [since W is linear] Δ⁢Δ Δ s =′−=1β⁢W−1⁢(′−)absentsuperscript′1superscript1superscript′ =s -s= 1β\,W^-1(o% -o)= s′ - s = divide start_ARG 1 end_ARG start_ARG β end_ARG W- 1 ( o′ - o ) (18) We observe that the edit direction Δ⁢Δ Δ s needs to be magnified to achieve good edit efficacy. In our experiments we magnify Δ⁢Δ Δ s by β (or set β=1.01.0β=1.0β = 1.0 in Equation 18). Δ⁢Δ Δ s =W−1⁢(′−)absentsuperscript1superscript′ =W^-1(o -o)= W- 1 ( o′ - o ) (19) D.2 Why low-rank inverse W†superscript†W W† instead of full inverse W−1superscript1W^-1W- 1 is necessary? In practice, we need to take a low rank approximation W†superscript†W W† instead of W−1superscript1W^-1W- 1 for the edit depicted in Figure 2 to be successful. With β set to 1.0, ′−superscript′ -oo′ - o =W⁢(′−)absentsuperscript′ =W(s -s)= W ( s′ - s ) Δ⁢Δ Δ o =W⁢Δ⁢absentΔ =W = W Δ s If we take U⁢Σ⁢VTΣsuperscriptU V^TU Σ Vitalic_T as the SVD of W, then Δ⁢Δ Δ o =U⁢Σ⁢VT⁢Δ⁢absentΣsuperscriptΔ =U V^T = U Σ Vitalic_T Δ s UT⁢Δ⁢superscriptΔ U^T _T Δ o =Σ⁢VT⁢Δ⁢absentΣsuperscriptΔ = V^T \;\;= Σ Vitalic_T Δ s Considering UT⁢Δ⁢superscriptΔU^T _T Δ o as usubscripto_uoitalic_u and VT⁢Δ⁢superscriptΔV^T _T Δ s as vsubscripts_vsitalic_v, u=Σ⁢v⁢or⁢Σ−1⁢u=vsubscriptΣsubscriptorsuperscriptΣ1subscriptsubscript _u= _v\;\;or\;\; ^-1% o_u=s_voitalic_u = Σ sitalic_v or Σ- 1 oitalic_u = sitalic_v (20) Here, Σ Σ maps vsubscripts_vsitalic_v to usubscripto_uoitalic_u. This Σ Σ is a diagonal matrix that contains the non-negative singular values in it’s diagonal and zero otherwise. The greater the singular value the more its effect on vsubscripts_vsitalic_v. However, if we take the full rank inverse of Σ Σ while mapping usubscripto_uoitalic_u to vsubscripts_vsitalic_v then the inverse becomes dominated by noisy smaller singular values and they wash out the contribution of meaningful singular values. Thus, it is necessary to consider only those singular values greater than a certain threshold τ or take a low rank inverse W†superscript†W W† instead of a full inverse W−1superscript1W^-1W- 1. The significance of different ranks on causality is depicted on Figure 9. We see that the causality increases with a rank up to some point and starts decreasing afterwards. This suggests that, we are ablating important singular values from Σ Σ before an optimal rank ρrsubscript _rρitalic_r is reached, and start introducing noisy singular values afterwards. Figure 9: Initially, faithfulness and causality of LRE improve with higher rank. However, after r⁢a⁢n⁢k=28superscript28rank=2^8r a n k = 28 causality starts declining whereas faithfulness remains stable. Appendix E Selecting Hyperparameters (β, ℓrsubscriptℓ _rℓitalic_r, and ρrsubscript _rρitalic_r) We need to select a scalar value β per LM since the slope W of the first order approximation underestimates the slope of F⁢(,c)F(s,c)F ( s , c ) (Appendix C). Additionally, we need to specify two hyperparameters per relation r; ℓrsubscriptℓ _rℓitalic_r, the layer after which ss is to be extracted and ρrsubscript _rρitalic_r, the rank of the low-rank inverse W†superscript†W W†. We perform a grid search to select these hyperparameters. For a specific β, hyperparameters ℓrsubscriptℓ _rℓitalic_r and ρrsubscript _rρitalic_r are selected to achieve the best causal influence as there is a strong agreement between faithfulness and causality when the hparams are selected this way. However, when ℓrsubscriptℓ _rℓitalic_r and ρrsubscript _rρitalic_r are selected to achieve best faithfulness there is a weaker agreement between our evaluation metrics (Figure 10). Section E.2 provides an insight on this. Figure 10: When hparams are selected to achieve best faithfulness there is a weaker correlation of 0.74 between our evaluation metrics unlike in Figure 6 where hparams were selected to achieve best causality. For both this figure and Figure 6 the LM is GPT-J and β=2.252.25β=2.25β = 2.25. In sweep over layers we notice that LRE performance per relation increases up to a certain layer and drops afterwards, suggesting a mode-switch in later layers (Figure 11). And, in the sweep over ranks we notice that both edit efficacy and faithfulness increases up to a certain rank. After that edit efficacy starts dropping while faithfulness remains stable. The reasoning behind how causality is affected by higher rank is discussed in Section D.2. E.1 Selecting β Table 6: Scores achieved by LRE on different values of β on GPT-J. β=2.252.25β=2.25β = 2.25 shows the best correlation between our evaluation metrics. Faithfulnessμ_μstart_FLOATSUBSCRIPT μ end_FLOATSUBSCRIPT means the average faithfulness across all the relations (same for Causalityμ_μstart_FLOATSUBSCRIPT μ end_FLOATSUBSCRIPT). β Faithfulnessμ_μstart_FLOATSUBSCRIPT μ end_FLOATSUBSCRIPT Causalityμ_μstart_FLOATSUBSCRIPT μ end_FLOATSUBSCRIPT Corr 0.00 0.17 ± 0.20 0.81 ± 0.22 0.30 0.25 0.21 ± 0.22 0.33 0.50 0.28 ± 0.24 0.38 0.75 0.36 ± 0.25 0.48 1.00 0.43 ± 0.26 0.58 1.25 0.50 ± 0.26 0.67 1.50 0.54 ± 0.25 0.76 1.75 0.57 ± 0.25 0.80 2.00 0.59 ± 0.25 0.83 2.25 0.59 ± 0.25 0.84 2.50 0.59 ± 0.25 0.84 2.75 0.59 ± 0.25 0.84 3.00 0.59 ± 0.25 0.83 3.25 0.58 ± 0.25 0.81 3.50 0.57 ± 0.25 0.80 3.75 0.56 ± 0.25 0.78 4.00 0.55 ± 0.25 0.76 4.25 0.54 ± 0.25 0.75 4.50 0.53 ± 0.25 0.74 4.75 0.53 ± 0.25 0.72 5.00 0.52 ± 0.25 0.71 Table 6 represents how performance scores of LRE change with respect to different values of β for GPT-J. In our experiments, causality is always calculated with β set to 1.01.01.01.0. So, the average causality score remain constant. β is selected per LM to achieve the best agreement between our performance metrics faithfulness and causality. For GPT-J optimal value of β is 2.252.252.252.25. E.2 Layer-wise LRE performance on selected relations (GPT-J) If LRE remains faithful to model up to layer ℓf⁢a⁢i⁢t⁢hsubscriptℓℎ _faithℓitalic_f a i t h it is reasonable to expect that LRE will retain high causality until ℓf⁢a⁢i⁢t⁢hsubscriptℓℎ _faithℓitalic_f a i t h as well. However, an examination of the faithfulness and causality performances across layers reveals that the causality scores drop before ℓf⁢a⁢i⁢t⁢hsubscriptℓℎ _faithℓitalic_f a i t h (Fig. 11). In fact Fig. 14 from Section G.2 indicates that all the intervention baselines exhibit a decrease in performance at deeper layers, particularly our method and the oracle method at similar layers. This might be a phenomenon associated with this type of intervention in general, rather than a fault with our approximation of the target ′superscript′s s′. Notice that, in all of our activation patching experiments we only patch a single state (at the position of the last subject token after a layer ℓ ℓ). If the activation after layer ℓ ℓ is patched, layers till ℓ−1ℓ1 -1ℓ - 1 retains information about the original subject ss and they can leak information about ss to later layers because of attention mechanism. The deeper we intervene the more is this leakage from previous states and it might reduce the efficacy of these single state activation patching approaches. It is not reasonable to expect high causality after ℓf⁢a⁢i⁢t⁢hsubscriptℓℎ _faithℓitalic_f a i t h, and causality can drop well before ℓf⁢a⁢i⁢t⁢hsubscriptℓℎ _faithℓitalic_f a i t h because of this leakage. This also partly explains the disagreement between the two metrics when the hyperparameters are chosen to achieve the best faithfulness (Figure 10). Figure 11: LRE performance for selected relations in different layers of GPT-J. The last row features some of the relations where LRE could not achieve satisfactory performance indicating a non-linear decoding process for them. Appendix F Varying n and prompt template Figure 12 shows how l⁢r⁢elrel r e performance changes based on number of examples n used for approximation. For most of the relations both faithfulness and efficacy scores start plateauing after n=55n=5n = 5. In our experiment setup we use n=88n=8n = 8 as that is the largest number we could fit for a GPT-J model on a single A6000. However, Figure 12 suggests that a good LRE estimation may be obtained with less number of examples. Figure 12: LRE performance across different n. We also test how LRE performance on a relation r changes when the same relation r is contextualized with different prompt templates. Table 7 shows minimal change in faithfulness and causality scores when LRE is calculated with different prompt templates. Table 7: LRE performance on different prompt templates. The subject s is inserted in place of . Performance scores presented as mean and standard deviation across 24 trials with different sets of training examples. Relation Prompt Template Faithfulness Causality country capital city The capital of is 0.84±0.09plus-or-minus0.840.090.84± 0.090.84 ± 0.09 0.94±0.04plus-or-minus0.940.040.94± 0.040.94 ± 0.04 The capital of is the city of 0.87±0.08plus-or-minus0.870.080.87± 0.080.87 ± 0.08 0.94±0.04plus-or-minus0.940.040.94± 0.040.94 ± 0.04 The capital city of is 0.84±0.08plus-or-minus0.840.080.84± 0.080.84 ± 0.08 0.94±0.04plus-or-minus0.940.040.94± 0.040.94 ± 0.04 What is the capital of ? It is the city of 0.87±0.07plus-or-minus0.870.070.87± 0.070.87 ± 0.07 0.92±0.05plus-or-minus0.920.050.92± 0.050.92 ± 0.05 plays pro sport plays the sport of 0.78±0.07plus-or-minus0.780.070.78± 0.070.78 ± 0.07 0.90±0.03plus-or-minus0.900.030.90± 0.030.90 ± 0.03 plays professionally in the sport of 0.78±0.09plus-or-minus0.780.090.78± 0.090.78 ± 0.09 0.90±0.03plus-or-minus0.900.030.90± 0.030.90 ± 0.03 What sport does play? They play 0.81±0.06plus-or-minus0.810.060.81± 0.060.81 ± 0.06 0.90±0.03plus-or-minus0.900.030.90± 0.030.90 ± 0.03 person occupation works professionally as a 0.41±0.08plus-or-minus0.410.080.41± 0.080.41 ± 0.08 0.55±0.09plus-or-minus0.550.090.55± 0.090.55 ± 0.09 works as a 0.44±0.11plus-or-minus0.440.110.44± 0.110.44 ± 0.11 0.58±0.07plus-or-minus0.580.070.58± 0.070.58 ± 0.07 By profession, is a 0.46±0.14plus-or-minus0.460.140.46± 0.140.46 ± 0.14 0.58±0.08plus-or-minus0.580.080.58± 0.080.58 ± 0.08 adjective superlative The superlative form of is 0.93±0.02plus-or-minus0.930.020.93± 0.020.93 ± 0.02 0.97±0.02plus-or-minus0.970.020.97± 0.020.97 ± 0.02 What is the superlative form of ? It is 0.92±0.02plus-or-minus0.920.020.92± 0.020.92 ± 0.02 0.96±0.03plus-or-minus0.960.030.96± 0.030.96 ± 0.03 Appendix G Baselines G.1 Faithfulness In Figure 13 we examine whether our method can be applied to ss extracted from zero-shot prompts that contain only the subject s and no further context. It appears that even when LRE is trained with few-shot examples, it can achieve similar results when applied to ss that is free of any context specifying the relation. Figure 13: LRE faithfulness on GPT-J compared with different linear functions baselines(same as Figure 4). Each of the functions were approximated with n=88n=8n = 8 samples, each prepended with n−11n-1n - 1 few-shot examples (Table 2). Dark blue bars indicate faithfulness when evaluated on ss extracted in a similar setup. Light blue bars represent how LRE (trained on few-shot examples) generalize when applied on ss extracted from zero-shot prompts that contain only the subject and no further context. G.2 Causality Figure 14 shows the LRE causality performance in comparison to other baselines for selected relations and across different layers. If LRE is a good approximation of the model computation, then our causality intervention should be equivalent to the oracle baseline, which replaces ss with ′superscript′s s′. The graphs demonstrate the similarity between our method performance and oracle performance across all layers. This provides causal evidence that LRE can reliably recover ′superscript′s s′ for these relations. While model editing is not the primary focus of this work, it is still worth examining how our intervention may affect multiple token generation, since it may reveal unexpected side effects of the method. Based on a qualitative analysis of the post-edit generations, it appears that the edits preserve the fluency of the model. Table 8 presents examples of generated texts after intervention. Figure 14: LRE causality across different layers of GPT-J. The causality curve closely matches the peaks and valleys of the oracle baseline, replacing ss with ′superscript′s s′, suggesting that LRE is a good approximation of the model computation F. Table 8: Generated texts, before and after our causal intervention on GPT-J to change its prediction to o′. Prompt o→o′→superscript′o→ o o → o′ Before After Miles Davis plays the trumpet → guitar trumpet in his band. guitar live with his band. Siri was created by Apple → Google Apple as a personal assistant. Google and it has become a huge success within Google Maps. Chris Martin is the lead singer of Coldplay → Foo Fighters Coldplay and a man of many talents. Foo Fighters, one of the most successful and popular rock bands in the world. What is the past tense of close? It is closed → walked closed. It means it gets closed or is closed. walked. What is the past tense of read? Appendix H LRE on GPT2-xl and LlaMa-13B We provide further results for GPT2-xl and LLaMA-13b to show that autoregressive LMs of different sizes employ this linear encoding scheme for a range of different relations. Figure 15 present LRE performances for each of the three models grouped by relation category. Figure 15: LRE performance in different relation categories on different LMs. Figures 15(a) and 15(a) illustrate the high correlation between our two metrics on GPT2-xl and LLaMA-13b respectively. These findings are consistent with the results reported for GPT-J (Figure 6). (a) GPT2-xl (β=2.25,c⁢o⁢r⁢r=0.85formulae-sequence2.250.85β=2.25,corr=0.85β = 2.25 , c o r r = 0.85) (b) LLaMA-13B, (β=8.0,c⁢o⁢r⁢r=0.83formulae-sequence8.00.83β=8.0,corr=0.83β = 8.0 , c o r r = 0.83) Figure 16: High correlation between faithfulness and causality in both GPT2-xl (R=0.85) and LLaMa-13B (R=0.83). Each of the dots represent LRE performance for a relation. Bold dots indicate relations for which LRE is evaluated on ≥30absent30≥ 30≥ 30 test examples. Similarly, Figures 16(a) and 16(b), compare the faithfulness of our method with other approaches of achieving a linear decoding scheme for GPT2-xl and LlaMA. (a) GPT2-xl (b) LLaMA-13B Figure 17: LRE faithfulness on GPT2-xl and LLaMA-13B compared with different baselines. Refer to Figure Figure 4 for details on these baseline approaches of achieving a linear decoding function. Lastly, Figure 18 depicts LRE faithfulness in GPT2-xl and LLaMA-13b for each of relations in our dataset. According to Spearman’s rank-order correlation, GPT-J’s relation-wise performance is strongly correlated with both GPT2-xl (R=0.850.85R=0.85R = 0.85) and LLaMa-13B (R=0.710.71R=0.71R = 0.71), whereas GPT2-xl and LLaMa-13B are moderately corelated (R=0.580.58R=0.58R = 0.58). Figure 18: Relation-wise LRE faithfulness to the LM relation decoding function F. Horizontal red lines per relation indicate accuracy of a random-guess baseline. Relations are ordered according to their LRE faithfulness in GPT-J. We do not calculate LRE estimation for GPT2-xl on the relation president birth year as GPT2-xl can not accurately answer for that relation (Table 4). We also do not calculate LRE of LLaMa-13B for the relations where the o is a year (i.e. president birth year and president election year) as tokenizer of LLaMA models splits the digits of a year. Such behavior of LLaMA tokenizer makes the relation decoding function trivial, since most of the answers start with "1" for these relations. Those cases were included in this plot (grayed out) to align LRE faithfulness of different LMs by relations. Appendix I Limitations Our analysis of linear relational embeddings has several core limitations. Dataset size We have only tested a small set of 47 relations; although we have covered several categories of relations, there are many types of relations we have not explored, such as numerical or physical relations, or logical inferences or multi-hop reasoning. First-token correctness criterion During all of our experiments, we consider a predicted object correct if the first predicted token matches the first token of the true object. This introduces a risk for false positives: if multiple different objects share a first token, we might erroneously label a prediction as correct. For example, in the relation person university, many of the university names start with "University of" and such cases would inflate our measurements. We quantify the risk for such false positives in Table 9 and find that many relations have few to no collisions. Nevertheless, the first-token evaluation scheme limits the relations that can be reliably evaluated for linear decoding using this approach. Table 9: Relation-wise r⁢a⁢n⁢g⁢eranger a n g e (count of unique o) in our dataset along with the percentage of o s uniquely identified by their first token when tokenized with different LM tokenizers. The number inside parenthesis is the actual count of unique first tokens. *^*start_FLOATSUPERSCRIPT * end_FLOATSUPERSCRIPT Tokenizer for LLaMA models splits numeric words by digits. For the relation president birth year it only finds one unique first token 11\1\ 1 , since all the years in this relation start with 1. For president election year LLaMA tokenizer finds only two unique first tokens 1,212\1,2\ 1 , 2 . Counts for these relations were ignored while calculating average first token coverage for LLaMA-13B, Relation |r⁢a⁢n⁢g⁢er|subscript|range_r|| r a n g eitalic_r | GPT-J GPT2-xl LLaMA-13B adjective antonym 95959595 100.0%⁢(95)percent100.095100.0\%\;(95)100.0 % ( 95 ) 100.0%⁢(95)percent100.095100.0\%\;(95)100.0 % ( 95 ) 98.9%⁢(94)percent98.99498.9\%\;(94)98.9 % ( 94 ) adjective comparative 57575757 100.0%⁢(57)percent100.057100.0\%\;(57)100.0 % ( 57 ) 100.0%⁢(57)percent100.057100.0\%\;(57)100.0 % ( 57 ) 93.0%⁢(53)percent93.05393.0\%\;(53)93.0 % ( 53 ) adjective superlative 79797979 97.5%⁢(77)percent97.57797.5\%\;(77)97.5 % ( 77 ) 97.5%⁢(77)percent97.57797.5\%\;(77)97.5 % ( 77 ) 98.7%⁢(78)percent98.77898.7\%\;(78)98.7 % ( 78 ) city in country 21212121 95.2%⁢(20)percent95.22095.2\%\;(20)95.2 % ( 20 ) 95.2%⁢(20)percent95.22095.2\%\;(20)95.2 % ( 20 ) 95.2%⁢(20)percent95.22095.2\%\;(20)95.2 % ( 20 ) company CEO 287287287287 72.5%⁢(208)percent72.520872.5\%\;(208)72.5 % ( 208 ) 72.5%⁢(208)percent72.520872.5\%\;(208)72.5 % ( 208 ) 67.6%⁢(194)percent67.619467.6\%\;(194)67.6 % ( 194 ) company hq 163163163163 100.0%⁢(163)percent100.0163100.0\%\;(163)100.0 % ( 163 ) 100.0%⁢(163)percent100.0163100.0\%\;(163)100.0 % ( 163 ) 93.3%⁢(152)percent93.315293.3\%\;(152)93.3 % ( 152 ) country currency 23232323 100.0%⁢(23)percent100.023100.0\%\;(23)100.0 % ( 23 ) 100.0%⁢(23)percent100.023100.0\%\;(23)100.0 % ( 23 ) 91.3%⁢(21)percent91.32191.3\%\;(21)91.3 % ( 21 ) landmark in country 91919191 100.0%⁢(91)percent100.091100.0\%\;(91)100.0 % ( 91 ) 100.0%⁢(91)percent100.091100.0\%\;(91)100.0 % ( 91 ) 97.8%⁢(89)percent97.88997.8\%\;(89)97.8 % ( 89 ) person father 968968968968 41.3%⁢(400)percent41.340041.3\%\;(400)41.3 % ( 400 ) 41.3%⁢(400)percent41.340041.3\%\;(400)41.3 % ( 400 ) 38.9%⁢(377)percent38.937738.9\%\;(377)38.9 % ( 377 ) person lead singer of band 21212121 85.7%⁢(18)percent85.71885.7\%\;(18)85.7 % ( 18 ) 85.7%⁢(18)percent85.71885.7\%\;(18)85.7 % ( 18 ) 85.7%⁢(18)percent85.71885.7\%\;(18)85.7 % ( 18 ) person mother 962962962962 39.5%⁢(380)percent39.538039.5\%\;(380)39.5 % ( 380 ) 39.5%⁢(380)percent39.538039.5\%\;(380)39.5 % ( 380 ) 31.9%⁢(307)percent31.930731.9\%\;(307)31.9 % ( 307 ) person occupation 31313131 100.0%⁢(31)percent100.031100.0\%\;(31)100.0 % ( 31 ) 100.0%⁢(31)percent100.031100.0\%\;(31)100.0 % ( 31 ) 93.5%⁢(29)percent93.52993.5\%\;(29)93.5 % ( 29 ) person university 69696969 53.6%⁢(37)percent53.63753.6\%\;(37)53.6 % ( 37 ) 53.6%⁢(37)percent53.63753.6\%\;(37)53.6 % ( 37 ) 50.7%⁢(35)percent50.73550.7\%\;(35)50.7 % ( 35 ) pokemon evolution 44444444 90.9%⁢(40)percent90.94090.9\%\;(40)90.9 % ( 40 ) 90.9%⁢(40)percent90.94090.9\%\;(40)90.9 % ( 40 ) 81.8%⁢(36)percent81.83681.8\%\;(36)81.8 % ( 36 ) president birth year 15151515 60.0%⁢(9)percent60.0960.0\%\;(9)60.0 % ( 9 ) 60.0%⁢(9)percent60.0960.0\%\;(9)60.0 % ( 9 ) 6.7%⁢(1)*percent6.7superscript16.7\%\;(1)^*6.7 % ( 1 )* president election year 18181818 77.8%⁢(14)percent77.81477.8\%\;(14)77.8 % ( 14 ) 77.8%⁢(14)percent77.81477.8\%\;(14)77.8 % ( 14 ) 11.1%⁢(2)*percent11.1superscript211.1\%\;(2)^*11.1 % ( 2 )* product by company 30303030 100.0%⁢(30)percent100.030100.0\%\;(30)100.0 % ( 30 ) 100.0%⁢(30)percent100.030100.0\%\;(30)100.0 % ( 30 ) 86.7%⁢(26)percent86.72686.7\%\;(26)86.7 % ( 26 ) star constellation name 31313131 93.5%⁢(29)percent93.52993.5\%\;(29)93.5 % ( 29 ) 93.5%⁢(29)percent93.52993.5\%\;(29)93.5 % ( 29 ) 87.1%⁢(27)percent87.12787.1\%\;(27)87.1 % ( 27 ) superhero archnemesis 90909090 84.4%⁢(76)percent84.47684.4\%\;(76)84.4 % ( 76 ) 84.4%⁢(76)percent84.47684.4\%\;(76)84.4 % ( 76 ) 81.1%⁢(73)percent81.17381.1\%\;(73)81.1 % ( 73 ) superhero person 100100100100 89.0%⁢(89)percent89.08989.0\%\;(89)89.0 % ( 89 ) 89.0%⁢(89)percent89.08989.0\%\;(89)89.0 % ( 89 ) 84.0%⁢(84)percent84.08484.0\%\;(84)84.0 % ( 84 ) task done by tool 51515151 98.0%⁢(50)percent98.05098.0\%\;(50)98.0 % ( 50 ) 98.0%⁢(50)percent98.05098.0\%\;(50)98.0 % ( 50 ) 90.2%⁢(46)percent90.24690.2\%\;(46)90.2 % ( 46 ) Relation where all o is uniquely identified by the first token (×26absent26× 26× 26) — 100%percent100100\%100 % 100%percent100100\%100 % 100%percent100100\%100 % Average — 93.17%percent93.1793.17\%93.17 % 93.17%percent93.1793.17\%93.17 % 92.17%percent92.1792.17\%92.17 % Single object assumption In some relations, there may be more than one correct answer (e.g., fruits often take many different outside colors). Our dataset only catalogs one canonical related object in each case. This does not impact our results because for each subject we focus on the single object that the LM predicts, but future work could extend our evaluation scheme to measure how well LREs estimate the LM’s distribution of candidate objects. Table 10: Relation-wise LRE performance on GPT-J, and respective hyperparameters. relation, r |r⁢a⁢n⁢g⁢er|subscript|range_r|| r a n g eitalic_r | ℓrsubscriptℓ _rℓitalic_r β ρrsubscript _rρitalic_r Faithfulness Causality adjective antonym 95 8 2.25 243 0.69 ± 0.07 0.86 ± 0.04 adjective comparative 57 10 121 0.98 ± 0.01 0.94 ± 0.04 adjective superlative 79 10 143 0.93 ± 0.02 0.99 ± 0.01 characteristic gender 2 1 74 0.77 ± 0.11 0.97 ± 0.04 city in country 21 2 115 0.44 ± 0.10 0.89 ± 0.09 company CEO 287 6 173 0.06 ± 0.03 0.31 ± 0.05 company hq 163 6 126 0.21 ± 0.06 0.49 ± 0.04 country capital city 24 3 68 0.88 ± 0.07 0.99 ± 0.02 country currency 23 3 88 0.58 ± 0.08 0.98 ± 0.03 country language 14 1 63 0.88 ± 0.09 0.99 ± 0.03 country largest city 24 10 74 0.92 ± 0.05 0.99 ± 0.02 food from country 26 3 113 0.51 ± 0.12 0.97 ± 0.05 fruit inside color 6 7 107 0.65 ± 0.15 0.93 ± 0.07 fruit outside color 9 5 160 0.78 ± 0.15 0.83 ± 0.12 landmark in country 91 6 97 0.36 ± 0.06 0.68 ± 0.02 landmark on continent 5 4 158 0.56 ± 0.13 0.91 ± 0.02 name birthplace 8 7 91 0.92 ± 0.05 0.96 ± 0.07 name gender 2 emb 17 0.80 ± 0.16 0.94 ± 0.04 name religion 5 4 57 0.80 ± 0.10 0.99 ± 0.02 object superclass 10 7 91 0.85 ± 0.05 0.93 ± 0.03 occupation age 2 5 34 0.68 ± 0.03 1.00 ± 0.00 occupation gender 2 4 34 0.98 ± 0.04 1.00 ± 0.00 person father 968 8 217 0.07 ± 0.03 0.28 ± 0.04 person lead singer of band 21 8 163 0.64 ± 0.09 0.84 ± 0.09 person mother 962 6 170 0.14 ± 0.04 0.39 ± 0.05 person native language 30 6 92 0.65 ± 0.16 0.87 ± 0.03 person occupation 31 8 131 0.49 ± 0.08 0.66 ± 0.06 person plays instrument 6 9 198 0.59 ± 0.10 0.76 ± 0.04 person sport position 14 5 97 0.42 ± 0.15 0.74 ± 0.05 person university 69 4 153 0.64 ± 0.11 0.91 ± 0.04 plays pro sport 5 6 117 0.76 ± 0.06 0.94 ± 0.01 pokemon evolution 44 7 206 0.15 ± 0.05 0.25 ± 0.08 president birth year 15 6 106 0.54 ± 0.14 0.84 ± 0.08 president election year 18 emb 84 0.52 ± 0.20 0.91 ± 0.09 product by company 30 4 158 0.31 ± 0.14 0.54 ± 0.05 star constellation name 31 8 152 0.41 ± 0.08 0.27 ± 0.04 substance phase of matter 3 7 60 0.87 ± 0.09 0.97 ± 0.03 superhero archnemesis 90 11 192 0.30 ± 0.08 0.60 ± 0.10 superhero person 100 8 228 0.44 ± 0.06 0.71 ± 0.07 task done by tool 51 5 145 0.29 ± 0.10 0.76 ± 0.07 task person type 32 8 109 0.49 ± 0.10 0.77 ± 0.10 verb past tense 76 11 182 0.95 ± 0.03 0.97 ± 0.02 word first letter 25 7 121 0.58 ± 0.09 0.91 ± 0.02 word last letter 18 6 61 0.57 ± 0.13 0.83 ± 0.10 word sentiment 3 4 94 0.63 ± 0.16 0.93 ± 0.03 work location 24 5 112 0.55 ± 0.09 0.94 ± 0.06

Instant research discovery