What Algorithms can Transformers Learn? A Study in Length Generalization

Hattie Zhou; Arwen Bradley; Etai Littwin; Noam Razin; Omid Saremi; Josh Susskind; Samy Bengio; Preetum Nakkiran

Source-linked AI summary

What Algorithms can Transformers Learn? A Study in Length Generalization

Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Josh Susskind, Samy Bengio, Preetum Nakkiran

arXiv:2310.16028v1Oct 24, 2023cs.LGcs.AIcs.CLstat.ML

TL;DR

Transformers can perform length generalization on algorithmic tasks, but it remains unclear when they learn the underlying algorithm rather than merely fit the training distribution. This paper uses RASP-L to formulate and test a conjecture based on Transformer representability, finding that short all-length programs predict generalization while also improving hard tasks such as parity and addition. The authors frame the result as a phenomenological predictive tool rather than a mechanistic account, with limitations from RASP-L's incompleteness and the absence of perfect arbitrary-length generalization.

Problem
The paper asks when standard decoder-only Transformers learn algorithms that generalize systematically beyond their training lengths, given prior failures on arithmetic and other algorithmic tasks.
Method
The authors operationalize Transformer-algorithm simplicity with RASP-L and test a three-condition RASP-Generalization Conjecture through algorithmic-task experiments and redesigned input or scratchpad formats.
Results
Short RASP-L programs capture most known instances of length generalization, and the approach yields strong length generalization on parity and addition for Transformers trained from scratch.
Takeaways & Limitations
Transformer length generalization is associated with algorithms that are simple to represent in the architecture, offering a unifying perspective on algorithm learning and compositional generalization.
Takeaways & Limitations
RASP-L is incomplete, does not efficiently represent some numerical Transformer algorithms, and the conjecture predicts behavior phenomenologically rather than explaining learned internal weights.

Abstract

from arXiv · show

Large language models exhibit surprising emergent generalization properties, yet also struggle on many simple reasoning tasks such as arithmetic and parity. This raises the question of if and when Transformer models can learn the true algorithm for solving a task. We study the scope of Transformers' abilities in the specific setting of length generalization on algorithmic tasks. Here, we propose a unifying framework to understand when and how Transformers can exhibit strong length generalization on a given task. Specifically, we leverage RASP (Weiss et al., 2021) -- a programming language designed for the computational model of a Transformer -- and introduce the RASP-Generalization Conjecture: Transformers tend to length generalize on a task if the task can be solved by a short RASP program which works for all input lengths. This simple conjecture remarkably captures most known instances of length generalization on algorithmic tasks. Moreover, we leverage our insights to drastically improve generalization performance on traditionally hard tasks (such as parity and addition). On the theoretical side, we give a simple example where the "min-degree-interpolator" model of learning from Abbe et al. (2023) does not correctly predict Transformers' out-of-distribution behavior, but our conjecture does. Overall, our work provides a novel perspective on the mechanisms of compositional generalization and the algorithmic capabilities of Transformers.

1 Introduction

The paper asks when Transformers learn algorithms that generalize to longer inputs, and argues that representability by short RASP-L programs predicts this behavior. Experiments show strong generalization on suitable tasks and demonstrate improvements through task and scratchpad design.

Motivation: Length generalization tests whether models trained on shorter problems learn algorithms that transfer to longer, harder inputs.The paper uses this setting as an indicator of whether the correct underlying algorithm was learned.
Evidence and implications: Figure 1 reports that tasks admitting short RASP-L solutions exhibit strong length generalization, while modified input and scratchpad formats can make traditionally hard tasks generalize.The paper also uses the counting task to illustrate near-perfect generalization for models trained on sequences of length 60 or more and evaluated up to length 150.
Motivation: Prior studies found little or no length generalization on addition and related arithmetic tasks, leaving the conditions for systematic generalization unclear.The paper contrasts these findings with evidence that some algorithmic tasks generalize far beyond the training distribution.
Evidence and implications: A counting algorithm works for sequences of all lengths, and a Transformer can represent it uniformly with the same weights across different input lengths.The paper reports supporting evidence that trained models implement this algorithm.
Framework: RASP-L provides a human-readable language whose programs can be compiled into Transformer weights, linking algorithmic descriptions to Transformer-representable computation.Each RASP-L line compiles into at most one Transformer layer.
Framework: The RASP-Generalization view predicts that Transformers tend to learn the shortest RASP-L program fitting the training data, provided it generalizes correctly.This is presented as a toy model related to minimum-description-length principles, with complexity tailored to Transformer information flow.

2 Main Conjecture

The RASP-Generalization Conjecture identifies realizability, simplicity, and diversity as conditions under which completed decoder-only Transformers are likely to length-generalize. Experiments support these predictive relationships, while the authors qualify the conjecture as phenomenological and note important representational and optimization limits.

Conjecture: The RASP-Generalization Conjecture predicts length generalization when a task satisfies realizability, simplicity, and diversity.These conditions concern a single all-length Transformer solution, a short RASP-L representation, and sufficiently diverse training data.
Conjecture: The three conjecture features correlate empirically with longer-length generalization and faster in-distribution training optimization.Simplicity implies realizability because an RASP-L solution guarantees a Transformer representation.
Scope and assumptions: Realizability is stringent because it requires one Transformer to solve the task at all lengths, although unbounded-length representability is used as a heuristic for bounded contexts.Tasks requiring a different Transformer at each length may be unnatural for the architecture and unlikely to generalize well.
Scope and assumptions: The conjecture is primarily phenomenological: it predicts external out-of-distribution behavior without claiming that learned weights resemble compiled RASP-L weights.The mechanism suggested by the toy model remains an open question.
Experimental validation: The experiments vary task-level difficulty and training-length range while holding the study focused on task characteristics rather than architectural factors.Architectural innovations, positional embeddings, and training methodology may also affect generalization and how robustly solutions are learned.
Experimental validation: Perfect generalization over arbitrary lengths is not expected because optimization noise, continuous weights, finite precision, and approximate learned programs can cause degradation.The paper therefore studies non-trivial external length-generalization performance and leaves mechanistic analysis for future work.

3 RASP-L: What Algorithms Can Transformers Learn?

RASP-L is a restricted, human-readable language designed to capture Transformer-representable algorithms in a learnable, length-uniform form. Its causal, straight-line operations support parallel computation, while bounded integers and restricted index operations exclude difficult-to-represent behaviors.

RASP is a domain-specific language for specifying Transformer weights through human-readable programs.
Every RASP program maps sequences of any length to same-length outputs using straight-line programs without branching or loops.
Causal RASP interprets the final output position as the next-token prediction for autoregressive decoder-only Transformers.
RASP’s parallelizable operations make arbitrary sequential computation difficult because Transformers have constant depth.
RASP-L restricts variables to bounded integers, limits token-index operations to simple comparisons and predecessor or successor relations, and adds min and max aggregation.
The RASP-L standard library includes helpers such as induction, while excluding decimal parsing because operating on decimal representations is nontrivial.

4 Experiments: Length Generalization

Experiments compare four tasks with simple RASP-L programs against three without them. Tasks with simple programs generally length-generalize well, whereas addition, parity, and copying with repeated tokens generalize poorly under the experimental setup.

Four tasks with simple RASP-L programs—count, mode, unique-token copy, and sort—are compared with three tasks lacking such programs.
Successful Length Generalization: Count models generalize near perfectly to twice their training lengths when training includes diverse lengths from 1 through the maximum.
Successful Length Generalization: Mode models generalize strongly beyond training lengths, with models trained up to length 10 reaching median accuracy of 50% on length-60 sequences.
Successful Length Generalization: Unique-token copy generalizes perfectly to length 50 after training on sequences up to length 40, aided by induction heads.
Successful Length Generalization: Sort models show strong length generalization to length 50 when trained on sequences of length 35 or more.
Unsuccessful Length Generalization: Addition and parity show no length generalization under these settings: performance remains at random chance when test examples are at least five tokens longer than training examples.
Unsuccessful Length Generalization: Repeating-token copy fails to generalize because induction heads no longer help and precise index arithmetic is required.

5 Application: Improving Length Generalization

The paper improves length generalization by reformulating tasks so shorter RASP-L programs can solve them, or by diversifying training data when complexity remains. Index hints and reverse-order outputs help addition, while scratchpads help parity but can hurt mode.

Addition: Index hints simplify addition by addressing the indexing operations required to select summand digits for each output digit.The reformulation adds matching index symbols to the prompt and answer.
Addition: Reverse-order addition makes carry propagation causal through the previous output digit, whereas forward-order addition requires a more complex calculation over remaining inputs.The forward algorithm consequently has a much longer RASP-L program than the reverse algorithm.
Addition: On length-50 hard-carry examples, forward addition shows no length generalization, while most reverse-addition runs generalize perfectly when training length exceeds 40.The comparison uses index hints in both addition settings.
Addition: Balanced carry sampling increases training-data diversity so shortcut programs cannot fit forward addition as easily, enabling learning of its more complex RASP-L algorithm.Forward addition admits a RASP-L program, but the paper motivates more careful training distributions because that program is complex.
Scratchpads: Scratchpads can simplify next-token prediction by unrolling sequential loops into locally solvable steps, making tasks more amenable to short RASP-L programs.The paper constructs good and bad scratchpads to test whether this predicted complexity difference changes generalization.
Scratchpads: For parity, index hints and a running-state scratchpad yield perfect length generalization on length-50 inputs for all models once training length reaches 40.Some runs already generalize perfectly from training lengths up to 30, and training speed varies with scratchpad difficulty under RASP-L.
Scratchpads: For mode, a frequency-sorting scratchpad performs significantly worse than no scratchpad on exact match and final-answer accuracy, and also slows training.The scratchpad adds explicit frequency sorting and decimal conversion, increasing the task's RASP-L difficulty; a nonsorting variant shows a similar effect.

6 Comparison to Min-Degree-Interpolators

The paper compares Transformer-specific RASP-L complexity with minimum-degree interpolation and finds that RASP-L better predicts the observed unseen-distribution behavior in a Boolean-AND task.

Comparison: The comparison illustrates that Transformer-specific minimum RASP-L program length can be more predictive than architecture-agnostic complexity measures in certain settings.The authors explicitly limit this conclusion to the studied settings and do not claim universal applicability.
Experimental comparison: 100% test accuracy was achieved on unseen Boolean-AND inputs, including inputs where the previously irrelevant final five bits could contain 0s.The Transformer correctly computed conjunction beyond the training distribution.
Minimum-degree interpolation: The minimum-degree interpolator ignores input bits that are constant on the training set, so it does not depend on the final five bits in this setting.This follows because dependence on constant training-set bits can be removed without changing training outputs.
RASP-Length: A one-line RASP-L program computes Boolean AND by searching the prior context for a 0, and this program is representable by a one-layer Transformer.The RASP-L construction predicts the experimental behavior because directly detecting any 0 is simpler than learning to ignore constant bits.

7 Discussion, Limitations, and Open Questions

The discussion identifies limitations in RASP-L's expressiveness and in using program length as a complexity measure, while noting that the conjecture remains informal and several theoretical questions are open.

Limitations of RASP-L: RASP and RASP-L are incomplete and poorly suited to some efficient Transformer algorithms, including numerical procedures involving high-dimensional floating-point operations.They also support only deterministic outputs and binary attention masks, and RASP-L imposes restrictions such as no floating-point operations.
Limitations of Complexity Measures: RASP-L program length does not perfectly match Transformer complexity because multiple program lines can sometimes fit within one Transformer layer.The paper suggests parallel depth and robust weight norms as potentially refined alternatives, although such measures may be intractable.
Limitations of Scope and Strength: The main conjecture is not fully formal: it cannot precisely predict generalization extent or determine the data diversity needed for strong length generalization.The authors present formalization and more precise conjectures as open directions.
Open questions: Whether tasks simple for autoregressive Transformers admit a natural complexity-theoretic characterization remains open.The paper contrasts this open question with the known connection between recurrent models and regular-language tasks such as parity.
Relations to Mechanistic Interpretability: The work is related to mechanistic interpretability but operates at a different abstraction level, using a toy model to provide intuition for the RASP conjecture.Mechanistic interpretability instead studies internal algorithms or circuits by inspecting model internals.

8 Conclusion

The paper studies algorithm learning through length generalization and proposes short RASP-L programs as a Transformer-specific lens for predicting when it occurs. It concludes that this framework captures many observed cases while leaving formal characterization and broader applicability open.

8 Conclusion: The paper frames length generalization as a controlled test of whether standard decoder-only Transformers can learn general problem-solving strategies from examples.The experiments use Transformers trained from scratch on synthetic algorithmic tasks.
8 Conclusion: The central perspective treats Transformers as programs defined for inputs of all lengths rather than functions with a fixed input size.This motivates analyzing algorithms by how simply they can be represented.
8 Conclusion: The RASP-Generalization Conjecture predicts that tasks expressible by short RASP-L programs are more likely to length-generalize, and it captures most known instances in the paper's scope.The authors suggest the framework may extend to compositional generalization beyond length generalization.
8 Conclusion: The conclusion connects the framework to prior work on arithmetic reasoning, scratchpads, formal languages, positional design, and other approaches to Transformer generalization.These connections position the paper within several lines of research on algorithmic reasoning and compositional generalization.
8 Conclusion: The discussion of “Faith and Fate” distinguishes settings where Transformers compositionally generalize from settings where they may rely on subgraph matching, while noting that error propagation can theoretically be reduced by scaling.The paper reports strong length generalization in its relatively small experiments.
8 Conclusion: RASP programs can also compile to linear-attention Transformers, suggesting computational restrictions stronger than those of generic Transformers.The paper notes that linear-attention Transformers can evaluate their next-token function in O(n) time in their input context.

B Additional Experimental Details

The appendix specifies tokenization, sampling, training, evaluation, and optimizer settings for the algorithmic-task experiments, including distinct procedures for each task.

General setup: The experiments tokenize characters individually except in Count, where each integer is one token, and use online iid training batches with 20 random seeds reported by median.Test performance generally uses 5× the batch-size number of samples.
Count: Count uses an alphabet of 155 and evaluates sequences up to length 150, with training lengths sampled uniformly from 1 to the maximum.Evaluation enumerates all possible sequences for each test length.
Mode: Mode uses 52 tokens and examples containing five unique sampled tokens, with ties resolved to ensure a unique answer.Example lengths are sampled uniformly up to the maximum training length.
Copy: Copy uses alphabet size 100 for unique-token sequences and alphabet size 2 when repeated tokens are allowed.Lengths are sampled uniformly from 1 to the maximum training length.
Sort: Sort trains on sequences sampled without replacement from an alphabet of 100 tokens, with lengths sampled uniformly up to the maximum.
Addition: Addition samples number lengths uniformly and distinguishes standard-carry sampling from balanced-carry sampling based on the carry-chain structure.The balanced-carry setting samples carry-chain length and position explicitly.
Parity: Parity samples binary sequences with lengths from 1 to the maximum and defines length by sequence length rather than by the number of 1s.
Boolean-AND: Boolean-AND trains and tests at length 20, with training inputs consisting equally of all-ones sequences and sequences containing one 0 among positions 0 through 15.The final five positions are constant during training, matching the comparison setting.

C Appendix: Counterfactual Analysis on Count

The appendix tests whether count models learned the intended algorithm by evaluating their behavior on counterfactual random sequences, including far-out-of-distribution lengths. Models generally follow start-token and increment rules, but termination behavior is less robust, especially for models trained on shorter sequences.

Counterfactual tests: The intended count algorithm searches for the prompt’s start marker, reads boundary values, and either starts, terminates, or increments the sequence.The next-token procedure outputs the start value after the start marker, emits EoS at the ending value, and otherwise increments the previous token.
Counterfactual tests: Counterfactual tests use random sequences to probe whether models apply the expected algorithm independently of training-distribution statistics.The tests cover start-token output, termination at the ending value, increment-by-1 behavior, and premature termination.
Results: All models perform near perfectly on starting with the prompt’s start token and incrementing by 1, both in-distribution and out-of-distribution.These tests evaluate behavior on random sequences rather than only examples resembling training data.
Results: Models trained on shorter lengths are less robust at emitting EoS when the ending condition is met on long test sequences.Models trained on longer sequences show more robust end-on-last behavior and generally have better length generalization.

D Appendix: Additional Ablations

This appendix section reports additional experiments supporting the paper’s main results.

Additional Ablations: The appendix includes additional experiments to support the results reported in the main paper.
Additional Ablations: These experiments are presented as supplementary evidence rather than as a separate primary analysis.
Additional Ablations: The section provides no further experimental details in the supplied passage.

D.1 Training Speed

The appendix examines whether scratchpad simplicity relates to training speed and length generalization, using parity, count, sort, and mode-related experiments. Results indicate that easier scratchpads can improve parity learning, while scratchpad design and training-length coverage affect generalization across tasks.

Training Speed: The RASP Generalization Conjecture motivates testing whether simpler programs are learned faster and generalize better.The training-speed analysis compares scratchpad variants according to the simplicity of their corresponding RASP-L programs.
Count: Models trained on count can increment correctly on counterfactual sequences but sometimes terminate prematurely, particularly beyond training lengths.The premature EoS behavior is identified as a likely explanation for observed count length-generalization failures.
Parity: The alternative parity scratchpad computes a sum-mod-10 intermediate sequence but remains more complex because it does not process the task autoregressively.It is nevertheless described as simpler than the harder parity formulation.
Sort: Sort models trained with lengths of 35 or more show strong generalization to length 50, with broader training-length coverage improving further extrapolation.
Parity: Easy Scratchpad generalizes significantly better than Hard Scratchpad on parity, while both optimize better than parity without a scratchpad.Parity without a scratchpad cannot fit the training set and shows no length generalization.
Mode: An appearance-ordered mode scratchpad performs much worse for length generalization than using no scratchpad.The result contrasts with the paper’s other scratchpad experiments and indicates that scratchpad design matters.

D.3 Performance with rotary embedding

The rotary-positional-embedding experiments compare length generalization across count, mode, addition, and parity, including scratchpad and index-hint variants. Across these tasks, rotary positional embeddings yield worse length generalization, while balanced-carry training supports strong addition extrapolation.

Rotary positional embedding: Models with rotary positional embeddings show worse length generalization across count, mode, addition with index hints, and parity.This observation is reported as consistent with prior findings on learned positional embeddings.
Task coverage: The rotary-embedding evaluations cover count, mode, addition with index hints, and parity.The corresponding results are shown in Figures 11–13, while additional addition experiments use index hints and carry distributions.
Task coverage: The rotary-embedding appendix also compares scratchpad variants for addition and parity, and count and mode task performance.
Addition with index hints: Figure 14 compares forward and reverse addition with index hints on easy-carry examples of length 50 across 20 training runs.The diamonds show the spread across runs; the caption identifies separate panels for forward and reverse addition.
Addition with index hints: With balanced carry training, forward and reverse addition with index hints show strong length generalization on hard-carry examples of length 50.Each point in the experiment represents the median test performance over 20 runs.

E RASP Specification

RASP-L restricts Transformer-like sequence programs to fixed-length, straight-line operations over finite-valued sequences and constrained selectors, while providing library primitives for attention-style selection, aggregation, indexing, and sequence manipulation. The specification includes task programs such as addition, whose next-token function is evaluated autoregressively in several output formats.

RASP-L core: RASP-L maps arbitrary-length input sequences to same-length outputs using sequences and binary selectors, without branching or loops.Its programs are straight-line implementations intended to match the Transformer computational model.
RASP-L core: RASP-L restricts non-index values to int8, limiting tokenwise functions to small finite domains and excluding floating-point or unbounded-number operations.The restriction is presented as addressing infinite precision and learnability concerns.
Index restrictions: Index values support only order comparisons and successor or predecessor operations; adding two indices or casting indices to int8 is disallowed.These restrictions are enforced through the IndexInt type.
Index restrictions: Selector width returns counts of selected prior elements, but its outputs inherit restricted IndexInt behavior and distinct calls produce incomparable types.These constraints prevent selector width from bypassing the intended restrictions on index arithmetic.
Index restrictions: The selector-width restrictions are motivated by its ability to simulate indices using an all-ones selector, while their formal adequacy remains an open question.The paper also notes that selector-width implementations may require weights growing linearly with sequence length.
Library functions: The RASP-L library supplies primitives for sequence mapping, causal selection, selector-width counting, aggregation, masking, shifting, occurrence lookup, indexing, and induction.These functions are used to construct programs for counting, mode, copying, and addition.
Addition programs: The addition program computes a next-token function and must be called autoregressively to generate the complete output sequence.The specification covers reverse-order addition, standard-order addition, and index-hinted inputs.