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

Understand

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.