Unveiling Transformers with LEGO: a synthetic reasoning task

1.3 Related works

• The study probes architecture and data-composition questions in LEGO, seeking clearer insight into what Transformers learn than real-data analysis can readily provide.The authors identify globally attending to identical tokens as a novel association pattern in their study.Source: §1 · p. 2, §1.3 · p. 4

1.1 LEGO: A synthetic reasoning task

• LEGO represents variable assignments and group operations, testing long-range identity matching and short-range operations across chains of reasoning steps.Training can supervise only an initial portion of the sequence, enabling tests of generalization to longer sequences without positional-encoding errors.Source: §1.1 · p. 2

1.2 Some takeaways

• Classical generalization is reliable across architectures and data regimes, but length extrapolation depends on architecture and data composition.BERT-like models without special data preparation do not extrapolate, while ALBERT and BERT trained with diverse sequence lengths or pretrained BERT do.Source: §1.2 · p. 2
• Extrapolating models develop attention heads for long-range identity association or short-range manipulation; seemingly unrelated pretraining may help by providing these heads.The authors also suggest that pretraining and fine-tuning across different tasks may benefit from better initialization rather than complex knowledge encoded during pretraining.Source: §1.2 · pp. 2–3
• Non-extrapolating models can solve classical generalization through a shortcut that exploits group operations to jump to the chain’s end and reason from both ends.The authors caution that classical generalization can therefore be deceptive, and observe that training-data diversity seems crucial for generalization.Source: §1.2 · p. 3
• ALBERT’s repeated layer appears suited to algorithmic “for loop” problems, and it extrapolates in data regimes where BERT does not.The authors also motivate LEGO attention from learned short- and long-range patterns, reporting that it matches or outperforms its baseline in large-scale pretraining at lower computational cost.Source: §1.2 · p. 3

2 Learning equality and group operations (LEGO)

• Each clause a = gx assigns a variable the action of g on an element or on a previously assigned variable, and the task outputs the assigned element for every variable.Sentences contain a fixed number of clauses in arbitrary order, with a formal restriction ensuring each variable is assigned exactly once.Source: §2 · p. 4
• Valid sentences form directed line graphs rooted at an element of X, with all non-root vertices in A, so following the chain yields a unique assignment.A clause a = gx is represented by a directed edge from x to a labelled g.Source: §2 · p. 4
• Task 1 uses G = {+, −} acting on X = {1, −1}, with clauses ai = ±ai−1 and fixed root a0 = 1; training and test chains randomize clause order.For n ≥ 10, this task has well over a billion unique valid sentences.Source: §2 · pp. 4–5
• Task 2 extends the setup to the permutation group on N = 3 elements, which is non-abelian, unlike Task 1’s two-element group.The paper focuses on Task 1 in its main text and reports Task 2 experiments in Appendix E.Source: §2 · pp. 4–5

3 Transformers for LEGO

• The task predicts variable assignments from input clauses using token classification, with supervision restricted to the first ntr clauses to assess out-of-distribution generalization.The study focuses on BERT and ALBERT as representative NLP transformer architectures.Source: §3 · p. 5, Figure 2 · p. 5
• With n=12 and ntr=6 or 12, both BERT and ALBERT achieve good classical generalization, but only ALBERT generalizes to slightly longer sequences.The authors caution that classical generalization may be a deceptive metric for evaluating logic and reasoning tasks.Source: §3 · p. 5

Q1. Do iterative architectures indeed exhibit better length extrapolation?

• ALBERT extrapolates near-perfectly to variables #6 and #7 and nontrivially to #8, while BERT is only somewhat better than random on #6.This supports a length-extrapolation advantage for iterative architectures.Source: Q1. Do iterative architectures indeed… · p. 6, Figure 3 · p. 6

Q2. Does the ALBERT architecture actually implement the for-loop?

• ALBERT extrapolates to variable #6 as soon as it fits training variables, while BERT improves gradually; probing suggests ALBERT’s layerwise information flow is closer to linear.The probing classifiers predict variable values from representations at individual layers, consistent with a for-loop inductive bias.Source: Q2. Does the ALBERT architecture… · pp. 6–7, Figure 4 · p. 7

Q3. How can we incentivize models to learn iterative solutions?

• Training ALBERT with stochastic depth, sampling depths 6–12 and testing at depth 12, improves generalization to longer chain lengths.The approach incentivizes a solution whose output remains readable at layers deeper than the reasoning chain.Source: Q3. How can we incentivize models to… · p. 7, Figure 5 · p. 8

4.2 Rand-Init vs. Pretrained: Structural advantages from pretraining

• Pretrained BERT generalizes to unseen LEGO sequence lengths much better and in far fewer epochs than randomly initialized BERT.Figure 6 also shows that mimicking initialization narrows the performance gap.Source: §4.2 · p. 8, Figure 6 · p. 8

4.2.1 Why does pretraining help in LEGO?

• Pretrained BERT attention heads exhibit global association between occurrences of the same token and local, convolution-like manipulation patterns.The authors hypothesize that these learned information-transfer patterns largely explain pretraining’s LEGO advantage.Source: §4.2.1 · pp. 8–9, Figure 7 · p. 9

4.2.2 Verifying the hypothesis with Mimicking

• BERT initialized by mimicking association and manipulation heads generalizes better than randomly initialized BERT, supporting the proposed explanation without real-data pretraining.The attention patterns are trained to match target matrices using random data.Source: Figure 6 · p. 8, §4.2.2 · p. 9

4.3 Shortcut solutions and their effect on generalization

• Random-init BERT and ALBERT first learn a shortcut that resolves the last variable, whose value is determined by the parity of minus signs.Training with the shortcut available significantly degrades extrapolation to intermediate variables #5–#10.Source: §4.3 · p. 10

5 LEGO Attention: faster and better

• Pretrained association, manipulation, and special-token broadcasting patterns motivate LEGO attention; BERT uses pathway (b), v0 uses (a) and (c), and v1 uses all three (Figure 8).The paper notes that [CLS]/[SEP] broadcasting may matter for NLP tasks beyond LEGO, although those tokens do not play a role in LEGO itself.Source: §5 · p. 10, Figure 8 · p. 10
• LEGO v1 reaches 1.47 final pertaining validation loss, versus BERT’s 1.49 and LEGO v0’s 1.66, while learning faster initially (Figure 10).The paper describes v1’s final validation loss as slightly lower than BERT’s; v0 learns quickly at first but falls short later.Source: Figure 10 · p. 11, §5 · p. 11
• LEGO attention constructs token-ID association and [CLS]/[SEP] broadcast maps without learning Q/K or attention probabilities, and implements local manipulation with temporal convolutions.The global maps are row-normalized and shared across layers; only a value map per layer is learned.Source: §5 · p. 11
• LEGO v1 reduces model size and accelerates inference relative to the base model, with Figure 9 comparing inference Flops and model size.Figure 9 measures Flops on a batch of 64 sequences of 512 tokens.Source: §5 · p. 11, Figure 9 · p. 11
• BERT, LEGO v0, and LEGO v1 achieve 88.2, 82.5, and 88.1 Dev F1 on SQuAD v1.1 [RZLL16], respectively; comprehensive evaluation is left for future work.The reported values are ordered BERT/LEGO v0/LEGO v1.Source: §5 · pp. 11–12

• The LEGO attention mechanism accelerates inference while achieving comparable or even better performance.Source: §6 · p. 12
• LEGO, a controllable synthetic logical reasoning task, provides insights into Transformers’ inductive biases and the role of pretraining.Source: §6 · p. 12
• Many attention heads beyond manipulation and association remain to be understood, and LEGO may inspire better algorithms and architectures for tasks beyond LEGO.Source: §6 · p. 12

A Shortcut solutions and their effect on generalization

• Randomly initialized models first learn to predict the last variable by counting minus signs, as shown by variable #11 improving earlier than most others.This pattern appears in the top plots of Figure 3, where models train on all 12 variables.Source: A Shortcut solutions and their effect… · p. 17
• The intended iterative solution resolves variables linked to known ones, while the shortcut resolves a variable appearing once using the parity of all minus signs.For a=+1; d=-c; b=-a; c=b, iteration resolves b, then c, then d; the shortcut directly resolves d to 1.Source: A Shortcut solutions and their effect… · p. 17
• The shortcut resembles learning a spurious feature and is brittle under simple task variations that the iterative solution would handle.The passage motivates studying whether shortcut reliance harms generalization.Source: A Shortcut solutions and their effect… · p. 17

Q1. How does reliance on shortcut solutions affect the ability of the network to generalize?

• Early learning of the shortcut coincides with drops in previously learned variables’ accuracy, and the models later abandon it as accuracy improves across variables.The observed drops affect #2 in ALBERT and #3 in BERT; #11 accuracy falls before improving again.Source: Q1. How does reliance on shortcut… · p. 18
• Models supervised on variables #0–#4 and #11 generalize worse to variables #5 in BERT and #6 in ALBERT than controls without #11 supervision, despite receiving less training information.Both groups were tested on all 12 variables; the comparison is shown in Figure 11.Source: Q1. How does reliance on shortcut… · p. 18, Figure 11 · p. 18

Q2. What are effective ways to prevent models from learning shortcuts, and do they result in better generalization?

• Pretrained BERT and ALBERT appear to avoid the shortcut solution, with #11 accuracy improving later than or alongside the other variables.The models were finetuned on all 12 variables; the authors speculate that pretraining lets them learn the iterative solution before a randomly initialized model learns the shortcut.Source: Q2. What are effective ways to prevent… · pp. 18–19, Figure 12 · p. 19
• Avoiding the shortcut may partly explain pretrained BERT’s superior out-of-distribution generalization over randomly initialized BERT.The authors present this as a possible explanation, not a confirmed causal result.Source: Q2. What are effective ways to prevent… · pp. 18–19
• Mimicking BERT and ALBERT show accuracy patterns similar to their pretrained counterparts, suggesting that they also avoid the shortcut solution.The mimicking technique directly reproduces pretrained attention patterns without training on any data.Source: Q2. What are effective ways to prevent… · p. 19, Figure 12 · p. 19

B.1 Data generation

• For each i in [n], the generator samples variable ai uniformly from A and assigns its corresponding label yi, which is +1 or −1 with probability 0.5.Source: §B.1 · p. 19
• Each clause has the form a_i = g_i a_{i-1}, with a_0 = r = 1 and group elements chosen so the clauses agree with assignments a_i = y_i.The procedure generates n clauses.Source: §B.1 · p. 19
• The sentence is formed by placing the n clauses in a random order and separating them with semicolons.Source: §B.1 · p. 19
• Each generated sentence is padded with [BOS] and [EOS] tags to mark its beginning and end.Source: §B.1 · p. 19

B.2 Training

• LEGO sentences are embedded as token vectors and processed by Transformer blocks; experiments use pretrained BERT/ALBERT or identically configured randomly initialized models.The vocabulary includes variable symbols, group operations, a root node, equality, semicolon, and BOS/EOS tokens; the setup draws on [VSP+17] and pretrained weights from.Source: §B.2 · p. 20
• Training and test data are i.i.d. samples from D(n); models receive supervision on the first ntr graph clauses but are evaluated on all n clauses.Randomized clause positions allow positional encodings to be trained for longer sequences than those seen in training; the datasets contain 10^4 × n training and 10^3 × n test examples.Source: §B.2 · p. 20
• Models are trained for 200 epochs with cross-entropy over ntr clauses, Adam, batch size 1000, a 5×10−5 learning rate, and cosine scheduling.All experiments use the pretrained BERT tokenizer; runs use four A100 GPUs.Source: §B.2 · p. 20
• Across runs, iterative ALBERT and pretrained BERT generalize better to unseen lengths, although absolute test accuracy varies substantially.The authors attribute the variance partly to small datasets and differences in train/test samples, initialization, and training seeds.Source: §B.2 · p. 20, Figure 14 · p. 21

B.3 The Mimicking procedure

• The mimicking procedure assigns two heads per layer to match manipulation and association attention patterns using KL divergence on independent, uniformly sampled token sequences.The association target marks identical-token pairs with a zero diagonal; the manipulation pattern uses localized, shift-invariant filters, and Adam reduces the loss nearly to zero quickly.Source: §B.3 · pp. 21–22

C Details of the LEGO v0/v1 models

• LEGO v0 and v1 retain BERT-base’s layer count, hidden dimensions, and feed-forward structure while replacing its attention layers with convolutional and hardcoded-attention pathways.Both use three hardcoded heads with Aasso, Acls, and Asep patterns; v1 additionally has three ordinary attention heads.Source: C Details of the LEGO v0/v1 models · p. 22

Effect of length of LEGO chains and depth of model

• Longer training chains improved generalization, especially monotonically for pretrained BERT; n=8 (ntr=2) learned training lengths but not unseen ones, while n=20 generalized strongly beyond training.Depth-12 models were trained with supervision on ntr = n − 6 clauses; the strong n=20 result may reflect a compact representation rather than natural iteration.Source: Figure 15 · p. 23, D Effect of length of LEGO chains and… · pp. 23–24
• Greater depth improved generalization for randomly initialized BERT and ALBERT on n=12, with a minimum depth threshold that appeared sublinear rather than D ≥ n.No pretrained models were available at these smaller depths.Source: Figure 16 · p. 24, D Effect of length of LEGO chains and… · p. 24

D.1 Effect of number of parameters

• At matched parameter count, thin BERT lacks ALBERT’s length extrapolation and in-distribution generalization; variable #2 test accuracy stays near 50% after 200 epochs.The model used hidden size 96 instead of 784, with other hyperparameters unchanged.Source: §D.1 · p. 24, Figure 17 · p. 25

E LEGO Task 2: dihedral group

• Task 2 learns D3, a group of order 6 isomorphic to S3; unlike Task 1’s D1, its chain-reasoning shortcut solution is invalid.The task uses six categorical classes and otherwise repeats Task 1’s model configurations and training hyperparameters.Source: E LEGO Task 2: dihedral group · p. 25
• Pretraining benefits are larger on Task 2 than Task 1, while mimicking matches and can outperform pretraining; randomly initialized models also face optimization issues without further tuning.Figure 18 presents the out-of-distribution generalization results.Source: Figure 18 · p. 25, E LEGO Task 2: dihedral group · pp. 25–26
• None of the tested models non-trivially generalizes to more than one extra variable, making suitable architecture or training adaptations an open direction.Source: E LEGO Task 2: dihedral group · p. 26
• Mappings to memorize grow quadratically with group size, motivating future analysis of where models store this information, including at the neuron level.Source: E LEGO Task 2: dihedral group · p. 26

F Attention maps of pretrained BERT on LEGO

• Attention maps for pretrained BERT on LEGO sequences show that almost every attention head implements either a manipulation or an association operator.The visualizations complement Figure 7.Source: F Attention maps of pretrained BERT on… · p. 27
• The sequence assigns a=+1, b=+a, c=-b, d=+c, e=-d, f=-e, g=+f, h=-g, i=-h, j=-i, k=+j, and l=-k.Source: F Attention maps of pretrained BERT on… · p. 27
• The displayed chain defines successive signed equalities from a = +1 through l = -k.The sequence is repeated throughout the supplied passages.Source: F Attention maps of pretrained BERT on… · p. 28
• The sequence assigns a=+1 and defines b through l by signed relations to the preceding symbol, ending with l=-k.The displayed chain includes both positive and negative transitions.Source: F Attention maps of pretrained BERT on… · p. 28