Your Context Is Not an Array: Unveiling Random Access Limitations in Transformers
Authors: MohammadReza Ebrahimi
Sunny Panchal
Roland Memisevic
Venue: COLM (2024)
Keywords: Transformers, Large Language Models, Length Generalization, Arithmetic Tasks, Context Addressing
TL;DR: We provide analysis that suggests failure of Transformers in length generalization is intricately linked to the model's inability to perform random memory access within its context window.

Abstract: Despite their recent successes, Transformer-based large language models show surprising failure modes. A well-known example of such failure modes is their inability to length-generalize: solving problem instances at inference time that are longer than those seen during training. In this work, we further explore the root cause of this failure by performing a detailed analysis of model behaviors on the simple parity task. Our analysis suggests that length generalization failures are intricately related to a model's inability to perform random memory accesses within its context window. We present supporting evidence for this hypothesis by demonstrating the effectiveness of methodologies that circumvent the need for indexing or that enable random token access indirectly, through content-based addressing. We further show where and how the failure to perform random memory access manifests through attention map visualizations.

Published as a conference paper at COLM 2024
Your Context Is Not an Array:
Unveiling Random Access Limitations in Transformers
M.Reza Ebrahimi
Sunny Panchal
University of Toronto, Qualcomm Al Research" ebrahimi@qti. qua..uullllcco.cc....cc.. qualcomm...c.... comm.
Qualcomm AI Research
sunnpanc@qti qualcomm. com
Roland Memisevic
Qualcomm AI Research
memisevv@tii. qualcomm. com com
Abstract
Despite their recent successes, Transformer-baseed large language models show surprising failure modes. A well-known example of such failure modes is their inability to length-generalize: solving problem instances at inference time that are longer than those seen during training. In this work, we further explore the root cause of this failure by performing detailed analysis of model behaviors on the simple parity task. Our analysis sug- gests that length generalization failures are intricately related to a model's inability to perform random memory accesses within its context window. We present supporting evidence for this hypothesis by demonstrating the effectiveness of methodologies that circumvent the need for indexing or that enable random token access indirectly, through content-based addressing. We further show where and how the failure to perform random memory access manifests through attention map visualizations.
1 Introduction
The evolution of "Transformer-based large language models (LLMs) has marked new era in how machines understand and interact with human language. Their capabilities extend far beyond natural language tasks, encompassing instruction following (Ouyang et al., 2022), code generation (Zhang et al., 2023), theorem proving (Wu et al., 2022), and common sense and multi-step reasoning (Yu et al., 2023). This has made LLMs play a pivotal role the backbone of Al agents (Xi et al., 2023), and even has sparked discussions around their ability to exhibit glimpses of general intelligence (Bubeck et al., 2023).
Despite these remarkable capabilities, surprisingly, the same models struggle with seemingly simple arithmetic tasks, such as multi-digit addition and multiplication (Dziri al., 2024). Specifically, the models fail to learn simple algorithms to perform these arithmetic operations. This becomes apparent when models are applied problems of greater length than those encountered during training (Hupkes et al., 2020), problem setting generally referred to as length generalization.
Arithmetic tasks fundamentally differ from natural language tasks in two key aspects. First, unlike natural language, responses to arithmetic tasks are objective and unambiguous, corresponding to the exact execution of sequence of algorithmic steps. The second dif- ference, and the focus of our work, is their reliance on formatting: arithmetic expressions represented using limited vocabulary, such digits, with each token holding equal significance.
Crucially, in the representation of arithmetic tasks, token's position is important its value. This stands in stark contrast to natural language expressions, in which the coupling
"Qualcomm AT Research is initiative of Qualcomm Technologies, Inc.

---

Published as a conference paper at COLM 2024
between token word positions on the one hand and the meaning of the expression the other is much weaker and much more flexible. In the context of language modeling this has been demonstrated, for example, by Sinha et al. (2021), who show that permuting word orders has a surprisingly small effect on the performance of BERT models in natural language processing tasks.
In other words, the meaning of natural language utterances depends largely on the meaning of their constituents (e.g,, words) and only partially on their positions. This well-known influence of meaning (semantics) over pure syntax is exemplified in expressions, such as "He saw the cat with the binoculars", in which the phrase "with the binoculars" is more likely subordinate to "He", even though syntactically it could equally be subordinate to "the cat". The precise position of individual words becomes even less informative when references stretch over larger distances, such as across sentences.
As illustrated in Figure 1, when predicting the next token in a natural language task, to- ken references which are "contenl-based in this way are well represented by the common attention mechanism prevalent the Transformer, and they are further reinforced through pre-training natural language. This is in contrast to arithmetic tasks, which rely exclu- sively on "index-based addressing" (random access memory) into the context window to retrieve the information necessary for generating the next algorithmic step.
Example: Content-based vs. Index-based Addressing
He
prepared
delicious
chocolate
cake
and
everyune
com compliment ed
him
on
the
dessert
12
Figure 1: Top: Prediction in natural language tasks. To predict the pronoun him, the model needs to access previously used pronouns in the context, among other tokens, regardless of the exact position of the token He in the context (content-based addressing-- Bottom: Prediction in an arithmetic task. The model returns the running parity of the binary sequence after For the third output, the model must precisely attend to the token in position 3 of the context window (index-based addressing).
In this work, we provide in-depth study of this addressing dichotomy and present evidence for its role in the failure of Transformer language models in algorithmic tasks. We focus on the binary parity task as it is, arguably, the simplest sequential arithmetic task, making it well-suited to study the underlying computational requirements of Transformers applied it. When properly formatted, the state needed to carry over at each step is only one bit, and the key operation required to learn is XOR. Yet, Transformer-based models struggle to learn length generalizable algorithm for this task (Anil al., 2022).
Our detailed empirical study of the parity task across models with various positional embedding methods strongly supports the hypothesis that Transformers pre-trained on natural language learn to retrieve tokens using content-based addressing, leading them to fail on algorithmic tasks which, as discussed, depend on random memory access.
In Sections 3 and 4, we further demonstrate how the addition of "mmemonics" to leverage content-based addressing as a workaround for index-based addressing allows models to learn length generalizable algorithms for the parity and addition tasks, both of which were previously shown to be hard for Transformer language models. While the introduction of mnemonics is not proposed our hypothesis. Our work suggests that equipping models with effective index-based addressing mechanisms could be a key to learning algorithms that can length-generalize.
practical fix, it highlights the underlying issue and reinforces
2

---

Published as a conference paper at COLM 2024
2 Related Work
Length generalization is a well-known problem in the context of Transformer-based se- quence models (Qian et al., 2022; Newman et al., 2020; Zhang et al., 2022b; Zhou et al., 2024; Xiao & Liu, 2023). Notably, Anil et al. (2022) conducted careful empirical studies exploring the length generalization capabilities of Transformer-based LLMs with a focus on the boolean variable assignment and binary parity task. They demonstrated that models, even when fine-tuned on these tasks using a scratchpad format, struggle significantly with generalization, regardless of a model's scale.
The study by Dziri et al. (2024) examines the ability of Transformers to rength-generalize in compositional tasks, such as multi-digit multiplication, and highlights their generalization failures across zero/fes /few-shot and fine-tuning regimes, both with and without the use of scratchpad. It suggests that Transformers may approach compositional tasks by simplifying multi-step reasoning into a form of linearized subgraph matching, rather than developing systematic problem-solving skills.
The work by Zhou et al. (2022) examines the extent of in-context learning for algorithmic tasks through the strategic use of meticulously designed prompting techniques, called algorithmic prompting. As we shall show, our work suggests an alternative interpretation for the results of that work based indexing. Similarly, Zhou et al. (2023) build on the RASP computational model proposed by Weiss et al. (2021), and focuses on identifying algorithmic tasks learnable by transformers. conjectures that Transformers demonstrate strong length generalization for tasks that can be solved by a concise RASP program across various input lengths.
The work presented in Kazemnejad et al. (2024) involves a systematic comparison of length generalization performance across Transformers with various positional encoding schemes. It reveals that none of the commonly used positional embedding methods effectively solve the length generalization problem in downstream tasks. Surprisingly, having no positional embedding outperforms these methods, echoing a finding previously identified by Shen et al. (2023). This observation further indicates that current positional embedding approaches fail to equip the model with the capability for proper index-based addressing. Moreover, Shen et al. (2023) propose a modification to the positional embedding itself, by marking tokens with random tags. This allows the model to distinguish identical tokens appearing in different positions, offering slight improvement in generalization.
A study similar in spirit to our work is Dubois et al. (2019), albeit using recurrent sequence- to-sequence models instead of Transformers. That work hypothesizes that models equipped with separate content and location- extrapolate. It evaluates this hypothesis through variants of the Lookup Table task, designed to directly assess a model's performance in index-based addressing.
attention mechanisms are more likely to be able to
The work by Mohtashami & Jaggi (2024) proposes a method for handling long contexts by using sparse learnable "landmark tokens" to retrieve relevant token blocks. These landmark tokens bear some similarity with our use of "mnemonics" we shall discuss below.
3 Random Accessing in LLMs A Case Study
In this section, we focus on the binary parity task as a case study on learning algorithmic tasks with Transformers. We chose the parity task for its simplicity as one of the most basic sequential arithmetic tasks. With the correct scratchpad format, it requires carrying over just one bit of state at each step, and the primary operation to learn XOR. However, is known that Transformer-base solution fails for sequences longer or shorter than those seen during training (Anil et al., 2022).
models struggle to learn the correct algorithm as their
We begin with a brief note on the usage of scratchpads. When the model is asked to directly output the final answer, such as the parity of a sequence, we encounter potential complica- tion: Transformers execute a fixed amount of computation for each token generated, yet the problem size can vary. In other words, the model must simulate for-loop over the entire
3

---

Published as a conference paper at COLM 2024
sequence in single forward pass. Note that this represents distinct contaminating issue that falls outside the scope of this work. This challenge can be addressed by incorporating "scratchpad" (which is also referred to as chain-of-thought) (Nye et al., 2021; Wei et al., 2022). The scratchpad enables the effective use of the context window to explicitly simulate for-loop and output intermediate results.
Adopting the format used in Anil et al. (2022) for the parity task, we begin with a start-of- sequence symbol >>>, followed by binary sequence, end-of-sequence symbol ===, and the sequence's running parity. For instance:
No Scratchpad Standard Scratchpad
>>> 0 100100000 11 1 ===
>>> 1 0 1 1 === 1 1100011
Throughout the paper, blue bold tokens are used to indicate tokens over which the loss is calculated during training, and thus also the tokens that the model predicts during inference. Meanwhile, other tokens are added externally into the model's context during generation (via "environment forcing" (Recchia, 2021)). Also, we ensure that the start- -end-of-eeuuence symbols are converted to single tokens and bits within the sequence are represented by single fixed tokens, preventing any merging due tokenization.
3.1 Interleaved scratchpad
In essence, length generalizable solution generate the running parity in the specified format involves three steps: 1) Reading the current active bit; 2) Reading the current running parity, and; 3) Performing XOR between the active bit and the current parity. We hypothesize that the failure of Transformers can be attributed to the first step, since the subsequent two steps are straightforward: the current running parity is the last token generated, and the XOR operation is trivial to learn.
To support this claim with empirical evidence, we implement an interleaved scratchpad format where sequence bits and running parities are alternated, ensuring that at each step, the current active bit is the last token, and the current running parity appears immediately before the last token in the context. This arrangement dramatically simplifies the first step (reading the current active bit), which, as we will see shortly, lets the model learn a length generalizable solution.
Interleaved Scratchpad
>>> 1 1 11011
We fine-tuned several small Transformer models with different positional embedding meth- ods: BLOOMZ-560M with AliBi (Muennighoff et al., 2022; Le Scao et al., 2023; Press et al., 2021), Pythia-410M with RoPe (Biderman et al., 2023; Su et al., 2024), and OPT-350M with learned positional embedding (Zhang et al., 2022a). All models were initialized with their pre-trained weights and fine-tuned on task sequences of length 10 to 20 bits. They were tested on sequences of up to 60 bits. Refer to Section for experiment setup information.
Figure 2 illustrates the length generalization performance of fine-tuned BLOOMZ models using both standard and interleaved scratchpad formats, using training sequence lengths indicated by the shaded region. While the standard scratchpad method exhibits minimal improvement over not using a scratchpad, the interleaved version demonstrates perfect generalization. Notably, the sole difference between the two formats lies in the placement of the tokens in the context. The standard scratchpad format requires the model to perform index-based addressing to fetch the value of the current active bit, while the interleaved format eliminates this requirement. Section B.1 shows similar results for other models.
The observation above supports the hypothesis that the models' inability to learn arithmetic tasks stems from their failure to accurately perform index-based addressing of the input bits. In contrast, content-based addressing is inherently natural for Transformers through the attention mechanism and natural language pre-training Next, we will further reinforce this hypothesis by introducing another modification to the standard scratchpad.
4

---

Published as a conference paper at COLM 2024
1.0
Acouracy
0.4
No Scratchpad
0.2
Standard Scratchpad
Interleaved Scratchpad
Training Sequence Length
35
40
45
50
55
60
Test Sequence Length
Figure 2: Length generalization performance of fine-tuned BLOOMZ-560M models on sequences of length 10 to 20 bits, using standard and interleaved scratchpad formats, as well as without a scratchpad.
3.2 Mnemonics
We can leverage content-based addressing in Transformers indirectly perform index-based addressing, by adding matching "anchor" tokens before every pair of corresponding tokens in the standard scratchpad format. As they allow model revisit earlier information in the context window, we shall refer to these as mnemonics. Similar approaches are discussed in Bueno et al. (2022), Qian et (2022) and Zhou al. (2023).
During training and inference, for each example of length n, we first randomly sample tokens from pool of mnemonic tokens', then add the mnemonics before each bit in the input sequence and the running parity bits:
Mnemonics
>>> M₁ 1 M2 0
Mg 1 M4
Ms O Ms === Mj
1
M2 Mg 0
Mg Ms 0 Mg 1
Mnemonics (Environment Forced)
M₂ Ns M4 Ms O Ms Mj
Ma
Ms
Ms
Note: Mnemoonic tokens Mj, M2. are randomly sampled without replacement from the mnemonics pool, every problem instance.
Note that in the non-environmenn-foreed version, the model is trained to first place the matching mnemonics from the input sequence, and then use them address the active bit at each step. Conversely, in the environment-foree version, at each step, we first append the matching mnemonic from the input sequence the context, after which the model predicts the running parity.
Figure compares the length generalization performance of fine-tuned BLOOMZ models with and without using mnemonics in the scratchpad. The results illustrate that adding mnemonics enables the model to learn the correct algorithm for solving the task, leading to perfect length generalization for sequences of up to 60 bits, while being trained on sequences of only 10 20 bits. Additionally, Appendix Section B.3 investigates the in-context learning performance of the parity task using mnemonics.
These results suggest that equipping a model with effective index-based addressing could be key to enabling it to learn correct arithmetic algorithms. Interestingly, the performance
We used all space-preceded tokens containing only English characters for the mnemonics pool.

---

Published as a conference paper at COLM 2024
1.0
Acouracy
0.4
No Scratchpad
Standard Scratchpad
0.2
Mnemonics
Mnemonics (Env. Porced)
Training Sequence Longth
35
40
45
50
55
60
Test Sequence Length
Figure 3: Length generalization performance of fine-tuned BLOOMZ-560M models with and without using mnemonics in the scratchpad.
of the model using non-environment-fored mnemonics is nearly identical to that of the environment-foreed version, indicating the model's capability to both place and utilize mnemonics for indexing effectively. Similar results are reported in Appendix Section B.1 for other models. Additionally, we explore the effects of varying the interval between mnemonic tokens in Section B.2.
Using these scratchpad strategies, we also trained the same model initialized randomly instead of pre-trained on natural language. The results shown in Figure 4. Notably, when training from random initialization, mnemonic scratchpads are ineffective. This could be attributed to the fact that successful utilization of mnemonics requires the model to perform both, global addressing of the relevant mnemonic, followed by local addressing of adjacent tokens. The latter may be an ability that persists in the length generalization setting only due to pre-training on natural language.
BLOOMZ-560m (Random Initilization)
No Scratchpad
Standard Scratchpad
Mnemonics
Mnemonics (Env. Forced)
Training Sequence Length
Accuracy
0.6
0.4
50
35
Test Sequence Length
Figure 4:
BLOOMZ-560M models trained from random initialization on the parity task using twice the number of epochs.
3.3 Analysis of attention patterns
To further analyze how the model's attention changes with and without mnemonics, we present input attribution visualizations in Figure using the gradientx input method

---

Published as a conference paper at COLM 2024
191 000 LOnOnAT
Intedte
didbe
Mouniains
stnc
Kantba
reclbence
protocolls
rocaui
gouerr
Lave
sunto
nelghbourg
a fllrasciing
auropat
Intedd
Poramn
Mountaiio
shnngg
transler
neighoows
edtrautties
Dile'
amropott
Figure Input attribution visualized through the gradient xniput method during perform- ing the parity task. Models were trained sequences of 10 to 20 bits while predicting the parity of a 40-bit sequence, shown with (right) and without (left) mnemonics. Columns represent output tokens (after ===) and rows represent all tokens in the context window. Observe the scrambled attention pattern in the left figure, after the 20th output.
(Shrikumar et al., 2016). These visualizations show aggregate attention maps, with columns representing output tokens (after ===) and rows showing all tokens in the context window. Since the model's 's task is to produce the running parity of the input sequence, at step i, it only needs attend to the current bit (bit of the input) and the previous running parity (the last bit generated). Thus, the ideal attention map would show two diagonal lines, corresponding to these two relevant tokens. The attention maps are calculated on a sequence of length 40 for model trained sequences of length 10 to 20 bits.
As shown Figure on the left, immediately following the 20th bit (in-distribution length), the model fails to attend to the current bit when calculating the parity. In other words, the model has not learned length generalizable method for indexing the correct bit at each step, thus failing at indexing outside of its training regime. In contrast, plot of Figure 5, when mnemonic bits are added, near-perfect attention map is observed beyond the training regime.
in the right
7

---

Published as a conference paper at COLM 2024
3.4 Mnemonics variations
Finally, we study several variations of the introduced mnemonic tokens, which further support our hypothesis, as discussed below:
Numeric
>>>
a
3
4
5
6
2
3
4
5
6
b
Constant
>>>
# 0
#
#
#
#
== #
#
#
#
#
Non-aligned
>>> Mj 1 M2
0
Ms
M4
0
Ms
0
Ms 1 === My 1 Ms 1
Mg
0
Mi0
0
M11
M12
Cyclic
red
green yellow red green yellow
=== red
1
green 1 yellow
red 0 green 0
yellow 1
Numeric Mnemonics: We use consecutive numeric indices (1,2,3,...) as mnemonic tokens for all samples. To avoid confusion between mnemonics and binary values in the sequence, we use a, instead of o, 1 to represent the bits. Note that this form of mnemonics corresponds to absolute positional encoding.
Constant Mnemonics: A single fixed character (#) is used as the mnemonic token for all samples, during training and testing. This approach allows us test whether the effectiveness of mnemonics is related to the attention sink phenomenon (Xiao et al., 2023), or the model uses the mnemonic tokens as "placeholders" allowing it store intermediate calculations in their activations.
Non-aligned Mnemonics: This variant is similar to the original mnemonics with the difference that the random tokens used in the input and output do not match. Specifically, for sequence of size и bits, we sample 2n tokens to serve as mnemonics. We use this variant to test whether the impact of mnemonics results from making each digit unique for the model, rather than acting as positional anchors.
Here, we cycle through a predetermined array of mnemonic tokens, fixed across all samples in training and testing. Specifically, we used 10 color names as mnemonics in our experiment.
Cyclic Mnemonics:
Figure shows the failure of the aforementioned innemoncc variants at length generalization. Note that in the environment-force versions, mnemonic tokens are placed in the context window of the model externally. Compared to the original randomly sampled aligned mnemonics, each variation corrupts the mnemonics' utility as positional anchors.
In the numeric mnemonics variant, the model is exposed to mnemonic tokens 1,2, 20 during training, while at test time, it encounters unseen mnemonics 22, We further explore the impact of unseen mnemonics at test time in Appendix Section B.4. Additionally, the fixed nature of numeric mnemonics across training examples may hinder length gener-
1.0
0.8
Numeric Indices
Numeric Indires(Env. Forred)
Constant
Accurary
0.6
Indices(Eav. Forced)
Cyclic Indices
10 Cyelic In Indices(Env: Forcnd)
0.4
Non-aligmed Indices
Non aligned Indices(Env, Forced)
Indices Only Input
0.2
Training Sequence Length
0.0
15
20
30
35
45
50
Test Sequence Length
Figure 6: Length generalization performance of fine-tuned BLOOMZ-560M models the parity task, trained on sequences of length 10 to 20 bits, using different variants of mnemonics.
8

---

Published as a conference paper at COLM 2024
alization: in contrast to the original mnemonic scheme, which randomly selects mnemonics from a large pool of tokens for each training instance, the numeric variant uses the same mnemonics for all training samples.
In the constant and non-aligned variant, anchor-based alignment between the sequence and scratchpad is eliminated entirely. Finally, cyclic mnemonics are repeating and thereby create ambiguities regarding the correct next bit to read.
Overall, these results further support our hypothesis that Transformers struggle with per- forming random token accesses, and demonstrate how random mnemonics can mitigate this by facilitating random access through content-based addressing of the relevant mnemonic.
Solving the Multi-digit Addition Task
This section extends results to another arithmetic task: multi-digit addition. This task has been explored extensively in the literature with different scratchpad formats (Qian et al., 2022; Nye et al., 2021; Kazemnejad et al., 2024; Zhou el., 2024; Xiao & Liu, 2023; Zhou et al., 2022), among others. We focus on the length generalization performance of the addition task with mnemonics in three different formats.
In our format, the addition result is initially presented in reverse order, from the least to the most significant digits. Following the symbols ###, the model then reverses this produce the final addition result. It is important to mention that every single digit is converted to an individual token. We fine-tuned the BLOOMZ-560M model on the addition task using the specified format, training on operands with 5 to 10 digits and testing operands with up to 14 digits.
We use the same mnemonics for corresponding digits in both operands, as demonstrated below:
Digit-aligned Mnemonics
Mnemonics
>>>
###
0
Mnemonics
>>> Mj
1
Mg
2
M3
Mg
9
Ma
M2
1
M3
0
###
M3
M1
2
Mg 1
Env. Forced
>>> M;
1
Mg
2
K3
M2
9
Mg
--- M2
1
M₁
M3
0
###
M3
M1
2
M2
In another format, we first zero-pad the operands to ensure they have the same number of digits, then insert digit-aligned mnemonics:
Digit-aligned Mnemonics + Zero Padding
No Mnemonics >>> 1 2 + 0 9
-=- 1
0
#8B
Mnemonics >>> M2
2
Mj
M2
Mg
M2
1
M3
0
### M1
Env. Forced
M2 2 Ma Mi
M2
Ma
M2
M3
Ma
M1
2
Lastly, we explore format in which the mnemonics for corresponding digits of the two operands are not identical, as depicted below:
Non-aligned Mnemonics
No Mnemonics Mnemonics
>>>
###
M2 2
M3
M3
M2 2
0
M3
Env. Forced
>>> Mj
M2 2 Mg
M3
M2
M
2 O
###
0
M₁
M3
M2
1
The length generalization performance of the addition task, both with and without the specified mnemonic formats, shown in Figure As expected, aligned mnemonics guide the model in selecting the correct digits for addition at each step. Furthermore, zero-padding simplifies the task's format by ensuring an equal number of mnemonics and digits in both operands. Overall, findings show that similar to the simpler case of binary parity, by utilizing content-based addressing to enable index-based addressing via mnemonics, Transformer models successfully learn the correct algorithm for the addition task.

---

Published as a conference paper at COLM 2024
No Mnemonics
Mnemonics
Env. Forced Mnemonics
66
Operand 1 Digits
3300
Sugit- aligned
066
1o
a2
ta
10
12
LE
-24
13 31
E7
CO 03
10
19
10 19
25
53
16
10
D2
12
0.0
24 13,00
Digit- Digit-aligned aligned + Zero Padding
22100
Operand Digits
13 10
36
10
33,00
19 TE
10 15
-38 10
10
83
19:0
23.000
18
00
00 04
10
28
35000
13
Operand Digits
Non-aligned
10
00 16 99
10
00:00
41
1
12 31
E2
10
622
29
ce
80
Operand 2 Digits
Operand Digits
Operand 2 Digits
Figure 7: Accuracy of the addition task tested on operands with up to 14 digits, with models trained and evaluated with and without digit-aligned, zero-padded, and non- aligned mnemonic formats. The red box indicates the number of digits used during training.
5 Conclusions
We argue that, while the attention mechanism of Transformers is well-suited to perform content-based addressing into the context window, it struggles with random token accesses- crucial capability in virtually all algorithmic reasoning tasks. We present supporting evidence for this hypothesis by demonstrating the effectiveness of methodologies that either circumvent the need for indexing, such as the interleaved scratchpad, or enable indirect random token access through content-based addressing via mnemonics. Additionally, we illustrate where and how failures in index-based retrieval manifest using attention map visualizations.
Our work demonstrates that Transformers can in fact learn to length-generalize in algo- rithmic tasks, such as parity and addition, as long as they are able to perform random memory access. This suggests that equipping these models with the ability to perform such index-based addressing- either into their own context window, or into an external memory- mmy be key to enabling them to learn algorithmic tasks more generally.
10

---

Published as a conference paper at COLM 2024
References
Cem Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowycz, Vedant Misra, Vinay Ra- masesh, Ambrose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. Exploring length generalization in large language models. Advances in Neural Information Processing Systems, 35:38546-38556, 2022.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pp. 2397-2430.
PMLR, 2023.
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: Early experiments with gpt-4. ar Xiv preprint ar arXi::2303.122 127122, 2023.
Mirelle Candida Bueno, Carlos Gemmell, Jeff Dalton, Roberto Lotufo, and Rodrigo Nogueira. Induced natural language rationales and interleaved markup tokens enable extrapolation in large language models. In Proceedings of the 1st Workshop Mathematical Natural Language Processing (MathNLP), PP- 17-24, 2022.
Yann Dubois, Gautier Dagan, Dieuwke Hupkes, and Elia Bruni. Location attention for extrapolation to longer sequences. ar arviv preprint ar Xiv:1191 91.09987,, 2019.
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jiang, Bill Yuchen Lin, Sean Welleck, Peter West, Chandra Bhagavatula, Ronan Le Bras, et al. Faith and fate: Limits of transformers on compositionality. Advances in Neural Information Processing Systems, 36, 2024.
Dieuwke Hupkes, Verna Dankers, Mathijs Mul, and Elia Bruni. Compositionality decom- posed: How do neural networks generalise? Journal of Artificial Intelligence Research, 67: 757-795, 2020.
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding length generalization in transformers. Advances in Neural Information Processing Systems, 36, 2024.
Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Iliè, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, et al. Bloom: A 176b-parameter open-access multilingual language model. 2023.
Amirkeivan Mohtashami and Martin Jaggi. Random-access infinite context length for transformers. Advances in Neural Information Processing Systems, 36, 2024.
Niklas Muennighoff, Thomas Wang, Lintang Sutawika, Adam Roberts, Stella Biderman, Teven Le Scao, M Saiful Bari, Sheng Shen, Zheng-Xin Yong, Hailey Schoelkopf, et al. Crosslingual generalization through multitask finetuning, ar arxiv preprint arXiv:2211.01788, 2022.
Benjamin Newman, John Hewitt, Percy Liang, and Christopher Manning. The eos decision and length extrapolation. arXiv preprint arXii:2010.07174, 2020.
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowyez, Maarten Bosma, David Luan, et al. Show your work: Scratchpads for intermediate computation with language models. arXio preprint ar Xxi::2112.00011 2021.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, al. Training language models to follow instructions with human feedback. Advances Neural Information Processing Systems, 15527730-22744 2022.
11

---

Published as a conference paper at COLM 2024
Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arx preprint wxXi:22108.12400 2021.
Jing Qian, Hong Wang, Zekun Li, Shiyang Li, and Xifeng Yan. Limitations of language models in arithmetic and symbolic induction. ar arxio preprint arXiv:2208.0500 05051, 2022.
Gabriel Recchia. Teaching autoregressive language models complex tasks by demonstration. CoRR, abs/2109.02102, 2021. URL https://aaxii: http://raivv.oraaaaa////aaa httggabb/21000 021022.
Ruoqi Shen, Sébastien Bubeck, Ronen Eldan, Yin Tat Lee, Yuanzhi Li, and Yi Zhang. Posi- tional description matters for transformers arithmetic. arXiv preprint ar arXio:2311.14737, 2023.
Avanti Shrikumar, Peyton Greenside, Anna Shcherbina, and Anshul Kundaje. Not just black box: Learning important features through propagating activation differences. Xrv preprint ar arXie:1605.01713, 2016.
Koustuv Sinha, Robin Jia, Dieuwke Hupkes, Joelle Pineau, Adina Williams, and Douwe Kiela. Masked language modeling and the distributional hypothesis: Order word matters pre-training for little. arXiv preprint arXiv:2104.066: 06644, 2021.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputiing, 568:162882 1666122763, 2024.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. ar Xrv preprint Xix::2330 Xii::2307.09288, 2023.
Jason Wei, Xuezhi Wang. Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V
Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 35:24824-24837 2022.
Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In International Conference on Machine Learning, pp. 11080-11090. PMLR, 2021.
Yuhuai Wu, Albert Qiaochu Jiang, Wenda Li, Markus Rabe, Charles Staats, Mateja Jamnik, and Christian Szegedy. Autoformalization with large language models. Advances in Neural Information Processing Systems, 35:32353-32368 2022.
Zhiheng Xi, Wenxiang Chen, Xin Guo, Wei He, Yiwen Ding, Boyang Hong, Ming Zhang, Junzhe Wang, Senjie Jin, Enyu Zhou, et al. The rise and potential of large language model based agents: A survey, ar preprint ar siv:2309.078644 2023.
Changnan Xiao and Bing Liu. Conditions for length generalization in learning reasoning skills. Xiv preprint arXiv:2311.16173, 2023.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient stream- ing language models with attention sinks. arXiv preprint ar 2rX1to2309.17453, 2023.
Fei Yu, Hongbo Zhang, and Benyou Wang. Nature language reasoning, a survey. arXiv preprint ar Xix:::2333.14225 2023.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint ar 2rx::2205.000668 2022a.
Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: a synthetic reasoning task. ar preprint arXiv:2206.04300 2022b.
Ziyin Zhang, Chaoyu Chen, Bingchang Liu, Cong Liao, Zi Gong, Hang Yu, Jianguo Li, and Rui Wang. Unifying the perspectives of nlp and software engineering: A survey on language models for code. arXiv preprint ar ari: 2023.
12

---

Published as a conference paper at COLM 2024
Hattie Zhou, Azade Nova, Hugo Larochelle, Aaron Courville, Behnam Neyshabur, and Hanie Sedghi. Teaching algorithmic reasoning via in-context learning. aiv preprint $tt::2211009066, 2022.
Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Josh Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? a study in length generalization. Xrv preprint arXiv:2310.16025 2023.
Yongchao Zhou, Uri Alon, Xinyun Chen, Xuezhi Wang, Rishabh Agarwal, and Denny Zhou. Transformers can achieve length generalization but not robustly, arXiv preprint 4rX1::2402.0937,, 2024.
13

---

Published as a conference paper at COLM 2024
A Details of Experiments
We initialized with the pre-trained weights for training, except when specified otherwise. We used a learning rate of le - 6 for parity and 2e for addition, with a 1000-step warm- up. The training consists of 4 epochs, each containing 8000 training steps, with batch sizes of 64 for parity and 32 for addition tasks. We ensured an equal number of training examples for each problem length, reserving 200 samples for parity and 32 for addition from each length for evaluation. When training from random initialization, we used 8 epochs, twice the number of epochs used in our fine-tuning settings.
During training, the loss is calculated only for the target tokens (indicated by bold blue tokens in the main text). During inference, when the next token is a target, we perform greedy decoding from the model; otherwise, place the correct token into the context window.
B Additional Experiment Results
B.1 Additional models trained on the parity task
Here, we present results similar to those shown in Figure for Pythia-410M with RoPe, and OPT-350M with learned positional embeddings.
OPT-350M
OPT:350M (Random Initilization)
0.8
Acotracy
0.6
0.4
No Scratchpad
No Scratchpad
Standard Scratchpad
Slandard Scratchpad
0.2
Mnemonia
Maemonics
Moem onics (Euv. Forced)
Moemonics (Bov. Forced)
0.0
Training Sequence Length
Training Sequence Length
40
Test Sequence Length
Test Sequence Length
Figure 8: Length generalization performance of the OPT-350M model on the parity task using different scratchpad strategies. Left: fine-tuning; Right: training from scratch.
Pythia-410M
Pythia Initilization)
1.0
0.8
Accuracy
0.6
0.4
No Scratchpad
No Scratelipdd
Standard Scratohpad
Standard Scratchpad
0.2
Moemm onies
Mnemonics
onces (Bov.
Mnemonics Forced)
Training Sequence Length
Training Sequence Length
15 Tast Sequence Length
20
40
45
60
20 25
45 50
Test Sequence Length
Figure 9: Length generalization performance of the Pythia-410M model on the parity task using different scratchpad strategies. Left: fine-tuning; Right: training from scratch.
14

---

Published as a conference paper at COLM 2024
B.2 Exploring mnemonic intervals
Here, investigate the effectiveness of reducing the number of mnemonics within the parity scratchpad. At a innemonic interval of i, mnemonic tokens are inserted before every i-th bit in the input and output sequences. Therefore, mnemonic interval of 1 token corresponds to the original mnemonic format described in the main text. For instance, with mnemonic interval of the format would be as follows:
Mnemonics with interval of
>>> M₁
M2 1 0 M3 0
1 M4 0 === M₁
M2
0
M4 11
As shown in Figure 10, length generalization performance remains largely unaffected with mnemonic intervals of up to 3 tokens. However, when the interval exceeds 5 tokens, the impact of mnemonics begins to diminish.
0.8
Accuracy
0.6
Mnemonic Interval:
Mnemonic Interval:
Mnemonic Interval:
Mnemonic Interval:
Mnemonic Interval:
Mnemonic Interval: 10
35
40
45
50
55
60
Test Sequence Length
Figure 10: Length generalization performance of fine-tuned BLOOMZ-560M models with non-environmentffored mnemonics of different intervals in the scratchpad.
B.3
In-context learning with mnemonics
1.0
No Scratchpad
Standard Scratchpad
Mnemonics
0.8
Mnemonics (Env. Forced)
Length of Context Examples
Accuracy
0.4
0.2
5
20
25
35
40
50
55
Test Sequence Length
Figure 11: Length generalization performance of a Llama2-7B model on the parity task, with in-context examples (3 examples per length) with and without mnemonics.
15

---

Published as a conference paper at COLM 2024
We investigate the in-context learning capabilities, without fine-tuning, of larger Trans- former model, Llama2-7B (Touvron et al., 2023), in performing the parity task with and with- out mnemonics. We use examples of lengths 10 to 20, with three examples for each length. Additionally, we preface the examples with the problem statement prompt: "Cal culate the running parity of the sequence after ===". Figure 11 illustrates the model's per- formance with and without the use of mnemonics (refer to Section 3.2). Similar results were also observed with Llama2-7B-chat and BLOOMZ-7.18 models.
B,4 Unseen (OOD) mnemonics at test time
In this section, investigate whether the model treats mnemonics merely positional anchors, disregarding their values, or learns to memorize the mnemonic tokens for indexing. Following the methodology described in Section 3.2, we fine-tune model using single-token English words as mnemonics. In contrast, at test time, we use single-token integers as mnemonics.
Figure 12 presents the results of length generalization performance for models evaluated on in-distribution and out-of-distribution mnemonics. It shows that performance degrades when a model is evaluated on unseen, semantically novel mnemonics. This suggests that the learned approach to using mnemonics still relies on token values.
0.93
Accuracy
0.90
0.80
Mnemonics
Mnemenics (Env. Forced)
Mnemonics
0.75
Mnemonics (Env. Forced)
Training Sequence Length
5
10
25
30
45
50
55
Tost Sequence Length
Figure 12: Length generalization performance fine-tuned BLOOMZ-560M models, tested using in-distribution (ID) and out-of-distribuiion (OOD) mnemonics. Note that the y-axis is truncated, with values ranging from 0.7 to 1.
16


Reviews:
[{"review_id": "uFO9LzMNNu", "reviewer": "Reviewer_Uz62", "summary": "The paper is focused around the issue of length generalization in LLM models in the arithmetics tasks. The authors’ hypothesis is that this problem is strongly connected to the ‘random’ addressing problem, i.e. the ability to exactly address a token by its position. Unlike for natural language, for arithmetic tasks the positional information is equally important to the token content; hence this is less of a problem for the NL tasks than for the arithmetics.\n\nAs a motivating example, the paper starts with a task of calculating XOR of a bit sequence. Arbitrary input-length calculation is not solvable by a Transformer as it would have a fixed computational power; yet, by adding the chain-of-thought (scratchpad) output with the running partial answer and an incoming bit, that should be alleviated. However, the experiments show this not to be the case. In contrast, in the case where the incoming bit is externally fed, in the same format, the model works perfectly across lengths. From that, the authors conclude that the main issue is the inability to retrieve the correct input bit.\n\nIn the following experiments, the authors improve the addressing capability by adding special mnemonic tokens in the sequence.. Those are randomly sampled (w/o replacement) tokens that serve as addressing anchors. This modification seems to generalize well across sequences even when the mnemonic tokens are not externally forced. They also work better than numeric, constant, non-aligned, and cyclic anchors. \n\nIn the next steps the authors also provide attention pattern analysis, experiment with the addition tasks. Appendix gives bits of additional insights, e.g. experiments with the mnemonic-encoded intervals.", "questions": "", "limitations": "", "rating": 7, "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": "Thank you for your constructive feedback and suggestions.\n\n>What I think is missing ...\n\nWhile we study these specific tasks, the problem should affect any task that requires random memory access, which could include many NLP tasks, such as counting occurrences of references. We will clarify this point in the paper.\n\n>Authors highlight that their ...\n\nOur work highlights that this is a fundamental problem for which there is currently no general solution. It also suggests that learning to procedurally search in memory could be a fruitful direction toward a potential solution.\n\n> Zooming in on the ablation ...\n\nWe acknowledge your observation and note that, besides being monotonic, numerical indices are also fixed across all samples, making them 1) fixed, 2) monotonic, and 3) partially OOD during test time.\n\nWhile our experiments suggest that your idea of using a randomized starting point improves length generalization, it remains unclear whether this is due to alleviating the OOD problem or disrupting fixed index patterns.\n\nTo further isolate these factors, we propose using OOD tokens at test time, while using a randomized starting position for numerical indices during training. This will assess if the model has learned to disregard the specific values of the indices and treat them merely as positional anchors.\n\nWe will include a discussion and corresponding experiments in the final version of the paper.\n\n> The paper says that “.. the introduction of ...\n\nWe consider this strategy impractical for two reasons. It doubles the token count required for tasks, and, as seen in the scratchpad design for addition, it requires careful placement of mnemonics to align the digits in the scratchpad format. This could become tedious for more involved algorithms, such as multiplication.\n\n> The paper deals with two arithmetics-based ...\n\nAs we discussed above, a more principled solution than mnemonics, although out of scope for this work, would be a learned procedural search capability, which should be applicable even in tasks without strict alignment patterns. \n\n> I assume that the interleaved ...\n\nYes, the input bits are environment-forced, and the model predicts the running parity at each step.\n\n> Suggestion: Please consider using ...\n\nThank you. We will apply this in the final version. \n\n> Comment: the introduction has ....\n\nThe point made in the paper indeed supports our argument on content-based vs. index-based addressing, which we will include in the final paper.", "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": "* I really like that the paper provides a focused, well-explained and in-depth study of a single problem. It tries to motivate the problem, nail down the root cause and find a grounded solution for it.\n* To my taste, the experiments are clever, illuminating and nicely support the hypothesis/conclusions.\n* The paper allows us to get a better inner model of Transformers", "reasons_to_reject": "Playing a devil’s advocate:\n* What I think is missing is a motivation why/when this is important beyond arithmetics. If that is only a problem for the arithmetics, this can be, in principle, solved by tool-using a calculator/python interpreter.\n* Authors highlight that their proposed mnemonic solution is not a practical one. What does those findings tell us about a potential practical solution?", "questions_to_authors": "* Zooming in on the ablation where ‘mnemonics’ anchors are replaced by the numerics (which are akin to absolute positional embeddings). Now, there are two differences between those two: (a) randomized order of the mnemonics vs strict order sequential of numerics, and (b) ‘out-of-vocabulary’-ness of the numerics on longer sequences (e.g. 22 is not seen at training).\n\nNow, can we nail down which one makes the mnemonics work? For instance, the randomized order can make the model be really careful to pay attention to the input, while with monotone numbers it can just put the a priori known order w/o looking at the input even. Alternative explanation is that the model indeed is lost when it gets an OOD anchor in the input.\n\nCan we design an experiment where we differentiate the two above cases? Eg, always use monotone numerical anchors from the range of 0…100, but starting at a random position (10, 11, 12, …, 30 for one sequence and 0, 1, 2, …, 20 for another). This way the OOD issue disappears (for sequence of length 30 we can go w/ 0…29) but the monotonicity is preserved.\n\n(Maybe this exact proposal is wrong, but I am interested if we can come up w/ a way to distinguish the causes for the behavior).\n\n* The paper says that “.. the introduction of mnemonics is not proposed as a practical fix, it highlights the underlying issue and reinforces our hypothesis.” Could you please expand (a) why you would deem it impractical, and (b) what could be a practical step from here?\n\n* The paper deals with two arithmetics-based experiments where the alignment is very much fixed. Can we carry over any ideas from here to tasks with less strict alignment? Or even would it be any useful beyond arithmetics, for eg NMT tasks?\n\n* I assume that the interleaved scratchpad has forced input bits, is that correct?\n\n* Suggestion: Please consider using different labels for points in the figures, not solely relying on the color for differentiating them.\n* Comment: the introduction has a bit of a discussion on relative importance of words vs their positions. I feel that [1] would be a fitting citation to show that one can go a very long way in natural language w/o even having correct positional information.\n\n [1] https://arxiv.org/abs/2104.06644", "extra_scores": {}}, {"review_id": "ammeXBSP4b", "reviewer": "Reviewer_yGJY", "summary": "This paper argues that certain reasoning tasks, even with a scratchpad/chain-of-thought, are difficult because they depend on precise position information in a way that natural language does not. It studies two example tasks, parity and (more briefly) addition, and explores several variations on these tasks to see what helps and what doesn't.\n\nThe three things that help are:\n\n- For parity, forcing the model to copy the input and fine-tuning it to generate the state interleaved with the copied input.\n- For both tasks, prompting it with the input interleaved with \"mnemonics\" and fine-tuning it to generate the state interleaved with the mnemonics.\n- Same, but forcing the second generation of the mnemonics.\n\nMany contrastive experiments are performed to show that it is accessing information at a precise, non-local position that is difficult.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> The paper seems to be using the wrong font.\n\nWe thank the reviewer for their attention to detail. The incorrect font usage resulted from an issue with importing a package, which we have now corrected.\n\n\n> For parity, why didn't you try the interleaved scratchpad without environment-forcing?\n\nWhile mnemonic tokens can optionally be environment-forced, in the interleaved scratchpad version, the input sequence bits are always environment-forced, as without the model would not know the problem input. We hope this clarifies the reviewer's question.\n\n\n> Although you tried several different established position embeddings, why not encode the \"mnemonics\" into the position embedding instead of interleaving them? That is, for your parity example, you could prompt the model with word/positions:  \n1/1 0/2 1/3 0/4 0/5 1/6 1/7 ===/8  \nand then the completion could be  \n1/1' 1/2' 0/3' 0/4' 0/5' 1/6' 0/7'  \nwhere 1' is the same position embedding as 1, but possibly modified somehow.\n\nThank you for the suggestion. If you are referring to an embedding-level mnemonic, such as appending real-valued vectors to positional embeddings, it's an interesting idea. We believe this could potentially help with length generalization, similar to the effects observed with token-level mnemonics.\n\nHowever, if you mean adding \"/position\" as token-level mnemonics to prompt the model with 'word/positions,' this approach aligns with the numeric mnemonic format discussed in Section 3.3.\n\nIf we have misunderstood your suggestion, please let us know so we can address it further.", "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": "This is an interesting phenomenon and the experiments are fairly convincing.", "reasons_to_reject": "The experiments are limited to two tasks, and the addition experiments are not as thorough as the parity experiments.", "questions_to_authors": "The paper seems to be using the wrong font.\n\nFor parity, why didn't you try the interleaved scratchpad without environment-forcing?\n\nAlthough you tried several different established position embeddings, why not encode the \"mnemonics\" into the position embedding instead of interleaving them? That is, for your parity example, you could prompt the model with word/positions:\n\n1/1 0/2 1/3 0/4 0/5 1/6 1/7 ===/8\n\nand then the completion could be\n\n1/1' 1/2' 0/3' 0/4' 0/5' 1/6' 0/7'\n\nwhere 1' is the same position embedding as 1, but possibly modified somehow.", "extra_scores": {}}, {"review_id": "EPju66LcTb", "reviewer": "Reviewer_B5h5", "summary": "In this paper, the authors study transformer's ability to perform length generalisation when the task to be solved involves index-based search.\n\nTo do so, the authors first focus on the binary parity task and fine-tune three models (which use three different positional encoding mechanisms: AliBi, RoPE, and learned positional embeddings). The three models are able to perform the task almost perfectly for lengths similar to ones seen during training but completely fail when increasing the length.\n\nThe authors then simplify the task by allowing the models to keep a scratchpad. However, as this still involves performing index-based addressing, the models still fail. But when having an interleaved scratchpad, as the model only needs to attend to the two previous tokens, at each step, its accuracy its almost perfect. \n\nTo further illustrate this, the authors also experiment inserting mnemonics in the sequence (adding matching random anchor tokens before pairs of corresponding tokens) so that the model can solve the task by performing content-based addressing. The transformer models are able to solve this task with very good accuracy, showing that the attention mechanism is well suited for content-based addressing but not for index-based addressing.\n\nThe authors also show through input attribution plots, that the attention mechanism clearly fails to attend to the correct bits when the length of the sequence is bigger than the ones seen during training.\n\nTo conclude the authors, perform similar experiments on a different task: multi-digit addition, which shows similar results.", "questions": "", "limitations": "", "rating": 7, "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": "> I think the paper could be even more interesting if the authors explored possible solutions to improve trasnformer's random token access ability or at least proposed some possible solutions for future work.\n\nThank you for your suggestion. \nOur work identifies random-accessing as a fundamental issue that currently lacks a general solution. A more principled approach, although out of scope for this work, would be a learned procedural search capability. We will also include this as a direction for future research.\n\n> The paper could also explore \"real-world\" tasks for which transformers might struggle due to the lack of index-based search capabilities.\n\nWhile the purpose of this work is to identify the core issue, we acknowledge that solving random memory access could benefit many real-world tasks, including most algorithms. We will include a discussion on this topic in the paper. As previously mentioned, we consider learning to procedurally search in memory as an important direction for future research.", "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": "- The paper is very well written and is easy to follow.\n- The authors study an interesting flaw of transformers (or current positional encoding mechanisms) that might prevent their ability to perform algorithmic tasks.\n- The authors present interesting evidence that show that transformer models (with the tested positional encoding mechanisms) are very good at performing content-based addressing but fail to length generalise when having to perform index-based addressing.\n- The paper provides in-depth analysis of the experiments.", "reasons_to_reject": "- I think the paper could be even more interesting if the authors explored possible solutions to improve trasnformer's random token access ability or at least proposed some possible solutions for future work.\n- The paper could also explore \"real-world\" tasks for which transformers might struggle due to the lack of index-based search capabilities.", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "02GygtIaAj", "reviewer": "Reviewer_VRhj", "summary": "This paper studies length generalization of transformers using parity and addition as prototypical tasks and identifies _index-based access_ as a cause of failures of generalization.  In particular, they show that a scratchpad does not help with parity, but that an _interleaved_ scratchpad enables near-perfect length generalization.  To fully test the index-based access idea, they introduce \"mnemonic\" tokens into the input and scratchpad, which provide a direct anchor between output tokens and the corresponding input token.  They show that this method also enables perfect generalization (NB: when the mnemonic token pool is large enough and they are randomly sampled without replacement each time; see questions below).  Other forms of mnemonic do not work, suggesting that it really is the model's ability to know which position to attend to that hinders length generalization in the normal case (an attention pattern analysis also supports this).  Similar results are shown for addition as well.  These results are a clever use of scratchpad format to do detailed analysis of the factors which prevent and enable length generalization in transformers; while the authors do not explicitly recommend these formats for specific tasks, they are to be commended for the careful and detailed analysis.", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> The introduction of position- vs content-based addressing was a bit fast and loose (e.g. the discussion of a PP-attachment ambiguity at the bottom of page 1).\n\nWe will add more clarification on positional vs content-based addressing in natural language tasks. This is also related to [this reference](https://arxiv.org/abs/2104.06644) suggested by Reviewer 1, which we will include in the final version. \n\n>Only reports results of fine-tuning; I'd be curious to see if their methods also work in an in-context learning setup.\n\nThe models in the size range we explored do not perform well with in-context learning on these tasks (even though they do very well when trained with mnemonics).\nWe are exploring the possibility of in-context learning for mnemonics using larger models.\n\n> is it possible that some longer digits are getting split up?\n\nWe can confirm that this is not the case for numeric mnemonics. We specifically ensured that all mnemonics presented in the paper were encoded as single tokens.\n\n> did you try an \"O.O.D. mnemonic generalization\"\n\nThank you for the interesting suggestion. Based on your feedback and also comments from Reviewer 1, we will include discussions and results for scenarios where unseen mnemonic tokens are introduced during test time. This will assess whether the model has learned to disregard the specific values of the indices and treat them as mere positional anchors. We also refer you to our response to Reviewer 1 for a related discussion.\n\n>I think some wording in 3.1 could help clarify this\n\nFollowing your suggestion, we will include clarifications on the OOD nature of numerical mnemonics.", "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": "* Detailed analysis of failure and success modes of language models on length generalization parity and addition.\n* Identifies random/position-based access (as opposed to content-based access) as a difficulty.\n* Construction of careful scratchpad formats using \"mnemonic tokens\" both enables length generalization and helps to diagnose failure modes.", "reasons_to_reject": "* The introduction of position- vs content-based addressing was a bit fast and loose (e.g. the discussion of a PP-attachment ambiguity at the bottom of page 1).\n* Only reports results of fine-tuning; I'd be curious to see if their methods also work in an in-context learning setup.", "questions_to_authors": "* I really enjoyed the various mnemonic ideas, but was puzzled by this: \"For example, with the numeric indices, the model is exposed to mnemonic tokens 1, 2, · · · , 20, but encounters unseen mnemonics 21, 22, · · · in test time.\"  How is this different for numerical mnemonics than for the mnemonics results reported in Figure 3?  I think I eventually figured this out: the $M_i$ in Fig 3 are sampled from the same overall pool for each instance, so it is possible that every $M_i$ will be seen during training even if no inputs of a certain length have been seen.  I think some wording in 3.1 could help clarify this, or at least-referencing the **Note:** in the pnemonics figure when discussing the numerical indices later on (it at least took me awhile to unpack all of this).  Also: I guess you wouldn't be able to length generalize beyond the _total_ number of mnemonics for the same reason.  E.g. if you have 100 total $M_i$ tokens and train on seqs of length up to 20, you'd expect length generalization up to 100 but then not beyond; is that right?\n\n* Numeric mnemonics: it strikes me as interesting and a little perplexing that the models do not generalize to longer lengths with numerical index mnemonics.  (If the model learned a general rule about \"look to the right of the same token\" [or the equivalent generation rule in the non-environment-forced case], it shouldn't matter which tokens it has and hasn't seen as mnemonics.)  Do you know if all of the digits you use are full tokens in the relevant vocabs, or is it possible that some longer digits are getting split up?  Related to this and my above question: did you try an \"O.O.D. mnemonic generalization\", where the mnemonic tokens at inference time can be different, but are the same ones as in the first experiment (i.e. single tokens with whitespace first)?", "extra_scores": {}}]

Author rebuttal:

Meta-review: The reviewer appreciated this analysis of length generalization, and found the results interesting and useful.
The authors may find this paper on markups somewhat relevant (and possibly related papers which use different markings to "help the model count"): https://arxiv.org/abs/2208.11445
Decision: Accept