Abstract:With the attention mechanism, transformers achieve significant empirical successes. Despite the intuitive understanding that transformers perform relational inference over long sequences to produce desirable representations, we lack a rigorous theory on how the attention mechanism achieves it. In particular, several intriguing questions remain open: (a) What makes a desirable representation? (b) How does the attention mechanism infer the desirable representation within the forward pass? (c) How does a pretraining procedure learn to infer the desirable representation through the backward pass? We observe that, as is the case in BERT and ViT, input tokens are often exchangeable since they already include positional encodings. The notion of exchangeability induces a latent variable model that is invariant to input sizes, which enables our theoretical analysis. - To answer (a) on representation, we establish the existence of a sufficient and minimal representation of input tokens. In particular, such a representation instantiates the posterior distribution of the latent variable given input tokens, which plays a central role in predicting output labels and solving downstream tasks. - To answer (b) on inference, we prove that attention with the desired parameter infers the latent posterior up to an approximation error, which is decreasing in input sizes. In detail, we quantify how attention approximates the conditional mean of the value given the key, which characterizes how it performs relational inference over long sequences. - To answer (c) on learning, we prove that both supervised and self-supervised objectives allow empirical risk minimization to learn the desired parameter up to a generalization error, which is independent of input sizes. Particularly, in the self-supervised setting, we identify a condition number that is pivotal to solving downstream tasks.

A Mathematical Theory of Attention

On the Regularity of Attention

Attention: Marginal Probability is All You Need?

The Lipschitz Constant of Self-Attention

The Quarks of Attention

An Analysis of Attention via the Lens of Exchangeability and Latent Variable Models

Neural Attention Models in Deep Learning: Survey and Taxonomy

No one knows what attention is

Rethinking the role of attention mechanism: a causality perspective

An Attentive Survey of Attention Models

Attention, please! A survey of neural attention models in deep learning

Understanding More about Human and Machine Attention in Deep Neural Networks

How Smooth Is Attention?

Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth

Understanding Attention: In Minds and Machines

You Need to Pay Better Attention: Rethinking the Mathematics of Attention Mechanism

On the Learning Dynamics of Attention Networks

A General Survey on Attention Mechanisms in Deep Learning

Attention mechanism in neural networks: where it comes and where it goes

Dissecting the Interplay of Attention Paths in a Statistical Mechanics Theory of Transformers