Research map / Architectures and efficiency / Attention and transformers
Attention Mechanism, Max-Affine Partition, and Universal Approximation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that a single softmax attention layer (single head, with a prepended sum-of-linear layer) is a universal approximator for continuous and Lebesgue-integrable sequence-to-sequence functions, via a max-affine partition interpretation of attention, and proves the first such result for cross-attention.
Problem
Whether attention alone (single layer, single head, no positional encodings, FFNs, or multi-head) can universally approximate sequence-to-sequence functions, under L_infinity and L_p norms, for self- and cross-attention.
Approach
Interpret single-head softmax attention (with a linear layer before it) as a max-affine partition of the input domain that assigns distinct values to each cell; engineer W_K, W_Q, W_V so the softmax acts as a near-hard argmax indicator and the value matrix outputs f on each cell, yielding universality with attention only.
Key results
- Single-layer single-head self-attention with a prepended sum-of-linear layer approximates continuous functions on compact domains in L_infinity norm and integrable functions in L_p norm.
- Single-head cross-attention with a prepended linear layer is shown to have the same universal approximation guarantees, claimed as the first such result.
- In proof-of-concept experiments on MNIST, CIFAR-10 and Fashion-MNIST, the norm of W_K and W_Q decreases monotonically as injected label noise increases.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase Gaussian noise in the training data | Single-head self-attention on MNIST, CIFAR-10, and Fashion-MNIST · MNIST, CIFAR-10, Fashion-MNIST training sets | W_K and W_Q weight scale | Hurt |
Limitations
- Constructions need many partition regions, inflating parameter count for high-dimensional inputs or long sequences.
- Proofs are constructive, not prescriptive; gradient training may not find the required weights.
- Does not address distribution shift or generalization beyond the compact domain.
Experimental setup
Datasets: MNIST CIFAR-10 Fashion-MNIST
Metrics: l2-norm of W_K and W_Q versus label-noise level
Baselines: Yun et al. transformer universal approximation Kratsios et al. universal approximation under constraints
Builds on: max-affine functions (convex function approximation) Lusin's theorem transformer universal approximation results
Closest papers by meaning
- Are Transformers with One Layer Self-Attention Using Low-Rank Weight Matrices Universal Approximators? (ICLR 2024 · similarity 0.87)
- Towards Interpretable and Efficient Attention: Compressing All by Contracting a Few (NeurIPS 2025 · similarity 0.87)
- ZeroS: Zero‑Sum Linear Attention for Efficient Transformers (NeurIPS 2025 · similarity 0.86)
- Mixture-of-Top-$k$ Attention: Efficient Attention as Scalable Fast Weights (NeurIPS 2026 · similarity 0.86)
- Fast attention mechanisms: a tale of parallelism (NeurIPS 2025 · similarity 0.86)
- Geometry of Lightning Self-Attention: Identifiability and Dimension (ICLR 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.