Universal Approximation with Softmax Attention

We prove that with linear transformations, both (i) two-layer self-attention and (ii) one-layer self-attention followed by a softmax function are universal approximators for continuous sequence-to-sequence functions on compact domains. Our main technique is a new interpolation-based method for analyzing attention’s internal mechanism. This leads to our key insight: self-attention is able to approximate a generalized version of ReLU to arbitrary precision, and hence subsumes many known universal approximators. Building on these, we show that two-layer multi-head attention or even one-layer multi-head attention followed by a softmax function suffices as a sequence-to-sequence universal approximator. In contrast, prior works rely on feed-forward networks to establish universal approximation in Transformers. Furthermore, we extend our techniques to show that, (softmax-)attention-only layers are capable of approximating gradient descent in-context. We believe these techniques hold independent interest.

Modern AI systems (e.g., foundation models like ChatGPT or Gemini) often use Transformer models. Yet we still lack a clear theory of what makes them expressive. This paper studies attention, the component that lets each part of an input look at other parts. We prove that attention alone can approximate any well-behaved rule from one bounded sequence to another. The proof shows that attention can act like a selector. It compares an input with many reference values and chooses a nearby one. This turns attention into a simple building block for representing complex sequence rules. Our results show that the extra neural-network layers commonly paired with attention are not necessary for this kind of expressive power. We also show that attention can imitate a basic learning procedure inside its input context. Overall, these results clarify why attention is powerful and may help guide simpler theories and designs of Transformer models.