Linearity of Relation Decoding in Transformer Language Models

Thanks for the review and your many helpful comments! We are glad to hear that you found our work insightful and interesting, and we think we have addressed your concerns in our revision.

---

Limitations are not discussed in enough detail.

Thanks for the suggestion. We have added a Limitations discussion in the appendix (Appendix I) that discusses dataset size, the risk for false positives when looking only at first tokens, and the assumption that every subject maps to a single object.

---

For example, given a fixed subject and relation, it is possible that there are multiple true objects. This has consequences both for the generation of the datasets (where examples are filtered out if they're not generated by the model) and questions to what extend an invertible function can be a reasonable approximation of the true function, since invertibility implies that the function/relation is injective. Discussing these aspects would improve the paper.

You are correct that some of the relations we investigate can have one-to-many mapping where our dataset only records one object for each example. We have discussed this issue in Limitations (Appendix I).

Our measurements are based on the following observation: the role of the LRE is not to recover the idealized relation, but the LM’s predictions for that relation. An LM presented with any specific sentence will assign the highest probability on one specific token, so the approach we have taken asks: can an LRE recover that same token? By its nature, the top-token prediction is well-defined (each subject maps to one object).

---

Table 4 (Appendix A) shows that GPT-J gets zero examples from the "occupation gender" category right.

Good catch! It turns out there was a bug in the generation of Table 4: instead of collating our full experiment setup on $n = 8$ examples across 24 trials, we generated it based on $n = 5$ examples on a single trial, which erroneously resulted in some tabulated zeros. We have fixed the script and updated Table 4, and now there are no cells with zeros.

---

For example, the LLaMA-13B achieves zero percent on the task of predicting president's birth year and election year.

We appreciate the thorough error analysis! Indeed, LLaMA’s tokenizer separates each digit in the years, Since we only look at the first token of both the true object and the predicted object to determine correctness, it is trivial in the opposite direction: the first token of the true object is always "1" for president birth year, so LLaMA need only predict "1" (which it does). This is also true for most of the examples in president election year. The reason LLaMA was scoring 0 in Figure 18 is because we weren’t accounting for spacing in front of the "1".

Note that for these cases specifically, we had already excluded them from the faithfulness and causality results in the original draft. We realize this was very unclear in the text, so we have added extra wording to Tables 4, 9 and Figure 18 to explain why LLaMA was not evaluated on relations where the object is a year. Additionally, in our new limitations section we discuss this as a pitfall of the first-token evaluation procedure.

Thank you for the positive review! We are pleased to know that you found our work novel and insightful.

---

In experiments, n = 8 examples are used to estimate W and b. Would it be difficult to use more examples? They seem to be too few to reliably estimate W and b. I would also be interested in the variance as well as the mean.

We find that for most of the relations LRE performance saturates after $n = 5$. We have included a Figure (Figure 12) in the appendix which shows our measurements up to $n = 11$, and the plateau after 5. The reason that our LRE estimation procedure is so sample efficient is because it is estimated as a first-order approximation of the LM computation itself. We don’t train a linear classifier (or affine transformation) to match a set of input-output pairs. Our choice of $n=8$ for large-scale experiments is largely a practical one, determined by the memory consumption of GPT-J on the A6000 GPU available for our research.

We agree the variance is also important to visualize: Table 10 (previously Table 8) in our appendix shows the standard deviation in LRE performance across 24 trials, each time calculated on different sets of training examples.

---

Does a “sample” mean a single example? In statistics, a sample usually means a collection of examples (data points).

Yes! In our draft by sample we mean a single example or a datapoint. To avoid confusion, we have updated the paper to call these “examples”.

---

various typos

Thanks for catching these typos. We have addressed all of them in the revision.

Thank you for your constructive review. We are pleased to know that you found the question we are investigating to be “interesting” and our findings to be “useful”!

---

It is unknown that how the context prompt (e.g., [s] plays the) affects the conclusion. For example, will the conclusion hold if we change some other contexts that express same meaning of relation?

During early versions of our experiments, we averaged results over multiple different prompts to account for potential difference in LRE performance. We later decided to use only one prompt per relation due to time constraints and because we noticed similar LRE performance across different prompts. We realize this confound is important to account for, so we’ve added Table 7 to the appendix which shows how faithfulness and causality scores vary when LRE is estimated with different prompt templates. Ultimately, we find that most scores do not change dramatically.

---

In term of handling objects that have multiple tokens, I saw in Table 4: only the first token was used to determine correctness. Do we apply this strategy to all experiments? Will this lead to false positives? For causality experiment, do we also use the first token of o’ for experiment?

We always use the first token for the objects that get tokenized to multiple tokens. This is true for the causality experiments as well. For some specific relations this strategy can indeed lead to false positives, though we find in practice that such collisions are rare.

To clarify this, we have made two adjustments to our submission. First, we now discuss this issue explicitly in the limitations. Second, we have included a new table (Table 9) in the appendix which shows that first token collisions are rare: on average, over $93$% of the objects per relation in our dataset can be uniquely identified with their first token. It is worth noting that this number is skewed downwards by relations such as person mother, person father, and president birth year. And, the LRE performs poorly on these relations anyway.

---

Do we have some examples showing that when LRE can not fully capture the LM's computation of the relation, the linear approximation can still perform a successful edit?

Here is one illustrative example where the LRE is not faithful yet we have observed a successful causal edit. Given the prompt "Malcolm MacDonald works as a", GPT-J correctly predicts the expected answer, " politician". However, an LRE (estimated based on the following examples) will assign the highest probability to " writer".

[
    (subject='Janet Gunn', object='actor'),
    (subject='Tony Pua', object='politician'),
    (subject='Valentino Bucchi', object='composer'),
    (subject='Christopher Guest', object='comedian'),
    (subject='Giovanni Battista Guarini', object='poet'),
    (subject='Menachem Mendel Schneerson', object='rabbi'),
    (subject='Cristina Peri Rossi', object='novelist'),
    (subject='Nils Strindberg', object='photographer')
]

Although the correct answer " politician" is within the top-5 prediction of the LRE, our faithfulness metric stringently measures top-1 accuracy and will consider it a failure case.

In contrast, when we compute the edit direction as $\Delta \mathbf{s} = W_r^{-1}(\mathbf{o}' - \mathbf{o})$, we find that $\mathbf{s} + \Delta \mathbf{s}$ is often enough to change the LM output from $o$ to $o'$ even when LRE($\mathbf{s'}$) does not assign highest probability to $o’$.

Continuing this example: LRE failed to predict the correct occupation of $s’$ = "Malcolm MacDonald". However, we can use GPT-Js object representation $\mathbf{o'}$ corresponding to its encoding of Malcom Macdonald’s profession to successfully change the profession of $s$ = "Walter Hines Page", from $o$ = " journalist" to $o’$ = " politician", which matches the correct profession of "Malcom Macdonald", even though LRE($\mathbf{s’}$) incorrectly predicted " writer".

When causality exceeds faithfulness, it reflects relations for which such examples are prevalent.

---