Optimization and generalization in the Θ(nω) parameter regime

Investigate the optimization and generalization behavior of one-layer multi-head decoder-only transformers in the Θ(nω) parameter regime where the number of model parameters scales on the order of nω, with n the number of distinct contexts and ω the vocabulary size.

Background

The paper establishes matching upper and lower bounds (up to constants) on next-token prediction capacity, showing that Θ(nω) parameters are necessary and sufficient for memorization. It also provides numerical evidence that training can reach the entropy lower bound near this regime.

However, the theoretical understanding of optimization dynamics and generalization properties in precisely this Θ(nω) parameter regime is not developed in the paper, and the authors explicitly point to it as future work.