Fundamental limits of learning in sequence multi-index models and deep attention networks: High-dimensional asymptotics and sharp thresholds
This work provides foundational theoretical insights into the learning dynamics of deep attention networks, which is crucial for understanding their scalability and efficiency in AI applications.
The paper tackles the problem of learning deep attention neural networks by mapping them to sequence multi-index models and deriving sharp asymptotic performance limits and sample complexity thresholds in high-dimensional settings. It characterizes the optimal Bayesian performance and the performance of approximate message-passing algorithms, revealing sequential learning across layers.
In this manuscript, we study the learning of deep attention neural networks, defined as the composition of multiple self-attention layers, with tied and low-rank weights. We first establish a mapping of such models to sequence multi-index models, a generalization of the widely studied multi-index model to sequential covariates, for which we establish a number of general results. In the context of Bayesian-optimal learning, in the limit of large dimension $D$ and commensurably large number of samples $N$, we derive a sharp asymptotic characterization of the optimal performance as well as the performance of the best-known polynomial-time algorithm for this setting --namely approximate message-passing--, and characterize sharp thresholds on the minimal sample complexity required for better-than-random prediction performance. Our analysis uncovers, in particular, how the different layers are learned sequentially. Finally, we discuss how this sequential learning can also be observed in a realistic setup.