TL;DR: Language models sometimes seem to ignore parts of the chain of thought, and larger models appear to do this more often. Shapley value attribution is a possible approach to get a more detailed picture of the information flow within the chain of thought, though it has its limitations.
Project status: The analysis is not as rigorous as I would prefer, but I'm going to be working on other directions for the foreseeable future, so I'm posting what I already have in case it's useful to others. Code for replicating the Shapley value results can be found here.
Thanks to Jacob Hilton, Giambattista Parascandolo, Tamera Lanham, Ethan Perez, and Jason Wei for discussion.
Chain of thought (CoT) has been proposed as a method for language model interpretability (see Externalized Reasoning Oversight, Visible Thoughts). One crucial requirement for interpretability methods is that they should accurately reflect the cognition inside the model. However, by default there is nothing forcing the CoT to actually correspond to the model’s cognition, and there may exist theoretical limitations to doing so in general.
Because it is plausible that the first AGI systems bear resemblance to current LMs with more sophisticated CoT and CoT-like techniques, it is valuable to study its properties, and to understand and address its limitations.
Shapley values have been used very broadly in ML for feature importance and attribution (Cohen et al, 2007; Štrumbelj and Kononenko, 2014; Owen and Prieur, 2016; Lundberg and Lee, 2017; Sundararajan and Najmi, 2020). Jain and Wallace (2019) argue that attention maps can be misleading as attribution, motivating better attribution for information flow in LMs. Kumar et al. (2020) highlight some areas where Shapley value based attribution falls short for some interpretability use cases.
Madaan and Yazdanbakhsh (2022) consider a similar method of selectively ablating tokens as a method of deducing what information the model is dependent on. Wang et al. (2022) find that prompting with incorrect CoT has surprisingly minor impact on performance.
We use a method similar to Kojima et al. (2022) on GSM8K (Cobbe et al., 2021) with GPT-4 to first generate a chain of thought and evaluate the answer, and then for all chains of thought that result in a correct answer we perform an intervention as follows: we choose a random numerical value found in the CoT, and replace it with a random number in a +/-3 range about the original. We then discard the remainder of the CoT and regenerate it. If the LM is following strictly the CoT described, this intervention should almost always result in an incorrect answer, the same way one would if they made a mistake in one calculation and propagated the error through to the answer (with occasional rare cases where the new value happens to also result in the correct answer, though from qualitative inspection this is very rarely the case).
Some cherrypicked examples (red = intervention, blue = correct continuations that are seemingly non-sequiturs):
We test how frequently this occurs in several different settings (n=100):
Interestingly, if we condition on the CoT answer being correct and the single forward pass answer being incorrect (i.e the LM could only solve the problem with the CoT), the intervened accuracy for GPT-4 is still 0.65.
We would like to get more granular information about the causal structure (i.e which tokens cause which other tokens). One thing we could do is look at how an intervention at each token affects the logprob of each other token. However, one major problem with this is that especially in the larger models, it turns out there’s lots of cases where a token depends on multiple previous tokens in some complicated way. In particular, if a model looks at multiple different places in the context and takes a vote for the most common value, then intervening on any one of them doesn’t change the output logprob a lot, even though there’s a lot of information flow there.
To get around this problem, we instead estimate Shapley values, which take into account all the interactions (in the case where the model takes a vote among 3 values in the context, those three values would each get ⅓ of the attribution). We also normalize the attributions to sum to 1 for each token, clamping negative Shapley values to 0. We do this to make the attributions more comparable across different models.
Here's an example chain of thought in GPT-4:
Here, we can see patterns like the 23 and 20 being copied, or the 3 depending heavily on the preceding 23 - 20. We can also look at some other models:
Interestingly, we notice that the more capable the model, the more spread out the attributions become. We can quantify this as the mean entropy of the parent attributions across all tokens to get a measure of how spread out this attribution is, at least on this particular data sample:
Methodological footnote: the Shapley experiments actually blank out the numbers with underscores, rather than doing the +-3 perturbation of the last section.
Negative shapley values did not occur very often, but this is still somewhat unprincipled. This was primarily to make the entropy calculation work.
We only look at the shapley values for the numbers, because shapley values take exponentially longer to attribute for more tokens under consideration.
Alternative visualization style:
When doing Shapley value attributions for every pair of tokens, there is a dynamic programming trick that we can use to prevent this from becoming n * 2^n: because the model is autoregressive, we can run attributions to only the last token, and, if we take care to save the logprobs for the correct number token at each underscore, compute all other attributions for free.
Speculative hypothesis: Maybe some of the cases in which editing one number doesn't ruin the bottom line result, are not as bad as it sounds, for the following reason: The system is surprised to see the wrong number there, and mentally writes it off as a typo, and proceeds with the true number that it expected in mind. (This happens to me a lot. Sometimes I'll be reading someone's argument about how X leads to anti-Y which correlates with Z, and I'll notice an obvious typo like "don't they mean anticorrelates here?" and then I'll just keep reading assuming they meant what I think they meant instead of what they actually said.)Does the structure of the transformer allow for this sort of cognition?