Interpretability of Language Models for Learning Hierarchical Structures

We thank the review for the comments. As for the comments regarding the paper being very condensed, do you have some concrete suggestions regarding how we can improve the writing? It’s hard to please everyone, and this current paper format is after us communicating with ~20 readers and taking into account their taste of reading/writing styles.

Also, many of the findings are not conventional, including our data, our probing method, so we tried to spend the main body explaining why we designed the experiments this way, and how exactly we certify that the GPT2 is learning to implement a DP. It becomes hard to remove any of these without confusing the readers.

Of course, we'd love to hear more suggestions and wouldn't mind revising this paper further. Thank you!

We sincerely thank the reviewer for supporting this work, and as you have pointed out, transformers work much much beyond fitting causal logic. We wish to use this study as a first step towards (more or less) understanding how LLMs learn “formats” or “structures” before we delve into more meaningful tasks, such as sorting or more complex forms of reasoning, etc. Also, a particular goal of this paper is to refute the argument that LM is a stochastic parrot, so clearly it is not relying on some simple (local) pattern recognition protocol to generate, but can rely on some more involved algorithm (such as DP).

As for your question regarding works that “explicitly compose tokens via dynamic programming to learn latent structures and sentence representations”, we have two thoughts.

A simple thought is that, adding task-specific designs (including a CKY-style encoder, or a p(si | s1..{i-1} s{i+1}..n) type of pretraining objective like 2203.00281), might make us a little farther away from the grand goal of AGI. We could be wrong, but a more desirable goal is to use some unified model / objective for all possible AI tasks whenever possible.

Also, such MaskedLM type of objective (studied in 2203.00281) still makes me worry once again of how deep their method can go — we argued in this paper that BERT-like models fail to learn deep structures because of MaskedLM, and that that PennTree Bank (and possible most real-life language CFG) may not be very “deep”, and this might explain the success of 2203.00281. For a more “complex” CFG like ours, we are curious what’s their performance.

A more important thought can connect to Chain of Thought (CoT). In this submission, we demonstrated that the DP process is learned implicitly. More generally (including our own follow-ups, not to cite due to anonymity), GPT2 can be good at learning other implicit tasks, without CoT. Whenever adding the intermediate steps (such as explicitly teaching the model to do DP), then the learning task may become easier.

Ultimately, we perhaps want a delicate balance: we don’t want the AGI model to do A+B+C step by step; we also don’t expect it to solve a hard math problem without CoT. We think it is interesting to study what’s the balance here, in particular, what are the tasks that can be learned implicitly, in order to better understand what type of model/data we need in tomorrow’s large model training.

We hope this to some extent answers your questions! Thanks again!

We sincerely thank the reviewer for the time raising the questions. Let us try to follow-up with your main comments.

It’s worth noting that long-range dependencies don’t require context-free languages; even simple regular languages (e.g., “first symbol = last symbol”) can capture long dependencies.

Our response: That is correct. This is also the motivation of our paper, as mentioned on Line 48 to study “how language models perform hard tasks, involving deep logics / reasoning / computation chains”. We designed the CFG to be exactly not of such simpler form (like “first symbol = last symbol”).

The CFGs defined in the paper technically describe finite languages (and are therefore regular) because recursion depth is limited by L. This limitation might make it hard to tell whether the models are learning true recursion or just memorizing up to the required depth.

Our response: Can we kindly ask the reviewer what you concretely mean by “memorizing up to the required depth”? If it’s memorizing the strings then we said in Line 181 that there are at least 10^140 possible strings (of our chosen depth L); the model can't even see all such possibilities, not to say memorize them. The model only has 10^8 parameters.

If what you actually mean is that the model has memorized the small set of “rules”, then we believe this is not memorization but rather it has learned to parse/generate CFGs according to the right algorithm?

RNNs are often compared to finite-state automata or pushdown automata. Did you consider comparing transformers to simple RNNs or LSTMs in this study? In the Conclusion: While you might not want to discuss ongoing or future work in detail, it would be helpful to hear a bit more about potential directions or applications of this research.

Our response: Thanks and not in this study. We choose decoder-only models due to their popularity and maturity in the AI community. Another very interesting candidate is Mamba, one of the most popular RNN-type of model, but there’s one concrete tradeoff issue. Mamba’s SSM layer has (3D)xD dimension (serving as long-term memory), where in full attention one can read from LxD dimension data (all the previous L tokens times hidden dimension D). This means, Mamba will be more desirable only when L>>3D. In this regime, the most interesting study IMHO would be whether Mamba can learn a very long CFG with small dimension D. This setup will be significantly different from this current submission, so we are exploring this as a follow-up.

We sincerely thank the reviewer for supporting this work, and let us try to address your questions.

Question:

The paper suggests that Transformer models learn CFGs by emulating dynamic programming (DP) algorithms. Would it be possible to test this hypothesis more directly, perhaps by comparing the model’s internal representations to the states of known DP algorithms? The current evidence feels somewhat indirect.

Answer: The difficulty is that there are “infinitely” many ways to implement even a specific dynamic programming. While humans typically write code like “for k = i to j” and break whenever it suffices to certify that the subsequence from i to j can be generated from a CFG node, language models do this (kind of) in an arbitrary order. This is why we can only verify that those DP states in the backbone (i.e., DP states shared by essentially all DP algorithms, regardless of the implementation) are reached by the language model.

Question:

The paper also notes BERT’s weaker performance in predicting deep non-terminal (NT) structures but doesn’t delve into the reasons behind this. While attributing it to the locality of the MLM task seems plausible, more direct evidence would strengthen this claim.

Answer: Our original way of thinking is to contrast the NT prediction accuracies on lower/higher NTs to show that BERT “does not reason very deeply” on such CFG data. We’ll be more than happy to discuss with you regarding what you think can be “more direct evidence”.

One proposal is to contrast two very close CFGs, one has root->a b b and the other uses root -> a a b, where all other CFG rules stay the same. In the most extreme case, let a = 11..1122..2211..11 and b = 11..1133..3311..11 where “11..11” means 100 copies of 1 and similarly for “22..22” and “33..33”. Then if you use masked LM then the BERT model will certainly learn 1 from its surrounding 1’s, and 2 from its surrounding 2’s, 3 from its surrounding 3’s. It will not distinguish root -> a b b vs root -> a a b. (One can make this more complex by designing more complex / deeper a and b CFG nodes.)

But how to test this? Unless you mask out the 100 consecutive 2’s or 3’s you cannot tell the difference; but once you do that, it’s not a typical MLM data (since in MLM, one randomly masks tokens out) and it’s not surprising that BERT models cannot perform well on this.

Another proposal is to use BERT to do generation, to put all the tokens to the right as <MASK> and use the first <MASK> to do generation. This is in the same spirit as above, and the performance is poor as expected.

Overall, in our mind, a “direct evidence” sounds like directly applying BERT models to see the difference, but this seems to require masking out a large consecutive number of tokens and demonstrate that BERT’s performance is poor. While this can be easily done, it does not seem convincing, and BERT is not trained on such type of data. — This also explains exactly why BERT prefers local reasoning.

Question:

Could the authors discuss potential limitations? Acknowledging the method’s constraints would be helpful for future research directions.

Answer: Certainly, for instance we designed the CFG to be “canonical” in order for the probing to better succeed. If there are global ambiguities then (meaning there’s no unique answer to the probing, such as if a->b c; b->1 2 or 1 2 3, c-> 3 1 or 1, then when a-> 1 2 3 1 there’s no unique answer to where this 3 comes from) then it will naturally fail.

Also, we used “infinite” training data and the performance shall degrade if one controls the amount of training data. We intentionally left these out because we want to see, in a controlled environment (with infinite data + canonical CFG) what’s the capability of GPT2. We want to for instance refute the argument that LM is stochastic parrot, so clearly it is not relying on some simple (local) pattern recognition protocol to generate, but can rely on some more involved algorithm such as DP.

Another technical limitation is our probing technique, where we did linear probing to extract the NT information, where the NT has just a dozen choices. When probing from dimension D networks, an information that has M possible choices, there is a linear head of dimension D x M. When the information we wish to probe is diverse (e.g., M>>100), this can be a problem — too many trainable params, etc. In our follow-up work we actually developed additional probing techniques to tackle situations like this, and let us not go into details to remain anonymous.

Thanks again for your time!