Unveiling Transformers with LEGO: A Synthetic Reasoning Task
Authors: Yi Zhang
Arturs Backurs
Sebastien Bubeck
Ronen Eldan
Suriya Gunasekar
Tal Wagner
Venue: Submitted to ICLR 2023 (2023)
Keywords: transformers
logical reasoning
role of pretraining
attention pattern

Abstract: We propose a synthetic reasoning task, LEGO (Learning Equality and Group Operations), that encapsulates the problem of following a chain of reasoning, and we study how the Transformer architectures learn this task. We pay special attention to data effects such as pretraining (on seemingly unrelated NLP tasks) and dataset composition (e.g., differing chain length at training and test time), as well as architectural variants such as weight-tied layers or adding convolutional components. We study how the trained models eventually succeed at the task, and in particular, we are able to understand (to some extent) some of the attention heads as well as how the information flows in the network. Based on these observations we propose a hypothesis that here pretraining helps for LEGO tasks due to certain structured attention patterns, and we experimentally verify this hypothesis. We also observe that in some data regimes the trained transformer finds ``shortcut" solutions to follow the chain of reasoning, which impedes the model's robustness, and moreover we propose ways to prevent it. Motivated by our findings on structured attention patterns, we propose to replace certain attention heads with hardcoded patterns. This architectural change significantly reduces Flops and maintains or even improves the model's performance at large-scale pretraining.

Under review as a conference paper at ICLR 2023
UNVEILING TRANSFORMERS WITH LEGO: A SYNTHETIC REASONING TASK
Anonymous authors Paper under double-blind review
ABSTRACT
We propose synthetic reasoning task, LEGO (Learning Equality and Group Op- erations), that encapsulates the problem of following a chain of reasoning, and we study how the Transformer architectures learn this task. We pay special atten- tion to data effects such as pretraining (on seemingly unrelated NLP tasks) and dataset composition (e.g., differing chain length at training and test time), as well as architectural variants such as weight-tied layers or adding convolutional com- ponents. We study how the trained models eventually succeed at the task, and in particular, we manage to understand some of the attention heads as well as how the information flows in the network. In particular, we have identified a novel association pattern that globally attends only to identical tokens. Based on these observations we propose hypothesis that here pretraining helps for LEGO tasks due to certain structured attention patterns, and we experimentally verify this hy- pothesis. We also observe that in some data regime the trained transformer finds "shortcut" solutions to follow the chain of reasoning, which impedes the model's s robustness, and moreover we propose ways to prevent it. Motivated by our find- ings structured attention patterns, we propose to replace certain attention heads with hardcoded patterns. This architectural change significantly reduces Flops and maintains even improves the model's performance at large-scale pretraining.
INTRODUCTION
The deep learning revolution is about training large neural networks on vast amount of data. The first field transformed by this methodology was computer vision, crucially leveraging the convolutional neural network architecture LeCunn et al.] (1989); Krizhevsky et al. (2012). More recently natu- ral language processing was revolutionized by the Transformer architecture Vaswani et al. (2017). Transformers are designed to process input represented as "set of elements" (e.g., the words in sen- tence with their positional encoding). This is of course an incredibly generic assumption, and thus Transformers can be applied to wide variety of tasks, including vision Dosovitskiy et al. (202 I), re- inforcement Bearning/Chen et al. (2021a), and protein structure prediction Rives et al.
et al. (202T) among others, even jointly across domains to produce generalized agents Reed et al. (2022). In fact, learning with Transformers is rapidly becoming the norm in deep learning.
Transformer models display excellent performance on the standard criterion "training error/test er- ror" (e.g., for masked language prediction translation). However, what makes them particularly noteworthy, is that large-scale Transformer models seem to exhibit unexpected emergent behaviors, such basic reasoning ability Thoppilan al. (2022) [Brown et al. (2020); (hhowdhryy al.
2022):
Du et (202T): Rae et al. (2021): Hoffmann et al. (202)) Smith et al. (2022)) [ Zhang et al. 2022);
Wei al. 2
al. (2022). excellent fine-tuning performance Hu al. (2022) ; Thoppilan
et al. (2022); Nye et al. (2022);1 Rae et al. (202 )); Polu et al. (2022). or zero-shot learning Brown
( Smith et al. (022)) Zhang et al. (2022). ). Currently, there is a remarkable community effort towards at-scale experimental investigation of Transformers, essentially trying to find out what such models can do when they become large enough and are trained on large/diverse enough datasets. The suc- cesses are striking and capture the imagination Brown et al. (2020); Ramesh et al.) (2022). Yet, for all of these wonders, there is very little understanding of how these models learn, in fact what do they learn. Answering such questions in the at at-scale experiments is particularly challenging, as one has little control over the data when hundreds of billions of tokens are harvested from various
et al. 2020): Chowdhery et al. (2022): Du al. (2021); Rae et al. (202T); Hoffmann et al., (022);

---

Under review as a conference paper at ICLR 2023
sources. In this paper, we propose to take step back, and try to understand how learning occurs and what is being learned in a more controlled setting that captures important aspects of "reasoning".
The benefit of such a controlled setting that we can try to understand some of the most pressing questions in learning with Transformers, particularly around (i) the architecture and (ii) the impor- tance of training data. For (i) we probe the role of multiple heads and depth, and we show that we successfully understand them in our controlled setting. For (ii) we investigate how much the dataset composition matters, as well as how pretraining on merely vaguely related tasks makes fine- tuning successful. In turn, these insights can guide our thinking for large-scale experiments, and we give some of the lessons learned below. In particular, our insights crystallize into an architectural change BERT for faster inference with matching even better performance (Section 5).
1.1 LEGO: A SYNTHETIC REASONING TASK
Core components of reasoning include the ability to associate concepts, and to manipulate them. We propose simple task that captures these two aspects, which we call LEGO (Learning Equality and Group Operations). In LEGO, the input describes a sequence of variable assignments as operations on these variables by a fixed (mathematical) group. One needs to be able to deal with both long-range assignments (the same variable appearing in different parts of the input should be viewed as being equal to same quantity), as well as short-range operations (describing what group element is applied to which variable). A key parameter of an input sequence will be its length, which is proportional to the number of sequential reasoning steps one has to do in order to resolve the value of each variable. We will mostly train with a fixed sequence length (say 12). We often
well
provide supervision only on part of the sequence (say the first 6 variables). We do in order to test the generalization capabilities from smaller length sequences to longer length sequences without introducing potential errors due to the positional encoding in Transformers.
1.2 SOME TAKEAWAYS
In LEGO, we are interested in both classical generalization (i.e., training and test distribution are the same) and vary the length of the chain of reasoning, and thus we refer to this type of generalization as length extrapolation. Specifically, the setting for length extrapolation is to train with supervision on shorter sequence lengths (c.g., supervision on only the first 6 variables) and test on a long sequences (c.g., accuracy computed on 12 variables). A summary of our empirical observations is follows:
generalization. For the latter we focus on distribution shifts that
1. First, classical generalization happens reliably for all architectures and data regimes.
2. More interestingly, length extrapolation seems to depend on architectural/data composition choices. Specifically, BERT-like models without special data preparation do not extrapolate to longer sequences, while other models like ALBERT, or BERT with carefully selected data (such as diverse sequence lengths, or pre-trained BERT) do extrapolate.
3. The extrapolating models all seem to evolve attention heads dedicated to either association (long-range identity matching) or manipulation (short-range operations). We provide evidence that pre-trained BERT (which is pre-trained on seemingly unrelated dataset) generalizes be- cause it has learned such heads.
4. The non-extrapolating models seem to solve the classical generalization problem using a cer- tain shortcut-like solution, whereby using the specificity of the group operations they are able to jump to the end of the chain of reasoning, and then complete the rest of the variables by following the reasoning both from the start and the end of the chain.
We interpret our findings as follows:
Classical generalization can be deceptive metric, as there might be unexpected ways to solve the problem. This famously related to the issue of embedding machine learning systems with common sense reasoning. Namely, we hope that when an ML system solves task, it does so in "the way humans do it", but of course, nothing guarantees that this will happen. Our findings are consistent with the current methodology of increasing the diversity of the training data, which seems crucial for generalization.
(ii) ALBERT-like models, where a layer repeated several times, seem to be an ideal structure for problems that could be descrihed algorithmically "for loop" (as is the case with following
2

---

Under review as a conference paper at ICLR 2023
chain of reasoning). Indeed we find that ALBERT extrapolates in data regimes where BERT does not, clearly separating these two architectures.
(iii) The success of pyvraininngfinetuuiing in vastly different tasks might actually come from a "simple" better initialization, rather than complex knowledge encoded during pre-training.
(iv) The interplay between short-range (close-by information in a sentence) and long-range (the same concept appearing different places in the sentence) is relevant more broadly than in our synthetic task. We observe that the networks effectively learn to deal with short-range/long range information by implementing specific attention patterns. This motivates us to study a new architecture where these patterns are hardcoded, and we show it matches or even outperforms its baseline on the large-scale pretraining but with significantly less computational cost.
1.3 RELATED WORKS
In (Zhang et al. 202I), the PVR (Pointer Value Retrieval) task is introduced, with a similar high- level goal to ours in introducing the LEGO task, namely to study how neural networks learn to reason in a controlled setting. In PVR task, part of the input indicates another part of the input where a function of potentially varying complexity has to be computed. Like us, they use distri- bution shift investigate how various network architectures learn this task, and they observe that networks can learn the task at hand ("classical generalization") yet fail to extrapolate to mild dis- tribution shift. They then ask the following questions: "Are there architectural changes that can enforce better priors and withstand distribution shift? Can novel learning objectives prevent these adversarial correlations? Progress on these questions holds promise for greater robustness." Our study attacks these questions directly in the context of the LEGO task (e.g., ALBERT versus BERT, and training set composition investigations), and our preliminary results indicate that this is indeed a fruitful direction to obtain better models in some aspects (e.g., more interpretable). Other examples of recent synthetic benchmark with a similar philosophy include SCAN (Simplified version of the CommAl Navigation) (Lake & Baroni 2018). CFQ (Compositional Freebase Questions) Keysers jt al. (2020) LIME Wu et al. (2021). PCFG SET Hupkes et al.] (2020). and BONGARD-LÖGO (Nie al.. 2020). In SCAN for example, one has to "translate" command of the form "turn left twice and jump" into a sequence of actions "LTTURN LTURN JUMP" (see Patel et al. (2022) for more recent progress on this dataset). Again, similarly the PVR tasks, these works focus derstanding generalization (in these cases, compositional generalization). Another related line of works on studying Transformers recognize various formal languages, see e.g.. [Bhttamishra et al. 2020);; Yao et al. (202I). A contemporary work (Csordás et al. 2021) proposed modifications to Transformer architectures to achieve significantly better length extrapolation (other works study- ing this important class of distribution shifts include Anil et al. (2022)) ). As far as we know, none of these works try to probe the inner workings of the networks in the same depth as we do here. On the other hand, networks trained real data are being extensively scrutinized, see for example Rogers et al. (2020) where they try to understand some of the attention heads of BERT (see also Saha et al. (2020b). However, making sense of these real-data-traing networks is daunting task, and a key contribution of ours is to show in a limited setting one can obtain a clearer picture of what Transformers learn
The LEGO task is also naturally related to the growing literature on testing mathematical/codingg abilities of Transformers (e.g.. [Saha et al. (2020a)), specifically the simpler tasks of checking the correctness of a proof (or simplifying one, such as Agarwal al. (2021) which studies simplifica- tion of polynomials), or executing code for a given input Chen et al. (2021b). It would be interesting to see if some of the insights we derive in the present paper apply to currently challenging mathe-
matical tasks such as Hendrycks et al. (202T) and IsarStep Li et al. (202T)
There are an abundance of studies on attention heads that have identified the importance of local, convolutional, attention patterns Voita et al. (2019); Correea et al. (2019); Clark et al. (2019); Ra- ganato et al.) (2020); You et al. (2020). However, to the best of our knowledge, we are the lirst to demonstrate the importance of the association pattern that globally attends to identical tokens, thanks to the LEGO task.
LEARNING EQUALITY AND GROUP OPERATIONS (LEGO)
We propose the following synthetic task, which we call LEGO. Let G be a finite (cemiggroup acting on finite set X, and denote g(x) for the action of g E G on x E X. We define formal language
3

---

Under review as a conference paper at ICLR 2023
a
1
d
Figure 1: The graph representation of the sentence +1; d; +b; = -s: fie +d; /
using the symbols from G and X as well as symbols from finite alphabet A which refer as the variables. sentence in our formal language is made of clauses separated by a semi-colon. A A. If E X, such
clause is of the form clause means that the variable a assigned the element ((x) E X. On the other hand if A and X through another clause (or chain of clauses) in the
ga with
A, E G and either
X
the variable was assigned an element sentence, then the clause a = ge assigns variable a to the element g(y) E X. The task's s goal is to take in input sentence with a fixed number of clauses, given in arbitrary order, and to output the assigned element to each variable that appear in the sentence (the formal language will have further restriction that ensures that each variable assigned one and only one element).
We can view sentence as a directed graph on the vertex X with labelled edges as follows: a clause a = gr corresponds to a directed edge from the vertex I to the vertex a, and the edge is labelled with g. We restrict our attention to sentences corresponding to line graph directed away from some fixed root vertex E X, and whose non-root vertices are all in A, see Figure | for an example. particular such sentences are "consistent", meaning that a variable is assigned a unique element (the assignment is obtained by simply "following the chain").
Task 1. The most basic instantiation of LEGO is when G' is the unique group of elements acting on set X also of elements, that of n clauses of the form
{ } and X = {1, -1]. Our sentences thus consists +o;-1 where ag A for i = 1. ....... T and
(we fix 1). Note that in this case our formal language has well over billion unique valid sentences when
10. Example of sentence with = 6 is (see Figure for the graph depiction): a +1; b -a; c = +b; d = f;:
+d; +e. Our task's goal is to report the elements or values from X assigned the variables appearing in the sentence. In the above example, assignments for variables a, b, c, d, e, are 1, 1, -1, 1,,1,
Task 2. One can think of Task as the case of LEGO for the permutation group on N = elements (acting itself). Our second task will correspond to N 3, which is qualitatively different since the permutation group on 3 elements non-abelian.
We will focus on Task in the main paper and include in Appendix E experiments on this Task 2. Our training and test data for the task consists of n length chains described above with the order of clauses in the sentence randomized. A sample input sentence to a transformer looks like [BOS] j=-f; f=-b; y=+t; o=te; d=ty; v=+d; h=-o; b=-i; i=+1;
t=+1; e=-1; 1=-h; [EOS] See appendix for further data generation details.
TRANSFORMERS FOR LEGO
or
elassifieation head
representations
f
nodel
Transformen
f
f
f
T
fi
! (POA]
pans
Figure 2:
Illustration transformer model applied LEGO Task on input sentence d=-a; b=-a; c=+b; a=+1;. We apply a linear classification head to the output representations of each clause's first token to generate predictions for the variables assignment.
We apply transformer models in the token classification pipeline to predict the assignments of the variables in the input sentence, depicted in Figure 2 To evaluate the out-of-distribution generaliza- tion (referred to simply generalization), training. supervision is provided only on the first n,, clauses (first in the graph representation the input sentence). We mainly foeus BERT Devlin al. (2018) and AI ALBERTT Lan al. (2019)
introduce the notation of nur S n, such that during
4

---

Under review as a conference paper at ICLR 2023
Rand-Init
Rand-Init ALBERT
variable
100%
90%
80%
(a)
70%
60%
50%
10
10a
101
102
10a
10
10c
11
epach
epach
Rand-Init BERT
Rand-Init ALBERT
variable
100%
BT
90%
BCC
80%
(b)
70%
60%
50%
10
10
10a
10a
10a
10a 10a
11
epoch
Figure 3: Solving LEGO (Task 1) using BERT and ALBERT, trained from random initialization. Each corresponds to the test accuracy of single variable appearing in the sentence the course of train- ing. The variable numbers in the legend are their position the reasoning chain (or graph representation) of the input sentence, rather than the position in the sentence itself. For example, on input sentence: b--a; d=-c; c-+b; a=+1; i, variable #0 is a, #1 is b, #2 is and #33 is d. Top a): models are trained to fit all variables (n = 12, variables (n 12, Bar 6). Dashed curves represent variables not supervised during training.
12). Bottom b): models are trained to fit the first 6 variables but test on all 12
chitectures.
These two models are representative large transformer architectures for NLP tasks, and observe they exhibit intriguing behavior difference See appendix for training hyper-parameter and dataset construction details.
tasks which we will detail in Section 4
In Figure 3 we report initial results on LEGO with n = 12 and Лür = 6, 12. Both BERT and AL- BERT able achieve good classical generalization, while only ALBERT appears to generalize even to slightly longer sequence length. We observe similar behavior across different lengths of inputs This suggests that classical generalization might be deceptive metric evaluate learn- ing of true logic/reasoning tasks. Motivated by these initial results, in the next section we focus on breaking down the learning dynamics of BERT and ALBERT for the LEGO task towards carefully understanding their strengths and weaknesses.
4 UNVEILING TRANSFORMERS WITH LEGO
4.1 BERT VS. ALBERT: ITERATIVE REASONING ITERATIVE ARCHITECTURES
salient feature of many reasoning tasks an iterative component, meaning they can (or must) be solved by sequentially repeating certain operations. In this section, compare Transformer architectures through the lens of iterative reasoning.
LEGO study and
natural solution to LEGOOanm arguably the go-to solution for a human-is to implement a "for- loop", where each iteration resolves step in the reasoning chain. The iteration could look for the next unresolved variable token whose value could be resolved in one step. Iterative Transformer ar- chitectures such as ALBERT and Universal Transformers Dehghani et al. (2018), where the weights are shared across different layers, inherently implement a for-loop with a number of iterations equal the number of layers. If the model manages learn to implement one such iteration during train- ing, the network would immediately be capable of performing length extrapolation. If this indeed occurs, would point clear advantage of ALBERT over BERT in our setting. This leads the following questions.
Q1. DO ITERATIVE ARCHITECTURES INDEED EXHIBIT BETTER LENGTH EXTRAPOLATION?
The bottom plots of Figure 3 display the length extrapolation result for BERT and for ALBERT. They show the clear advantage of recurrence: While the non-iterative BERT achieves only some- what better-than-random accuracy for one variable (#6) beyond the ones accounted for during train-
5

---

Under review as a conference paper at ICLR 2023
tuned BERT
ture-ttddddd ALBERT
03
Variatie
Venaobe
Figure 4: Visualization of information percolation within the fine-tuned models. The color indicates the test accuracy of the probing classifier at each layer. Brighter higher. We observe ALBERT's information perco- lation is linear than BERT's, which implies ALBERT is biased towards learning a for-loop.
ing (#0- #5), the iterative ALBERT reaches near-pen pear-perfect accuracy two additional variables (#6 and #7). and nontrivial accuracy on the third (#8). These results clearly support that iterative archi- tectures do generalize better in the iterative LEGO reasoning task.
Q2. DOES THE ALBERT ARCHITECTURE ACTUALLY IMPLEMENT THE FOR FOR-LOOP?
To lesser extent, Figure also hints at positive answer to Q2. Observe that ALBERT exhibits length extrapolation to variable #6 immediately (in terms of epochs) as soon as fits the training variables (#0 - #5), whereas for BERT, the corresponding plot (#6) climbs gradually even after the training variables are predicted perfectly. This suggests that once it manages to learn the operations required for one step of reasoning. it can immediately implement those operations over a few more iterations not required in training.
In order to gain stronger evidence, we measure the dependence between the location of a variable token in the chain and the layer in which its value typically resolved. To this end, given a trained model, we train one linear classifier per layer which predicts the value of variable token based only on token representation at the corresponding layer (without using other information), while keeping the original model unchanged. This allows to gauge the rate of information percolation along the reasoning chain in terms layers per reasoning step. the model indeed implements for-loop in its forward pass, one expects linear relationship between the number of layers and the number of reasoning steps already completed. We visualize in Figure f the test accuracy of prediction as function of the layer in the network and depth in the chain. While not perfectly linear, the relation clearly looks closer to linear in ALBERT, suggesting that the ALBERT model has an inductive bias towards learning implement the "natural" for-loop with its forward pass.
Q3. HOw CAN WE INCENTIVIZE MODELS TO LEARN ITERATIVE SOLUTIONS?
We attempt to incentivize the model to implement the "natural" for-loop solution. We rely on the observation that if each iteration of the for-loop simply percolates the information one more step (assigning a value to the next variable), then adding more layers with the same weights should not affect the output, and in fact, one should be able read out the output of the calculation from any layer of the neural network, as long as its depth exceeds the length of the chain. With this observation in mind, train ALBERT model with stochastic depth Haath et al. (0016). We uniformly sample depth between 6 and 12 per batch during training while fixing it at 12 during test. Figure shows clear improvement in generalization to longer lengths using stochastic depth.
Rand-Init ALBERT
Rand-init ALBERT
w/ stochastic depth
variable *
100%
90%
80%
50%
10
10
10a
10a
101
10a
epoch
epoch
Figure 5: Generalization of ALBERT trained with stochastic depth. The stochastic depth improves the length extrapolation to longer sequence lengths.

---

Under review as a conference paper at ICLR 2023
Rand-init BERT
Pre-Trained BERT
Mimickinn BERT
variable
100%
82%
test
70%
304
1pl
10a
10a
10
100
10
10a
opoch
opoch
epoch
Figure 6: Pretrained BERT exhibits significant performance advantages over its Rand-Init counterpart, while the mimicking procedure (a simple initialization scheme we describe below) heads closes the gap.
4.2 RAND-INIT VS. PRETRAINED: STRUCTURAL ADVANTAGES FROM PRETRAINING
Pretraining large models has emerged as a prominent and highly successful paradigm large-scale deep learning. It advocates first training the model followed by task-specific fine-tuning on the task at hand. Our goal here is to use LEGO as a testing ground for this paradigm. To this end, compare (a) training the BERT architecture for LEGO from random initializations to (b) fine-tuning the standard pre-trained BERT model to solve LEGO. Figure 6 (left and center plots) shows that pretraining helps generalization in LEGO dramatically: the pre-trained model generalizes to unseen sequence lengths (the dashed plots) much better, and within far smaller number of epochs, than the randomly initialized model.
large dataset perform generic task,
4.2.1 WHY DOES PRETRAINING HELP IN LEGO?
One simple explanation is that pre-trained BERT is already aware of the semantics of tokens like '=' or ' We have easily ruled out this possibility, by replacing those tokens with arbitrary ones that do not encompass the same semantics; this does not affect the performance of pre-trained BERT. A more intriguing explanation pertains the attention mechanism itself. At its basis, LEGO requires two fundamental types of information transfer:
Asooiiaionn encoding long-range dependencies that transfer value between two occurrences of the same variable. For example, if the input contains the two clauses "a +1" and -a" (with arbitrary separation between them), the architecture must associate the two occurrences of the variable a in order to correctly set to -1.
Manipulation: encoding short-range dependencies of transferring value from the right-hand to the left-hand side the clause. For example, to successfully process the clause "b -a", the architecture must associate these particular occurrences of a and with each other, in order to transfer the value of a (after applying to the group element -1) into b.
Association corresponds to purely global attention pattern, completely reliant on the identity or content of the tokens and oblivious their positions in the input sequence. Manipulation, in trast, corresponds to a purely local attention pattern, where nearby positions attend to each other.
Head
Layer 1 Head
1.0
Figure 7: Visualization of representative attention maps from pre-trained BERT model not yet fine-tuned on LEGO. complete visualization all attention patterns the pre-trained BERT in Appendix [F] On the LEGO input sequence, certain heads implement local, convolution-la lliee manipulation operators (left), while some others implement global, long-range association operators (right). Note that the sample input sequence is presented in the reasoning chain order for visualization purposes only.

---

Under review as a conference paper at ICLR 2023
It is natural to ask whether they are indeed manifested in the pre- pre-trained model's attention heads in practice. Indeed, Fig. [7 shows two exemplar attention heads of pre-trained BERT on an input LEGO sequence without any fine-tuning. The right head clearly depicts association: each token attends to all other occurrences of the same token in the input sequence. This motivates us to make the following hypothesis: the advantage of pre-trained models LEGO can be largely attrihuted to the association and manipulation heads learned during pretraining.
Note that merely the existence of the heads does not fully validate the hypothesis yet. To rule out other factors, we carefully design controlled experiments test this hypothesis the section below.
4.2.2 VERIFYING THE HYPOTHESIS WITH MIMICKING
To test this hypothesis, we conduct the following mimicking experiments.
Mimicking BERT We 'initialize' certain attention heads to perform association and manipulation, without access pretraining data. We achieve this by specifying the target attention matrices (one for association and one for manipulation), and training the model on random data to minimize a "mimicking loss" that measures how well the actual attention matrices at every layer match the target matrices. The precise mimicking loss and training protocol are specified in the Appendix [B.3] The rightmost plot Figure 6 shows that BERT with mimicking initialization attains significant advantage in generalization over randomly initialized BERT, despite not being pre- trairri on any real data (and thus not having learned to "reason"). This confirms that much of the advantage of pre-trained BERT stems from having learned these information transfer patterns.
4.3 SHORTCUT SOLUTIONS AND THEIR EFFECT ON GENERALIZATION
As discussed in Section 4.1, a natural solution to LEGO is to resolve variables iteratively by the order of their depth in the chain. Surprisingly, we find that the Rand-Init BERT and ALBERT models first learn a "shortcut" solution: they immediately resolve the last variable in the reasoning chain, perhaps by counting the total number of minus signs. Indeed, the last variable can be easily identified as it appears only once whereas every other variable appears twice, and its value is fully determined by the parity of the number of minus signs. This behavior is observed in Figure 3a the randomly initialized models are trained to fit all 12 variables: the last variable (#1 1, indicated by the brightest green curves) improves earlier than almost all other ones.
This behavior may be related to the well-observed phenomenon of spurious features: model suc- ceeds in training not relying on any actual features of cows and circumventing the intended solu- McCoy et al. (2019): Srivastava et al. (2020); Gururangan al. (0118); Nguyen et al. (202T).
We use LEGO as case study of shortcut solutions and their effect on generalization. Instead of training the model to fit the first six variables (as in bottom Figure 3 in Appendix), we train it to fit the first five (#0-#4) and the last variable (#1 I). This allows
measure length extrapolation (to #5-#10) in a setting where models can learn the shortcut. The results show significantly degraded performance, implying that shortcut solutions impede generalization. We then study ways to prevent models from learning them, by pretraining and mimicking. The full section appears in Appendix A
ATTENTION WITH HARDCODED PATTERNS
Our analysis in Section 4.2 reveals that the advantage of the pre-trained BERT model on LEGO orig- inates from two specific types of attention structures emerging from pre-training - the association and manipulation patterns. quick examination of all the attention heads depicted in Appendix suggests that there is one more clearly identifiable attention pattern: broadcasting on the [CL.]] token or the [SEP] token (sometimes both). Namely, it "broadcasts' the value inside the special tokens to the others. Even though [CLS]. [SEP] play no role on LEGO per se, they are vital to the pretraining objective as well as many downstream tasks. Thus the broadcasting attention pattern is presumably important for many real-life NLP tasks beyond LEGO. Association, manipulation, and broadcasting consist of a considerable portion of the pre-trained BERT's attention heads, and they are so structured that we can in fact hardcode them efficiently. Hardcoding can be viewed as an extreme version of the previous mimicking procedure.
BERT with hardeoded attention For the association, manipulation, and broadcasting heads, we can efficiently construct the sparse attention matrix based on the input token IDs only, without learn-
8

---

Under review as a conference paper at ICLR 2023
13933
eent
6023
hardccde
hytel
5903
Concel (b)
1943
2003
L near
Altention
Hardoods Attartiinn
mode
Porameter Count
Allention
weythwiss Cor
Ausse;
93
hybri
compune hardcoded
Fiddan
pattams
FOW takenss
23
model
Attention layers
Figure 8: Left: Our proposed architecture in a generic form consisting of 3 pathways. BERT has pathway (b)
only; the Hardcode model has (a) and (b); the Hybrid model has (a), (b), and (c). See Appendix C Right:
Comparison of inference Flops and model size. Flops are measured on batch of 64 sequences of 512 tokens.
ing Q and K or the expensive attention probability computation. For manipulation maps, due their intrinsic locality, we decide to implement them directly with temporal convolutions (along the time dimension). For the other global maps, given raw input sequence of 7 tokens, U1, U2, M2,.... WT N, we manually construct the association and broadcasting maps Aasso. Acts. Asep E IR TXT such that (Jasso)ij [t, ūÿ] (Aeis)s
[u, [CLS]], (Arep) = μ = [SEP] ]] where 1 [-] is the indicator function which outputs 1 if the argument is true and 0 otherwise. In the end, we nor- malize them to have row-wise unit & norm. Notably, the latter three hardcoded attention patterns
require no training (other than a value map for each layer) and are shared across all layers.
vanil a BERT
venille BERT hardccde Attn hybrid
10
hardcode hybrid
training loss
validation loss
2000
4000
6000
8000
2000
4000
6000
8000
step
step
Figure 9: Training and validation performance BERT pertaining task (Masked Language Modelling+Next Sentence Prediction). As a standard, the training sequence length increases from 128 to 512 around the 7k-th step, where the BERT training loss exhibits a sudden bump in response, while the hardcode/hybrid models exhibit remarkable resilience. The hybrid model learns faster and (slightly) outperforms BERT in validation.
On the standard BERT pertaining benchmark, we compare the following three models: BERT-base model, Hardcode model, and Hybrid model. We use convolutional kernel size for the latter two. In Figure 9 we show that the Hardcode model learns fast in the beginning but falls short later However, the Hybrid model not only reduces model size and accelerates inference, but also renders models that are extremely competitive with the base model in terms of the final performance of large-scale pertaining. We follow precisely the training pipeline and hyperparameters of Devlin et al. (2018) See Appendix Aoppendx[fforro architecture details of the Hardcode and Hybrid models.
We observe that the Hardcode model learns faster but gradually falls short, while the hybrid model achieves the best of both worlds: it learns faster at the beginning and achieves even (slightly) lower validation loss at the end. The Hybrid model's validation loss curve appears be lower envelope of the others. The BERTT/arddcoee/Hyrrii models achieve 1.49/1.66/ 1..6/1..7 final pertaining vali- dation loss and 88.2/82.5/ 888.//88 18.1 Dev Fl score on SQuAD v1.1 Rajpurkar et al. (2016). We leave comprchensive evaluations for future work.
CONCLUSION
In this work, we study Transformers by constructing LEGO, a controllable synthetic logical reason- ing task. With LEGO, we have gained insights into their inductive bias, the role of pertaining, etc. Based on these insights, we proposed the hardcoding attention mechanism which both accelerates inference and leads comparable or even better performance. There are many important attention heads beyond just manipulation and association, and their roles remain to be discovered. We believe LEGO will continue to deepen our understanding on Transformers' inner working and to inspire better al algorithms/rrchitecture for tasks beyond LEGO.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Vishesh Agarwal, Somak Aditya, and Navin Goyal. Analyzing the nuances of transformers' poly- nomial simplification abilities. arXiv preprint arXiv:2104.14095, 2021.
Cem Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowycz, Vedant Misra, Vinay Ramasesh, Am- brose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. Exploring length generalization in large language models. arXiv preprint arXiv:2207.0490), 2022.
Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the Ability and Limitations of Trans- formers to Recognize Formal Languages. Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 7096-7116, 2020.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Πla Sutskever, and Dario Amodei. Language models are few-shot learners. In Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901, 2020.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34, 2021a.
Xinyun Chen, Dawn Song, and Yuandong Tian. Latent execution for neural program synthesis beyond domain-specific languages. Advances in Neural Information Processing Systems, 34, 2021b.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.023111 2022.
Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning. What does BERT look at? analysis of BERT's attention. In Proceedings of the 2019 ACL Workshop Black- boxNLP: Analyzing and Interpreting Neural Networks for NLP, pp. 276-286, Florence, Italy, August 2019. Association for Computational Linguisties. doi: 10.18653/11/111444283 URL
4828
Gonçalo M. Correia, Vlad Niculae, and André F. T. Martins. Adaptively sparse transformers. In EMNLP, 2019.
Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber. The neural data router: Adaptive control flow in transformers improves systematic generalization. arXiv preprint 1rXiv:2110.07732, 2021.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Fukasz Kaiser. Universal transformers. arXiv preprint 1807.038 2018.
Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: /810.0480, 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. Glam: Efficient scaling of language models with mixture-of-experts. arXiv preprint arXiv: 2tt::2//2.00 /2.06905 2021.
Suchin Gururangan, Swabha Swayamdipta, Omer Levy, Roy Schwartz, Samuel R Bowman, and Noah A Smith. Annotation artifacts natural language inference data. arXiv preprint arXiv: (80::880300000 02324, 2018.
10

---

Under review as a conference paper at ICLR 2023
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the MATH dataset, 2021.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Train- ing compute-optimat large language models. arXiv preprint arXiv:2203. 15552660..555,, 2022.
Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Con- ference on Learning Representations, 2022. URL pettps://openrevieee net met/foruu forum?
iddnnevveeFYff
Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Q Weinberger. Deep networks with stochastic depth. In European conference on computer vision, pp. 646-661, 2016.
Dieuwke Hupkes, Verna Dankers, Mathijs Mul, and Elia Bruni. Compositionality decomposed: How do neural networks generalise? Journal of Artificial Intelligence Research, 67:757-795, 2020.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger,
Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-588) 2021.
Daniel Keysers, Nathanael Schärli, Nathan Scales, Hylke Buisman, Daniel Furrer, Sergii Kashubin, Nikola Momchev, Danila Sinopalnikov, Lukasz Stafiniak, Tibor Tihon, Dmitry Tsarkov, Xiao Wang, Marc van Zee, and Olivier Bousquet. Measuring compositional generalization: A com- prehensive method on realistic data. In International Conference on Learning Representations, 2020.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep con- volutional neural networks. In Advances in Neural Information Processing Systems, volume 25, 2012.
Brenden Lake and Marco Baroni. Generalization without systematicity: On the compositional skills of lequence-to-sequenee recurrent networks. In International conference on machine learning. pp. 2873 2882.. PMLR, 2018.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Sori- cut. Albert: A lite bert for sule-supe learning of language representations. arXiv preprint arXiv:/909 1942, 2019.
Y. LeCun, B. Boser, J. S. Denker, D. Henderson, R. E. Howard, W. Hubbard, and L. D. Jackel. Buckpropagation applied to handwritten zip code recognition. Neural Computation, 1(4):541- 551, 1989.
Wenda Li, Lei Yu, Yuhuai Wu, and Lawrence C. Paulson. Isarstep: benchmark for high-level mathematical reasoning. In International Conference on Learning Representations, 2021.
Thomas McCoy, Ellie Pavlick, and Tal Linzen. Right for the wrong reasons: Diagnosing syntactic
heuristics in natural language inference. arXiv preprint 192.v:190..01007, 2019.
Thao Nguyen, Vaishnavh Nagarajan, Hanie Sedghi, and Behnam Neyshabur. Avoiding spurious correlations: Bridging theory and practice. In NeurIPS 2021 Workshop on Distribution Shifts: Connecting Methods and Applications, 2021.
Weili Nie, Zhiding Yu, Lei Mao, Ankit B Patel, Yuke Zhu, and Anima Anandkumar. Bongard- logo: A new benchmark for human-level concept learning and reasoning. Advances in Neural Information Processing Systems, 33: 18::16668-1644 16480, 2020.
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowyez, Maarten Bosma, David Luan, Charles Sutton, and Au- gustus Odena. Show your work: Scratchpads for intermediate computation with language models. In Deep Learning for Code Workshop, 2022. URL /tttts:///peeeeeeeeee http////pppeeeieeeeee/ee net/ forum?? id=HB xtel1122dbbbe

---

Under review as a conference paper at ICLR 2023
Arkil Patel, Satwik Bhattamishra, Phil Blunsom, and Navin Goyal. Revisiting the compositional generalization abilities of neural sequence models. arXiv preprint arXiv:2203.07402, 2022.
Stanislas Polu, Jesse Michael Han, Kunhao Zheng, Mantas Baksys, Igor Babuschkin, and Ilya Sutskever. Formal mathematics statement curriculum learning. arXiv preprint arXiv:2202.01334, 2022.
Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint 2r 2021.
Alessandro Raganato, Yves Scherrer, and Jörg Tiedemann. Fixed encoder self-attention patterns in transformer-base machine translation. arXiv preprint arXiv:2002. 10260, 2020.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ questions for machine comprehension of text. In Proceedings of the 2016 Conference on Empirical Methods
in Natural Language Processing, pp. 2383 2392,, Austin, Texas, November 2016. Association for Computational Linguistics. doi: 10. 10.11663/v1/DI6- 1264.
URL /rtts://aclattholooyy
Drg/D16- 1264
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen, Hierarchical text- conditional image generation with clip latents. arXiv preprint arXiv:2204006125, 2022.
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint arXiv:2205.06177 2022.
Alexander Rives, Joshua Meier, Tom Sercu, Siddharth Goyal, Zeming Lin, Jason Liu, Demi Guo, Myle Ott, C Lawrence Zitnick, Jerry Ma, et al. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. Proceedings of the National Academy of Sciences, 118(15), 2021.
Anna Rogers, Olga Kovaleva, and Anna Rumshisky. A primer in bertology: What know about how bert works. Transactions of the Association Computational Linguistics, 8:842-866, 2020.
Swarnadeep Saha, Sayan Ghosh, Shashank Srivastava, and Mohit Bansal. Prover: Proof generation for interpretable reasoning over rules. In Bonnie Webber, Trevor Cohn, Yulan He, and Yang Liu (cds.). Proceedings of the 2020 Conference Empirical Methods in Natural Language Process- ing, EMNLP, pp. 122-136. Association for Computational Linguistics, 2020a.
Swarnadeep Saha, Sayan Ghosh, Shashank Srivastava, and Mohit Bansal. Prover: Proof generation for interpretable reasoning over rules. arXiv preprint arXiv:2010.02830. 2020b.
Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, et al. Using deep- speed and megatron to train megatron-turing nlg 530b, a large-scale generative language model. arXiv preprint 2 11900, 2022.
Megha Srivastava, Tatsunori Hashimoto. and Percy Liang, Robustness spurious correlations via human annotations. In International Conference on Machine Learning, pp. 9109-9119. PMLR, 2020.
Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, Heng-T Tze Cheng, Alicia Jin, Taylor Bos, Leslic Baker, Yu Du, et al. Lamda: Language models for dialog applications. arXiv preprint arXiv:2201. 00.8399 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Infor- mation Processing Systems, volume 30, 2017.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. arXiv preprint arXiv: 19050099000 2019.
12

---

Under review as a conference paper at ICLR 2023
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoc Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. arXiv preprint arXiv:2201.11903, 2022.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gug- ger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the-ar natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 38-45, Online, October 2020. As- sociation for Computational Linguistics. http::/ http://wwwwwww......... acwweb. wrg.anthooogg...
2020 emn w demos.
Yuhuai Wu, Markus N Rabe, Wenda Li, Jimmy Ba, Roger B Grosse, and Christian Szegedy. Lime: Learning inductive bias for primitives of mathematical reasoning. In International Conference on Machine Learning, pp. 11251-11262. PMLR, 2021.
Shunyu Yao, Binghui Peng. Christos Papadimitriou, and Karthik Narasimhan. Self-attention net- 2021.
works can process bounded hierarchical languages. arXiv preprint
Weiqiu You, Simeng Sun, and Mohit lyyer. Hard-coded gaussian attention for neural machine translation. arXiv preprint arXiv:2005.00742. 2020.
Chiyuan Zhang, Maithra Raghu, Jon Kleinberg, and Samy Bengio. Pointer value retrieval: A new benchmark for understanding the limits of neural network generalization. arXiv preprint 1rXiv:2107.12580, 2021.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christo- pher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, Opt: Open pre-trained transformer language models. arXiv preprint arXiv: 200500000, 2022.
13


Reviews:
[{"review_id": "q0bBMB__t1d", "reviewer": "Reviewer_s6py", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "", "summary_of_the_paper": "LEGO is a synthetic reasoning task that targets chain reasoning. Think Markov chain, where the current state only depends on the previous state. They specifically evaluate BERT and ALBERT. ALBERT is particularly interesting since weights are shared across all layers in ALBERT and is hypothesized that it may be particularly suited for this type of reasoning. Indeed, it outperforms BERT. Various ablations are carried out to understand why.", "summary_of_the_review": "I think the paper adds to the understanding of the various inductive biases that different transformer architectures have. There is much interest int the mathematical ability of transformers. I would like to see more models benchmarked on this task.", "strength_and_weaknesses": "The task and setup are framed well. Ablations such as stochastic depth length extrapolation, pretraining, and probing classifier are creative ways at understanding the inductive biases of the networks.\n\nThe task seems not that specific to transformers. I do wonder how other RNNs would fare.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "aGlnIp-IUy", "reviewer": "Reviewer_7SQ8", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper provides empirical investigation of transformer by constructing a synthetic reasoning task -- LEGO (Learning Equality and Group Operations). The LEGO tasks consist of group operations and a chain of reasoning. The major purpose of the synthetic experiment is to provide a controlled setting and unveil what transformer learns. To do so, the paper conduct a series of experiments on the both classic setting (training and testing distribution are the same) and the OOD setting (out of distribution) and they found\n\n(1) Classical generalization happens all the architecture.\n\n(2) OOD generalization depends on the architecture and data preparation.\n\n(3) The non-extrapolating models (architecture fails for OOD) learns short-cut solution.\n\nThe author also propose hardcode attention layer that incorporates knowledge over the dataset that seems to work well on LEGO. ", "summary_of_the_review": "The paper provides solid contribution to the literature and I vote for acceptance.\n\n\nSome minor questions over \"shortcut\":\nThe paper finds transformer sometime learns short-cut solution to the reasoning tasks. However, it is unclear to me why this is a bad thing. I feel it is not equivalent to spurious feature because some reasoning tasks involves short-cut solution or parallelizable solution, e.g., it is possible that a linear size circuit only has depth $O(\\log n)$ (i.e., do not follow the chain of reasoning). Some further discussions should be a good addition on the finding!\n\n\n\n------------\n\nPost rebuttal\n\nI have read the rebuttal and my evaluation remains.", "strength_and_weaknesses": "Strength. The paper provides solid contribution towards understanding the Transformer architecture. By constructing a synthetic benchmark, it provides a series experiment that provide new insight and confirm existing observations. \n\nWeakness. There is no major weakness.\n\nMinor comment. The related work section can be improved, in particular, there is a previous line of work \"unveils\" transformer over regular language ( in particular dyck language) that should be cited and discussed. Examples include [1] (and see the reference therein), which studies the dyck language and provide some theory and empirical investigation on what transformer learns.\n\n[1] Self-Attention Networks Can Process Bounded Hierarchical Languages. Shunyu Yao, Binghui Peng, Christos Papadimitriou, Karthik Narasimhan. ACL'2021\n\n\n\n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "OxFq6AWS-C", "reviewer": "Reviewer_drii", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper proposes a synthetic task that aims to capture aspects of multi-hop reasoning, and uses it to characterize the behavior of Transformer models. It finds that pre-trained transformers learn the task more rapidly, and attributes this to specific patterns of attention in the network. Introducing mechanisms that emulate these specific attention patterns helps close this gap in performance without the need for expensive pre-training. Models struggle with extrapolating to longer chains of reasoning than required during training. This deterioration can be mitigated by repeating layers with the same parameters as well as training with stochastic depth. ", "summary_of_the_review": "My overall recommendation is based on weighing the thorough and well-written exposition/investigation of the LEGO task with the potential downsides of any synthetic benchmark, namely that it may or may not capture the key elements of actual real-world tasks, and can therefore be limited in its impact.", "strength_and_weaknesses": "Strengths:\n- The proposed task is useful as a tool for understanding some ways in which a Transformer can perform multi-hop reasoning\n- The proposed operationalization of requiring length extrapolation is effective at disentangling issues with positional encoding from the question of whether the model can extrapolate to longer chains of reasoning \n- Experiments help provide an understanding for how pre-training, as well as stochastic depth, help improve certain aspects of the models' capabilities\n- Hardcoded attention patterns are shown to be just as effective for solving this specific task\n- The writing of the paper is clear and well-structured\n\nWeaknesses:\n- One of the strengths of the attention mechanism is that it can learn any attention pattern a task may require, and the multi-head version can learn when (at what layer) and how much to use each pattern. The hardcoded attention patterns seem sufficient for the proposed synthetic task -- almost by construction of the LEGO task. There are some preliminary findings that they may also be effective for masked language modeling and at least SQuAD, but this direction is underexplored and perhaps not all of the necessary components for real NLP tasks are captured.\n- The proposed task is perhaps not the best test-bed of *pre-training* specifically, because the approach used is to pre-train on natural language but then test on synthetic expressions. It would not be surprising for language-specific reasoning to not undergo transfer in this case, potentially leaving only the most general of structural patterns, like the attention patterns revealed in this work.\n- Overall, it's not clear how much of the Transformers' capabilities are stressed by the proposed synthetic task. It could be that understanding multi-hop computation in the style of the LEGO task is the key for explaining current behaviors and proposing improvements. Alternatively, multi-hop reasoning may just not be a crucial component of many NLP tasks as currently formulated, or LEGO might not be capturing the actual character of chain-based reasoning as required for real-world task.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "FKAskJ-vD9O", "reviewer": "Reviewer_3VMG", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper proposes a synthetic task, LEGO, and uses it to understand how Transformers operate. The task is in following a chain of simple reasoning, in particular for simple mathematical operations with variables valued either +1 or -1 and expressions of the kind a=+b or a=-b. With this task, the authors consider BERT’s and ALBERT’s abilities to generalize to longer inputs. They find specific attention heads (e.g., local attention patterns) and propose modification with some of the heads fixed to these patterns.", "summary_of_the_review": "The paper introduces a synthetic task to analyse Transformers. In the analysis, finds specific attention patterns and proposes model modification with some of the heads being fixed. Most of the observations done in the paper and model modification have already been done before without resorting to synthetic tasks.", "strength_and_weaknesses": "While the paper introduces the new task to help understanding Transformers, the task does not lead to new observations not widely known before. Specifically, it’s well-known that the most important Transformer attention heads play interpretable roles, including positional ([1], [2], [3], etc). This was used extensively to improve/simplify models: [4] fixes attention window size for some of the heads and gets improvement for low-resource settings (similar to this work), [5] develops hard-coded attention without learned parameters, to name a few. From this perspective, (1) neither finding these attention patterns nor proposing fixed attention is novel; (2) doing it on a synthetic task when this has already been done before for different real-world tasks does not make much sense.\n\n\n[1] ACL 2019: Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned.\n\n[2] EMNLP 2019: ADaptively Sparse Transformers.\n\n[3] BlackBoxNLP 2019: What does BERT Look at? An Analysis of BERT’s Attention\n\n[4] EMNLP findings 2020: Fixed Encoder Self-Attention Patterns in Transformer-Based Machine Translation\n\n[5] ACL 2020: Hard-Coded Gaussian Attention for Neural Machine Translation\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "1: The main claims of the paper are incorrect or not at all supported by theory or empirical results.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject