Reproduce “Bayes optimal learning of attention-indexed models” on OpenResearch

About this paper

Fabrizio Boncoraglio, Emanuele Troiani, Vittorio Erba, Lenka Zdeborová

Researchers from EPFL introduced the Attention-Indexed Model (AIM), a theoretical framework for analyzing extensive-width self-attention layers in a high-dimensional setting. The model yields closed-form predictions for Bayes-optimal generalization error, revealing distinct phase transitions and learnability behaviors depending on activation functions and architectural parameters.

Key insights

• AIM establishes a solvable theoretical framework for attention layers with full-width key and query matrices, aligning more closely with modern transformer architectures.
• Bayes-optimal learning in AIM exhibits sharp phase transitions in generalization error, quantifying critical thresholds for learnability as functions of sample complexity, sequence length, and attention width.
• The choice of activation function critically impacts learning behavior: softmax leads to a finite strong-recovery threshold where error can vanish, while hardmax results in strictly positive error for any finite sample complexity.

Summary by alphaXiv.