A Formal Framework for Understanding Length Generalization in Transformers

2024-10-03

arXiv.org

AI Summary

Transformers exhibit limited generalization under sequence length extrapolation, and existing theoretical explanations remain inadequate. Method: We establish the first rigorous formal framework characterizing the class of functions recognizable by causal Transformers with learnable absolute positional encodings, along with precise length extrapolation bounds. Our approach integrates formal language theory, function recognizability analysis, norm-regularized idealized models, and explicit positional encoding modeling, complemented by theory-guided empirical validation. Contributions/Results: (1) We derive the first provably sound theoretical criterion for length generalization; (2) we uncover fundamental connections between positional encoding structure and learnability of target functions; and (3) we achieve *a priori* prediction—confirmed empirically—of generalization success or failure on canonical tasks including bracket matching, counting, and sequence copying. Our theory explains well-known empirical phenomena and provides provably grounded design principles for length-generalizable architectures.

Technology Category

• Natural Language Processing: (Large) Language Models
• Machine Learning: Causal Learning
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

A major challenge for transformers is generalizing to sequences longer than those observed during training. While previous works have empirically shown that transformers can either succeed or fail at length generalization depending on the task, theoretical understanding of this phenomenon remains limited. In this work, we introduce a rigorous theoretical framework to analyze length generalization in causal transformers with learnable absolute positional encodings. In particular, we characterize those functions that are identifiable in the limit from sufficiently long inputs with absolute positional encodings under an idealized inference scheme using a norm-based regularizer. This enables us to prove the possibility of length generalization for a rich family of problems. We experimentally validate the theory as a predictor of success and failure of length generalization across a range of algorithmic and formal language tasks. Our theory not only explains a broad set of empirical observations but also opens the way to provably predicting length generalization capabilities in transformers.

Problem

• Understanding length generalization in transformers theoretically
• Analyzing transformers with learnable absolute positional encodings
• Predicting length generalization capabilities in transformers

Innovation

• Theoretical framework for transformer length generalization
• Absolute positional encodings with norm-based regularizer
• Experimental validation on algorithmic tasks