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

arXiv: https://arxiv.org/abs/2212.14852 · HTML: https://mlsanity.com/paper/2212.14852/

## Key facts

- Cited by 16 papers in the map (data through 14 Oct 2026).
- Builds on 60 of its 92 references that are in the map.
- First posted on arXiv 30 Dec 2022.

## Builds on (most cited first, 30 of 60)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Relational inductive biases, deep learning, and graph networks](https://mlsanity.com/paper/1806.01261/) (2018)
- [Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges](https://mlsanity.com/paper/2104.13478/) (2021)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [A Convergence Theory for Deep Learning via Over-Parameterization](https://mlsanity.com/paper/1811.03962/) (2018)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [Spectrally-normalized margin bounds for neural networks](https://mlsanity.com/paper/1706.08498/) (2017)
- [On Lazy Training in Differentiable Programming](https://mlsanity.com/paper/1812.07956/) (2018)
- [Gradient Descent Provably Optimizes Over-parameterized Neural Networks](https://mlsanity.com/paper/1810.02054/) (2018)
- [Gradient Descent Finds Global Minima of Deep Neural Networks](https://mlsanity.com/paper/1811.03804/) (2018)
- [Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networks](https://mlsanity.com/paper/1901.08584/) (2019)
- [Fantastic Generalization Measures and Where to Find Them](https://mlsanity.com/paper/1912.02178/) (2019)
- [A Mean Field View of the Landscape of Two-Layers Neural Networks](https://mlsanity.com/paper/1804.06561/) (2018)
- [Kernel Mean Embedding of Distributions: A Review and Beyond](https://mlsanity.com/paper/1605.09522/) (2016)
- [Learning and Generalization in Overparameterized Neural Networks, Going Beyond Two Layers](https://mlsanity.com/paper/1811.04918/) (2018)
- [Learning Overparameterized Neural Networks via Stochastic Gradient Descent on Structured Data](https://mlsanity.com/paper/1808.01204/) (2018)
- [Generalization Bounds of Stochastic Gradient Descent for Wide and Deep Neural Networks](https://mlsanity.com/paper/1905.13210/) (2019)
- [Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernel](https://mlsanity.com/paper/1908.11775/) (2019)
- [Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit](https://mlsanity.com/paper/1902.06015/) (2019)
- [Universal Invariant and Equivariant Graph Neural Networks](https://mlsanity.com/paper/1905.04943/) (2019)
- [E(n) Equivariant Normalizing Flows](https://mlsanity.com/paper/2105.09016/) (2021)

## Built on by (most cited first, 16 of 16)

- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding](https://mlsanity.com/paper/2303.04245/) (2023)
- [Test-time regression: a unifying framework for designing sequence models with associative memory](https://mlsanity.com/paper/2501.12352/) (2025)
- [Nested Learning: The Illusion of Deep Learning Architectures](https://mlsanity.com/paper/2512.24695/) (2025)
- [ATLAS: Learning to Optimally Memorize the Context at Test Time](https://mlsanity.com/paper/2505.23735/) (2025)
- [Reason for Future, Act for Now: A Principled Framework for Autonomous LLM Agents with Provable Sample Efficiency](https://mlsanity.com/paper/2309.17382/) (2023)
- [Approximation and Estimation Ability of Transformers for Sequence-to-Sequence Functions with Infinite Dimensional Input](https://mlsanity.com/paper/2305.18699/) (2023)
- [Understanding Scaling Laws with Statistical and Approximation Theory for Transformer Neural Networks on Intrinsically Low-dimensional Data](https://mlsanity.com/paper/2411.06646/) (2024)
- [Sequence Length Independent Norm-Based Generalization Bounds for Transformers](https://mlsanity.com/paper/2310.13088/) (2023)
- [A Mechanism for Sample-Efficient In-Context Learning for Sparse Retrieval Tasks](https://mlsanity.com/paper/2305.17040/) (2023)
- [Federated In-Context Learning: Iterative Refinement for Improved Answer Quality](https://mlsanity.com/paper/2506.07440/) (2025)
- [Learning Linear Attention in Polynomial Time](https://mlsanity.com/paper/2410.10101/) (2024)
- [DF2: Distribution-Free Decision-Focused Learning](https://mlsanity.com/paper/2308.05889/) (2023)
- [Tight Sample Complexity of Transformers](https://mlsanity.com/paper/2606.09731/) (2026)
- [Deep spatial context: when attention-based models meet spatial regression](https://mlsanity.com/paper/2401.10044/) (2024)
- [Spectrum-Adaptive Generalization Bounds for Trained Deep Transformers](https://mlsanity.com/paper/2605.07297/) (2026)

## Abstract (arXiv)

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.