Paper: How do Language Models Bind Entities in Context? (ICLR-2024)

Reviewer comment (Theory):
The paper addresses a simple binding problem between a pair of entities. It is unclear how this can be generalized to the binding behavior on more than binary relations and entities involved in multiple binding.

Author response:
> It is unclear how this can be generalized to the binding behavior on more than binary relations and entities involved in multiple binding. 

We have conducted new experiments that test binding with three terms. Please see our [response to reviewer H1Kh](https://openreview.net/forum?id=zb3b6oKO77&noteId=BYsiAMpu7m).

Reviewer comment (Reproducibility):
When the authors substitute activations in Figure 3 or calculate median-calibrated accuracy in Figure 5 and Table 2, which part of the model is substituted and which is frozen? LLMs are autoregressive models, and during decoding, when and how are the parameters substituted?

Author response:
> When the authors substitute activations in Figure 3 or calculate median-calibrated accuracy in Figure 5 and Table 2, which part of the model is substituted and which is frozen? LLMs are autoregressive models, and during decoding, when and how are the parameters substituted?

The internal activations corresponding to the context (i.e. the part of the prompt that contains the binding information) is frozen (denoted as $Z_{\text{context}}$ in the paper). We do not intervene on the _parameters_, but rather we substitute a subset of $Z_{\text{context}}$, and observe how the new context activations change the behavior of the LM. This procedure is explained in greater detail in section 2.2 of the initial submission, which, in the current version, is the second half of section 2.

Reviewer comment (Reproducibility):
Among what did the authors calculate the mean of log probs in Figure 4?

Author response:
> Among what did the authors calculate the mean of log probs in Figure 4? 

In Fig. 4, we sampled 100 2-entity contexts from the CAPITALS task, with different values of entity/attributes. For each context, we perform the position interventions, and the resulting log probs are averaged across the 100 samples.

Reviewer comment (Writing):
In equation (2), they do sampling, but from what did the authors sample?

Author response:
> In equation (2), they do sampling, but from what did the authors sample?

In general when obtaining estimates for $\Delta(k)$ we draw 500 samples of different contexts (with each entity/attribute chosen uniformly at random without repetition), but with the same number of entity/attribute pairs.

Reviewer comment (Presentation):
The lines in the legend do not match those in the graph in Figure 4. In the bottom graph, the mean log probs increase or decrease with the positions, but the authors say they are small changes. How can they be said small?

Author response:
> The lines in the legend do not match those in the graph in Figure 4. In the bottom graph, the mean log probs increase or decrease with the positions, but the authors say they are small changes. How can they be said small?

The changes in log prob are small in that they do not change the prediction of the language model. It may have become less confident in its predictions, but it still outputs the correct answer.

In addition, we interpret the overall trend in the mean log probs as a _position dependent bias_. This is is a bias because regardless of which entity is being queried, shifting an attribute to a later position always makes it have a higher log probability. We speculate that this is due to the long-term decay property of [Rotary Position Embedding (RoPE)](https://arxiv.org/abs/2104.09864) that the original authors discussed. 

We account for this bias using the median-calibrated accuracy instead of accuracy in our experiments. In appendix C we discussion this bias further, and we have also updated the main text to contain some of this reasoning.

Reviewer comment (Presentation):
Section 5 supports the generality of the mechanism across other similar tasks, which further strengthens the central contribution. The existence of the alternative _direct binding_ mechanism was interesting to see and another good contribution of the paper, although it was not super clear without reading the corresponding Appendix.

Author response:
We agree that this work has identified a new theory that explains emergent in-context binding behaviors - we believe that this will be an active and fruitful line of research.

> alternate grammatical structures

The PARALLEL task in our original submission explores this. An example prompt from it looks like:

```
Answer the question based on the context below. Keep the answer short.

Context: Alice and Bob live in the capital cities of France and Thailand respectively.

Question: Which city does Bob live in?

Answer: Bob lives in the city of
```

All of our binding experiments work for PARALLEL tasks. This shows that the binding ID is unlikely to be assigned based on simple heuristics like proximity or whether the terms are part of the same sentence, but rather involves some amount of lexical parsing.


> what about > 2 attributes

We have performed an additional experiment where we bind 3 terms together. 

```
Answer the question based on the context below. Keep the answer short.

Context: Carol from Italy likes arts. Samuel from Italy likes swimming. Carol from Japan likes hunting. Samuel from Japan likes sketching.

Question: What does Carol from Italy like?

Answer: Carol from Italy likes
```

In this binding task, there are three terms that have to be bound together: name, country, and hobby. We can query any two of the three terms, and ask the language model to retrieve the third. The above example shows how we query for the hobby, given the name and the country. We query for country and name by asking instead:

To systematically study three-term binding, of the three attributes we choose one to be the _fixed_ attribute, one to be the _query_ attribute, and one to be the _answer_ attribute. Altogether, there are $3!=6$ possible assignments. For every such assignment, we can perform the same set of analysis as before. 

To illustrate, suppose we choose country as the fixed attribute, name to be the query attribute, and hobby to be the answer attribute. An example prompt for this assignment will look like:

```
Answer the question based on the context below. Keep the answer short.

Context: Carol from Italy likes arts. Samuel from Italy likes swimming. 

Question: What does Carol from Italy like?

Answer: Carol from Italy likes
```

We can then vary the query attribute (Carol or Samuel) and check if the LM returns the corresponding answer attribute (arts or swimming), analogous to how in 2-term binding we vary the entity and check if the LM returns the corresponding attribute. To test the binding ID mechanism, we apply mean interventions and verified that the accuracies flip completely when we apply mean interventions on either query or answer attributes, but not both. 

The full results, where the mean interventions are performed for all 6 assignments, are shown in the new Appendix I.

Reviewer comment (Writing):
"If there no context consistent...say that the LM is _confused_" - I didn't fully understand this why or how this would happen, and also it didn't seem to come up elsewhere in the paper. Is it necessary to include?

Author response:
> "If there no context consistent...say that the LM is confused" - I didn't fully understand this why or how this would happen, and also it didn't seem to come up elsewhere in the paper. Is it necessary to include?

We meant confusion as a catch all for when model behavior doesn’t seem consistent with any belief (perhaps because of a blunt intervention that broke things). 

In section 4.2 in particular, we constructed an experimental setting where we deliberately use mean interventions to make the model give equal predictions to the attributes regardless of which entity is queried (i.e. the model is completely ambivalent between the two attributes). This is also an example of confusion.

Reviewer comment (Presentation):
Figure 3 b seems like the wrong plot? It doesn't show what the text claims but rather seems to be a copy of Figure 3 a.

Author response:
> Figure 3 b seems like the wrong plot? It doesn't show what the text claims but rather seems to be a copy of Figure 3 a.

Yes, this was a plotting mistake. The paper has been updated.

Reviewer comment (Writing):
In Section 4.1, binding ID $k$ is represented by some pair of functions $\[b_E(k),b_A(k)\]$.

Author response:
We’ll start by addressing your “central confusion”. You asked
> So, is the hypothesized mechanism that:
> 1. LLM, internally, have some weird analogue of LISP gensym that generates arbitrary but unique abstract binding IDs like "square" and "triangle" (but of course are really lie in a linear subspace of vectors) when an entity/attribute is seen in-context?
> 2. OR that LLM have some stable/consistent binding function

The answer is option 2. We find the binding ID is unchanged when we change the value of any entity or attribute. i.e. if we have:

> “Alice lives in Thailand. Bob lives in France.”

and

> “Charlie lives in Singapore. David lives in Spain.”

Then, Alice and Charlie will have the same binding ID, as will Bob and David. 

To generalize this, for a fixed binding task, the $k$-th entity will always have the same binding ID and the $k$-th attribute will always have the same binding ID. In section 3, we arbitrarily label the binding ID that the $k$-th entity/attribute pair has as $k$, and in section 4.1, we showed that the binding ID labelled $k$ is actually represented as a pair of binding vectors $[b_E(k), b_A(k)]$.

We have updated the paper to clarify this. This is a very subtle point that you have raised and we are happy to clarify any doubts you have.

Reviewer comment (Theory):
The $\Delta$ calculation scheme in Eqn 2 and the expectation sampling seem to crucially depend on, e.g., $b_E(k)$ being a stable/consistent function of (the seemingly arbitrary index?) $k$ across _any context_ $\bf{c}$.

Author response:
We’ll start by addressing your “central confusion”. You asked
> So, is the hypothesized mechanism that:
> 1. LLM, internally, have some weird analogue of LISP gensym that generates arbitrary but unique abstract binding IDs like "square" and "triangle" (but of course are really lie in a linear subspace of vectors) when an entity/attribute is seen in-context?
> 2. OR that LLM have some stable/consistent binding function

The answer is option 2. We find the binding ID is unchanged when we change the value of any entity or attribute. i.e. if we have:

> “Alice lives in Thailand. Bob lives in France.”

and

> “Charlie lives in Singapore. David lives in Spain.”

Then, Alice and Charlie will have the same binding ID, as will Bob and David. 

To generalize this, for a fixed binding task, the $k$-th entity will always have the same binding ID and the $k$-th attribute will always have the same binding ID. In section 3, we arbitrarily label the binding ID that the $k$-th entity/attribute pair has as $k$, and in section 4.1, we showed that the binding ID labelled $k$ is actually represented as a pair of binding vectors $[b_E(k), b_A(k)]$.

We have updated the paper to clarify this. This is a very subtle point that you have raised and we are happy to clarify any doubts you have.

Reviewer comment (Theory):
Confusion about whether the hypothesized mechanism is:
1. LLMs, internally, have some weird analogue of LISP `gensym` that generates arbitrary but unique abstract binding IDs like "square" and "triangle" (but of course really lie in a linear subspace of vectors) when an entity/attribute is seen in-context,  
2. OR that LLMs have some stable/consistent binding function $b_E(k)$, where the first time `llm-gensym` gets called, the result for $k=0$ is returned, and some internal entity counter (?) executes $k++$,  
3. OR some altogether different process.

Author response:
We’ll start by addressing your “central confusion”. You asked
> So, is the hypothesized mechanism that:
> 1. LLM, internally, have some weird analogue of LISP gensym that generates arbitrary but unique abstract binding IDs like "square" and "triangle" (but of course are really lie in a linear subspace of vectors) when an entity/attribute is seen in-context?
> 2. OR that LLM have some stable/consistent binding function

The answer is option 2. We find the binding ID is unchanged when we change the value of any entity or attribute. i.e. if we have:

> “Alice lives in Thailand. Bob lives in France.”

and

> “Charlie lives in Singapore. David lives in Spain.”

Then, Alice and Charlie will have the same binding ID, as will Bob and David. 

To generalize this, for a fixed binding task, the $k$-th entity will always have the same binding ID and the $k$-th attribute will always have the same binding ID. In section 3, we arbitrarily label the binding ID that the $k$-th entity/attribute pair has as $k$, and in section 4.1, we showed that the binding ID labelled $k$ is actually represented as a pair of binding vectors $[b_E(k), b_A(k)]$.

We have updated the paper to clarify this. This is a very subtle point that you have raised and we are happy to clarify any doubts you have.

Reviewer comment (Writing):
Eqn 2: is there any significance to the notation change of introducing $\alpha$?

Author response:
> Eqn 2: is there any significance to the notation change of introducing $\alpha$?

Yes, this is a typo.

Reviewer comment (Evaluation):
"These random vectors have no effect" - I would have expected it to break things. So I guess this is even further evidence that the binding vectors are treated in a "special" way by the LLM machinery? If so, maybe this reasoning and conclusion could be emphasized.

Author response:
> "These random vectors have no effect" - I would have expected it to break things. so I guess this is even further evidence that the binding vectors are treated in a "special" way by the LLM machinery? If so, maybe this reasoning and conclusion could be emphasized.

Yes, one interpretation is that the binding vectors point in a specialized subspace reserved for binding IDs. Thus, random vectors outside this subspace don’t affect things that much. We also conducted a sanity check by scaling up the random vectors - at a large enough scale, the random vectors do in fact break the model's ability to bind.

Reviewer comment (Reproducibility):
For any of this to make sense, I guess you are using a word-level tokenizer and not something like BPE/etc? Except it seems like the direct binding experiment uses sub-word tokens? Again, for clarity it might help to lay this out explicitly.

Author response:
> For any of this to make sense, I guess you are using a word-level tokenizer and not something like BPE/etc? Except it seems like the direct binding experiment uses sub-word tokens? Again, for clarity it might help to lay this out explicitly.

We are using sub-word tokens. In most of our tasks, the entities and attributes are chosen to be a single token wide, since a lot of common words are encoded as a single token.

Reviewer comment (Evaluation):
It is not entirely clear whether the authors’ conclusions are supported by experimental evidence.

Author response:
> It is not entirely clear whether the authors’ conclusions are supported by experimental evidence.

We are happy to discuss any question you have about our experiments and their conclusions.

Reviewer comment (Writing):
Section 2.1: What exactly do you mean by “the k-th entity” and “the k-th attribute”? Is the entity that appears first in the sentence called the 0-th entity?

Author response:
> Section 2.1: What exactly do you mean by “the k-th entity” and “the k-th attribute”? Is the entity that appears first in the sentence called the 0-th entity?

Yes, we use 0-indexing. This is also consistent with the notation of referring to the entities in an $n$-entity context as $E_0, E_1, \dots, E_{n-1}$ and the attributes as $A_0, \dots, A_{n-1}$.

Reviewer comment (Experiments):
Section 2.1: What did you sample N=100 contexts from?

Author response:
> Section 2.1: What did you sample N=100 contexts from?

Given a binding task with a set of possible entities $\mathcal E$ and attributes $\mathcal A$ (e.g. $\mathcal E$ is a set of first names and $\mathcal A$ is a set of countries in the CAPITALS task), we can sample an $n$-entity context by drawing $n$ entities uniformly at random without repetition from $\mathcal E$, and $n$ attributes uniformly at random without repetition from $\mathcal A$.

Reviewer comment (Experiments):
Section 3.2: What is the experimental setting in the experiment?

Author response:
> Section 3.2: What is the experimental setting in the experiment?
We evaluate LLaMA-30b on the 2-entity CAPITALS task, under the position intervention described in section 3.2. The plotted mean log probs are averaged across N=100 samples from the 2-entity CAPITALS task.

Reviewer comment (Evaluation):
Section 4.2: Does the experimental result really mean that the vectors that are not linear combinations of binding ID vectors do not work?

Author response:
> Section 4.2: Does the experimental result really mean that the vectors that are not linear combinations of binding ID vectors do not work?

No. Our experiments in 4.2 only show the converse is true: vectors that are linear combinations of binding ID vectors are also binding vectors.

However, from section 4.1, the fact that random vectors do not affect binding information suggests that random vectors do not work as binding IDs.

Reviewer comment (Presentation):
Figure 5: What are the x and y axes? Are they the coefficients of the basis vectors?

Author response:
> Figure 5: What are the x and y axes? Are they the coefficients of the basis vectors?

The x-axis corresponds to one basis vector. The y-axis is the orthogonalized second basis vector (i.e. $v_2 - \frac{v_2 \cdot v_1}{||v_1||\,||v_2||} v_2$). The basis vectors are plotted as black arrows in each of the heatmaps.

Reviewer comment (Writing):
p. 8: What exactly is a cyclic shift?

Author response:
> p. 8: What exactly is a cyclic shift?
We use a cyclic shift to refer to a cyclic permutation. When $n=3$ as in this case, there are only two cyclic permutations --- either k -> (k+1) mod 3 or k -> (k-1) mod 3. 

In the experiments in fig. 6 we intervene on the context activations so that $E_0$ becomes binded with $A_1$, $E_1 \leftrightarrow A_2$, $E_2 \leftrightarrow A_0$. This can be achieved by either applying mean interventions on the binding IDs on the entities or the attributes. This corresponds to the 'entity' and 'attribute' labels in Fig. 6.

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

Author response:
> p. 2: Does a “sample” mean an example? In statistics, a sample usually means a collection of examples (data points)

Yes, a "sample" here refers to sampling in the traditional sense. A sample of an $n$-entity context is obtained by sampling $n$ entities uniformly at random without repetition from $\mathcal E$ (a universe of entities), and $n$ attributes uniformly at random without repetition from $\mathcal A$ (a universe of attributes).

Reviewer comment (Experiments):
While the ideas presented in this work were novel, it was unclear how generalized they are. The authors presented a series of experiments based on somewhat synthetic datasets. Had the task been reading comprehension, we might not have observed the same mechanism. I think adding more tasks did not provide meaningful results unless they required different reasoning complexities. In addition, the experiments presented in Section 3 only provided anecdotal evidence from a few samples (Factorization and Position Independence claims).

Author response:
__Generalizability__

> I think adding more tasks did not provide meaningful results unless they required different reasoning complexities. 

We agree that reasoning is an important test that we have understood the binding representations. Note that the CAPITALS task already requires a small reasoning step: the LM has to infer the capital city given the country. In our original submission we showed that our experiments preserve this reasoning ability.

To further push on reasoning, we now construct a "one-hop" binding task. Here is an example from it:
```
Answer the question based on the context below. Keep the answer short.

Context: Elizabeth lives in the capital city of France. Ryan lives in the capital city of Scotland.

Question: The person living in Vienna likes rose. The person living in Edinburgh likes rust. The person living in Tokyo likes orange. The person living in Paris likes tan. What color does Elizabeth like?

Answer: Elizabeth likes
```
Here, based on the binding information in the context the LM has to perform two inference steps. The first is to infer that the capital city of France is Paris. The second is to, based on the additional information in the question, infer that the person living in Paris likes tan, and output tan as the correct answer.

This is a more challenging task than our other tasks, and we thus conduct experiments on LLaMA-65b instead of LLaMA-30b. Overall, we found that all of our results still hold: the representations factorize and are position independent in the same way as in the CAPITALS task, and our mean interventions also successfully switched the model beliefs as predicted by the binding ID mechanism. The results are in a new appendix (Appendix H).


> In addition, the experiments presented in Section 3 only provided anecdotal evidence from a few samples (Factorization and Position Independence claims). 

To clarify, the factorization and position independence figures in section 3 show the aggregated results for 100 samples from the capitals task (i.e. 100 prompts with different names and countries filled in), and not "a few" samples. 

To further address "anecdotal"-ness, we have conducted the same experiments for the other binding tasks, namely PARALLEL, SHAPES, FRUITS, and BIOS. The plots are in a new appendix (Appendix J). The experiments do not show significant difference from the CAPITALS results, except for BIOS - for BIOS, attribute factorizability does not hold. This is because BIOS attributes are many tokens wide, whereas the factorizability test only substitutes the first attribute token.

Reviewer comment (Presentation):
Figure 3 shows the same image for (a) and (b). I was not sure why, or was it a mistake?

Author response:
> 1. Figure 3 shows the same image for (a) and (b). I was not sure why, or was it a mistake?

Yes, thanks for catching this. The paper has been updated with the correct plot.

Reviewer comment (Writing):
I found the notion of the *binding function* rather confusing in Eq (1). I had many questions about it. If the binding function specifies how an entity binds to an attribute, why is it a sum of entity representation and the binding vector? What is the relationship between the binding function and the residual stream?

Author response:
> 2. I found the notion of the binding function rather confusing in Eq (1). I had many questions about it. If the binding function specifies how an entity binds to an attribute, why is it a sum of entity representation and the binding vector? What is the relationship between the binding function and the residual stream?

Thanks for the feedback! We’ve updated our paper to improve clarity in section 4.1 where Eq(1) is discussed. 

- The main feature of the binding ID mechanism that we proposed is that instead of _directly binding_ entity to attribute, the language model _indirectly represents binding_ by binding the entity and attribute to the same binding ID (Fig. 1). The entity/attribute binding function is thus a function of entity/attribute and binding ID.
- Binding functions are defined (Section 3) so that we can substitute the entire residual stream of either entity or attribute with the output of the corresponding binding function without changing the behavior of the LM. They should capture all information that is relevant for the LM to solve the binding task.

Reviewer comment (Reproducibility):
How many samples did you use to generate Figure 3? (100 or 500)

Author response:
> 1. How many samples did you use to generate Figure 3? (100 or 500)

We use 100 samples by default. We use 500 samples when estimating the means in the mean interventions used in sections 4 and 5.

Reviewer comment (Experiments):
Were the samples you used to estimate the *differences* the same samples for testing in Table 1? If so, why?

Author response:
> 2. Were the samples you used to estimate the differences the same samples for testing in Table 1? If so, why?

No, we split the dataset into a “train/test” set so that the means are estimated from the train subset and the evaluation of the intervention is done on a test set. More specifically, we split the set of entities $\mathcal E$ and the set of attributes $\mathcal A$ each into two sets, so that the contexts from the train subset and the test subset will not have any entities or attributes in common.