Your Context Is Not an Array: Unveiling Random Access Limitations in Transformers

Authors: m reza ebrahimi
roland  memisevic
Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-20 22:15:11

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.

Main paper: 
 introduction: 
 The evolution of Transformer-based large language models (LLMs) has marked a 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 as the backbone of AI 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 et al., 2024) . Specifically, the models fail to learn simple algorithms to perform these arithmetic operations. This becomes apparent when models are applied to problems of greater length than those encountered during training (Hupkes et al., 2020) , a 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 a sequence of algorithmic steps. The second difference, and the focus of our work, is their reliance on formatting: arithmetic expressions are represented using a limited vocabulary, such as digits, with each token holding equal significance.Crucially, in the representation of arithmetic tasks, a token's position is as important as its value. This stands in stark contrast to natural language expressions, in which the coupling between token or word positions on the one hand and the meaning of the expression on 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, token references which are "content-based" in this way are well represented by the common attention mechanism prevalent in the Transformer, and they are further reinforced through pre-training on natural language. This is in contrast to arithmetic tasks, which rely exclusively on "index-based addressing" (random access memory) into the context window to retrieve the information necessary for generating the next algorithmic step. In this work, we provide an 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 to 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 a length generalizable algorithm for this task (Anil et 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 "mnemonics" 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 as a practical fix, it highlights the underlying issue and reinforces 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. 
related work: 
 Length generalization is a well-known problem in the context of Transformer-based sequence 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 length-generalize in compositional tasks, such as multi-digit multiplication, and highlights their generalization failures across zero/few-shot and fine-tuning regimes, both with and without the use of a 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 on 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. It 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 a slight improvement in generalization.A study similar in spirit to our work is Dubois et al. (2019) , albeit using recurrent sequenceto-sequence models instead of Transformers. That work hypothesizes that models equipped with separate content and location-based attention mechanisms are more likely to be able to extrapolate. It evaluates this hypothesis through variants of the Lookup Table task, designed to directly assess a model's performance in index-based addressing.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. 
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 is XOR. However, it is known that Transformer-based models struggle to learn the correct algorithm as their solution fails for sequences longer or shorter than those seen during training (Anil et al., 2022) .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 a potential complication: Transformers execute a fixed amount of computation for each token generated, yet the problem size can vary. In other words, the model must simulate a for-loop over the entire sequence in a single forward pass. Note that this represents a distinct contaminating issue that falls outside the scope of this work. This challenge can be addressed by incorporating a "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 a for-loop and output intermediate results.Adopting the format used in Anil et al. (2022) for the parity task, we begin with a start-ofsequence symbol >>>, followed by a binary sequence, an end-of-sequence symbol ===, and the sequence's running parity. For instance: >>> 1 0 1 0 0 1 1 === 0  >>> 1 0 1 0 0 1 1 === 1 1 0 0 0 1 0 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-sequence symbols are converted to single tokens and bits within the sequence are represented by single fixed tokens, preventing any merging due to tokenization.In essence, a length generalizable solution to 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.>>> 1 1 0 1 1 0 0 0 0 0 1 1 1 0We fine-tuned several small Transformer models with different positional embedding methods: 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 A 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. We can leverage content-based addressing in Transformers to indirectly perform index-based addressing, by adding matching "anchor" tokens before every pair of corresponding tokens in the standard scratchpad format. As they allow a model to 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 al. (2022) and Zhou et al. (2023) .During training and inference, for each example of length n, we first randomly sample n tokens from a pool of mnemonic tokens1 , then add the mnemonics before each bit in the input sequence and the running parity bits:Mnemonics >>> M 1 1 M 2 0 M 3 1 M 4 0 M 5 0 M 6 1 === M 1 1 M 2 1 M 3 0 M 4 0 M 5 0 M 6 1 Mnemonics (Environment Forced) >>> M 1 1 M 2 0 M 3 1 M 4 0 M 5 0 M 6 1 === M 1 1 M 2 1 M 3 0 M 4 0 M 5 0 M 6 1 Note: Mnemonic tokens M 1 , M 2 ,• • • are randomly sampled without replacement from the mnemonics pool, for every problem instance.Note that in the non-environment-forced version, the model is trained to first place the matching mnemonics from the input sequence, and then use them to address the active bit at each step. Conversely, in the environment-forced version, at each step, we first append the matching mnemonic from the input sequence to the context, after which the model predicts the running parity.Figure 3 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 to 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 a key to enabling it to learn correct arithmetic algorithms. Interestingly, the performance of the model using non-environment-forced mnemonics is nearly identical to that of the environment-forced 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 are 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. To further analyze how the model's attention changes with and without mnemonics, we present input attribution visualizations in Figure 5 , using the gradient×input method ( 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 task is to produce the running parity of the input sequence, at step i, it only needs to attend to the current bit (bit i 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 a model trained on sequences of length 10 to 20 bits.As shown in Figure 5 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 a length generalizable method for indexing the correct bit at each step, thus failing at indexing outside of its training regime. In contrast, as seen in the right plot of Figure 5 , when mnemonic bits are added, a near-perfect attention map is observed beyond the training regime.Published as a conference paper at COLM 2024Finally, we study several variations of the introduced mnemonic tokens, which further support our hypothesis, as discussed below:Numeric >>> 1 b 2 a 3 b 4 a 5 a 6 b === 1 b 2 b 3 a 4 a 5 a 6 b Constant >>> # 1 # 0 # 1 # 0 # 0 # 1 === # 1 # 1 # 0 # 0 # 0 # 1 Non-aligned >>> M 1 1 M 2 0 M 3 1 M 4 0 M 5 0 M 6 1 === M 7 1 M 8 1 M 9 0 M 10 0 M 11 0 M 12 1 Cyclic>>> red 1 green 0 yellow 1 red 0 green 0 yellow 1 === red 1 green 1 yellow 0 red 0 green 0 yellow 1We 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, b instead of 0, 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 to test whether the effectiveness of mnemonics is related to the attention sink phenomenon (Xiao et al., 2023) , or if the model uses the mnemonic tokens as "placeholders" allowing it to store intermediate calculations in their activations.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 a sequence of size n 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.Cyclic Mnemonics: 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.Figure 6 shows the failure of the aforementioned mnemonic variants at length generalization. Note that in the environment-forced versions, all 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, 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 performing 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 our 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 et al., 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 to 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 on operands with up to 14 digits.We use the same mnemonics for corresponding digits in both operands, as demonstrated below:Digit-aligned MnemonicsNo Mnemonics >>> 1 2 + 9 === 1 2 0 ### 0 2 1 Mnemonics >>> M 1 1 M 2 2 M 3 + M 2 9 M 3 === M 2 1 M 1 2 M 3 0 ### M 3 0 M 1 2 M 2 1 Env. Forced >>> M 1 1 M 2 2 M 3 + M 2 9 M 3 === M 2 1 M 1 2 M 3 0 ### M 3 0 M 1 2 M 2 1In 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 PaddingNo Mnemonics >>> 1 2 + 0 9 === 1 2 0 ### 0 2 1 Mnemonics >>> M 1 1 M 2 2 M 3 + M 1 0 M 2 9 M 3 === M 2 1 M 1 2 M 3 0 ### M 3 0 M 1 2 M 2 1 Env. Forced >>> M 1 1 M 2 2 M 3 + M 1 0 M 2 9 M 3 === M 2 1 M 1 2 M 3 0 ### M 3 0 M 1 2 M 2 1Lastly, we explore a format in which the mnemonics for corresponding digits of the two operands are not identical, as depicted below:Non-aligned MnemonicsNo Mnemonics >>> 1 2 + 9 === 1 2 0 ### 0 2 1 Mnemonics >>> M 1 1 M 2 2 + M 3 9 === M 3 M 2 1 M 1 2 0 ### 0 M 1 2 M 3 M 2 1 Env. Forced >>> M 1 1 M 2 2 + M 3 9 === M 3 M 2 1 M 1 2 0 ### 0 M 1 2 M 3 M 2 1The length generalization performance of the addition task, both with and without the specified mnemonic formats, is shown in Figure 7 . 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, our 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 can successfully learn the correct algorithm for the addition task.  
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 accessesa 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 algorithmic 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-may be key to enabling them to learn algorithmic tasks more generally. 
 
 Appendix: We initialized with the pre-trained weights for training, except when specified otherwise. We used a learning rate of 1e -6 for parity and 2e -6 for addition, with a 1000-step warmup. 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, we place the correct token into the context window.Here, we present results similar to those shown as in Figure 3 for Pythia-410M with RoPe, and OPT-350M with learned positional embeddings. Here, we investigate the effectiveness of reducing the number of mnemonics within the parity scratchpad. At a mnemonic interval of i, mnemonic tokens are inserted before every i-th bit in the input and output sequences. Therefore, a mnemonic interval of 1 token corresponds to the original mnemonic format described in the main text. For instance, with a mnemonic interval of 2, the format would be as follows:Mnemonics with interval of 2 >>> M 1 1 0 M 2 1 0 M 3 0 1 M 4 0 0 === M 1 1 1 M 2 0 0 M 3 0 1 M 4 1 1As 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. We investigate the in-context learning capabilities, without fine-tuning, of a larger Transformer model, Llama2-7B (Touvron et al., 2023) , in performing the parity task with and without 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: "Calculate the running parity of the sequence after ===". Figure 11 illustrates the model's performance with and without the use of mnemonics (refer to Section 3.2). Similar results were also observed with Llama2-7B-chat and BLOOMZ-7.1B models.B.4 Unseen (OOD) mnemonics at test timeIn this section, we investigate whether the model treats mnemonics merely as positional anchors, disregarding their values, or if it learns to memorize the mnemonic tokens for indexing. Following the methodology described in Section 3.2, we fine-tune a 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. We used all space-preceded tokens containing only English characters for the mnemonics pool. 


Decision: 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

Reviews: review: paper_summary: paper_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.

As 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.

In 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. 

In 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.
strength_weakness: strength_weakness: Playing a devil’s advocate:
* 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.
* Authors highlight that their proposed mnemonic solution is not a practical one. What does those findings tell us about a potential practical solution?
questions: questions: * 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).

Now, 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.

Can 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.

(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).

* 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?

* 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?

* I assume that the interleaved scratchpad has forced input bits, is that correct?

* Suggestion: Please consider using different labels for points in the figures, not solely relying on the color for differentiating them.
* 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.

 [1] https://arxiv.org/abs/2104.06644
paper_summary: paper_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.

The three things that help are:

- For parity, forcing the model to copy the input and fine-tuning it to generate the state interleaved with the copied input.
- For both tasks, prompting it with the input interleaved with "mnemonics" and fine-tuning it to generate the state interleaved with the mnemonics.
- Same, but forcing the second generation of the mnemonics.

Many contrastive experiments are performed to show that it is accessing information at a precise, non-local position that is difficult.
strength_weakness: strength_weakness: The experiments are limited to two tasks, and the addition experiments are not as thorough as the parity experiments.
questions: questions: The paper seems to be using the wrong font.

For parity, why didn't you try the interleaved scratchpad without environment-forcing?

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:

1/1 0/2 1/3 0/4 0/5 1/6 1/7 ===/8

and then the completion could be

1/1' 1/2' 0/3' 0/4' 0/5' 1/6' 0/7'

where 1' is the same position embedding as 1, but possibly modified somehow.
paper_summary: paper_summary: In this paper, the authors study transformer's ability to perform length generalisation when the task to be solved involves index-based search.

To 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.

The 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. 

To 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.

The 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.

To conclude the authors, perform similar experiments on a different task: multi-digit addition, which shows similar results.
strength_weakness: strength_weakness: - 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.
- The paper could also explore "real-world" tasks for which transformers might struggle due to the lack of index-based search capabilities.
paper_summary: paper_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.
strength_weakness: strength_weakness: * 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).
* Only reports results of fine-tuning; I'd be curious to see if their methods also work in an in-context learning setup.
questions: questions: * 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?

* 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)?

Comments: title: Thanks for the reply!
comment: I thank the authors for the detailed and helpful reply, and enjoyed reading the related discussion with another reviewer on O.O.D. mnemonics.  My score is already very high, indicating my belief that this paper should definitely be accepted, and so I will leave it as is.
Thanks for the reply.

Having read the authors' rebuttal, I've decided to maintain my score.