Attention layers provably solve single-location regression

2024-10-02

arXiv.org

AI Summary

This work addresses the lack of theoretical foundations for attention mechanisms—particularly Transformers—by investigating token sparsity and the learnability of internal linear representations. We introduce the “single-position regression” task, where only one token in the input sequence determines the output, and its position is implicitly encoded via a linear projection of the input. We establish, for the first time, that under a suitably simplified self-attention architecture, nonlinear attention layers can asymptotically achieve Bayes-optimal performance. Moreover, we demonstrate that such architectures reliably recover the sparse positional structure even under nonconvex training dynamics. Through asymptotic statistical analysis and nonconvex optimization modeling, we theoretically prove that attention mechanisms can consistently extract token-level sparse signals and underlying linear structures, thereby attaining optimal generalization on this task.

Technology Category

• Natural Language Processing: Learning & Optimization for NLP
• Machine Learning: Other Foundations of Machine Learning
• Search and Optimization: Non-convex Optimization

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

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.

Problem

• Understand attention-based models theoretically
• Analyze token-wise sparsity in Transformers
• Explore linear representations in attention layers

Innovation

• Attention layers solve regression
• Simplified non-linear self-attention
• Handles sparse token information