Learning Theory of Transformers: Local-to-Global Approximation via Softmax Partition of Unity
Authors: Zhongjie Shi, Wenjing Liao
Summary
The authors develop a theoretical framework for understanding how Transformers approximate functions. Their key insight is that Transformers can build local approximations of a target function and blend them together using softmax as a "partition of unity"—the attention mechanism creates spatial localization and softmax stitches the pieces into a coherent global output. They prove that shallow-but-wide Transformers with just two encoder blocks can approximate smooth functions efficiently and achieve near-optimal generalization.
Main takeaways:
- Transformers work by learning many local approximations and using softmax attention to weight-average them based on position
- Two encoder blocks plus simple feed-forward layers are enough to uniformly approximate smooth (Hölder continuous) functions with O(ε^(-d/α)) parameters
- Generalization error is near minimax-optimal at O(n^(-2α/(2α+d)) log n) for n training samples
- The architecture studied is shallow and wide (not deep), uses softmax and sinusoidal positional encodings like real Transformers
- Softmax plays a dual role: it's both the aggregation mechanism and the key to proving uniform approximation bounds
Relevance
Tangential to my work on behavioral installation and persona switching—this theory shows how attention localizes and aggregates, which might inform how personas or markers could be spatially localized in representation space, but it's a foundational ML theory paper rather than directly about conditioning.
Abstract
arXiv:2605.08811v1 Announce Type: new Abstract: This paper investigates the learning theory of Transformer networks for regression tasks on the compact Euclidean domain $[0,1]^d$ and $d$-dimensional compact Riemannian manifolds. We propose a novel constructive approximation framework for Transformers that builds local approximations of the target function and aggregates them into a global approximation via softmax partition of unity. This approach leverages the attention mechanism to achieve spatial localization through affine transformations of the input. The softmax activation plays a crucial role in aggregating local approximations to a global output. From an approximation perspective, we prove that a dense Transformer equipped with only two encoder blocks and standard single-hidden-layer point-wise feed-forward networks can achieve a uniform $\varepsilon$-approximation error for $\alpha$-H"older continuous functions with $\alpha \in (0,1]$ using $\mathcal{O}(\varepsilon^{-d/\alpha})$ total parameters. Building upon this approximation guarantee, we establish a near minimax-optimal generalization error bound of order $\mathcal{O}\big(n^{-\frac{2\alpha}{2\alpha+d}} \log n\big)$ for the empirical risk minimizer, where $n$ is the training data size. The Transformer architecture studied in this paper is dense, shallow and wide, and employs softmax activation and sinusoidal positional encodings, closely reflecting practical implementations.