How Transformers Learn Structured Data: Insights From Hierarchical Filtering

We thank the reviewer for their valuable feedback, and address their comments and questions.

On the referee’s concerns towards the ‘surprise’ of our results: Let us clarify why we believe that our findings are not trivial. While one of the paper’s conclusions is indeed that, in the end, the transformer learns what it’s trained to learn, we think that the novelty and surprise lie in uncovering how it does that. In particular:

1. The model shows excellent generalization performance, pointing to the fact that it has learned the complex data model without overfitting the training set, which is not trivial.
2. It does so in a way that almost perfectly reproduces the output of the exact algorithm at the logits level, without ever being explicitly trained to do so, as it is only given hard labels in training and no calibrated supervision. This demonstrates that the model closely mimics the exact algorithm, even on entirely out-of-sample inputs.
3. It shows some very interpretable characteristics that are closely related to the natural implementation of Belief Propagation that we propose in the appendix, spontaneously organizing the computation in a hierarchical fashion, progressively going up the tree, which is not a priori required in our overparametrized context.
4. It progressively includes higher and higher levels in the hierarchy during training instead of e.g., suddenly ‘grokking’ to the optimal outcome after having seen enough examples.

Points (2), (3), and (4) go beyond what the model is trained to do via SGD, which just points at predicting the right class label or a masked symbol. We think that this highlights our contribution towards model interpretability and a genuine understanding of how and what the architecture learns beyond pure performance; as argued in the response to referee ivGY, an important aspect of our work is mechanistic interpretation. Note that points (3) and (4) also allow us to make contact with recent works on ‘simplicity bias’ in successful machine learning architectures; see also our answer to referee Rk87. Finally, in our context, we can also understand why, e.g. an insufficiently trained model would provide sub-optimal performance due to it incorporating only some of the spatial correlations in the data to make its prediction, which we believe is a rather rare occurrence in a complex data setting such as this one.

We appreciate the reviewer’s feedback and address here the weaknesses and questions they have raised.

• On the weakness about the lack of theoretical analysis: We are indeed unable to derive precise analytical results in our paper (as the complexity of our data models implies that we do not even have access to a closed-form expression of the distribution). However, we would argue that our setting does carry very significant theoretical justification and control to understand what is (or is not) being learned. Having access to the optimal oracle in the form of the full BP algorithm, as well as all its factorized counterparts, is indeed what allows us to show how the transformer progressively learns to solve the task during training, and that it does so with correctly calibrated predictions. Moreover, it is the knowledge of the exact inference algorithm that allows us to propose a plausible implementation within the transformer architecture and to verify that what is being learned in practice is compatible with it. As such, while we do agree with the referee that the type of structure we consider is indeed quite narrow, our choice is motivated by solid theoretical grounding. As argued in the response to referee Rk87, our work should be taken in the context of mechanistic interpretation, for which the state of the art is centered on simple tasks such as histogram counting. Our setting offers an important stepping stone toward the understanding of models dealing with complex data and lies closer to natural language processing, where transformers are ubiquitous but not well understood (even in simplified models of language, see the discussion on context-free grammars with reviewer Rk87).
• Clarity of the introduction: We are definitely interested in feedback in order to improve the clarity of our work towards a wider audience, and thank the referee for pointing it out. If possible, we would greatly appreciate precisions from the referee regarding which part of our introduction was particularly difficult to understand. We assume that it is with regards to results around Context Free Grammars (i.e. second paragraph), is this correct? If so, we would be glad to attempt a rewriting in the next iteration of reviews.
• On the possibility of working with data with other types of hierarchies: We agree with the referee that leveraging a similar type of analysis beyond fixed topology binary trees would be of great interest. As mentioned in our conclusion, we believe this is a clear direction for future work. Indeed, our current paper is, in our opinion, already quite dense for the 8-page format of ICML.
• On the link to broader structure learning: We thank the reviewer for pointing towards motif learning, which we were not familiar with. It indeed appears that our hierarchical model may be an interesting setting to explore this idea, as blocks of symbols of common ancestors can naturally be interpreted as motifs. Therefore, the fact that in our model learning essentially takes place through the identification of larger and larger clusters, to reconstruct higher and higher levels of the hierarchy, does support the idea that motif learning may facilitate sequence memorization for instance. We will add a mention to some references on the topic, such as Wu, S., Thalmann, M., & Schulz, E. (2023), in our conclusion for the next iteration of our paper.

We thank the reviewer for their feedback, and answer the points that they have raised.

• On the weakness point about the data being far from real-world one: We do agree that the data we used is far from, say, natural language. However, the complexity of real data strongly limits one’s understanding of what the model does given the lack of an objective ground truth. By putting ourselves in a simplified and controlled setting, we believe we uncovered a nontrivial way transformers can learn (see also the responses to reviewers Rk87 on context-free grammars and ivGY on the theoretical grounding of our work).
• Regarding the possible comparison with other models apart from BP: We agree that a comparison with other machine learning models is possible and interesting, but believe it would fall in the category of performance comparison. While there is no reason to believe that other architectures could not optimally solve the problem given enough data, the attention mechanism offers a significant advantage in terms of mechanistic interpretability, which is one of the central focuses of our work. Moreover, the transformer architecture has been chosen for its relevance, as it is ubiquitous in applications toward the analysis of sequences (text, amino acids…), and is known to be able to effectively implement algorithms. Note that, as mentioned in the response to reviewer Rk87, the works by Wyart and co-workers and Mei have demonstrated the ability of CNNs to implement belief propagation in similar hierarchical data models. However, the fact that the architecture and its convolutional filters mirror the tree structure of the data model limits the generality of their findings. Finally, note that the versions of BP obtained from the factorized graphs can be seen as other approximate algorithms that we compared the transformer to. In the root prediction task, for instance, the fully factorized BP corresponds to the well-known Naive Bayes estimator.
• On the reviewer’s comment about the term “structured data”: We in fact agree with them, although we had hoped that the subsequent precision towards “hierarchical filtering” clarified the more narrow context of our work. Nonetheless, we would be willing to change our title for e.g. ‘structured sequences’ if they believe that it better describes the scope of the paper.

We now answer the specific questions formulated by the referee:

1. Robustness across different random seeds and model weight initializations: We indeed did experiments with several different random seeds and we saw no qualitative differences. Given the large number of training epochs, we did not find particular sensitivity to learning rates and initializations, and did attempt efficient learning rate schedules without finding any significant differences. As our paper already includes a large number of figures, we did not think it was judicious to also include learning dynamics for different instances, but we are open to including them if the referee deems it necessary.
2. Using larger than needed number of layers: Experiments were carried out with up to 6 transformer layers (for $\ell = 4$), yielding the same qualitative results in terms of training dynamics and sample complexity. Taking $n_L = \ell$ though leads to the most interpretable attention maps, which can be seen in Fig. 4. There, we indeed show the attention maps resulting from models trained on $k$-filtered data models, that are trees with $\ell-k$ layers. As such, all intermediate cases in Fig. 4 fall into this category of larger-than-needed number of transformer layers. As shown in the maps, what occurs is that the computation may be "diluted" over more layers than necessary, while some attention maps remain close to unused due to the presence of skip connections in the architecture. Consider e.g. the second-to-last row of Fig. 4: the required mixing is carried out in the first two transformer layers, while the last two do not contribute.
3. Order dependence of the leaves: The hierarchical model producing the data is in fact fully order-dependent, as the parent-to-children production rules are described by a transition tensor that is not symmetric (with overwhelming probability). Our data therefore behaves like natural language, in which order counts. On the transformer side, we employ standard positional embedding to explicitly add this information to the representation of each leaf, allowing the architecture to accommodate for the order dependence of our model just like in natural language processing.

We thank the reviewer for carefully reading our work and providing valuable feedback.

On the efficiency of the BP implementation within the transformer and the MLP role: We had omitted to include an additional point in the Appendix, which is a precise proposition for performing the update of Eq. 22 through a two-layer fully connected network with $\mathcal{O}(q^3)$ hidden neurons. We have added back the explanation in the revised version, which we cannot yet reupload. Unfortunately, the character limit prevents us from detailing the construction here, but should the referee request it we could add it in our next reply.

We now answer their two questions directly.

1. The reviewer is correct in pointing out that here BP is equivalent to a simplification of the inside-outside algorithm (IO) for context-free grammars (CFGs) on a fixed topology, and is not more efficient per se. We argue that this is a feature of our setting, as this difference importantly leads the algorithmic complexity of optimal inference to be linear in the sequence length, whereas it is cubic for the IO. As pointed out in Khalighinejad & al. (2023), this cubic complexity means that there must be some approximation in the transformer implementation of the IO, self-attention having a complexity that is only quadratic in the sequence length, or the network depth must be scaled with the context length (see the implementation proposed in Zhao et al. (2023)). Therefore, the evidence brought forth by Zhao et al. (2023) and Allen-Zhu & Li (2023) points toward transformers learning something very close to the IO for CFGs, but there is still a major open question as to what they do precisely. Relative to CFGs, the fixed topology is advantageous as it allows us to understand more precisely how transformers closely align with the exact inference algorithm, notably by leveraging the filtering procedure described in our work, which is not easily generalizable to CFGs. While we agree that our assumption is not as realistic in practice, it thus allows us to go much further in the mechanistic interpretation of the transformer. On the other hand, we believe our setting is at least as realistic as many mechanistic interpretation works (considering, e.g., histogram counting on integers).
2. We believe that our work is at the crossroads between different bodies of literature. First, as described above, it studies a complex yet completely controlled task that allows mechanistic interpretation. Second, it builds upon the body of work of Wyart et al. on hierarchical models with, in our opinion, a central improvement through the introduction of the hierarchical filtering procedure. This filtering uniquely allows us to study the learning dynamics and therefore to also understand how insufficiently trained architectures might fail—here using an incomplete correlation structure. In doing so, it allows us to make contact with the expanding literature on staircases in learning dynamics and simplicity biases in machine learning architectures (Refinetti et al., 2023; Rende et al., 2024; Bardone & Goldt, 2024). In a nutshell, we believe that our most significant scientific contribution is to provide a truly comprehensive study of how a markedly non-trivial task, which shares similarities with practical natural language processing problems, is implemented in transformers.

Finally, on their comments:

the authors should clarify (...) new insights beyond previous CNN-based results (…)

In these CNN architectures, the mechanistic interpretation is somewhat trivial, as the convolution filters are made to mirror the tree structure. This is not the case in the transformer architecture, which has to learn this structure. We demonstrate that it implements it incrementally through its attention layers and progressively in training. Transformers are also ubiquitous for sequence-like data (text, amino acids…) whereas CNNs are seldom employed in this context, highlighting the importance of understanding how transformers digest long-range correlations.

Cagnetta & Wyart (2024) also study (...) theoretically the progressive learning of hierarchical correlations with transformers (...) while the present submission empirically studies the problem (...)

While the study of Cagnetta & Wyart (2024) is very interesting, we would highlight that, as it relies on a signal-to-noise ratio analysis that is tractable for their uniform transitions, it does not provide predictions for any specific architecture. While more empirical in spirit, we would again emphasize that our filtering strategy allows one to probe the behaviour of transformers, be it as a function of the sample complexity or of the training time. We therefore believe that the two approaches are complementary and not redundant.