Principled Understanding of Generalization for Generative Transformer Models in Arithmetic Reasoning Tasks

TL;DR

This paper presents a theoretical framework to understand how transformer models generalize in arithmetic reasoning tasks, emphasizing the role of task structure and positional encoding in length generalization.

Contribution

It introduces a unified theory linking positional encoding and task structure to transformer generalization, validated through experiments on GPT models.

Findings

Translation invariance in addition aids generalization

Base mismatch in modular operations causes generalization failure

Framework accurately predicts transformer behavior in arithmetic tasks

Abstract

Transformer-based models excel in various tasks but their generalization capabilities, especially in arithmetic reasoning, remain incompletely understood. Arithmetic tasks provide a controlled framework to explore these capabilities, yet performance anomalies persist, such as inconsistent effectiveness in multiplication and erratic generalization in modular addition (e.g., modulo 100 vs. 101). This paper develops a unified theoretical framework for understanding the generalization behaviors of transformers in arithmetic tasks, focusing on length generalization. Through detailed analysis of addition, multiplication, and modular operations, we reveal that translation invariance in addition aligns with relative positional encoding for robust generalization, while base mismatch in modular operations disrupts this alignment. Experiments across GPT-family models validate our framework, confirming its ability to predict generalization behaviors. Our work highlights the importance of task structure and training data distribution for achieving data-efficient and structure-aware training, providing a systematic approach to understanding of length generalization in transformers.