Paper: Linearity of Relation Decoding in Transformer Language Models (ICLR-2024)

Reviewer comment (Theory):
Limitations are not discussed in enough detail. 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.

Author response:
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.

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).

Reviewer comment (Evaluation):
Table 4 (Appendix A) shows the number of correctly predicted relations per language model. It is quite remarkable that some models in some categories get zero percent of the facts right, whereas other models get a much larger percentage right. Zero percent hints at systematic errors that remain undiscussed in the paper. For example, the LLaMA-13B achieves zero percent on the task of predicting president's birth year and election year. This is quite curious, given that these are quite well-known facts that appear a lot in the pretraining data, and the other models get them right. There is a reason for the poor performance of LLaMA on this particular task: Their tokenizer represents every digit as its own token. However, the caption of Table 4 states that only the first token was used to determine correctness. Since the LLaMA model would require 4 tokens, it had no chance of predicting the "correct" token. This is a shortcoming of the proposed measurement, and accepting this bias should at least be discussed. Similar reasons might cause zero percent accuracy on some of the other tasks.

Author response:
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.

Reviewer comment (Evaluation):
Table 4 (Appendix A) shows that GPT-J gets zero examples from the "occupation gender" category right. Hence, none of these examples from this category should remain in the dataset for GPT-J according to the description in Section 4 (Dataset). Curiously, Figure 3 shows that "occupation gender" has a near 100% faithfulness success rate. This raises the question of how that is possible if there are no such examples.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.

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

Author response:
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”.

Reviewer comment (Writing):
p. 4: of LM’s decoding -> of the LM’s decoding?

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

Reviewer comment (Writing):
p. 5: by (Merullo et al., 2023) -> by Merullo et al. (2023)?

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

Reviewer comment (Writing):
p. 7: is a higher -> is higher?

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

Reviewer comment (Writing):
p. 8: visualizes a next -> visualizes next?

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

Reviewer comment (Evaluation):
In terms 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 the causality experiment, do we also use the first token of o’ for the experiment?

Author response:
> 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.

Reviewer comment (Theory):
In which level of faithfulness can we say that the linearity holds for the relation?

Author response:
> In which level of faithfulness, we can say that the linearity holds for the relation?

Our LRE faithfulness measure is intended to be a fine-grained measure of linearity that is more informative than a binary assessment of “linear” vs “nonlinear”. In the paper, we say that relations for which the LRE attains at least 60% faithfulness are linearly decoded (meaning 22 out of 47 in our dataset are), but this is intended to help provide a summary of findings more than it is to suggest a hard cutoff. The answer is ultimately relation dependent, relying on the size of the range and how well the LM encodes that relation at all.

That said, while it’s hard to provide an exact cutoff for when a relation is linearly decoded, we observe multiple cases where the relation is **unequivocally not** linearly decoded, including *person mother*, *person father*, *company CEO*, etc. LRE performance is close to random guessing on these relations, so we conclude that the model must employ some non-linear decoding mechanism instead.