Bayes optimal learning of attention-indexed models

2025-06-02

AI Summary

Transformer attention lacks a tractable learning-theoretic foundation. Method: We propose the Attention Index Model (AIM), a theoretically rigorous framework characterizing token-level outputs as hierarchical bilinear interactions within high-dimensional embeddings; it is the first analytically solvable model supporting full-rank key/query matrices, aligning closely with practical architectures. Our analysis integrates statistical mechanics, random matrix theory, and approximate message passing to perform the first complete learning-theoretic characterization of full-rank attention—deriving a closed-form expression for the Bayes-optimal generalization error and revealing a sharp phase transition governed jointly by sample size, model width, and sequence length. Contribution/Results: Theory–experiment agreement is strong, confirming that gradient descent achieves Bayes-optimal performance. AIM establishes the first solvable, analytically tractable, and architecturally realistic learning theory for attention mechanisms.

Technology Category

• Machine Learning: Learning Theory
• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP

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
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

We introduce the attention-indexed model (AIM), a theoretical framework for analyzing learning in deep attention layers. Inspired by multi-index models, AIM captures how token-level outputs emerge from layered bilinear interactions over high-dimensional embeddings. Unlike prior tractable attention models, AIM allows full-width key and query matrices, aligning more closely with practical transformers. Using tools from statistical mechanics and random matrix theory, we derive closed-form predictions for Bayes-optimal generalization error and identify sharp phase transitions as a function of sample complexity, model width, and sequence length. We propose a matching approximate message passing algorithm and show that gradient descent can reach optimal performance. AIM offers a solvable playground for understanding learning in modern attention architectures.

Problem

• Analyzing learning in deep attention layers theoretically
• Deriving Bayes-optimal generalization error predictions
• Understanding learning in modern attention architectures

Innovation

• Introduces attention-indexed model (AIM) framework
• Derives Bayes-optimal generalization error predictions
• Proposes matching approximate message passing algorithm