Systematic Generalization and Emergent Structures in Transformers Trained on Structured Tasks
Authors: Yuxuan Li
James McClelland
Venue: Submitted to ICLR 2023 (2023)
Keywords: systematic generalization
transformers
representation
multi-task learning

Abstract: Transformer networks have seen great success in natural language processing and machine vision, where task objectives such as next word prediction and image classification benefit from nuanced context sensitivity across high-dimensional inputs. However, there is an ongoing debate about how and when transformers can acquire highly structured behavior and achieve systematic generalization. Here, we explore how well a causal transformer can perform a set of algorithmic tasks, including copying, sorting, and hierarchical compositions of these operations. We demonstrate strong generalization to sequences longer than those used in training by replacing the standard positional encoding typically used in transformers with labels arbitrarily paired with items in the sequence. By finding the layer and head configuration sufficient to solve the task, then performing ablation experiments and representation analysis, we show that two-layer transformers learn generalizable solutions to multi-level problems and develop signs of systematic task decomposition. They also exploit shared computation across related tasks. These results provide key insights into how transformer models may be capable of decomposing complex decisions into reusable, multi-level policies in tasks requiring structured behavior.

Under review as a conference paper at ICLR 2023
SYSTEMATIC GENERALIZATION AND EMERGENT STRUCTURES IN TRANSFORMERS TRAINED ON STRUCTURED TASKS
Anonymous authors Paper under double-bling review
ABSTRACT
Transformer networks have seen great success natural language processing and machine vision, where task objectives such as next word prediction and image classification benefit from nuanced context sensitivity across high-dimensional inputs. However, there is an ongoing debate about how and when transformers acquire highly structured behavior and achieve systematic generalization. Here, we explore how well causal transformer can perform set of algorithmic tasks, including copying, sorting, and hierarchical compositions of these operations. We demonstrate strong generalization to sequences longer than those used in training by replacing the standard positional encoding typically used in transformers with labels arbitrarily paired with items in the sequence. We search for the layer and head configuration sufficient to solve these tasks, then probe for signs of systematic processing latent representations and attention patterns. We show that two-layer transformers learn generalizable solutions to multi-level problems, develop signs of systematic task decomposition, and exploit shared computation across related tasks. These results provide key insights into how stacks of attention layers support structured computation both within task and across tasks.
INTRODUCTION
Since their introduction (Vaswani et al. 2017 transformer-based models have become the new norm
of natural language modeling (Brown al. 2020 Devlin al. 2018) and are being leveraged for
machine vision tasks as well as in reinforcement learning contexts (Chen et al. 2021 Dosovitskiy cet al. 2020 Janner et al., 2021 Ramesh et al. 2021). Transformers trained on large amounts of data under simple self-supervised, sequence modeling objectives are capable of subsequent generalization to a wide variety of tasks, making them an appealing option for building multi-modal, multi-task generalist agents (Bommasani et al. 2021 Reed et al. (0022).
Central to this success the ability to represent each part of the input in the context of other parts through the self-attention mechanism. This may be especially important for task objectives such as next word prediction and image classification at scale with naturalistic data, which benefit from nuanced context sensitivity across nigh-dimensional inputs. Interestingly, transformer-based language models seem to also acquire structured knowledge without being explicitly trained to do so and display few-shot learning capabilities (Brown et al. 2020 Linzen & Baroni 202T Manning et al. 2020). These insights have led to ongoing work exploring these models' potential to develop more broad reasoning capabilities (Binz& Schulz 2022
Dasgupta et al. 2022).
Despite success in learning large-scale, naturalistic data and signs of acquisition of structured knowledge or generalizable behavior, how transformer models support systematic generalization remains to be better understood. Recent work demonstrated that large language models struggle at longer problems and fail to robustly reason beyond the training data (Anil et al. 2022 Razeghi et al.) 2022). Different architectural variations have been proposed to improve length generalization in transformers, highlighting the role of variants of position-baa based encodings (Csordás et al. 202 b Ontanón et al./ 2021 Press et al. 2021 Indeed, whether neural networks will ever be capable of systematic generalization without building in explicit symbolic components remains an open question
(Fodor & Pylyshyn. 1988 (Smolensky et al. 2022.

---

Under review as a conference paper at ICLR 2023
Here, we approach this question by training causal transformer model perform set of algorithmic operations, including copy, reverse, and hierarchical group or sort tasks. We explicitly sought the minimal transformer that would reliably solve these simple tasks and thoroughly analyze such minimal solution through attention ablation and representation analysis. Exploring how a transformer with no predefined task-aligned structure could adapt structures in these algorithmic tasks provides starting point for understanding how self-attention can tune to structures in more complex problems, c.g.. those with the kinds of exceptions and partial regularities of natural datasets, where the exploitation of task structures may occur in more approximate, graded manner. Our main contributions are:
1. We present of two-layer causal transformers that are capable of learning multiple algorithmic operations, and show that putting more attention heads at deeper layers has advantages for learning multi-level tasks.
2. We show that the attention layers in these models reveal signs of systematic decomposition within tasks and exploitation of shared structures across tasks.
3. We highlight simple label-bane basel- order encoding method in place of the positional encoding methods typically used in transformers, and show that it helps to achieve strong length generalization performance.
METHOD
item head
Input soquence:
norm
output sequerice
feod furward
copy
cousal transformer
2
add norm
manked mulli-had atfection
group: [shapo]
group: (color)
task embed
Item
label
item embed
latell embed
item ambed
label ambed
task
KDOSS
sort:
foolos shapp, (ectune)
sequerice
quary for 18t ousput takan
Figure 1: Task and model design.
Dataset. We created item pool covering all combinations of 5 shapes, 5 colors, and 5 textures, and generated a sequence dataset by sampling 100k sequences 5-50 items randomly selected from the item pool. The tasks we used to train the models are shown in Fig IT.. Each task corresponds to one of the following rules, which relies on item feature and/or item order information to rearrange an input sequence (grouping sorting items by a particular feature is with respect pre-defined feature sort order, e.g., circles < squares < pentagons, or red < purple < blue):
COPY (C); copy the input sequence. REVERSE (R): reverse the input sequence.
GROUP[SHAPE] (G[s]): group the items by shape, preserve the input order within each shape group. GROUP[COLOR) (G[c]): group the items by color, preserve the input order within each color group. SORT[SHAPE,COL.OO TEXTURE] (S[s]): sort the items first by shape, then by color, then by texture. SORT(COL.OR,, SH.....,,S,.PEEE TEXTURE (S[C)): sort the items first by color, then by shape, then by texture.
We instantiated the token vocabularies as onehot or multihot vectors. The task tokens were onehot vectors with the corresponding task category set to one, with one additional task dimension corre- sponding to the end-of-scquence (EOS) token. The item tokens were multihot vectors whose units indicated its value in each feature dimension (equivalent to concatenated onehot feature vectors). As such, the model receives disentangled feature information in the input, though in principle it can learn to disentangle feature information given onehot encodings for each unique item.
Label-baset order encoding. Using position-bac based order encodings (whether absolute or relative), models trained with sequences up to length I. encounter out-of out-of-distribution problem when tested
2

---

Under review as a conference paper at ICLR 2023
on longer sequences, as position encodings beyond L are unfamiliar to the model. We thus introduce label encoding, which allows longer sequences of items to be encoded with familiar labels. Instead of using item positions, we paired items in each sequence with ascending random integer labels to communicate order information (Fig B). The sampled labels were represented as onehot vectors which the models then embed with learnable label weights. In all reported results, we pre-generated item labels sampled from range up to the maximum generalization length (50) for all sequences in the dataset, and these labels were shared across training steps and model seeds. In practice, the labels for each sequence can be sampled online and from a larger range to accommodate generalization to even longer sequences. We contrast our random label encoding scheme with sinusoidal and learnable encodings based on position indexes.
Model. The main model architecture is shown in Fig BB. Each input sequence consisted of task token and the paired item and label tokens, with the EOS token serving as the first query for tokens in the output sequence. The input tokens were first embedded to the model's latent representational space through set of embedding layers depending on the token type (task, item, or label). The item and label embeddings were then added to form a composite item embedding. These embedded tokens were fed into a causal transformer, which contained one or two layers of alternating future-masked attention sublayers and MLP sublayers. Residual connections and layer normalization were applied after each sublayer as in Vaswani et al. (2017). We tested architectural variations in the number of attention heads in different layers of the model while controlling for the total number of learnable parameters (see detailed hyperparameters Appendix BB. The state of the query token the output of the causal transformer was passed through two linear heads to predict the next output token (the task token, or an item and its associated label).
Training and evaluation. The models were trained using full teacher forcing (where we always feed the model the correct input tokens) on all sequences of lengths 5 to 25 in the dataset (~46k) and evaluated for length generalization on sequences of lengths 26 to 50 (~54k). We trained models in both single-task and multi-task settings. ordered items and their labels given the task being trained, followed by an FOS token. In single-task learning, we did not include the task token in training or evaluation. In multi-task learning, the task token was used and the models were trained to first output the task token before predicting the output sequence. The training sequences used in multi-task learning remained the same ones between lengths 5- 25, but each sequence corresponded to a different output sequence under different tasks.
both cases, the output sequence consisted of the correctly
The models were trained using softmax cross-entrone loss on the prediction of feature classes, labels, and task/FOS categories for tokens in the output sequence. Item predictions were treated average feature prediction accuracy, i.e., if the model predicted 2/3 features correct, its token-level item accuracy is 2/3. Training stopped at 32k gradient updates for single-task models and 38k gradient updates for multi-task models. Below, we report both token-level and sequence-level accuracy, under both teacher forcing and topl rollout (i.e., greedy decoding). Results were aggregated over four random seeds for each task type x architecture pair. Unless otherwise specified, results were taken from the checkpoint with the highest generalization accuracy within each seed. Error shades and error bars indicate standard error of the mean across models.
3 RESULTS
3.1 SINGLE-TASK LEARNING
Two-layer models with label encoding learn the SORT task and generalize to longer sequences. We first trained the model with the SORT[SHAPE.COLON REX]TUU task. Using our label encoding method, models with two single-headed layers (indicated as [1,1]) were able to achieve near- near-ceiling accuracy on training sequences and generalize to longer sequences (Fig 2 also see quantitative results in Appendix C The predictions of the EOS token were also highly accurate in these models (see Fig STA in Appendix I. Item prediction was more accurate than label prediction in this task, reflecting that the models represented item feature information more accurately in order to sort the input tokens. The two-layer models showed some degradation in sequence- accuracy as a function of sequence length, but the failures longer sequences were not catastrophic, as these models scored very well on longer sequences when up to 5% prediction errors were allowed (Fig p also see Fig STB, and Fig S2 for accuracy under rollout in Appendix A.T. In contrast,
3

---

Under review as a conference paper at ICLR 2023
A
B
bwmn precittion
labe predition
D
Item prediction
label prediction
brwiing Hergth 251
genoraiizoo on flrgghh 25 500
3,11
0.4-
[141-4 anaadd postion
12]
[2]
6.2-
bureaa laud
10
20
30
40
50
20
99
50
sezuercc lnrgh
16.
X
cc--
0.8
0.8
60-
30c 236
ra.iing 5000
50p
146
tolkingg Stop
Ex prodittio
secuecce protics
20
41
secuerce prsfice
Jabe predetion
10
35 SE00009CC knggt)
20
40
50
12
soquaroe longtn
20
70
40
55
Figure 2: Token- and sequence-level accuracy for dingle-task models. A. Token-level accuracy training and generalization sequences over learning. B. Token-level accuracy over sequence length. C. Token-level accuracy over sequence positions. D. Proportion of sequences that the model predicted 100% tokens correct (upper) or predicted greater than 95% tokens correct (lower). In B, C, and D. results were taken from 5k training sequences (in and D) and 5k generalization sequences (B, C. and D). Legends indicate the number of attention heads in each layer and the order encoding used (in A). Gray shades indicate the sequences used in training.
two-layer models trained with sinusoidal or learnable position embeddings performed worse across both training and generalization sequences AA)
The two-layer models were also much better than single-layer models with either one or two attention heads. While these single-layer models were able to exploit some correlations between items and output positions (e.g., item [0,0.0] always came first, and item [4,4,4] always came last), they failed to sort items in the middle positions (Fi In contrast, a single-layer, single-headed model was sufficient to learn the COPY or the REVERSE task (see Fig [3A in Appendix AIII, suggesting that multiple layers strongly benefit successful learning of level problems.
A
laye-- hsad=0
layer= head=0
2.
01-
Bor roder Peteria
ouigits
24
3.0
2
145
+43
1epodared IESUENO [] catpul UEQUATAC
Rum E1 Wit S7-
Sources
Sources
Figure 3: Attention patterns in two-layer models. A. Attention maps for an example generalization sequence. Items in the input sequence were reordered to match their output order for visualization purposes. Numbers 1-5 mark the beginning of each shape group. Label indicates the EOS token. B. First-layer attention from query items to input items within shape groups. C. Attention to EOS as a function of item index within each shape group (indicated by labels s1-s5). Results in B and C were aggregated across Ik generalization sequences and across seeds.
Distinct two-stage processing across attention layers. The attention weights in the two-layer models revealed signs of task decomposition 3A). The attention head in the first layer tended to distribute attention to the unsorted items that share the same shape as the current query item. The attention head in the second layer then almost exclusively attended to the next output token in preparation for feature and label readout. This pattern appeared robustly across sequences and across different seeds (Fig B). Interestingly, there was an increase in the attention weights to the BOS token as the model received query items towards the end of each shape group. This attention to EOS increased to similar degrees early or late shape groups (Fig 3C), again suggesting that the model learned to systematically process items within each shape group, token was only relevant after sorting all items. We also found similar attention patterns in two-layer, single-headed models learning the GROUP SHAPE] task (see Fig S33 in Appendix A.
though generating the EOS
GROUP[SHAPE]
4

---

Under review as a conference paper at ICLR 2023
The single-layer models displayed some attention to subsequent items in the output sequence but lacked consistent structures across different shape groups (see Fig Fig/S4jin in Appendix This could be due to the burden for the attention head(s) within a single layer to implement mixture of item contextualization and readout of the correct item or label, and reflects an advantage of the two-layer models in injecting inductive bias for multi-stage solution.
Acquisition of weature order information. To solve the sort task accurately, the models needed learn the invariant sort order within each feature. We tested if the learned order information can be parsed out from the input embeddings. We found that the embedding weights associated with shape and color feature units reflected within-feature order similarities, with weights associated with different features appearing mostly orthogonal (A) The weaker structure and smaller magnitude associated with texture weights may be due to lower demand in sorting multiple texture values, as texture was the third-level sort feature and thus had fewer consecutive values in single output sequence under limited sequence length.
A
TWA GOT product
0.04
2.5
2.2
2.5
0.02
1.0
20k
225
22k
20k
Braniig Skop
Basiing Stop
wiggh Indes
Figure 4: Input item embeddings capture information about feature sort order. Dot product similarity measure of item embedding weights. B. Fitting Gaussian representations track learned feature sort order in the input embeddings. Left, fit objective (mean L2 loss) over pairwise similarities of feature embeddings. Right, best-fitting Gaussian spread parameter for each feature.
We quantitatively tracked how information about feature sort order was acquired in the input embed- dings over learning, using an L2 loss over the difference the pairwise similarities from synthetic Gaussian feature representations and that from the models' feature embeddings. This analysis sug- gested that the two-layer models initially began to acquire feature sort order information concurrently in all three features (Fig 4B, left panel). In later stages during learning, the sort order informa- tion in color embeddings more quickly and better corresponded to similarities between Gaussian representations, with shape embeddings deviating from Gaussian-like, monotonic order similarities.
3.2 MULTI-TASK LEARNING
Multi-task learning and length generalization in two-layer, multi-headed models. To explore the ability for the causal transformer to simultaneously learn multiple algorithmic tasks, trained models to predict different output sequences on the same input sequence conditioned on a task token. The two-layer, single-headed model used in single-task learning was unable to learn all six tasks, while two-layer, multi-headed models achieved good training and generalization performance (Fig also see quantitative performance in Appendix Increasing the number of attention heads in the model did not lead to drastically different learning curves under teacher forcing, but more attention heads supported much better performance under topl rollout 5B).
We tested whether headed attention served for better learning when it occurred in the first layer (attention-frontlee
or the second layer (attention- (attention-backload)) Attention Attention-backload models generally achieved higher accuracy across all tasks compared to their attention-froottone counterparts (Fig D)). Performance of the iirention-backload models was even comparable with their attention- balanced-eeeddd counterparts, which signals that a bottlenecked architecture may be particularly suited for multi-task learning in our setting, considering that some tasks in the task suite share the first-level grouping feature.
Consistent with the single-tast model, six-task models demonstrated strong generalization to longer sequences (Fig (A). Even though sequence-leve accuracy degraded sequence length increased, the models only made less than 5% prediction errors for long extrapolation sequences (Fig 6C). Token-level accuracy among generalization sequences was also relatively stable until the last few
5

---

Under review as a conference paper at ICLR 2023
output positions ()) We also observed that accurately predicting item features in long sequences was easier in the SORT tasks but harder in the COPY, REVERSE, and GROUP tasks (Fig 6 and (C). This echos results from the single-task models and again suggests that the models more accurately represented information directly used for sorting the items in each task. See additional results on EOS prediction and accuracy under rollout in Appendix A.2
1.4
1,
Uihentwe bom
0.2
11.31
14-41
13.50
14.40
0.0
32x
324
424
20k
30k
4000
4012
360
101
Guiningg
traning SHP
attention frontledd
attertion backtoud
3.49
09
0ld 5la S0c1
46
Oct Sol
5d
Stal 04 Fa4 S1c
031 9c 201 94
Figure 5: Token-level item and label accuracy in multi-task models. Legends indicate the number of attention heads in the two layers. A. Accuracy across all tasks learning. B, C, and D. Top1 rollout accuracy for models with attention balanced, attention-fromm frontload, and iicention-backlow architectures. In A, results were taken from 75k generalization sequences (12.5k per task). In C, and D, results were taken from 1.2k generalization sequences (200 per task).
3.2
0.09
6.4
6PC AOUI LOMA CH ueperr
6.2
0.00
requence losgth
sequence length
30
1.00
Baar WAN 1904 scz
0.90
Seegeeet Kar
0.94
0.2
S65
564)
25
20
30
30
3
20
20
20
40
50
S0QU0TOS
20quenic peaitio
Figure 6: Length generalization across tasks. A and B. Token-level item and label accuracy over sequence length and sequence position. C. Proportion of sequences that the model predicted 100% tokens correct (upper) or predicted greater than 95% tokens correct (lower). Results were taken from Ik training sequences (in A and C) and Ik generalization sequences (in A, B, and C) for each task. Performances were aggregated across the five oor-performing runs. Gray shades indicate the sequences used in training.
Learned task embeddings recover task similarity. We examined the input task embeddings to understand the basis of task-shared and task-specific computation in the model. Fig shows that similarities among the learned task embeddings reflected shared computation across tasks. For example, the COPY task was recognized to be similar to the REVERSE task and the two GROUP tasks, potentially reflecting the shared stronger attention to label information in order to accurately sort the input items. Representations for the two SORT tasks were also highly similar, as they both rely item features to sort the input items. The models also recovered similarities between the pairs of GROUP and SORT tasks that share the same first-level grouping feature.
Task-shared and task-dependent computation across attention layers. Structures in the task embedding similarities and strong performance from models with fewer attention heads than there

---

Under review as a conference paper at ICLR 2023
A
ablation head
alow attartion tolems
Sexc lnnoel lotamneeve fiem ecc
Hoache-foocsi totandieze lebel acc
Tokor lewd itom RCC
scc
4x1-
13:10
164
14
501
Tast 9000
50
K eianld 860 ski
Oblidd stll Sic 1ack
tank: Gle) 40-4
tadt: CCS) L1-4
twat: Rir)
taxe: Spk
bade: Stel
L1-1
H1
LI-I
querte
20urc0s
Figure 7: A. Pairwise similarities of task embeddings. B. Teken-level accuracy from ablating single attention heads. C. Token-level accuracy from preserving attention to certain tokens across attention heads. D. Attention maps from selected attention heads in an example generalization sequence (see Fig[S7for for all attention maps). Items corresponding to the input sequence were reordered to match their output order for visualization purposes. Numbers 1-5 mark the beginning of each first-level feature group (shape color). Label indicates the task token. Label indicates the BOS token. In B, C, and D, results were taken from the top-performing model (architecture: [1,4]).
tasks already hinted that the model could be exploiting shared processing across tasks. Because task-conditioned computation can only occur across the attention layers, we next sought to understand the role of multi-headed attention in implementing such shared computation. We performed two ablation experiments: ablating a single attention head entirely and preserving attention in all other heads, or preserving attention certain tokens across all attention heads. Attention weights were ablated by masking the attention weights as zero after softmax.
Fig B and C show the ablation results for the top- cop-perforning model with one attention head in the first layer and four attention heads in the second layer. The attention heads did not exhibit strong selectivity for single tasks as they usually contributed to multiple tasks, and they also showed equal contribution to item and label prediction (Fig 7B, also see results from other models in Fig [S8 in Appendix p. One attention head appeared redundant as ablating it resulted in little impact on the performance of any task under teacher forcing, but it significantly improved accuracy under topl rollout.
Ablating attention to certain items in the sequence further indicated that the models were heavily relying on the learnable task embedding to contextualize items under different tasks (Fi also see Fig (S). When the models were only allowed to attend to the task token in the first layer and the next output token in the second layer, performance was largely preserved in the GROUP tasks and the SORT tasks (ablation type mossLL++nexLI")) This is different from the two-layer single-task models, whose first layer relied on attention to other items in the first-level feature group to cross-contextualic items. However, there were still some signs of task decomposition in the multi-task model. For example, one attention head in the second layer of the top-performing model consistently attended to the first item the next feature group across all multi-level tasks (Fig 7P; see detailed attention maps in Fig S7 in Appendix A.2), which proved useful for accurate predictions across multi-level tasks under both teacher forcing and rollout (Fig B).
Task-conditioned item encoding. To better understand the effect of fask-conditioned computation on encoded item representations, we analyzed how the items were represented under different tasks after each layer. We computed average item representations of all 25 shape x color feature pairs in each task, aggregating across texture and label values in 200 generalization sequences. Projecting these fask-conditioned item representations to lower dimensions revealed different embedding structures across the two layers (Fig For the two GROUP tasks, representations in the first layer encoded all items by their relevant first-level grouping feature (shape or color). The item representations
7

---

Under review as a conference paper at ICLR 2023
did not consistently encode other feature information across both layers, presumably because label information became more relevant for sorting items within each feature group.
shape and orders
eapo
proup jhapp)
greas Isder
seri jahaae colorooosared sgm celeramneecereue
layer 0 ercoding
ioutt wribed
layer 1 ancoding
20
Figure& 8: Task-conditioned encoding of the same sequence of items from the top-performing model. PCA was performed for average item representations across 200 generalization sequences within each layer x task pair. The legend indicates the sort order of the shape and color features.
Interestingly, for the two SORT tasks, the encoded representations after the first layer also showed a high degree of clustering in the corresponding first-level grouping feature, with the secondary sort feature represented in ring structure. The subsequent layer then expanded the item representations along the secondary feature in a way that was consistent across all first-level feature groups. Notably, the clustering effect and the ring structure in the first-layer representations in the SORT tasks were more strongly observed in the attention-backlon suggest that room for flexible multi-headed transformation downstream relax first-layer encoding, enabling a more consistent solution to emerge across the GROUP and SORT tasks.
models (see Fig [S8 in Appendix A..)) This could
4 RELATED WORK
There is a growing interest in analyzing small models to better understand the capabilities and detailed computations in transformers in more controlled settings. For example. Power et al. (2022) explored learning and generalization dynamics in two-layer causal transformers, and Elhage et al. (2021) explored mechanistic interpretability in one- and two-layer transformers without MLP sublayers. Our work contributes to these efforts in beginning to shape some understanding of the computation and representation dynamics in amall-scale transformers.
Ablation experiments and variety of representation analyses have been applied to understand the role of the attention mechanism in NI.P tasks as well as in transformer-basee vision models (Chefer et al. 2021 Manning et al.) 2020 Michel et al.) 2019 Voita et al.) 2019). However, it is often debated to what extent attention weights afford model interpretability in these settings (Jain & Wallace 2019) Wiegreffe& & Pinter 2019 Vashishth et al. 2019). especially considering head redundancy and the difficulty in correctly attributing relevance over high-dimensional inputs. We show that at least in simple settings, the attention heads can exhibit some level of interpretability consistent with known task structures. Similar methods have also been applied to understand unit-level and layer-level dynamics that support multi-task computation in small-scale RNNs (Driscoll et al. 2022 Yang et al.) 2019, which revealed some patterns that are consistent with our findings here, as we discuss below.
Recent work examining systematic generalization in transformers or pre-trained language models highlighted that length generalization remains a challenge and observed that positional encoding can have significant impact on the models' ability systematically generalize (Anil al., 2022
Csordás et al., 2021a la Delétang et al., 2022 Ontanón et al. 202T). A concurrent work explored the randomized position encoding method (equivalent to our label-based encoding method) and showed that it outperforms a variety of sequence order encodings (Anonymous) 2022). Other types of architectural modifications have also been proposed to help transformers achieve better length generalization (c.g. CSordás et al. 202 Dehghani et al. 20188 Press al.) 2021 I Su et al. I)2 Our work adds to the continued effort to enable better length generalization in transformers by demonstrating the potential in formulating sequence modeling tasks with more general item-label binding approach rather than item- posiiion binding.
Outside of the context of transformers and language models, the of synthetic, algorithmic tasks has enabled much understanding of the core capabilities of many models (c.g. Graves et al. (0044). There has also been some interest in performing algorithmic reasoning with neural networks for its
8

---

Under review as a conference paper at ICLR 2023
own sake (Veličković & Blundell 2021). Correspondingly, Veličković et al. (2022) recently proposed a benchmark for algorithmic reasoning, and evaluated a variety of graph neural network architectures, all of which struggled to extrapolate algorithms to longer or larger inputs.
5 DISCUSSION
We sought to understand how transformers can solve a set of highly-structured tasks and systematically generalize. We presented two-layer causal transformers that can learn copying, reversing, and hierarchical sorting operations and generalize to sequences longer than seen during training. We found that these models learned to exploit shared task structures and exhibited interesting signatures of task decomposition. The first layer tended to contextualize the input items with the task token or with other items in a feature group, and the subsequent layer was more responsible for within-group or task-specific processing.
Our findings the contributions of individual attention heads in a multi-headed layer in solving these algorithmic tasks echo results from analyses of language models, For example, single attention heads are rarely responsible for a particular task or syntactic relationship and can often appear redundant (Manning et al. 2020; Michel et al. 2019 Voita et al. (019)). We do see some degree of selectivity, with conceptually distinct task components shared across subset of attention heads. Recent work studying multi-task computation in RNNs has similarly found that multi-task learning led to the exploitation of reusable computation across related tasks (Driscoll et al. 2022 Yang et al.. 2019). Interestingly, these recurrent models also exhibited high degrees of task-selectivin at the level of individual units in the hidden layer. The degree of task-sclcttivity in different architectural components may vary depending being learned. Future work is needed to fully characterize the degree of computational modularity in multi-headed attention in relation task structures.
the inductive biases of different model families and the tasks
We highlight that the label-based order encoding method was key to enabling our models to generalize the learned tasks longer sequences. The key insight is to sample random labels to communicate sequence order information instead of relying on sequence positions. This simple extension exposes models a large range of possible labels evenly during training so they can learn to encode items in longer sequences with familiar labels. Compared to this approach. learnable position embedding leaves untrained embedding weights unseen positions and thus would result in poor length generalization. However, learnable position encoding also resulted in slower learning on the training sequences in our tasks- - potentially because the models had fewer sequences to learn to encode later positions compared to early ones. Sinusoidal position embeddings have been noted to have limited length generalization capabilities (Ontanón et al., 2021 Csordás et al., 10)la). In our setting, it may have additionally suffered from the demand to read out the absolute item positions, which is not common across NLP tasks. A concurrent work similarly noted the out-of out-oo-diistibution problem with position-based encoding and demonstrated the superiority of the randomized position approach (equivalent to our label encoding method) for length generalization in a range of different algorithmic tasks (Anonymous) 2022). Although label-based encoding alone would fail represent true item distance information, in light of recent work showing that transformers can learn positional information without explicit positional encodings (Haviy al. 2022). it is promising that the advantages of label-based encoding may transfer to natural language inputs.
The flexibility to learn and perform multiple tasks is a key desired capability for machine learning. Our work provides insights into the possible structures of within-task and cross-task computations that stacks of attention layers can afford through a thorough analysis of the dynamies in a simple transformer on a set of highly-structured tasks. Recent work has explored multi-task learning in transformers at scale and achieved impressive results (Lee et al. 2022, Reed et al. 2022). As transformers are increasingly being leveraged for multi-task multi-modal learning in domains with richer, more complex task structures, it is possible that these models may implicitly learn to decompose complex decisions into reusable, multi- multi-level policies. In future work, we hope to explore these learning and generalization dynamics in transformer based agents to understand the acquisition of condcondit,,n,,, multi-level behavioral policies in structured environments.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Cem Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowycz, Vedant Misra, Vinay Ramasesh, Ambrose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. Exploring length generalization in large language models. arXiv preprint S.Kiv:2077009901,, 2022.
Anonymous. Randomized positional encodings boost length generalization of transformers. 2022.
Marcel Binz and Eric Schulz. Using cognitive psychology to understand gpt-3. arXiv preprint arXiv:2206. 14576, 2022.
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportuni- ties and risks of foundation models. arXiv preprint arXiv:2108.07258, 2021.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Hila Chefer, Shir Gur, and Lior Wolf. Transformer interpretability beyond attention visualization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 782-791, 2021.
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:1 5084-15097, 2021.
Róbert Csordás, Kazuki Trie, and Jürgen Schmidhuber. The devil is in the detail: Simple tricks improve systematic generalization of transformers. arXiv preprint arXiv:2108. 12284, 2021a.
Róbert Csordás, Kazuki Trie, and Jürgen Schmidhuber. The neural data router: Adaptive control flow in transformers improves systematic generalization. arXiv preprint 2rXi::211000732,22222 2021b.
Ishita Dasgupta, Andrew K Lampinen, Stephanie CY Chan, Antonia Creswell, Dharshan Kumaran, James L McClelland, and Felix Hill. Language models show human-like content effeets on reasoning. arXiv preprint arXiv:2207.07053 2022.
Mostafa Dehghami, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Eukasz Kaiser. Universal transformers. arXiv preprint / 2018.
Grégoire Delétang, Anian Ruoss, Jordi Grau-Moya, Tim Genewein, Li Kevin Wenliang, Elliot Catt, Marcus Hutter, Shane Legg, and Pedro A Ortega. Neural networks and the chomsky hierarchy. arXiv preprint arXiv:2207.022: 02009,, 2022.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010. /1929, 2020.
Laura Driscoll, Krishna Shenoy, and David Sussillo. Flexible multitask computation in recurrent networks utilizes shared dynamical motifs. bioRxiv, 2022.
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Satfield-Dodds Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. http::/transformee--crcuiii / html.
10

---

Under review as a conference paper at ICLR 2023
Jerry A Fodor and Zenon W Pylyshyn. Connectionism and cognitive architecture: A critical analysis. Cognition, 28(12)) 1988.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv:1410.5401, 2014.
arXiv preprint
Adi Haviv, Ori Ram, Ofir Press, Peter Izsak, and Omer Levy. Transformer language models without positional encodings still learn positional information. arXiv preprint arXiv:2203. 16632 2022.
Sarthak Jain and Byron C Wallace. Attention is not explanation. arXiv preprint arXiv: 1018. 2019.
Michael Janner, Qiyang Li, and Sergey Levine. Reinforcement learning as one big sequence modeling problem. In ICML 2021 Workshop on Unsupervised Reinforcement Learning, 2021.
Kuang-Huei Lee, Ofir Nachum, Mengjiao Yang. Lisa Lee, Daniel Freeman, Winnie Xu, Sergio Guadarrama, Ian Fischer, Eric Jang, Henryk Michalewski, et al. Multi-game decision transformers. arXiv preprint arXiv:2205.1522: 15241, 2022.
Tal Linzen and Marco Baroni. Syntactic structure from deep learning. Annual Review of Linguistics, 7:195-212, 2021.
Christopher Manning, Kevin Clark, John Hewitt, Urvashi Khandelwal, and Omer Levy. Emergent linguistic structure in artificial neural networks trained by self-supervision. Proceedings of the National Academy of Sciences, 117(49):30046330055 2020.
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? Advances in neural information processing systems, 32, 2019.
Santiago Ontanón, Joshua Ainslie, Vaclav Cvicek, and Zachary Fisher. Making transformers solve compositional tasks. arXiv preprint 3xXx:2100.04377 2021.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gen- eralization beyond overfitting on small algorithmie datasets. arXiv preprint arXiv:2201.02177, 2022.
Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108. / 12409, 2021.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning, pp. 8821-8831. PMLR, 2021.
Yasaman Razeghi, Robert L Logan IV, Matt Gardner, and Samcer Singh. Impact of pretraining term frequencies on few-shot reasoning. arXiv preprint arXiv:2202.07206 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.06/17 2022.
Paul Smolensky, Richard Thomas McCoy, Roland Fernandez, Matthew Goldrick, and Jianfeng Gao. Neurocompositional computing: From the central paradox cognition to new generation of ai systems. Al Magazine, 2022.
Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. arXiv preprint S10iv::2004009644 2021.
Shikhar Vashishth, Shyam Upadhyay, Gaurav Singh Tomar, and Manaal Faruqui. Attention inter- pretability across NLP tasks. arXiv preprint /909. 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones. Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Petar Veličković and Charles Blundell. Neural algorithmic reasoning. Patterns, 2(7):100273, 2021.

---

Under review as a conference paper at ICLR 2023
Petar Veličkovié, Adrià Puigdomênech Badia, David Budden, Razvan Pascanu, Andrea Banino, Misha Dashevskiy, Raia Hadsell, and Charles Blundell. The CLRS algorithmic reasoning benchmark. arXiv preprint arXhv:2205. 2022.
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:1905.00418. 2019.
Sarah Wiegreffe and Yuval Pinter. Attention is not not explanation. arXiv preprint Xiiv:1908.04626. 2019.
Guangyu Robert Yang. Madhura R Joglekar, H Francis Song, William T Newsome, and Xiao-Jing Wang. Task representations in neural networks trained to perform many cognitive tasks. Nature neuroscience, 22(2):297-3066 2019.
12

---

Under review as a conference paper at ICLR 2023
A ADDITIONAL RESULTS
A.1
SINGLE-TASK LEARNING
B
1,000
YV YVW
1.0
50sievel ope ast (owi)
0.0
=002
0.996
0.4-
0.992
22
0 990
30
40
position
length
length
Figure S1: Additional teacher forcing accuracy. A. Single-task models predict BOS token near perfectly. B. Sequence-level accuracy when up to 10% errors were allowed. Results were taken from 5k training sequences (in B) and 5k generalization sequences (in A and B).
A 1.-
P
teNen loved Jabe
D.B
0.6
0.8
0.4
0.4
Wost
11.11
0.0
0.0
1.0
500 bus (ra4190)
1.0
0.8-
0.8
0.6
0.4
tahutskeed
0.4
022
0.0
20
20 position
40
longth
Figure S2: Topl rollout accuracy for two-layer models. A. Token-level (upper) and sequence-level (lower) accuracy across sequence length. Token-level accuracy across sequence positions. Results were taken from 5k training sequences (in A) and 5k length generalization sequences (in A and B).
A
training
generalization
B
LD-HC
10
1.4
00
0.6
02
0.2
d
3
L1 HO
10kaeo abet accoeEy
0..8
0.
8.4
0.4
10k
30x
20x
32k
0.0
tra ming deep
training deop
Figure S3: Learning other tasks in the single-task setting. A. Single-layer, Engle-headed models can learn the COPY task and the REVERSE task, while two-layer, single-headed models learn the GROUP[SHAPE] task. B. The attention maps from the GROUP[SHAPE] model resemble that from the SORT[SHAPE model (visualization in Fig A
13

---

Under review as a conference paper at ICLR 2023
Figure S4: Attention maps from single-layer models in an example generalization sequence. A. Single-layer. single-headed model. B. Single-layer, two-headed model.
A.2 MULTI-TASK LEARNING
2.0
M0
822
seauence engh
oeguenice bength
2.1
2000 boonteeebomaa
dee SpC
14
Sta
Gl4l
10
20
30
50
20
30
aREuANTCe langth
MMQUETCR Bngth
but item and label predictions were obtained using topl rollout. Figure S5: Visualization in Fig 6
A
tokln lord -002 SC2
toter-iond <001
0.59
0.38
0.6
0.27
0.4
0.96
G[3]
Si6]
30
20
30
50
10
20
30
50
sequence positoo ttom
sequence posttion
Figure S6: A. Six-task models predict the EOS token perfectly under teacher forcing. B. Prediction of the EOS token deteriorates under topl rollout in six-task models. Results were taken from 1k training sequences and 1k length generalization sequences across the five top-performing runs.
14

---

Under review as a conference paper at ICLR 2023
LG
task: copy
Lask: reNNne
Lonk: shopps
task grouptcolor]
task: sontshapp]
task: sontfcolor
Figure S7: Attention maps for an example generalization sequence (for the top-performing multi-task model with [1,4] architecture). Each row shows the attention maps for one task, and each column corresponds to a single attention head denoted by the layer and head index. Tokens corresponding to the input sequence were reorderd to match their order in the output sequence for visualization purposes. Label indicates the task token. Label e indicates the EOS token. In the bottom four rows, number labels ranging from 1-5 mark the beginning of the items in each first-level feature group.
15

---

Under review as a conference paper at ICLR 2023
[4.4]
architacture:
los alta moan tolann
Jtaroo Text
62
DHOWY www.aateccoeee
architecituee [6.6]
architecture: [4,4]
T
recoding
architocture: bocee
architecture:
poop bcheec byeid eeooding
(raoor grodding
boysi sronding
Figure S8: Attention ablation and task-conditioned item representations in other top-performing models. Visualization in Fig 7 and Fig 8
16

---

Under review as a conference paper at ICLR 2023
B HYPERPARAMETERS
Table S1: Model and experiment hyperparameters.
Hyperparameter
Single-task learning
Six-task learning
Number of layers
or 2
2
Number of attention heads
(see paper)
Embedding dimension
128 (two-layer) or 184 (single-layer)
192
MLP hidden layer size
64
Activation function
ReLU
Batch size
128
Training teacher forcing rate
1.0
Optimizer
Adam
Learning rate
10-1
5 10-4
C
QUANTITATIVE PERFORMANCE
Table S2: Token-level item and label accuracy across 54k generalization sequences in single-task models learning the SORT ISHAPE,COLOR TEXTURE] task. Mean and standard deviation across four random seeds are shown for each architecture.
task
architecture
position encoding
item prediction
label prediction
C
label
99.03=+0.41
100000++0.00
R
label
99.75=+0.21
100.00++0.00
S[S]
label
76.49:12.99
51.75
S
2
label
91.011+2.76
78.02-45.89
S[s]
[1,1]
label
99.60=0.07
98.4110.233
G[S]
[1,1]
label
98.41::0.32
99.3110.16
S[s]
[1,1]
sinusoidal
57.04+11075
16.50+6.52
S[s]
[1,1]
learnable
13.15=111.55
40.51
17

---

Under review as a conference paper at ICLR 2023
Table S3: Token-level prediction accuracy across 75k generalization sequences in six-task models (including 500 unique sequences for each length evaluated within each task). Mcan and standard deviation across four random seeds are shown for each architecture. Model architectures are denoted by [# heads in first layer, # heads in second layer].
(a) Attention models.
Tasks
[1,1]
[2.2]
[3.3]
[4,4]
[6,6]
item
13.08
97.04=0.88
98,55=00,66
98.79-10.94
99.11:10.32
label
87.21
9..4611.24
97.98=0.42
98.68-10.4
99.17:10.22
item
89.28+11.44
97
98.8120.77
98.98:1.95
99.13:1.97
label 97.39:10.84
98.5910.99
99.46-10.39
99.6410.34
99.79:10.17
item 64.30:124.46
96.9511.21
98.95=L0.45
98.69-10.4
99.20.10.60
R
label 94.56:8.226
98,53+0.80
99.6410.16
99.39-10.29
99.73-10.19
item 86.47413.51
96.5110.93
98.4210066
98.52:11.15
98.44
G[s]
label 91.54++.30
9..83+1.07
99.07+0.33
9..2
99.05+0.30
item
86.30+13.39
96.72+0.80
98.68
98.45+1.13
98.62+0.47
G[c]
label
91.6947.79
97.93+1.06
98.98+
99.21+0.46
99.22:10.24
item
86.97+8.80
9.028+1.00
98.48:0029
99.01+0.82
99.59=10.25
S[s]
label
74.75+11.47
93.21+2.23
95.53+1.01
97.09+1.99
98.62:10.22
item
86.7718.66
97.20=
98.37=0.33
99.08
99.62-10.13
S[c]
label
73.57=10.59
92.73+2.05
95.23+0.90
97.53+1.71
98.65
(b) Ceention-frontlaad models.
Tasks
[2.1]
[3.1]
[4.1]
[6.1]
item
69.45:117.65
69.05
17.86-22.11
0..76.10.16
label
79.95418.53
18.35+11.60
86.63
97.52.10.1
item
57.93-127.09
54.074.30.36
42.014330:48
98.50
label 9..6
95.18-3.28
98.48+1.110
99.61
item 57.97+26.94
55.33+29952
70.49
98.78
R
label 95.74+13.34
94.90-3.50
97.97+2.17
99.56-10.23
item 70.16+:1..79
16.29
79.20=19.34
96.85-10.23
G[s]
label 80.56:111.79
12.57
82.52
98.96:10.26
item
70.33+16.54
16.99+16.21
78.86
96.50-1..47
G[c]
label
19.18
78.71
87.29
98.82-10.22
item
79.94=11.66
18.00=112449
83.99-15.22
98.03100.23
S[s]
label
62.76+11.85
11.17
16.69
94.1711.774
item
80.03-111.15 78.22=11.211
82.46=16.50
97.9510.29
S[c]
label
12.88
61.27+18.88
74.27=24.19
94.08-10.22
(c) Attention-bacee backload models.
Tasks
[1,2]
[1.3]
[1.4]
[1,6]
item
96.06:10.44
98.65
0.80
99.40 0.22
99,29
0.17
all
label
95.48-10.30
97.37-0.89
99.21
99.12-0.36
item
96.86=0.57
99.10-0.74
99.46
0.34
19.22-0.55.
label 97.78
99.62
99.85
99.59-0.55
item 95.99-0.59
98,82=1.08
99.47=0.37
99.56-0.14
R
label 98.09+0..2
99.55 0.39
99.83
0.00 0.04
99.82 0.09
item
12.02.8
98.28-1.00
98.9810.43
98.96-0.16
G[s)
label 97.70+0.42
98.77=0.66
99.29 0.27
99.25 0.07
item
95.82-1.72
98.53-0.76
99.1710.34
98.78-0.40
G[c]
label
97.51=0.51
99.06=0.42
99.51
99.13=0.32
item 95.86.10.77
98.51
99.68 0..26
0.13
S[s]
label 91.1211.83
93.54=2.44
98.40+1.11
98.41-0.89
item 95.72+0.54
98.70-0.71
0.28
99.63 0.29
S[c]
label
90.7440.66
93.76=2.15
98.36+1.05
98.54=1.10
18


Reviews:
[{"review_id": "megJ_DHKBWK", "reviewer": "Reviewer_yZ2V", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": 5, "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": "In this paper, the authors explore Transformer's systematic generalization in algorithmic tasks, including copy, reverse, and hierarchical group or sort operations on an input sequence. The authors create a set of tasks and show that a two-layer Transformer successfully learns these tasks and generalizes to sequences longer. Particularly, a random label-based order encoding method, in place of the positional encoding, improves the systematic generalization of Transformer on the studied tasks.", "summary_of_the_review": "The paper is of limited novelty and lacks comparison with previous methods. Therefore, I recommend rejection.", "strength_and_weaknesses": "# Strengths\nThe paper provides a thorough analysis of a small Transformer on a set of algorithmic tasks.\n\n# Weaknesses\n1. The paper is of limited novelty. The major technical novelty is the proposed random label-based order encoding method. However, there lacks an explanation and insights why it works.\n\n2. The authors do not conduct any experiment on the well established datasets, like SCAN. Besides, there is no baseline in this work. The authors should have comapred with Transformers using relative positional encoding (Csordas et al, 2021).", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "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": {}}, {"review_id": "JvhMstv8w0", "reviewer": "Reviewer_1ohG", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "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 introduces a novel positional encoding that improves transformers' systematic (length) generalization capabilities on a set of algorithmic tasks, including copying, reversing, sorting, and grouping. Concretely, the paper replaces the learnable position embedding method with random labels from longer sequences, thus allowing the model to encode the longer sequences with familiar labels. The paper empirically demonstrates that the proposed method achieves better length generalization performance than learnable and sinusoidal position embeddings. Finally, the paper analyzes the attention matrices of the trained models and shows that they implement the algorithmic operations required to solve the corresponding tasks.", "summary_of_the_review": "Given the insufficient comparison to related work, which calls into question the practicality and validity of the results, I do not recommend accepting the paper in its current form.", "strength_and_weaknesses": "**Strengths**\n\nThe paper proposes a novel position embedding method that increases length generalization on a series of algorithmic tasks compared to learnable and sinusoid position embeddings. To that end, the paper introduces a new suite of algorithmic tasks and considers the single-task and multi-task settings, which are interesting. Moreover, the paper conducts an in-depth analysis of the model behavior on the different tasks and shows that, e.g., a two-layered transformer trained on the sorting task learns to first sort shapes and then the other features.\n\n**Weaknesses**\n\nThe paper fails to discuss and compare to relevant related work:\n* The paper does not compare its position embedding method to ALiBi encodings (Press et al., 2021), which were proposed to increase the length generalization capabilities of transformers on natural language processing tasks.\n* The paper does not compare its position embedding method to the Neural Data Router (Csordás et al., 2022), which is an extension of the transformer architecture (consisting of shared layers, gating, geometric attention, and directional encodings) that significantly improves the systematic generalization capabilities of transformers on a wide range of algorithmic tasks.\n* The paper does not discuss the length generalization benchmark introduced by Delétang et al. (2022), which evaluates transformers on a highly related set of tasks, including reversing, copying, and sorting. In particular, Delétang et al. (2022) assess the length generalization capabilities on significantly longer sequences (training on lengths up to 40 and evaluating on lengths 41 to 500). Given the accuracy decrease over sequence lengths (see Figures 2 D) and 6 C)), I am not convinced that the proposed encodings meaningfully increase length generalization for substantially longer sequence lengths.\n* It would be interesting if the paper would discuss recent work showing that transformers can learn positional information without position embeddings (Haviv et al., 2022). \n\nMoreover, while the paper shows that its proposed position embedding increases the transformer's length generalization on a set of algorithmic tasks, it does not evaluate its effectiveness on natural language processing tasks. In particular, it would be interesting to see if the proposed embedding can also be used to \"train short and test long\", as demonstrated by the ALiBi encoding (Press et al., 2021).\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "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": "C9D5Cf-T72u", "reviewer": "Reviewer_ygnx", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "This paper proposes an encoding method that replaces the standard positional encoding in transformers. The label-based order encoding method achieves strong generalization to sequences longer than training ones.\n\nThis paper also analyzes two-layer causal transformers to learn multiple algorithmic operations. It finds that multi-level task learning improves when more attention heads at deeper layers.\n\nThe paper also finds that these models have signs of systematic decomposition within tasks and exploitation of shared structure across tasks.", "summary_of_the_review": "This paper proposes a new encoding method that outperforms a widely used one and has informative analyses and findings.\n\nHowever, the reviewer does not recommend acceptance because of the abovementioned weaknesses.", "strength_and_weaknesses": "***Strength***\n\n- This paper proposes a simple yet effective new method that has improved over a state-of-the-art and widely used method.\n- It analyzes how transformers address the generalization problem.\n- It finds how to improve the performance (e.g., putting attention heads at deeper layers) and signs of systematic phenomenons.\n\n***Weakness***\n\n**The weakness mainly comes from the concern of whether the results and findings in the particular setting still hold in more complicated cases.**\n\n**1. The advantage over conventional positional encoding**\n\nThe conventional positional encoding was designed for natural language inputs.\nFor example, it contains information on relative position or distance in input, as the context window is helpful for natural language.\nThe experiment in this paper does not have natural language input, and the tasks are not much related to relative position or distance in input.\nSo it is not very convincing that the advantage still exists in other tasks, e.g., natural language processing.\n\n**2. Disentangled item representation**\n\nThe items contain shape, color, and texture.\nHowever, the item units indicate the value in each feature dimension.\nSo the item representations are already disentangled, e.g., color and shape values are in different input nodes.\nAddressing entangled data (e.g., the images of the items) is a critical expectation for systematic generalization.\nThe model behavior on disentangled data may not be the same as on entangled data.\n\n**3. Model design**\n\nIt is not very convincing that findings in two-layer transformers will naturally extend to deeper transformers.\nThe limit of depth may prevent the model from having complicated behaviors.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made 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": {}}]

Author rebuttal:

Decision: Reject