Exploring Length Generalization in Large Language Models
Abstract
The ability to extrapolate from short problem instances to longer ones is an important form of out-of-distribution generalization in reasoning tasks, and is crucial when learning from datasets where longer problem instances are rare.
These include theorem proving, solving quantitative mathematics problems, and reading/- summarizing novels.
In this paper, we run careful empirical studies exploring the length generalization capabilities of transformer-based language models.
We ﬁrst establish that naively ﬁnetuning transformers on length generalization tasks shows signiﬁcant generalization deﬁciencies independent of model scale.
We then show that combining pretrained large language models’ in-context learning abilities with scratchpad prompting (asking the model to output solution steps before producing an answer) results in a dramatic improvement in length generalization.
We run careful failure analyses on each of the learning modalities and identify common sources of mistakes that highlight opportunities in equipping language models with the ability to generalize to longer problems.
Introduction
Many natural problems, such as theorem proving and program synthesis, have a notion of length that strongly correlates with the difﬁculty of the task.
However, in these domains, the number of available problems typically drops rapidly as a function of problem length (e g. Figure 2).
Hence, it is desirable to learn from examples of shorter lengths to generalize to longer ones or at least reduce the number of samples required for longer examples.
We refer to this type of problem as length generalization .
Recent work on large language models (LLMs) has shown consistent improvement in their performance by scaling model and dataset size.
However, such models are still incapable of length generalization.
For example, [1] shows that even though scale helps with solving arithmetic problems, scale alone is likely insufﬁcient for learning to solve instances of arbitrary lengths.
This implies that models fail to learn the general algorithms that would enable this kind of generalization.
Indeed,
Razeghi et al.
[2] showed that the performance of LLMs on mathematical calculations correlates with term frequency in the training data.
This suggests that LLMs might have gained their current performance from surface-level memorization instead of learning to apply the correct algorithm.
A recent line of work proposes to use a scratchpad, or chain-of-thought reasoning, when prompting
LLMs [3, 4, 5] on multi-step tasks.
Breaking down tasks into multiple small steps and presenting these steps to the model leads to improved performance across a variety of reasoning tasks including word problems, arithmetic, and code execution.
We perform a systematic study of length generalization with transformer-based large language models.
We consider problems in which learning an algorithm can in principle enable a model to extrapolate
Techniques
In
[Figure]
distribution
Out
[Figure]
of
[Figure]
distribution
Improves with scale
[Table]
(cid:51)(cid:51)
[Table]
(cid:51)(cid:51) from short examples to problems of arbitrary length.
In particular, we focus on two simple algorithmic tasks, parity and variable assignment , in which the model needs to keep track of a state in order to extrapolate to longer lengths (see Figure 1).
These problems are illuminating because their simplicity allows us to probe the failure modes as well as contrast the learned solutions with the ground truth algorithm.
They provide us with a setting to study how/when these large language models start to fail.
We study combinations of three kinds of techniques for LLMs: ﬁnetuning, few shot prompting (also referred to as in-context learning), and use of a scratchpad (also referred to as chain-of-thought), to understand the role of each method and the interplay among the three in length generalization.
Interestingly, we observe non-trivial interactions among the three techniques; see Table
Contributions Our main contributions are as follows
: • We deﬁne and characterize the problem of length generalization using notions such as state tracking, execution depth, and per-step error rate.
We study and carefully design two tasks, parity and variable assignment , that measure length generalization (Section 2).
• We ﬁnd that in the ﬁnetuning regime, scaling data, model sizes, and compute does not improve length generalization (Section 3.1).
We also observe that even when the model attains perfect in-distribution accuracy, it performs poorly in out-of-distribution domains.
Surprisingly, different hyperparameter choices for ﬁnetuning have a large effect on length generalization performance, while having minimal effect on the ﬁnal in-distribution performance (Section 3.3). • We establish ﬁnetuning with scratchpad also fails to generalize to longer problems, in contrast to what is suggested by previous works [3].
We look into three potential failure cases: positional encoding, the presence of distractors, and end of token prediction, and conclude that distractors are the main culprit of failures for length generalization (Section 4). • We show that in the in-context learning regime, use of a scratchpad shows a qualitatively different behavior and signiﬁcantly alleviates the decay of performance on longer problems.
This capability is signiﬁcant, as it implies that for LLMs, there are certain skills, like length generalization, that can be learned through in-context learning rather than through ﬁnetuning even in the presence of inﬁnite data.
This is in stark contrast to the common norms of machine learning (Section 5).
Length Generalization
Many sequence tasks—especially ones that require reasoning capabilities—have problem instances that differ in terms of their lengths .
Shorter instances are often easier to state, process, and handle, and require less compute to ﬁnd the answer.
By contrast, longer instances are more challenging to parse and require more compute to solve.
Tasks that have a reasoning component are especially well represented in this category — multi-hop reasoning [7], program execution [8], deductive reasoning [9] and theorem proving [10], to name a few.
Note that having to deal with differing problem lengths poses two signiﬁcant challenges.
First, it is often the case that one encounters longer problem instances than the ones ever encountered during training, and is required to extrapolate.
Second, even though longer problem instances have much more variety, real-world datasets often contain few long instances (see Figure 2).
Both of these challenges are exacerbated if learning agents are not able to generalize across and beyond the lengths they learn from during training.
This paper is about investigating to what extent transformer based language models are able to observe short problem instances and extrapolate to longer ones.
Instance Length as Number of Steps in a Markov Process It is possible to deﬁne problem length in many different ways to capture different aspects of problem difﬁculty.
Does there exist a notion of length that would expose the same length-generalization-related problem structure observed in qualitatively very different settings?
Such a framing would enable researchers to design algorithms and interventions that have the potential to generalize across a broad range of tasks.
To this end, we take the approach of characterizing length in the context of a deterministic Markov process .
From this perspective, length is simply the number of state transitions experienced by an initial world state.
In other words, the data-generation process can be described as sampling an (1) initial state and a (2) variable number of state transformations to be applied sequentially on the initial state.
The agent is provided both the initial state and the transformations, and is asked to predict the ﬁnal state.
This framing applies to a wide range of sequence problems, if not all of them—ranging from more mechanical tasks such as code and algorithm execution and theorem proving, to less structured tasks, such as solving math problems and summarizing novels.
In our empirical investigation we focus on two synthetic tasks: parity and variable assignment .
These tasks avoid problem-speciﬁc subtleties that could mislead our analyses, while strongly capturing the deterministic Markov process structure.
2.1 Tasks
Parity: The parity task is an age-old learning problem that requires the trained agent to predict whether a bit-string has an even or odd number of ones in it.
For example, the parity of the bitstring [0 , 1 , 1 , 0 , 1] is “odd" (or 1 ) as opposed to “even" (or 0 ), because there is an odd number of 1 s in the bit-string.
The parity task admits a sequential solution that enables length generalization in a straightforward way: simply process the bits left-to-right and record the parity of the bits processed so far as the state.
The default notion of length in the parity task is the number of bits in the input.
However, we also experiment with a version where the number of bits is kept constant, and the number of 1 s (i.e. the parity ﬂipping bit) is systematically varied.
The number of 1 s stands for the number of state changes contained in the input bit-string, and actually appears to capture a more relevant notion of length for transformer models (see Section 3.1). [Figure]
[Figure]
Boolean Variable Assignment Task:
The Boolean Variable Assignment task is designed to capture arbitrarily long, potentially branching unidirectional execution ﬂows.
An instance of this task can be seen in Figure 1.
The inputs consist of semantically correct (i.e. bug-free)
Python programs in which each line contains a boolean variable assignment operation.
The output is simply the value of the variable presented in the ﬁnal line of the program.
The sequential solution to this task is to simply execute the program line by line while keeping track of the state of all variables.
The data generation procedure involves randomly generating execution ﬂows that involve Boolean operations; see Supplementary Material (SM) for details.
We focus our evaluations on two variants of this dataset.
(1) The diverse variable assignment split consists of a wide range of boolean operators available and is intended to contain maximally diverse programs.
The chain-like variable assignment split consists only of operations that compose the values of already deﬁned variables.
This results in long chains of dependencies between the initial values of the variables and the queried one, ensuring that there are almost no redundant operations in the program (i.e. operations that can be removed without affecting the output of the program).
This split emphasizes the sequential nature of the variable assignment problem.
Standard Finetuning Fails at Length Generalization
We begin by demonstrating that ﬁnetuning transformer models on length-generalization tasks results in poor out-of-distribution performance.
In experiments we use LaMDA 2 decoder-only models.
These checkpoints were trained using general natural language data.
We use the AdaFactor optimizer [11] during ﬁnetuning, and tune the learning rate, batch size and dropout.
We trained the networks until the in-distribution validation accuracy settles (20000 gradient steps for parity and 18000 gradient steps for variable assignment).
The loss was only computed on the target tokens (i.e. the model wasn’t trained to model the input questions).
3.1 Scale Doesn’t Improve Length Generalization
Parity: We ﬁnetuned four pretrained LaMDA models with 244m, 422m, 1b and 64b parameters on the parity task, where the training distribution included randomly sampled bitstrings of length 10 to 21.
We then evaluated the performance on bitstrings of length 3 to 40; see Figure 3.
We ﬁnd that model scale has a little effect on length generalization.
Variable Assignment: We ﬁnetuned the same models on the chain-like Variable Assignment Task, described in Section 2.
We kept the in-distribution lengths at 3 to 11, and evaluated the test performance on lengths 3 to 19.
The results can be seen in Figure 3.
Just like in the parity task, while the in-distribution performance is (near) perfect, out-of-distribution performance degrades rapidly as length increases.
To get a sense of just how weak the out-of-distribution performance is, we also trained a 422m model on the same dataset, except we shufﬂed the operations before feeding it to the
[Figure]
[Figure] [Figure]
model.
This removes the sequential dependency between the operations, and helps us establish a strong baseline that only predicts the answers based on non-sequential, spurious correlations.
The accuracy-length curves for the baseline can be found in SM.
3.2 Transformers Prefer Parallel Strategies over Sequential Ones
The results presented in Section 3.1 establish that, when presented with sequential length generalization problems, transformers are biased toward learning non-sequential “shortcut” solutions that fail at longer problem instances.
We ran additional experiments to gain a better understanding of the nature of this generalization pattern.
On parity, we ran ﬁnetuning on a different distribution of bit-strings: Instead of ﬁrst randomly sampling the number of bits in the input bit-string, then sampling the values of the bits, we ﬁxed the total number of bits in the input, and only varied the number of ones in the bit-string uniformly.
We trained with 10 to 20 ones in the input distribution and tested on an interval containing 1 to 30 ones.
This makes sure that the number of tokens (now ﬁxed at 30 ) is now disambiguated from number of state changes, which for parity is equal to the number of ones.
The difference between in and out-ofdistribution performance is even starker for this data distribution (Figure 5): while in-distribution performance was 100 %, OOD performance was roughly equivalent to random prediction 3 .
This suggests that the transformers are learning a non-sequential solution that involves counting the number of ones in the input, and then thresholding the output.
This is not surprising, given that self-attention is an equivariant transformation capable of performing pooling operations like max-pooling [12].
This strategy doesn’t allow for knowledge transfer between problems of different lengths.
Note that this bottom-up counting behaviour is complementary to the left-to-right counting behaviour displayed by recurrent models Suzgun et al.
On the variable assignment dataset, we ﬁnetuned a 255m LaMDA model on the diverse split of the variable assignment dataset of programs up to 16 lines, and evaluated on the same data generating distribution up to 32 lines.
We measured the evolution of the model’s accuracy with respect to training iterations on different program lengths (quantiﬁed by number of lines).
The results are in SM.
We again observed that a different notion of length (which we call computational graph depth ) captures the difﬁculty of problem instances better than number of program operations.
A variable assignment program can be represented as a computational graph where each node corresponds to a variable, and each edge corresponds to an operation.
Computational graph depth is the length of the longest dependency chain that connects to the queried variable node.
This notion of length corresponds to the highly parallelizable strategy of executing programs by iteratively resolving computational graph dependencies.
We present two results that suggest that computational graph depth is a more relevant notion of length for transformers.
(1) Inspecting the order of problem instances in which the trained transformer correctly solves this task, we ﬁnd that performance is
[Figure]
strongest on examples with small computational graph depth, even if these examples are long in terms of number of operations.
The transformer does a good job of handling programs with an out-of-distribution number of operations, but for which computational graph depth is in-distribution.
3.3 In-Distribution Generalization Doesn’t Predict OOD Generalization on Length Generalization Tasks
Prior work on out-of-distribution generalization establishes that in many tasks, in-distribution loss is a strong predictor of out-of-distribution generalization [14].
Our experiments on the parity task indicate that the distribution shift induced by changing problem lengths falls outside of the this category.
Figure 5 shows how the same model trained on the same data achieving roughly the same in-distribution cross entropy loss behaves on OOD data, where the difference is solely induced by the choice of different hyperparameters.
Scratchpad Finetuning Still Fails at Length Generalization
It has been shown in prior work that it’s possible to get pretrained LLMs to solve a given task by not only outputting the answer, but also the solution steps behind it.
Nye et al.
[15] use scratchpad ﬁnetuning to achieve strong in-distribution performance on execution based tasks such as code execution and computing polynomials.
While they also report modest length generalization results on integer arithmetic, we ﬁnd that scratchpad ﬁnetuning suffers from similar length generalization pathologies than vanilla ﬁnetuning does.
The results on parity and variable assignment tasks can be seen in Figure 4.
The precise scratchpad strategies used for these tasks are described in detail in SM.
Error analysis: To understand the causes of failure in training scratchpad strategies, we focused on two architectural choices that could account for the poor performance: (1) how transformers encode position information, and (2) whether the transformers are trained to predict an end-of-sequence (EOS) token.
LaMDA models use T5 position biases [16] to handle position information.
If the network is only trained with short instances, position biases that handle longer positional distances might not be trained, explaining poor length generalization.
Similarly, Newman et al.
[17] report that networks trained with EOS token prediction often suffer from generalizing to longer problem instances, because of the models’ tendency to emit EOS tokens prematurely, as well as the EOS tokens’ effect on the representations that get learned.
We tested the extent to which these effects can explain lack of length generalization as follows.
We padded both the input bit-strings and the scratchpad content with dummy padding tokens to make the token count the same.
We also augmented the input and scratchpad targets with the same number of padding tokens on the left and right so that the relevant bit to attend to when executing the sequential scratchpad strategy corresponds to the same T5 position bias bin.
Examples of the updated input-target pairs can be seen in SM.
While this intervention helps, the trained models still display signiﬁcant length generalization issues.
To gain further insight about the source of the problem, we plotted how the scratchpad target prediction error rates change as a function of (1) how far along one is in constructing the scratchpad, and (2)
[Figure]
a cc u r a c y
Partial sequence accuracy training       i n s t a n c e l e n g t h
the length of the input bit-string.
The results can be seen in Figure 6.
The fact that the model makes mistakes in in-distribution scratchpad steps when the input has an OOD length implies that the attention mechanism isn’t capturing the relevant part of the input to form the scratchpad output.
See
SM for additional analysis.
Scratchpad Prompting Signiﬁcantly Improves Length Generalization
Wei et al.
[4], Nye et al.
[15] and Lewkowycz et al.
[5] showed that combining prompting (i.e. in-context learning) with scratchpad strategies present a powerful combination.
They demonstrate that pretrained LLMs, without the help of any ﬁnetuning, can solve grade school math word problems and execute pieces of code with nontrivial correctness [15], when prompted with the right scratchpad strategy.
We corroborate these ﬁndings, and report that scratchpad prompting endows pretrained
LLMs with the capability of variable length template matching (see Figure 8).
That is, in-context learning enables the model to “learn" solution steps from a small number of short instances, and apply the same template on signiﬁcantly longer instances with a high degree of accuracy.
5.1 Few-shot scratchpad
Contrary to vanilla and scratchpad ﬁnetuning, we ﬁnd that under the right conditions, few-shot scratchpad strategies sometimes signiﬁcantly improves LLMs’ capability to extrapolate to lengths much further than what pretraining weights grant them.
To evaluate the performance of few-shot conditioning with scratchpad inputs without any ﬁnetuning, we phrase the parity problem in natural language as a coin ﬂipping task.
An example for the few-shot prompts we used can be seen in Figure 8.
Wei et al.
[4] also report results on the coin-ﬂip task: the scratchpad format we used differs from theirs in that while ours respects the sequential nature of the task (i.e. each coin ﬂip corresponds to a step in the scratchpad solution), Wei et al.
[4]’s scratchpad strategy involves summing up the number of coin ﬂips, then deciding on the ﬁnal output based on the evenness/oddness of the sum.
Also, while they only test up to 4 ﬂips, we go up to 20 ﬂips while still attaining highly nontrivial accuracy levels.
For the variable assignment task, our scratchpad strategy involves copying over the program that’s being executed, with comments added in between lines specifying the value of the variable that was assigned in the line above.
Instances of this scratchpad strategy can be seen in SM.
Figure 7 shows the performance of the pretrained LaMDA 128b model on the coin-ﬂip version of the parity task.
Figure 8 shows an instance of how a length 3 prompt can induce the model to correctly output a 20 step scratchpad.
We ﬁnd that with the right scratchpad prompt, LLMs are able [Figure] to generate correct scratchpad solutions.
This reduces the problem to simply ﬁlling in the content of the generation correctly by inferring the right state transitions without having to ﬁgure out how to extrapolate the solution template.
Few-Shot Finetuning with Scratchpad Strategies: Does combining ﬁnetuning, few-shot prompting, and scratchpad strategies improve length generalization?
We ﬁnd that the answer is yes in the case of parity.
As seen in Figure 7, few-shot ﬁnetuning performs signiﬁcantly better than the baseline model, both on in- and out-of-distribution lengths.
Note that the vanilla (i.e. no shot) ﬁnetuning baseline also outperforms the no-ﬁnetuning baseline, it actually does worse on the larger lengths — a pathology that doesn’t appear with few-shot ﬁnetuning.
The results point to a qualitatively different picture for the variable assignment task.
Both few-shot ﬁnetuning and vanilla ﬁnetuning result in similar length generalization behavior (Figure 7).
We hypothesize that this distinction is caused by the different pretrained performances that the model displays on these tasks: while length generalization is already strong with no ﬁnetuning on parity, that’s not the case for variable assignment.
In the latter case, the model is forced to acquire a new skill via ﬁnetuning, which displays the same pathologies as zero-shot ﬁnetuning with scratchpad.
As a sanity check, we evaluated the (few-shot) ﬁnetuned performance of the pretrained model on an alternative, synthetic prompt style that yields poor performance without any pretraining: As expected by the aforementioned hypothesis, we observed that the few-shot ﬁnetuned model on this task also shows signiﬁcant length generalization pathologies.
The results can be found in SM.
We leave a more rigorous evaluation of this hypothesis as future work.
Related Works
There have been many attempts to study generalization from shorter/easier to longer/harder examples.
Challenges in length generalization: Several existing works have investigated pathologies that arise when models are asked to generalize to processing and generating longer (measured by number of tokens) sequences.
Newman et al.
[17] ﬁnd that sequence models trained with and in the absence of the end-of-sequence token display qualitatively different length extrapolation behaviour and learn different representations.
Dubois et al.
[19] proposes modiﬁcations to the commonly used dot-product attention to improve the models’ ability to extrapolate to longer sequences.
Murray and Chiang [20] demonstrate that neural machine translation models tend to have a bias towards generating shorterthan-desired translations.
Yehudai et al.
[21] show that length generalization issues are also present in training graph neural networks, where extrapolating across graph size presents a challenge.
Ju et al.
[22] propose a new attention mechanism to facilitate recurrent processing in transformer models.
Press et al.
[23] propose modifying transformer attention biases to facilitate generalization beyond the training context length.
Concurrent work [24] propose a synthetic dataset named LEGO (Learning
Equality and Group Operations), an instantiation of which resembles our [Table] where the only boolean operations allowed are assign and negate and assign , and overriding the values of variables is not allowed.
Their analyses on OOD generalization largely complement ours: while we focus on decoder-only architectures and scratchpad strategies as a way of carrying over state, they focus on encoder-only architectures, and investigate the effect of weight-sharing.
Easy-to-Hard generalization: Schwarzschild et al.
[25] and Bansal et al.
[26] use weight-tied neural networks to generalize from easy to hard examples.
Schwarzschild et al.
[25] also provide three tasks to benchmark easy-to-hard generalization.
Dehghani et al.
[27] and Kaiser and Sutskever [28] assess the capabilities of their proposed architectures on easy-to-hard generalization problems.
Inductive Biases Related to Lenght Generalization: McCoy et al.
[29] study the inductive bias of seq-to-seq learners on English question formation and English tense reinﬂection tasks and ﬁnd that LSTM and GRU networks often display differing strategies, caused by the use of differing activation functions.
Suzgun et al.
[13] ﬁnd that recurrent networks can perform dynamical counting, and encode hierarchical representations, which enables them to solve nontrivial Dyck tasks using k-counters.
Kharitonov and Chaabouni [30] also study the inductive bias of different architectures, and conclude that transformer and LSTM architectural have a tendency to learn hierarchical strategies, whereas CNN based strategies display more linear structure.
He et al.
[31] propose a method to learn natural inference models that are not biased on spurious correlations.
McCoy et al.
[32] show that transformer models that display strong performance in natural language inference can have superﬁcial biases that fool them in systematic ways and proposes a framework to think about these biases.
Conclusion
The ability to learn from shorter/easier problem instances to generalize to longer/harder ones is a key capability in a large number of tasks, especially ones requiring reasoning.
We deﬁned the concept of length generalization and measured language models’ length generalization capabilities.
After conducting careful experiments using ﬁnetuning, scratchpads, and few-shot prompting, we reached the following conclusions: (1) Generalizing in length is a challenge for language models at least up to the 100B parameter scale.
Both vanilla ﬁnetuning and ﬁnetuning with scratchpads suffer from a lack of length generalization caused by models’ tendency to pick up non-sequential pattern that don’t apply to longer problem instances.
(2) Few-shot scratchpad prompting enables pretrained large language models to pick up scratchpad-templates that extrapolate to arbitrary lengths, leading to dramatic improvements on longer problem instances.
Unlike raw ﬁnetuning, this approach does scale with model size [4].
(3) Trying to further enhance the performance of few-shot scratchpad prompted LLMs via ﬁnetuning yields mixed results, depending on the non-ﬁnetuned performance of the base model at the target task.
We emphasize that the aforementioned few-shot variable length pattern matching capability — something that doesn’t require changing model architecture — offers a qualitatively different approach to handle length generalization in contrast to prior art that introduced architectural modiﬁcations to achieve the same goal.
This capability is also signiﬁcant in that it implies that for LLMs, there are certain skills, like length generalization, that can be learned better through in-context learning rather than through ﬁnetuning, even in the presence of inﬁnite data.