Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Attention Mechanism, Max-Affine Partition, and Universal Approximation

Hude Liu, Jerry Yao-Chieh Hu, Zhao Song, Han Liu

TL;DR

This work proves that a minimalist attention block—comprising a single-layer, single-head self- or cross-attention preceded by a linear layer—has universal sequence-to-sequence approximation power under $L_\infty$ (and $L_p$ for $1\le p<\infty$) norms, even without positional encodings or FFNs. It introduces a max-affine partition interpretation of attention, where attention scores encode partition indicators and reassign values across partition cells, enabling piecewise-constant and piecewise-linear approximations of target functions. The results establish a constructive, theory-grounded view of Transformer expressivity, showing cross-attention shares the same universality as self-attention under a minimalist design. Practical extensions and proof-of-concept experiments further illustrate how these partition-based mechanisms behave under noise and in data-limited regimes, highlighting both the potential and current limitations of attention-centric universal approximation in real-world settings.

Abstract

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_\infty$-norm. Furthermore, we extend this construction to approximate any Lebesgue integrable function under $L_p$-norm for $1\leq p <\infty$. Lastly, we also extend our techniques and show that, for the first time, single-head cross-attention achieves the same universal approximation guarantees.

Attention Mechanism, Max-Affine Partition, and Universal Approximation

TL;DR

This work proves that a minimalist attention block—comprising a single-layer, single-head self- or cross-attention preceded by a linear layer—has universal sequence-to-sequence approximation power under (and for ) norms, even without positional encodings or FFNs. It introduces a max-affine partition interpretation of attention, where attention scores encode partition indicators and reassign values across partition cells, enabling piecewise-constant and piecewise-linear approximations of target functions. The results establish a constructive, theory-grounded view of Transformer expressivity, showing cross-attention shares the same universality as self-attention under a minimalist design. Practical extensions and proof-of-concept experiments further illustrate how these partition-based mechanisms behave under noise and in data-limited regimes, highlighting both the potential and current limitations of attention-centric universal approximation in real-world settings.

Abstract

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 -norm. Furthermore, we extend this construction to approximate any Lebesgue integrable function under -norm for . Lastly, we also extend our techniques and show that, for the first time, single-head cross-attention achieves the same universal approximation guarantees.
Paper Structure (63 sections, 16 theorems, 210 equations, 4 figures)

This paper contains 63 sections, 16 theorems, 210 equations, 4 figures.

Table of Contents

  1. Introduction
  2. Organization.
  3. Provable Capabilities of Transformer.
  4. Preliminaries
  5. Self-Attention and Cross-Attention Layers.
  6. Layer of Sum of Linear Transformations.
  7. Attention as Max-Affine Value Reassignment
  8. Max-Affine Partition
  9. Attention Scores Encode Indicators for Max-Affine Partition
  10. Attention Reassign Value to Each Part of the Max-Affine Partition
  11. Single-Layer, Single-Head Attention Achieves Universal Sequence-to-Sequence Approximation
  12. Single-Head Self-Attention as a Universal Seq-to-Seq Approximator
  13. Single-Head Cross-Attention as a Universal Seq-to-Seq Approximator
  14. Extension to Practice
  15. Proof-of-Concept Experiments
  16. ...and 48 more sections

Key Result

Proposition 3.1

Following def:max-affine_f, consider a max-affine function $\mathop{\rm MaxAff}(x) = \max_{i\in [N_{\rm ma}]} \{a_i^\top x+b_i\}$, and let $\mathcal{X} \subset \mathbb{R}^{d_x}$ be its input domain. Then $\mathop{\rm MaxAff}$ generates a partition on $\mathcal{X}$: We call the partition $P_{\rm ma}$ the max-affine partition of $\mathcal{X}$ induced by $\mathop{\rm MaxAff}$.

Figures (4)

  • Figure 1: Max-Affine Partition on a $2$-Dimensional Domain. The different affine components are dyed with different colors. The domains corresponding to the different affine components are also dyed with the according color.
  • Figure 2: Value Reassignment of Max-Affine Function in \ref{['fig:2d-ma-partition']}. Each part of the max-affine partition in \ref{['fig:2d-ma-partition']} was assigned with an affine function different from the original affine function that generated this partition.
  • Figure 3: Scale of Attention Weights vs. Training noise. For MNIST, CIFAR-10, and Fashion-MNIST we plot the $\ell_2$-norm of $W_K$ and $W_Q$ against the injected label-noise ratio. In all three datasets the weight scale declines monotonically as noise increases, corroborating \ref{['prop:indicator']}: higher noise hampers precise partitioning, so the model reduces the magnitude of weights that form the attention score matrix.
  • Figure 4: Result of using a single-head attention to approximate a step function. The max-affine function generated in the attention score matrix turns at points close to the switching points in the step function.

Theorems & Definitions (48)

  • Definition 3.1: Max-Affine Function
  • Remark 3.1: Technical Assumption
  • Proposition 3.1: Max-Affine Partition
  • Definition 3.2: Indicator of Max-Affine Partition
  • Proposition 3.2: Attention Approximates Indicator of Max-Affine Partition
  • proof : Proof Sketch
  • Proposition 3.3: Attention Reassigns Value to Max-Affine Partition
  • proof : Proof Sketch
  • Remark 3.2: Extension to Function on All Tokens
  • Remark 3.3
  • ...and 38 more