Next-token prediction capacity: general upper bounds and a lower bound for transformers

2024-05-22

AI Summary

This paper investigates the maximum interpolatable context length—the largest sequence length for which a decoder-only Transformer can exactly interpolate the next-token probability distribution in next-token prediction. Method: We derive tight asymptotic upper and lower bounds (within constant factors) on this capacity, covering both arbitrary true distributions and empirical distributions. Our analysis leverages the injectivity of self-attention, information-theoretic arguments, parameter-complexity analysis, and numerical experiments. Results: We prove that a single-layer multi-head Transformer achieves the lower bound; the bounds are tight up to constant factors; and the minimal memory parameter count suffices for convergence to the entropy lower bound. This work provides the first rigorous, quantitative characterization of the modeling capacity of Transformers in terms of interpolatable context length.

Technology Category

• Natural Language Processing: (Large) Language Models
• Machine Learning: Large Multimodal Models (LMMs)
• Computer Vision: Large Vision Models

Application Category

• Web Mining and Content Analysis: Large pretrained models with web data
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Large language models for search

Abstract

Given a sequence of tokens, such as words, the task of next-token prediction is to predict the next-token conditional probability distribution. Decoder-only transformers have become effective models for this task, but their properties are still not fully understood. In particular, the largest number of distinct context sequences that a decoder-only transformer can interpolate next-token distributions for has not been established. To fill this gap, we prove upper and lower bounds on this number, which are equal up to a multiplicative constant. We prove these bounds in the general setting where next-token distributions can be arbitrary as well as the empirical setting where they are calculated from a finite number of document sequences. Our lower bounds are for one-layer multi-head decoder-only transformers and our proofs highlight an important injectivity property satisfied by self-attention. Furthermore, we provide numerical evidence that the minimal number of parameters for memorization is sufficient for being able to train the model to the entropy lower bound.

Problem

• Establishes bounds on transformer's next-token prediction capacity
• Proves interpolation limits for decoder-only transformers
• Analyzes memorization requirements for optimal training

Innovation

• Proves upper and lower bounds for transformer interpolation
• Focuses on one-layer multi-head decoder-only transformers
• Demonstrates memorization via minimal parameter training