Attention layers provably solve single-location regression

Pierre Marion (1) , Raphaël Berthier (2, 3) , Gérard Biau (4) , Claire Boyer (5)

1 Département de Mathématiques - EPFL
2 Inria de Paris
3 SU - Sorbonne Université
4 LPSM (UMR_8001) - Laboratoire de Probabilités, Statistique et Modélisation
5 LMO - Laboratoire de Mathématiques d'Orsay

Attention-based models, such as Transformer, excel across various tasks but lack a comprehensive theoretical understanding, especially regarding token-wise sparsity and internal linear representations. To address this gap, we introduce the single-location regression task, where only one token in a sequence determines the output, and its position is a latent random variable, retrievable via a linear projection of the input. To solve this task, we propose a dedicated predictor, which turns out to be a simplified version of a non-linear self-attention layer. We study its theoretical properties, by showing its asymptotic Bayes optimality and analyzing its training dynamics. In particular, despite the non-convex nature of the problem, the predictor effectively learns the underlying structure. This work highlights the capacity of attention mechanisms to handle sparse token information and internal linear structures.

Domaines: Machine Learning [stat.ML]

Pierre Marion, Raphaël Berthier, Gérard Biau, Claire Boyer. Attention layers provably solve single-location regression. 2024. ⟨hal-04720799⟩