Attention Mechanism, Max-Affine Partition, and Universal Approximation
Authors: Hude Liu
Jerry Yao-Chieh Hu
Zhao Song
Han Liu
Venue: NeurIPS 2025 poster ()
Keywords: Attention
Universal Approximation
Primary area: deep_learning
TL;DR: We establish the universal approximation capability of single-layer, single-head self- and cross-attention mechanisms

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
Hude Liu" Jerry Yao-Chieh Hu*t
Zhao Song
Han Liut
CCenter for Foundation Models and Generative & Department Computer Science,
Simons Institute for the Theory Computing, UC Berkeley, Berkeley, CA 94720, USA Department of Statistics and Data Scienee, Northwestern University, Evanston, IL 60208, USA
hudel hudelluu2080gmail com, wow.bnnnnthvveten.ccm edu, magic. 1 limaxkde com, hanliu@horthrestern.cc
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 domain-partition mechanism that assigns distinct values subregions. This allows us to engineer the attention weights such that this assignment imitates the target function. Building on this, we prove that single self-attention layer, preceded by sum-of-linear transformations, is capable of approximating any continuous function the -orm. Furthermore, we extend this construction to approximate any Lebesgue integrable function under Ly-norm for S p co. Lastly, we also extend our techniques and show that, for the first time, single-head cross-attention achieves the same universal approximation guarantees.
compact domain under
1 Introduction
We establish the universal approximation capability of single-layer, single-head self- and cross- attention mechanisms. Departed from prior studies, results demonstrate that the expressive power of transformers arises from only the (softmax) attention module and an attached linear layer, without additional components such positional encodings feed-forward networks (FFNs). More importantly, our proofs show that sentennc-to-sequene universal approximation requires only minimalist configuration; single-layer, single-head attention with lincar transformations.
In this era, the power of transformers [Vaswani et al., 2017] is undeniable, given their dominance in modern machine learning. They drive models such BERT [Devlin, 2018], ChatGPT [Brown et al., 2020, Achiam et al., 2023], and LLaMA [Touurron et al., 2023a,b, Dubey et al., 2024] for language; VIT [Dosovitskiy al., 2021] and DiT [Peebles and Xie, 2023] for image and video; DNABERT [Ji et al., 2021, Zhou et al., 2023] for genomics; and Moirai [Woo et al., 2024, Liu et al., 2024] for time series, among many others. Central these successes the allention mechanism. While numerous variants and implementations exist [Tay et al., 2022], the softmax-based vanilla attention
'Equal contribution. Version: December 15, 2025. Future updates are on https: /ettps://an h/tp://aaxvvvoo// abgg/ba///00..
19901.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

[Vaswani et al., 2017] remains mainstay in both research and industry communities (e.g., ChatGPT and Llama).
However, despite its practical importance, theoretical insights into why softmax attention is so powerful remain incomplete. Moreover, the extent which softmax attention alone drives perfor- mance is unclear. Empirical [Tay et al., 2022] and theoretical [Keles et al., 2023, Deng et al., 2023, Alman and Yu, 2024] evidence suggests that deviating from softmax attention (e.g., via sub-quadratic approximations) often degrades performance, indicating that softmax attention may be a central engine in Transformer architectures. At the same time, growing body of work explores its memory capacity [Mahdavi et al., 2023, Kim et al., 2023, Kajitsuka and Sato, 2024], universal approximation properties [Yun al., 2019, Kajitsuka and Sato, 2023, Jiang and Li, 2023], representation learning [Sanford et al., 2024b, Chen and Li, 2024], and task-specific theoretical performance [Gurevych et al., 2022, Edelman al., 2022]. However, these studies often rely on additional components, such as feed-forward networks (FFNs) or multi-head setups customized assumptions, as they target the entire Transformer architecture rather than isolating the role of attention module.
To this end, this work presents attention-only expressiveness results; seftmax-based attention alone already suffices for universal approximation of coguuence-t--seqeenee functions. We operate under three key premises for investigating the expressiveness of attention:
1. We focus on softmax-based attention,
2. We seek a minimalist design (a single layer of single-head attention plus a linear transforma- tion),
3. We impose minimal assumptions on the data distribution or network architecture (no posi- tional encodings, multi-head expansions, no FFNs).
We provide new proofs that single self-attention layer approximates any continuous sequence-to- sequence function on compact domain, in both the L and T, norms. Furthermore, we show, for the first time, parallel result for cross-attention, revealing its universal approximation capability under the same minimalist setting.
Contributions. Our contributions are as follows:
Interpreting Attention as a Max-Affine Partition. We show that single-head softmax attention, combined with a linear layer, implicitly partitions the input domain using a max- affine construction. This partitioning allows attention to assign distinct outputs to each partition cell. This perspective clarifies how softmax-based attention enables a powerful piecewise-linee linear approximation scheme.
Siygle-Layer, Single-Head Self-Attention Universality. We prove that a single self- attention layer is universal approximator for continuous on compact domains. Our results cover both Ly- and -noms guarantees and require minimal assumptions on data and architecture, highlighting the inherent expressive power of attention alone.
functions
Single-Head Cross-Attention Universality. We establish, for the first time, that the same approach also endows a single-layer, single-head cross-attention with universal approxima- tion capabilities. This result further underscores that much of a Transformer's expressiveness can reside solely in its attention block, even when the queries and keys come from distinct input sequences.
Organization. Section presents the ideas built on. Section 3 shows our interpretation of Atten- tion as a Max-Affine Partition in a simplified setting. Section 4 presents our universal approximation results for single-layer, single-head self- and tross-attentions.
1.1 Related Work
Universal Approximation. Farly works of universal approximation theorems focuses on the ex- pressiveness of feed-forwart networks (FFN) [Cybenko, 1989, Hornik, 1991, Carroll and Dickinson,

---

1989]. Since Vaswani et al. [2017] propose the transformer architecture and the scaled dot dot-product attention module, there is series of research aiming to explain the expressiveness of transformer. Yun et al. [2019], Kajitsuka and Sato [2023] offer explanation from the perspective of contextual map- ping. Among them, Yun et al. [2019] are the first to prove the universal approximation capability of transformer. Yet since the network in [Yun et al., 2019] requires excessive layers (O(n(1 /8) /n!)), Kajitsuka and Sato [2023] make more careful estimation upon the numerical results of contextual mapping and proves that with skip connections, one-layer transformer is capable of approximat- ing any permutation equivariant continuous function. Takakura and Suzuki [2023] add positional encoding to lift the restriction of permutation equivariance, and demonstrate one-layer transformer approximates equiveqquiariintt a-smoothness function with an error independent of input and output dimension. Jiang and Li [2023] give non-constructive proof using Kolmogorov representation theorem on the Jackson-type approximation rate of a two-layer transformer. While prior works have achieves diverse and extensive result regarding the expressive capability of transformer, their results require the feed-forward network (FFN) to add expressiveness to the attention module in order to achieve universal approximation, which differs from our results derived from attention-only network. Concurrently, Hu et al. [2025a] give an interpolation-ban -based proof that softmax attention alone (no FFN) is a universal approximator for continuous squuence-to-sqqeence maps on compact domains.
Provable Capabilities of Transformer. Recent theoretical studies also shed light on the practical behavior of attention mechanism. Olsson et al. [2022] show that induction heads help models learn patterns in context. Sanford et al. [2024a] prove that Transformers can do complex computations with few layers because they work in parallel. In contrast, Luo et al. [2022] find that some Transformer designs lose expressivity when using relative positional encodings. Kim and Suzuki [2024), Chen et al. [2025) provide Transformer's hardness results learning constrained boolean functions. Building on [Hu et al., 2025a]. Hu et al. [2025b] show that a fixed two- evo-attention-lat -ateeetion-lyyer softmax Transformer prompp-programmabll it emulates any algorithm implementable by single attention layer (cf. [Bai et al., 2023]). providing a constructive account of ane-modell-mny-tte (not ReLU) Transformers. To add on these ideas, we prove that a single-layer, single-head softmax attention with a simple linear layer can approximate any continuous function on a compact domain. This shows that attention alone can learn arbitrary sequence-tor sequuueett--squuece mappings.
behavior with softmax
2 Preliminaries
We now present some ideas built on.
Notation. For a vector v, we denote its i-th entry by v, and its subvector from the i1-th to the ig-th entry (inclusive) by Uigsig with i, < ig. For matrix M, we use Mi, for the entry in the i-th row and j-th column, Mr.; for the i-th row, and M., for the j-th column. The submatrix spanning rows through i2 and columns f1 through f2 is denoted by Mf 2i. with i < ,2,,11 < j2- We define Caxb as an a X matrix with constant entries c, and abbreviate cax1 as ca. For norms, we define ||- lxc as the maximum absolute element in a vector or matrix. The p-norms are given by ||ll| = <E; |oa (e/lo)/p for a vector and |MI| (Ex, (o) for matrix M. For function norms, we define the L norm as || = sup=ex, [x))) where Xy is the input domain of f. and the Lp norm as IIf,,,, == ULEx, [=))) dz) 1/P for 1 S p < co. For functions, when function f vector/matrix (i.c..exp((e, ag)) = [exp(ar), exp(a)))
+ R applied on a vector or a matrix, it means to apply f on every entry of the
Self-Attention and Cross-Attention Layers. For self-attention Attn, DDN +R DxNous, and any input Z E DDN, we define its output as:
Attn, (Z) = Wv Dooftmax((Wxx22 T
where Wx, We E RdAuLDD Wv E DxD, Wo E N.NxNou Here d'ttn stands for the hidden size of the attention block. Nont stands for the output sequence length.
For a cross-attention Attn. : RDNN x RDxN + DxNom and any input Zk. Zq E DxN we define its output as:
Wv Zι Softmax T Attne (K.. Zo)
Here WK. Wo, Wy, Wo are defined as those in self-attention.
3

---

Since we provide separate discussions for self attentitnnion and cross- attss-attnntoon in this work, we omit the subscript and denote them as Attn when this causes no ambiguity.
Layer of Sum of Linear Transformations. We use Linear: PRDxx + D2xN2 to denote layer of sum of linear transformations. For any input Z E RRDxx xN; define its output follows:
Linear (Z)
Σ PIZQ R.
where P; E DDxxD, Qi E N for E [M], R E IR PaxN2, Here H is positive integer which denotes the number of linear transformations sum.
3 Attention as Max-Affine Value Reassignment
In this section, we introduce new interpretation of attention as value reassignment max affine function. Essentially, we show that attention prepended with Linear layer is able to reassign values partition generated by a max max-affine function. We start with the below definition.
Definition 3.1 IMa--fffiee Function). Let X' d be domain, and fix positive integer Nua- For each i [Nma], define an affine function : X + for all x E X::
(ı(x) = a. bi, where E R and by E ER.
The max-affine function MaxAff - corresponding to affine functions (yi()) is defined as
MaxA(f(x) .nax {a x b:). iE(Num)
Intuitively, a max-affine function selects, at each point x E X', the largest output among Nma affine functions. Geometrically, each affine function yı(x) a b; defines hyperplane in Rd++1 Thus, MaxAff follows the highest hyperplane at each æ; forming piecewise linear, convex surface the upper envelope of the given affine hyperplanes.
Remark 3.1 (Technical Assumption). For simplicity of presenting our interpretation, we make the following technical assumption
results this section:
Assumption 3.1. For any max-affine function MaxAff, we exclude situations where the difference between its largest and second-larg largest affine components smaller than specified threshold. (Please see proofs for explicit definition.)
We do not apply this assumption in other sections.
3.1 Max-Affine Partition
We show that max-affine function MaxAff(.) induces partition of its input domain X. Specifically, the input domain is divided up according to which affine function is the maximum at each point To be conerete, define this partition follows:
Proposition 3.1 (Max-Affine Partition). Following Definition 3.1, consider max-affine function Max (x)) max {a, x + bi), and let X c IRds be its input domain. Then Max generates a partition on X:
Puna = {U; [Nmal}, U= & MaxAff(x) = a, x + b,), i [Nma].
We call the partition Pi.. the max-affine partition of X' induced by MaxAff.
Intuitively, U, is the set of all point for which the i-th affine function a. x + b, achieves the same value as the max-affine output. Since MaxAff(-) is the maximum of all the affine components, the i-th component is (one the highest among all components. Hence, the input domain X becomes partitioned "regions" {Ush=(Nma)- That is, if point x belongs to a region U,, the corresponding affine function a; + by is (tied for) the largest. Please Appendix D. for detailed proof.

---

Set Overlaps and Boundaries. By construction, every E X lies in at least one of the sets {U;]. but may belong to multiple if several affine components attain the same maximal value. Hence, the collection (Vi} is generally "partition" in an informal sense: while each U; is typically associated with a distinct region, their pairwise intersections are non-empty on boundary hyperplanes. We address these overlaps in detail within our theorems, where boundary regions do not affect the main approximation arguments but require careful handling to ensure mathematical rigor.
Indicator Encoding of the Partition. For certain analytical and algorithmic tasks, it is helpful to embed the notion of "which affine part is active" into a vector-valued indicator. Formally, we define the indicators for max max--afine partitions.
Definition 3.2 (Indicator of Max-Affine Partition). Following the notations in Proposition 3.1, for a max-affine partition { E [Nmal]: we define b := argmax Nam (vs(z)) to the label of the maximal affine component. Then, we define the indicator E : Rd. - *R" Noa as:
(Nmm)
F(x)
which is one-hot vector whose only non-zero entry is the i one.
Namely, each component of ((x) is zero unless it corresponds to an index achieving the maximum, in which has the value of 1. In Figure la, we show an example of the max-affine partition.
X
(a) Max-Affine Partition on a 2-D Domain. Colored regions show where each affine component is active.
(b) Value Reassignment of Figure la. Each region reassigned different affine function.
3.2 Attention Scores Encode Indicators for Max-Affine Partition
We discuss the connection between self-attention and max-affine partition. We show that self-attention with Linear layer attached before can generate a max-affine partition. Further, for every input token, the attention belongs to. We state this result as follows:
matrix approximately indicates which part of the partition it
Proposition 3.2 (Attention Approximates Indicator of Max-Affine Partition). Let X
X2,, .x,.] d denote any input sequence. We use X' to denote the domain all X,, E in]. Let MaxAff be any max-affine function on X' with Nma components, and let > 0 be any positive real number. We define P Pm. = {U,li E [Numa]] as the max-affine partition generated by MaxAff in Proposition 3.1. Then, there exists Linear layer and self-attention Attn whose attention matrix satisfies:
Softmax ((W (aneaeeWWWWWnrr/rrr Linear(X)) W WoLincar(X))Ve
[F(X)), E(X2). F(Xx))llo SC,
with exception of region of arbitrarily small Lebesgue measure in R Here WK. Wo the attention weights within Alln. Wo only truncates the irrelevant part of the attention score matrix.
5

---

Proposition 3.2 shows that the attention matrix is able to approximate a vector denoting the position of the input token, by indicating which part the max-affine partition contains the input token.
3.3 Attention Reassign Value to Each Part of the Max-Affine Partition
In the work of [Kim and Kim, 2022], they prove that max-affine functions are universal approximators for convex functions. In order to turn them into universal approximators, a possible solution is to reassign value to each part of the max-affine partition generated by the original max-affine function. In the following theorem, we show that single-heal self-attention is capable of completing this task.
Proposition 3.3 (Attention Reassigns Value Max-Affine Partition). Following the notation in Proposition 3.2, Let F _² - Rou be piece-wise constant function which is separately constant on each U,, E [Nma]. We show that for any there exists an self-attention Attn such that
[X)ttn(X)= (F F(X.), ll S,
for every X in X' with exception of region of arbitrarily small Lebesgue measure in R".
Proposition 3.3 shows that attention is able to output different values according to the indicator generated in Proposition 3.2.
We conclude this section with two remarks.
Remark 3.2 (Extension Function on All Tokens). In this section, for the conciseness, in demon- stration method, we adopted a token-wise function F as the example function. Yet since affine functions on all tokens can be easily obtained by adding token-wisc affine functions, this simplified version of our method generalizes well on functions taking all tokens as input and leads us to results shown in Section 4. Remark 3.3. Lastly, emphasize that here the approximation excludes small area for overall simplicity in this demonstration of our method. We address this issue in the proofs of the universal approximation theorems in the section.
Figure 1b provides us example of Proposition 3.3.
Single-Layer, Single-Head Attention Achieves Universal Sequence-to-Sequeece Approximation
In this section, we present main results:
A single layer of single-head self-attention preceded by one linear layer sequence-to-sqquence universal approximator for continuous functions on any compact domain.
A single layer of single-head cross-attention preceded by one linear layer is likewise sequence-to-sequenee universal approximator for continuous functions on any compact domain.
Importantly, we achieve attention-only universal approximation for both the L, Ln-norm and L norm, whereas most existing results apply only the -ormn and require additional auxiliary components in the transformer block (e.g., multiple attention or feed-forward layers). Moreover, our universality result for cross-attention is the first of its kind. Specifically, we present results for self-attention in Section 4.1 and for cross-attention in Section 4.2.
4.1 Single-Head Self-Attention as Universal Seq-to-Seq Approximator
We now present our main result: single-layer, single-head self-attention module, combined with linear transformation, is sufficient approximate any continuous map f _dxx gdxa on a compact domain U C [-D, Djdx", We present the result first in terms of the Loo norm for continuous and then extend it to Lp integrable functions.

---

Theorem 4.1 (Loo Universal Approximation) Let f _dx= + Rd denote any continuous function on a compact domain U C Rdx= and let be any positive real number. There exists self-attention Attn with a prepended Linear layer, such that
f Attn o Linear| S.
Theorem 4.1 indicates that a single-layer self-attention block, combined with a linear preprocessing layer Linear, approximates
in the -nom....
Overview of Proof Strategy. We adopt proof strategy based on key observation: self-attention is capable of approximating target functions via implicit MaxAff operations. Our proof consists of the following steps:
Step 1: Partition Input Domain U via MaxAff. Construct a max- afx-nffine function MaxAff
over U (i.e., input domain of target function f) such that this MaxAff induces partition of size-Nma U.
Step 2: Configure Linear and Attn Imitate MaxAff over U. Use Linear and WK, Wa
in Attn to map the input Z & U to values of the affine components {vi(Z) = 0
bi of MaxAff. Here we flatten the input sequence Z E Rdx= Z E R dr to
compute MaxAff.
Step 3: Engineer Attn to Generate an Indicator of Which Partition Cell the Input Belongs To. Within self-attention Attn, design K Q so that Softmax Q) produces near-one-hot vector an indicator to the max-affine partition induced by MaxAff (as defined in Definition 3.2). This indicator (approximately an one-hot vector) shows which part (i.e., partitioned cell) of the partition contains the input sequence Z.
Step 4: Map the Indicator to the Target Value (()). Map each partition cell's indicator to the corresponding value of f. By continuity of f, refining the partitioned cell ensures [[ Attn Linear S
Proof Sketch. We elaborate above in detail. Consider a continuous function U Djdxn - gdx= on a compact domain U. Let >> 0. We aim to construct a single-layer, single-head self- attention mechanism Attn (prepended with a linear transformation Linear) such that
|| Attn Linear|
Step 1: Partition Input Domain via MaxAff.
Flattening Input. Each input e Rdx= is reshaped into single vector Ž Idd by stacking its rows or columns. This unifies the domain as D)
Grid / Max-Affine Construction. Since f uniformly continuous on the compact set U, choose such that
2 2211 < = I <.
G-1 We subdivide [-D, Djdn into cubes of side 6, yielding - pdn grid centers { We treat MaxAff piecewise (max-)affine or piecewise-constant partition: for each Z, there's s a nearest v within 6/2.
Technical Highlight. This partition- based approach leverages uniform continuity to dis- cretize U. The number of partitions can be large but finite, ensuring only need single-layer attention to "select" the correct grid cell.
Step 2: Configure Linear and Attn to Imitate MaxAff over U.
Trannaaar-TTnnffrmattion Map Linear. Design Linear : Rdxx + RM (for some dimension M) to capture the dot products (vy, Z). Essentially, Linear(Z) arranges these
7

---

{o] 2} in a form accessible to attention. This ensures each grid center by can be individually
"queried."
Encoding Affine Components. Observe that max; max, Z) - =11,;11²} is akin to max- affine function. We store terms o) Z, plus - ~2,,,,1?, into K and Q for later use in Softmax( ( Q).
Technical Highlight. This step demonstrates how we embed {(0,, 2)} into a single-head attention setting no extra feed-forward layers required. The linear map Linear is carefully constructed so that each "component" is individually addressable.
Step 3: Engineer Attn to Generate an Indicator of Which Partition Cell the Input Belongs to.
Construct KT Q. the self-attention block, let KTQ = R((,, 2) large. This makes Softmax (K' (Q) favor the row 5" maximizing
where 2
(v), 2) |1yl1
Near-One-Hot Distribution. Hence the j'-thh row obtains probability close to 1, effectively identifying which grid center vy- is nearest Z. We interpret this as near-one-hot "indicator" vector for the correct partition cell.
Technical Highlight. This is the crux: attention's ' softmax can act as a continnous arg max by scaling the scores with R. As R oo, the distribution becomes more peaked, approxi- mating hard partition.
Step 4: Map the Indicator to the Target Value f(Z).
Assigning Values. We place f(u;) in the "value matrix" Wv, v, that once row j is selected, the attention output is z f(uy-). Since Z is within 8/2 of Uyz, uniform continuity implies
llf(z) (1y--)) )II S
(for suitably chosen 8).
Final Reshaping (If Needed). A small linear projection M can reshape the output back Rdxn, The essential logic is that the correct f(u;) is "routed" to the final output via the near-one-hot attention distribution.
Technical Highlight. This reveals how single-head attention layer, armed with linear preprocessing, suffices replicate the entire function f. No feed-forward sub-layer or multiple heads are needed to achieve universal approximation.
that: (i) A finite grid subdivides U to handle uniform continuity. (ii) Linear encodes ((uy 3) (iii) Large-R Softmax K Q) selects the best anchor uy-- (iv) A "value matrix" translates that selection into f(uz-). We conclude that single-layer, single-head self-attention block approximates within in the Los norm. Please see Appendix E. for proof. O
In sum, combining these steps,
Our result in Lee norm can be easily extended to Lp norm, where it applies to not just the continuous functions but all Lebesgue integrable functions with compact support. Please see Corollary E. for more details.
4.2 Single-Head Cross-Attentiin
Universal Seq-to-Seq Approximator
Here we extend self-attention universal approximation results from Section 4.1 cross-attention. Importantly, we establish the first known universal approximation in cross-attention setting. First, we state our main result in nnrm.
Theorem 4.2 -Nom Universal Approximation). Let f Ux Uq - Rdxx denote any continuous function on compact domain UK x Uo and let be any positive real number. Here Uk, Uq E xdxa stands for the compact domain the two input sequences of cross-attention. Then

---

there exists cross-attention Attn prepended with Linear layer such that
|/ Attn Linear|| S
Theorem 4.2 indicates that single- layer cross-attention block, prepended with linear preprocessing layer Lincar, approximates f UK + Uq + Rdxa in -nomm
-norm.
Proof Sketch. Our proof follows that of Theorem 4.1 except one additional step: use Attn to aggre- gate the max-affine functions on UK, Ug and merge into MaxAff function on UK x Uq. The proof consists of the following steps:
Step 1: Partition the Input Domain UK and Uo with MaxAff and MaxAffo Respectively, Con-
struct two max-affine function MaxAffx over Up and MaxAffo over Uo such that this MaxAff induces partition of size-Nma of U and MaxAflo same size partition Uq.
Step 2: Configure Linear and Attn to Imitate MaxAff K, MaxAff over WK, Up Respectively. Use Linear and WA K. Wq in Attn to map the input ZK. ZQ U to values of the affine components
(v)(Z) = a! 2 + bi Fee[Naua] of MaxAff and MaxAffe respectively. Here we flatten the input sequence Z E gdx= to Z E RRd to express MaxAff concisely.
Step 3: Use Attn Aggregate MaxAff and MaxAffq Form a MaxAff UK X UQ + I on Both Input Sequences. Use Attn to generate MaxAff (Z Zo) := MaxAff x (Zk) - MaxAff (Zz). This max-affine function merges the partition on UK and Uq to generate a unified partition on UK x UQ.
Step 4: Use Attn to Indicate the Position of the Both Input Sequence in the Partition. Use Attn to generate an indicator to the max-affine partition generated by MaxAff (as defined in Definition 3.2). This indicator (approximately one-hot vector) shows which part of the Maxxfff-eneratee partition contains the Cartesian product of both input sequences Zx x Zo.
Step 5: Map the indicator to the Corresponding Value of f. Map the indicator to the corresponding value of the target function f by adding terms related to f to Attn.
Please see Appendix E.2 for detailed proof.
5 Concluding Remarks
We introduce a novel interpretation of attention as a mechanism for reassigning values to a partition induced by a max-affine function. This unique perspective allows us to show that prepending single linear layer before either self-attention or cross-attention enables the network to (i) generate indicator functions representing max-affine partitions (Proposition 3.2) and (ii) selectively reassign values to each partition cell (Proposition 3.3). As a result, we prove that both single-head self-attention and single-head cross-attention, when combined with a single layer of sum of linear transformations, achieve universal approximation of compactly supported continuous functions under Loe norm, or integrable functions under Lp norm. Numerical validations backup our theory in Appendix B.
Key Insights and Results.
Max-Affine Partition. A max-affine function naturally partitions its input domain, and attention (with appropriate transformations) can approximate the indicator functions of these partitions.
Value Reassignment. Self-attention reassigns output values based on partition indicators, capturing a broad class of piecewise-detee definee functions.
Universal Approximation. With only single linear layer and a single-head attention module, one can approximate arbitrary souuence-to-equeeee maps in both the Loc and Lp senses, for both self-attention (Theorem 4.1 and Corollary E.1. and cross-attention (Theorem 4.2 and Corollary 1.2. architectures.
Limitations. While our results highlight the surprising representational power of single-head attention with linear preprocessing, several limitations warrant discussion:

---

Large Dimensions and Network Size. Our minimal -assupption design needs many partition regions to cover diverse targets. This follows naturally from the general setting we study. High-dimensional inputs or long sequences then inflate the parameter count and hinder practice. Appendix eases the burden but does not eliminate it entirely.
Training Complexity. Our proofs are constructive rather than prescriptive for training, meaning standard gradient-based methods may not (always) efficiently find the required weight configurations.
Data Distribution Shifts. Like many universal approximation results, our approach does not account for distribution shifts or generalization beyond the compact domain used for training.
Implications and Future Work. Our findings explain why transformers excel at modeling heteroge- neous data: attention can create flexible partitions of the input space and assign context-dependent outputs. This perspective raises open questions for future research: Can multi-hear or deeper atten- tion layers simplify representational requirements or reduce approximation constants? How might learned partitions or specialized positional encodings improve efficiency in practice? Can adaptive or data-driven strategies automatically discover near-optimal partitions for specific tasks?
Overall, results establish theoretical foundation for understanding attention- based architectures as universal function approximators. They illustrate how token-wise information is partitioned and reassigned to represent complex eqquence-t--equeen functions with minimal assumptions and structural requirements on data and model.
Acknowledgments
The authors thank Mimi Gallagher, Sara Sanchez, Dino Feng and Andrew Chen for useful discussions; and Weimin Wu, Hong-Yu Chen and Jennifer Zhang for collaborations on related topics. JH also thanks the Red Maple Family for support. The authors also thank the anonymous reviewers and program chairs for constructive comments.
Lastly, JH dedicates this work to the memory of his aunt, Lily Cheung, who passed away during its preparation (March 2025). Her loving and caring spirit will always inspire him.
JH supported by the Walter P. Murphy Fellowship and the Terminal Year Fellowship (Paul K. Richter Memorial Award) of Northwestern University. HL is partially supported by NIH ROILM1372201, NSF AST-2421844 Simons Foundation MPS-AL-00010513, AbbVie, Dolby and Chan Zuckerberg Biohub Chicago Spoke Award. This research resources and staff contributions provided for the Quest high performance computing facility at North- western University which is jointly supported by the Office of the Provost, the Office for Research, and Northwestern University Information Technology. The content is solely the responsibility of the authors and does not necessarily represent the official views of the funding agencies.
supported in part through the computational
References
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.088 0887774744 2023.
Josh Alman and Hantao Yu. Fundamental limitations on subquadratic alternatives to transformers. arXiv preprint arXiv:2410.04271, 2024.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36:57125-572. 2023.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Gireeja Ranade Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems, 13:1877-1901, 2020.
10

---

Carroll and Dickinson. Construction of neural nets using the radon transform. In International 1989 joint conference on neural networks, pages 607-611. IEEE, 1989.
Bo Chen, Zhenmei Shi, Zhao Song, and Jiahao Zhang. Provable failure of language models in learning majority boolean logic via gradient descent. arXiv preprini arXiv:2504.0477 2025.
Sitan Chen and Yuanzhi Li. Provably learning a multi-head attention layer. arXiv preprint arXiv:2402.04084, 2024.
George Cybenko. Approximation by superpositions sigmoidal function. Mathematics of control, signals and systems, 2(4):303-314, 1989.
Yichuan Deng, Zhao Song. and Tianyi Zhou. Superiority of softmax: Unveiling the performance edge over linear attention. arXiv preprint 1oXi:22310116683 2023.
Jacob Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.0488 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition scale. In International Conference on Learning Representations, 2021.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024.
Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning, pages 5793-5831. PMLR, 2022.
Iryna Gurevych, Michael Kohler, and Gözde Gül Şahin. On the rate of convergence of a classifier based on a transformer encoder. IEEE Transactions on Information Theory, 68112):813,-8155, 2022.
Kurt Hornik. Approximation capabilities of multilayer feedforward networks. Neural networks, 4(2): 251-257, 1991.
Jerry Yao-Chich Hu, Hude Liu, Hong-Yo Yu Chen, Weimin Wu, and Han Liu. Universal approximation with softmax attention. arXiv preprint arXiv:2504.15956 2025a.
Jerry Yao-Chieh Hu, Hude Liu, Jennifer Yuntong Zhang, and Han Liu. In-context algorithm emulation in fexght transformers. arXiv preprint arXiv:2508. 17550, 2025b.
Yanrong Ji, Zhihan Zhou, Han Liu, and Ramana V Davuluri. Dnabert:
pre-trained bidirectional
encoder representations from transformers model for dna-language in genome. Bioinformatics, 37 (15):2112-2120, 2021.
Haotian Jiang and Qianxiao Li. Approximation theory of transformer networks for sequence modeling. CoRR, 2023.
Tokio Kajitsuka and Issci Sato. Are transformers with one layer self-attention using low-rank weight matrices universal approximators? arXiv preprint arXiv:2307.1400) 2023.
Tokio Kajitsuka and Issei Sato. Optimal memorization capacity of transformers. arXiv preprint arXiv:2409. 77677, 2024.
Feyza Duman Keles, Pruthuvi Mahesakya Wijewardena, and Chinmay Hegde. On the computational complexity of self-attention. International Conference on Algorithmic Learning Theory, pages 597-619. PMLR, 2023.
Jinrae Kim and Youdan Kim. Parameterized convex universal approximators for decision-making problems. IEEE Transactions Neural Networks and Learning Systems, 35(2):24482244 2022.
11

---

Junghwan Kim, Michelle Kim, and Barzan Mozafari. Provable memorization capacity of transformers. In The Eleventh International Conference on Learning Representations, 2023.
Juno Kim and Taiji Suzuki. Transformers provably solve parity efficiently with chain of thought. arXiv preprint arXiv:2410.08553 2024.
Xu Liu, Juncheng Liu, Gerald Woo, Taha Aksu, Yuxuan Liang, Roger Zimmermann, Chenghao Liu, Silvio Savarese, Caiming Xiong, and Doyen Sahoo. Moirai-moe: Empowering time series foundation models with sparse mixture of experts. arXiv preprint arXiv:2410. 240.:4990,0046, 2024.
Shengjie Luo, Shanda Li, Shuxin Zheng, Yan Liu, Liwei Wang, and Di He. Your transformer may not be as powerful as you expect. Advances in Neural Information Processing Systems, 35: 4301-4315, 2022.
Sadegh Mahdavi, Renjie Liao, and Christos Thrampoulidis. Memorization capacity of multi-head attention in transformers. arXiv preprint arXiv:2306.02010, 2023.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint ar arXiv:2209. 11895 2022.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4195-4205, 2023.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Transformers, parallel computation, and logarith- mic depth. arXiv preprint arXiv:2402.0922 092668, 2024a.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational strengths and limitations of transformers. Advances in Neural Information Processing Systems, 36, 2024b.
Shokichi Takakura and Taiji Suzuki. Approximation and estimation ability of transformers for sequence-to-sequeree functions with infinite dimensional input. In International Conference on Machine Learning, pages 33416-33447. PMLR, 2023.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys, 55(6):1-28, December 2022. ISSN 1557-7341, doi: 10. 00.1145/3530811.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv: 13771, 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023b.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Gerald Woo, Chenghao Liu, Akshat Kumar, Caiming Xiong, Silvio Savarese, and Doyen Sahoo. Unified training of universal time series forecasting transformers. arXiv preprint arXiv:2402.02592, 2024.
Chulhee Yun, Srinadh Bhojanapalli,, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? arXiv preprint arXiv: 2019.
Zhihan Zhou, Yanrong Ji, Weijian Li, Pratik Dutta, Ramana Davuluri, and Han Liu. Dnabert-2: Effi- cient foundation model and benchmark for multi-species genome. arXiv preprint arXi::2306.15006 2023.
12

---

Appendix
A Extension to Practical Settings
13
Proof-of-Concepp Experiments
14
Additional Experimental Results
14
C.I Numerical Justifications for Theoretical Results in Section
14
D Proofs of Results Section 3
16
D.I Proof of Proposition 3.1
16
D.2 Proof of Proposition 3.2
16
D.3 Proof of Proposition 3.3
19
E Proof of Results in Section 4
20
E.1 Proof of Theorem 4.1
20
E.2 Proof of Theorem 4.2
34
F Proof of Results in Appendix A
47
F.I Proof of Theorem A.1
47
G Extended Related Work
55
H Extension to Transformer Setting
56
Impact Statement
By the formal nature of this work, we do not expect any immediate negative social impact.
A Extension to Practical Settings
In practical scenarios, despite defined high dimension input domain (R4xa), attention is often considered to approximate a function defined upon small input domain X' C Rdx=.
To this end, extend our method the approximation rate of L-Lipschitz functions with relatively small input domain. We state our result as the following theorem.
Theorem A.1. Let f Rdxx -+ Rdxa denote an L--Lipshitz function (in terms of 2-norm) whose input domain is X'. For any > 0, assume X' is contained in N₂ spheres by the radius of c/(3L) in 2-norm. Then, there exists a Linear layer and Attn layer such that:
||Attn Linear
flloc S.
Furthermore, Attn and Linear have total number of OddnN, trainable parameters.
This proof only differs from the proof of Theorem 4.1 on the choice of partition. For universal approximation, we choose partition that evenly partition the whole space. In this theorem, we change this partition have each part centered on different sphere described in the Theorem A.I. By characterizing our partition, we achieve more precise approximation result.
Proof Sketch. Sketch.
Please Appendix F.1 for detailed proof.
Theorem A. states that when the input domain contained N₂ spheres e-level radius, there exists single-head self-attention layer that approximates the target function with precision of
13

---

B Proof-of-Concept Experiments
00
Figure 2: Scale of Attention Weights Training noise. For MNIST, CIFAR-IO, and Fashion-MNIST we plot the f2-norm Wx and Wq against the injected label-noise ratio. all three datasets weight scale declines monotonically as noise increases, corroborating Proposition 3.2: higher noise hampers precise partitioning, so the model reduces the magnitude of weights that form the attention score matrix.
Proposition 3.2, we demonstrate domain- partitiio mechanism of attention. In this mechanism, the temperature of the Softmax function affects the precision of the max-affine partition generated by attention, which crucial to the complex approximations accomplished in Theorem 4. and Corollary E.1.1.
Since the temperature of Softmax is equivalent the scale of the matrix involved in computing the attention score matrix (WK. WQ), our theory suggests the scale of WK, Wo decreases when the input data contains more noise, as result of the rise in difficulty to form a clear partition, and an approximation based on this partition.
To verify this conjecture, we the correlation between the scale of WK. K, We and the noise level in the training data.
Objectives. Examine the relationship between scale matrix involved in computing the attention score matrix in attention (WK, WQ) and the noise level (using Gaussian noise) in the dataset.
Data. We perform separate experiments on the training set of the noised MNIST, CIFAR10 and FashionMNIST datasets with noise level (the coefficient multiplying the standard Gaussian noise) gradually adding from 0 to 0.72 by the step size of 0.03.
Network setups. Our network consists of a single-head self-attention followed by a feed-forward network. Due the complexity and different characteristics the selected datasets, the size of the feed-forward network slightly differs between datasets.
Results. Figure 2 presents our results. As the noise level increases, a decrease in the scale of weights in WA K, Wq becomes evident in all settings. This aligns with theory.
C Additional Experimental Results
In this section, we present additional experimental results to support our theoretical results.
C.1 Numerical Justifications for Theoretical Results in Section 3
To validate our results in Proposition 3.2, we conducted the following experiment to examine whether the max-affine function generated within the attention of the form in Proposition 3.2 can learn to separate the input domain according to the values of the target function.
Specifically, use attention to approximate step function and observe the max-affine function generated by the weights K and Q matrices in the attention score matrix. The result of this experiment is shown in Figure 3.
14

---

of target fundton
offic
17000
4:00
-15
Figure 3: Result using 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.
The max-affine function generated in the attention matrix turns at points close to the switching points in the step function. This generates a partition in the input domain that resembles the distribution of the flat parts in the step function. This result aligns with theory.
15

---

D Proofs of Results in Section 3
D.1 Proof of Proposition 3.1
Proposition D.1 (Proposition 3.1 Restated: Max-Affine Partition). Following Definition 3.1, con- sider max-affine function MaxAff(z) [axae[Nma] {a x + b,}, and let & C Rd be its imput domain. Then MaxAff generates partition on X::
pm {U i [Nma]];
U; == { E X (axAff(z) = aj x + bi}, E [Nma]-
We call the partition Pms the max-affine partition of X induced by MaxAff.
Proof. If an zo is not grouped to any U4, i [NNaxAu]]- Since MaxAff is define over X' and thus defined on To. we have:
MaxAff(xo) & a 2o + bi, iE [NMaaAAt]
This is contradictory to the definition MaxAff.
Since in Section 3 we exclude the discussion on the overlapped regions of the affine components
{ys- a. x b:), {U; [NMaxA#]] form a partition on X. This completes the proof.
D.2 Proof of Proposition 3.2
Proposition D.2 (Proposition 3.2 Restated: Attention Approximates Indicator of Max-Affine Par- tition). Let X [X1, X2..... X.,) Rdxx denote any input sequence. We use X' denote the domain of all X,, i E [n]. Let MaxAff be any max-affine function on X' with MaxxA components, [[NMaxAf]]] as the max-afline partition generated by MaxAff as in Proposition 3.1. Let E' be the indicator of PMaXA as defined in Definition 3.2. Under the above definitions, there exists Linear layer and self-attention Attn whose attention matrix satisfies
and let 0 be any positive real number. We define PMaxAff {U:
SSotmax((WW Linear(X)) (VLLinea(X)))aa -X(X) ....
E(x,)]l I
with exception an arbitrarily small region. Here WK. Wq are the attention weights within Attn.
Proof. We first denote that according the premise of Section 3, the intersection region of dif- ferent affine components are omitted. This means for an arbitrarily small 8 >> 0, this proposition malfunctions on any points within a radius neighborhood of the intersecting lines of max-affine partitions.
Our proof consists of two parts:
1. Construct Linear and Attn.
2. Estimate the error between the attention score matrix of Alln o Linear and the target indicator.
For the max-afline function MaxAff, we denote it as follows.
Definition D.1 (Max-Afline Function). Let a, IRd b, E i MaxAA denote the coefficients of the affine components of MaxAff. In this definition, MaxAff writes out
MaxAff(Z)
max [E(NstxAAs)
{a Z b.).
(D.I)
for any E d
Remark D.I. For conciseness of presentation, we assume the top component of MaxAff exceeds the largest by fixed A 0, independent of the input and arbitrarily small.
16

---

Construction of Linear.
Without loss of generality, assume NMaxA 2 n. We construct Linear
(the layer of linear transformations) be
Ia
O MaxAA
Linear (2)
=
Onxd
TT Onx(NssneAR-n)) +
The the output of Linear(X) is
la
Odx NMtUAE MMaxAfff
Lincar(X)
X IIn
Opx(Nssanaann)) +
Onxa
(N
Oax Nasuaait
Onxn
O (NataNAA -n)
INMAAA
Odx(NstnrAt-
Onx ONoiuaazAHE-
(D.2)
0 NNMaAA Xn
NatuxAa-n
Construction of Atin. Since we only use the attention score matrix Softmax (KTA) TQ), we only have to construct the W and Wq matrices.
We construct them be follows
WK =
Oaxd
01
a2
M.N
0
bj
b2
la
O1xd
0
NatuxEE dd
Wq =
01xd
lixd
01x NytuAat
where a coefficient to control the precision the approximation. Specifically, as R increases, Soltmax is closer to maximum function, and the approximation is more precise.
In this construction, we now calculate the K and Q matrices of attention
K Wx Linear(X)
X
Oaxd
R
a1
02
ONNcaaA bN
Ln
O Onx( (NMMXAAI NM
(By (D.2))
0
b
b2
ONNsaaaa xnn
[a
=
R
b1
b2
and
Q = WoLinear(X))
X In
Oax (NNtsseAAa
la
O1xd
01x NotaAAE-d 01x NxLOAAA d
On (NNNooxAAA NNMaXAII- -
(D.2)
01xd
lixd
D (Nstaan x
X La
Oax NstanA
lixd
Nstnaa
OdxNscaeAR
Lixd
0 OIxNsAAAAR
Calculation of Softmax K T Q)-
We now calculate the attention score matrix as
Softmax K QQ)
Softmax
R
a2
X
Odx NNtuRA
b
b2
Lixa
01xNstuxAR NstaxA
by
02
ba
X
Oax Nxtuaa O1x NaxaaaA
Softmax
R
lixd
Nstaxa
bN
17

---

a 01 + b
a In b1
x1x (NNtana d) (NNNSssAA
a2 11 b2
0g In ba
Softmax
R
NMonA + b UNStAAA
Natuna
+ b UNSsAAA
01 01x(NasaaaH-d)
Estimation of Approximation Error. For i [m], we have
a 24 + b1
a2 22 + b2
Softmax (KTQ),, -Softmax
R
a NNMaAAa
It + NN
exp Ra x; + Rbb) exp Rag H; Rlby)
En= NM exp( Raj a4 Rbn)
exp Ra Na By R RbNSSSAAn)
This yields the entry on the k--th row of Softmax KT Q. to be
exp(Ra 24 + Rbx)
Softmax (KTQ)s:
NA exp Ra, + Rb,,)
m=
When aix, by is the maximal affine component and axx + bw is the second largest, we have
E [NMas exp Ray x; + Rbn))
Softmax (K (KTQ)kt
EnEt MMssAI exp [a,] + Rbb,))
2 Σ ne [NuaAAr) exp RR,, x; + Rb))))
Ene
NMaxAll exp [aj xi + Rbn)
exp( RuT + Rbw)
2 1 (NMaxAff
exp (Ra +Rbx)
NMaxAf -
=1 exp( Ra H,
Rby (Raf. + Rbw))
2 NMaxAA -1
exp(RA)
Thus when
R > A (In(NMaxAff 1) In Ine),
we have
NMaxAR -1
Se,
exp(RA)
which means
Softmax K" Qki 2 21-6
(D.3)
Moreover, since the sum of all entries in Softmax K Q., is have
Softmax (K )) n. < 1 - Softmax K Qk < 1 - (1 - c) =c,
h #
(D.4)
(D.3) and (D.3) are equivalent to
/Sottmax K Qh 11/00 Se
18

---

Softmax K Qn - s, h h k.
This yields
||oftmmax (KT Q) - - [xx 6
Thus, by the nature of ||- 1l.0
|Softmax (KTQ),: E(Xa), ... E(Xn))llx Sc.
construct Wo to discard Softmax KJ KTQ I: in Softmax KJ KTQ:
In
((NaaaAa n)xn
Thus
|Softmaxx (KTQ) Wo [E(()),, E(x2), .M (x,)ll
|Softmax
(K') Q) imm [E(X)), E(x2), =lx,))l
S c.
This completes the proof.
D.3 Proof of Proposition 3.3
Proposition D.3 (Proposition 3.3 Restated: Attention Reassigns Value to Max-Affine Partition). Following the notation in Proposition 3.2, let F IRd - IR be a piece-wise constant function which is separately constant on each U4, E [AMaxAtr] We show that for any 0, there exists an self-attention Attn such that
[Xttn(X)) [F())). .... F(X.)ll Se,
for every X in X' with exception of region of arbitrarily small Lebesgue measure in R"
Proof. Let Linear and the WK, WQ and Wo matrices be the same as in Appendix D.2. Then by Appendix D.2, we have
[|oottmxx (K Wo -
F(X,))]lo S c0.
for any 60 > 0.
Let V; denote the value of F on U,
Construction of Wv.
We construct Wy to be
Wv = [1xx
Vi 1/
VNusvaul
Thus V equals to
V : Wy Lincaar(())
X
0 (N
= 10xd
Vi V/2 ... VNSemaat
In
O (Nx(NstaaAttn))
O(NeenAa n)xn
IN
= [V1
1/2
Thus we have
|oottmax (K (TQ) Wo F(x),. F(Xn)ll
= [ V22 VN Softmax (K' Q) Wo F(X1), F(X2), ... F(Xa)]l1x
[[
V2
VNMAAA Softmax (K' T Q) Wo Vi V/
VNuamAm E(X,),)(()), ... Elx
< 1o
Let VVllo050 I e yields the final result. This completes the proof.
19

---

E Proof of Results in Section 4
E.1 Proof of Theorem 4.1
In this section we give the proofs of our universal approximation theorems of self-attention. We first prove the Loe version whose target function are continuous. Then combine this result with the well known Lusin's s theorem and extend our result to Lebesgue integrable functions terms of Ly norm.
Theorem E.1 (Theorem 4.1 Restated: - Universal Approximation of Self-Attention). Let f: RRdx= + __ denote any continuous function on compact domain U C R and let be any positive real number. Then, there exists a self-attention Attn with a prepended Linear layer, such that
| Attn Linear]
Proof Sketch. Our proof consists of four conceptual steps.
Step 1: Partition Input Domain U via MaxAff.
Flattening Input. Each input E Rdxx is reshaped into a single vector ğ E IRdd by stacking its rows or columns. This unifies the domain as E D-
Grid / Max-Affine Construction. Since f is uniformly continuous on the compact set U, choose such that
[ 72 <
- ((2))) <.
We subdivide [-D, D]dn into cubes of side 8, yielding = =pp grid centers 10 We treat MaxAff piecewise (max-)affine piccewise-constant partition: for each Z, there's s a nearest v within 8/2.
Step 2: Configure Linear and Attn to Imitate MaxAff over U.
Map Linear. Design Lincar : Rdxxn R (for some dimension M) to capture the dot products (vy, Z). Essentially, Linear(Z) arranges these {u] Z} in a form accessible to attention. This ensures each grid center vy can be individually "queried."
Encoding Affine Components. Observe that (13, ž) [1 is akin max- affine function. We store terms vf Z, plus - 11, into K and Q for later use in Softmax(K Q).
Step 3: Enginner Attn to Generate an Indicator of Which Partition Cell the Input Belongs To.
Construct Q. the self-attention block, let KTQ ss R((,,, 2) - 11,,, where > 0 is large. This makes Softmax K KTQ) favor the row j maximizing
(v,, 2) 111y,11
Near-One-Hot Distribution. Hence the -*tth row obtains probability close to 1, effectively identifying which grid center Vy- is nearest to Z. We interpret this as near-one-hot "indicator" vector for the correct partition cell.
Step 4: Map the Indicator to the Target Value f(Z).
20

---

Assigning Values. We place f(0}) in the "value matrix" Wv. so that once row j is selected, the attention output is f f(0-)) ). Since Z is within 6/2 of Dig, uniform continuity implies
If( f(öz-) )1l Σ e, (for suitably chosen 6).
Final Reshaping (If Needed). A small linear projection M can reshape the output back Rdxn, The essential logic is that the correct (ūū) is "routed" to the final output via the near-one-hot attention distribution.
Thus, single-head attention block with minimal linear layer can approximate any continuous function on the domain. This completes the proof.
Proof. We divide our proof into two parts:
Part Construction of Attn and Linear. We construct Attn and Linear accordance with the steps shown in the proof sketch, and calculate the precise output of our construction.
Part 2:
Estimation of Approximation Error between Attn Linear and f. We calculate the difference between the output calculated in previous part and the target function to
Part 1: Construction of Attn and Linear.
We first construct the grid points in [-D, Djdn used the construction Linear and Atln.
These grid points are used to construct the max-affine partition. Specifically, the max-affine partition we use is rrrid-arrittion and these points are the center points of these grids.
Construction of Grid Centers in D)D, Let Z [21, 22. 2nd R denote the input to Linear. Define Z = [ 2 ..... IT
, E N parameter that controls the size of the attention block and the error of approximation.
Definition E.1 (Grid Centers in Di Define Uk. A2... kan E Ra as
2Dk1 DD 2Dk2 DP
2Dkan DP
Uki Aar.. Aen
for k; 20, P 1), i [dn].
Remark E.1 (Scalar-Labeled Grid Centers). For each multi-index (k1,..., kan) with e {O...., P - 1}. we define
kg
pi-1
{0.
pdu 1}.
This base- P expansion gives one-to-one map between the tuple and the scalar. This notation allows us to define another representation of the grid center:
U Uki......
For every E V, we define
U1dd Ud4-1:2d ... U(n l)ndland
dxn
We now construct functions and T. They are linear functions f RRdx - Rdxn playing erucial roles in the constructions of Wx and Wo in Attn(-).
21

---

Construction of E and T. We first show that f is bounded. Because f is continuous within closed region, its output value is bounded co-norm. Let Bn denote this bound
Bg = |
We now construct functions E(), T() related Their sum is constant while their subtraction is scaled f. For any Z E R2x= we define
f(Z)
E(Z) 1dxn
(E.1)
Bo
f(2)
T(2) = Laxn +
(E.2)
Bo
and
(E+ E(Z)+T(2), (E - ()) == T(Z).
By the definition of E()) and (.), we have
(E + (Z) 2dxn
(E.3)
21(7)
(E - ()) Bu
(E.4)
for any Z Rdd
Construction of the Layer of Sum of Linear Transformations. layer to be
We now construct the Linear
(n-1) I)
Linear (2near(2)= =
Σ
Ie
k+1
(n)
el
(2dG+1)
Σ
@j+41 (j++dd++1)
(24G) (24G)
+
01x2aG
12aG
R=0
dx1
(cta-da41 shifting the (2dG) down dG rows.) 34441
where == pdn.
This layer multiplies the flattened input with the grid centers in Definition E. and append 2dG- dimensional identity matrix below the matrix containing these multiplications.
We now express the output of Linear in simpler form in the following discussion. First, we show that
(n- 1)
Ze
(1, kdd
Σ
2 (2s)
0=
k=0
retrieve the (Ic + 1) token
= , 113
(By Z being the flattened input)
= e 2 E R, j 10.1,2.-... .G GG 1}.
This yields
G-1
Linear(Z)= =
Σ
7 Σ
(2dG)
(+++dd++) )
(2dG)
(2dG+1)
01x2aG
lIAG
1=0
Xo
Xo
Lag
OaGxdG
(E.5)
OsGxaG
LaG
Explicitly, the last line is by
G-1
2
(+s+1)
Xoo
==
s=0
22

---

which implies
G-1
d-1
Σ j=0
0/j
2
+0+++do+1) "=
(2dG)
=(x Xol.
Here
Xo = 1lxxd la Zlixd
0/ Zlixd
u1-12l1xd
To summarize, in the output of the first layer of linear transformations, the first row consists of linear transformations of the flattened input, while the other rows are together identity matrix (l2ac).
Construction of K and Q Matrices. We now construct the Wi and Wo matrices in the self- attention block and calculate the output of Softmax (T)...
We define Wx follows
0 Ie 113 2 (T(T
0
1 lix
1
I 111
WK == 0
2
1xx
2
1ix
2 In
1xx
Op,
(n()
(n))(Tn))
E (E
The definition of WK yields that
K = Wx Linear(2 ())
0
Ioall2 1ixd
I| 12
I l2 (T(((Ta-1))
Iolll 1 1ixd
l1 11xd In In(E(()))
| | 1 1 1xd
0
lixd
1ixd
On
In (TVo))
(n()(
In
(n))())))
In In(E((G-1)
No
Xa OaGxaG
TaG
OaGxaG
Lag
to Ālixd
ut Z1ixd
of Zlix
1o Zl1xd
o. Z1xxd 1 lixd In (n(F)(n))
_o_ Elixd
1 lixd
_|ealli lixd
1Pa- lixd
1 f 1ixd (n(()(o))
llvo 1ixd
In (T(o)))))
In (ni)
In(T((G-1))
In (n(((GG-1))
(By (E.5)
where the last line follows from Xo being multiplied by and thus appearing in the first row of the output.
Next, we construct Wq to be
0
Rlixn
0 (2dGGnn)
Wi =
0 Rlixn
(2dGGnn))
On
In
Onx
This yields that
Q Wo WVLineaa(2 2)
0
1lixn
O 01x(2dG-n)
Xo
Xo OaGxaG TaG
Rlixn
O (2dGGnn))
Lag
On
T
Onx(2dG-n)
Oacxaa
Rlixn
0 01x(2d(-n))
= Rlixn
01x(2dG-n)
In
0 (pd(2d))))
We now calculate the attention matrix Softmax (K (TQ)..
23

---

Calculation of Softmax (TQ).. First, K Q writes out as
00 La
la
In(T(T))) In(T(Ti))
01 Zla
V Z1d
| la I 2 |ella la
In(T(va-1))
Rl1xn 01x(2dG-n) m)
K KTQ-
Rlxxn
(1x(2dG-))
0/ 21,
la
In(E())))
In [n)
0j 21d
In(E(Ti))
| 1a
In In(F(ro-1))]
R(v 2
|| 1
dxn +In)T(o)))
(2dGGnn)
2
R(v 2
13
Idxn ln(T(()))
0 (2dG2i))
R(vo-17
||00 l12
1axn + In In((((a-1)) (1)
0 (2dG-n))
(E.6)
R(vo 7
|
1dx ln(E(
(2xG2d))
2
2 R(v
|
1dxn
(n()
(2dG-n))
R(v6_17 1l 11 Taxa In(F((G-1)) (2xG-n))))
where the last line follows from the multiplication of block matrices. This multiplication between KT and Q is equivalent to first multiplying the first columns K" with and then broadcasting their sum to the first columns, and then adding the result with T and E related blocks. Columns are all filled with except for the first n columns.
24

---

Remark E.2 (Interpretation of K' Q). The non-zero entries of K" Q an aggregation of two matrices
R(u Z
1lo dxn
R(v 2
|
1dxn
1 1
1axx
2
(E.7)
R(v 2
ddxn
R(v] 7
) dxn
R(VC-12 ,2
1 112 Ildxn
and
IT((()))
(n(((()))
In(T(Cc-1) In ((o) In (E)
(E.8)
In In(B((c-1)) I))
In these two matrices, (F.7) is identical between columns and has the precision coefficient R free of our choice. In later discussions, we set R to be sufficiently large that the Softmax approximates a maximum function, and "selects" the of the maximal R(v; Ž |0, 1dxn for i {0, 1, ... G 1}. By "select" mean only the entries with the selected label has value not close to 0 in each column of Softmaxx K Q).
(E.8) does not include related terms. Thus when set to be sufficiently large in our later discussions, (E.8) does not affect the selection made by (E.7).
If we exclude the (E.8) in the attention score matrix Softmax [K (TTTTT Q), the output approximates matrix whose columns are all-zero except for two sub-vecto equal to 1/2d 14. This writes out as (here we only show the first n non-constant columns)
(s 1)
Laxn
0 ()e)ddxn
(E.9)
0 (s 1)dxn
Laxn
0 ())ddd
for any E [G]. The addition of (E.8) change the 1 in (F.9) to
0 1)dxn
2 T (7.-) 1)
0 GG- 2)dxn
(E.10)
0 1)dxx
24 (V.-1)
(G-e))dxn
In later discussion, we use transform (E. to T(ū.-1) obtain the final output.
22(6))))2dB to
Now, we divide the calculation of Softmax (KT Q) into two parts: the calculation of exp(K" [K Q) and the calculation of the denominator of every column of Softmax (K' (TQ)). This denominator explicitly writes out as 2dG exp(K K (KTQ); for each E [2dG].
25

---

For exp(KTQQ) Q), by (E.6), we have
exp
R(o0 2
|| T(wo)
(2dG-GG)))
exp
R(v]
1
T(ū))
(2dGGG)))
exp
(
| 112
T(Ta-1)
ldx(2dG-n)
xxp(kTQ)=
(E.11)
exp
R(v0 2
[ull2
E((o)
(2dG
exp
R(v1 7
Iv
12
E(7i)
(2dGGG))))
2
exp
R(v [
Ilve ( mllir
E(UG-1) ldx(2dG-n)
For the denominator, we calculate it in columns. Let denote the column which we calculate the denominator in Softmax. When E {n + 1,n ., 2dG), there are 1 -2GG = 2dG columns. And when i [n]], we denote that
2dG
100-112 2
exp(K K (T)),
Σ
(lix (1g-1)as + E(t-((())):a)) exp R (o]_,
5=1
(By (E.1I)
1
Σ
+ixa(E+T)(o;--i)) exp [0]_
2
(lixa(2xxn)ic exp (n(1)_12 ||0
Σ
2
==1
Σ 2d exp (R (1/
| 12 ))
i [n].
(E.12)
2
Observing from (E.12), Σ j=1 exp(K KT Q),5 is invariant of i for i E [n]. In this case, we define
2nG
2dG
Σ exp R [u || 12 )) E R, i [n].
a(z) exp( (xp(KTQ)) Q)
j=1
From (I. and (E.12), we have
Softmax (K") Q)
ex 3 Σ
adG j-1 exp(K T Q),
12d6xn
2dG 12GGx(2GG-))
(By
invariant of e [m])
Sadd exp (xp(xTQ),
exp
A(v 7
|| T(ūo)
1dx
R(o exp
||
l
T(ū1)
L 12dG
2
exp [R(v]
| 11
((wa-1)
(2dG
2
2dm(7) 12dGxn 2da (2aG
exp [(v))
7
||
E(7o)
(2d
exp
R(v]
|all
E(Tr)
Lax
R ž | B(Ja-1)
exp
(2x(2a))))
26

---

exp
R(v0 2
!ngll2)
T(ūp)
t (2dGG-))
(22)
exp
R(v)
2
T(ūi)
T (2dG-n))
x(2)
:
exp
10 6
2
1:0 120-112)
((06-))
き (2dG--n)
((2)
(F.13)
2d
exp
1100
1ool
2
E((o)
t (2dGG-)))))
(22
2
1011
exp
E(01)
dx (2dG-n)
o12)
exp
R
7
|va-1113
E(ūc-1)
(ax(22a)
((Z)
Construction of Wy and Wo. We now construct the Wy matrix and calculate the V matrix of the self-attention.
We define Wv as:
Wy 10a X1
-X1lax(142ac))
where
X1 = [La la
Lalaxaa
a matrix formed by stacking G La matrix horizontally. With this definition, we compute V matrix as follows
V == Wy Linear(Z)
Xo Iac OaGxaa
Xo OaGx
00x
X1
-X1]
Lag
X
-X1].
(E.14)
After the construction and calculation of V, we on to construct Wo as:
d 022GGnnxxx
Wo
The sole purpose of Wo is to extract the non-zero entries of the final output.
Calculation of the Output of Attn Linear. We now compute the final output of the self-attention block
Attn Linear(Z)
27

---

R(v0 2
1voll
exp
T(wb)
t d2dGGGn)
((2)
2 1
exp
T(
G dxddGGn))
a(2)
:
11100 2
11ve 122)
exp
T(TG-1)
t dd
[1
((2)
Wo
2d
exp
11(1g
11vo
E(Tn)
G (2
a(2)
[[vu] 2
||u1
exp
E(ūi)
1 dx
a(2)
:
R(va 2
exp
E(0G-1)
(dx
a(2)
(By (F.14) and (F.13))
2 1vo12)
exp
Riog
(T(06) E(ōo))
Oax (2dG-n)
(2
[t(e] 10113
exp
1
X1
a(2)
(T(0i) E(Gi)
(2dG-n))
Wo
2d
exp B(v 7 Ive-1l1
(T(Wc-1) -((ic-))-((io-i)) Oax (2dG-n))
a(Z)
(Sum of X1 multiplied by T related blocks and X1 multiplied by E related ones)
RRup 2 1
20(r))
(2dGGnn)
o(1)
Bo
exp R(u] 2 I 12
2f(1))
X1
a(2)
Bo
0 (2dG-nn)))
Wo.
(By (E.3))
2d
A(v 2
IvG 113
2f(wa 1)
(2dGGn))))
(
Pa
Let Ia denote the d-dimensional identity matrix. We have
exp
R(v 2
|
2 (el)) Bo
o(Z)
7 |1 1
exp
R(v]
2 2 Bo
X1
a(%)
exp
R(v
2
2
2
a(Z)
Ba
28

---

exp
Riv0 2
15012
2f(00)
((2)
Bu
exp
Rie]
2
10g12
2f(01)
=Ia
la
La] dxag"
a(2)
Bu
fl(vo 2
lIvG-
exp
2f(v
a(2)
Bo
Eqqivalent summing blocks in S)
G-1
Σ
La"
exp
R(U)
T 2
111-111
2
2f(0))
a(Z)
Bo
j=0
G-1
||,ll 2f(2,)
Σ
R(u -u,T2 7
exp
aZ
Bo
jm0
This yields
R(vj 2- N1912
Attn o Linear(2) Z)
G-1
exp
2f(2g)
Wo
=
a(2)
Bo
Odx
Zj=0
G-1
exp
B(v
a(Z)
2
11:12
2f(w))
Bb
0dx(2dG-n)
( (2)mm
d.o
1
exp
R(U 2
(1,)118) f(ū,). 2
(E.15)
a(
5=0
Part 2: Estimation of the Approximation Error between Atin o Linear and
With above calculations of the output of Attn o Linear, we now demonstrate how this output approximates our target function.
Essentially, demonstrate that each term in the summation of (E.15), given by
1
a(z) exp R(v
Ž 12y)))
approximates max-affine indicator as R becomes sufficiently large. They are each multiplied with f(ū,), which is the value of the target function at the center point the indicated region.
Definition E.2 (Max-Affine Function on 7). Let Aff, E Rdda + R with j E (0, ..-1,,......11 -G denote group of affine functions defined as:
Afl ; (2) = 11 7 1
jE ...
Then let MaxAff E gd + R denote max affine function whose affine components are (Aff,)
j E 1,,1,2..... G 1}}. Explicitly defined as:
MaxAff(2) JE(0.122.- G- { (2f
Because the target function f is continuous function on closed domain, the function f is uniformly continuous. Thus for there exists a 8 > 0 such that for any Z1, Zz, long 1 Z2lloo < 8,
we have []]( [(2)) I e/33.
According to this 6, we divide the affine components of MaxAf into three parts:
1. The maximal component, which has the smallest label jm.
29

---

2. All affine components that match the maximal component or fall within ò of it (Jo as defined below).
3. The remaining , for to, 13 (J1 defined below).
We write out the labels of these groups of components as follows
Im ==
(0,1, G-11 (Af (2) = (2)).
min
Jo = (j MaxAff (2) (2) S 8).
J1 = (i |MxXff( 2) Aff, > 6}.
For any pair of i,j E {0, 1, .. ,, G - 1}. we have
|0: 2
( Aff (Z) = 11 7
11
1
7
2
=-12/13
|1:
( 11
10
+ 0/ oT7 7
+
7
2
2
2
2
0s 1
+ 2
Thus for Jm. have
1 im
+
11112
03 11 AlT (2) ATT, 20, {0 (0.1,......11
This yields
1 - 1jm 12 - 1 - -,
for all j E 10. G
This denotes Jm is also the label of the closest es to 7 among all n, i E (0,1,..,, (G-1 1}. Thus we have
11 Ž112
ief0,1,-- GG1) flles
min
2112).
(E.16)
Now, we prove u. (the grid point nearest Z) has distance to Z smaller than half of the grid width (e.g.. D/g) in infinite norm.
Let D 2D/g 1,),0,1)d denote differences Vim. from the set of all v, (i (0, 10,,,..,, - 1}) neighboring "j..- For any A in D, from (F) we have
|| 3 < |0jm + A 21
This yields
2A T 7im S ||
This means that, for any k E [dn], by selecting A to be +2D ge we have:
(dn)
2D
4D
+2
7 )/m)) = 21 T (ž Vsm) - || 1
g
g²
Thus we have
( 13 12 <
D
k [dn].
g
which implics
D
11 2im 11,0 <
g
30

---

Set g to be larger than 2D/d; we have
8 2
12 U/ 2/mlls 11:00 S
thus
[1(()) - f(0mm)ll S 3
(E.17)
where the inequality holds by 8/2 < 6.
Calculation of |Atnn o Linear - flooo+ (E.15) and target function
We now calculate the difference between the output in
Attn o Lincar(Z) - =[(
G-1
exp
R(u 2
11,,112
2
((5,) 1(
a(Z)
1=0
G-1
exp
R(v] 2
exp Aii, 2- 11
(f(t,) - f(2))1I
(By G-1
a(Z)
a(2)
==0
exp R(v] 7 11)g13)
G-1
(1((T)) f(2)
(By property of infinite norm)
a(Z)
1=0
exp
[R(U] jm
Z 1 12
[l((T)m) f(2)ll
a(Z)
exp [Ro,] 2 l
(5(T,) 1(2)l1
a(Z)
JEJp
exp [R(0] T 2
f
(E.18)
a(Z)
JEЛ
We now calculate each part in (E.18).
As previously stated, for any Z1. Zz, as long as 1| 12 S 8, we have - S e/3.
Thus when we designate Z1 Uy for any j E Jo and Zg U. along with (E.17) we have
exp [(e] Ž 1|g,11
(17))-((2)11)
a(Z)
jeJo
[ivo] 7
exp
(11s(T,) - f(e 1llc + IIf((T,..) l(z)
a(Z)
jeJo
Riv 2
||,
exp
G
a(Z)
jEJo
R(v 7
exp
26
(E.19)
a(Z)
3
jEJo
For Jm. we have
exp [Rlu] 7 7- 111 13)
exp [R(u, 2 1 u 112)
[
f(2)]l -
(E,20)
a(2)
alZ)
3
31

---

When Rss larger than
8 In
8 BuGe),
we have
362
2
R(u 7
exp
2
[1f(0T)
f(2))lo0
a(Z)
jeJ
R(1y 7
[1g1112
exp
2
BBo
(By 1ffle Bo)
a(Z)
JEJI
Σ Ejen, exp
R(U1]
2
1
2
< B0
a(7)
Ejen exp R(U1] 2
|1v,
2Bo
exp CRp(R()T.
[R(o] T
Ž 11 113
jm
(a(Z) is the sum of all exp R(u3 7 [w ) thus larger than exp [R 7 11 12
= =20
Σ
exp
212m
- ||0 -212)
JEJ
R
8 )2
82) 8²))
$ 2BBol|J[ll exp
(2)
-3R82
22ooGexp
8
In 2 BgGe)
332.
$ exp
382
(By R2 8 In BoGG))
3 =3'
(E.21)
Combining (E.19) and (E.20) yields
R(v 7
|
exp
11((T)) f(2)ll
a(Z)
jeJau(im
exp
[(v,] 2
[ug113
2
exp [(v)) 7
|1 11
(By (E. and (E.20)
a(Z)
3
a(Z)
3
jeJo
exp [[t %
2
a(2)
2e
(B.22)
3
where the last line is by ΣSelu a(Z) exp [Ro3 - )111;11) $ 1.
1
We plug (E.21) and (E.22) to 18) and get
R(v Ž I1 12
exp
- -[(()] < | Attn o Linear (Z)-
jm
[1((0).)) -
a(2)
exp [(v] 7 1
I1((o)) f(2)
a(Z)
jEJo
exp R(v Ž 11,,
[1(T,)) f(2)
a(Z)
jEJ,
32

---

26 < +
6.
This completes the proof.
We also extend this -Nrm result we just proved to Ly-Norm.
Corollary E.1.1 (Ly-Norm Universal Approximation). Let f : Rdxn + Rdx= denote any Lebesgue integrable function on a compact domain Igdxn and let > 0 any positive real number. Then, there exists self-attention Attn prepended with a Linear layer such that
f Attn o Linear
6.
Proof Sketch. The same partition- based construction applies almost everywhere; outside negligible set, f is continuous (Lusin's ss theorem). Thus the Loc. argument extends.
Proof. Since f is Lebesgue integrable on compact set, f is bounded almost every where. Let Bp denote the bound of Il
By Lusin's theorem, for f compact domain U. there exists continuous function which is equal to f U except for a region Dj such that su(s)) & A. This can be written as
D3 = (Z) # g()), m(Ds) A.
(E,23)
(E.24)
Here stands for the Lebesgue measure of a set.
By Theorem 4.1, there exists net work Attn Linear, consists of self-attention Attn and layer of sum linear transformation Linear such that
ALtn Linear - gllr. 60,
for any 60 >
This denote that for any Z
|Atta Lincar(Z) (Z) 9(Z)llp I (dn. e2) ca(dn)
Combine this with (E.23) and (E.24), we get
# Atttt Linear(2) (2) - g(2) > co3) $ ((U(()) + (g(2))
1 A
(E.25)
since that f(Z) =(),, ||Attn Lincar(Z) - - = ||Attn o Linear (2) (Z)-f(Z) s co.
This yields
- =1ʃ ZEU
Ilf - Attn o Linear
|l Attn a Lincar|p dz)
ZEUUDS |lf Attn Lincar| dx +
||f Attn Lincar|[ī dz)#
IZDDs
=I llg - Attn o Linear|, dr + / ZEDs
I Attn o Linear || dæ)
(d) + A B s so(dn)(U)) + ABB,. Bp-
(By (E.25)
Set
60 5
(d)
e
A <
Bp 2p
33

---

We have
|lf Attn Linearlla, I co(dny(UU)) + A Bp
e
s (dn
2(dn)(U))) + Bp
} Bp
This completes our proof.
E.2 Proof of Theorem 4.2
Theorem E.2. Let UK C Rdxx and UQ C Rdxa be two compact domains, and let f: UK x UQ + Rdx= be any continuous function that takes input from both domains. We use Zx. Zq E Rdxn to denote the two inputs of from UK and UQ respectively. Without loss of generality, suppose both input domains be [-D, D]ax", where D R., Then for any 0, there exists single-head cross-attention Attn and two layers of sum of linear transformations, Linearx and Linear such
that:
Attn (Linearx (Zx), Linearq(Za)) ( (2)x sc.
for any Zk. Zo E DjD,
Proof. Without loss of generality, assume UK = UQ = (-D, for a D E R4.
Construction of Grid Centers in UK, Same in Appendix E.I, define Ž [&] 22 ..... 2 IT of our approximation. Define Uk, ka. Jaan E IR to be
P N4 is parameter that controls the size of the attention block and the eπror
2D - DP 2Dk2 DP
2Dh DP
hs E 10,1,2, -.. P-1), i [dn].
.2.... Jadn
P
P
P
Let V = (D),, kg... ka."" Jan [R: E 1,, ...
....P-11
ie [dn]} be the set of all K2... han We also
define another way to refer to a vector V, denoted as
Eih kPP A2. Jan
Please see Remark E. for the reason for the feasibility such expression.
Following the notation in Appendix 1. for every E V, we define
[Uidds Ud41:2d: .. Vnn 1)d+
dxn
d matrix-form representation v.
Construction of Related Function and T. The continuity of ʃ within a closed region
guarantees be bounded in oo-norm. Let Bg denote this bound. For any ag, aq E Rdxa define
Elak, aQ) = Idxn flax,a0)
Bo
T(ak, 02) = 1dxn + ((axaa))
Bo
We define (E + [ax aq) Elak, + T(ax, aQ). By the definition of E and T. (E (ax, a2) 2axn n for any ak, aq _ddn.
34

---

Remark E.3 (Intuition behind E and T). E and T are constructed to satisfy conditions:
Tz... Zo) + E KZ Za) = 2.
T(z ZQ)
(ZK Zo) 2f (Z Zo
E.EE. for any input.
The first condition is used to configure the denominator in the Softmax expression of attention to a constant value. The second condition is used to form the value of f in the
For simplicity, same as in Appendix 1, define
G= pdn.
We now construct the Lincark and Lincarq layers be
(n-1) 11
dG
Linear x x<(x)
=
Σ
((
(u5)
edG+++++1
12dG2+
Σ
(2409) + (5+0+dG441)
==
k=0
I200
+
DGX2dG3
6-1
(n-1)
Linearg(Zo) = Σ
(n) (w))) (Zee+t)
(n+C) 4t+
[lixa 02x(2ac2-n))
jm0
Oaxn
02
+ In
(nd
Same that in Theorem E.I, we have
(2ee(t)) (vi)aaaaaa = 2k, ZK,
(E.26)
Σ
= Zo,
(E.27)
for 20.1,
We now calculate the output of Linearg and Linearq- For Lineark, we have
dG-1
Linear x (Zk)
Σ 1 Zx 22G2
+G)
Σ
(2d02)
+e (2dc9)
(+o+00441)
OGx2doa
I2aGe
12dG2
G-1 T. Zx Edc-1 Σ
(
(2d0*) (jttt (ftttd0++1) )
1240%)
(by (E.26)
=== 11
Iaga
04G2xdGa Iace
0acr xdC²
0 1 7x DdC-1
(2AG2)
E 1/ 7x
Edc-1
(2dG2)
[+++]
Laoz
Oa6² xd
OacaxaGa XK
Ldc²
(E.28)
XK
Linearg (Zq) Σ 0 Zoeeftte (n+G) [lixa (1x(2ac=_n)
00xx 0Gx(2dG2-n)
+
I
(xx(2dG=nn)
35

---

Σi= oj Zgez+1 XQ 06x(2d02-n)
(2)02+0)
11x
01 (2dG22
(by (E.27))
0 (2x((2)==n)
(E.29)
In
(2x(2aGe-n))
in which Xx and XQ defined as
Xk =
10 k 1 IxdG V1 7 lxxG
0/ 15xlixxl
10-17x 7x 1xda
1xdG²
Un 1xl
1 Zo 1xx
XQ =
122 Zol 1xn
_ 1 Zo 1ix
Gxn
We now construct the We and Wq matrices in the self-attention block and calculate the output of Softmax (K Q).
In the following, we define Wx parts. First, we present as block matrix
[0
xdG
0
xxaGa
1
Wo
Wo
0
WK
(E.30)
Wi
Wi
0
WT
WE
0
then define the submatrices (E.30) as follows
Wo
|
lixa + Wo
11
1xda + Wo
Iell
lixa + W
1xd + W
2
2
W7 ==
Wr(0)
w
W(G-1)
Wy
T
WE =
W-(0)
w
W(C-1)
E
E
E
W1 = [W
W
W1
W
WilGxaGe
in which
Wo
11
1xd
1xd
lixd
l 11
1xd
2
2
2
w T
[In(T ())
(T)
j
dxGa
W = (o))
In(E(T,, (() T((((,,
j E (0,1,2, ",
dxGn
W. =
Re (G)
lixd
Rez
(G)
lixd
(0) lixd
Gxd
The definition of Wx yields that
K = WxLinear (Zk)
0
O Oxxaaa
IaGa
OaG xGG I Jaga X k
Wo
Wa
Odc xda Xx
(By (F.28)
W
W.
WT
WE
Xx Xx
Wo
Wo
W
Wi
WT
WE.
36

---

Next, we construct the We matrix as
01xG
Rlixn
OxG la
Rl1xn
Wq=
01xn
01x
In
In this definition, the Q matrix in attention can be calculated as follows
Q =
WoLinearq @ (Zq)
T1xG
11
01xG
Rl
XQ 02x
(By (E.29)
TG
01xn
I (nx(2dG2-n)
0Ix
In
Rlixx
(2x(2aC2-n)
Rl1xn
0 (2x22a2=_n)
Xe
02dG2=
In
Onx (2dG m))
Now caleulate the attention matrix Softmax (K (KQQ Q).
KJ Q can be calculated as follows
Xx
X
Rlixa
01x(2d02-1)
KTQ=
Wo
Wo
Rlixn
01x(2dG=-n)
W1
W1
XQ
0Gx (2dC2-n))
WT
WE
In
Onx (2dG
[RX) K + RW 111x + W XQ + W.T
024Px(24GP-n)
(RX) + RW }[1xn + W XQ + W
O2d2x
The W,T Xq in the expression of K Q matrix is further calculated as
WiT XQ W, W,
W. Wilexaor Xq
W XQ W XQ
WT XQ_ xG*xG
We define Q : WT XQ. then W,T Xo can be denoted as stacking this block vertically for G times. In this definition, Q matrix can be expressed as
Q = W.T. XQ
vi Zq 7olix lxn
v Z Zqlixn
-[Re Re
(R)
Rez
(G) lixd
Reg lixd
(G)
2y ZQlixn
L Zg 1 1xm
Re (G) ld
2/ Zolixn
Rez (G)
1 lix
2/2 1ql
RE
(G) ld
T Zo 1 1xnd
37

---

Reg Zq ldxn Rvj Zo Laxn Rv2 Zo 1dn
(E.31)
Rog- 2e,22o11000 1dx
The calculation of Softmax (K') Q) can be disassembled into two parts, the numerator exp(Softmax (KT))) in the expression of Softmax and the denominator of every column of Softmax (T)),, as in the expression of Softmax, explicitly written out as Σ)= exp (K" (KTQ)) for
each i E [2dG].
We calculate exp [ (KQ) Q) as follows
exp(KTQ)
Texp( (RX + RW, Lixn + RW, X4) o exp (WT)
Lag= xadex
expl (RX K + RWo 11ıx + RW, X2) o exp (W)
E Lag x (2dG3 mn)
(E.32)
For the denominator, we calculate it in columns. Let denote the column which we calculate the denominator Softmax. When (n + 1, + ...... 2dG ), the i-th column has in every entry. Thus the sum of all entries in this column equals to 1 2dG 2dG.
And when i [n], we have
2dC
Σ
exp(K T Q)&
j=1
(lix -1/ + Lixab(U) -1).))
10
11/22110 11
exp
T 2x R
+ -
2
2
G
G
1)28 exp (
|
|1 11
Σ
lixa(E
((1)=_ 1/1-1
-1 Zo
2
G
G
(1ixa(20xm)as) - exp ( R
11
11
1112-1112 112
Σ
Σ
T
1/
+ Ujj
2
2
11
G
G
11/21112 1 ||1
10j 132-1113 12 ))
Σ
2d
exp
2
i [n]. (E.33)
2
31
We observe from (E.33), that Σ exp(KTog Q) is invariant of i for i E [n]. In this case, we define
2AG²
1
alZ,, 2o)
exp
K TQ))
2d
j=l
| 1 + 0TT
|| 10
112
Σ
Σ
exp
R
Zx
Zo
2
12
denote the 1/2d of this value invariant of i for simplicity. Because
1 2dGa
alZ,. Zo)
Σ
exp K (QTT22,.)) [))
2d
38

---

from (E.32) and (E.33) we have
Softmax (K") Q)
(xp(KTQ)
2U12 exp(KTQg 1)1,
12dGxn
20G 22dxx(2dG-1)
jm1
nominator Softmax
denominatoe of By
is invariant of for E [ri]
exp (T))
exp( K Q) ④
Edea ZK 20) 1dd 20G 22Gx(2
Texp(( [RX]
11xa + RWT Xo) exp [ R
Lacs x (2d02 aa (2dG3-
exp (R.))
RWO
1xx + RWT XQ) o exp ( E
K
2/017x %
2d
24C x2dG
(By (E.32))
denominator Softmax
2d 7q)
exp exp (RX + RW, 1x + RW1 XQ) 3 exp (WT) exp (RX + RW, 111xn + RW1 XQ) 3 exp (WE)
207 x
20G x x(20G
2da 2x 2q)
Now we' Wy matrix and calculate the result of multiplying V Wy Linear (2x) to the attention score matrix.
we've ve defined and calculated the attention score matrix Softmax (K Q). we go to construct
We define Wv as
Wv
[X2
-x₂
Odl
dx(2G+++1)
where
X₂ = Ia Is
La]
dxdGa
This yields the V matrix to be
V WyLinnar (Zx)
Iaga
Oaga xdG"
X2
-X2
O.]
x
Laga
XK
XK
Ax (x(G2+++1)
(2dG++1)x2de)
Taga
DaG=xaGa
X2
X2]
(since Xx multiplied by 0)
Oace xdG=
Taca
2x2
2dG x2dG2
[X2
-X2]
With V. compute the output of "Sotmma(K" T TQ) follows
V Softmax (K Q))
2dmí 7x 1g 7g) exp (R + RW, 11xx + RW, XQ) o exp(W77) Xo) @ exp(WWT)
2d 1ass ((203=
=[2
-X2]
2de 7x %00 exp RX + RW, 1ixa @
शाल idps (2dG2
=X²
exp (RXT + RW, Llxx + RW, XQ) 2 exp(WWF)
शे LaG3x
2drr(Txx 7a)
X2
2da(2x.20 exp(RXX. + RW, lix O cxp
RR, cxp XQ) o exp (WT) 2ng2 1a0 (200x
-x in (X2 -Xa]
39

---

1 2da ZQ)
X2 [exp((XXX + RW, 11xx + RW1 XQ) o [exp(WF) - -xp Oaga 02cRx(2ac2-n))
(E.34)
To further calculate V "Sottma((K'' (Q), we now calculate the result of its non-trivial part (the part beside Oag x (2dG2 -x))
X2 exp (RXX + RW; 11xx + RWT XQ) @ [exp(WF7) - exp(WE)]]
(E.35)
We now calculate each part in (E.35)
(EXp(WE) -(exp ([w20) w
exp(W.T)
W7 (
exp ([w(
w
W w(G-1))
exp (W(())
(())
exp
exp (w()) T
W
exp
(E.36)
exp (wrG-11)² exp
W
W(G-1))
In (E.36), we have
exp(In(E(in) m))) exp(In(T(e. vo))) (T)))))
exp(In(T(o,, 0i)) exp(In( E(U,, oo))
exp (w/)) -exp (
(ep(w!?)
exp((In((ū;, va- ()) exp(In 00
T(ü un) - E(U, vo) T(U v1) E(ū;, vo)
T(7,, vG-1)
E(U, vo)
2f(0),0) Ho 2f(":.1)
Bo
Ba
Thus (E.36) is equal to
2f 1,00) Bo
2f(2,-1
exp (
exp(WW (W
Bo
i [G].
(E.37)
(1-1)G+ING:
2f 1.0--1.00-0 vG-11) Bo
We also calculate the other part X + RW 1lx + RWT XQ)
in separate parts
1|e0
exp
00
1dGxn
2
0. 2xx
|| 1
cxp((RX RRW,J 11xx)) (a0+55411100+++)1
exp
2
LdGxn
exp
1/-
7x
vo
2
113
ldaxx
idG+jd+ 100+))+1sd0+++1)
||4113
exp
8:
0xx
Laxns
2
40

---

and
Q1
Q1
e (00+)4+1100++
Q1 +
(This is a stack of G Q1 in (E.31))
(Q (Q)))4426++)d:
1j ZQ laxn-
Thus
cxp( + RW,
11x + RW, XQ) iac 10++)4+1:H0++++1D
12
1
12
exp R(v ZK
0 (
Laxns ú (0,1, G 1}.
(E.38)
2
Combing (E.37) and (E.38), we have
(eXp + RW, 0 T
) 11
+ RW, XQ) @ (exp(W.]) exp(
1
1
12
exp R(v, T 2x
+ + Zo))
Laxn 2
i.j (0,1,...,,G1
2
2
Bo
Thus compute (E.35)
(X2 ( + RW, )ixa + RW, X2) @ (exp(WFF) eex((WT)))
G-16-1
= ΣΣ (X2) (.
1=0 5=
R [exp((XXX RW, 111x + RW, X2) |exp(WFF) - exp (
id (a
G-16-1
||:
|uy
11
daxn C 2 13-1)
Σ
exp
R(u, 7x
+ 11 710)
2
Bu
0=
(Because X2 is horizontal stack of Ia)
G-1G-1
|1:112 1
|1yy112
2f( 1,
Σ Σ exp
R(v! ZK
+ 0J Za) Laxn
Bo
j=0
We now put back the 1/2da(Zx. Zo) in (E.34) and calculate the final output as
V Softmax KTQ Q
1
X2
Texpp( + RWW, 11xx + RW, XQ) 2 [exp(WTT) - exp((WT) 026*x(206*-n)
2do (K,, Zo)
1
G-1--1
exp [R(u 2x
| I 12
UJ Zo))
Ldxn
2f 1/111)
(d
Bo
2do 2d Zo)
1=0
1=0
exp(W.7 (WT)
G-1
exp
R(u.
7x
11::13 10,13 tu] Za)
laxa 10 -1
Bu
dBg
Li=0
Ei=0
=
alZxx
02
Next, construct Wo to be
dB In
Wo =
01200=_n)xx
41

---

This yields the final output of Attn o Linear to be Attn o Linear ())
V Softmax K" QWo
G-1 G-1 dBg
exp
Alle
2K
11,;1 2
1 4 Zo)
laxa
fin 1.y 1
d (2dcr4-n))
2==
===
(11x,oo)
(2dG2
Wo
G-1 G-1
exp
R(v
Zx
110:11
11uy112
12 Za))
Idxn
(01-1185-1) Ro
(E.39)
ZZ ZQ)
1=0 1=0
Estimation of Error between Attn o Linear and f We now calculate the loss between the result
in (E.39) and the target function f. For simplicity, we first define Z = 11 _T to accommodate to the expression of affine functions.
Definition E.3 (Max-Affil -ffinffine Function on Ž.). Let Aff., E Rd + R. j E f0.1.2. ... .., .G. -1 denote group of affine functions defined
Aff,, (Z) = v. ZK + 13 Zo
2 1
2 |u E úje 10,1,2... .. ,G- 1}.
Then let MaxAff E Rdn
R denote max affine function whose affine components (Aff..ll.j E {0, 10,1,2,... ... G - 1}}. Explicitly defined as:
max MaxAff(2) (0, 1G
{ Aff (2)).
In the following discussion, use 17 10, G to refer to pair of coefficients (i.j). and denote Ai,j as A, for the corresponding y,. Furthermore, we denote the two labels encapsulated in 1) i, and jn
1
Because the target function f is a continuous function on a closed domain, the function f is uniformly continuous. Thus for c, there exists
0 such that for any 2 1 (1) = 2
[2 18º1. long | - [I 5, we have Ilf(Z(()) f(z S c/3.
According to this 8, we divide the affine components of MaxAff into three parts, the maximal component (and also with the smallest label on both entry), whose label is denoted as 1mm the group of affine components equal to the maximal component or smaller than it by no more than d, and finally, the other Aff,, We write out the labels of these groups of components as follows
min Mm 20,1,,2 GG-122 (Z) MaxA((()))
En = {n MaxAff( - (Z) - 8).
Ei {1) MaxA(l( (Z) > 8).
For any pair of m, 1)2 E 10
132, we denote that
Aff m (Z)
Aff 122 (Z)
|| 11
||
12
10in
|1 112
Um 7k
+
T
7Q
U
7K
T
7
2
2
1
114 113
||
1/
112
Uen
2x
+ 1/m
Zo
(E.40)
2
2
2
2
||
|
| 112
||
|| 112
1/402
7x
+ 3 7g
2
2
2
2xx
Umm
Vin
+ 11
Viny
112 +
11/2
2
1
1j 11
42

---

2112
1
Vin2 Vjnz
11
2
Via1 Vini
11.
(E.41)
Let Un = I ",, denote a flatten stack Vi, and U),, Same as vs define टy := 1 "/
Then the above expression denotes Tmm is also the label of the uy closest Z among Upp 1/ {0, (0,1,".,,0 1)2. Thus we have
1l 2112- = nE(011,-. G-132 flle Z112].
min
(E.42)
This means that Vym. is the grid center closest to 7 in 2-norm.
We now prove this closest grid center has a distance to Z smaller than half of the grid width (D/g) in infinite norm.
Let D = 2D/g2/2//// (1///x(-11, neighboring Vy.. For any A in D, from (E.42) have
1-1, 0, lldr denote a set of differences to Umn. all the (i E {0, 1, ... ,G-13)
2 S 11:m + 213.
This yields
2A (Z V/m) S |
which means for any & E [dn], by selecting A to be + 22 E
(dn)
we have
2D ( Um )a = 21 2T
4D²
#2 x
Um
S |
g²
Thus we have
D
(12
Un 1i
h: E [dn].
g
This is equivalent to
D 1 Uy 1100 S
k E [dn].
Set g to be larger than 2D/5, 6, we have
8 2
11 Um 1000
thus
llf(Z) - f( 11o00 S
(because 6/2 < 8)
3
Calculation of ||Attn o Linear flle.. We now calculate the difference between the output in (E.39) and target function
G-1 exp
R(u, Ž
1
2
Attn o Linear(Z) (
((
a(Z)
n=
6 exp
|n
=
2
(f(T,)) - f(Z)))I
a(2)
0=
[io - Isy 12
By m=
a(2)
43

---

G-
exp
R(E 2
11( 1(2)
a(Z)
7
(By property of infinite norm)
exp
R(u mm
11 11
[f( - f(2)
a(Z)
exp R((,) T 7
(1f(8)) f(
a(z)
exp R(v T 1 (1((T))
2
f( (E.43)
a(Z)
MEM
The last row is simply a separation of the summation the row above. We now calculate each part in (E.43).
As previously stated, for any Zı, Zz, as long as -2 S 6, we have | s e/33. Thus when designate Zi Un for any yo and Z: Umm along with (E.40) we have
R 21 T 2 |
exp
(1((,)) 1x
(z))
7
exp R T 7
2
(11((T)) - f(ōmm 11 + (lf((Tnm))))))))))))))()))) f(2))loo)
a(Z)
exp R(u, J 7
b. 1
2
C
+
a(Z)
3
1
exp R(u, T 2
12
2
2e
(E.44)
a(Z)
3
For any 7ms we have
exp R(C T 7
1mmm 112
exp (((RRTTTT.. 7
Vmm 11
7m
2
2
[lf(0
-ff(2)llo 1
(E.45)
(Z)
a(z)
3
When iss larger than In (3/2. BoGG)//(38)) we have
exp [R 2 |,
exp Riv 2
|5(T,) - f(Z)
2B
a(Z)
a(2)
mEm
nEm
(y that f is bounded)
Σ Σnem exp R(u, 7
11,,112
< 2Bo
a(Z)
Σ Σnen exp R(v, 2
112
2B
exp
R(vT 7
(a(Z) is the sum of all exp
[R(,] _ 11yl113
larger than any element in the sum)
R (1)
2Bp Σ exp
- ||.., 2
Em
44

---

R
6
2Bollmalll exp
(2)
)2
82)
-3R82
2BoG exp
8
8
$ BoGe)
-302
2B,G exp
(B.46)
3
Combing (E.44) and (E.45) yields
exp R( _ 7
|
2
(l((o))
f(2)
a(z)
neno (1mm)
exp
R(v, 7 |1ull12
2e
exp
R(v1 7
Vumn 11
2
(By (E.44) and (E.45))
a(Z)
3
alz)
1/n00
exp R((1]] T 2
1
2
2c
a(Z)
3
2e
(F.47)
3
exp [((, , 2- 1=111
where the last line is by EmEEou(Am. By (E.47) and (E.46), we have
$1.
a(7)
exp
R(U T
Tm
7-
2
Attn Linear(2) (2)- 5
- f(2)
a())
exp [((x 7
||:112
[1f(Tr) f(2)
a(Z)
R(u, 7
exp
[l f(2)1l
a(Z)
nEΣ
26
3
3
6.
This completes the proof.
Theorem 4.2 can be easily extended Lebesgue integrable functions in Lp norm in the following result.
Corollary E.2.1 (L,-Norm Universal Approximation). Let J' Ux Uq + Rdx= denote any Lebesgue integrable function on compact domain UK x Uo and let be any positive real number. Here UQ E ddxx stands for the compact domain of the two input sequences of cross-attention. Then, there exists cross-attention Attn prepended with a Linear layer such that
| - Altn o Linear | <
Proof. Without loss of generality, assume Uk Uq [-D, D, Djaxn for D Rt.
Since Lebesgue integrable on a compact set, / bounded almost every where, Let B, denote
the bound of |[
45

---

By Lusin's s theorem, for f on compact domain U, there exists a continuous function which is equal U except for region D such that p(Ds) I A. This can be written as
D3 (Z) + o(Z)),
(E.48)
(F.49)
where j stands for the Lebesgue measure of set.
By Theorem 4.2, there exists network Attn o Linear, consists of a cross-attention Attn and layer of sum of linear transformation Linear such that
Attn Linear gllu> ED,
for any > O.
This denote that for any Z X
|Atn o Linear(Z) g(2)lpp - (dn - .P) / = co(dn)
Combing this with (E.48) and (E.49), we get
p Z ||Attn Linear(2) (Z) (())l > t0)) S ( S A,
(E.50)
since if f(Z) = 9(2), ||Attn o Linear(Z) - g(2)ll ||ttn o Linear(Z) f(2) s co
This yields
|| - Attn o Linear]| r. = - IZEUXU
|| Attn o Linearlp d
/ ZZEUXUDD |l - Attn o Linear| dx + / I2EDs ||f Attn Lincar dx)
-I lg Attn o Linear de + /
I/ Attn o Linearll dx)
ZEDs
(M(U ((Dd)(c ) + A B s co(dnu(U x U)} + A Bp-
Set
05 2(dnp(U x U
A < Bp 2"
We have
|| Attn o Linearl| S <o(dnu(U x U))} + A
(dnp(U U))
BBp
Bp
2(dnu(U x U))
Bp Bp.2"' 2D
é
2
2
e.
This completes the proof.
46

---

Proof of Results in Appendix A
F.1 Proof of Theorem A.1
Theorem F.1 (Theorem A. Restated), Let f Rdxa - d denote an T--Lipschitz function (in terms of 2-norm) whose input domain X. For any > O, assume X' contained in N₂ sphere by the radius of e/(3L) in 2-norm. Then, there exists Linear layer and a Attn layer such that:
Attn o Linear - fl S e.
Furthermore, Attn and Linear have total number (ddnN,, trainable parameters.
Proof sketch. This proof is identical with Theorem 4.1, except for alteration the of uy- Proof. We follow the proof of Theorem 4.1.
Notation of Sphere Centers. Let Z = 12,, 22, ... ,,,,,,
2n.] E Rd denote the input to Linear. Define
Z [2] 2 ... IT P N+ is parameter that controls the size of the attention block and the error of our approximation.
Let E [Na] denote the centers of the N₂ spheres that covers X. Let V == (vili [N.]] denote set of all Vz.
For every E V, we define T. == [u]a: T 1241:24 " ((n-1)d+Imd!
Construction of Related Functions. Because J is continuous within a closed region, its output value bounded co-norm. Let Bg denote this bound, we now construct two functions that. For any a Rdxn, we define F((a) == 1axn f(a)/ / Bo and T(a) 1axn + [a)/ Bp- We define (E' (a) = E(a) + (a). By the definition of E and T, (E + (a) = 2dx for any E Rg
Construction of the Layer of Sum of Linear Transformations. layer to be
We now construct the Linear
No-1
(n-1) 1)
Linear (2) =
Σ
Σ
(m)
[2dN.+1) +1) e
(2dN.) j+s+1 + (2dN,)
) +
01x2aN. 12aN.
1=0
k=
where N₂ pin
We now express the output of Linear simpler form in the following discussion. First, we show that
(n-1) 1)
(n-1)
(Zekr))
(1U))
=
Σ
2x+1 (2)
T
k
km0
=[2ī : 22 1 ; |
= r.. 7 E R, j E 1,,1,2.... ... N₂ - 1}-
,
This yields
Linear(2)
Nr. -1
Σ
2
d-1
Σ
( + (2dN.) T (2) + @zo0+dN+11
)
01x2aN= TzaN
1=0
==0
Xo
Xo OaN xdN, IdN.
LaN=
Laa XdN,
in which Xo is defined as follows
Xo = 10 E1ixd
ou 71xx
To Zlix
"N_-1 71 lxd
47

---

Construction of K and Q Matrices. We now construct the Wy and Wq matrices in the self- attention block and calculate the output of Softmax (Q) We define Wx follows:
01 x
O1 x - =1lد lixd In In(T(uN.-1).
O1 N lixd
0 xd
WK =
| lixd
U 1xd
2 In (T)))
2 In (
In (Too)
The definition of WK yields
K
W Linear(2 (
1
01xd
01xd
01xd
0 xx
Xo LaN,
Xo OaN xdN, IaN
Ivall
| - 11
I 1
0
2
1ix
2
1ix
lixd
1xx
0
In (T()))
In (2N=-1)
In In))(o)))
In E UN
OaN. xdN,
Zlixd
N.
Z1xxd
110 Z1xxd
1 N, -1 1xd
wull 1xd
lex, 1 11
- 11
1xd
2 1xx
2 1xx
2
2
In (TT)))
(n(T(NN.1))
In (n(B(To))
In in(B(TN.->
Next, we construct WQ to be
Rl Rlixe Ix
0 01x -n)
Wº =
Rlixn
0 (2x(2.
On
I
(2x(2d...
This yields that
Q
W WoLinear(Z)
Rlixn
0 (1x(2dN,-))
Xo
Xo OaN, xdN. LaN.
0
Rl
0 (2x(2
Lana
On
In
O2dNN
OaN xdN=
Rlixn
01x(2dN.-1)
Rlixn
01x(2aN,
(nx(2dN.n))
We now calculate the attention matrix Softmax (K (TQ).
Calculation of Softmax K (TQTTTQQ Q)-
First, KT Q can be expressed as follows
|eall la 2
In(T(T)) In(T(T)))
vi] Z1a
| 12
1a
1N 71d
|| 1a |ealll 1a |e 1a
In(T In(F(66))
Rlixn (Ix(2aN.-1)
kTQ=
Rlixn 01x(2dN.-n)
10 71d
In
Onx (2dN.
o, 21d
In(E(T)))
UN Z1,
| la
In In(E(on,-1))
48

---

R(og Z
|
dxxn (T))))))
0dx(2dN,-n)
2
R(v] Z m
Idxx
In (1))
(2xNNdN
R(ON, -1 ,7-
||ux,-
dxn + In((T(NN.-))
(2x(2
R(v0 2
||
11 1ddn + ln(
Oax(2dN,
R(u 2
|11112
2 1dxn (n()
(2x(2dN.-n)
R(UN R( |للد- )
Laxn In (n(E(öN,-I) (2x(
Now, divide the calculation of Softmax (K Q) into two counterparts, the calculation of exp( K Q) and the calculation of the denominator of every column of Softmax (KTQ) in the expression of Softmax, explicitly written out Σjey oxp(K* Q) for each E [2dN2].
For cxp(K" Q), we have
exp R(vd
||e0 2 T(ūo)
ldx(2aN,-n)
2
exp
R(o 2
I|en i T(ū1)
Lax(2dN,
exp
[ T -1
Iuy,-1
T(TN.-1)
lax(2aN.
exp exp(KTQ)=
(F.I)
exp
E(To)
Lax(2dN,
exp
R(us Ż
E(71)
Lax(2aN,
R(uN expa
1
| 111
F((
1 1dx(2aN_-))
exp
2
1100 12 T(To)
Lax(2aN,
exp
R(o]
7
Illl 1 T(ū2)
(ax(2aN_-))
exp R(UN. -
Iun, -1
((UN.-1)
Lax(2aN,
(F2)
exp
Rivy 2
woll?
E(dg)
Lax(2dN,
exp
R(uT 7
ivall
E(w1)
Lax(2aN,
exp - ( E(WN.-1) Lax(2aN,
For the denominator, we calculate it in columns. Let denote the column which we calculate the denominator in Softmax. When i E {n 1, 2,. 2dN,), it obviously equals to 1 2dN, 2dNz. And when E [n], we denote that
24N.
N₂
1 112
exp( (K T QQ),,
=
(1) + (xxab(tg-1)ad)) exp
R
(Lixa(E +T) exp R (uj1 || 113
(1xxa(2dxn)aa)) exp [R R0 |1/z- 112 ))
exp
49

---

N₂
R [0]1 1 112 )
Σ exp
i [n].
(F33
j=1
We observe from (F.3,, that 2ラ exp K Q) is invariant of for E [n]. In this case, we define
(QQQ)
2dN;
N₂
100- 113 )) E R, i E [n].
a(Z) == 2d ==1
exp
K (TT)),g
=
Σ
exp
2
==1
From (F.1) and (F3) we have
Softmax (K"6 Q)
=xxp(K Q) の
aaN oxp(KT 1)15 j=1
12dN.x0 xn
2dN, 12dN x x22dN.
(By 1/2n
2dN= exp K (KTT) Q) is
is invariant of é for á E [ri]
exp
R(o 7
11v0 13)
T(ūb)
1dx(2dN=-x)
exp
R(o) 7
ull &
T(ū))
(dx(2dN.-))
R(vA. exp(R(u). Sollll -12
exp
Ivy. 111.
T(NN.-1)
1dx
2
2da(Z) 12dNxxn xn 2dN. 1aaN, x(2dN,-))
exp A(vD 2
[
E(7o)
Ldx(2dN.-o)
2
exp
|| 113
E(ū1)
1dx(2dN,-o) -n)
exp
R(vN
-1 2
1I alle
E(UN,-1)
12x
RR(v) 2 1
exp
((0))
N. Lax(2dN, -n)
o(2)
[R(11 7 101113
exp
T(ū1)
N. Lax(2dN,- m1)
o(Z)
-1 7 11 11
exp
R(Ux,
T(UN--1.
N. Lax(2dN,-
(17)
2d
exp
A((1
7 l1vol2
E(ūn)
N. Lax(2dN, -1)
% 11v1 13
exp
A[v1]
E(τ1)
N. 1dx(2dN,-1)
a(2)
2 I 113
exp
R(vy
ElUN, -1)
N. [xx(2dN.-n)
a[2)
Construction of Wy and Wo. We now construct the Wy matrix and calculate the V matrix of the self-attention
We define Wy
Wv = 10a
X1 - xX1],
where
X1 [T
La
Ialaaaa
is matrix formed by stacking N₂ Ia matrices horizontally. In this definition, matrix can be calculated as follows:
V == Wy-Lenra Z)
50

---

Xo IaN. OaN, xdN,
Xo OaN xdN, IdN,
Od
X1 -X1]
X
-XI].
After the construction and calculation of V. go construct Wo
d
Wo =
(2ddN=-n)x n)xn
The sole purpose of Wo is to extract the non-zero entries of the final output.
Calculation of the Output of Attn o Linear. We now calculate the final output of the self-attention block.
exp R(v0 2
-
T(op)
N, Lax (2dN,-N
R(v 2 1
exp
T(ú1)
N, Lax (2dN, m))
a(2)
2 11
R(UN,
1
T(UN-
1
Lax m)
Attn Linear(Z)
[X1 -X1]
(22)
N,
Wo
2d
exp
R(v]]
2 1vo1]
E(0o)
N. Lax(2dN,
(())
[[(p] 7 1o113
exp
E(71)
N₂ Lax
a(2)
exp
Blux,
2
11
E(UN.-1
No 1ax m))
a(%)
exp
R(u)
(T(T)) E(7o))
(dx(2dN
((Z)
2 1011
1
R(u]
X1
(T(T)) E(Ti))
Odx
Wo
2d
a(2)
|1oNa-111 111
R(M
(T(UN- B(ūN.-1))
Odx(2aN.-n)
a(2)
exp
R(v) Z
2f(m0)
0ax (2dNa-n))
o(Z)
Bo
1
R(vi]
Z
1= 2
2
X1
0 Odx(2
Wo.
2d
a(Z)
Bo
7 1l0 11
R(u
2f 2f(
Odx(2dN=-n)
a(2)
13
We have
A(vv] 2-
[1vgl1
[1(vg] 2
11
exp
2f((n)
exp
2f(w))
a(2)
Bo
a(
Bu
X,
exp
A(u]] 2 11
2
2
[I Ia
IdlaxaN.
exp
[((v]
2
1|u] 13
2f(01)
a(2)
Bo
a(2)
Bo
R(vT,
Z
1'N,
1'N
2f( 1)
exp
R(vT,
2
2f(UN= 1)
alZ)
Bo
a(2)
Bo
51

---

N, -1
exp R(v)
2
1ng 1113)
2l(0,)
la
a(2)
Bu
1=
N.
exp
[[(v] 2
2
2f((;)
=
a(Z)
Bo
j=0
This yields
Attn Linear(Z) (2)
N.-1
exp
A(v3 2
2f(1z)
Wo
(ax(2aN.-n)
0=
a(2)
Ba
N
exp
[[(v]
2- 11g11
25(1g)
Odx(2aNs
dBoh O(2dN, -n)xn
0=
c(2)
Ba
N. exp
R(v] 2
110g113
2
f(ūz).
(F.4)
a(Z)
0=0
Estimation of the Error between Attn o Linear Z) and f(Z). After the above calculations of the output of the network, we can now demonstrate how this output approximates our target function.
Definition F1 [Max-Affine Function on Ž). Let Aff, E Rd + R. E f0, ... ,N - 13 denote group of affine functions defined
ArT, u, ley
E 10, ..,.,...,,,,111 1).
Then let MaxAff E Rd > R denote max affine function whose affine components are (Aff, | j E 1,,
N₂ 1}}. Explicitly defined
MaxAff(2) =
max 2E .o (Z).
Because the target function f is continuous function on closed domain, the function f is uniformly continuous. Thus for e, there exists a 6 such that for any Z1, Zz, as long as [2, Zelloo S 8, we have [Z, f(2))ll S e/3..
According to this 8, we divide the affine components of MaxAff into three parts, the maximal component(and also with the smallest label), whose label is denoted as Jm. the group of affine components equal to the maximal component or smaller than it by no more than 8, and finally, the other Aff j E {0,1,2, follows
Na - 1]. We write out the labels of these groups of components as
Jm = min
JE N.-11 (Aff, = MaxAff (2)}.
Jo = (i MaxA(T(Z) (Z) s 5),,
J1 = (j MaxAff( (z) > 8).
For any pair of i,j {0, 1... ... .o 1}. we have
|10:11
10y12
- = u.= Ž
1
2
2
12112
(-112)
[112
|
+ v/
T 7
+ 11 2J Ż
2
2
2
2
1
1 - v/
1/12
+ 2
12 1
This denotes jm is also the label of the closest Us to 7 among all 0ú i E 10,1, ,N - 1}. Thus we have
1/1m - 2112
min
{llo - 2112)
(F5)
(0 No -1-11
52

---

Thus, when considering the Z in the input domain of f. which by definition is contained in N₂
spheres, the closest center to Z is the sphere containing Z. This gives
c
1 -imiil 122 S 3L
(F.6)
Then, with the L Lipschitzness of L we have
[I( f(mjm)]lo S
I. E 3
(F.7)
3L
Difference between Attn Linear and f. We now calculate the difference between the output in (F.4) and target function f
Attt o Linear(22 (2) =((2)l
N-11 exp M[v; T 7 |1
f(ö,) ol
a(Z)
exp R(v) T 2
B(v] 2 11,
(f(0)) - ((Z))
N--1 ==0
1)
((Z)
a
exp [ 0
11 12)
2
11((8)) f(2)
(property of infinite norm)
alZ)
j=0
R(v]]. T 7
1
exp
2
[l((0)m) - f(2)1lo
a(Z)
exp R(v] T 7
2
11f(;)) - f(2)
a(Z)
JEJo
exp [(v] _ 7 1اy13
+
11(T,)) f(2)]loo.
(F.8)
a())
JEJ1
We now calculate each part in (F.8).
For the L-Lipschitzness f, for any Z1, Zo. as long as 11 - -elloo S i
37. we have [[(())
f(Z2) I e/3. Thus when we designate Z1 vy for any j E Jo and Zz Vym. along with (F.7)
we have:
exp Rio 2
-(11)--(Z)1))
(F.9)
a(z)
jeJo
exp
[(v]] 7
(ll((T)) - f(o 1l + +I]((T,,.)) f(2)
a(Z)
jEJo
exp
R(o, 7
||e,11
E
(
a(2)
jEJo
exp
Aloo] 2
||
2c
(F.10)
a(Z)
3
JEJo
For Jm. we have
exp
[R(u]__ 2
|| 11
exp [B(u]. 7
||0 11
2
[|
f(2)]loo
<
(F.11)
a(Z)
53

---

When is larger than
8
In( Bo.Nzs), we have:
362
[((2 -- 1لايl
exp
11s(T))
f(2)
a(z)
jcN
exp R(Uj
7
leglli
2
2Bo
(by the bounded nature of
a(z)
jeh
Σjen, exp
R(v,
1
< 2o
o(Z)
Ejen exp
R(v]] %
le,y
<Bo
exp R((1) T %
1m3m l12
2
(a(Z) is the sum of all exp
R(u,] 7
) thus larger than any element within the summation)
=B
Σ exp
-
R (1 -12 110 - -1
2
jEJ
< 2 exp
R
(-)²
8²
2
-3R82
< 2BoN exp
8
8ln 3 Na
352
2Bo exp
8
(F.12)
3
Combing (F.10) and (F.11) yields
exp [(( [[oo] |
|1((T)) - f(2)
a(Z)
jEJuU(mm
R(v; 7 [|u,113]
expl R(o,J jm
7- || 11
exp
2
2e
2
(By (F.10) and (F.11))
a(Z)
3
a(Z)
jEJo
exp
R(v..2 7
||
2e
a(Z)
3
JEJ
2c
(F.13)
3
exp [1(,] 2- 11g13))
where the last line is by Σjebu((i))
1.
a(2)
We plug (F.12) and (F.13) to (F.8) and get
exp A(r jm
2 |1 12
|| Attn o Lincar(Z) (2))))) <
[1((rm)--(22
a(Z)
exp R(wy] 7 11,11
[[((o))) f(2))l
a(Z)
jEJo
54

---

R(u
2
1r,11 [1((;)) 5(2)
exp
a(2)
JEJI
+
=.
This concludes result on the approximation error
Estimation of the Number of Trainable Parameters. We now estimate the number of trainable parameter in the network constructed to verify our claim number of trainable parameters in the main text this theorem.
Remark F.1 (Meaning of Trainable Parameters). By trainable parameters we denote the parameters that differs according to f. This includes the parameters related to the input domain of X, and excludes the constants (i.e., 0 and 1) in the network.
We estimate the number of trainable parameters by each layer in the network.
First, we do the estimation for the Linear layer. It consists of sum over N. v2 Z, j E [Na]. and thus contain dn N₂ trainable parameters.
Then we do the estimation for WK and WQ. We restate the construction of Wx and Wo:
R
01xd
01xd PNN..1
01xd |leelll lixd
0 01xd
W =
0
R Poll
lixd
R
2 1ixd
R
R Ju. 1 1ixd
0
(n))
In ((N,
In In(E(To))
In In(D(UN,-1)
Rlixn
(2dN2
Wo==
Rl1xn
01x(2aN--1)
On
In.
(nx(2dN,-n)
From this, we observe they combined together have 2d N₂ 2dn N2 trainable parameters.
Finally, For Wv and Wo, we restate their definition:
Wv = 00 X1
-11,,
dBoLn
Wo = 022a
where
X1 = [la la ... Lal Ialaxaa
Wo contains n trainable parameters (dBo).
In conclusion, the whole network contains total of dnNa, + 2dN, + 2dnN + 4dnN + 2dNN + nn,
trainable parameters, which is of OddNN,) level.
This completes the proof.
G Extended Related Work
Recent empirical studies also shed light on the practical behavior of attention mechanism. Olsson al. [2022] show that induction heads help models learn patterns in context. Sanford al. [2024a] prove that Transformers can do complex computations with few layers because they work in parallel. contrast, Luo [2022] find that some Transformer designs lose expressivity when using relative positional encodings. Our work builds these ideas. We prove that single-layer, single-head softmax attention with a simple linear layer can approximate any continuous function on a compact domain. This shows that attention alone can learn arbitrary pesuencc-to-sqqeeee mappings.
55

---

H Extension to Transformer Setting
In this section, we demonstrate how our proof technique attention extends to the transformer setting. Specifically, we show it is possible to three-layer multi-head attention-only network to achieve universal approximation in Lp norm.
Lemma H.1 (Attention Simulate Column- Wise Transformation). For am input X E R Let ((X) AXB, A Rouiput xd B R=x" be the linear operation (both boken-wise and sequence- wise included) we wish to get. Without loss of generality, suppose B has positive entries (since all-zero padding token. Let n++ be the
subtraction of positive matrices yields all matrix). We positional encoding below X. Then for any e > O, there exists an single-head attention Attn such that
X
0g
Attnn
In
O, ((X) Odoutput 11 S C.
0 01xn
Proof. Let 8, be the of all entries of B. (i-th column) and M maz Define
[s1 181...Spl Construct Wy 3M * [A Oduappa x(n+ WK == [0nxd In(B) In(3M 1, ] (T will be turned to arbitrarily large), then
S] ] (In entry-wise), and WQ == [Onxd In T 1n]
the output attention 3M Oy]Softmar((n) (annttar(I((BB (ln(BT) In(M 1, - ST
[ T-1.).
This equals
T T-H In(3M 1ixm S) T T-H
(n(B)
3MA{X So]Softmarl
in which H1, H2 are the sum of all columns in In( (B), In(3M lixn S) respectively. This further equals
B
au
3MA[X
Oa]
3M
= [AX
CJ. (fao ai] is the T'related column),
1ixn
3M
S
a1
in which C 3MAX ag, when T +00, do Od and the padding token is preserved as well.
L
Lemma H.2 (Preservation of Identity Matrix). For any c and N+, there exists an attention head Attn that satisfies
[Attn(T,) Jnllo S.
Proof. Construct the WK. Wo, Wy matrices of this attention head as
Wx == RI
We = In
Wv : In-
The output of this layer is
Softmax(RI,,
Hence when R is sufficiently large, the output is within an arbitrarily small error to 4n
Corollary H.0.1. For any > 0 and input of form
A
In
where A E jg denotes any matrix, there exists an attention head Attn such that
Attnn(
A
0mxn
alxo S e.
1n
In
56

---

The proof is obvious. We only have to zero out A in WK. Wo. Wy and do the rest as Lemma H.2. We now prove the universal approximation result in the transformer setting.
Theorem H.1. Let X IR denote the input sequence. Let f be continuous function on compact support. For any yj O, there exists three-layer multi-head attention Attn (9), 1 [3] such that
X
0₂]
Attn (3) Attn (2) Attn (
I
On 1
),1mm f(X) S n.
01xn
The First Layer
Let U; i E [Pdn] be as defined in Remark E,
Let X represent the flattened input X and let ö represents the Rdx= sequence form of E Rdv
Let d;, i [d] be defined as
0;-1
5; ==
1
E
Rd.
P
Od-
Let Attul (1) h E [n]] denote the h-th head of Attn exists such set of Attn,, such that
According Lemma H.1, for any co 0,there
X
0a
|Atto
In
On
(X)
[(4+im) 1l S co.
O 01xx
1
in which
0d(h-1)xx
4. (x))) =
2h lixn
h E [n]
Oja(a- (o((n-A)+4)]
And construct the n. + 1-th head to preserve the identity matrix using Corollary H.0.1 (A having the
dimension of dn (n + 1) and the identity matrix preserved being Int1 this case). So the output approximates
X lix
Oan On
In
(H.1)
01xe
arbitrarily small error in the infinite norm.
The first attention layer is then defined as
n+1 Σ h=1
Attn (1)
Attn (1)
h
The Second Layer
Label the heads in the second layer with (v.g, s) pair, v [pdd], [n]. (+1}. Construct
parameters in Attn (2)
(n.gs) to be
Onx(g-1)d
D
Onx(n-ghe
-D -D)) lixn
We =
01x(g-1)4 1)d
01xd
01x(n-9)d
27-
1ixn
WK = R Onxdn
In
0,
01xdn
11xn
57

---

in which D is
61
D
and IR a parameter that will set later process.
Then the attention score matrix of Attn (2)
is
(v.g.5)
La
Og
[ (Wo)..)) -Ixn
Oz
Softmax
R
1ixn
0
Lixn
2P
Softmax
R
D(ag
(Vo):.)
lixa
+
2P
Inxn
On
01xn
D
(U))..g)
1ixn
0g 0
Softmax
R
2P
1ixn
0
When R taken sufficiently large, for any given 0 and an arbitrary A, 0 invariant of and
R. if
67 (zg (f))..g)
+ A. 2P2
for specific i, then in the attention matrix of Attn (2) the last row by A., and this difference after amplified by R makes the last row to be lesser than the given 61- Similarly, when
(e. the i-th row should be larger than
(x. - (Fo)..) < 2P2 A
for specific i, the (n + 1, entries in Attn (2)
(0.9.-1 (j = 1,... are smaller than 61 when setting R to be sufficiently large. And the sum the rest of the th column is larger than 1 1- c1.
Finally, when 16, (x, (15)).g) < 2p2 A., for all E [n], it makes the last row largest in every entry (except the last column). Thus there exists sufficiently large Rt such that the last row of the attention score matrix larger than 61 in every entry. This also makes the sum of the rest of the rows less than or equal to in every column.
We show that the last situation is equivalent the input token Tg falling into the grid centered (6r kg in Rdd, (notice v. is grid point in Rdn and its column is grid point in Rd))
1
6, (xg (Fo).o)
6, (xg
(1%) (gg) - 1
222
1
[|| (1% n || ((T)). + 5,) 12)
2
- 6. (Igg - -(To):g)
= 5. (5. (Fo)..) 1/1 1:::112
2P2
2 11 (o.) ||I ((,)).. 5.) 6.)
Therefore
167 (xg (To)..g) < 2 A.
58

---

is equivalent to Ig being closer to ( than its neighboring grid points in R (which are {(T)): +
5,j2 E [d]), this concludes be within the grid centered at (6.)..g
This means that when bg is within the grid centered at G. both heads in Attn At (=.g.s)) E [+1] has their last row in the attention score matrix larger than 1 (which makes their sum larger than 2 - 261) and g
in the grid, their sum would be smaller than + 6 (because at least one smaller and all entries should be smaller than 1)
Let Wv of head labeled (v, 9, 8) be
O x 0ixda Ot,xen
0s.
lixn
0
Oc
Oz.
in which == vd + (g 1), t. pdnn 8o + 1. We note that here head (v.g. 1) and (v. -1) have the same head.
Then V is
lix
Odn On
V Wv
In
01xn
Xauxdn
Osaxn
X lix
Oan On
Onxdn
Lixn
In
Ot xdn
0t.xn
O xn lixn
E (R(PPd xtn)x((+)))
0
The output of the (v, (0.g.s)-th s)-th head in the second attention layers is
0s.xn (S(()
0,.
V ((....))
lixn
1+1111
n+1
(x+x()+1)
O,,
in which
s(u.a))) = Softmax
R
D(x, (Uo):.)) lixn
O₂
- 2P 1ixn 0
the previously constructed attention score of the (v.g. vogg. 8) head.
Finally, add head Attn| (2) to preserve the identity matrix as that done in the first layer using Corollary H.O.I.
This head outputs (with arbitrarily small error)
T0 Opdenxn
0p On 1
In
01xn
Define Attn (2) as the sum of previously defined heads
Attn (2) = Σ
Attn
(e.g.)) + Attn,
v.g.a
Then we have that when the input is
1ixm
Odn
In
On 1
01xn
59

---

the output of this layer is within an arbitrary small error to
21xn S(1
21xn
S(1,n)
(H.2)
21xn S(Pán, n)
In
O,
01xn
in which
S(v.g) = (5( 111.1. + (S(1...))) (a..0.))++111. 1+11mm
We note that only when input is within the grid centered u, 21xn S(v,9) is smaller than 2 -
261) 26, for every g in all entries.
The Third Layer
The third layer consists of pdn blocks, the v-th block identifies the input falls inside the v-th grid. For the v-th block, construct its WK. WQ as
WK = O (ax(v=1)n 1)m
Ouxn
O (nx(nen-u)a O1x(pan-u))
2
1
In
On pan u+1
0 (1x(111)
11x
01xn
We= == R1 10(=+1)x
In+1]
in which R3 is a scaler whose size we will determine in later process. When the input has the form of that in (H.2). The K and Q matrices are
21xn - 5(1
m+1
21xn S),,
K : WK
21xn S(Pdn n)
n+1 O,
In
01xn
d EgE(n) S(v.g)
0n
R (m+ (x++))x((+1)
E
n 21xn
Q = R Lat1-
This yields the attention score matrix to be
1 Egeta S(v. 9)
)
Softmax K T Q) Soottma Ri n 21xn
By n 21xn - -gge(n) S(v,g g) is smaller than 2ne; if input is in the grid centered at v and larger than 1 e1 if not. Set Ej = 1/4n, get that for any eg O, when R₁ is sufficiently large,
I
1
Softmax K Q)
1 S cz.
01 xn
n+1
when input within the grid centered at u.
And when the input is not the grid centered at v, we have
Toxx
1p
Softmax (K (KTQ) Q)
a+1
alo S e2
Lixn
mt1
60

---

Construct Wy as
Oax {1() Oa]
This makes the V matrix
V (f()) Oa]
when input has the form in (H.2).
Thus the third layer's v-th block outputs
In V Softmax K Q) = (()) 0.] Sif--oofmaxAh R1 n 21xn Egetn S(v, 9)
0g
Then for any Es > 0, there exists an R such that the output of the third layer satisfies
||V Softmax(K Q) (f(%) *] ll E3.
when the input falls inside the grid centered at Up. And
V Softmax( (K -) [0axn *] llco S cs.
when the input does not fall inside the grid centered up. Here denotes the padded column that not considered in the final output and is hence omitted.
Now we sum all these heads and get the final output of Attn (3), The final output is within error of pdn e3 to
(S(E))
*] + [(Pán 1) Odxn *]
with o, being the center of the grid that has the input sequence in it.
Analysis of Total Error
Let Ω be the compact support of f with finite measure p(2) V.
Let F(x) (Attn Attn Attn (z) denote the output of the network given ideal inputs Z, and let F* (z) denote the output given perturbed inputs Z (which are the actual inputs).
We have formerly proven the following conditions for any 0:
Ideal Approximation: [|F(x) (x)]l <.
Perturbed Inputs: The input to the third layer, (z),,, satisfies [|Zz(x) Zz(z)]loc S on a set Ω Bs, where Bş is a region of arbitrarily small measure μ(Bs) < ò. On Bs, the етror exceeds but is still bounded.
Given the conditions above, for any precision > and any p E (o) the actual output F approximates the target function f in the Lp norm such that | fllp < 17.
We aim to bound the Lp error || fllp- By the Minkowski inequality (triangle inequality for Lp norms), we can separate the error into an approximation component and stability component:
| - < || - Fllp + +IFF fllp-
(H.3)
We analyze these two terms separately.
Bounding the Approximation Error
From the problem statement, the ideal network approximates f uniformly with error Using the definition of the Lp norm on the domain
F F-1 = ( [I((z) f(x)) du(x)) du(z))
1/2
(f = =V//p
1/v
(H.4)
Since can be chosen arbitrarily small, this term converges to 0.
61

---

2. Bounding the Stability Error
Consider the term F Fllp. which represents the deviation caused by the perturbed intermediate
inputs. Recall that F(x) = Attn (3) (z)( and F* (x:) = Attn (Zz (z)). To evaluate the integral
over 22 partition the domain into the "good" set 52 B3 and the "bad" set Bs:
||F - FIB = / ||Attn (8) (Zz) Attn (3) (Z))llP du + / Ba
||Attn (3) (7z) Attn (8) (2)) dj..
(H.5)
The Multi-Head Attention layer Attn (3) consists of linear projections and the Softmax function, which smooth and differentiable operations. Consequently, Attn is locally Lipschitz continuous compact domains. Let K be the Lipschitz constant for Attn On the S, we are given that || Zsll S e. Therefore:
Amalysis on S (Lipschitz Continuity):
||ttn (3) (Z3) Attn (3) (Za) I K Zsll S Ke. The integral over is bounded by:
(H.6)
/ ||Aun (3) (Zg) Autn (8) (Za)) dp (()) V(KK))
(H.7)
Analysis on B3 (Uniform Boundedness): On the region Bs, the Lipschitz bound fails. However, since Ω compact and the functions comprising the network continuous, the image of the network is bounded. There exists a constant M > 0 such that [[ < M and |F*()) s M for all Ω. By the triangle inequality, the pairwise difference is bounded by 2MM Thus:
/ Ba
||Attn (3) (2g) Attn (3) (Z du & (2M)² du 8(2M)).
(H.8)
Bss
Combining these results, the stability error satisfies:
- Fllp +
(H.9)
3. Conclusion
Substituting the bounds back into the original inequality, we obtain:
|| fII S <vV + M(2
(H.10)
To ensure || fllp 1), we select parameters as follows:
Choose sufficiently small such that both eV and Vl/p are smaller than n/33
Choose the region of perturbation B3 sufficiently small such that d < 12/3))/(2M)))
Since e and can be made arbitrarily small, for any 1y > 0, the condition holds, and we get the global Ip approximation error converges to
62

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility. transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes]. [No]. or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked use the checklist as one of the factors their evaluation.
While "[Yes] is generally preferable to "[No] it is perfectly acceptable to answer "[No] provideed proper justification is given (c.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist".
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The Abstract and Introduction explicitly list the three main contributions Max-Affine interpretation, self-attention universality, and cross-attentis universality and additional claims are made elsewhere.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
63

---

Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: See 5 "Concluding Remarks", paragraph "Limitations" (lines 336-350) details assumptions partition granularity, boundary effects, training difficulty, and distribution shift.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't I acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: All assumptions are stated next to each theorem in Section 3-4; complete formal proofs are given in Appendices C-D (pp. 16).8).
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
64

---

Answer: [Yes]
Justification: Appendix B specifies datasets, noise-injection protocol, model size, and training procedure; code will be released anonymously with the supplementary material.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, it may be that access to the model limited in some way (e.g.. registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: All datasets (MNIST, CIFAR-10. -MIIT) are public; an anonymized Py PyTorch implementation and run scripts will be included in the supplemental ZIP.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.ce
cublic/guidess Sew.dddee/CodeSumissssesee wommiissiooooicyyyy for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: //nips. wttp:cccccccccccccttcccccccccclt httpc/ggiddeeooeeuuiiiieseeeeees wctp/ooeuubbissssoooooooo.ooo for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
65

---

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: Appendix B lists optimizer, batch size, epochs, and random-seed protocol for each dataset.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Figure 2 (page 14) shows mean + 1s.d.v. over five random seeds for each noise ratio; the caption clarifies what the shaded region represents.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative eтror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [No]
66

---

The proof-of cooof-f-ccoceppl experiments run single commodity GPU, but exact hardware specifications and wall-clock times are reported.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't oodn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: /tt:://nnuurrpi http:////c.ppppccc/ccccccoo...d.
Answer: [Yes]
Justification: The work is purely theoretical/empirical public data; ethical concerns were identified (Impact Statement, line 361).
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: The Impact Statement (lines 361-363) argues that the work is foundational and poses no immediate societal risk.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymeet of technologics that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
67

---

11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: No high-risk models or new datasets are released; only small proof-of-concept code is provided.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Standard datasets are cited in Appendix B; each dataset carries a permissive academic license.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The work does not introduce new datasets, models, benchmarks.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
68

---

The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: No human subjects or crowdsoureing involved.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No human-subject study carried out.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMSs if it is important, original, or non-standard component of the core methods this research? Note that the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: No LLMs are used in the method, only standard writing support tools.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: / //nereuups...// Ceeurips.cc/ooofeeeeeeee for what should should be described.
69


Reviews:
[{"review_id": "QMVzZKTd8x", "reviewer": "Reviewer_bcbT", "summary": "The paper aims at establishing a proof for the Universal Approximation property of (self/cross) Attention. Unlike previous work on UAP, this paper focuses on a minimalistic architecture, where one single attention layer, prepended by a linear application, is considered. The main proof relies on highlighting that this minimal architecture can perform a max-affine partition of the input domain: this, matched with some clever definition of the attention weights and usage of the softmax operation, turns attention output into a piece-wise linear function, which is a universal approximator for continuous functions.", "questions": "On the impact of the Linear layer\n- The Linear operator you’re introducing in [l117] (and that your proof heavily relies on), is not an orthodox one. On a first quick read, I was hoping for it to be a classical Linear layer acting on the embedding dimension - which is what is commonly interwoven with Attention layers in a Transformer. Checking your construction in (D.2) (and even more evidently so in (E.5)), it appears clear that it morphs the input in ways which are not allowed within a typical Transformer network: you’re duplicating the input along the sequence length, which is just not a feasible thing (as it would require dynamically adjusting the context length). Moreover, you’re also adding a bias term which varies along the sequence length (even though this might be justified, for example, if we considered positional encodings). Incidentally, this should’ve raised an alarm: if Attention is permutation-invariant, then that should be reflected as a limitation on the class of functions it can approximate.\n- In my opinion, this dramatically reduces the impact of your results, for two reasons:\n    - It makes the main claim of the paper misleading: the assumption is that your results refer to a single layer of Attention, as it is commonly used within a Transformers network - indeed, throughout the Introduction and Contribution sections, you do remark on how your results help explaining the expressivity of Attention in Transformers. However, that’s not the case, as the architecture you’re considering is not akin to a Transformer’s.\n    - If you allow the prepending of a sequence-wise linear layer, then your framework becomes more akin to that of a generic FFN acting on the whole (unrolled) input tensor, skipping much of the constraints of a Transformer. The results might still be relevant (as you’re still imposing a specific structure to this operator via the Attention mechanism), but much less remarkable, given the plethora of UAP results for MLPs with various nonlinearities.\n- The remainder of the proof seems otherwise solid: I admit I mainly skimmed through the main points in the appendix, without applying too rigorous a scrutiny, but the main steps seem reasonable and I’m expecting them to hold.\n\nThe one highlighted above are the main reason behind my evaluation of the paper, and I see two main ways this can be improved: either the authors raise a valid argument against the need of their linear layer being a strong limitation; or the authors accept this as a strong limitation, and adapt the paper and its message accordingly, by properly highlighting the framework considered - although again, the impact would remain rather contained.\n\nMinor\n- On a similar (but possibly less concerning) note to the above, also the construction of your K/Q tensors relies on non-orthodox procedures in Attention: Wk and Wq are generally low-rank projections of the input, mapping features from \\mathbb{R}^d \\to \\mathbb{R}^{d’} with d’<d. In your proof, you need instead to expand the feature size.\n\nCorrections\n- In the future, I invite you to thoroughly number equations, even though you don’t reference them directly in the paper - especially for highly theoretical work! Being able to refer to Eq(x) rather than “the equation right after comment xxx in section yyy” makes reviewing much easier!\n- [l6] preceded by sum-of-linear transformations -> sum of linear transformations is still a linear transformation?\n- [l8] under the L_\\infty-norm -> in L_\\infty-norm\n- [l39] role of attention module -> role of the attention module\n- [l40] the sentence reads weirdly. Maybe dropping the first part: “presents the results that softmax-base …”?\n- [l61] It is not your approach that endows the architecture: it’s an inherent property of the architecture itself, as shown by your approach “we establish […] with the same approach, that also […] cross-attention is endowed with”\n- [l68] “early works […] focuses” -> work\n- [l70] propose -> proposeD\n- [l71] “a series of research” … works?\n- [l75] “make more careful estimation upon the numerical results of contextual mapping” -> I’m not sure what this means? “Provide a more careful estimate of the … ?”\n- [l75 vs l91] Highlight better how your result differs from Kajitsuka and Sato [2023]? It’s not just because you don’t make use of skip-connections, right?\n- [l76] and prove\n- [l79] \\alpha-smooth functions\n- [l80] dimensions\n- [l81] have achieved\n- [l82] results […] transformers\n- [l84] derived FOR attention-only networkS\n- [l86] of the attention\n- [l101] the maximum absolute element -> maximum (in absolute value) element\n- [l105] apply f on -> apply f to\n- [l107 - l111] for a self/cross-attention…layer?\n- [l111] I found the notation Zq Zk rather confusing: usually cross attention is written for X and Y, but it’s a personal preference\n- [l114] separated -> separate\n- [l129] For simplicity of presenting -> to simplify the presentation of\n- [l140] Why “is (one of) the highest”? Shouldn’t it be THE highest, according to Assumption3.1? In general, I’ve found this assumption a bit confusing: at the boundary regions between partitions, this cannot hold, no?\n- [l141] partitioned in “regions”\n- [l142] “is (tied for) the largest” tied to?\n- [Prop3.2] with the exception of\n- [Prop3.2] If I got it right, the Lebesgue measure should be small in \\mathbb{R}^d, not \\mathbb{R}^n?\n- [l232] approximately an one-hot -> A one-hot", "limitations": "The authors correctly highlighting some of the limitations of their approach (particularly, the inflation in parameter count the architectures described in their proof needs). At the same time, they fail to point out what in my opinion is the strongest limitation, namely that the linear operator they’re considering invalidates the transformer architecture.", "rating": 2, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 4, "significance": 1, "originality": 2, "strengths_and_weaknesses": "Strengths\n- Despite the technicalities involved, and the heavily theoretical aspect of the paper, the relevant theorems are outlined in a clear manner, their main ideas cleanly fleshed out, and the role of each step involved in the proof properly highlighted\n- In particular, that Attention can build max-affine partitions is an interesting insight\n\nWeaknesses\n- I am highly dubious about the generality of the overall procedure. While the reliance on the preceding Linear operator might seem like a minor technical detail, in practice it ends up dramatically reducing the impact of the proof, as it wouldn’t hold for classical frameworks where Attention is applied", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review. We have addressed all your comments and questions in the following responses.\n\n---\n### **W1, “Linear layer” part of Q1 & Q2** *The reviewer is concerned with the linear layer that is applied before the attention layer. The main reasons are as follows:*\n\n1. (In Weaknesses) *The linear layer wouldn’t hold in classical frameworks where attention is applied.*\n2. (In Questions) *The linear operator is not orthodox, which does not align with the structure of the linear layer in a transformer. Permitting such linear operations means allowing the model to duplicate the input sequence. Also, a bias term is added to the input sequences that varies along the sequence length. The Reviewer also raised concerns about the target function class not reflecting the permutation-invariant nature of attention.*\n3. (In Questions) *Based on the second point, the reviewer thinks that our result on attention might not transfer to the setting of transformers. The reviewer also argues that allowing such a sequence-wise linear layer to precede the attention layer makes the overall architecture more akin to a generic FFN, since it skips much of the constraints of a transformer.*\n---\n**Response:**\nThank you for your thoughtful review. Our work focuses on the approximation ability of attention and thus uses a sequence-wise linear transformation mainly to keep the structure concise and the method easy to follow, which is also the style we strive for in section 3. However, we don’t think it is *critically harmful* to the *genuineness and transferability* of our work.\n\n1. **Resemblance to FFN**\nWe wish to clarify that the linear layer in our network is *completely affine*, that is, it doesn’t incorporate any non-linear activations. This also means our approach of theoretical analysis is *vastly different* from ones used in UAP of FFN. In FFN, the non-linear activations are separate and hence the whole model can be dealt with in parts. In the structure considered in our work, the softmax function is the *only* non-linear activation and it cannot be broken into smaller counterparts.\n\n2. **The Necessity of the Linear Layer**\nIn our UAT for single-head attention, we do need to duplicate the input sequence. However, that does not harm the transferability of our technique. In fact, duplicating the sequence is only necessary if we wish to use one head to approximate the whole target function. In practical settings with multiple heads, we can configure the network into transformer-style setting. For example, if we divide the target function into cutoff functions on separate supports, when these supports are small enough, each cutoff function can be approximated by a head that preserves the sequence length under a given precision. We admit that in the single head setting sequence-wise operations are necessary, but it does not defy our work’s purpose of demonstrating a method to analyze attention. Also, it does not harm this method’s transferability into transformer settings as discussed above.\n\n3. **Permutation Invariance and the Position‑Dependent Bias**\nWe believe you mean permutation equivariance. Because we added a position-dependent bias, the permutation equivariance is thus removed. The position-dependent bias can be seen as positional encoding and is widely used in theoretical works like\n---\n### **Questions (minor).** *K/Q projections expand input dimension (\\$d' > d\\$), unlike low-rank projections in practice.*\n---\n**Response:** We acknowledge that our construction uses key/query projection matrices that expand the input dimension (\\$d' > d\\$), which deviates from the typical low-rank (\\$d' < d\\$) projections used in practice. However, **this choice is a deliberate and mathematically valid design for the purpose of our universality proof.** By expanding the feature space, the attention softmax can partition the input domain into a larger number of distinct regions. This enables more fine-grained max-affine partitions of the input. This finer partitioning is crucial for accurately representing complex target functions in our theoretical construction. Importantly, using a higher \\$d'\\$ does not compromise the theoretical integrity of our results. It remains within the standard definition of attention and is essential for achieving the universality guarantee.\n\n---\n### **Corrections.** *Fix scattered grammar/word‑choice issues, sharpen notation, clarify novelty vs. Kajitsuka & Sato [2023], tighten Assumption 3.1 wording and boundary cases, correct minor typos in Propositions 3.1–3.2.*\n---\n\n**Response:** Thanks for pointing these out! We have made the following modifications in our latest draft accordingly.\n\n- Grammatical edits include consistent tense, correct plurals (“dimensions”, “results”), and article use (“the attention module”) \n- Ambiguous phrases have been replaced (e.g., “sum‑of‑linear transformations” → “a linear transformation”; “series of research” → “research works”)\n- Notation is unified: \\$Attn\\_s\\$, \\$Attn\\_c\\$ are explicitly labeled as *layers*\n-  \\$Z\\_Q, Z\\_K\\$ are defined as query and key/value sequences \n- \\$\\lVert\\cdot\\rVert\\_\\infty\\$ is “largest absolute entry” \n- We also emphasize that our one‑layer, attention‑only network achieves universal approximation **without** skip‑connections or permutation‑equivariance, unlike Kajitsuka & Sato 2024 \n- Assumption 3.1 now rules out measure‑zero ties, so “the highest” is used throughout\n- Boundary comments are added \n- Proposition 3.2 references Lebesgue measure in \\$\\mathbb{R}^{d\\times n}\\$ and now reads “with the exception of a region of arbitrarily small measure.” \n- All remaining typos (e.g., “a one‑hot”) are fixed.\n\n\n---\n\nThank you once again for your constructive feedback, attention to detail, and encouraging kind words. Your review, particularly your suggestions regarding presentation and organization, has made this work more robust and accessible to a broader audience.\n\nPlease feel free to reach out if you have any further questions or need additional clarification about our work. Thank you!"}, {"review_id": "OtA0QN1nUv", "reviewer": "Reviewer_VjNz", "summary": "This paper presents a rigorous theoretical result showing that a single-head Softmax attention model, when combined with a single layer consisting of a sum of linear transformations, can approximate any continuous function defined on a compact domain to arbitrary accuracy. The approximation is established in the $L_\\infty$\n  norm and is further extended to integrable functions under the $L_p$ \n  norm. The core insight underlying the proof is that max-affine functions can partition a compact domain (which is expected) and that Softmax attention can be used to approximate the indicator functions for these partitions.", "questions": "See Weaknesses above.", "limitations": "As the authors acknowledged, the approximation results in this paper are entirely constructive. That is, the results show that there exists a specific single-head transformer with a specific weight configuration that can approximate any compactly supported target function. However, the result does not provide guidance on how to find or learn such a configuration in practice. This limitation is common in approximation theory, where existence  do not translate into practical algorithms. More important for learning-based applications is the question of generalization, that is, for a general class of transformers, the empirical risk minimizer can converge to the true regression function.  As far as I am aware, establishing generalization bounds for softmax-based architectures remains challenging, since many classical tools (such as the covering number argument) are difficult to apply due to the nonlinearity and normalization constraints inherent to Softmax attention.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strengths:\n\nA key strength of this paper is that, to the best of my knowledge, it is the first to establish a universal approximation theory for Softmax transformers. While the theoretical analysis is limited to a simplified setting, that is a transformer with one Softmax attention head, the result is nonetheless significant. It offers a valuable contribution to the growing body of work on the approximation power of attention mechanisms and provides a foundation for future studies to build upon in multi-head, and multi-layer architectures.\n\nWeaknesses:\n\n1. I was surprised that the paper does not mention or discuss a highly relevant prior work: [1]. That paper establishes that a ReLU attention head (as opposed to the Softmax one considered here) can be interpreted as a cubic spline, and leverages classical results on spline approximation to show that ReLU attention heads are universal approximators. As far as I know, the arguments used there do not directly extend to Softmax attention due to the fact that Softmax output sums up to one, which ReLU attention does not impose. Including a discussion of this distinction would highlight the contribution of the current paper, particularly because this work adopts an entirely different approach to establishing the approximation capabilities of softmax attention.\n\n2.  Related to the point above, I believe the paper underemphasizes an important aspect of its contribution, that is, to the best of my knowledge, the first to establish a universal approximation result for s\\Softmax-based transformers. All prior work cited in the paper focuses on ReLU-based attention mechanisms.\n\n3. I would hope the authors can discuss the possibility (or the challenges) of extending their results to multi-layer transformers with stacked Softmax attention heads. Since practical transformer architectures typically involve multiple layers, understanding whether the theoretical guarantees hold or how they scale in deeper settings would be a valuable addition to the paper.\n\nReference:\n[1] Lai, Zehua, Lek-Heng Lim, and Yucong Liu. \"Attention is a smoothed cubic spline.\" arXiv preprint arXiv:2408.09624 (2024).", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review. We have addressed all your comments and questions in the following responses.\n\n\n---\n### **W1.** *Missing citation of Lai et al. (2024) for ReLU-based attention UAP. Discuss differences vs. softmax attention.*\n---\n\n**Response:** Thank you for pointing this out. **We will add a citation to *Lai et al. (2024)* and discuss the distinction in our related work section.** Lai et al. prove that with ReLU-based attention, the transformer’s attention module can be viewed as a piecewise-polynomial (cubic spline) function, and under certain assumptions this implies a universal approximation property. \nIn essence, their analysis leverages the *max-affine spline* structure arising from ReLU activations, which partition the input space into linear regions. \n\nIn contrast, **softmax** attention introduces a normalization constraint that yields a *nonlinear convex weighting* of values. For example, a single-head softmax attention produces output \\$y\\$ as a convex combination of value vectors \\$v\\_j\\$:\n\n$$\ny =\\sum\\_{j=1}^n \\alpha\\_j v\\_j, \\qquad \\text{where } \n\\alpha\\_j = \\frac{\\exp(\\langle q, k\\_j\\rangle)}{\\sum\\_{l=1}^n \\exp(\\langle q, k\\_l\\rangle)},\n$$\n\nand thus \\$\\sum\\_{j=1}^n \\alpha\\_j = 1\\$. This smooth softmax mechanism does **not** create the hard, piecewise-linear partitions that ReLU-based attention (or a “hardmax” selection) would. Consequently, our proof required new techniques: we approximate a max-affine partition using softmax weights despite the coupling caused by normalization. We agree that clarifying these differences will highlight the novelty of our approach, which tackles the softmax-specific challenges (normalization and smoothness) absent in the ReLU-based spline analysis.\n\n---\n\n### **W2.** *Contribution underemphasized: first UAP result for softmax attention (distinct from prior ReLU-based results).*\n\n---\n\n**Response:** We appreciate the reviewer emphasizing this point. To the best of our knowledge, our work is indeed the **first** to establish a universal approximation theorem for a single-head *softmax* attention layer (with minimal linear pre/post-processing). We will make this significance more explicit in the revised paper. Prior universal approximation results in the literature either focused on attention variants with ReLU/hardmax-style activations (piecewise-linear splines as in Lai et al.) or required adding extra network components (e.g. **Kajitsuka & Sato (2023)** needed one attention layer *plus* two feed-forward networks to achieve universality). In contrast, our result shows that the standard *softmax attention mechanism alone* (with a single head and one layer) has sufficient expressive power. This novel finding matters because it theoretically validates the expressiveness of the vanilla softmax attention used in practice, without resorting to simplifications or additional architectures.\n\n---\n\n### **W3.** *Multi-layer extension: discuss if/how the result extends to stacked softmax attention heads (deep Transformers).*\n\n---\n\n**Response:** We agree that extending the result to *multi-layer* (stacked) softmax attention is an important consideration. In principle, if a single-layer, single-head attention is a universal approximator, then a multi-layer architecture (with multiple attention heads per layer, as in real Transformers) should be **at least as expressive.** Additional layers cannot reduce expressive power, and could potentially approximate complex functions with fewer units per layer. Indeed, one could always let one layer carry out the construction from our proof while the other layers perform identity mappings. Thus, we expect that *deep transformers* also retain universal approximation capability.\n\nThat said, formally **extending the proof** to multiple layers is non-trivial. The challenge is that the composition of softmax attentions may not preserve the same neat max-affine partition structure we exploited in the single-layer case. Each softmax layer produces a smooth convex combination of its inputs rather than a strict partition. Stacking them can yield highly entangled representations where the “partitioning” of the input space from one layer is *distorted* by the next. This means a direct induction on layers would require careful new analysis. Additionally, while depth can theoretically be traded for width, a multi-layer constructive proof would need to coordinate attention heads across layers (and ensure the errors don’t compound). This is a significant technical leap beyond our current scope.\n\nIn summary, we believe multi-head, multi-layer Transformers are also universal approximators (since one layer already suffices), but **proving** it rigorously would demand new ideas to handle the composition of softmax layers. We will add a brief discussion to acknowledge this extension. We also note that analyzing deeper architectures remains an interesting open challenge (e.g., how successive softmax layers might refine or interfere with the learned partitions). This point will be clarified in the paper, and we thank the reviewer for the insightful suggestion.\n\n---\n### **Limitations.** *Result is constructive but not learnable in practice. Generalization bounds for softmax Transformers are still challenging.*\n---\n\n**Response:** Thanks for pointing these out. To clarify, our goal is **expressivity**, not learnability. The theorem shows that a single‑head softmax Transformer *can* approximate any compactly supported function. It does **not** claim gradient descent will find those weights or provide sample‑complexity guarantees.\n\nRecent progress begins to address these gaps:\n\n* **Learnability.**  Hu et al. (NeurIPS ’24) give a covering‑number analysis of softmax attention. Li et al. (’24) prove GD learns a 1‑NN rule with one softmax head.\n* **Generalization.**  Capacity and covering‑number bounds for softmax Transformers appear in 2110.10090, 2404.03828, 2407.01079, and 2411.17522\n\nWe will cite these works and note that extending universal approximation to provably learnable and generalizing algorithms is an open, promising direction.\n\n\n---\nThank you again for the detailed review! We're open to any further questions or clarifications you might have about our work."}, {"review_id": "xM7M3dVZRv", "reviewer": "Reviewer_eNoW", "summary": "The paper considers the expressivity of a sequence to sequence map modeld by an affine linear transformation and a attention layer and finds that this 'shallow' architecture already achieves the universal approximation property for continuous functions on a compact set of input sequences of length $n$ and dimension $d$ in each step and output sequences of the same dimension with respect to the $\\infty$-norm. This is achieved by dividing the input space into a grid with sufficiently small 'cube' cells (in high dimension) and constructing a set of linear affine functions $a^\\top_j\\tilde Z+b_j$ where $a_j$ is the center of the $j$-th cube and $b_j=-\\frac{1}{2}\\|a_j\\|^2$. $\\tilde Z$ is the flattening of $Z$ to a vector. As this expression is the minus linear affine part of $\\|\\tilde Z-a_j\\|^2$, and the quadratic part $\\frac{1}{2}\\| \\tilde Z\\|^2$ of this expression is the same for all $j$, the linear affine function with index $j$ is minimal whenever $\\tilde Z$ lies in the cube centered by $a_j$. Therefore, the authors construct a grid-partition of the input space. They then design the key and query matrices in a way that essentially $K^TQ_{:,j}=-R\\frac{1}{2}\\|\\tilde Z-a_j\\|^2$ where $R>0$ is a parameter. This is then inserted in the softmax function, where again the quadratic term drops out, so that, if a input point is inside the $j$-th cube and $R$ is large, essentially only one column of the softmax matrix is getting activated. By another linear transformation for V, the authors then take care that these 'close to indicator function' values get multiplied with the correct results for the output sequence. using the fact that the softmax effectively implements a partition of unity, the authors prove the UAP for continuous functions. Also UAP in $L^p$ is stated (although this trivially follows from UAP in $C$ as $C$ is dense in $L^p$). Also, a similar argument is given for the cross attention.  Proofs are given in an extensive  appendix, which also contains some numerical examples..", "questions": "Are you able to revise the paper carefully, write it much more clearly and delete unnecessary and faulty parts in the rebuttal phase? Then i would support the paper. Otherwise I would recommend a careful revision and a probably successful re-submission elsewhere. Probably the revision can only be successful if you already worked on this in the mean time.", "limitations": "Nothing.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 2, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Streangths:\n* Given the relevance of transformer architectures for contemporary machine learning, the UAP for a single layer composed with linear affine transformation has a high relevance. Other results on the UAP of transformers have been published, but the core content of this work is certainly worth publishing also at very high level.\n* Underlying idea is intriguing and very insightful, which also strengthens the case of this paper.\n* The cross attention 'bonus' result is also novel and can be proven with minor adaptions to the strategy of the self-attention case.\n* The authors mastered quite some computations.\n\nWeaknesses:\n* While the paper's result is very strong and deserves publication, the present state of the paper's presentation is not equally good - unfortunately by a large margin. This starts with a number of completely unnecessary parts as essentially all of Section 3 which could just be deleted without any loss for the core messages of the paper (piecewise constant functions are in $L^p$). The proofs for this section given in the appendix have technical problems as assumptions on $a_j$ and $b_j$ are missing, under which the 'results' immediately break down (take e.g. all these quantities to be identical). Such easy errors are embarrassing and unnecessary as this 'shaky style' does not influence the main theorems logically. \n* The proof of the main theorem is quite unstructured and fails to explain the main ideas behind the proof effectively. The readers have to follow lengthy parts of the appendix which are either irrelevant (like proofs of section 3) or tivial/well known (like the proofs of the UAP for $L^p$ where the authors repeat well known stuff based on Lusin's theorem, forget to apply Tietze's extension theorem as an example for an extremely complicated an incomplete proof of something that could just be cited form almost any book on function spaces) or lines like lines 704 and following which elaborate trivial stuff over almost a page.\n* This said, also the main core of the proofs - as far as I can see containing beautiful ideas and being essentially correct - could be worked through in a much better and more structured fashion except for just computing.  \n* Unfortunately the attitude of the paper prefers  convoluted key facts, hardly understandable descriptions of steps to poofs which only make sense when one has read and understood the actual proof in the appendix (and thus does not need further explanations).\n* Layout of the appendix is messy as well with many formulae exceeding margins, gaps etc.\n* English sometimes is faulty and requires further proofread. \n \nMinor remarks (no mention in the rebuttal required)\n*l 108: Define the way the softmax function acts on matrices precisely\n*Section 3: This entire section is rather misleading and contains imprecision like the $\\Delta$-story, which is  not properly resolved here and properly in section 4. Partition, Def. 3.2, Proposition 3.2 are all not entirely wrong but faulty and only repaired with ad hoc assumptions. this unnecessarily draws down the level of the paper.\n* Proposition 3.2 - is there a mismatch between the dimensions in this Proposition where $X_i\\in\\mathcal{X}$ whereas later it seems to be $X\\in\\mathcal{X}$?!\n*l. 147: This is a well known result from convex analysis/theory of Legendre transforms and not due to Kim and Kim 2022. Cite properly.\n*l. 178: You can't re-define the $\\|\\cdot\\|_\\infty$ norm. Saying that this norm only holds up to an exceptional set means that you say it does not hold at all. Give a proper formulation with proper assumptions on $a_j$ and $b_j$ or better leave this out like the rest of Sec 3.\n*l.181 and following. Explanations cant be understood. Better just refer to proof 8or even better just leave out).\n*l. 218: The repetition of the content of Theorem 4.1 is superfluous and redundant - delete.\n* l.223: Leave out the steps and refer to a concise proof.\n* l.246: $P$ not defined\n* l. 248: technical Highlight is not needed - delete (like other technical highlights). Also the steps don't help. Focus on good and worked though proofs. If they are well written, your pedagogy isn't needed. \n* Corollary 4.4.1 why does all over sudden the numbering change from tow to three digits. The Corollary is a trival consequence of Thm 4.1 and can be omitted or put into one sentence in plain text. Also Lusin's Thm is not enough (see above).\nl. 488: I understand the temperature thing but other readers might not. Express this simpler. The entire paragraph is over complicated.\nl. 507: There is no need to validate a mathematical theorem experimentally. rephrase as \nl.517: Proposition D1 as stated is wrong/lacking assumptions, see above\nl.524: Proposition D2: same unfortunate usage of the $L^\\infty$ norm with exceptions  - does not male sense\nl.535: Remark D1 is illegitimate - leave out      \nl.539: Check Matrix dimensions between matrices and bias in (D.2), there seems to be a typo $n->N_{MaxAff}$\nl.534 and following: Would be helpful to read if the dimensions of the matrix objects always would be mentioned. It is not evident if $a_j$'s are row or column vectors etc. \nl. 534 and following: Treatment of the exceptional set always sloppy.\nl.578: Just prove Thm 4.1, a second list of steps is really superfluous.\nl. 591: akin?\nl. 607: Proof of what?\nl. 624: This is just well-known and does not require a 'poof' (which then is not given)\nl. 631: As the e-notations does not work on the boundaries of the cubes- rather avoid using it\nRemark E2: I find this unnecessary\nl. 704 and following: This is too lengthy for a trivial result.\nE.2 can be deleted and cited\nRemainder: I didn't check the Cross attention results line by line - revise yourself.", "was_revised": true, "final_justification": "The changes in the presentation that the authors announce take into account my main concern. The paper's main result is interesting, as it points towards the expressivity of attention layers in connection to an affine transformation proving an 'sequence to sequence ' UAP theorem.  I agree that the differences to standard attention layers should be made clear in the text.", "rebuttal": "Thank you for your detailed review. We have addressed all your comments and questions in the following responses. \n\n---\n### **Q1.** *Are you able to revise the paper carefully, write it much more clearly and delete unnecessary and faulty parts in the rebuttal phase?*\n---\n**Response:**\nYes. We promise and have already conducted, and will continue to refine a **line‑by‑line revision**:\n\n* **Main proof:** fully re‑structured into concise conceptual steps. All lengthy algebra moved to the appendix.\n* **Appendix:** redundant material deleted, English polished, typos corrected, and every formula resized to fit the margins.\n* **Section 3:** rewritten as a self‑contained intuition module with the tie‑breaking Assumption 3.1 made explicit. It can stay in the main text, be moved to the appendix, or be removed entirely.\n* **Ongoing updates:** the draft has been under continuous revision since submitted because the results serves as the basis of several active projects in our group.\n\nAlthough policy prevents us from uploading the new PDF at rebuttal time, these revisions are complete and will appear in the final version.\n\n***Since no PDF sharing is allowed, all revisions will only appear in the final version. Thus, our responses below are all in future tense (even if the revisions had been made).***\n\n---\n### **W1.** *Section 3 is unnecessary. Appendix proofs have missing assumptions on affine coefficients.*\n---\n\n**Response:** We respectfully clarify that **Section 3 is included to provide intuition** for our main results. In this section, we illustrate how attention can act as a max-affine partitioning mechanism and approximate an indicator function over those partitions. We wrote section 3 in consideration of the diverse potential readers of this conference. We hope people with relatively weak theoretical background are still capable of understanding the intuition of our work. That said, we acknowledge that Section 3 is not strictly required for the proofs. **We are willing to streamline or relocate Section 3** (for example, moving some of its content to an appendix) to maintain focus on the core results, if the final publication format demands it.\n\nFor the assumption that regards  \\$a_j \\$ and \\$b_j \\$,  we do have noticed the assumption in the appendix isn't brought up in a standard and straightforward way. We do not think it is missing (as mentioned in remark D.1 in the submitted version), but we do acknowledge that it is not presented in a rigorous manner (that is, explicitly stated with mathematical expression and addressed as \"assumption\" instead of \"remark\"). We are sincerely sorry for any confusion caused. We have refined it's form and fix the above mentioned presentation problems. \n\nFinally, we are willing to **trim or relocate the proofs from Section 3** to the appendix (or further condense them in the appendix). So that Section 3 takes a smaller proportion and the main text remains focused. \n\n\n---\n### **W2.** *Main proof is unstructured. Derivation to \\$ L_p\\$ is trivial and redundant.*\n---\n\n**Response:** We thank the reviewer for this valuable feedback. We agree that the proof of Theorem 4.1 can be presented more clearly. In the revised version, we will **reorganize the proof of Theorem 4.1** to highlight its key ideas in a step-by-step, modular fashion. Specifically, we will break the proof into clearly-labeled steps focusing on: \n- (1) constructing the max-affine partition of the domain, \n- (2) approximating indicator functions on each partition region, and \n- (3) reassigning function values on these regions to complete the approximation. \n\nThis structured approach will make the logical flow transparent and ensure that important ideas do not get buried in technical details.\n\nFor the \\$L_p\\$ derivation, we do admit it is obvious to people with a basic knowledge of the measure theory. The proofs of these derivations are simple and short so we don't think they cause very serious redundancy problems. The proofs are meant for people who aren't familiar with the measure theory, but we are open to cutting them off and make the \\$L_p\\$ derivations self-evident corollaries given in remarks.\n\n\n---\n### **W3.** *The core proof ideas are strong but obscured by long symbolic computations. They should be reorganized around conceptual steps.*\n---\n\n**Response:** We agree that the proof of our main theorem is currently obscured by excessive symbolic computations. In the revision, ***we will reorganize this proof around its conceptual mechanism and key ideas***, rather than lengthy algebra. Detailed computations will be either summarized or moved to the appendix, allowing the main text to focus on the clear, structured development of the core argument. This change will highlight the underlying concepts and make the proof more accessible to readers.\n\n---\n### **W4.** *Key steps are described in a convoluted way and assume readers have already parsed the full appendix. The exposition is circular and hard to follow.*\n---\n**Response:** We acknowledge that the original exposition of some key steps was convoluted and hard to follow without the appendix. ***We will rewrite these parts of the proof to be self-contained and easily understandable***, so that readers do not need to consult the appendix to grasp the logic. Every crucial step will be explained clearly in the main text, and we will remove any circular references. This revision will make the argument flow in a straightforward, reader-friendly manner.\n\n---\n### **W5.** *The appendix layout is messy: formulas overflow margins, formatting is inconsistent, etc.*\n---\n\n**Response:** ***We will address the formatting issues in the appendix.*** \n\nAll formulas will be adjusted to ensure they fit within the margins, and we will enforce consistent formatting throughout the appendix. This includes fixing any overflow, aligning equations properly, and using uniform styles for text and math. The updated appendix will have a clean layout with no margin overflow and improved consistency.\n\n---\n\n### **W6.** *The English is sometimes faulty and needs proofreading.*\n\n---\n\n**Response:** We have conducted 3 extra rounds of proofreading the entire manuscript for grammar and clarity. All grammatical errors and awkward phrasings will be corrected in the next iteration. By the final version, the paper will read smoothly and meet high standards of English usage.\n\n---\n### **Corrections.** *Section 3 lacked rigor, Proposition 3.2 was ill‑specified, norms were re‑defined, notation (\\$X\\_i,\\mathcal{P},|\\cdot|\\_\\infty\\$) inconsistent or undefined, Theorem 4.1 and corollaries were repeated, trivial “technical highlights” and proof steps cluttered the text, matrix dimensions in (D.2) unclear, exceptional sets handled informally, validation logic for Lusin’s theorem and experiments opaque.*\n---\n\n**Response:** To reiterate, we have made the following modifications:\n\n- Simplify the proof sketch part of Section 3 to avoid over-specifying.\n- Redundant norm definitions have been removed, and notation is unified: \\$X\\_i\\in\\mathcal{X}\\$ and a single partition symbol \\$P\\_{\\text{ma}}\\$ replace ambiguous uses of \\$\\mathcal{P}\\$. \n- Repetitive material around Theorem 4.1 has been consolidated\n- corollaries that follow immediately are now mentioned only in prose. \n- Superfluous technical highlights and step‑by‑step derivations have been deleted or moved to an appendix. \n- Appendix D labels every block matrix with its shape, resolving dimension ambiguities in \\$(D.2)\\$. \n- Exceptional sets are further checked rigorously and revamped several obscure descriptions. \n- Assumptions on Proposition 3.2 in the appendix is now stated in a more rigorous fashion. \n- Extra rounds of proofreadings \n- Finally, the use of Lusin’s theorem and the experimental validation are clarified, showing how partition error converges as predicted. \n\nThese edits address all reviewer concerns while making our draft shorter and clearer.\n\n--- \nThank you again for the kind words, and detailed review and proofreading (wow!).\nWe believe the above changes and clarifications address all the points raised. The draft has been updated to reflect these corrections.  Your efforts have helped us improve the clarity and quality of our paper. We are eager for further discussions!"}, {"review_id": "kyO0X0Fuk1", "reviewer": "Reviewer_uiVo", "summary": "This paper shows the universal approximation capability in continuous function spaces under both $L^\\infty$ and $L^p$ (for $1\\leq p<\\infty$) norm of single-layer, single-head self- and cross-attention mechanisms with minimal attached structures.", "questions": "- In practice, the input sequence is first processed by a self-attention layer, finally followed by a linear layer. However, in this work, the universality is established for Attn∘Linear(X), which is the opposite order compared to practical implementations. How to comment on this?\n- This paper shows the universality of the Attention mechanism by approximating piece-wise constant functions because the set of piece-wise constant functions is dense in the continuous function space. By such, it seems the number of the parameters of the constructed network is of order $\\mathcal{O}(\\epsilon^{-d})$ where $d$ is the input dimension? \n\nMinor.\n- I understand the definition of the softmax function is well-known but providing it in Sec. 2 will make the paper smoother. \n- Typos. In Prop. 3.3, $\\mathbb{R}^d \\to \\mathbb{R}^{d_{out}}$ rather than $\\mathbb{R}^d \\to \\mathbb{R}^{d}_{out}$?", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**. This paper is well-written and well-organized. The result is very cute and interesting with potential significant contribution to the machine learning theory field.\n\n**Weakness**. While this paper establishes the universality of attention mechanisms with minimal attached structures, the approximation rate (or the upper bound of the number of parameters) is not discussed, which is theoretically important. See **Question** part for details.\nOverall, there is no major weakness in my opinion.", "was_revised": true, "final_justification": "After the rebuttal and discussions, I will keep my positive evaluation for the paper.", "rebuttal": "Thank you for your detailed review. We have addressed all your comments and questions in the following responses.\n\n---\n### **Q1.** *Why does the theoretical result use Attn∘Linear instead of Linear∘Attn as in practice?*\n---\n**Response:** The paper deliberately reverses the order for theoretical convenience. It applies a linear (feed-forward) layer first and then attention (`Attn∘Linear`), whereas practical transformers do attention then linear. This reversal simplifies the universality proof by giving the attention mechanism a richer, “position-mixed” input to work with. In effect, the initial linear layer acts like an embedding or positional encoding step. It injects positional information and mixes the sequence tokens before self-attention is applied. This is crucial because vanilla self-attention is permutation-invariant (it can’t distinguish token order without position info). By encoding position and content via the linear layer first, the attention module can attend over a sequence where each token representation carries positional context, granting the attention full access to order-specific patterns. \n\nIn short, placing the linear layer first is a theoretical choice that ensures the attention sees an input enriched with positional cues (much like adding positional encodings), which is essential for expressive power. It also clarifies the roles in the proof: the pre-attention linear layer provides necessary position-dependent mixing, and then the attention layer performs contextual weighting over the entire sequence. This makes the one-layer transformer mathematically easier to analyze while retaining universality. Namely, the attention can leverage the position-encoded sequence to approximate any function (rather than being limited to order-invariant functions), and the final feed-forward can then map the attended context to the desired output.\n\n---\n### **Q2.** *Does the parameter count scale as $\\mathcal{O}(\\varepsilon^{-d})$?*\n---\n\n**Response:** Yes, you are correct. To approximate a general continuous function in $d$ dimensions within error $\\varepsilon$, the construction essentially partitions the input space into fine grid cells of side length about $\\varepsilon$. The number of such small regions grows on the order of $(1/\\varepsilon)^d$ (since in \\$d\\$ dimensions an \\$\\varepsilon\\$-grid yields roughly \\$\\varepsilon^{-d}\\$ cells). \n\nThe proof indeed leverages this kind of partitioning (approximating the target function by piecewise-constant outputs on a very fine grid), so the total parameters needed scale as \\$\\mathcal{O}(\\varepsilon^{-d})\\$ for a worst-case \\$d\\$-dimensional function. This is in line with standard uniform approximation results. It reflects the “curse of dimensionality.” \n\nIn other words, without additional structure or smoothness assumptions, the number of parameters (or regions) required to achieve error \\$\\varepsilon\\$ typically grows exponentially in \\$d\\$. This mirrors other universal approximation schemes (e.g. using ReLU networks or Fourier series), which also require on the order of \\$(1/\\varepsilon)^d\\$ basis functions or neurons to approximate an arbitrary continuous function on a \\$d\\$-dimensional compact domain to accuracy \\$\\varepsilon\\$. Therefore, the parameter count in the transformer’s universality proof does scale as \\$\\mathcal{O}(\\varepsilon^{-d})\\$, consistent with expectations from uniform approximation theory and existing UAP proofs.\n\n---\n### **Q3 (minor).** *Consider including the definition of the softmax function in Section 2, even if it's well known, to improve flow.*\n---\n**Response:** We have added a formal definition of the softmax function in Section 2, as suggested. This inclusion ensures clarity and improves the flow of the presentation. Specifically, we now define the softmax function as \\$\\mathrm{softmax}(z)\\_i = \\frac{\\exp(z\\_i)}{\\sum\\_j \\exp(z\\_j)}\\$ when it is first introduced in Section 2.\n\n---\n### **Q4 (minor).** *Typo in Proposition 3.3: should the output space be written as \\$\\mathbb{R}^{d\\_{\\text{out}}}\\$ rather than \\$\\mathbb{R}^{d}\\_{\\text{out}}\\$?*\n---\n**Response:** Yes are absolutely correct. This was indeed a typo. We have corrected the notation in Proposition 3.3. The output space is now written as \\$\\mathbb{R}^{d\\_{\\text{out}}}\\$, as suggested by the reviewer.\n\n---\nThanks again for your kind words and detailed review! We have taken all comments into careful consideration and have made corresponding revisions to address the concerns raised.\n\nWe look forward to further feedback and discussion!"}]

Author rebuttal:

Meta-review: This work shows that a single-head Softmax attention model, when combined with a single layer (sum of linear transformations), can approximate to any desired accuracy any continuous function on a compact domain. With the exception of one reviewer, the reviewers appreciated the contribution, both technical and conceptual, and recommended acceptance. Based on my own reading and understanding of the work, the remarks (and associated score) of the outlier reviewer are unjustifiably adversarial. Consequently, I believe this paper merits acceptance.
Decision: Accept (poster)