Representations and Computations in Transformers that Support Generalization on Structured Tasks
Authors: Yuxuan Li
James McClelland
Venue: Accepted by TMLR ()

Abstract: Transformers have shown remarkable success in natural language processing and computer vision, serving as the foundation of large language and multimodal models. These networks can capture nuanced context sensitivity across high-dimensional language tokens or image pixels, but it remains unclear how highly structured behavior and systematic generalization can arise in these systems. Here, we explore the solution process a causal transformer discovers as it learns to solve a set of algorithmic tasks involving copying, sorting, and hierarchical compositions of these operations. We search for the minimal layer and head configuration sufficient to solve these tasks and unpack the roles of the attention heads, as well as how token representations are reweighted across layers to complement these roles. Our results provide new insights into how attention layers in transformers support structured computation within and across tasks: 1) Replacing fixed position labels with labels sampled from a larger set enables strong length generalization and faster learning. The learnable embeddings of these labels develop different representations, capturing sequence order if necessary, depending on task demand. 2) Two-layer transformers can learn reliable solutions to the multi-level problems we explore. The first layer tends to transform the input representation to allow the second layer to share computation across repeated components within a task or across related tasks. 3) We introduce an analysis pipeline that quantifies how the representation space in a given layer prioritizes different aspects of each item. We show that these representations prioritize information needed to guide attention relative to information that only requires downstream readout.

Published in Transactions on Machine Learning Research (08/2023)
Representations and Computations in Transformers that Sup- port Generalization on Structured Tasks
Yuxuan Li
liguzuan Sytaafon@sttaoord.coo
Deportment of Psychology
Stanford University
James L. McClelland
jince ford. edu
Department of Psychology Stanford University
Reviewed on OpenReview: https: // openreview. net/ fort/ id= oFC2LAgS6Z
Abstract
Transformers have shown remarkable success in natural language processing and computer vision, serving the foundation of large language and multimodal models. These networks capture nuanced context sensitivity across high-dimensional language tokens image pixels, but it remains unclear how highly structured behavior and systematic generalization can arise in these systems. Here, we explore the solution process a causal transformer discovers as it learns to solve set algorithmic tasks involving copying, sorting, and hicrarchical compositions of these operations. We search for the minimal layer and head configuration sufficient to solve these tasks and unpack the roles of the attention heads, well as how token representations are reweighted across layers to complement these roles. Our results provide new insights into how atlention layers in transformers support structured computation within and across tasks: 1) Replacing fixed position labels with labels sampled from larger set enables strong length generalization and faster learning, The learnable embeddings those labels develop different representations, capturing sequence order if necessary, depending on task demand. 2) Two-layer transformers can learn reliable solutions to the multi-level problems explore. The first layer tends to transform the input representation to allow the second layer to share computation across repeated components within task across related tasks. 3) We introduce analysis pipeline that quantifies how the representation space in a given layer prioritizes different aspects of each item. We show that these representations prioritize information needed to guide attention relative to information that only requires downstream readout.
1 Introduction
Since their introduction, transformer-basse based models (Vaswani et al. 2017 have become the new norm of natural language modeling (e.g. BBronn al. 2020 Devlin et al. 2018) and are being leveraged for machine vision tasks as well as in reinforcement learning contexts (e.g. Chen et al. 2021 Dosovitskiy et al. 2020
Janner et al.) 2021). Transformers trained on large amounts of data under simple self-supervised, sequence modeling objectives
capable of subsequent generalization to wide variety of tasks, making them appcaling option for building multi-modal, multi-task, generalist agents (c.g. [Bommasani et al. 2021 Recd el al. 2022). In particular, large-size Lransformer-based language models have shown impressive genera inneralizatinn beyond language tasks, in arithmetic tasks, code generalization, and sometimes reasoning and planning
problems (c.g. Bubock al., 2023 Chowdhery et al. 2022 OpenAl 2023). These findings have lod to ongoing
work that evaluates the broader reasoning capabilities in these models (c.g. Binz and Schulz 2022 Dasgupta
et al.
2022 Dziri et. al. 2023).

---

Published in Transactions on Machine Learning Research (08/2023)
Central to this success is 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 with naturalistic data such as next-word prediction and image classification, which requires the exploitation of nuanced context sensitivity across high-dimensional inputs. Interestingly, analyses of transformer- based language models suggest that they can acquire knowledge of syntactic structures without being explicitly trained to do so [Lnzen and Baroni, 2021 Manning et al. 2020p.
Despite success in learning large-scale, naturalistic data, how well transformer models support systematic generalization and their internal computation remains to be better understood. Recent work has demonstrated that transformer-ban
language models struggle with 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-bast based encodings (Csordás et al.] 2021a Ontanón et al. 2021 Press et al., 2021). Indeed, whether neural networks will ever be capable of systematic generalization without building explicit symbolic components remains an open
question (Fodor and Pylyshyn 1988
Smolensky et. al. 2022
Our goal in this paper is to explore what solutions transformers develop to solve structured tasks, and assess the roles attention heads and token representations play in supporting such solutions. We approach this question by training a causal transformer to perform a set of algorithmic operations, including copy, reverse, and hierarchical grouping and sorting tasks. We explicitly seek the minimal transformer that would reliably solve these tasks either in a single-task setting or in a task-cued, multi-task setting. We then analyze such minimal solutions through attention ablation and representation analysis. Exploring how a transformer with pre-defined task-aligned structure can adapt to structures in these algorithmic tasks provides starting point for understanding how self-attention can tune to structures in more complex problems, e.g., those with the kinds of exceptions and partial regularities of natural datasets, where the exploitation of task structures may occur in a more approximate and graded manner. Our main contributions are:
1. We introduce simple label-based encoding method that samples random integer labels from label pool larger than the maximum sequence length and uses encodings of these labels in place of position- based encodings. We show that this method helps transformers achieve faster in-distribution learning and strong length generalization compared to position-based encoding methods, and that learnable label encodings come to represent sequence order in tasks requiring sequence order information.
2.
We show that, in support of robust and generalizable solutions to the algorithmic tasks, two-layer cansal transformers develop signs of systematic decomposition within tasks and exploitation of shared structures across tasks. For example, when learning lo solve a lask that sorts items by ilem features, two-layer transformers learned a two-step algorithm where attention patterns focused on one item group (e.g, all circle items) at a time. When two-layer models learn multiple tasks at once, each attention head contributes to multiple related tasks (e.g., all tasks that rely on item features).
3. We propose a novel representation analysis pipeline, using singular value decomposition and variance partitioning to quantify how the representations evolve across layers. Using this method, we find that transformers use sequence contextualization to re- weight task-relevant information to dominate the representation space. For example, when only item feature information is needed to sort the sequence, item representations in the intermediate layer prioritize representing item feat over item order, while in tasks where only order is required for sorting, item order is prioritized over item features.
2

---

Published in Transactions on Machine Learning Research (08/2023)
2 Method
Datasct. We created an item pool covering all combinations of 5 shapes, 5 colors, and 5 textures and generated a sequence dataset by sampling 100k sequences of 5 50 items randomly selected from the item pool. The tasks we used to train the models are shown in Figure [A. Each task corresponds to one of the following rules, which relies on item feature /nd/or item order information to rearrange itemss the input sequence (grouping or sorting by item feature is with respect a pre-defined feature sort order, c.g., circles < squares < pentagons, or red purple < blue):
COPY (C): copy the input sequence.
REVERSE (R): reverse the imput sequence.
GROUP[SHAPE) (G[s]): group the items by shape, preserve the input order within each shape group. GROUP GROUP[COLOR] [COLOR] (G[cl): group the items by color, preserve the inpul order within each color group. (SRA(SHAPR,COMM TEXTURE] (S(s,c,T)): sort the items first by shape, then by color, then by texture. TEXT,,
SORT COLLR,,OLO SOOOLORR,SIAEE,TETTURRE (S(c,s,T)): sort the items first by color, then by shape, then by texture.
We instantiated the token vocabularies onehot and multihot vectors. Task tokens were onehot vectors with the corresponding task category set to one, with one additional task dimension corresponding to the end-of-sequence (BOS) token. 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.
A task suite
B model architecture
item head
input sequence:
attention layer
task
output sequence
attartion layer D
capy
feed forward
reverse
norn
group: [shape]
future-maskad attantion
group: [color]
label embed
label embed
embed
embed
embed
embedi
embed
Ishapo, color, taxture)
task
<eos>
sort: [oolor, shape, texture]
imput sequence
1st output probe
Figure 1: Task and model design.
Label-based encoding. Models using position-based encodings trained with sequences up to length L encounter out-of-distribution problem when tested longer sequences (in either absolute relative position-based settings). We conjectured that this problem arises because useful internal representations have not been learned for positions beyond those encountered in sequences up to length L. To overcome this problem, we introduce label-based encoding, which instead pairs items in each sequence with random integer labels sampled from a set labels larger than sequence positions. Specifically, we initialize a pool of K available labels. For a sequence of N tokens, we randomly sample out of K labels (with K > N) and use their corresponding label encodings encode the sequence. In our experiments, the sampled labels are always assigned in ascending order to the items in each sequence. Though this scheme does not capture exact position information, it is sufficient to convey sequence order information, since each successive item is always associated with a larger integer label than its predecessor. The sampled labels could, alternatively, be assigned to items in random order, in which case they would function simply as arbitrary item labels. Either way, randomly sampling labels allows the model to learn to form representations during training for set of labels larger than the maximum training sequence length. As we show below, this scheme supports length generalization in the tasks we consider with very mild degradation as function of sequence length up to the number of available integer labels.
3

---

Published in Transactions on Machine Learning Research (08/2023)
In experiments, we use learnable embedding weights the label encodings, and contrast the label encoding method with sinusoidal and learnable embeddings based on fixed position indices. Concurrent work also explored the same method as our label-based encoding method (termed randomized positional encodings, see Ruoss et al. 2023 and tested it with other types of embeddings. In all reported results, we used the 100k sequences mentioned above and pre-generated item labels from a range up to the maximum generalization length in the dataset (K = =50). In practice, the labels for each sequence can be sampled online and from a larger range to accommodate generalization to even longer sequences.
Model. The main model architecture is shown in Figure DB. Each imput sequence consisted of task token, sequence of the paired item and label tokens, and an EOS token. The input tokens were first embedded in the model's latent representational space through a 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 layers of alternating future-masked attention sublayers and MLP sublayers. Residual connections and layer normalization were applied after each sublayer in Vaswani et. al. (017). We lested 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 in Appendix B). The EOS token served as the first probe for tokens in the output sequence. The state of the probe item at 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. All models were trained using teacher forcing, i.e,, we always feed the model the correct prior sequence of tokens as inputs during training. The training set consisted of all sequences of lengths to 25 in the dataset (~46k), and models were evaluated for length generalization of lengths 26 to 50 (~54k) and novel sequences in the training length range. We trained models in both single-task and multi-task settings. In both cases, the output sequence consisted of the correctly ordered items and their labels given the task being trained, followed by an EOS token. In single-task learning, we did not include the task token in training or ovaluation. In multi-task lcarning, 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 were the same length 5 25 sequences from the dataset, but the corresponding output sequence changed depending on the task.
sequences
The models were trained using softmax cross-entropy losses on the prediction of feature classes, labels, and task/Eos categories for tokens in the output sequence. Item predictions were treated as 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 top1 rollout (i.e., greedy decoding). We ran four random seeds for each task architecture pair. Unless otherwise specified, results were taken from the checkpoint with the highest generalization accuracy within cach seed. Error shades and error bars indicate the standard error of the mean across models.
3 Results
3.1
Label-based encoding supports length generalization and order representation
We first demonstrate the effectiveness of the label-based encoding method. We trained two-layer, single-headed models on the AP,,HAA,,, method, these models were able to achiove ncar-ceiling accuracy on training sequences and generalize to longer sequences with comparable performance, but models trained with sinusoidal or learnable position encodings performed worse across both training and generalization sequences (Figure 2 A4; also quantitative results in Table S2h.
task in the single-task setting. Using the label-based encoding
We found that label embeddings learn different. representations depending on whether sequence order information is needed to solve the task. In the (SAPEESH,E,,OLON COLOT,, task, item positions in a sequence do not affect their output position. However, in the COPY, REVERSE, and GROUP tasks, the model needs to rely on item order in the input sequence correctly sort the sequence. We thus compared the label embeddings learned by the single-task model trained on the SORr (ORr/SHAAR,,OO.OH TEXTURE] task and the
4

---

Published in Transactions on Machine Learning Research (08/2023)
A
training: length 5-25
generalization: length 26-50
U S of label embedding (sar([scc. madel)
C
U of label embeddiing (multi task model)
tatendloval Item acc
1.0
0.8
0.6
dimension
30
dimension
10
0.4
learnable pas
0.2
sinusoidal
20
20
learnable
0.0
10k
20k
30k
10k
20k
30k
10
20
30
40
50
30
training step
training step
label
label
Figure 2: A. Label-based encoding enables fast learning and strong generalization compared to other methods. B and C. Left singular vectors (scaled by singular values) from the singular value decomposition of the learned label embeddings. The decomposed top dimensions of the label embeddings in the multi-lask model show Fourier modes from low to high frequency.
multi-task model that learned to solve all six tasks. We used singular value decomposition to inspect the most prominent modes in each set of embeddings. The label embeddings in the multi-lask model learned a strong order representation, marked by Fourier modes of increasing frequency in the top decomposed dimensions (Figure 2C). This typical for a chain- or ring-like structural form (Saxe et al., 2019). In contrast, the label embeddings associated with the single SORT task model resemble more independent labels and do not contain explicit order information (Figure 23). In this case, length generalization still occurred with these unstructured representations as input order information is not needed to specify the sort order, and the model relied on developing structured attention patterns to solve the task (discussed below).
3.2 Successful solution relies on multiple layers and shared attention heads
Label embedding is not the only factor that contributes finding robust solution for these tasks that generalize to longer sequences. As shown in Figure D architectural variations in the number of model layers and the placement of multiple attention heads can also have significant impacts on the resulting performance. When learning only the (URT[SH
TEXTURE] lask, two-layer, single-headed models (designated as [1,1]) were much better than single-layer models with either one or two attention heads (designated as [1] and [2), Figure 3). The singlo-layer models were able to exploit some correlations between items and output positions (e.g., that erpt-c-eep always came first, and blue-stripedlcross always came last), but they failed to sort items in the middle positions (Figure A1,, right). In contrast, single-layer, single- headed model was sufficient to learn the COPY or the REVERSE task (see Table S2 in Appendix Cl, suggesting that multiple layers strongly benefit successful learning of multi-level problems.
When only appended with a task token, the two-layer, single-headed model that succeded in single-task learning was unable to learn all six tasks at once. Instead, two-layer, multi-headed models achieved good performance and were able generalize all tasks longer sequences (Figures 3 see quantitative performance in Table S3). In addition, models with multi-headed attention in the second layer generally achieved higher accuracy across all tasks compared to models with multi-headed attention in the first layer (Figure Models with multi-headed attention in the second layer were comparable with models that have multi-headed attention in both layers. This signals that bottlenecked architecture where the attention head is shared in the first layer may be particularly suited for multi-task learning in these tasks, considering that some tasks share the first-level grouping feature.
In both single-task and multi-task settings, the fop-performing two-layer models showed some degradation in sequence-level accuracy as function of sequence length, but the failures on longer sequences were not catastrophic, as these models scored very well on longer sequences when up 5% prediction errors were allowed (Figures A2 and] 3 C2)).
5

---

Published in Transactions on Machine Learning Research (08/2023)
Performance of sort[s.c.t] models
Performance of multi-task models
A1
token-level accuracy
multt-headed attention in
C1
token-level accuracy
token- evel item ucc
1.0
1.0
both layers.
token-leve levee item ace
1.00
singletast
mult-task
0.98
0.6
architecture
multiitask token-ieved item (rallout)
0.96
[3.1]
[4.4)
0.94
2-
Gi-]
G[c]
S[s,0,t]
S[c.8.t]
10
30
50
25
50
1.0
1st layer only
0.90
10
30
40
50
25
50
sequence length
sequence position
sequence length
sequence position
A2
sequence-level accuracy
0.5
C2
sequence-level accuracy
1.0
erro sequence
erfor within sequence x3%
[3.
12.11
[4,1]
16-1
see level item ace
1.00
singletast
0.8
0.6
[1.1]
1.0
2nd layer only
muls task
0.75
sequence +3%
seglevel
0.4
0.50
R Stc.s.t1
0.5
0.25
Gis]
0.0
[1:15]
[1:4]
0.00
within Sequence
error
S1s,6,0]
10
40
20 40
40 sequence length
20
30
50
10
20
30
40
50
sequence langth
sequence lngtth
0.0
sequence length
Figure 3: Performance of single (SRT(,,C,T] task models and six-task models. A1 and C1. Token-level item accuracy over sequence lengths and sequence positions. A2 and C2. Sequence- item accuracy the proportion of sequences where models make no error (left) or make less than 5% errors (right) within sequences. B. Multi-task models' token-level item accuracy from topl rollout on 1.2k length generalization sequences across all tasks (200 per task). Models are arranged in ascending order of performance. See more accuracy under rollout in Figures S1] Legends in A and B indicate the number of attention heads in each model layer. Legends in C indicate the six tasks. Gray shades indicate the range of lengths used in training. Results in A and C within the training length range are from novel sequences.
3.3 Early layers transform and reduce the task to support repeated computation in later layers
Are these models performing different computations across layers solve multi-level problems? We closely inspected the attention patterns of the top-performing models to understand the solutions they found. As discussed below, this analysis revealed distinct multi-stage processing across successive attention layers. In both single-task and multi-task settings, we found that earlier layers lended to transform token representations such that subsequent layers can deploy similar computation across items and across tasks.
Within-task dccomposition in the single-task model. Figures fA and AB show the learned solution process of the top-performing two-layer, single-headed model that successfully solved the (OAAPESAAE, TEXTURE] task in the single-task setting. When given a probe item after the pre- sentation of the input sequence, the attention head in the first layer distributed attention to the unsorted items in the same shape group the current probe item. The attention head in the second layer then almost exclusively attended to the target next item in preparation for feature and label readout. This pattern appeared robustly across sequences and across different seeds. Interestingly, there was also an increase in the attention weights to the EOS token as the model received probe items towards the end of each shape group (Figure (C). This attention to FOS increased similar degrees early or late shape groups, again suggesting that the model Icarned to systematically process items within each shape group, even though generating the BOS token was only relevant after sorting all items.
We found similar hierarchical attention pattern in two-layer, single-headed models trained to perform the GROUP[SHAPE) task (see Figure C2C in Appendix A). In contrast, the single-layer models lcarning the SORT|SHAPE,COL (OT,, TEXTURE] lask displayed some atlention to subsequent items in the output sequence but lacked consistent structures (see Figures S2A and S2 This could be due to the burden for the attention head(s) within single layer to implement mixture of item contextualization and readout of the correct target information, and reflects an advantage of the two-layer models in providing the capacity develop a structured, multi-stage solution.
6

---

Published in Transactions on Machine Learning Research (08/2023)
A Illustration of learned solution process
B
8
laye=1 head=0
1.0
shape and color orders:
layer encoding
item head
queries
outoit outoi pitbes
0 1 R
0.6
0.8
MLP
0.4
022
#
6
0.0
layer encoding
layer=0 head=0
MLP
queries
cutpue pinsas
4
0.0
Weights
antention
811
0.1- 0
*
0.0
Input embed
0.1
neordored Input
curput
P-10 n-5
<EOS>
sources
probe within shape
TEXTURE]. A. An illustration of a single step the learned solution process across layers on an example sequence, oquivalent to one row in B. Item textures are visualized as edge colors. B. Attention maps for predicting all target items in a generalization sequence. Items in the input sequence were reordered to match their output order for visualization purposes. The five shapes mark the beginning of cach shape group. Label e indicates the EOS token. C. Attention to EOS increases as models generale items in each shape group. Results in C were aggregated across 1k generalization sequences and across all four two-layer models.
Figure 4: Single-task model's solution to
Shared computation across tasks in the multi-task model. Consistent with the single-task models, multi-task models also developed distinct multi-stage solution. Strong performance from models with fewer attention heads than there are tasks already hinted that the model could be exploiting shared processing across related tasks. Indeed, the attention heads in the top-performing model (with one attention head in the first. layer and four atlention heads in the second layer) show little lask selectivity: ablating each attention head resulted in impaired performance in multiple tasks, and the drop in performance is usually to similar degrees for each pair of related tasks (Figure 5A; see ablation results from other models Figure S3 in Appendix A). The attention heads also contributed equally to target feature and label prediction (Figure AA left and middle). Some heads appeared less important for accurately predicting the next item under teacher forcing, but the performance score here is mislcading. In fact one of the heads (L1-H1) was crucial for shifting from the current top-level group to the next group; its ablation led to systematic failure to generate the correct first item in the next group. Teacher-forcing the first item of each group provided this information to the model, allowing it to sort the remaining items in each group correctly, but without this head the accuracy for the multi-level problems is significantly impaired under topl rollout (Figure 5 right). We next discuss the weight patterns observed in these attention heads.
The attention head in the first layer was heavily relying the task token to transform item representations from the same input sequence differently under each task (Figure LO-HO). As basis of this task-dependent contextualization, the pairwise similarities between the learned task embeddings indeed reflected similarities on the piece of item information (feature or label) that is used for sorting the items in each task (Figure 5
Although the attention pattern in the first layer of the multi-task model differs from that in the two-layer single-task models, it effectively allowed the attention heads in the second layer to perform similar processing across differrnt tasks (see Figure S4 for full attention maps across all heads). For example, some heads in the second layer collectively attended to the target next item in the input sequence (marked by the diagonal attention weights distributed across multiple heads), together resembling the attention pattern in the second layer of the single-task model (Figure 5C, L1-HO and L1-H3). The multi-task model also developed a different signature of within-Lask decomposition. As mentioned above, one atlention head in the second layer played a crucial role in correctly shifting to the next top-level feature group, Its attention pattern showed consistent attention to the first item in the next feature group as the model generated target items the GROUP and SORT tasks (Figure 5c, L1-H1). These results suggest that the multi-task models are simultancously transforming tasks into similar ones and reducing tasks that are multi-level.
7

---

Published in Transactions on Machine Learning Research (08/2023)
A attention head ablation
G[s]
S[s.c.t)
S[c,s,]]
learned task embeddings
L1-H1 L1-H2
Lo-Ho L1-H0 L1-H1 L1-H2 L1-H3
L1-H0 LI 11 L1-H2 L1:H3
1.0
drop in accuracy from origina perfommance
0.8
R-
G(+)
0.6
-0.4
6(c)
-0.6
S[6]
0.2
0.8
Stc]
-1.0-
G ([s] Std]
C attention patterns
task: G[s]
G[c]
task: G[c]
L1-H3
task: G[s]
S(c, t]
Q-
irppA bequence
P
paopat--oieee grotees
querias
0.6
queries
ourjast
# O
0.2
O
to À
U
X
ipplt sequance
sequence
rearbered Inpur Sources
SOUrCes
Sources
Figure 5: Unpacking the solution process of the multi-task model. A. Token-level accuracy across all tasks from ablating single attention heads. B. Pairwise similaritics in the learned task emboddings. C. Partial allention maps from generalization sequence in selected atlention heads (see Figure (44 for full attention maps in all heads). Items corresponding to the input sequence were reordered to match their output order for visualization purposes. The five shapes and colors mark the beginning of each feature group. Label indicates the task token. Label e indicates the EOS token. In A and C, results were taken from the top-performing model (architecture: [1,4]). Results in B were aggregated across the five top-performing models (across architectural variations).
3.4 Token contextualization prioritizes relkvatlevant information in the representation space
How are token representations transformed as result of and in support of the structured attention patterns in these models? As shown in Figure A there were systematic discrepancies in how accurately and quickly the multi-task model predicted output features and labels under different tasks, suggesting that the model prioritized representing the sort-relevant information more robustly over information only for downstream readout. We next examined the transformation of item representations over layers to confirm this hypothesis and to explore the extent to which the tasks are functionally treated as one canonical item rearrangement task. As we discuss in more detail in the next few paragraphs, the results reveal that a stronger represental antienttaiion of sort- relevant information in the model reflects a fundamental feature of its learned solution: when contextualizing items differently under different tasks, the model re-weighted task-relevant information to dominate the representation space in support of downstream shared computation.
Partitioning variance in the representation space. We devised representation analysis pipeline to quantify how different pieces of information compete in the representation space. For a given layer, we performed singular value decomposition (SVD) on the collection representations associated with the output state of each item at that layer. This included all 6250 possible token embeddings (125 items the input embedding layer and contextualized item representations from sampled sequences in subsequent attention layers. Representations in each layer were first zero-centered separately for each unit in the layer, that the dimensions of the SVD correspond to dimensions from principal component analysis (PCA). Then for each dimension in the SVD, we computed the proportion of the variance in its unit-variance left singular vector that. was explained by each feature type (color, shape, or lexture) and by label information. These explained portions of variance were then scaled by the variance associated with that dimension. This allowed us quantify which aspect of item information dominated cach dimension, and then sum these quantities across dimensions to assign a total variance-explained ratio to each piece of item information in a layer's representation space.
labels)
8

---

Published in Transactions on Machine Learning Research (08/2023)
A differential learning of feature/label prediction
task-cond token representations
copy
groupts]
grouptc]
sortt.c..]
oogy
reverse
groupte]
groupto)
sorge, c.ll
sont|c...
sortic,s.t]
0.98
0.99
1.00
0K
10k
20k
30k
token level accuracy
stops reach 95% acc
B
Inout embedding sorl[s.c.]] model
sppo
mult-llsk model
02
verfence
0
10-
2000
30
60
30
60
80
30
90
30
60
30
60
90
30
60
90
dimension
10-
1888
within-group item representation similarity
input embee
layer 0
layer 1
5558
1000
cosine similaity
shape
prggoa
wethe
60
S(s.c.)]
Glel
dimonsion
dimonsion
Figure 6: Models prioritize representations for sort-relevant information. A. Token-level feature and label prediction accuracy on generalization sequences (left) and the time taken to reach 95% accuracy (right). B. Learned imput embeddings from the top-performing single-task model (trained on S(s,C,T]) and the multi-task model. Top, the lower-dimensional projection of input item feature embeddings. Ttem textures are visualized as edge colors. Bottom, variance explained in each dimension of the full decomposition of all 6250 (125 items x 50 labels) imput embedding representations. C. Representaaions of input and probe items in layer 0 of the multi-task model. Top, the dimer-dinansional projection of these representations. Bottom, the variance-cxplainee profile. D. Mean within-group pairwise similarity of input item representations in the multi-task model. In B and C, blue lines represent the total variance explained in category across all dimensions. Weaker dimensions beyond those displayed are excluded for visualization purposes.
Feature/label information in the input embeddings. Our results suggest that item representations as early in the task-agnostic input embeddings are sensitive to the general relevance of different pieces of information. Figure 6B shows the parianco-explainee profile across the decomposed dimensions of the input embeddings. The multi-task model weighted shape, color, and label information (more relevant for sorting the items) more than texture information (mostly for readout), with shape and color information occupying the top-most dimensions. This variance-explained profile stands in stark contrast to that of the input embeddings in the single SORT task model, where item feature information together dominated the top dimensions of larger variance, and label information is strictly separated from feature information and occupiod weaker dimensions.
Task-conditioned token representations in the intermediate layer. Applying the same method to items sampled from 200 length-30 sequences, obtain full view of how representations evolved under cach task in layer 0. In this layer, the output state of the source items and the probe items jointly reflect a re-weighting of the sort-relevant information ion, with the total amount of variance explained by each feature or label reflecting their relative importance for sorting items in any given task (Figure GC). For example, in the GROUP[SHAPE] and OOR[(SHAPE,COLOR COLT,,,TURRE)) tasks, shape information occupies the top-most dimensions, whereas color information occupies the top-most dimensions in the GROUP[COLOR, |ORT[CoL.OR. (COLOR, TEXTURE] tasks. Meanwhile, label information brought top dimensions
and SORI

---

Published in Transactions on Machine Learning Research (08/2023)
in the four tasks that rely on item order to sort the items (COPY, REVERSE, and the GROUP tasks). The stronger representations of the grouping features in the GROUP and SORT tasks resulted in an increase in the within-group item similarity in the intermediate layer, i.e., a stronger sense of item clustering following contextualization (Figure 6D). Notably, within-group item similarity also increased for the secondary grouping Teature in the SORT lasks to reflect an elevated but relatively weaker sense of item clustering.
We used a cross-decoding analysis to explore the possibility that the model is mapping different feature information to a shared functional encoding for related tasks. We trained linear classifier to extract shape category from layer 0 probe item representations under the GROUP GROUP|SHAPE] task, and tested whether this classifier perfectly transferred to classifying item color in the GROUP [COLOR] task (chance level: 20%). The results show that such a classifier was above chance but only somewhat effective when classifying the color of the probe items in the GROUP[COLOR] task (accuracy: 57%),, equaling the transfer to decoding color category in the SORT|SHAPE, COLOR, TEXTURE] task (accuracy: 59%). Instead, the classifier best transferred to extracting shape category in the SORT|SHAPE, COLOR, TEXTURE] task (aceuracy: 100%). These results indicate that the model still largely preserved fealure-unique information and may only be partially features to a shared encoding. Although specific attention heads
transforming these funetionnllly-equivalee could in principle rely more on the shared encoding information, at the level of item output states in this layer, it seems that the model remains sensitive to the different tasks and is only partially transforming related tasks into the exact same one.
Probe and target representations in the final layer. The re-weighted representational characteristics of the probe items in layer 0 carried over to layer 1, although representation for the target item dominated the layer 1 representation space (Figure S5h. For this analysis, we turned to set of 2000 randomly-sampled two-item sequences to decouple feature correlation between the probe and target items. Here, the variance- explained profile of the layer target item representations is more similar than that in layer 0, suggesting that the model projected target item representation into a more shared space for readout. Notably, the model performance on these two-item sequences was lower than on longer sequences. The performance on the single-level COPY and REVERSE lasks remained at 99% and 98%, but performance on the multi-level tasks (GROUP and SORT tasks) dropped to around 90%. This stark contrast in performance reflects that two-item soquences are challenging for the model, the model developed a strategy to rely the multi-level grouping structure to identify the sort order of given sequence, which is more suitable for sorting longer sequences.
4 Related Work
There is growing interest in analyzing small models in more controlled task settings to better understand the capabílitics and detailed computations in transformers. For example, Power et al. (2022) explored learning and generalization dynamics in two-layer causal transformers learning binary operations, and EThage el 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 representations and computations transformers can develop to support generalization structured tasks.
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 a significant impact on the extent to which models can systematically generalize (Anil et al. 2022 Csordás et al. 2021a
Delétang et al. 2022 Ontanón et al., 2021 Many types of architectural modifications have also been proposed to help transformers achieve better length h generalization, including different ways to represent positional
information (e.g. Csordás et al. 2021b Dehghani et al. 2018 Press et al. 2021 Su et al. 2021 Our label-
based encoding method adds to this effort by demonstrating the potential of formulating sequence modeling tasks under a more general item-label binding approach rather than itcm-position binding. Concurrent to our work, Ruoss et al. (2023) developed randomized position encoding (equivalent to our label-based encoding method) and showed its advantage over a variety of position- position-based encodings. We view our work as complementary to theirs, in that they demonstrated strong length generalization to longer sequences (with up to 500 elements) and over a range of tasks, while we have uncovered how these labels are differentially encoded and used in service of performing different tasks.
10

---

Published in Transactions on Machine Learning Research (08/2023)
Outside of the context of transformers and language models, the use of synthetic, algorithmic tasks has enabled much understanding of the core capabilities of many models (e.g. Graves et. al.] 2014). There has also been some interest in performing algorithmic reasoning with neural networks for its own sake (Veličković
and Blundell 2021). Corrspoondnggy; /Velčkovié et al. (2022) recently proposed a benchmark for algorithmic reasoning, and evaluated a variely of graph neural network architectures, all of which struggled to extolll algorithms to longer or larger inputs. Our work shows that self-attention has the potential to adapt to structures in the sequence and find reliable solutions algorithmic tasks.
The representation analysis method we employ here builds on the use of singular value decomposition in Saxe ct al. (2019) to visualize multi-dimensional representational structure and to relate the learning dynamics in deep neural networks to the strengths of the singular dimensions in the SVD of the training data. We have extended this analysis by partitioning the variance in each singular dimension as associated with representing item feature and order information, then summing the variance across dimensions to shed light on how different kinds of information are prioritized in solving tasks with different structures. Below we consider how the approach in Saxe et. al. (2019) and related more recent work might be further extended to address representation and learning dynamics in more naturalistic tasks and datasets.
Ablation experiments and a variety of analysis methods have been applied to understand the rolc of the attention mechanism in NLP tasks as well as in transformer-based 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 and Wallace 2019 Wiegreffe and 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, 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 computations that support multi-task computation in small-scale RNNs DDissooll et al. 2022; Yang et al. 2019). which revealed some patterns that are consistent with our findings here, as we discuss below.
5 Discussion
We sought to understand how transformers can solve a set of highly-structured algorithmic tasks and systematically generalize. We presented two-layer causal transformers that can learn copying, reversing, and hierarchical sorting operations that generalize to sequences longer than seen during training. Detailed analysis of the attention patterns suggest that these models tended to rely on similar computation to sort items in different feature groups (within-task decomposition) and also to sort items in different tasks (across-task) shared computation). Using a novel representation analysis pipeline to quantitatively partition variance in the representation ion space, we also find that the atiention layers learned to re-prioritize task-relevant information in task-dependent manner, which helps attention heads in subsequent layers multitask.
We highlight that the use of random sample of ordered position labels was key to enabling our models to learn robust solutions to these tasks and generalize the learned tasks to longer sequences. The key insight is to sample random labels to communicate sequence order information instead of relying on fixed position indices during learning. This simple extension exposes models to a large range of possible labels during training so that longer sequences can be encoded with familiar labels, and is shown to be effective both in our tasks and in a range of different algorithmic tasks (Ruoss et al. 2023). Compared to this approach, learnable encodings tied to fixed sequence positions resulted not only in poor length generalization, but also in slower learning on the training sequences in our tasks. We suggest that this likely occurred because the models experienced relatively fewer opportunities to learn the position encodings that only appeared in longer sequences (e.g., position encodings 1-5 always appeared in all training sequences, but position encodings 20-25 appcared much less frequently). This poses a problem even in cases where position information is not required and the position labels only serve as arbitrary labels (as in the SORT tasks in our task suite). Sinusoidal positional encodings are known to be limited in their ability to support length gencralization @Ontanón et al.] 2021 Csordás et al. (021a). In our setting, the demand to read out the corresponding item positions, which not common across NLP tasks, may have further limited the usefulness of this scheme.
11

---

Published in Transactions on Machine Learning Research (08/2023)
Our findings and those of Ruoss et al. (2023) are important in indicating that length per is not an intrinsic limiting factor in transformers - up to, of course, the total length of the transformer's context memory. This makes sense given that the query-based attention mechanism treats all positions equivalently except insofar as it is sensitive to the actual positional information encoded in the keys, queries and values. It is worth noting that using a random ordered subset of position labels would not represent absolute item distance information. Such a scheme can, however, capture sequence order, which appears to be sufficient for our tasks and those of Ruoss al.] (2023). This scheme may also prove sufficient in natural tasks, and this will be an important question to explore in future research. For example, in natural language, it may be sufficient to represent the relative proximity of items in the context rather than their absolute distance, since constituents such as noun phrases can be of arbitrary length. Of course, it may be possible to combine label-based encoding and other distance-preservingg position encoding schemes to allow a model to achieve length generalization ion while remaining sensitive to absolute token distances. We further note that transformers can acquire sensitivity to position information without any explicit positional information (Haviv et al.) 2022). This and other recent work suggest that additional research is needed to fully understand the capacities and relative merits of the full range of approaches to providing sequence and positional information to transformers and other architectures.
Our synthetic task suite and and representations used to solve our tasks in the models we have studied. Like others, we are ultimately interested in understanding models trained on naturalistic tasks, which have rich and only partially systematic structures and are more challenging to fully understand. In our view, the study of synthetic tasks and naturalistic tasks can mutually inform each other, and we hope our work leads to more discourse between research in these complementary domains. Specifically, we hope that the rosults show here provide useful insights and analysis tools that can be applied and extended to enhance our understanding of the computations and mechanisms that emerge in larger models. Some of our findings already show convergence with other studies using more complex tasks or different model architectures. For example, our findings on the contributions of individual attention heads in solving our algorithmic tasks echo results from analyses of language models. 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.) 2019). 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-selectivity at the level of individual units in the hidden layer. The degree of task-sclectivity in different architectural components may vary depending on the inductive biases of different model families and the tasks being learned. Future work is needed to fully characterize the degree of possible computational modularity in multi-headed attention, and to confirm whether our specific findings generalize to broader settings.
approach have allowed us to elucidate the computations
The representational analysis approach we have applied here is stepping stone toward uncovering how cradient-based lcarning leads to the development of structured solutions in neural network models. The analysis tools we have exlended from [axe et al. (2019) were developed in the context. of exact analytic solutions to the non-linear learning dynamics in linear multi-layer networks in terms of the strengths of the singular modes that characterize the structure of the training data. [aax et al. (2022) have recently extended this analytic reduction be used to analyze a wide range of feed-forward network architectures. In future work, we foresee the possibility of using this approach to clarify how the gradient signals flowing through our transformer network serve to prioritize task-relevant information, and even how they allow the task embeddings to modulate this prioritization process in a task-dependent manner. We also expect the approach to be useful for uncovering the priority given to different kinds of information in more naturalistic task settings, including image-based tasks or natural language and reasoning tasks in transformer- based language models. [Saxe et al. (2022) have taken a small step in this direction. One natural application of the approach will be to understand the relative prioritization and synergistic integration of in-context vs. in-weights information as explored experimentally in Chan et al. (2022). We consider the prioritization and integration of these two sources of information to be central to understanding the general ability of transformers to behave as multi- task systems.
12

---

Published in Transactions on Machine Learning Research (08/2023)
The flexibility to learn and perform multiple tasks is key desired capability for machine learning and artificial intelligence. Our work here provides insights into the structured within-task and cross-task computations that stacks of attention layers develop when learning 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-lask and nulli-modal learning in domains with richer task structures, it possible that these models may implicitly learn to decompose complex decisions into reusable, multi-level policies. In future work, we hope explore these learning and generalization dynamics in transformer- based systems to understand the acquisition of task-condiiioned, multi-level behavioral policies in structured environments.
Acknowledgments
We would like to thank Andrew Nam, Satchel Grant, and members of the Stanford PDP lab for useful discussions. We also thank Andrew Lampinen and anonymous reviewers for helpful feedback on the manuscript.
References
Cem Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowyez, Vedant Misra, Vinay Ramasesh, Ambrose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. Exploring length generalization in large language models. arXiv preprint WVXiv:2207.04901 2022.
Marcel Binz and Eric Schulz. Using cognitive psychology to understand gpt-3. arXio preprint or arXiv:2206.14577 14570, 2022.
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emna Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXic:2108.07258 2021.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neclakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neurol information processing systems, 33:1877-1901, 2020.
Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, et al. Sparks of artificial general intelligence: Farly experiments with gpt-4, arXiv preprint XrXi::2303.1211 2023.
Stephanic Chan, Adam Santoro, Andrew Lampinen, Jane Wang, Aaditya Singh, Piorre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learning in transformers. Advances in Neural Information Processing Systems, 35:18878- 18891, 2022.
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, pages 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:15084-15097, 2021.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sabastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXie preprint arXiv:2204.0233 2022.
Róbert Csordás, Kazuki Tric, and Jürgen Schmidhuber. The devil is in the detail: Simple tricks improve systematic generalization of transformers. arXiv preprint or Xrxiv:2108. 12284, 2021a.
Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber. The neural data router: Adaptive control flow in transformers improves systematic generalization. arXiv preprinl 10.00:210 2021b.
13

---

Published in Transactions on Machine Learning Research (08/2023)
Ishita Dasgupta, Andrew K Lampinen, Stephanic CY Chan, Antonia Creswell, Dharshan Kumaran, James L McClelland, and Felix Hill. Language models show human-like content effects on reasoning. arXiv preprinl arXiu:2207.07051, 2022.
Mostafa. Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. arXiv preprint (rXxc:1807.03811 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 hicrarchy. arXiv preprint arXiv:2207.02098 2022.
Jacob Devlin, Ming-Weeeettee Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding, orXiv preprint arXiv:1810.04800 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 at scale. ai preprint arXiv:2010.11929, 2020.
Laura Driscoll, Krishna Shenoy, and David Sussillo. Flexible multitask computation in recurrent networks utilizes shared dynamical motifs. bioRziv, 2022.
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jian, Bill Yuchen Lin, Peter West, Chandra Bhagavatula, Ronan Le Bras, Jena D Hwang, et al. Faith and fate: Limits of transformers on compasitionality. Xrv preprint Xxio:2905.18644 2023.
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 Hatfield-Dodee Danny Hernandez, Andy Jones, Jackson Kermion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. https://ttansforrerrdddrenttsss
harammooork/ideexex html.
Jerry A Fodor and Zenon W Pylyshyn. Connectionism and cognitive architecture: critical analysis. Cognition, 28(1-2)33-71, 1988.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv preprint arXim:1410.5401, 2014.
Adi Haviv, Ori Ram, Ofrr Press, Peter Izsak, and Omer Levy. Transformer language models without positional encodings still learn positional information. arXiv preprint w:Xiv:2203.16334 2022.
Sarthak Jain and Byron C Wallace. Attention is not explanation. arXiv preprint arXiv:1902.10158: 2019.
Michael Janner, Qiyang Li, and Sergey Levine. Reinforcement learning one big sequence modeling problem. In ICML 2021 Workshop on Unsupervised Reinforcement Learning, 2021.
Kuang-Hun Huci Loc, Ofir Nachum, Mengjiao Yang, Lisa Loe, Daniel Froeman, Winnic Xu, Sergio Guadarrama, Ian Fischer, Eric Jang, Henryk Michalewski, al. Multi-game decision transformers. ar preprinl arXiv:2205.152.4 2022.
Tal Linzen and Marco Baroni. Syntactic structure from deep learning. Annual Review of Linguistics, 7: 195-212, 2021.
Christopher D 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, 107(44))30044630000 300544, 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 Cvicck, and Zachary Fisher. Making transformers solve composi- tional tasks. arXio preprint arxim:2108 0499888 2021.
14

---

Published in Transactions on Machine Learning Research (08/2023)
OpenAL. Gpt-4 technical report. ArXiv, abs/ abs/2333.08877, 2023.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets. arXiv preprint arXiw:2211.022 2022.
Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. aix preprint ari:22100.. 12409, 2021.
Yasaman Razeghi, Robert L Logan IV, Matt Gardner, and Sameer Singh. Impact of pretraining term frequencies on few-shot reasoning. arXiv preprint or Xit:2202.07200 2022.
Scott. Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth- Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, al. A generalist agent. arXiv preprint ar arXie22255.011 0000000061115 2022.
Anian Ruoss, Grégoire Delétang, Tim Genewein, Jordi Grau-Moya, Róbert Csordás, Mehdi Bennani, Shane Legg, and Joel Veness. Randomized positional encodings boost length generalization of transformers. arXiv preprint TrXfe:2305.16844 2023.
Andrew Saxe, Shagun Sodhani, and Sam Jay Lewallen. The neural race reduction: Dynamics of abstraction in gated networks. In International Conference on Machine Learning, pages 19287-19309, PMLR, 2022.
Andrew M Saxe, James L McClelland, and Surya. Ganguli. A mathematical theory of semantic development in deep neural networks. Proceedings of the National Academy of Sciences, 116(23):115377-15566, 2019.
Paul Smolensky, Richard Thomas McCoy, Roland Fernandez, Matthew Goldrick, and Jianfeng Gao. Neuro- compositional computing: From the central paradox of cognition new generation of ai systems. AI Mogazine, 2022.
Jianlin Su, Yu Lu, Shengleng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. arXiv preprint orXiv:2104.0980 09864, 2021.
Shikhar Vashishth, Shyam Upadhyay, Gaurav Singh Tomar, and Manaal Faruqui. Attention interpretability across NLP tasks. arXiv preprint 1909: 11218, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is 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.
Petar Veličković, Adríà Puigdomènech Badia, David Budden, Razvan Pascamu, Andrea Banino, Misha Dashevskiy, Raia Hadsell, and Charles Blundell. The CLRS algorithmic reasoning benchmark. arXic preprint 220. 15659, 2022.
Elena Voita, David Talbot, Fedor Moiseev, Rico Semnrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. arXiv preprint 1900109050941 18, 2019.
Sarah Wiegreffe and Yuval Pinter. Attention is not not explanation. arXie preprint arXio: (Ef:::99000466,,, 2019.
Guangyu Robert Yang, Madhura R Joglekar, H Francis Song, William Newsome, and Xiao-Jing Wang. Task representations in neural networks trained to perform many cognitive tasks. Nature newroscience, (2):297 306, 2019.
15

---

Published in Transactions on Machine Learning Research (08/2023)
A
Additional Results
A
B
1.0
1.00
token-ievel tem acc
0.8
sespiee item acc
(roilout, >=0 95)
0.75
treiloutt
0.50
0.4
R
0.2
G[s]
G[c]
0.25
G[s]
G(c]
S[s,c,t]
S[c,s,t]
S[s.c,t]
S[c,s,t]
0.0
0.00
10
20
30
40
50
10
20
30
40
50
10
20
30
40
50
sequence length
sequence position
sequence length
1.0
1.00
tolen- evel <OOS>> acc
token-level <eop> acc
0.99
0.8
0.5
(re)
0.4
G[s]
G[c]
G[s]
G[c]
0.96
S[s,c,t]
S[c,s.t]
0.2
S(s.c.t]
s[c,s,t]
0.95
0.0
10
50
10
30
40
position
sequence
Figure S1: Multi-lask models' accuracy under lopl rollout. A. Token-level rollout accuracy over sequence length and sequence positions. B. Sequence-level rollout accuracy where the models predict more than 95% tokens correct within sequence. A and B, visualization as Figure 3C. C. Six-task models predict the token perfectly under teacher forcing. D. Prediction of the EOS token deteriorates under topl rollout in six-task models. Results aggregated across the five top-performing runs and taken from novel sequences within the training range.
A
single lingle-layer, single-head model sort(shape, color, texture]
B
single-layer, two-head model
C
two-layer, single-heade model on group[shape)
sort[shape, color, texture]
Jayer=0 h
layer=1 head-0
0.8
0.8
layer-0 h
10
queries
2
queries
0.4
0.4
004
1.8
0.0
layer=0 =ead=1
layer-d head=D
02
0.8
2
0.8
000
quaries
2-
Sources
0.4
0.4
0.0
sources
sources
Figure S2: Attention maps from other single-task models. Visualization as in Figure 43. Number labels indicate the beginning of each shape group. Label e indicates the EOS token.
16

---

Published in Transactions on Machine Learning Research (08/2023)
architecture: [4,4]
LO HO LO H1
LO H2 LO H3 L1 HD L1 H1 L1 H2 L1 HB
LO HO LO H1 LO H2 LO H3 L1 HO L1 H1 H1 L1
0.0
drop in accuracy from original performance nee
-0.2
-0.6
-0.8
teacher-forced item accuracy
top1 rollout Item accuracy
G[s]
G[c]
S[s.c,t]
S[c,5,t]
architecture: [1,6]
LO-HO
L1-H1 L1-H2 L1-H3
Lo-Ho L1-H0 L1-H1 L1-H2
L1-H4 L1-H5
0.0
drop in accuracy from original performance
-0.2
-0.4
0.88
teacher:forced item accuracy
topl rollout item accuracy
-1.0
Figure S3: Attention ablation for two other top-performing models.
17

---

Published in Transactions on Machine Learning Research (08/2023)
LG-HO
L1-H0
11-H1
L1-H2
11-85
task: copy
task:
1.
Lask: groppe
e up
mm
task: groupEcolor]
23
sort[s.c.t]
4-
234
152
task: sort/c.s.t)
OUSSUE INCCONRS
QUeriEs
rscrdderdd Sources
SIQUETICS
Figure S4: Full attention maps for example generalization sequence (for the top-performing multi-lask model with [1,4] architecture). Each row shows the attention maps for one task, and each column corresponds to 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 ROS loken. the bottom four rows, number labels ranging from 1-5 mark the beginning of the items in each first-level feature group (shape color).
18

---

Published in Transactions on Machine Learning Research (08/2023)
multi-task layer 1 representations
10
0
10
9C2
copy
reverse
group(s]
group(c]
sort(s,c,t]
sort(c,s,t]
target
probe
-10
Adppa
0.1
variance explained (%)
0.3
color
0.1
0.5
-10
texture
0.3
0.1
0.5
-10
label
0.1
0.5
residual
0.3-
0.1-
0
60
120
180
60 120
180
60 120 180 60 120180 dimension
60 120 60 120
Figure S5: Multi-task model's layer representations. Top, the lowcr-dimensional projection when represen- tations are viewed as target items. Bottom, the variance-explainer profile when representations are viewed as probe or target items. Blue lines indicate the total variance explained in a category across dimensions.
19

---

Published in Transactions on Machine Learning Research (08/2023)
B Hyperparameters
Table S1: Model and experiment hyperparameters.
Hyperparameter
Single-task learning
Six-task learning
Number of layers
or 2
Number of attontion heads
(sec paper)
Embedding dimension
128 (two-layer) or 184 (single-layer)
192
MI.P hidden layer size
64
Activation function
ReLU
Batch size
128
Training teacher forcing rate
1.0
Optimizer
Adam
Learning rate
4
5
C Quantitative Performance
Table S2: Token-level item and label accuracy across 54k generalization sequences in single-task models. Mean and standard deviation across four random seeds are shown for each architecture. Model architectures are denoted by the number of attention heads in each layer.
task
architecture
position encoding
item prediction
label prediction
C
label
99.03-10.41
100.00
R
1
label
99.75=10.21
100.00++0.00
S
S|
label
76.49112.99
51.75:15.35
S(SCC,T
2
label
91.01 --2.7
78.02-+5.89
S|S,C,T
1.1
label
99.60=0007
98.41+:0.23
G[s)
1,1
label
98.41:00.32
99.311:0.16
S(S.C,T
1,1
sinusoidal 57.04+100775
16.50+6.52
S S|S,CTT
1.1
learnable 73.15+11.65
40.51=14.28
20

---

Published in Transactions on Machine Learning Research (08/2023)
Table S3: Token-level prediction accuracy across 75k generalization sequences in six-task models (including 500 unique sequences for cach length evaluated within each task). Mean and standard deviation across four random seeds are shown for each architecture. Model architectures are denoted by [# heads in the first layer, # heads in the second layer].
(a) Equal number of attention heads in each layer.
Tasks
[1.1]
2,2
[3.3]
[4.4]
[6,6]
item
83.38-113.00 13.05
97.04=10.88
98.55:10.46
98.79:t0.94
99.111:.32
all
label
87.21=7.86
96.46=1.24
97.98:10.42
98.68::0.84
99.17:1.32
item 89.28+11.444
97.57+1.19
98.81+00.77
98.98-+0.95
99.13+0.97
C
label 97.39:10.84
98.59-1.99
99.46:10.39
99.64=0.34
99.79-1.117
item 64.30-124..6
96.95-11.21
98.95:10.45
98.69:10.94
99.20-10.60
R
label 94.56+8.26
98.53=0..0
99.6440.16
99.39::0.29
99.7310019
item
86.47=13.51
96.51=0.93
98.42:+0.66
98.52=+1.15
98.54:10.47
G[s]
label
91.5448.30
97.83+1.07
09.07
99.23
99.05
item
86.30+13.39
96.72+00.8
98.28
98.45-+1.13
98.62-+0477
G[c]
label
91.69-7.779
97.93=1.06
98.9811.19
99.21=:046
99.22:10.24
item
86.97=8.80
97.28=1.00
98.48$0.29
99.01=0.82
99.59-0.25
S|s,c,]]
label
74.75+11.47
93.21+2223
95.53+1.01
97.09++1.90 1.99
98.62++0.22
item
86.77=8.63 97.20=10.95
98.37
99.08=10.70
99.62-10133
S[C,s,]]
label
73.57:::0.59
92.73-12.05
95.23++0.90
97.53
98.65+0.55
(b) Multi-head attention in the first layer.
Tasks
2,1
3.1
[.1
6,1
item
69.45-
68.05-119.27
77.86+122.15
97.76+0.11
all
label
79.95:18.53
78.35+11..00
86.53111.67
97.5210.31
item
57.93:1.7.09
54.074330.35
72.014:30.48
98.50
C
label 97.66+1.64
95.18+3.28
98.48+1.10
99.61=0.22
item 57.97+:26.94
55.33+-29.52
70.49+332.68
98.78:00.46
R
label 95.74113.34
94.90-13.50
97.97+2.117
99.56
item
70.16+116.99
71.29+16.08
79.20+19.34
96.85+0.23
G[s)
label
80.56+11..9
19.37+12.51
87.82+12.52
98.9610.26
G[c]
item
70.33:116.54
70.29
78.86+19.18
96.50=0.47
label
79.18++12.71
78.71+12.82
87.28+12.69
98.82+0.32
item
79.94=11.69
78.00=112.49
83.99+15.22
98.03=00233
S[s,C,T]
label
62.76+13.885
61.17+199.00
76.69+18.13
94.17-0.74
item
80.03+111.15 78.22412:21
82.46-116.50
97.9510.29
S[C,s,T]
label
64.28
61.27+18.82
71.27+224.9
94.08=0.82
(c) Multi-head attention in the second layer.
Tasks
[1,2]
[,3
[1,4]
[,6]
itern
96.06 0.244
98.65-10.80
99.40
99.29-10.17
all
label
95.48
97.37+0.89
99.2110.43
99.1210.36
item
96.86=0.577
99.10:10.74
99.46-10.34
99.22=0.52
C
label 9..7
99.62-10.24
99.8510.08
99.59-10.55
item 95.99+0.59
98.82+1.08
09.47+0..37
99.56--0.14
R
label 98.09:0.52
99.55+0.39
99.83:10.04
99.82=00099
item 16.1240.38
98.28+1.00
98.9810.43
98.96+0.16
G(s)
label 97.70++0.47
98.77+0.66
99.29+0.27
99.25+0.07
item 95.82+0.72
98.53+0.76
99.1710.34
98.78+0.40
G[c]
label 97.5110.51
99.06+0.42
99.5110.18
99.1310.32
item 95.8610.77
98.51
99.6810.26
99.59-0.13
S[s,c,r]
label 91.1240.83
93.54+2.44
98.40+1.16
98.4110.89
item 95.7210.54
98.7020.71
99.6410.28
99.63-0.29
S[C,s,T]
label
90.7440.66
93. 73.711
98.36:11.05
98.54+11.10


Reviews:
[{"review_id": "ohneqkHBwc", "reviewer": "Reviewer_zzw5", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall, I found this to be a very strong paper - I would like to thank the authors for creating such interesting and illuminating work!\n\nStrengths:\n- S1: The experimental setup is careful and well-motivated. The authors have chosen tasks that are simple enough to be interpreted yet complex enough to reveal interesting insights about how Transformers operate. This setting may be useful for future authors wishing to conduct mechanistic interpretability work.\n- S2: The results make important progress toward the goal of understanding “black-box” neural networks that are notoriously hard to understand - a goal that the community is increasingly recognizing as an important one.\n- S3: The paper uses an impressive range of analysis techniques (e.g., attention analyses, attention ablations, probing classifiers, and more). There are 2 reasons why this is a major positive: First, it provides converging evidence for the authors’ conclusions. Second, these analysis techniques (like the tasks/datasets) may be useful for future work.\n- S4: The paper is clear and well-written. Although there is a lot in it, the authors have explained everything very concisely, so I don’t think it could be condensed more without weakening the paper. There were also several occasions where I had a question that was immediately answered in the next paragraph - a sign that the authors have carefully thought through their paper to anticipate any confusion!\n- S5: I really enjoyed the figures - they have lots of content arranged in a careful way, so it was fun to dive into them.\n- S6: The paper seems like a great fit for TMLR: It has much more content than a typical conference paper, so the journal format is perfect for being able to include the extensive, detailed analysis. \n\nWeaknesses: There were no major weaknesses that I identified. The one partial weakness is that there is so much in the paper that it is a bit dense, and it requires some focus to process everything. But I am fine with that, and I think that the authors did a good job of balancing conciseness with clarity (see S4). Some readers may also view it as a weakness that the paper only uses synthetic tasks, but I do not think this is a weakness - in fact, I think it is a strength because it makes it possible to understand the phenomena thoroughly, paving the way for analysis of more complicated phenomena (see S1).\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors train Transformers to perform a range of synthetic sequence-manipulation tasks (e.g., reversing the sequence, or sorting a sequence based on features of the iterms). They then perform extensive and fine-grained analyses of the representations and processing of the trained models to understand how they have learned to perform these tasks. In more detail, the contributions include:\n- The set of tasks that models are trained on, as a useful setting for performing model analysis\n- A new positional encoding scheme that improves length generalization on these tasks. (As the authors note, another paper has concurrently introduced the same scheme, but the authors describe how their analyses of the scheme are complementary.)\n- A determination of the minimal number of heads and layers that are needed for Transformers to perform different individual tasks, or combinations of tasks in a multitask setting\n- Attention-based analyses showing how models use attention for different targeted purposes in different layers to compute the information needed to solve the task in a structured, multi-step way.\n- Analysis of attention head sharing, task embedding similarity, and cross-task decodability transfer of information supporting the view that multitask models have learned to partially but not completely share information across tasks\n- PCA analyses showing how representations induced for different tasks reflect the different types of information that are useful for the task in question (e.g., color is prioritzed in a color-based sorting task, while shape is prioritized in a shape-based sorting task)\n", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "None of these are necessary for securing my recommendation; they are just suggestions that could clarify some aspects of the paper.\n- RC1: I found the term “label encoding” to be unintuitive. I typically think of “label” as meaning “something that captures some information about an item” (e.g., a part-of-speech tag), whereas these labels are random and arbitrary. Perhaps another name would be clearer, such as “random ascending positions”? \n- RC2: The main contributions (1, 2, and 3 at the end of section 1) did not make much sense to me until I had read the rest of the paper. Given that this paper is dense, I expect that readers would really benefit from having a clearer summary at the start so that they can get the main points more clearly without diving into the details. Therefore, I would recommend expanding on these points a bit. Specifically:\n    - a. Under point 1, I recommend briefly summarizing the structure of the label based encoding\n    - b. Under point 2, I would recommend adding a sentence that says “For example, …” and then gives 1 or 2 examples of specific results, in order to give readers an intuition about what sorts of “systematic decomposition” and “exploitation of shared structures across tasks” are being discussed. E.g., you could say “For example, we find that 2-layer transformers implement a 2-step algorithm with the attention patterns in each layer performing one step of the algorithm; and we find that, in multitask models, each attention head is used roughly equally by all tasks.”\n    - c. Under point 3, similar to point 2, I would recommend adding a “For example…” to give readers intuition about what it means to “re-weight task-relevant information”\n- RC3: Figure 3B: In the caption, maybe mention that models are arranged in ascending order of performance? I assumed at first that it was in ascending order of size, so I missed that 1,6 is to the left of 1,4 in the bottom row.\n- RC4: Figure 4B, x-axis labels and y-axis albels: the “e” labels are hard to see because it touches the circle - could they be separated a bit? The same comment also applies to “t” labels in Figure 5C (for x-axes and y-axes)\nIn the middle of page 6, it says “reflects an advantage of the two-layer models in providing an inductive bias favoring a structured, multi-stage solution.” I slightly disagree with this characterization: It seems more like this task requires a multi-stage solution, such that any model with the capacity to perform a multi-stage solution would learn to implement one. Thus, it may not be that 2 layers gives an inductive bias for multi-stage solutions, but rather that it gives the capacity necessary to develop a multi-stage solution. In order to say that it has a bias, I think the proper type of evidence would be to show that it learns a multi-stage solution even when the task does not demand it; when the task does demand it (as it seems to do here), then the multi-stage nature seems likely to be driven by the task rather than by the model’s inductive biases.\n- RC5: When I first read the first paragraph in Section 3.4, I was confused and didn’t understand it. It may be worthwhile to say somewhere in that paragraph “As we will discuss in more detail in the next few paragraphs…” to signal to readers that this paragraph is just a summary. (My confusion came from the fact that I thought the paragraph was standalone, so I was confused about why it didn’t give numbers or explanations to support the statements about what the results reveal).\n- RC6: When it says “The results show that such a classifier was above chance“ (on page 9), it might be worth it to remind readers of what chance was\n- RC7: A related paper about RNNs that you might enjoy is: Processing of nested and cross-serial dependencies: An automaton perspective on SRN behaviour. Kirov and Frank 2012. https://www.tandfonline.com/doi/pdf/10.1080/09540091.2011.641939. \n", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "9fIAZVqzrP", "reviewer": "Reviewer_rtG2", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- The novel positional encoding scheme is an interesting proposal, and displays a promising ability to support systematic generalization to longer sequences.\n- Task space is simple yet sufficiently diverse to explore how models might learn to compose operations in support of structured multi-task reasoning.\n- Extremely thorough analyses yielding many insights concerning how transformers learn to perform algorithmic tasks. The role of inductive biases (especially attention and multi-layer composition of operations) is intuitively and effectively conveyed by the analyses and text.\n- An interesting analysis method is presented using PCA to show how representations are transformed based on task demands. This approach seems likely to be useful in many other contexts.\n\nComments:\n\n- Are there any drawbacks to the proposed positional encoding scheme? In particular, I'm wondering if this approach makes it harder to make precise positional judgments. For instance, how would this approach fare on a task such as 'given a sequence of objects, how far apart are objects A and B?'\n- As an alternative encoding scheme, have the authors considered simply randomly scaling the positional indices? This seems like it would convey the same benefit of generalization to longer sequences, but would also preserve the precise metric relations between positions.\n- The proposed encoding scheme is clearly beneficial in terms of generalization to longer sequences, but I was somewhat surprised that both learned and sinusoidal positional codes also underperformed on the training set. Do the authors have some sense for why these methods also struggled in an iid setting?\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work presents an in-depth analysis of transformer models trained to perform algorithmic tasks, incorporating a novel PCA-based analysis method to better understand how representations are transformed to support single- and multi-task learning. A novel positional encoding scheme is also introduced that supports better systematic generalization to sequences longer than those observed during training. ", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I have a few comments and suggestions listed above that may improve the paper further, but I think the paper is already in good shape.", "broader_impact_concerns": "There are no discernible ethical issues related to this work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "3OF5F9E4Se", "reviewer": "Reviewer_kxxZ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths \n\nThis paper addresses a very interesting and potentially impactful question. \n\nThe authors carried out extensive experiments. Most claims are supported. Some results (e.g., multi-task model) are very interesting. \n\nThe paper is overall well-structured and some visualizations (e.g., Fig-1) are good. \n\n\nWeakness \n\nThe main issue of this paper is that its conclusions seems anecdotal since the data and tasks are synthetic. It is not clear why the conclusions may generalize to more broad settings like text. \n\nThe paper is overall clear but vague/unclear at some key parts. \n\nE.g., the reason why random-label should work is not sufficiently justified, see Requested Changes. \n\nE.g., in 3.2, it is not clear why certain conclusions are drawn such as \"multiple layers strongly benefit...\"  The legends of mentioned figures are not defined (please correct me if I missed them) such as [1,1] and [2] in Fig-3.\n\nIt is not clear to me why the random label technique is connected to the other parts of the paper. \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper explores how a generative Transformer model discovers solutions as it learns to solve algorithmic tasks. Their results provide insights into how attention layers support structured computation within and across tasks, including how they learn generalizable layer-wise features. ", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I would like a better justification on why the proposed random label-based encoding can achieve the claimed benefit of \"reliably communicate relative position information if ... in ascending order\". The experiment results are positive, but I find it hard to be convincing unless the authors can show some good thought experiments. This may be a presentation issue: e.g., I am not sure if the sampled labels are always in ascending order. (If yes, it seems more reasonable to me since their embeddings may learn the notion like \"2 is before 5 and 9 whatever\" but it needs a good clarification from the authors.) \n\nCan you explain how the random label technique is connected to the other parts of this paper? ", "broader_impact_concerns": "No ethical concerns. ", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: All reviewers recommend acceptance of this paper and the AE agrees. Modifications made during the reviewing process improved the clarity of the proposed methodology.
Decision: Accept (poster)