PaperlayerView exact revision on arXiv

Source-linked AI summary

Exploring Length Generalization in Large Language Models

Cem Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowycz, Vedant Misra, Vinay Ramasesh, Ambrose Slone, Guy Gur-Ari, Ethan Dyer, Behnam Neyshabur

arXiv:2207.04901v2cs.CLcs.LG

TL;DR

The paper investigates whether transformer language models can generalize from short reasoning problems to longer ones, where long examples are scarce. It systematically compares finetuning, in-context learning, and scratchpad strategies, finding that finetuning fails broadly while few-shot scratchpad prompting can substantially improve extrapolation.

  • Problem

    Language models show limited evidence of learning general algorithms that transfer from short reasoning instances to arbitrary lengths, despite the importance of such transfer when longer examples are rare.

  • Method

    The paper empirically studies parity and variable assignment using finetuning, few-shot prompting, scratchpads, and their combinations.

  • Results

    Few-shot scratchpad prompting sometimes significantly improves pretrained language models’ ability to extrapolate beyond the lengths granted by pretraining, whereas vanilla and scratchpad finetuning fail to generalize reliably.

  • Takeaways & Limitations

    Scratchpad prompting can enable pretrained models to learn solution-step templates from short examples and apply them to significantly longer instances.

Abstract

from arXiv · show

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 first establish that naively finetuning transformers on length generalization tasks shows significant generalization deficiencies 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.

1 Introduction

Length generalization asks models to learn from shorter instances and handle longer, harder ones, a need that arises because many reasoning datasets contain far fewer long examples. The paper systematically studies this ability in transformer language models using finetuning, prompting, and scratchpad techniques.

  • Motivation: Longer problem instances are harder and rarer in domains such as theorem proving and program synthesis, motivating transfer from short to long examples.The number of available problems typically drops rapidly as problem length increases.
  • Problem: Model and dataset scaling improves some capabilities but does not by itself yield length generalization or arbitrary-length algorithm learning.Prior work indicates scale helps arithmetic while remaining insufficient for solving instances of arbitrary lengths.
  • Approach: The study examines parity and variable assignment, two algorithmic tasks requiring state tracking and enabling comparison between learned behavior and ground-truth algorithms.Their simplicity supports analysis of failure modes and extrapolation to longer lengths.
  • Approach: Finetuning, few-shot prompting, and scratchpad reasoning exhibit different failure modes and trade-offs across in-distribution and out-of-distribution performance.The paper studies both individual techniques and their combinations.
  • Contributions: The paper defines and characterizes length generalization through concepts including state tracking, execution depth, and per-step error rate.These notions support the design and analysis of the parity and variable assignment tasks.

2 Length Generalization

The paper frames problem length as the number of steps in a deterministic Markov-style execution process, then studies this structure using parity and Boolean variable assignment tasks. These tasks isolate state tracking and sequential execution while supporting extrapolation to longer instances.

  • Length Generalization: Reasoning tasks vary in length, with shorter instances generally easier to process and longer instances requiring more computation.This pattern appears across sequence tasks including multi-hop reasoning, program execution, and deductive reasoning.
  • Instance Length as Number of Steps in a Markov Process: The paper treats instance length as the number of steps in a Markov process to expose shared length-generalization structure across qualitatively different tasks.This framing is intended to support broadly applicable algorithms and interventions.
  • Parity: Parity requires predicting whether a bit string contains an even or odd number of ones, with a sequential left-to-right solution that records the running parity.The task therefore has a direct algorithmic route to length generalization.
  • Boolean Variable Assignment Task: Boolean variable assignment uses bug-free Python programs of Boolean assignments and asks for the final-line variable value, modeling potentially branching unidirectional execution.Its sequential solution follows the dependencies needed to evaluate the queried variable.
  • Boolean Variable Assignment Task: The evaluations include diverse programs with many Boolean operators and chain-like programs with long dependency chains and little redundancy.The chain-like split emphasizes dependencies between initial variables and the queried variable.

3 Standard Finetuning Fails at Length Generalization

Standard finetuning achieves near-perfect in-distribution performance but generalizes poorly to longer instances. Across parity and variable assignment, transformers favor non-sequential shortcuts, with performance shaped by task length, computational graph depth, and hyperparameters.

  • 3.1 Scale Doesn’t Improve Length Generalization: Finetuned LaMDA models show little scale effect on parity length generalization, despite training on lengths 10 to 21 and evaluating lengths 3 to 40.The evaluated models range from 244m to 64b parameters.
  • 3.1 Scale Doesn’t Improve Length Generalization: Variable assignment accuracy is near-perfect in-distribution but degrades rapidly as length increases beyond the training range.Training lengths were 3 to 8, while evaluation covered lengths 3 to 19.
  • 3.2 Transformers Prefer Parallel Strategies over Sequential Ones: Transformers are biased toward non-sequential shortcut solutions that fail on longer sequential problem instances.This pattern is established across the standard finetuning experiments.
  • 3.2 Transformers Prefer Parallel Strategies over Sequential Ones: On parity with a fixed 30-token input, 100% in-distribution accuracy falls to roughly random OOD performance when the number of ones changes.The result suggests counting ones and thresholding rather than sequentially tracking parity.
  • 3.2 Transformers Prefer Parallel Strategies over Sequential Ones: Computational graph depth predicts variable-assignment difficulty better than the number of program operations.The transformer handles OOD operation counts when computational graph depth remains in-distribution.
  • 3.3 In-Distribution Generalization Doesn’t Predict OOD Generalization on Length Generalization Tasks: Models with the same architecture, data, and in-distribution loss can have very different OOD length-generalization performance because of hyperparameter choices.Figure 5 highlights sensitivity to learning rate and batch size.

4 Scratchpad Finetuning Still Fails at Length Generalization

Scratchpad finetuning does not resolve length-generalization failures: its pathologies resemble vanilla finetuning, and OOD input lengths can corrupt even early scratchpad steps. Error analysis points to failures in attention patterns for longer inputs.

  • 4 Scratchpad Finetuning Still Fails at Length Generalization: Scratchpad finetuning displays qualitatively similar length-generalization pathologies to vanilla finetuning.The comparison is reported for the parity and variable-assignment tasks.
  • 4 Scratchpad Finetuning Still Fails at Length Generalization: The analysis tests whether positional encoding and end-of-sequence training choices explain poor scratchpad length generalization.LaMDA uses T5 position biases, motivating padding experiments that align relevant positions.
  • 4 Scratchpad Finetuning Still Fails at Length Generalization: When input length is OOD, even in-distribution scratchpad steps become inaccurate.This indicates that the model has not learned an attention pattern that generalizes to longer bit-strings.
  • 4 Scratchpad Finetuning Still Fails at Length Generalization: Few-shot finetuning on the coin-flip parity task yields roughly constant per-step error rates across scratchpad steps.This contrasts with zero-shot scratchpad finetuning, whose per-step errors increase abruptly at OOD lengths.

5 Scratchpad Prompting Significantly Improves Length Generalization

Scratchpad prompting lets pretrained language models extend solution templates from short exemplars to much longer problems, while finetuning adds task-dependent effects. The strongest gains occur when the base model already supports the task, especially parity.

  • 5 Scratchpad Prompting Significantly Improves Length Generalization: Scratchpad prompting combines in-context examples with intermediate solution steps to improve length extrapolation without finetuning.The study examines finetuning, few-shot prompting, and scratchpad strategies together.
  • 5.1 Few-shot scratchpad: Few-shot scratchpad prompting sometimes extrapolates to lengths far beyond those supported by the pretraining weights.This contrasts with vanilla and scratchpad finetuning, which show weaker length generalization.
  • 5.1 Few-shot scratchpad: A length-3 exemplar induced the pretrained LaMDA 128b model to generate a correct 20-step scratchpad.The scratchpad reduces the task to inferring state transitions within an established solution template.
  • 5.1 Few-shot scratchpad: On parity, few-shot finetuning with scratchpad significantly outperforms zero-shot finetuning with scratchpad.The improvement appears both because the non-finetuned model already performs well and because few-shot finetuning improves over the baseline.
  • 5.1 Few-shot scratchpad: For parity, vanilla finetuning improves over no finetuning in-distribution but performs worse on larger lengths, unlike few-shot finetuning.This larger-length degradation is described as a pathology absent from few-shot finetuning.
  • 5.1 Few-shot scratchpad: On variable assignment, few-shot and vanilla finetuning show similar length-generalization behavior.The authors hypothesize that the base model’s weaker pretrained performance forces finetuning to acquire a new skill, retaining finetuning pathologies.

6 Related Works

Related work studies length extrapolation, easy-to-hard generalization, and inductive biases in sequence models. These efforts span architectural modifications, weight tying, recurrent counting, and representation-learning analyses.

  • Challenges in length generalization: Prior work examines pathologies in processing and generating sequences longer than those seen during training.Reported differences include effects of end-of-sequence tokens and modifications to dot-product attention.
  • Easy-to-Hard generalization: Weight-tied neural networks have been used to generalize from easy examples to harder ones.Several studies also propose tasks or architectures for benchmarking easy-to-hard generalization.
  • Inductive Biases Related to Lenght Generalization: Studies of sequence-to-sequence inductive biases find that activation functions can produce different strategies on language tasks.Other work reports dynamical counting and hierarchical representations in recurrent networks for Dyck tasks.

7 Conclusion

The paper concludes that length generalization remains difficult for language models, whereas few-shot scratchpad prompting can substantially improve extrapolation. Its experiments define the problem through synthetic parity and variable-assignment tasks and identify mixed effects from subsequent finetuning.

  • 7 Conclusion: Length generalization is a key capability for reasoning tasks because shorter instances are often easier and more available than longer ones.The paper frames this as learning from shorter or easier problems to handle longer or harder examples.
  • 7 Conclusion: Vanilla finetuning and scratchpad finetuning lack robust length generalization through non-sequential patterns that fail on longer instances.This limitation persists at least up to the 100B-parameter scale.
  • 7 Conclusion: Few-shot scratchpad prompting enables pretrained models to extrapolate scratchpad templates to arbitrary lengths and improves performance on longer instances.Unlike raw finetuning, this approach scales with model size.
  • 7 Conclusion: Finetuning after few-shot scratchpad prompting has mixed effects that depend on the base model’s non-finetuned performance.The conclusion contrasts parity and variable-assignment behavior rather than reporting a uniform benefit.
  • 7 Conclusion: The parity experiments use synthetic bitstrings with training lengths of 3–20 and validation/testing lengths of 3–40.The datasets contain 1000000 samples, with additional splits varying the number of ones.
  • 7 Conclusion: Natural-language parity represents flips with “Then <NAME> flips.” and non-flips with “Then <NAME> doesn’t flip.”Each step also receives an integer identifier counting backward from the total number of steps.

A.2 Boolean Variable Assignment Dataset:

The Boolean Variable Assignment dataset generates semantically correct Python programs with varied operations, variables, and execution structures. Its scratchpad copies each program and annotates updated variable values while remaining valid Python.

  • A.2 Boolean Variable Assignment Dataset:: The dataset synthesizes qualitatively different programs by sampling boolean assignments, variables, and operations while preserving semantic correctness.Operations and variables are sampled from preset hyperparameters, then added sequentially under correctness constraints.
  • A.2 Boolean Variable Assignment Dataset:: The chain-like split permits 3–19 operations and 2–3 variables.Its operator set includes assignments using and, or, xor, and negation.
  • A.2 Boolean Variable Assignment Dataset:: The diverse split permits 8–32 operations and 4–10 variables.It includes conditional assignments and assignments involving boolean constants or other variables.
  • A.2 Boolean Variable Assignment Dataset:: The scratchpad copies the input program and adds comments after each line specifying the recently updated variable’s value.This format is itself a valid Python program and can leverage models pretrained on Python data.
  • A.2 Boolean Variable Assignment Dataset:: Both Boolean Variable Assignment splits contain 1500000 samples.The figure illustrates the comment-based scratchpad representation for a generated program.

B Baseline for Vanilla Finetuning on Variable Assignment

Vanilla finetuning fails to generalize from short to longer variable-assignment programs, even when operation order is shuffled. The evidence attributes this gap to transformers favoring parallel strategies over sequential algorithms.

  • Shuffled-operations baselines perform on par with clean-data models on out-of-distribution lengths.The shuffled setup removes the possibility of solving via a sequential algorithm while retaining some spurious answer-correlated features.
  • Transformers tend to prefer parallel strategies that do not generalize to larger lengths.This preference is identified as the primary source of the observed lack of length generalization.

C Experimental Conditions

The experiments use carefully selected optimization and decoding conditions across model sizes and datasets. Larger stable learning rates were preferred because they produced better observed length generalization.

  • The experiments sweep learning rates across model sizes and select the largest rates that maintain training stability.
  • Larger learning rates produced the best observed length generalization in the initial experiments.
  • AdaFactor is used for all finetuning experiments, with no dropout for parity and dropout rate 0.05 for variable assignment.
  • All experiments use greedy decoding, while temperature sampling with likelihood-based reranking produced no qualitatively different results.

D Computational Graph Depth is the Relevant Notion of Difficulty on Variable Assignment

Computational graph depth, rather than the total number of program operations, better captures variable-assignment difficulty for transformers. Models first learn shallow dependency structures before progressing to deeper ones.

  • Computational graph depth is a more relevant difficulty measure than the number of program operations on out-of-distribution lengths.Accuracy remains roughly unchanged across graph depths when evaluated at OOD program lengths, indicating that operation count is less informative.
  • Transformers initially learn to handle programs with small computational graph depth, then move to more difficult programs.
  • Few-shot finetuning with an alternative prompt style still shows significant length-generalization pathologies when non-finetuned performance is poor.This sanity check supports the hypothesis that useful base-model performance matters for few-shot finetuning.

F Distractor Analysis for Scratchpad Strategies

Scratchpad length-generalization failures persist beyond positional-encoding and EOS explanations because distractor tokens disrupt the sequential strategy. Removing distractors restores perfect generalization, with input distractors especially harmful.

  • Padded scratchpads do not eliminate length-generalization pathologies, ruling out untrained position encodings and EOS prediction as the sole explanation.
  • The experiment compares four masking conditions across in-distribution and out-of-distribution lengths while measuring scratchpad-token prediction accuracy.Conditions include masking scratchpad distractors, input distractors, both, or neither.
  • Removing all distractor tokens results in perfect length generalization.
  • Distracting input tokens hurt length-generalization performance more than distracting preceding scratchpad tokens.
  • The findings suggest that transformer architectures or training methods that alleviate distractor effects could significantly improve length generalization.
Loading 2207.04901v2…