How transformers learn structured data: insights from hierarchical filtering

We thank the reviewer for carefully reading our work and providing valuable feedback, which we will use to improve our presentation. We would first like to address the overall weaknesses identified by the reviewer:

(W1) I'm not sure whether the transformer accuracy matches that of belief propagation is surprising … transformer is simply learning the "optimal" prediction function. We thank the referee for his comment, and we agree, in a sense. As stated in our introduction, transformers assimilate natural language that is vastly more complex than our data model. What we think is key in our study is understanding how it does so, both in time during learning, capturing the longer range correlations through a clear staircase identified thanks to the filtering, as well as in ‘space’ within the architecture, organizing the computation throughout its layers in an interpretable way, strikingly similar to the most obvious implementation of BP in l layers. In the context of this data model, BP is the exact (and thus optimal) oracle, so we will rephrase “learning to implement BP” as “learning to perform optimal inference”.

(W2) I think the writing of this paper can be improved. We are currently working on streamlining the first half of our paper, in particular, to be more straight to the point. We will clearly highlight our contributions, which have been received with some confusion. Moreover, we added some additional experiments, that strengthen our claims in regards to the model approaching the implementation of the exact inference algorithm.

(W3) While the construction sec3.5 is interesting, it does not mean that the learned transformer is actually doing that. This is entirely true, and we have added a line highlighting that this is an existence proof for the implementability of BP in an l-layer transformer, but we are not claiming that it is exactly being implemented. In fact, the described implementation is meant to be understandable and requires, for example, full disentanglement between positional and semantic information, which is not forced in our experiments. Nonetheless, we believe the existence of this implementation is non-trivial, and that it strengthens our paper significantly relative to other works that have attempted to study the parsing of CFGs in transformers, since it can be used in practice as a tool for interpretability. The state-of-the-art work of Zhao et al. (2023), notably only provides a possible implementation that requires more attention layers than BERT has. We leave a precise study of the exact relation between the theoretical implementation vs the learned implementation for future work, but it is easy to notice that the organization of the attention layers qualitatively agrees between the two, suggesting a strong link.

To address the referee’s questions:

(Q1) Why the accuracy would have a drop in the middle of training. This behavior is very important as part of our understanding of the learning dynamics, so we regret that we have not successfully communicated this point (as pointed out also by reviewer VwvV). We are working on clarifying it. In a nutshell, the accuracy goes up and down on test sets generated from different (higher) filtering data than the training set, and therefore OOD from the perspective of the training model. The explanation comes from the fact that the transformer discovers the existence of higher hierarchical correlation levels (i.e., longer-range correlations) sequentially, during training. It starts by imputing a simplistic explanation for the data (e.g., a correlation length of 2), thus initially increasing its accuracy on the corresponding OOD data. But then an additional correlation level (e.g., correlation length of 4) is imputed, and the transformer stirs its predictions accordingly. At this point, the accuracy on the simpler, more factorized data decreases, since the model is learning to assume a richer correlation structure between the symbols, which is not present in the simpler OOD data. This staircase behavior is novel in our work and may be a more general feature of deep network learning dynamics (Refinetti et al., 2023; Bardone & Goldt, 2024; Rende et al., 2024).

(Q2) The writing for section 3.3 can be improved by giving more easy-to-understand intuition. We are making an effort to improve our presentation, putting stronger accents on the main results and on their discussion, and simplifying the technical details sections.

References:

Bardone, L., & Goldt, S. (2024). Sliding down the stairs: how correlated latent variables accelerate learning with neural networks. arXiv preprint arXiv:2404.08602.

Refinetti, M., Ingrosso, A., & Goldt, S. (2023). Neural networks trained with SGD learn distributions of increasing complexity. In International Conference on Machine Learning (pp. 28843-28863). PMLR.

Rende, R., Gerace, F., Laio, A., & Goldt, S. (2024). A distributional simplicity bias in the learning dynamics of transformers. arXiv preprint arXiv:2410.19637.

Zhao, H., Panigrahi, A., Ge, R., & Arora, S. (2023). Do transformers parse while predicting the masked word?. arXiv preprint arXiv:2303.08117.

(W4) Evidence of BP vs other algorithms: We appreciate this concern of the reviewer, and admit that further evidence, which we now provide in the updated version, allows us to more convincingly support our claim. Indeed, we have added in the revised version of our manuscript a more precise metric in the form of the Kullback-Leibler divergence between the BP marginals and the softmax (instead of argmax used for prediction) of the neural network outputs. We find that this quantity decreases following a staircase along the factorization levels during training, and eventually reaches small values in fully trained networks, showing that the ‘full’ (per-sample) output of the networks and BP match, and not just their accuracies. Nonetheless, we would like to stress several points that were perhaps insufficiently clear in the manuscript. In our data model, Belief Propagation is not simply an effective algorithm, it is the information-theoretically optimal oracle for any symbol inference task, whether it be the root or the leaves. As a result, matching its performance both in-sample (tested against the model it was trained on) and out-of-sample (tested against larger filtering levels) is a strong clue that the transformer implementation could be a close approximation of the exact algorithm. Moreover, it may be shown that the inside-outside algorithm is in fact identical to BP (Sato, 2007), when the topology of the parsing tree is given. Generic probabilistic CFGs may indeed be seen as tree-based graphical models with an additional probability distribution over the realization (i.e., topology) of the parsing tree. Finally, we would like to point out that the alignment between the output of the transformer and the optimal inference algorithm constitutes only part of the presented evidence. It has to be understood in conjunction with the ‘spatial’ organization of the attention maps, the interpretation of which is verified by our probing experiment at higher ancestry levels; this organization is compatible with the computational graph of BP. We feel that we have not managed to relate these points together convincingly in the current version of the manuscript, and we will address this issue by adding a comprehensive summary of our results, to paint a more complete picture of our analysis.

(W5) Major figures in the appendix: Unfortunately, we cannot easily comply with this request by presenting all our evidence in the main text, given the stringent page requirement of ICLR. On the other hand, we believe that the results presented in Fig. 7 and the content of Appendices C.3 and C.4 provide secondary evidence in support of our claims. Fig. 7, for instance, shows that also transformers trained on filtered data models can reach optimal inference performance, which in our opinion is mostly a sanity check (optimal inference is already reached, as shown for the full data model, in Fig. 1). In a similar vein, Appendices C.3 and C.4 are consistent with our findings but do not lead to any independent conclusions.

We thank the reviewer for carefully reading our work and providing valuable feedback, which we will use to improve our presentation. We would first like to address the overall weaknesses identified by the reviewer:

(W1) Novelty relative to Allen-Zhu & Li (2023): We understand the reviewer’s concern on this point, especially given the thorough experiments performed in this work on probabilistic CFGs. However, we believe that the objectives of our works are markedly different, albeit complementary. In a nutshell, their work demonstrates that GPTs, which are large and pretrained transformer-based models, have the ability to efficiently parse probabilistic context-free grammars. They notably show this by performing ‘invasive’ probing experiments on these large models, for instance predicting some ancestors in the (variable topology) trees characterizing the samples. However, they do not address (i) what learning dynamics the transformers may follow to achieve this–do they progressively discover the existence of ancestry through training? do they start by combining adjacent symbols through the attention and progressively include the relation between increasingly distant tokens in the sequences?–; (ii) how the trained transformers process the information within the architecture, and the possible relationship between the underlying parsing tree and the neural network; (iii) if the exact inference algorithm, here the inside-outside algorithm for probabilistic CFGs is even implementable in the neural network architecture they are studying, or if the transformers must necessarily rely on some approximation in their implementation due to some architecture related infeasibility. While Allen-Zhu & Li (2023) undoubtedly provide evidence that large transformer-based networks may implement the optimal algorithm for CFGs, we thus believe that they do not tackle the problem from a mechanistic interpretability standpoint. On the other hand, we propose a simplified setting that allows us to specifically tackle points (i)-(iii) stated above, and to gain valuable insights into the inner workings of transformers and how they learn from structured data. Finally, note that the proposed filtering procedure is only possible if tree topology (and the sequence length) is fixed, which is the case for our simplified setting, but not for general CFGs (see also Q3-Q4). Generalizing this tool, which was essential in understanding the sequential discovery of the hierarchical correlations, to the general parsing case is an interesting yet challenging future research direction.

(W2) Simplicity of the task: Somewhat related to the point above, we believe that this perceived weakness of our work strongly depends on the question one sets out to answer. We agree with the referee that it is indeed not a surprise in itself that, given enough samples, transformers manage to fit the training data distribution–after all, we know that actual language is perfectly mastered by LLMs. However, our work sets out to answer the question of how transformers are able to assimilate complex probability distributions, both in the sense of the training dynamics as well as the ‘spatial’ organization of trained transformers and how they perform computations throughout the architecture. In this context, we believe that our setting provides an intermediate complexity between probabilistic CFGs (see the point above related to Allen-Zhu & Li (2023)), and the typical state of the art mechanistic interpretability settings that often rely on simple mathematical tasks such as modular addition (Zhong et al., 2024) or else histogram counting (Behrens et al., 2024). As also mentioned in the answer to the previous point, the fixed depth allows us to introduce the filtering procedure, which is a key interpretative element in our analysis.

(W3) Clarification of the goals of the paper: We thank the reviewer for motivating us to restructure our abstract and introduction, which did not carry the goals and messages of our paper properly. In the revised version, we aim to recenter the introduction towards the title of the paper and our overarching goal of understanding how structure may be digested in transformer architectures and notably in the attention mechanism. The filtered data model is therefore a tool that allows us to achieve this goal in the context of well-controlled inference problems. We believe that applications of our insights towards practical problems, in NLP or other, are tangible. We added a discussion on the fact that the structure of data and wide-spanning correlations being progressively incorporated by the transformer during training is a feature that has been found in widely different contexts, and that may also have valuable implications for the practical design of training protocols.

(Q9) Why do you report validation accuracy but not test accuracy? We apologize for the abuse of terminology on our part. We do not have distinct test and validation sets in our experiments. The models are trained on a training set for a fixed amount of epochs, and we simply measure the test accuracy for the final configuration of model weights (not the best). We corrected this mistake by replacing validation accuracy with test accuracy throughout the paper. Why does the accuracy go up and then down? This behavior is very important for the interpretation of the learning dynamics. We will clarify this point in the revised version, given that the current presentation was ineffective. In a nutshell, the accuracy goes up and down on test sets generated from different (higher) filtering data than the training set, and therefore OOD from the perspective of the training model. The explanation comes from the fact that the transformer discovers the existence of higher hierarchical correlation levels (i.e., longer-range correlations) sequentially, during training. It starts by imputing a simplistic explanation for the data (e.g., a correlation length of 2), thus initially increasing its accuracy on the corresponding OOD data. But then an additional correlation level (e.g., correlation length of 4) is imputed, and the transformer stirs its predictions accordingly. At this point, the accuracy on the simpler, more factorized data decreases, since the model is learning to assume a richer correlation structure between the symbols, which is not present in the simpler OOD data. This staircase behavior is novel in our work relative to all the PCGF publications mentioned above. We believe that it may be a more general feature of deep network learning dynamics (Refinetti et al., 2023; Bardone & Goldt, 2024; Rende et al., 2024). Did you not use the best checkpoint when evaluating on OOD data? As stated above, we used the model weights obtained at the end of training. Our objective is to understand the inner workings of the learning, not to achieve the best possible performance on the OOD data.

(Q10) Why use accuracy instead of perplexity for MLM? Since the CFG can be ambiguous, there isn't only one correct answer, right? This is a fair remark, which was also indirectly raised by reviewer c8Kw. We have taken it into account in our revised version, including new experiments that we believe strengthen our evidence. We measured the Kullback-Leibler divergence between the BP marginals and the softmax (instead of argmax for a prediction) of the transformer output. This is not exactly the perplexity, however, we believe that it is more relevant for our point of understanding how the networks perform inference while taking into account the necessarily ambiguous nature of the prediction. We find that the KL divergence decreases in training following the same staircase scenario on the filtered models as described in Q9, and show that the full probabilistic prediction matches the exact one given by BP, and not only the accuracy.

References:

Allen-Zhu, Z., & Li, Y. (2023). Physics of language models: Part 1, context-free grammar. arXiv preprint arXiv:2305.13673.

Bardone, L., & Goldt, S. (2024). Sliding down the stairs: how correlated latent variables accelerate learning with neural networks. arXiv preprint arXiv:2404.08602.

Behrens, F., Biggio, L., & Zdeborová, L. (2024). Understanding counting in small transformers: The interplay between attention and feed-forward layers. In ICML 2024 Workshop on Mechanistic Interpretability.

Khalighinejad, G., Liu, O., & Wiseman, S. (2023). Approximating CKY with Transformers. arXiv preprint arXiv:2305.02386.

Krzakala, F., & Zdeborová, L. (2009). Hiding quiet solutions in random constraint satisfaction problems. Physical review letters, 102(23), 238701.

Mossel, E., Neeman, J., & Sly, A. (2014). Belief propagation, robust reconstruction and optimal recovery of block models. In Conference on Learning Theory (pp. 356-370). PMLR.

Refinetti, M., Ingrosso, A., & Goldt, S. (2023). Neural networks trained with SGD learn distributions of increasing complexity. In International Conference on Machine Learning (pp. 28843-28863). PMLR.

Rende, R., Gerace, F., Laio, A., & Goldt, S. (2024). A distributional simplicity bias in the learning dynamics of transformers. arXiv preprint arXiv:2410.19637.

Sato, T. (2007). Inside-Outside Probability Computation for Belief Propagation. In IJCAI (pp. 2605-2610).

Zhao, H., Panigrahi, A., Ge, R., & Arora, S. (2023). Do transformers parse while predicting the masked word?. arXiv preprint arXiv:2303.08117.

Zhong, Z., Liu, Z., Tegmark, M., & Andreas, J. (2024). The clock and the pizza: Two stories in mechanistic explanation of neural networks. Advances in Neural Information Processing Systems, 36.

(Q5) What is $O_a$? What is $q$? We start by acknowledging that our presentation was not effective, we will rewrite this paragraph for clarity. To address the question: $q$ is the size of the dictionary of symbols that can be taken by any element of the sequence, and by any node on the ancestry tree starting from the root (see Q4). $O_a$ represented the set of possible children pairs generated by a parent symbol $a$ (i.e. the allowed production rules). What we mean by $O_a \cap O_{a’} = \emptyset$ is that the production rules we consider are completely distinct for each parent symbol, making our unfiltered model non-ambiguous (at odds with more general PCFGs). Given a pair of children, there can only be one possible parent in the layer above. $|\cup_a O_a |=q^2$ on the other hand means that any possible children pair can be produced, and the rules are equal partitioned among the $q$ possible parent values. To summarize, our model is described by a transition tensor that has $q\times q \times q$ entries. It is then organized as q ‘slices’ populated by $q$ entries (the rest of the entries being zero), which are non-overlapping for each slice.

(Q6) l.144: It's not clear to me what this means. Can you express this in equations? This is related to the point above. What we mean is that in the unfiltered cases, transitions in the tree are of the form $M(a \to bc)$ where a is the parent symbol and b and c are the children, as shown by black factors in Fig.1a. By strongly correlated we mean that in each transition b and c are not drawn independently but as a pair. We hope this clarifies our statement.

(Q7) Can the root always be uniquely determined by the input symbols? Yes, this is the case with our choice of non-ambiguous production rules in the unfiltered (full hierarchical) model (see Q5). Given a pair of children symbols, one can exactly infer their parent. Combining this for all pairs in the sequence and going up the tree, the root can be determined with certainty. This difference with PCFGs is central, and justifies our framing in terms of a probability distribution described by a factor graph rather than in terms of CFGs, see Q3.

(Q8) How does the BP algorithm described in the main text relate to the experiments? Given any filtered/unfiltered hierarchical model, one can derive the associated BP algorithm (using explicit knowledge of the production rules and the corresponding transition rates), which represents the exact oracle for any inference task defined in the context of the graphical model. As such, it provides an information-theoretic bound for the accuracy of the reconstruction of hidden symbols (e.g, the root, or the masked tokens). This is what is shown by the black dashed line in Figs. 1b), 1c), 1d), Fig. 3 and Fig. 6; and by all dashed lines in Fig. 7. It also provides a comparison point in out-of-sample tests, i.e. used on data generated from different levels of filtering compared to the training data, which is shown by the colored lines on Figs. 1b), 1c), Fig. 3 and Fig. 6. As stated in the point above, the accuracy is trivially equal to 1 for the root classification given an entire sequence for fully hierarchical data.

To address the referee’s questions:

(Q1) What "knob" does the filtering parameter k represent? In simple words, k>0 creates a shortcut in the hierarchy of the generative tree (see Fig. 1). When we filter out the upper levels, the root directly generates 2^(k) ancestors, independently one from the other. This means that strong correlations, induced by the remaining (l-k) unfiltered layers, survive only within blocks of size 2^(l-k). Therefore, (as suggested by the reviewer) the lower the parameter k, the stronger the longer-range correlations in the data.

(Q2) Other relevant paper Khalighinejad et al. (2023). We were not aware of this work, indeed related to transformers on probabilistic CFGs. We thank the referee for bringing it to our attention, the reference was added in our introduction. Similarly to Allen-Zhu & Li (2023) and Zhao et al. (2023), the paper studies the ability of transformers to reconstruct parsing trees, but does not address (i) the learning dynamics, (ii) the implementation of the algorithm within the architecture and (iii) the existence of an exact transformer based implementation of the classic algorithm. This additional reference reinforces our interest in interpreting algorithm discovery within transformer architectures.

(Q3) Graphical models & BP vs CFG parsing algorithms. This is an interesting point, which we can answer along two main motivations. The first is simplicity and tunability. As mentioned above, we set out to go beyond the existing works on CFGs on the question of mechanistic interpretation and to understand more precisely how an inference algorithm is learned and implemented in transformers. This required a simplification of the data model, and notably relying on a fixed topology, allowing us to introduce the filtering parameter–i.e. to tune the range of correlations in sequences. In this setting, the standard CFG terminology would be misleading, as our data model lacks the central feature of probabilistic CFGs: variable sequence length and the related existence of terminal and non-terminal symbols. The other motivation is that our setting is slightly more model agnostic. Indeed, while not tuned to describe natural language, probabilistic graphical models can describe many other objects, from protein sequences to random graphs. In this respect, our work could be used to make contact with markedly different problems where transformers could also be effective, for instance community detection in stochastic block models (Mossel et al., 2014) or planted graph coloring problems on graphs (Krzakala & Zdeborova, 2009), for which Belief Propagation is still an effective algorithm.

(Q4) What would the equivalent PCFG be, incorporating the depth constraint and k? As already mentioned in W4 and Q3, it is unclear that a standard PCFG could be formulated to be equivalent to our model, since it would require one to remove the intrinsic ambiguity related to the topology of the parsing tree. Only in a fixed topology case (implying a fixed sequence length), the hierarchical tree can be cut at a specific level with the filtering procedure. In fact, in a general PCFG, the terminal symbols could be found at any level of the parsing tree, and their correlations to the other leaves is to be determined at inference time. In our model, the correlation between the elements of the sequence (i.e. the terminal symbols), is fully specified by the positions in the sequence. Therefore, generalizing the hierarchical filtering procedure to PCFGs is non-trivial. Moreover, we are not considering a separate dictionary for the non-terminal symbols and for the root symbol. While this choice may make little sense from the perspective of linguistics, it allows us to define a non-trivial root classification task, which would not be possible in a normal PCFG.

We thank the referee for taking the time to go through our response and engage in this dialogue.

On our contributions

Our original goal, which we hope is more clearly stated in the revised version, is to understand the learning process and the computation of a transformer trained on our data in a mechanistic way, in the light of the exact, known inference algorithm. To achieve that, on top of the point highlighted by the reviewer,

• We show that the transformer reaches calibration, i.e. approximates the exact output of the inference oracle, even on OOD data.
• We analyze how the exact computation is embedded in the transformer weights (learning “in space”).
• We provide a novel feasible implementation of the exact algorithm within the same architecture and show that probing experiments on the trained transformer qualitatively align with some of its properties.

Equivalent PCFG formulation

A probabilistic context-free grammar G can be defined by a quintuple $G=\left(M,T,R,S,P\right)$, where $M$ is the set of non-terminal symbols, $T$ is the set of terminal symbols, $R$ is the set of production rules, $S$ is the start symbol, $P$ is the set of probabilities on production rules.

In our model, we consider the special case where terminals and non-terminals coincide $M=T\overset{def}{=}\mathcal{S}$, and there is a set of possible root symbols $\mathcal{R}$. We pick $|\mathcal{R}|=|\mathcal{S}|=q$. The production rules $R$ and their probabilities $P$ are defined as:

• For $k>0$: $\forall r\in\mathcal{R}$, all the trasitions of the type: $r\rightarrow s_{1}...s_{2^{k}}$ are allowed, with probability $P\left(r\rightarrow s_{1}...s_{2^{k}}\right)=\prod_{k}P^{(k)}\left(s_{k}|r\right)$, computed as in Eq. (1) in the revised paper. Moreover, $\forall s\in\mathcal{S}$, we allow $q$ transitions of the type: $s\rightarrow s_{1}s_{2}$ with probability $P\left(s\rightarrow s_{1}s_{2}\right)=M_{ss_{1}s_{2}}$.
• For $k=0$: the production rules for the root correspond to those among the symbols: $r\rightarrow s_{1}s_{2}$ with probability $P\left(r\rightarrow s_{1}s_{2}\right)=M_{rs_{1}s_{2}}$.
• For $k=\ell$: we only have root-to-leaves production rules of the type: $r\rightarrow s_{1}...s_{2^{\ell}}$, with probability $P\left(r\rightarrow s_{1}...s_{2^{\ell}}\right)=\prod_{k}P^{(k)}\left(s_{k}|r\right)$.

Moreover note that, since we impose a non-ambiguity constraint over the transitions, we have that $\forall$ children pair $s_1,s_2$, $M_{p(s_1,s_2),s_1,s_2}\neq 0$ for a single parent symbol $p(s_1,s_2)$, i.e. any pair of children symbols can come from only one parent symbol.

Finally, we consider a fixed parsing tree: a full tree with $2^\ell$ leaves.

Question on BP

To answer the referee’s question, the time complexity of the BP algorithm on a tree is linear in the sequence length $n=2^\ell$, or exponential in the tree depth, $\mathcal{O}(n)$. Indeed, there is a large computational discount from knowing the topology of the parsing tree, compared to the inside-outside algorithm (which instead scales as $\mathcal{O}(n^3)$).

(Q3) Since the task is to predict the root, could the authors elaborate more on why they believe the weaker correlations among tokens within a sequence is causing difficulty for learning? We thank the reviewer for raising this question, as it pushed us to look a bit more attentively at this issue and to provide a new (simpler) interpretation that will hopefully convince them as well. To answer their question on why we thought this was related to token-token correlations rather than sequence-root correlations, this is because the k=l case is actually the one which has the smallest mutual information between the sequence and the root, while it is the easiest to learn from Fig.7 in the appendix. Moreover, this mutual information should be strictly increasing as k decreases, while the sample complexity as a function of k appears to present a non-monotonic behavior, being the smallest for k=l and the second smallest for k=0, therefore we ruled out the correlations between the root and the entire sequence as the cause of this phenomenon. However we missed an important aspect of the training, which is singular to the k=0 case. The deterministic nature of the root inference problem for k=0 makes it unique for two reasons: the first is that the logits outputted from the network need not be calibrated, so the accuracy can reach the optimum without the transformer having fully implemented an algorithm equivalent to BP, whereas the relative weights of prediction must be well understood to match the optimal inference in the ambiguous k>0 cases, this is visible on Fig.7; the other is that this being said, matching the BP is also easier in the k=0 case because the training cross-entropy loss corresponds exactly to that computed with the true BP marginals (that are also delta distributed due to the determinism again) whereas in the k>0 cases the training loss does not guide explicitly to the BP marginals, this is visible by re-plotting Fig.7 with the Kullback-Leibler divergence between the network logits and the BP marginals instead of the accuracy. Bringing these two points together, we clearly understand why the k=0 case requires less samples than intermediate k. At the other end of the spectrum, it is understandable that the k=l case is the easiest, as it is implementable in a single feedforward layer, as it requires only a Naive Bayes classifier and not an implementation equivalent to full BP. The intermediate cases then appear more or less equivalent, as they require an implementation of BP while not being guided towards the correct marginals during training and not benefitting from the argmax which washes away approximations of the correct marginals in the accuracy.

(Q4) [...] what do the r’s represent? It looks like all r’s are initialized to uniform distribution, and they each get updated with the same formula (eqn 22), so would ri(a,m) ever be different from ri(a′,m) for a≠a’? We thank the reviewer for reading carefully our appendix, and spotting the absence of the base case of the recursion for the “r” messages (which we are now reporting). Moreover, we agree that the current explanation is insufficient to fully understand the role of “r”. We added a new paragraph, showing how the “r” recursion is obtained from the standard BP recursion, by conveniently playing with the traced indices.

(Q5) Refer to figure 3 instead of 1c on lines 337: We thank the reviewer for spotting this mistake. However, we meant to refer to figure 1b. We fixed this typo in the revised version.

References:

Behrens, F., Biggio, L., & Zdeborová, L. (2024). Understanding counting in small transformers: The interplay between attention and feed-forward layers. In ICML 2024 Workshop on Mechanistic Interpretability.

Zhong, Z., Liu, Z., Tegmark, M., & Andreas, J. (2024). The clock and the pizza: Two stories in mechanistic explanation of neural networks. Advances in Neural Information Processing Systems, 36.

To address the referee’s questions:

(Q1) Have authors considered measuring the match between BP and transformer predictions on individual inputs? Again, we thank the reviewer for encouraging us to push the comparison further. Please see answer to W2.

(Q2) Similar behavior […] doesn’t imply similar implementation. Why couldn’t [...] transformers […] behave like BP on individual inputs […] on out-of-sample data as well […] without necessarily implementing BP? We thank the reviewer for this key criticism, which was similarly raised by reviewer VwvV. Given the relevance of this point in the present work, we decided to substantially rewrite the related paragraphs, clarifying what we meant by “learning an implementation of BP” (now rephrased as “learning to approximate the exact inference computation”) and how we reached this conclusion. Our claims are based on the following evidence:

• We obtain approximately identical outputs (in root prediction and MLM) from equal inputs, both in- and out-of-sample. Note that the model is “trained to match BP” only indirectly (we train on the correct values of the masked symbols, not on reproducing the BP marginals, which could differ) and on the training distribution. If the match was accidental, and purely driven by data-fitting, the predictions would differ in the OOD case. For our purposes, obtaining the same output on any sequence implies an equivalence in the computation.
• We observe the same hierarchical computational structure underlying BP arises in the transformer (see attention maps), with a sequential focus on longer correlation lengths, corresponding to the various levels of the tree. Indeed, the transformer just learns to “model the data well”. But in this case, it does this so well that the correct sequence of operations (underlying the exact inference procedure) is discovered. The computation doesn't need to be done precisely in the same way as in a standard BP implementation (in fact, it is embedded in high-dimensional space, and different combinations of non-linear and linear operations are employed). Still, there is a fundamental equivalence in how single token information is collected and integrated to eventually obtain the exact predictions for the masked tokens.
• Performing probing experiments on ‘ablated’ transformers, we confirm what was hinted by attention maps, that is that removing the k last attention blocks in the architecture leads a transformer that was trained on the full data model to achieve a similar performance at predicting the l-kth ancestor as one that was trained on a data model with a level of filtration k. This strongly reinforces our claim that the transformer goes up the generative tree as the tokens are passed through its architecture, as it would do if following our tentative transformer-based implementation of BP in l layers.

We thank the reviewer for carefully reading our work and providing valuable feedback, which we will use to improve our presentation. We would first like to address the overall weaknesses identified by the reviewer:

(W1) It’s not clear how much the observations made in this work generalizes to larger models, and more complicated data distribution: Considering somewhat prototypical tasks is customary and often necessary when the goal is to achieve a mechanistic interpretation of the transformer’s computation (Zhong et al., 2024; Behrens et al., 2024), and we believe it can still provide intuitive explanations that generalize to more complicated cases. For example, in our work, we can argue that the sequential discovery of higher levels of hierarchical correlations, involving longer-range token interactions, is likely to shape learning processes also in real-world NLP tasks. We clarified this point in the revised version.

(W2) The empirical evidence presented has alternative interpretations that haven't been ruled out: We thank the reviewer for motivating us to follow up on their question with some new experiments, which will be included in the revised version. In particular, we compare the full output probability distributions of the transformer and BP (via a KL divergence in the main text, and the Spearman correlation coefficient and scatter plots in the appendix), as a function of the training epochs. Not only do we find that the model learns to approximate the exact BP marginals, but we also show that, at intermediate stages of learning, the transformer sequentially aligns its predictions with the filtered-BP marginals, recovering the same stair-case behavior observed in Fig. 1(c). This evidence reinforces our picture of the consecutive discovery of hierarchical correlation levels in transformers. We address the reviewer’s more specific questions on this topic below.

(W3) The construction for implementing BP on depth-l tree using a transformer of only l-layers could benefit from a bit more details, especially on the root-to-leaf message passing part: We thank the reviewer for reading carefully our appendix and taking interest in this more challenging part of our work. We agree that the current explanation is insufficient to fully understand the role of “r”. We added a new paragraph, showing how the “r” recursion is obtained from the standard BP recursion, by conveniently playing with the traced indices. We address the well-spotted missing part of our recursion on “r” in the answer to the questions Q3 below.

Have the authors adequately addressed your concerns about generalizability and rigor? Are the presentation changes enough to adjust your score? Rebuttal period is ending soon.

We thank the reviewer for engaging in the discussion. We would like the reviewer to clarify their point in concern C1, if possible, before we try to answer their further questions. It seems that the reviewer is stating that the transformer can:

• not only solve the tasks in section 3
• but obtain the same input-output mappings as BP, both on the training distribution and out-of-sample (even on random inputs) but with a completely different, unrelated computation.

Is there some known example of a similar scenario, where two unrelated non-linear functions end up producing the same continuous outputs on all inputs, without this match being enforced explicitly?

In our understanding, for example, in standard situations if you train two neural networks on the same data, in the end they might align their outputs on the training data distribution. However, they will still provide different outputs on random data.

Is the reviewer suggesting there exists a completely different way of performing exact inference on a tree? And more generally, what evidence would be sufficient to support the claim that an architecture implements algorithm X?

We thank again the reviewer, we will follow up and address the technical points raised above.