Universal Approximation with Softmax Attention
This provides theoretical justification for attention-only architectures in Transformers, addressing a foundational problem in machine learning theory.
The paper proves that two-layer self-attention and one-layer self-attention with softmax are universal approximators for continuous sequence-to-sequence functions, showing they can approximate a generalized ReLU to arbitrary precision and subsume known universal approximators.
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 alone 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 various statistical models in-context. We believe these techniques hold independent interest.