Title: Universal Approximation with Softmax Attention
Primary area: foundation or frontier models, including LLMs
Keywords: universal approximation, attention, expressiveness, in-context learning

Abstract: 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.

TL;DR: We prove one- and two-layer attention with linear transformation are universal approximators for continuous sequence-to-sequence functions..