arXiv · 2405.13718
Next-token prediction capacity: general upper bounds and a lower bound for transformers
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.
Explore related subjects
Keep this discovery
Explore connections, maps & timelines
Liam Madden, Curtis Fox, Christos Thrampoulidis. 2024-05-22. Next-token prediction capacity: general upper bounds and a lower bound for transformers. https://doi.org/10.1109/tit.2025.3584013
Cite the original work for its findings. Save a collection to share your selection of sources.