How Transformers Learn Structured Data: Insights From Hierarchical Filtering
Authors: Jerome Garnier-Brun
Marc Mezard
Emanuele Moscato
Luca Saglietti
Venue: ICML 2025 poster (2025)
Keywords: Transformers
Belief Propagation
mechanistic explanation
structured data
hierarchical data model
attention
masked language modeling
Primary area: deep_learning->attention_mechanisms

Abstract: Understanding the learning process and the embedded computation in transformers is becoming a central goal for the development of interpretable AI. In the present study, we introduce a hierarchical filtering procedure for data models of sequences on trees, allowing us to hand-tune the range of positional correlations in the data. Leveraging this controlled setting, we provide evidence that vanilla encoder-only transformers can approximate the exact inference algorithm when trained on root classification and masked language modeling tasks, and study *how* this computation is discovered and implemented. We find that correlations at larger distances, corresponding to increasing layers of the hierarchy, are sequentially included by the network during training. 
By comparing attention maps from models trained with varying degrees of filtering and by probing the different encoder levels, we find clear evidence of a reconstruction of correlations on successive length scales corresponding to the various levels of the hierarchy, which we relate to a plausible implementation of the exact inference algorithm within the same architecture.

How Transformers Learn Structured Data: Insights From Hierarchical Filtering
Jérôme Garnier-Brun' 12
Mare Mézard
Emanuele Moscato
Luca Saglietti
Abstract
Due to the complexity of the standard transformer archi- tecture (Vaswani et al. 2017), understanding what strategy is precisely implemented via the attention mechanism to solve a given problem has been limited so far to very simple tasks (Weiss et al. 202 I Zhong et al., 2024] Behrens et al.] 2024). Nonetheless, significant results have been obtained by studying transformers on simplified models of language known as Context-Free Grammars (CFGs). Through prob- ing of the so-called parsing tree of CFGs, evidence has notably pointed towards transformers trained on predict- ing masked symbols implementing the optimal dynamic programming algorithm to reconstruct the hidden structure of the grammar, but alas without finding a fully plausible implementation within the architecture (Zhao et al. 2023 Allen-Zhu & Li 2023). On the other hand, when tasked with reconstructing the most probable parsing tree in the context of probabilistic CFGs, transformers may struggle to match the optimal algorithm if ambiguity is high (Kha-) lighinejad et al. 2023).
Understanding the learning process and the em- bedded computation in transformers is becom- ing a central goal for the development of inter- pretable AI. In the present study, we introduce a hicrarchical filtering procedure for data mod- els of sequences on trees, allowing us to hand- tune the range of positional correlations in the data. Leveraging this controlled setting, pro- vide evidence that vanilla encoder-only trans- formers can approximate the exact inference al- gorithm when trained on root classification and masked language modeling tasks, and study how this computation is discovered and implemented. We find that correlations at larger distances, cor- responding increasing layers of the hierarchy, are sequentially included by the network dur- ing training. By comparing attention maps from models trained with varying degrees of filtering and by probing the different encoder levels, we find clear evidence of a reconstruction of corre- lations on successive length scales corresponding to the various levels of the hierarchy, which we relate to a plausible implementation of the exact inference algorithm within the same architecture.
Beyond language models, the significance of data structure in machine learning applications is well recognized yet re- mains poorly understood. CFGs represent data structure characterized by hierarchical correlations (Mossel) (016)) In general, understanding how standard deep networks can take advantage of this hierarchical structure in their train- ing is an important research question. Towards this objec- tive, simplified hierarchical models of structured data on fixed trees have proved very useful in understanding the efffcctiveness of Convolutional Neural Networks (CNNs) (Cagnetta et al.] (2024), for which there are now formal results supporting the idea that the optimal Belief Propa- gation (BP) algorithm can be approximately implemented (Meil 2024). Unfortunately, while the implementation of the hierarchy in CNNs is made quite transparent by the hi- erarchical structure of their convolutional filters, this not true for transformers, and one can therefore not straightfor- wardly transpose this interpretation to other architectures (Cagnetta & Wyart 2024).
1. Introduction
Transformer- based large language models have revolution- ized natural language processing, and have notably demon- strated their capacity perfectly assimilate the grammati- cal rules of the languages they are trained on. While this ev- idence shows that transformers can handle and exploit the subtle long-range correlations that emerge in natural lan- guage, their inner workings remain largely unclear.
DDepartment of Computing Sciences, Università Boc- coni, Milan, Italy 'Department of Finance. Università Boc- coni, Milan, Italy. Correspondence to: Jérôme Garnier- jjerome.garnieree garnier@ unibocconi.ts>.
In this work, we present a complementary study to those described above, which allows us to understand further how transformers approach optimal inference in a struc- tured data model.
<jerome.
Proceedings of the 1nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

How Transformers Learn Structured Data: Insights From Hierarchical Filtering
(b)
(c)
(d)
k
In-sample
Transfrmme outpot
1.0-
10c
0.5
109
k
0.0
Det.
10-1
Din.
10-2
Random
Trinsformat output
1.0
10-2
10-2
10
10
10
102
108
10
10a
102
0.0
Epoch
Epochs
BP marginals
(f)
1. 1.00
Tost accuracy
k
0.75
MLM
Apope
0,50%
Layer 1
Layer 2
Layer
Layer
10
10a
Figure I. Synthesis of our main results. (a) The proposed filtered hicrarchical model, illustrated here with = 3 layers and with a truncate the hierarchy and generate data with more less structure. Scatter plot of the predictions of a trained transformer for a masked symbol (E = 4, k = 0, = 4 possible states) versus the corresponding exact marginals obtained with the BP oracle, in-sample on 10* sequences (top), and out-of-sample on uniformly generated sequenees (bottom). (c) Evolution along training, on root classification task with P = 2n examples (f = 4, k = 0, q 4) of the average Kullback-I.eibe divergence between transformer predictions and marginals obtained from matched (black) and mismatched (from light green k = to purple & = 4) identical in sample inputs. demonstrating the transformer learns increasingly structured representations, (d) Identical to (c) a MLM task on 2 data. (e) Attention maps averaged over 10* in-sample inputs, for transformer with my f = 1 layers of attention trained on the MLM task with fully hierarchical data, exhibiting structure that mirrors the organization of generative tree the sequence operations of BP. (f) Test accuracy on root classification on fully hierarchical data = 4, k = 0, 4) versus number of labeled training samples P with no pretraining (o) compared to MLM pretraining with frozen (D) and unfrozen (+) encoder weights during fine-tuning.
filtering parameter 0 I l: allowing
Our contributions. We propose a controlled hierarchical model of discrete sequences, which we can easily tune the strength of correlations between tokens thanks to "fil- tering" parameter kt illustrated in Fig. T[a), (). This tree- probabilistic graphical model gives us access to the exact inference algorithm for reconstructing any symbol the tree, Belief Propagation (BP) (Mézard & Montanari 2009).
simplified setting allows to understand how trans- formers learn from structured data time.
Well-trained transformers reconstruct the correct hi- erarchical structure through the succession of atten- tion blocks. Matching the number of transformer lay- ers to the number of layers in the generative tree, find that the attention maps compatible with natural implementation of BP within the architecture, see Fig. [Te.. We verify this affinity through prob- ing experiments, providing strong clues on how trans- formers learn from our structured data in "space", thereby explaining the effectiveness of unsupervised pre-training for supervised classification tasks, illus- trated in Fig. ft).
Leveraging this context, we show that
Transformers not only approach optimal performance in root classification and Mask Language Modeling (MLM) tasks, but they spontaneously do so cal- ibrated way-i.e. by predicting probabilities that ap- proximate those yielded by the BP oracle even on out- of-sample inputs, Fig. [b) provides evi- dence of equivalence in computation to the exact inference algorithm.
The paper is organized as follows. First, we provide detailed description of our tunable hierarchical model in Sec. 2 We then perform numerical experiments stan- dard transformer architectures in Sec. shedding light on the learning dynamics. The understanding of the imple- mentation learned by the transformer, and its compatibility with a possible implementation of the BP algorithm in the architecture that we propose, analyzed in-depth in Sec.
When trained with stochastic gradient descent, trans- formers sequentially discover the existence of higher hierarchical correlation levels (i.c., longer-range cor- relations), progressively aligning with the prediction of algorithms that impute only parts of the full corre- lation structure, see Fig. [c)- other words, our
2

---

How Transformers Learn Structured Data: Insights From Hierarchical Filtering
We finally conclude and diseuss the wider implications of results in Sec. 5
pairs from their respective parents and are strongly cor- related. When k > O, we instead generate the tree by drawing the children at level & conditionally independently given the root, with the same marginals as the full (k = 0) model. Then, for layers below layer k. the generative pro- cess is the standard one described above, inducing correla- tions within blocks of 2 tokens. The procedure is illus- trated in Fig. Ta), where dashed segments indicate condi- tional independence.
2. A model with filtered correlations
2.1. The full hierarchical model
We consider a tree-based generative process producing structured sequences of discrete symbols. We here focus on the fixed tree topology case, allowing for direct control over the effective range of the hierarchical correlations in- duced in the generated sequences ) and enabling exact and efficient inference through Belief Propagation (2.4).
In order to match the correct marginal probabilities in the truncated models, the conditional independent sampling at level & is done as follows. For each of the 2 variables at level k say 2, one considers the unique path that relates the root to this intermediate child in the original fully hier- archical tree, yielding probability
The "full" hierarchical generative process shown in the first row of Fig Ta) can be described as follows. The chain starts from an initial symbol wo, which we will refer to as the root of the tree, sampled with probability po from a vocabulary X' (1,....9).. Then, the first layer of the tree is drawn randomly using transition tensor M, which assigns the probability of generating some children-from the same vocabulary X' a parent (here zo). In this work, we will restrict ourselves to binary trees for simplic- ity. We therefore have M E R RTxaxq with Mate the proba- bility of generating the pair (b, c) given a parent a. Since its elements are transition probabílities, this tensor should satisfy Made [o, 1] V a, b, c and Exe Mabe 1 V The process, with the same tensor M, is then repeated indepen- dently for each of the newly created children nodes for a total of e generations, eventually yielding sequence of 2 symbols {- {==.i=1.....2. We will refer to the symbols in the sequence as the leaves of the generative tree.
P (zy = b xo = a) Mpo M.. ..M
a.b
with om (i) E {L, R)) indicating whether the path leading to the tree element considered at layer k takes a left or right branching at the previous layer m. The x tran- sition matrices M and M R are computed by tracing the original tensor
M
Σ Maber
MR R uc
Σ Mabe-
(2)
By constructing filtered trees in such a way, we ensure that the conditional correlations of the leaves capture up the hhhh level of the hierarchy. Note, however, that when k > 0 the root can no longer be recovered deterministically from the leaves.
The class of transition tensors M that we use is defined pre- cisely in Appendix A In short, will resort to randomly sampled log-normal transition probabilities, yielding com- plex long-range correlations along the sequences. Impor- tantly, we will only consider tensors with non-overlapping entries, such that: if Make > 0, then Va' * a Ma = O. As result, the production rules of our unfiltered generative model are non-ambiguous in the sense that pair of chil- dren symbols can only have single parent. Given all the symbols on the leaves, one can therefore deterministically reconstruct the underlying generative tree, all the way up to the root.
2.3. Related data models
Context-free grammars. Our hierarchical model can be considered as an instance of a simplified probabilis- tic context-free grammar (PCFG) with log-normally dis- tributed transition rates (De Giuli 2019). The simplifica- tion is two-fold. Standard CFGs typically include two dis- tinct sets of symbols, non-terminals and terminals, repre- senting parts of speech- -.e. nouns, verbs etc. and actual words respectively, plus root symbol. Here, instead, we consider single vocabulary X' for all the symbols in the tree, including the root-which allows us to define root classification task. Moreover, the parsing trees underlying CFGs are not fixed: terminals can be produced at differ- ent levels and the sequence length can vary. Instead, we assume a fixed parsing tree for our model, where the 24 leaves are collected from the last waycr-which allows us to define a filtering procedure based on removing layers of hidden symbols above the leaves.
2.2. Filtering hierarchical correlations
We develop a filtering tool that enables control over the cor- relation structure in the generated sequences. In particular, we consider a family of generative models, indexed by an integer k < f, with hierarchical correlations truncated at a given depth k of the tree.
In the k= = 0 case described in the previous paragraph, all children generated at any level of the tree are sampled in
'eere we take j > 2 to refer to the internal nodes of the tree, while zo remains the root and 2; with = 1,..., 2 are the leaves.
3

---

How Transformers Learn Structured Data: Insights From Hierarchical Filtering
The Random Hierarchy Model. Our model is closely related the recently introduced Random Hierarchy Model (RHM) of of'Cagntta et al. (2024), which was studied to improve the understanding of the effect of hierarchical structures on generative diffusion (Sclocchi et al. 2025) or last token prediction (Cagnetta & Wyart 2024). The main differences to our formulation are that in the RHM the al- lowed transitions have uniform transition while we consider log-normal distribution- and that the produc- tion rules depend on the layer-while we here consider single transition tensor throughout the tree. Correlations between the leaves arise in the RHM when some children pairs cannot be produced, leading to a reduced entropy of viable sequences. Having non-uniform transitions in our model similarly limits the entropy, while leading to sig- nificantly different correlation structure. One should for in- stance notice that the staircase decrease of the correlations as a function of the distance between leaves presented in Cagnetta & Wyart (202) is not visible in our case.
(I,..., . is first converted to a positionally-informed to- ken æ; E R2. For our experiments, we consider d = 128 and the standard sinusoidal positional encoding of (Vaswani et al. 2017). Each transformer block in the net- work then maps the previous encoded sequence onto a new sequence of tokens with the same length and embedding di- mension, through concatenation of self-attention layer and a fully connected layer, with residual connections and layer normalization. The self-attention layer importantly introduces some mixing between the different tokens in the sequence, represented by what we will refer to as an at- tention matrix A E R We take the fully connected layer to be a standard 2-layer network with relu activations and hidden dimension d' = 2048. Following these opera- tions, repeated ny times to obtain the full encoder, we ob- tain a position-dependen high-dimensional representation of each of the original symbols the sequence. What is finally done with this sequence of tokens depends on the task at hand: we consider root classification in Sec. 3.2 and masked language modeling in Sec. 3.3
2.4. Exact inference
Motivated by our focus on understanding the transformer's implementation, we will take the number of attention layers to match the depth of the unfiltered generative tree, n,, = E. Studying varying values of & for the training data will ef- fectively allow us explore cases where there are more at- tention layers than hierarchical levels in the generative tree, while we discuss the consequences of having ny smaller than the number of hierarchical levels in Appendix F.
A key advantage of generating sequences through tree- based process is that we can perform exact inference ef- ficiently using a dynamic programming approach. More- over, the fixed tree topology allows us to consider simpli- fied version of the general inside-outside algorithm (Baker) 1979), which can be written in a message-passing form within the Belief Propagation (BP) formalism (Sato) 2007: Mézard & Montanari 2009). Assuming that the transi- tion tensor M and root probabilities Po are known, with BP one can compute the exact marginal probabilities for all the symbols at any position in the tree, with a computa- tional cost linear in the size of the tree. More precisely, on the tree structures we consider, BP can be shown to con- verge in 2(e k + 1) steps i.e. an upwards and downwards pass along the tree. This procedure can be used to infer the root given the leaves and to find the most likely value of a masked leaf (or set of leaves) given the rest of the se- quence: we'll use it as an optimal solution against which to compare our numerical results. The details on the BP scheme for the filtered tree graphs we are considering can be found in Appendix B
In the following, all numerical experiments are performed on the same realization of the transition tensor, randomly sampled for q 4 using the parametrization described in Appendix A (see also our Reproducibility Statement be- low). While there may be quantitative differences for dif- ferent randomly generated tensors- partinuaarti at small q results remain qualitatively unchanged in experiments on different grammars, see Appendix E.2
3.2. Supervised classification
In the context of our model, a natural idea is to use the root of a tree xo as a label for the generated sequence { and to train a transformer encoder architecture on the associated classification task using dataset of P labeled sequences. To perform the root prediction, the tokens in the final layer are concatenated position -iise (forming a large d x 2 vec- tor) and fed to linear readout, which outputs g logits as- sociated with the possible root symbols. The network is trained by minimizing the cross-entropp loss between these logits and the correct one-hot encoding of the root.
3. How transformers learn to climb the hierarchy in time
3.1. Experimental setup
We will focus on the encoder-only variant (Devlin et al.)
2019) of the celebrated "vanilla" transformer architecture, introduced in (Vaswani et al. 2017). A full recap of this
parametrization is given in Appendix C
Optimal test accuracy. We find that given sufficient la- beled data P 2 P", transformers achieve perfect in sample root classification accuracy in the fully hierarchical model,
In a nutshell, each of the sequence elements 24 E
4

---

How Transformers Learn Structured Data: Insights From Hierarchical Filtering
1.0
1.0
Full prediction matching. So far, we have established that the trained transformers match the accuracy of the ex- act inference algorithm on the root prediction in- and out- of-sample. We can however go one step further, as the transformers output logits, which were passed through an arg max operation to yield a prediction. Taking the sofumax instead gives a normalized q-dimensional vec- tor, which we can interpret as the predicted probabilities of the root symbol given the input sequence, to be com- pared to the exact marginals obtained with BP. We find a close match at the end of training, as shown by the small Kullback-Leibler divergences averaged over in-sample in- puts in the h = O case in Fig. [e), and similarly for à 2 O. on both in-sample and entirely out-of-sample in- puts in Fig. of the Appendix. While such a match is not entirely surprising in the deterministic k = problem, as the one-hot encoding of the root label against which the transformer logits are compared at training corresponds to the exact marginal distribution yielded by BP the match is highly non-trivial in the ambiguous A: > 0 instances, where the transformer is never explicitly guided towards the cor- rect values during training, as the one-hot encoding of the root label does not correspond the exact marginals any- more. This calibration therefore provides a second strong piece of evidence that the transformers spontaneously im- plement exact inference.
(b)
Test accuracy
0.8
0.8
0.6
0.6
0.4
0.2
10a
10
10
0
10
10'
10a
Epoch
Figure 2. (a) Evolution of the root prediction accuracy on full hi- erarchical kieu = O test samples for transformers trained on P labeled samples generated with knain 0, 1, 2, 4 (top to bot- tom). Dashed lines indicate, the accuracy computed with the BP algorithm on unfiltered data. (h) Evolution of the root prediction accuracy of the Kuain 0 model computed on filtered test datasets. with Reest 0, 1, 2, 3, (top to bottom), for transformers trained on P = 2d Kuain = 0 data. Dashed lines represent the out-of- sample BP, prediction, both plots = 1, my=4
k = 0, as illustrated in Fig. 2aa). When the training data has filtering parameter & > 0, the networks approach the optimal in-sample accuracy predicted by BP.. see Fig. g of Appendix E.3] Notice that, while in the case k- the exact algorithm finds the value of the root with accuracy 1. this is no longer the case for k 2 where the optimal accuracy is < 1.
Supervised learning dynamics. Looking more specifi- cally at the learning dynamics of a network trained on the full hierarchy sheds some light on the learning process of the transformer encoder. Fig. 2(b shows the evolution of the test accuracy of the kirai 0 model both in- in-samlle with kiess data, and out-of-sample, on filtered data with Mest O. One can notice multiple stages in the learning procedure: in the first epochs, the network imputes a sim- plistic explanation of the training data, resolving the leaf- to-root correlations-aidee by the supervised signal-, as well as the short-range correlations between the leaves. As result, the test accuracy increases for all values of keat As time progresses and longer-range correlations are dis- covered in the training data, the accuracy on the most fil- tered datasets drops towards the mismatched BP, predic- tion, since the imputed higher correlation levels are not present in the out-of-sample kiest > 0 data. In the mean- time, the accuracy for the smallest values of hiest keeps in- creasing- In a limited number of epochs, as the network perfectly learns to infer the root on kisst data, the BP, oracle accuracy is reached test sets generated with all levels of factorization.
Different from the Random Hierarchy Model of Cagnetta et al. (2024), characterizing analytically the scaling of px with the parameters of the grammar with our non-uniform transition probabilities a challenging goal, and is left for future work. Still, discuss the role of the filtering pa- rameter k of the data model on the sample complexity in Appendix E.
Out-of-sample testing. In our data model, one can also test out-of-sample with respect to the filtering parameter k. For example, we test models trained on intermediate fil- tered data on a fully hierarchical dataset, i.e., Aminn > 0 and kest = 0. in Fig. 2(a)), or vice-versa, i.c., Arain = O and kiss > 0, in Fig. (b).. In both cases, the transformers achieve a performance that exactly matches that of BP in the presence of the same mismatch between the assumed inference model and the data generative model. We stress that, in this mismatched task, the BP predictii is no longer optimal, yet the trained networks systematically reach the same accuracy. This observation provides the first evidence that the transformers are implementing an approximation of the BP Br... algorithm matched to the training data distri- bution.
This picture can be further refined by considering the pre- dictions of a transformer trained on the full hierarchy and the evolution of their distance from the marginals predicted on the same data by the BPg oracles, for all k 2 As
5

---

How Transformers Learn Structured Data: Insights From Hierarchical Filtering
illustrated by the DKL in Fig. [c), we observe an ini- tial stronger alignment to BPg, which only considers leaf- to-root correlations. As training kimain = 0 data pro- gresses and the transformer shifts towards the correct pre- diction, the model predictions sequentially align to versions of BP that incorporate more and more of the correlation i..., BP& with decreasing values of k.
1.0
(a)
(b)
Test accuracy
0.6
0.8
0.4
0.6
-
2
0.4
3.3. Masked Language Modeling
0 100
10
102
102
We now turn to self-supervised training, where the model learns from a dataset of P unlabeled sequences. In sim- ple terms, the Masked Language Modeling (MLM) train- ing procedure consists of randomly masking parts of the sequences and asking the model to recover them from the context. This is closer to what is done in practice to train large language models, see c.g. Devlin et al. (2019); Liu et al. (2019). While in principle one could mask several symbols simultaneously in training, we focus on single- symbol masking- at a random position in the sequence- in the following, given the limited length of our sequences (a single symbol representing already 6.25% of the se- quence for L = 4). Contrary to the root inference task, in MLM perfect accuracy cannot be achieved even in the fully hierarchical case, because of the stochastic nature of the branching process in the generative tree. The optimal performance is still yielded by the BP matched to the test data.
Epoch
Tree layer
Figure 3. (a) Evolution of the MLM test accuracy computed on filtered test datasets, with kiss = 0, 2, 3,4 4 (from top to bot- tom), for a model trained on kiain = data and P 2 The dashed lines represent the in- and out-of-sample BPs predictions. (b) Test accuracy in the ancestor prediction task (layer 0 is the root) with Ruain = hise = 0 obtained by reading out the interme- diate transformer encoding levels (legend) of model pre-trained on the full hierarchy. The readout is trained on 2 labeled exam- ples. both plots = 4, =4 4.
transformer. As shown in the top panel Fig. [b), we find close match with the exact marginals obtained from BP when measured on in-sample inputs. To confirm the gener- ality of this correspondence, we extend the comparison to uniformly sampled data in the bottom panel of Fig. [bb. In this setting, we still observe high correlations between the outputs, albeit with more dispersion related to the markedly atypical nature of these test samples compared to the train- ing data distribution. Measuring the alignment using the Kullback-Leibler divergence, shown in Fig. [d),, or else the sample-specific prediction match and Spearman (rank- ing) correlation between the two discrete probability dis- tributions, shown in Fig. I0 of Appendix [E.4] confirms the near equivalence between transformer and BP compu- tation. Note again the remarkable calibration of the log- its, although the network is trained with hard labels for the masked symbols despite the probabilistic nature of the task.
To reconstruct the masked symbol, now feed single token, selected from the final transformer encoding at the positions associated with the masked element, to a linear layer producing a vector of logits. The network is then trained by minimizing the cross-entropy loss between these logits and the one-hot encoding of the masked element in the sequence.
Optimal reconstruction performance. Given sufficient data, find that transformers again approach optimal in- sample accuracy on data with any level of filtering. We 0 in Fig. Baa), where
show the case trained on kimin the transformer reaches the BP, accuracy also on out-of- sample test data with hiess > 0. Consistent with intuition, the required amount of training data pa is increased rel- ative to the supervised task, as the network must learn to resolve the weak long-range correlations in the sequence without any supervised signal from the top of the hierar- chy. Moreover, compared to root classification, the net- works trained for MLM require much longer training to approach optimal
Self-supervised learning dynamics. By analyzing the out-of-sample performance with different filtering levels, we also unveil the sequential nature of the MLM learn- ing process. Computing the test accuracy on all keest lev- els throughout the training dynamics, we observe a clean "staircase" behavior in the test accuracy, as shown in Fig. B(a). This picture confirms and clarifies the experi- ments in Fig. bb), showing that the network sequentially resolves the nested levels of the hierarchy, in bottom-up order. Note that the observation of the shorter-range corre- lations being learned first is consistent with the signal-to- noise picture exposed in Cagnetta & Wyart] (2024). More- over, the presence of a sequential mechanism of discovery and resolution of different moments of the data distribution
ceterrannnettppccallle ~ 103 epochs in place of a mere ~ 10 epochs for classification see Fig. a) vs Fig. bb)
Full prediction matching. To go beyond test accuracy, we also consider the full probabilities outputted by the
6

---

How Transformers Learn Structured Data: Insights From Hierarchical Filtering
has been studied in (Refinetti et al. 2023 Bardooee & Goldr 2024
Rende et al. 2024). Overall, the convergence of the transformer to both the in-sample and the out-of-sample token prediction accuracy of BP supports the claim that the model learns to implement a close approximation of the exact algorithm. The learning mechanism is also con- firmed by the behavior of DKI. along the training, shown in Fig. d)) analogous to the root inference case, but more qualitatively compelling, the predictions of transformer trained on the fully hierarchical data sequentially align with the marginals yielded by BPk, with decreasing k as training progresses and longer-range correlations are accounted for.
4. How transformers embed the exact inference computation
Attention map analysis. In the root inference task, the readout performing the prediction is fed with the entire se- quence of tokens. As result, there are many ways for the transformer encoder to distribute the computation across its layers, and no necessity for single tokens carry informa- tion on all the ancestry levels in the tree, making it non- ideal setting for mechanistic interpretation In the MLM task, on the other hand, single token encodings are used to predict the masked symbols. This requirement seems to guide the model towards more interpretable attention maps, shedding some light on how the model may approximate the optimal algorithm. They are shown in Fig. 4 each row referring to transformer encoder trained on data with dif- ferent filtering levels - increasing from top to bottom.
Figure 4. Visualization of the nı = 4 attention matrices (aver- aged over 10a input sequenees) for transformers trained the MLM task on different filtered datasets. with k = 0, 2, 3,4 (top to bottom rows). and fully factorized model, A 4, where the leaves independent conditional to the root attention matrix appears structureless. When & decreases size f~ 2 For = 0, 1, the trained attention matrices reflect all the hierarchies of the correlations.
2
q
For the
the emergence of attention blocks of
BP implementation includes downward pass, it is possi- ble to avoid this step if the token embedding dimension, d, is sufficiently large. To illustrate this point, we propose an existence proof of a plausible implementation of the BP a algorithm an architecture.
In the fully filtered case (bottom row) there is no need to combine the different elements of the sequence before the readout and the attention matrices are nearly uniform. Now, as we reduce the level of filtering in the generative process, clear patterns emerge in the attention map.
First, the model focuses short-ranged correlations be- tween nearest neighbors when &o = 3 and, as we decrease k, towards patterns of size ~ 2é-k, which is the exact size of the stronger correlated block with a filtering parameter k-see Sec. _ Note that similarity between the & = and ls cases (top two rows) natural, the tree topol- ogy in these cases being identical and with only the transition probabilities for this first layer differing.
Exact transformer embedding of BP. In a natural im- plementation of BP, inference for the MLM task requires the messages from the visible leaves to reach the top the hierarchy and descend back the masked symbol, effec- tively propagating through 2/ layers. A proposal in /hao et al. (023) for transformer embedding of the inside- outside parsing algorithm-4 generalization of the above- described BP to the unknown topology setting-requree as many transformer blocks as double the sequence length- here 2-_ and an attention head per hidden symbol the hierarchy. Thus, might seem surprising that single- head transformer encoder with é blocks could be sufficient to mimic the BP algorithm. To prove the feasibility of its implementation within these architectural constraints, we propose an idealized transformer implementation of the BP algorithm. Note that some of the key ingredients of this feasible implementation are introduced for the sake in- terpretability but are not imposed in our experiments, and
Interestingly, the network naturally organizes the attention layers hierarchically. This particularly visible when there are fewer redundant layers i.e. in the cases A = 0, (two top rows in Fig. 4. Such a layout is consistent with the BP algorithm on the full tree, where one combines ele- ments pairwise while going up the tree. While typical
Note that transformers trained on classification still present some patterns related to the hierarchical nature of the data model, alheit less clearly, see Appendix E.5
7

---

How Transformers Learn Structured Data: Insights From Hierarchical Filtering
therefore this does not represent an exact explanation of the trained transformer computation. The complete existence argument is deferred to Appendix [ while here we provide a high-level description of some key ideas.
nario, where the ancestors are reconstructed solely by the trained probes and the sequential reconstruction is tifact, can be ruled out by considering the gap between the accuracies achieved from different layers- the relative comparisons are fair since the readouts are trained on the same datasets-, and by training the probes only on some positions- see Appendix E.7
We consider a fully disentangled embedding of positional and semantic information in the vectorized tokens, con- tained in d q(9 + 2) + l dimensions. The isolation of the semantic information allows the implementation of simple position-based attention mechanism, inspired by the factor graph structure, and compatible with the attention matrices in Fig. 4 Then, going up the hierarchy requires the computation of a trace of products (see equation ( in Appendix
In Appendix E.7 we also conduct similar ancestor predic- tion experiments the last encoder layer of models trained with k > 0 data (lower rows of Fig. 4 where we again find that the ancestry information is consistent with the attention maps.
B which can be well approximated by the fully connected layers in the second part of the transformer blocks, provided the attention selects the right terms in the product. The less intuitive component of the implementa- tion is the computation of the messages directed towards the leaves, used in the MLM task. Given the limit on the number of transformer blocks, this computation must be done in parallel with the upward climb of the hierarchy, de- spite the missing downward messages. turns out that, by exploiting Ogg²) memory slots in the token embedding- and thus with an increased memory cost compared to BP- different recursion with the same result the standard message- paaagggg can be implemented, within the nı = L constraint for the number of transformers layers.
Synergy between tasks and MLM pre-training. In the context of our model, we can straightforwardly explain why self-supervised pre-training allows a large speed-up in the supervised training process, in line with many em- pirical observations on real-world data (Howard & Ruder 2018). We show in Fig. M MLM pre-trained model fine-tuned for root inference. A significant reduction in the labeled data required to achieve optimal root inference Pa in Sec. 3.2 is observed, both with frozen and with fine-tuned encoder weights.
5. Conclusions
By using a simple, tunable, hierarchical model of struc- tured sequences, we were able to shed some light on the inner workings of transformer encoders and better under- stand how they achieve optimal inference on both super- vised and self-supervised tasks. The modularity of our data model also allowed us to uncover how transform- ers sequentially implement longer-range correlations dur- ing the learning dynamics, compatible with similar con- trolled studies (Rende et al. (2024) and with the general un- derstanding of LLMs trained on natural language (Kaplan et al. 2020). This mechanism could perhaps be exploited to shape theory-driven curriculum learning strategies for NLP, where curating the presentation order of training examples was already proven effective (Campos) 202T). Moreover, because blocks of symbols inherited from common ances- tors are progressively integrated during training, learning our data model may perhaps be related to a form of motif learning, with increasingly longer motifs being identified over time (Wu et al. 2023
Probing the encoder representations. To confirm that the computation going up the tree is distributed sequen- tially in the transformer blocks, consistent with the pro- posed embedding of BP, we undertake a probing experi- ment similar to those performed e.g. in Zhao al. (2023). First, analyze the encoder trained for the MLM task on k- 0 data, ef. top row of Fig. 4 Keeping the en- coder weights about the ancestors of any leaf is contained in the succes- sive hidden representations of the corresponding token- see Appendix E.6 for implementation details. While in the exact embedding of BP the k-th level ancestor information must be available at layer A iterate the recursion for the downgoing messages, the MLM training does not set such a requirement. To probe the encodings, we employ a spe- cialized two-layer readout for each encoderr-ayyer/mccetty level pair-inie pair-indeppedeet of the token position trained on a supervised dataset with 214 examples. In Fig. 3b) we show that the prediction accuracy is high on ancestors up to the same level as the probed layer and deteriorates on higher levels of ancestry. Note that, unless the information about the entire block of 2 tokens is properly mixed in through the attention mechanism, a perfectly accurate prediction of the common 1.th level ancestor from a sin- gle token representation is impossible, as the mapping be- comes non-deteriiiistic. Moreover, the "overfitting" sce-
we investigate how much information
Generalizing our filtering based interpretative tool to the case of variable sequence lengths (Allen-Zhu & Li 2023 Zhao al. 2023) where the topology of the parsing tree is not known priori-is a challenging but promising di- rection for approaching a more detailed understanding of the learning dynamics and the embedded computation in transformers trained on natural language. On the other hand, while the idealized model of structured sequences
8

---

How Transformers Learn Structured Data: Insights From Hierarchical Filtering
studied in the present work might be less suited for mod- eling natural language compared standard CFGs, the ag- nostic nature of the approach could open connections to other related fields, like protein sequences analysis (Zhang ett al. 2023) and immunology (Meynard-Piganeau et al., 2024). It could finally be interesting to undertake a similar investigation on the way transformers learn in other prob- lems where optimal inference can also be achieved via dy- namic programning (Mossel et al. 2014 2023).
Bardone, L. and Goldt, S. Sliding down the stairs: How correlated latent variables accelerate learning with neu- ral networks. In International Conference Machine Learning, pp. 2024-3045. PMLR, 2024.
Behrens, F., Biggio, L., and Zdeborová, L. Understanding counting in small transformers: The interplay between attention and feed-forward layers. In ICMI. 2024 Work- shop on Mechanistic Interpretabilityy 2024.
Cagnetta, F. and Wyart, M. Towards a theory of how the structure of language is acquired by deep neural net- works. In The Thirty-eighth Annual Conference on Neu- ral Information Processing Systems, 2024.
Acknowledgments
The authors are grateful to Carlo Baldassi, Luca Biggio, Dirk Hovy and Gianmarco Perrupatto for fruitful discus- sions. JGB's research was developed within the MUSA Multilayered Urban Sustainability Action project CUP B43D2101101000 funded by the European Union NextGenerationEUU under the National Recovery and Re- silience Plan (NRRP) Mission 4 Component 2 Investment Line 1.5: Strenghtening of research structures and cre- ation of R&D "innovation ecosystems", set up of "terri- torial leaders in R&D".
Cagnetta, F., Petrini, L., Tomasini, U. M., Favero, A., and Wyart, M. How deep neural networks learn composi- tional data: The random hierarchy model. Physical Re- view X, 14(3):03100,, 2024.
Campos, D. Curriculum learning for language modeling. arXiv preprint arXiv:2108.02/70) 2021.
De Giuli, E. Random language model. Physical Review Letters, 122(12):128301, 2019.
Reproducibility statement
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. BERT: Pre-training of deep bidirectional transformers for language understanding. In Burstein, J., Doran, C.,, and Solorio, T. (eds.), Proceedings of the 2019 Confer- ence of the North American Chapter of the Association for Computational Linguistics: Human Language Tech- nologies. Volume 1 (Long and Short Papers), pp. 4171- 4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
We provide the source code used to perform our nu- merical experiments on the repository accessible at https::/gittub... wwt.///eaeeneeeeemmmcccaaaaaaaaa
trre--anguaae It includes a Python script generating the data, as well as the PyTorch implementation of the transformer and training scripts for both root inference and MLM. It finally provides efficient implementation of the Belief Propagation algorithm which can be used for both root inference and Masked Language Modeling. The data used to produce the figures in the main text corresponds to fixing seed 0 and sig = in the data generation script, see Appendix for details on the role of the latter.
Howard, J. and Ruder, S. Universal language model fine- tuning for text classification. In Proceedings of the 56th Annual Meeting of the Association for Computa- tional Linguistics (Volume 1: Long Papers), pp. 328- 339, 2018
Impact statement
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. 2020.
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
arXiv preprint arXiv:2001.088
Khalighinejad, G., Liu, O.,, and Wiseman, S. Approximat- ing CKY with transformers. In The 2023 Conference on Empirical Methods in Natural Language Processing. 2023.
References
Allen-Zhu, Z. and Li, Y. Physics of language mod- els: Part 1, context-free grammar. arXiv preprint 1rXiv:2305.13673. 2023.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint 14Xiv:441266980 2014.
Liu, Y., Ott. M., Goyal, N., Du, J., Joshi, M., Chen, D., Levy, O., Lewis, M., Zettlemoyer, L., and Stoyanov, V. Roberta: A robustly optimized bert pretraining ap- proach, 2019.
Baker, J. K. Trainable grammars for speech recognition. The Journal of the Acoustical Society of America, 65 (S1):S132-S132. 1979.

---

How Transformers Learn Structured Data: Insights From Hierarchical Filtering
Mei, S. U-nets Belief Propagation: Efficient classifica- tion, denoising, and diffusion in generative hierarchical models. arXiv preprint arXiv:2404.1844 2024.
Wu, S., Thalmann, M., and Schulz, E. Motif learning facil- itates sequence memorization and generalization. 2023.
Zhang, S., Fan, R., Liu, Y., Chen, S., Liu, Q., and Zeng. W. Applications of transformer- -aseed language models in bioinformatics: a survey. Bioinformatics Advances, (1):bad001,, 2023.
Peynard-Piganeau, B., Feinauer, C., Weigt, M., Walezak, A. M., and Mora, T. Tulip: A wansformer-basee unsu- pervised language model for interacting peptides and cell receptors that generalizes to unseen epitopes, Pro- ceedings of the National Academy of Sciences, 121(24): e2316401 121, 2024.
Zhao, H., Panigrahi, A., Ge, R., and Arora, S. Do trans- formers parse while predicting the masked word? In Proceedings the 2023 Conference on Empirical Meth- ods in Natural Language Processing, pp. 16513-16542. 2023.
computation. Oxford University Press, 2009. Mézard, M. and Montanari, A. Information, physics, and
Mossel, E. Deep learning and hierarchal generative models. arXiv preprini arXiv: /6000005 2016.
Zhong, Z., Liu, Z., Tegmark, M., and Andreas, J. The clock and the pizza: Two stories in mechanistic explanation of neural networks. Advances in Neural Information Pro- cessing Systems, 36, 2024.
Mossel, E., Neeman, J., and Sly, A. Belief propagation, ro- bust reconstruction and optimal recovery of block mod- els. In Conference on Learning Theory, pp. 356-370. PMLR, 2014.
Mossel, E., Sly, A., and Sohn, Y. Exact phase transitions for stochastic block models and reconstruction on trees. In Proceedings of the 55th Annual ACM Symposium on Theory of Computing, pp. 96-102, 2023.
A. Further details our data model
The transition tensor M-the "grammar" of our genera- tive model in CFG terminology- controls the prop- erties of the above-defined generative process. We define parametrized ensemble of random grammars, from which multiple transition tensors can be sampled independently. Two grammars generated with the same parameters are ex- pected to share some high-level features and produce data of comparable complexity, at least in the large vocabu- lary size limit. Elaborating on recent work on context- free grammars (see Sec. [2.3 of the main text), we generate tran- sition probabilities as
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L.,et et al. Pytorch: An imperative style, perforeefere deep learning library. Advances in neural information processing systems, 32, 2019.
Refinetti, M., Ingrosso, A., and Goldt, S. Neural networks trained with SGD learn distributions of increasing com- plexity. In International Conference on Machine Learn- ing. pp. (8843-28863. PMLR, 2023.
Sec.2..3oo
Rende, R., Gerace, F., Laio, A., Goldt, S., et al. A distribu- tional simplicity bias in the learning dynamics of trans- formers. Advances in Neural Information Processing. 37, 2024.
Mabe
(3)
Eve eharld
where the logits habe are generated
Sato, T. Inside-outside probability computation for belief propagation. In IJCAI, pp. 2605-2610. Citeseer, 2007.
SoEate -00
if (b,c) E otherwise,
habe
(4)
Sclocchi, A., Favero. A., and Wyart, M. A phase transi- tion in diffusion models reveals the hierarchical nature of data. Proceedings of the National Academy of Sci- ences, 122(1), 2025.
with Eade independent Gaussian random variables of zero mean and unit variance, and a controlling the probability fluctuations between likely and unlikely transitions. Here, the q sets O2 build a equal-sized partition of the q² possible children pairs (b, c), i.e. O, n n. 0 if * a' and Ua Oal q². This non-overlapping prescription implies that the broadcast from the root to the leaves has no ambiguity. Therefore, as stated in the main text, if the transition tensor M known, one can deterministically go up the hierarchy of the tree and infer the root given a set of leaves. We leave generalizations of this setting for future work.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, I.., Gomez, A. N., Kaiser, t.., and Polosukhin, At- tention is all you need. Advances in neural information processing systems, 30, 2017.
Weiss, G., Goldberg, Y., and Yahav, E. Thinking like trans- formers. In International Conference on Machine Learn- ing, pp. 1 10000 PMLR, 2021.
10

---

How Transformers Learn Structured Data: Insights From Hierarchical Filtering
(a)
(b)
M
connected to three variable nodes, representing the parent and two children in the standard branching process. The outgoing messages from the factor node should satisfy
Por 001-
P0- ag
Pa-su
Paga 124
Po-sl
a
Va-su(Tu) & Σ
MMassez Vi-aa (sa) (xr]-
2
Figure S. Illustration the two types of BP updates: (a) above:
(b) below the filter level k.
For all variable nodes except for the root detailed above, the single outgoing messages are equal to the single incoming messages in these variable nodes at the previous/next layer of the tree. For example, the upgoing messages V1->a, in Fig. 5(a) simply V where a is the full factor node lying below variable 1 (assuming k E). Efficient conver- gence to the fixed point is guaranteed if one starts from the leaves and updates the messages in an upgoing pass, and then performs downgoing pass from the root, for a total of 2 - + 1) steps. Once the messages have converged, any unknown variable can be optimally reconstructed by computing the marginals as
B. The BP inference algorithm
Here we present the exact Belief Propagation algorithm used the gold standard against which to compare the nu- merical results. For the full derivation, see Mézard & Mon- tanari 2009). We start by randomly initializing an upgoing and downgoing message-each one being a vector in IR' that represents probability distribution over the q possible symbols- for each edge in the generative tree. In the fol- lowing, we denote with Uj-+ message going from so- called variable node (shown by circle in the sketches) to factor node a (shown by a full or empty square in the sketches), and with Da-t the message in the opposite di- rection. Wherever there is a known variable one should then fix [z,+ = de where a is the known value e.g. of the leaf.
1420 x II Da-nloil, aEBi
(8)
where Oi is the set of factor nodes connected to variable node i. In our problem, this product will therefore typically be over single factor node when inferring masked leaves, or 2 factor nodes when inferring the root.
When the hierarchy is truncated, two distinet types of up- dates are possible, depending whether one lies in the filtered or unfiltered regions of the tree. In the former, the root is directly connected to 2* "empty" factor nodes, as shown in Fig. 5 (),, each connected to a single and distinct variable node below. In this case the BP fixed point equa- tions for messages from the root the empty factor are given by
C. Vanilla encoder-only transformer architecture
A sequence of leaves {z;} generated by the hierarchical model and represented by 2 integers is first converted into a sequence of one-hot vectors {æ;}, with æ, E B Then, we perform the first encoding step producing a sequence of tokens (x(")} E Rd, with arbitrary dimension d 2 9. obtained through a learnable projection to the embedding space and the inclusion of positional encoding p..
10--a [zol x II [ooooo
(5)
Egj
Le. outgoing messages are simply a product of the incom- ing messages from all the other edges. At each of the 2* factor nodes, both upgoing and downgoing messages sat- isfy
0o00) x Σ P(r5 | (o)) [ri],
(0) = WE,; + Pi: ;
(6)
[13) x P(z, 20)20-r0,(za)
with We &dd and p: Rd. For our experiments, we consider d 128 and the standard sinusoidal positional encoding of (Vaswani et al. 2017).
where P(ay 2o) is given by equation and is specific to the factor node considered. The notation means that the messages-the that are probabilities are to be normalized (e.g. Σn, Po 1).
As described in the main text, each transformer block in the
'For simplicity, the procedure described here does not con- sider special tokens. In practice, will take vocabulary of size q 1 account for masked symbols when doing MLM, see Devlin al. 2019).
We now consider the lower, unfiltered part of the tree. As illustrated in Fig. b), each of the "full" factor nodes is

---

How Transformers Learn Structured Data: Insights From Hierarchical Filtering
network then transforms the tokens as follows,
The Kullback-Leibler divergence between two discrete probability distributions encoded n-dimensional vectors u and is given by
layernorm
!1-1+
+ selfattention x
(1-1)
1),
W w()).
(u v)
log
(16)
(10)
aml
1 layernorm
( FC (1
W w
E. Additional figures
(11)
E.1. Influence of the number of attention layers
The single-head self-attention layer considered in this work entails the computation of three different quantities from each token: the query qi = Wo the key k; = wx and the value v, Wr For simplicity, we take We. WK and Wv in Rdxd The queries and keys are combined to compute the attention matrix
Establishing relation between the number of encoder lay- ers nL in the transformer and the ability to achieve this optimal classification on data generated from hierarchical models is also not straightforward. Indeed, given the con- catenation of operations involved in single transformer block and the presence of residual and normalization lay- ers, the effective number of computational layers in trans- former is not as explicit as in multilayer perceptron or a CNN architecture. As apparent in the main text, setting u = 6-or n 2 & for filtered data-enables the trans- former to converge towards a very interpretable parameter configuration. However, this natural choice does not ap- pear to be strictly necessary for the transformers to achieve optimal inference, at least when the number of embedding dimensions d is large.
9. Kj V
AG = softmax
(12)
then used to build a linear combination of the values,
2
sellattention(;; Wo. WK, Wv)
Σ
An
(13)
The fully-connected layer, instead, is standard 2-layer
network with relu activations:
More specifically, Fig. showss that the test accuracy on the root classification task on k unfiltered data can reach the optimal value for nı < While nı l is the most sample efficient, it is clear that nt, provides com- parable performance, and only nt = appears to lead to poor sample efficiency. In all the performed experiments, bigger value for nı corresponded to better sample effi- ciency, which seems to indicate that more flexible models require less data to reach the same performance level de- spite the increased number of parameters to train.
FC(:; W1, W2) = W2 relu (W
(14)
where W1 E 1xd W2 E Igd xd, and d' 2048 in our experiments. We refer the reader to the original paper by (Vaswani et al. 2017) for additional details on the trans- former encoder operations.
D. Further details on numerical experiments
All numerical experiments presented in this paper were per- formed using PyTorch (Paszke et al., 2019) version 2.3.0. We use the Adam (Kingma & Ba 2014) optimizer with batches size 32 and a fixed learning rate of 10-4, other parameters left as default. We did not find learning rate scheduling to provide significant benefits in experi- ments. All models were initialized randomly using the de- fault settings (Xavier uniform distribution).
In any case, the required complexity of the architecture is clearly related to the amount of structure in the data model. As an extreme illustration, in the case of fully fil- tered correlations k = L, the BP marginals for the root are just products of conditional probabilities on the leaves as P(zo (x;)) x II P(x; 2o = a), i.c. a "Naive Bayes" classifier optimal. Any layer of attention is thus superfluous since standard feed-forward network with single hidden layer is sufficient for this task. In fact, the analysis of the attention maps (trained this time on MLM) in Sec. 4 confirms this natural intuition, as most attention layers appear effectively unused by the transformer when nz > k.
In both root inference and MIM, the accuracy of the trans- former implementation and of the BP over M trials is mea- sured straightforwardly as
Accuracy
Σ Sa
(15)
E.2. Other grammars
where x, is understood as the ground truth and &, the sym- bol inferred using the network BP.
As expected from the log-normal nature of its entries, there may be significant sample to sample fluctuations in the
12

---

How Transformers Learn Structured Data: Insights From Hierarchical Filtering
bution is 8-distributed. The other is that this being said, matching the BP is also easier in the k = 0 case because it is the only case where the training cross-entropy loss corre- sponds exactly to that computed with the true marginals- that are also delta distributed due to the determinism of the ther-w in the > cases the training loss does not guide explicitly to the exact marginals. The latter clearly appears in Fig. 9 showing the Kullback-Leibe di- vergence between the transformer outputted logits and the BP marginals instead of the test accuracy.
10-
-
-6-
Test accuracy
0.8
0......0... ....000
0.6
10a
10
10
P
Note that the other case which has a singularly small sam- ple complexity is that of the fully filtered data, = L, as it is implementable in a single feedforward layer and does not require an implementation equivalent to BP.
Figure 6. Reproduction of Fig. bb with now ny S atten- tion layers in the transformer encoder and restricted to the "worst
case" k = 0 unfiltered dataset.
E.4. Additional comparison of the outputs
transition tensor M for a given value of a, which we ex- pect to (slowly) decay as q becomes large. All the results presented in the main text come from the same grammar 1 (corresponding to seed data generation script provided in the SM, see the Repro- ducibility Statement above), however we illustrate that all our conclusions should qualitatively hold for any realiza- tions of M in Fig. 7 Indeed, while there are some very clear differences in the "difficulty" of the grammars presented, the transformer architecture performs very similarly, here on the root inference task. All subsequent experiments can be reproduced on these different grammars, yielding un- changed phenomenology.
For completeness, we show the comparison between the full transformer predictions and the BP marginals through MLM training using the percentage of matches in the largest value (i.e. prediction match) and the spearman (or- dering) correlation in Fig. 10 These confirm the observa- tions described in the main text.
with q
in the
E.S. Classifier attention maps
Fig. I shows the attention maps resulting from the super- vised training for transformers achieving the optimal per- formance on datasets with different filtration levels. As in the masked language modeling task, one immediately notices the emergence of blocks of size ~ 2 In this prescription, where tokens are not required to be fully de- scriptive, it is however difficult to identify clear pattern relating to the distribution of the computation across the different layers.
E.3. In-sample classification performance on filtered datasets
Fig.8 8 shows the test accuracy computed in-sample for the facttrized datasets as a function of the training set size P. The optimal inference accuracy predicted by the Be- lief Propagation, which is not unity when l: >> 0, is reached by the transformers in all cases when trained on sufficient data.
E.6. Details on the probing experiments
In order to perform the experiments presented in Fig. B)), we replace the linear readout of a trained MLM transformer by a two-layer feedforward network with 64 hidden units, acting independently on all of the d-dimensional sequences (d = 128 all of our experiments, see Sec. 3 outputted by the frozen transformer encoder. The training of the readout is performed on 214 labeled sequences, the labels being, for each of the elements of the sequence, the symbol on the relevant ancestor in the generative tree. Here again, the loss is taken to be the cross-entropy between the logits outputted by the network for each token and their correct ancestor label, then averaged on all the sequence elements. We present another experiment, where the cross-entropy measured only with the first and the last token embeddings of the sequence, just below. The readout is trained on 100 epochs in all cases, which we found to be sufficient for the
It appears that the required amount of data P* for reaching optimal accuracy not only depends on the specific transi- tion tensor M (see Fig. 7 for an illustration for & 0), but also on the level of factorization. For intermediate values of k, PM is notably larger than with the = 0 full hier- archy. This is due to the fact that the = 0 case is quite unique for two (related) reasons. The first 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 un- derstood to match the optimal inference in the ambiguous > 0 cases in other words it is easier to match perfect accuracy with approximate weights when the true distri-
13

---

How Transformers Learn Structured Data: Insights From Hierarchical Filtering
1.0
Test accuracy
0.8
0.6
10a
10 P
10
10
109
P
Figure 7. Reproduction of Fig. [b) other realizations of the transition tensor M for the parameters L 4. 4. We remind that for the > 0 cases, the BP predictions (dashed lines) are not Bayes optimal, as the test accuracy is measured out-of-sample here. From left right, these grammars reproduced by fixing seed { 313 in the data generation code provided in the SM.
Prediction match to BP (%)
100
1.0
1.0
Spearman correlation
0.8
0.6
Test accuracy
0.8
0.4
0.2
0.6
0.0
0 100
10a Epochs
10a
10 0
10a
10a
10a
Epochs
0.4
104 P
10
Figure 10. Reproduction of Fig. ( with the prediction (i.e. arg max) match (left) and Spearman (i.c. ranking) correlation (right) between the transformer outputs and BP marginals.
Figure 8. Reproduction Fig. a) with the test accuracy com- puted (in-sample) factorized data, rather than the out-of- sample testing presented in the main text.
109
Dxs
10-1
10
10
10a
10a
10"
10a
10
10"
P
Figure 9. Reproduction of Fig. with the Kullback-Leibler diver- gence between the transformer outputs BP marginals for identical levels of factorizations for (Left) in-sample inputs, (Right) uni- formly randomly generated inputs.
Figure 11. Reproduction of Fig. 4 for the supervised task fil-
tered datasets size P = 217 for 0 and P = 22 for > 0.
14

---

How Transformers Learn Structured Data: Insights From Hierarchical Filtering
F. A possible transformer implementation of Belief Propagation
1.0
Test accuracy
0.8
We show here how the BP algorithm for leaf inference can be implemented using layers of transformers with token sizes which are compatible with what is used in our exper- iments. We consider the "worst case" scenario of a com- plete, unfiltered tree generative process depth
0.6
0.4
3
Tree layer
Tree layer
We propose an implementation that relies vectorized tokens with structure of the form
Token embedding.
Figure 12. (Left) Reproduction of the probing experiment pre- sented in Fig.Blb, B)b, with the readout trained only on the first and last token embeddings of the sequences and tested all elements. (Right) Test accuracy in the ancestor prediction task (layer is the root) with & ing out the complete transformer encoding of models pre-trained Atrain 2, 3, 4 (from top hottom), i.c. using the attention maps illustrated in Fig. The readout is trained on 2 labeled examples.
(1.a)
(m)
(g.m)
(17)
q
kreat 0. obtained by read-
m (m)
m; (m)
p;
where:
i E {l, 20 is the index of a leaf
m E (1,....
{) is the index of transformer layer
r (1.m) ....., (g.m) are g vectors of dimension g (q² el- ements total) storing the quantities needed to com- pute the final leaf marginals.
relatively small training set size we used.
E.7. Further probing experiments
(m) is vector of size q storing the up-going mes- m,
To complement and contextualize the probing experiments presented in the main text, we provide two additional ex-
sage for the ancestor of leaf at level
periments. In the left panel of Fig. [2| we perform the
m (m) is a vector of size q storing the up-going mes-
same experiment as 3b), but with probes trained only on two positions the token sequence (first and last) and tested across all positions. While some accuracy is lost, since the readout cannot fully disentangle the posi- tional information from the semantic one in positions that were never seen at training, the sequential effect still ev- ident. Morcover, we also performed the same procedure as Fig. B()) on the tokens' hidden representations, but with models trained on factorized data. As visible in the right panel of Fig. 12 model trained of filtered data can only accurately recover ancestors up to the level in which filter- ing kicks in. For example, in an = 4 tree, a model trained on Krain = 2 data can only predict ancestors up to level 2 (two ancestry layers above the filtered), while a model trained on Airii predict ancestors up to level (the ancestors right above the leaves for the same reason). This is exactly what could be expected from the attention maps of Fig. 4 As before, we are probing the hidden representations of individual tokens, so this happens because the attention must provide mixing between ~ 2 elements of the sequence in order for in- dividual tokens to carry information up to the level & of the fully hierarchical generative model.
sage for the mhth complementary ancestor of leaf see Fig. [3
P: é-dimensional binary vector containing posi- tional information on the full path from root leaf i (see below).
In this prescription, the total dimension of each token is therefore d = q² + 29 +
Initialization. We are going to consider the following initialization,
that, the tree can only
(1(000)
1
Va, b I.....g.
(18)
q
m (0)
0.
(19)
(0) while the messages m. should be initialized as in the standard BP given sequence, i.e. with a Kronecker 5 for known symbols and a uniform vector for masked leaves. The positional vector P; should finally be a binary +1 vec- tor representing the sequence of left/right turns from the root to leaf (as a in equation Tb.
15

---

How Transformers Learn Structured Data: Insights From Hierarchical Filtering
all the trajectories within the complementary sub-tree
m
m
selfattention (m))
(m)
W
[m)
W (m))
K
0
m
m; m;
(1)
E
"ES
(m)
m,
(m)
m (m)
(21)
Figure 13. Illustration of the upgoing messages embedded in the tokens of the transformer implementation of BP for a tree with t= 3. Complementary ancestors are shown with dashed lines.
where S. (mm)
the set of tokens belonging to the comple- mentary tree of token at layer é of the generative tree. Note that in principle it is not necessary to average since all of the paths should lead to the same message from the complementary tree, however keep in mind that in prac- tice some tokens will be masked. The averaging procedure therefore allows recovering the information (unless all of the tokens in 5 happen be masked). Thanks to the skip connections, this contribution is added to the initial token, populating the initially empty entries of these com- plementary messages while leaving the rest of the tokens unaffected.
Attention layer. In our implementation, the dot product
T (wimiaa (wim)asfmi) ) (wkk (m) (m))
entering the softmax and at the heart of the attention mech- anism only encodes positional information; more precisely, it combines the common ancestors of tokens and j down to layer f - m of the generative tree. This can be achieved w(() has elements equal to zero except in its lower right corner of size x E which has the following structure:
Fully connected feedforward layer. Following the ini- tialization and after the attention layer, the encoded token has the correct structure of equation [7 One must now up- date the relevant information in order to go the next at- tention layer and therefore the next layer in the generative tree. More precisely, we need to:
with query and key matrices such that (w(m)) Q
Compute the messages of the m + 1 ancestor,
B1 (t-m--))((e-ee
0
0
0
3
0
(20)
Update the quantities needed to compute the marginal for the leaf associated with the token considered,
0
0
17Xmm
Remove temporary or unwanted quantities stemming from the previous steps.
with 3 >>1 1. Let us detail the role of this x l sub-matrix. Its upper left terms proportional to will be relevant in the softmax, when >> 1, if they are positive, meaning these common ancestors to tokens i and j, and negligible if they are negative. The diagonal term proportional to -9 requires the two considered tokens to be in different posi- tions in the sequence to contribute to the softmax, ensuring there is no influence of the messages on themselves in the following steps. Its lower right corner, which is populated by a m m matrix of zeros, ensures that layers below é-m in the generative tree are longer considered.
All of these must be done with an identical operation for all tokens as the feedforward layer is applied independently for all positions in the sequence.
The first part is to update the messages following the equiv- alent of equation 6
(m+1) m
MMP,((0.)) mm)")), (m (m!")) (22) be
where Pi(b, c) is either be or cb depending on the topol- ogy of the factor node at which the update takes place- a piece of information fully contained in p, This type of operation should be implementable, at least approxi- mately, by a two-layer network since it is known to be universal approximator. A possible, non par-maasioo
On the other hand, the value matrix may be used select the correct messages in the token vector, with zeros else- where.
As a result, the total operation amounts to averaging the message incoming from the complementary sub-tree over
16

---

How Transformers Learn Structured Data: Insights From Hierarchical Filtering
way perform the above update with two-layer fully- connected network with (q²") neurons is the following. In the first layer, one can readily select the appropriate entries m (m)) 2
In particular, the idea is that we are tracing only over the index of the complement ancestor- is already avail- able from the first layer-but not on the index of the down- going message, which can only be computed after reaching the top of the hierarchy. Instead, we keep in memory all the separate contributions for each parent index. Then, we can obtain a recursion for the auxiliary messages:
(m ),,
2
in the embedding vector to output
and (mm{))) (174)
+ m )Tt for all pairs (b, c). Then,
for each transition MaP, (b.c) the argument of the sum in equation 22 can be obtained as it is equal to
r3 ((ammt)))
x Σ MoP((F.K) (r, (((m)))
)
(m)
m
1 Mapi(bac) ((mm.))) (m))
+ (ma) )²
(26)
2
(23)
with the base case given in Eq 25 treated in the transformer first layer. At the last transformer layer, one can also trace over the root index, completing the recursion. Doing so in the final feedforward layer notably yields, at the end of the transformer encoder,
2 -(m( ) -mm!mi) (mm)
The trace over b and is then performed by the second layer of the fully-connected block. For each transition, it reads the three corresponding hidden units and multiplies them by the same learned weights 1 (MeP.(b.e) (using the appro- priate positional embedding entry). while the summation is done as usual. In the sparse transition tensors we con- sidcring, this in fact only requires O(²²) hidden units. Note that this exact operation would require squared activations, but can be approximated with a ReLU network by means of a piece-wise linear approximation.
Σ
Σ h,k
Σ
MoP((hAA)
x
(27)
(r/ (1(aa6-1))
m 1
which is proportional to the incoming message the leaf and therefore to its marginal if it is to be inferred. The final linear readout may then select this relevant part of the outputted tokens to perform the masked language mod- elling. This requires the embedding of a negative identity operation for each the r and m component, which can still be done with (gg") hidden units in general (O(g") in our sparse case).
Now, we are to compute the actual leaf marginals. As men- tioned in the presentation of the standard BP implementa- tion (Sec. 2.4), the standard approach is perform both an upwards and downwards pass, which would require 2 attention layers.
Here, we instead wish to perform the computation in & step. as we have seen from experiments that the transformer can achieve perfect accuracy with l attention layers and that it does not appear to use all layers when k < f. To do so, we have included the q² elements of r1 ..... rg in the token and now show how to update these. Note that if we had 2 layers, we could instead only store q quantities.
Including intermediate layers. In principle, one could add q x (E - 1) new vectors entries in the token in or- der to store the marginals at intermediate layers. These would simply be used to store the intermediate values of the E. (a.1)
Accommodating for filtration. The implementation de- scribed above considered the case of l: O. unfiltered gen- erative trees, i.c. the most complex case from the BP stand- point. In the case of a dataset with filtering parameter k one can adapt the implementation by taking f - k layers. The central difference then lies in the l th block, which must then combine the 2* messages going up to the root in its feedforward layer (instead of two messages like at all other layers in the k 0 case).
As an example, consider the factor graph in Fig. 13 and as- sume the root is not pinned. We can start from the standard BP recursion for the down-going message received by leaf i:
(1) m; ),,
Σ
(3) m
Masases
bi
Ca
ag.ri
x
m (2)
Masbzas
m
(1)
Mazhies
by
(24)
and define an auxiliary message with a double index depen- dence:
= =
m
Magtie.-
(25)
bn
a
17


Reviews:
[{"review_id": "lMR174A9Sh", "reviewer": "Reviewer_Rk87", "summary": "The paper studies how simple transformer models learn to perform root and leaf inference (corresponding to classification and masked-language modeling tasks) on a synthetic generative hierarchical model of data on a regular tree of depth $\\ell$, belonging to the class of context-free grammars. For such a model, exact inference can be done using belief propagation (BP). The paper demonstrates that *(i)* Transformers learn the same marginals as BP; *(ii)* Increasingly deep levels of the grammar are learned sequentially by the transformer as training time increases; *(iii)*  Probing experiments suggest that the transformer reconstructs the grammar’s structure across its layers. Finally, the authors *(iv)* propose a theoretical implementation of BP in an $\\ell$-layer transformer architecture.\n\n## Update after rebuttal\n\nGiven the rebuttal and the additional results provided on the efficient implementation of the belief propagation algorithm within the considered transformer architecture, I have raised my score from 2 to 3 and now lean toward acceptance. I did not raise the score to 4, as I still find several of the paper’s contributions to be largely incremental.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for carefully reading our work and providing valuable feedback.\n\nOn 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.\n\nWe now answer their two questions directly.\n\n1. 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).   \n2. 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.\n\nFinally, on their comments:  \n\n> the authors should clarify (...) new insights beyond previous CNN-based results (…)  \n\nIn 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.  \n\n> Cagnetta & Wyart (2024) also study (...) theoretically the progressive learning of hierarchical correlations with transformers (...) while the present submission empirically studies the problem (...)  \n\nWhile 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.", "claims_and_evidence": "The claims in the submission are convincing and scientifically supported. The paper presents both theoretical and empirical evidence for transformers implementing BP-like inference.", "methods_and_evaluation_criteria": "The methods are well justified. The use of a synthetic model mode data enables controlled and interpretable experiments.", "theoretical_claims": "The paper provides a theoretical implementation of BP inside an $\\ell$-layer transformer architecture, which appears to be mathematically correct.", "experimental_designs_or_analyses": "The experiments are scientifically sound and adequately support the paper’s claims.", "relation_to_broader_scientific_literature": "The paper is of incremental nature with respect to previous literature in the area. Similar findings about learning probabilistic graphical models on $\\ell$-level trees with the transformer have been obtained for CNNs (Cagnetta et al., 2024; Mei, 2024; Cagnetta & Wyart, 2024). While the paper provides insights into transformers approximating BP, the authors should clarify whether their work offers new insights beyond previous CNN-based results or confirms those results in a different architecture. Cagnetta & Wyart (2024) also study numerically and theoretically the progressive learning of hierarchical correlations with transformers, although as a function of training points $P$, while the present submission empirically studies the problem as a function of training time $t$. Finally, Zhao et al., 2023; Allen-Zhu & Li (2023) study theoretically and empirically whether and how transformers implement the optimal inference algorithm when learning context-free grammars on non-regular trees (i.e., the inside-outside algorithm).", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Other strengths:\n\nThe paper is well written and organized. It considers an interesting and timely problem, namely how machine learning models learn the hierarchical structure of data such as language.\n\nOther weaknesses:\n\nOne of the paper’s main novel contributions is providing an implementation of the optimal inference algorithm for the considered model, i.e., BP, in an $\\ell$-layer transformer architecture. Previous results considering the implementation of the inside-outside algorithms – a generalization to BP that relaxes the assumption of fixed-tree topology – were requiring more layers. However, it remains unclear if the new proposed construction is more efficient only because of the fixed-topology assumption, which is then unrealistic in practice. \n\nFurthermore, the authors claim their solution to be “efficient”. However, from Appendix F, it seems that they do not control the number of neurons required in the MLPs to update messages, just leveraging the universal approximation property – which can, however, require an exponentially large number of neurons in the input dimension.", "other_comments_or_suggestions": "The appendix should be moved after the references.", "questions_for_authors": "1. Can you comment on the discussion above on BP vs inside-outside and the sufficiency of $\\ell$ layers in contrast to previous approaches?\n2. Can you elaborate more about how you place your work among the existing literature and what you think are the most important novel scientific contributions for publishing it at ICML?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gz593VE50Z", "reviewer": "Reviewer_GvGF", "summary": "This paper investigates how a vanilla transformer encoder learns to infer latent hierarchical structure from data. The authors introduce a *synthetic* hierarchical tree-structured data model with a tunable filtering parameter $k$ that controls the depth of correlations in the sequence. Using this controlled setting, they train vanilla encoder-only transformers on two tasks: *(i)* predicting the hidden root label (root classification) and *(ii)* masked token prediction (MLM), and compare the models’ behavior to the optimal Bayesian inference (belief propagation, BP) on the tree. The key findings are that transformers achieve near-optimal accuracy on these tasks and produce well-calibrated probability predictions closely matching the BP oracle even on novel inputs​. The network appears to learn hierarchical dependencies gradually during training, first capturing short-range (local) correlations and then longer-range ones​.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback, and answer the points that they have raised.\n\n- 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).  \n- 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.  \n- 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.\n\nWe now answer the specific questions formulated by the referee:\n\n1. 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.  \n2. 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.  \n3. 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.", "claims_and_evidence": "The central claims appear to be generally supported by convincing evidence, including: *(i)* a transformer can approximate the exact tree inference algorithm (BP) and produce calibrated posterior probabilities, and *(ii)* transformers learn hierarchical correlations in a progressive manner during training.No major claims appear unsupported.", "methods_and_evaluation_criteria": "The methods and evaluation setup are appropriate and well-designed for the research questions.", "theoretical_claims": "This paper is primarily empirical and appears to build upon established theoretical frameworks.", "experimental_designs_or_analyses": "No particular issues found.", "relation_to_broader_scientific_literature": "This work lies into the literature of the interpretability of transformers on structured tasks, which has been explored in the context of formal languages and syntactic trees", "essential_references_not_discussed": "Not found.", "other_strengths_and_weaknesses": "This work distinguishes itself by presenting a controlled framework that manipulates hierarchical structures through a tunable filtering parameter. However, a limitation is its divergence from the complexities encountered in real-world transformer applications, such as those used in language models. Expanding the discussion to explore potential implications and applications in practical settings could further strengthen the contribution. Additionally, the paper does not compare the transformer’s performance to alternative approaches (besides the BP oracle). For instance, could a simple feed-forward network that takes all tokens as input solve the root classification task? Or might an RNN, such as an LSTM with sufficient capacity, also approximate BP? Including such comparisons would help emphasise the unique advantages of the transformer architecture.", "other_comments_or_suggestions": "The term \"structured data\" in the title is quite broad and can imply a wide variety of complex, real-world patterns. However, in this work, the focus is on a synthetic binary tree scenario with a tunable filtering parameter.", "questions_for_authors": "- Did you observe the sequential learning of hierarchical correlations across multiple training runs or random seeds? In other words, is this progression reliably reproducible, and how sensitive is it to factors like learning rate or initialization?\n- What would happen if the number of transformer layers, $n_L$, is set larger than needed (i.e., over-parameterised)? This is interesting since, as seen in works like Mixture-of-Experts, not all parameters may be necessary to solve a given task.\n- How do you ensure that the leaf nodes are order-dependent? In natural language, the grammar tree depends on the actual words at the leaves, but in the synthetically crafted model, there is no inherent guarantee of order dependency. Could you elaborate on how this aspect is handled or justified in your framework?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ycHmjw9jLc", "reviewer": "Reviewer_ivGY", "summary": "Transformer architectures have become highly successful in deep learning, achieving state-of-the-art performance in various NLP and computer vision tasks. However, it is still not fully understood how transformers learn from different types of data. This paper takes a step toward a better understanding of transformers' ability to learn from structured data, specifically hierarchical structures, through attention heads.\nThe paper first introduces a complete hierarchical generative process, starting with a root symbol and iteratively generating child nodes using a probability function governed by a transition matrix. A filtering mechanism is applied by introducing a parameter k, where k=0 means that all pairs of child nodes are conditioned on their respective parents, implying strong correlations. For k>0, the children at level k are generated conditionally with respect to the root. This process of generating sequences allows for exact inference using a dynamic programming approach.\nIn the second step, the paper selects an encoder-only transformer model with l layers. The goal is to demonstrate that the encoder-only transformer can approximate the exact inference algorithm when trained on root classification and masked language modeling tasks. In the root classification task, each generated sequence is labeled with the root. The results show that, in any combination of training and testing with k=0 or k=0, 1,..., l, the transformer achieves the same accuracy as the Belief propagation (BP) method and approximates its inference. Similar results are observed in the masked language modeling task, where the model is pretrained by learning to predict masked tokens, optimizing a loss function. Visualization of the attention matrix further supports the claim that transformers can learn hierarchical structures effectively.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s feedback and address here the weaknesses and questions they have raised.\n\n- 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).\n- 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.   \n- 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.   \n- 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.", "claims_and_evidence": "The paper addresses its claims in several ways: experiments in two scenarios of root classification and MLM, and comparing the performance of transformers with BP. In addition, visualization of the attention layer illustrates the way hierarchy is captured by transformers. However, the approach is limited to one types of hierarchy and single dataset.", "methods_and_evaluation_criteria": "Despite being limited to one hierarchy type, the approach follows reasonable steps.", "theoretical_claims": "One of the main drawbacks of the paper is lack of theoretical analysis, given the fact that the experiments are narrow in terms of data type.", "experimental_designs_or_analyses": "The experimental results are valid and cover several aspects of the problem.", "relation_to_broader_scientific_literature": "The idea presented in this paper is closely related to graph-based machine learning and foundation models for graphs. Most existing graph foundation models rely on GCNs or graph transformers as their main components, which may not be as scalable as traditional transformers. This work highlights the potential of using transformers for structure learning,  thus can draw attention to this promising direction in the field.", "essential_references_not_discussed": "Nan", "other_strengths_and_weaknesses": "– The core idea of the paper is interesting and helpful for the community as proving the  structure learning capability of a transformer can lead to developing a general foundation model that learns from both structured and unstructured data. \n\n– the paper only studies a very narrow type of structure and at the same time lacks any theoretical justification. \n\n– for a reader who is not fully aware of the literature, it is a bit hard to understand the introduction. Thus, the writing can be improved.", "other_comments_or_suggestions": "If possible, studying other types of hierarchies would be interesting.", "questions_for_authors": "What is your opinion on structure learning by transformers from a broader aspect, like motif learning?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3XbIihW6FW", "reviewer": "Reviewer_oypm", "summary": "The authors propose a synthetic way to evaluate how Transformers learn interactions on trees that are generated with different positional correlations. They show that the Transformer learns to approximate the algorithm used to generate the synthetic data. They show that the Transformer learns \"longer range\" interactions in the deeper layers and more local structure in the earlier ones.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback, and address their comments and questions.\n\nOn 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:\n\n1. 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.  \n2. 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.  \n3. 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.  \n4. 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.\n\nPoints (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.", "claims_and_evidence": "I think the claims are very clearly backed up with convincing evidence. \n\nI believe that the title \"How transformers learn structured data\" is perhaps a bit of a strong claim. While the paper goes in this direction, it sounds like a pretty sweeping claim that this is what the paper solves. The paper is studying the behaviour in a very synthetic and controlled system. Of course I understand the need to study these kinds of behaviours under such constraints, but I find the title strongly worded.", "methods_and_evaluation_criteria": "The evaluation criteria is solid.", "theoretical_claims": "There are no theoretical claims.", "experimental_designs_or_analyses": "The experimental design is in fact a contribution of the work and I believe it is an interesting contribution.", "relation_to_broader_scientific_literature": "As this is slightly separate to my main area of research, I am unsure how it relates to literature in the surrounding area. There are however a number of studies (that are relatively dated) that show that machine learning models learn in deeper layers more complex dependencies. In some sense I feel like this work goes in this direction as well.", "essential_references_not_discussed": "I am not very familiar with this exact kind of literature that aims to learn probabilitistc models with Transformers so I cannot comment on essential references that are not discussed.", "other_strengths_and_weaknesses": "Strengths:\nI think the setup is quite clear and the results make sense. The authors provide an extensive amount of experiments\n\nWeaknesses:\nThe main weakness I see is that I do not find the conclusions particularly surprising. In fact the claim \"which provides evidence of an equivalence in computation to the exact inference algorithm.\" I would assume follows the the learning process was successful. In the sense that if your loss is low during SGD, I would imagine that the model has learnt to align itself with the underlying algorithm that generates the data.", "other_comments_or_suggestions": "This is a style suggestion but I would put the bibliography before the appendix as this is most standard as far as I am aware.", "questions_for_authors": "Regarding the weakness, why do you find the result that the model seems to align itself with the way the data is generated surprising? Is this not immediate from the fact that you have trained a model that manages to fit the data well using SGD? \n\nAre you suggesting that a Transformer could learn to fit this data in some other way? While I find this result still experimentally interesting I do not necessarily find it very surprising. \n\nRegardless, I still find the paper interesting and the methodology to make sense, for this reason I am leaning more towards acceptance.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The authors build a fully‑controlled playground in which sequences are generated by a hierarchical probabilistic grammar on a fixed binary tree.  Since exact inference on this model is belief‑propagation (BP), one can compare transformers to a known oracle and  propose an explicit BP‑in‑attention construction.  Two tasks: root classification and masked token prediction show that a 4‑layer vanilla encoder trained by SGD achieves BP‑level accuracy and produces calibrated logits that are similar to the oracle on unseen data.  Training‑time probes reveal a staircase dynamic: the network first captures leaf‑level correlations, then progressively longer‑range dependencies.  Attention maps make this process visible, and the authors’ hierarchical filtering technique (artificially limiting which levels carry correlation) confirms that deeper layers specialise to deeper tree levels. 

The paper presents a cohesive mechanistic interpretation by combining oracle comparisons, training dynamics, layer-wise probes, and explicit algorithmic hypotheses aligned with attention patterns. However, findings rely on fixed-topology trees and no benchmarks on  alternative architectures (CNNs, RNNs) are provided, limiting the demonstrated transformer-specific advantage. Initial reviewer concerns about novelty and efficiency were mitigated by additional clarifications and experiments provided in the rebuttal. Reviewers appreciate the extensive experimental analysis and useful insights, despite the synthetic scope, making the paper suitable for acceptance.
Decision: Accept (poster)