Explicitly Encoding Structural Symmetry is Key to Length Generalization in Arithmetic Tasks

TL;DR

This paper demonstrates that explicitly encoding structural symmetries in number formatting and positional encodings enables Transformers to generalize length in arithmetic tasks, outperforming traditional methods and highlighting the importance of structural awareness.

Contribution

The authors introduce a method to explicitly encode structural symmetries into Transformers, significantly improving length generalization in arithmetic tasks without extra data.

Findings

Transformers with explicit symmetry encoding generalize to 50-digit addition and multiplication.

Traditional absolute positional encodings fail to generalize to longer sequences.

Explicit structure encoding is necessary for out-of-distribution generalization.

Abstract

Despite the success of Transformers on language understanding, code generation, and logical reasoning, they still fail to generalize over length on basic arithmetic tasks such as addition and multiplication. A major reason behind this failure is the vast difference in structure between numbers and text; For example, the numbers are typically parsed from right to left, and there is a correspondence between digits at the same position across different numbers. In contrast, for text, such symmetries are quite unnatural. In this work, we propose to encode these semantics explicitly into the model via modified number formatting and custom positional encodings. Empirically, our method allows a Transformer trained on numbers with at most 5-digits for addition and multiplication to generalize up to 50-digit numbers, without using additional data for longer sequences. We further demonstrate that traditional absolute positional encodings (APE) fail to generalize to longer sequences, even when trained with augmented data that captures task symmetries. To elucidate the importance of explicitly encoding structure, we prove that explicit incorporation of structure via positional encodings is necessary for out-of-distribution generalization. Finally, we pinpoint other challenges inherent to length generalization beyond capturing symmetries, in particular complexity of the underlying task, and propose changes in the training distribution to address them.