Next-token prediction capacity: general upper bounds and a lower bound for transformers
This work addresses a foundational theoretical gap in understanding transformer limitations for language modeling, which is incremental as it builds on existing models without proposing new ones.
The paper tackles the problem of understanding the capacity of decoder-only transformers for next-token prediction by proving general upper and lower bounds on the number of distinct context sequences they can interpolate, which are equal up to a multiplicative constant, and provides numerical evidence that minimal parameters suffice for training to entropy lower bounds.
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.