Arithmetic Transformers Can Length-Generalize in Both Operand Length and Count

TL;DR

This paper demonstrates that arithmetic transformers can generalize to longer operand lengths and counts, achieving 2-3x length generalization on addition and multiplication tasks by using task-specific scratchpads and position coupling techniques.

Contribution

It introduces a novel approach with task-specific scratchpads and position coupling to enable length generalization in arithmetic transformers, achieving first-of-its-kind results.

Findings

Achieved 2-3x length generalization on addition and multiplication tasks.

Proved a 1-layer Transformer can solve multi-operand addition with exponential operand length and count.

Introduced task-specific scratchpads and multi-level position coupling methods.

Abstract

Transformers often struggle with length generalization, meaning they fail to generalize to sequences longer than those encountered during training. While arithmetic tasks are commonly used to study length generalization, certain tasks are considered notoriously difficult, e.g., multi-operand addition (requiring generalization over both the number of operands and their lengths) and multiplication (requiring generalization over both operand lengths). In this work, we achieve approximately 2-3x length generalization on both tasks, which is the first such achievement in arithmetic Transformers. We design task-specific scratchpads enabling the model to focus on a fixed number of tokens per each next-token prediction step, and apply multi-level versions of \Position Coupling (Cho et al., 2024; McLeish et al., 2024) to let Transformers know the right position to attend to. On the theory side, we prove that a 1-layer Transformer using our method can solve multi-operand addition, up to operand length and operand count that are exponential in embedding dimension.