Next-token prediction capacity for multi-layer and non–back-loaded architectures

Analyze the next-token prediction capacity of multi-layer transformers and of modified architectures such as Mixture of Softmaxes that allocate a greater proportion of parameters to embedding and self-attention than to the feedforward sub-layer (i.e., are not back-loaded).

Background

The paper provides lower bounds for one-layer multi-head decoder-only transformers with back-loaded parameterization, where most parameters are in the feedforward sub-layer. The capacity behavior of deeper transformer stacks and architectures with different parameter distributions remains unaddressed.

The authors explicitly identify extending the capacity analysis to multi-layer transformers and to architectures like Mixture of Softmaxes—where parameters are concentrated earlier in the network—as a future direction.