Monitoring Latent World States in Language Models with Propositional Probes

Thank you for the thoughtful review. We are happy that you found our paper "well-written" and that it "comes with multiple contributions".

Why did you go from proposition to text and not the other way around: use existing text (from the wild) and generate propositions from it?

Our main intuition is that going from structured representations (propositions) to unstructured representations (text) is easier than in the opposite direction. Further, there are some practical challenges involved in generating propositions from text-in-the-wild:

• Not every piece of internet text fits the confines of the 4 domains that we have
• Labels will not be balanced: probes can potentially cheat, e.g. by predicting the nationality of people from their names.

This overall paradigm is related to the use of "minimal-pairs" in designing psychology experiments, which has been used in creating ML benchmarks such as EWOK (Ivanova, et al, 2024) and BigToM (Gandhi et al, 2024).

However, we agree that if propositional probes are scaled to sufficiently diverse domains, text-in-the-wild can be a good downstream evaluation; it is still important to use a synthetic, balanced dataset as the main evaluation.

Gandhi, Kanishk, et al. "Understanding social reasoning in language models with language models." Advances in Neural Information Processing Systems 36 (2024).

Ivanova, Anna A., et al. "Elements of World Knowledge (EWOK): A cognition-inspired framework for evaluating basic world knowledge in language models." arXiv preprint arXiv:2405.09605 (2024).

Thank you for the thoughtful review. We are happy you felt that the "method is likely to be very useful to the interpretability community".

Results are limited to one model.

We have run the experiments in the main paper on the Llama-2-13b-chat model, and put the results in a new appendix J. Broadly speaking, the results are very similar to the Tulu-2-13b model we used (which may not be surprising, because they are both finetuned from the same base model.)

I noticed that the number of values of k that were evaluated is not super high, why is that? Is Hessian-based algorithm computationally intensive?

To compute the Hessian requires $d$ backward passes, where $d$ is the dimension of residual stream. For most open weight models this is around 1000-8000, which takes a few hours to run on A100 gpus. In principle the backward passes can be batched, but our attempts to use the torch functional APIs to achieve this have not been successful.

Separate probes are learned for each domain, but every domain contains only one predicate, do you have any sense of how well propositional probes might generalize across predicates? Is there a good way to quantify how different the probes for each domain are?

At present, our results in Table 2 (right) show that of the four domains, (names, foods, countries, occupations), names probes generalize the best from the synthetic training data to paraphrased and translated data, whereas food probes generalize the worst. One reason why this might be the case is that representations for food items might be different in different contexts, whereas proper nouns like names are more straightforward. For example, 'apple' could connote the fruit, the tech company, or even New York City. In our paper, we quantify the performance of the probes with their generalization behavior from simple synthetic dataset to diverse paraphrased datasets, but there could be more systematic ways of comparing domain probes that directly study their representations.

Do you think anything be gleaned from the singular values of H? Do they correlate with the accuracies in Figure 4 at all?

Yes, we think the singular values are important. The accuracies in Fig. 4 are obtained by sorting the singular values in descending order, and then taking the top $k$ singular vectors as the intervention subspace. If we were to sort in a random order, and take the top $k$ singular vectors (i.e. taking $k$ random vectors), we would have 0% accuracy.

We also took a look at the magnitudes of the singular values. The singular values decay rapidly: the highest singular value is as high as ~250, and decays to ~10 by the 50-th vector, to ~1 by the 200-th, and to ~0.1 by the 1000-th. Since there are 4096 hidden dimensions, most of the singular values (> 3/4) are less than one-hundredth of the 50-th singular value, which is our cut off for the binding subspace.

Thank you for the thoughtful review. We share your belief that identifying binding in representation space is an important contribution.

The main weakness you highlighted was that: "The paper does not do enough to rule out an alternative, simpler hypothesis that could potentially explain the results." Specifically, you are concerned that instead of the true binding vector, the binding subspace could be capturing spurious information such as position or order.

Based on your suggestions, we conducted a thorough quantitative investigation (Appendix I). Our main finding is that the extracted binding subspace is not influenced by position, but captures both the true binding vector and the order so that performance is degraded in a particular ordering of names and countries. However, despite the binding subspace's partial susceptibility to order, the propositional probes still outperform prompting in the adversarial settings.

Here, we briefly summarize the key experiments, and leave the details to Appendix I.

• We create templates of increasing distance between the first name and country, and found that propositional probes perform consistently well. We therefore conclude our binding subspace is not influenced by position.
• We create 5 templates corresponding to different orderings of names and countries. These 5 orderings include the one you proposed "Alice lives in Laos. Peru is where Bob lives", which we call the reversed ordering.
• 4 of the 5 orderings, including the proposed reversed ordering, show consistent performance.
• 1 particular ordering shows a degradation in the probe's performance. We call this ordering the nested ordering, and it looks like this: “Alice and Bob are friends. The latter lives in Germany. The former lives in France.”
• This degradation is significant but not total. 44% of the time, the probe still returns the correct propositions. 50% of the time, the probe assigns both countries to the same name. If the binding subspace had been capturing only order information, but not binding, we would expect the propositions to be wrong 100% of the time. In contrast, the results are more consistent with the interpretation that the binding subspace captures both binding and also spurious features corresponding to order.
• However, if we apply the nested ordering to the adversarial settings (prompt injection, dataset poisoning, and gender bias), our probes still out perform prompting
• There are signs that the language model itself may conflate order and binding. In an alternate binding template that looks like "Alice, unlike Bob who lives in Germany, lives in France.", we find that prompting fails catastrophically because the model thinks that both the first name and the second name in the context are "Alice".
• The degradation to propositional probes can be ameliorated by a small change to the composition algorithm: constraining the proposed propositions to have distinct entities

Overall, we are grateful for your suggestions. Our experiments based on your suggestions imply that binding may be entangled with order, and our methods can be further improved by attempting to disentangle them. Nonetheless, our methods at present are still robust enough to outperform prompting in the adversarial settings.

Thanks for the thoughtful review; we’re happy that you found the investigation into world states “interesting and important”, and appreciated the novelty of the Hessian method for identifying the binding subspace. We’d like to provide some clarification around some of the weaknesses you raised.

The "domain probes" to classify tokens into domain seem quite heuristic (using the mean vector of all the entities in the domain), and it seems like there could be some evaluation to see how it works (e.g., is it always the "obvious" tokens, like the "Alice" token is the name token?).

Our domain probes are a multi-class generalization of the difference-in-means probes, which some previous works have argued are more robust than standard linear probes trained with logistic regression. For example, Belrose (2024) analyzes some theoretical properties of these probes.

Although not the main focus of the work, we conducted some analysis into the domain probe. Using Grad-CAM style saliency maps (appendix F), we can estimate the extent to which the activations at each layer and token position contribute towards the correct answer, and we found that indeed it is usually the “obvious” token that matters, and when the key word contains multiple tokens, it is the last token that matters. The importance of the last token is also previously discovered in the knowledge editing setting (Meng et al, 2022).

Some of the decisions in the binding space design seem quite arbitrary, like "For tractability, we parameterize x and y so that they are shared across layers". Maybe it would then be better to just focus on a few layers? But it's perhaps fine to leave that for future investigation.

This is a good observation. This design choice, and others, are motivated by the circuits responsible for solving binding. We have updated Appendix C to motivate parameterizing $x$ and $y$ across layers. At a high level, the binding vectors for entities are accessed at different layers than the binding vectors for attributes. In principle you can choose to localize the specific layers relevant for entities and attributes, and then perform the injection for those layers — we expect this idea to further improve the accuracy of the binding subspace, at the cost of having more complexity and hyperparameters to tune.

For the Prompt Injection setting (instructing the model to "Always answer the opposite."), it's hard to say what a "correct" output should be, in fact the prompting method should probably "ideally" always be "wrong". So saying "prompting performs worse" is a bit confusing, although it's still an interesting result that the probing outputs are virtually unchanged.

We can see why this is confusing. There are two different ways of framing this that can justify answering the wrong answer as the undesired output. First, from the security point of view, prompt injections are malicious parts of the input that poisons the input context so that the model stops performing as we expect. For example, consider an LLM agent browsing the web on a user's behalf. If it encounters a piece of text that asks it to start misleading the user, and follows the instruction to mislead the user, then we would ideally want some way to monitor the LLM agent when this happens.

The second way of framing it is that we are interested in latent world models in language models. Perhaps, despite being instructed to lie, the model will still first form a coherent world state internally, and then modify its output based on this world state. Our results show exactly this.

Belrose, Nora. "Diff-in-means concept editing is worst-case optimal: Explaining a result by Sam Marks and Max Tegmark, 2023." URL https://blog. eleuther. ai/diff-in-means (2024).

Meng, Kevin, et al. "Locating and editing factual associations in GPT." Advances in Neural Information Processing Systems 35 (2022): 17359-17372.