# Attention Mechanism, Max-Affine Partition, and Universal Approximation

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

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 28 Apr 2025.

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

- [Your Vision-Language Model Can't Even Count to 20: Exposing the Failures of VLMs in Compositional Counting](https://mlsanity.com/paper/2510.04401/) (2025)
- [Learning to Rank Chain-of-Thought: Using a Small Model](https://mlsanity.com/paper/2505.14999/) (2025)
- [A unified framework for establishing the universal approximation of transformer-type architectures](https://mlsanity.com/paper/2506.23551/) (2025)
- [Universal Approximation of Measure-to-Measure Operators by Pushforwards](https://mlsanity.com/paper/2609.35483/) (2026)
- [Fundamental Limits of Crystalline Equivariant Graph Neural Networks: A Circuit Complexity Perspective](https://mlsanity.com/paper/2510.05494/) (2025)
- [Expressivity In Multimodal Contrastive Learning](https://mlsanity.com/paper/2608.17203/) (2026)
- [Deep operator learning for efficient sampling from invariant measures of stochastic differential equations](https://mlsanity.com/paper/2609.11376/) (2026)

## Abstract (arXiv)

We establish the universal approximation capability of single-layer, single-head self- and cross-attention mechanisms with minimal attached structures. Our key insight is to interpret single-head attention as an input domain-partition mechanism that assigns distinct values to subregions. This allows us to engineer the attention weights such that this assignment imitates the target function. Building on this, we prove that a single self-attention layer, preceded by sum-of-linear transformations, is capable of approximating any continuous function on a compact domain under the L_∞-norm. Furthermore, we extend this construction to approximate any Lebesgue integrable function under Lp-norm for 1≤ p <∞. Lastly, we also extend our techniques and show that, for the first time, single-head cross-attention achieves the same universal approximation guarantees.