Title: Attention Mechanism, Max-Affine Partition, and Universal Approximation
Venue: NeurIPS 2025
Accepted: yes
Keywords: Attention

Universal Approximation
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

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 ∞ -norm. Furthermore, we extend this construction to approximate any Lebesgue integrable function under L p -norm for 1 ≤ p < ∞ . Lastly, we also extend our techniques and show that, for the first time, single-head cross-attention achieves the same universal approximation guarantees.

1 Introduction

We establish the universal approximation capability of single-layer, single-head self- and crossattention mechanisms. Departed from prior studies, our results demonstrate that the expressive power of transformers arises from only the (softmax) attention module and an attached linear layer, without additional components such as positional encodings or feed-forward networks (FFNs). More importantly, our proofs show that sequence-to-sequence universal approximation requires only a minimalist configuration: single-layer, single-head attention with linear 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 as BERT [Devlin, 2018], ChatGPT [Brown et al., 2020, Achiam et al., 2023], and LLaMA [Touvron et al., 2023a,b, Dubey et al., 2024] for language; ViT [Dosovitskiy et 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 to these successes is the attention mechanism . While numerous variants and implementations exist [Tay et al., 2022], the softmax-based vanilla attention

[Vaswani et al., 2017] remains a 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 to which softmax attention alone drives performance 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, a growing body of work explores its memory capacity [Mahdavi et al., 2023, Kim et al., 2023, Kajitsuka and Sato, 2024], universal approximation properties [Yun et 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 et al., 2022]. However, these studies often rely on additional components, such as feed-forward networks (FFNs) or multi-head setups or 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: softmax-based attention alone already suffices for universal approximation of sequence-to-sequence functions. We operate under three key premises for investigating the expressiveness of attention:

We focus on softmax -based attention,

We seek a minimalist design (a single layer of single-head attention plus a linear transformation),

We impose minimal assumptions on the data distribution or network architecture (no positional encodings, no multi-head expansions, no FFNs).

We provide new proofs that a single self-attention layer approximates any continuous sequence-tosequence function on a compact domain, in both the L ∞ and L p norms. Furthermore, we show, for the first time, a 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 maxaffine construction. This partitioning allows attention to assign distinct outputs to each partition cell. This perspective clarifies how softmax-based attention enables a powerful piecewise-linear approximation scheme.

Single-Layer, Single-Head Self-Attention Universality. We prove that a single selfattention layer is a universal approximator for continuous sequence-to-sequence functions on compact domains. Our results cover both L p - and L ∞ -norms guarantees and require minimal assumptions on data and architecture, highlighting the inherent expressive power of attention alone.

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 approximation 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 2 presents the ideas we built on. Section 3 shows our interpretation of Attention as a Max-Affine Partition in a simplified setting. Section 4 presents our universal approximation results for single-layer, single-head self- and cross-attentions.

1.1 Related Work

Universal Approximation. Early works of universal approximation theorems focuses on the expressiveness of feed-forward networks (FFN) [Cybenko, 1989, Hornik, 1991, Carroll and Dickinson,

1989]. Since Vaswani et al. [2017] propose the transformer architecture and the scaled dot-product attention module, there is a 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 mapping. 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 /δ ) dn /n !) ), Kajitsuka and Sato [2023] make more careful estimation upon the numerical results of contextual mapping and proves that with skip connections, a one-layer transformer is capable of approximating any permutation equivariant continuous function. Takakura and Suzuki [2023] add positional encoding to lift the restriction of permutation equivariance, and demonstrate a one-layer transformer approximates shift-equivariant α -smoothness function with an error independent of input and output dimension. Jiang and Li [2023] give a 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-based proof that softmax attention alone (no FFN) is a universal approximator for continuous sequence-to-sequence 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 on learning constrained boolean functions. Building on [Hu et al., 2025a], Hu et al. [2025b] show that a fixed two-attention-layer softmax Transformer is prompt-programmable: it emulates any algorithm implementable by a single attention layer (cf. [Bai et al., 2023]), providing a constructive account of one-model-many-tasks behavior with softmax (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-to-sequence mappings.

2 Preliminaries

We now present some ideas we built on.

Notation. For a vector v , we denote its i -th entry by v i and its subvector from the i 1 -th to the i 2 -th entry (inclusive) by v i 1 : i 2 with i 1 < i 2 . For a matrix M , we use M i,j for the entry in the i -th row and j -th column, M i, : for the i -th row, and M : ,j for the j -th column. The submatrix spanning rows i 1 through i 2 and columns j 1 through j 2 is denoted by M i 1 : i 2 ,,j 1 : j 2 with i 1 < i 2 , , j 1 < j 2 . We define c a × b as an a × b matrix with constant entries c , and abbreviate c a × 1 as c a . For norms, we define ‖ · ‖ ∞ as the maximum absolute element in a vector or matrix. The p -norms are given by ‖ v ‖ p = ( ∑ i | v i | p ) 1 /p for a vector v and ‖ M ‖ p = ( ∑ i,j | M i,j | p ) 1 /p for a matrix M . For function norms, we define the L ∞ norm as ‖ f ‖ L ∞ := sup x ∈ X f ‖ f ( x ) ‖ ∞ , where X f is the input domain of f , and the L p norm as ‖ f ‖ L p := ( ∫ x ∈ X f | f ( x ) | p p , dx ) 1 /p for 1 ≤ p < ∞ . For functions, when a function f : R → R is applied on a vector or a matrix, it means to apply f on every entry of the vector/matrix (i.e., exp([ a 1 , a 2 ]) := [exp( a 1 ) , exp( a 2 )] ).

Self-Attention and Cross-Attention Layers. For a self-attention Attn s : R D × N → R D × N out , and any input Z ∈ R D × N , we define its output as:

where W K , W Q ∈ R d Attn × D , W V ∈ R D × D , W O ∈ R N × N out . Here d Attn stands for the hidden size of the attention block. N out stands for the output sequence length.

For a cross-attention Attn c : R D × N × R D × N → R D × N out and any input Z K , Z Q ∈ R D × N , we define its output as:

Here W K , W Q , W V , W O are defined as those in self-attention.

Since we provide separate discussions for self-attention and cross-attention 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 : R D 1 × N 1 → R D 2 × N 2 to denote a layer of sum of linear transformations. For any input Z ∈ R D 1 × N 1 , we define its output as follows:

where P i ∈ R D 2 × D 1 , Q i ∈ R N 1 × N 2 for i ∈ [ H ] , R ∈ R D 2 × N 2 . Here H is a positive integer which denotes the number of linear transformations to sum.

3 Attention as Max-Affine Value Reassignment

In this section, we introduce a new interpretation of attention as a value reassignment to a max affine function. Essentially, we show that attention prepended with a Linear layer is able to reassign values to a partition generated by a max-affine function. We start with the below definition.

Definition 3.1 (Max-Affine Function) . Let X ⊂ R d x be a domain, and fix a positive integer N ma . For each i ∈ [ N ma ] , define an affine function y i : X → R for all x ∈ X :

The max-affine function MaxAff : X → R corresponding to affine functions { y i ( · ) } i =[ N ma ] is defined as

Intuitively, a max-affine function selects, at each point x ∈ X , the largest output among N ma affine functions. Geometrically, each affine function y i ( x ) = a /latticetop i x + b i defines a hyperplane in R d x +1 . Thus, MaxAff follows the highest hyperplane at each x , forming a 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 for all results in this section:

Assumption 3.1. For any max-affine function MaxAff , we exclude situations where the difference between its largest and second-largest affine components is smaller than a specified threshold. (Please see proofs for explicit definition.)

We do not apply this assumption in other sections.

3.1 Max-Affine Partition

We now show that a max-affine function MaxAff( · ) induces a partition of its input domain X . Specifically, the input domain X is divided up according to which affine function is the maximum at each point x . To be concrete, we define this partition as follows:

Proposition 3.1 (Max-Affine Partition) . Following Definition 3.1, consider a max-affine function MaxAff( x ) = max i ∈ [ N ma ] { a /latticetop i x + b i } , and let X ⊂ R d x be its input domain. Then MaxAff generates a partition on X :

We call the partition P ma the max-affine partition of X induced by MaxAff .

Intuitively, U i is the set of all point x for which the i -th affine function a /latticetop i x + b i 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 of) the highest among all components. Hence, the input domain X becomes partitioned 'regions' { U i } i =[ N ma ] . That is, if a point x belongs to a region U i , the corresponding affine function a /latticetop i x + b i is (tied for) the largest. Please see Appendix D.1 for a detailed proof.

Set Overlaps and Boundaries. By construction, every x ∈ X lies in at least one of the sets { U i } , but it may belong to multiple sets if several affine components attain the same maximal value. Hence, the collection { U i } is generally a 'partition' in an informal sense: while each U i 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-affine partitions.

Definition 3.2 (Indicator of Max-Affine Partition) . Following the notations in Proposition 3.1, for a max-affine partition { U i | i ∈ [ N ma ] } , we define i x := argmax i ∈ N ma ( y i ( x )) to be the label of the maximal affine component. Then, we define the indicator E : R d x → R N ma as:

which is a one-hot vector whose only non-zero entry is the i x -th one.

Namely, each component of E ( x ) is zero unless it corresponds to an index achieving the maximum, in which it has the value of 1 . In Figure 1a, we show an example of the max-affine partition.

14

12

10

8

6

4

2

0

-2

0

2

3

2

2

-2

0

X

2

3.2 Attention Scores Encode Indicators for Max-Affine Partition

We now discuss the connection between self-attention and a max-affine partition. We show that self-attention with a Linear layer attached before it can generate a max-affine partition. Further, for every input token, the attention score matrix approximately indicates which part of the partition it belongs to. We state this result as follows:

Proposition 3.2 (Attention Approximates Indicator of Max-Affine Partition) . Let X = [ X 1 , X 2 , · · · , X n ] ∈ R d × n denote any input sequence. We use X to denote the domain of all X i , i ∈ [ n ] . Let MaxAff be any max-affine function on X with N ma components, and let /epsilon1 > 0 be any positive real number. We define P ma = { U i | i ∈ [ N ma ] } as the max-affine partition generated by MaxAff as in Proposition 3.1. Then, there exists a Linear layer and a self-attention Attn whose attention matrix satisfies:

with exception of a region of arbitrarily small Lebesgue measure in R n . Here W K , W Q are the attention weights within Attn . W O only truncates the irrelevant part of the attention score matrix.

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 of 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 a single-head self-attention is capable of completing this task.

Proposition 3.3 (Attention Reassigns Value to Max-Affine Partition) . Following the notation in Proposition 3.2, Let F : R d → R d out be a piece-wise constant function which is separately constant on each U i , i ∈ [ N ma ] . We show that for any /epsilon1 > 0 , there exists an self-attention Attn such that

for every X in X with exception of a region of arbitrarily small Lebesgue measure in R n .

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 to Function on All Tokens) . In this section, for the conciseness in demonstration of 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-wise 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, we emphasize that here the approximation excludes a 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 next section.

Figure 1b provides us an example of Proposition 3.3.

4 Single-Layer, Single-Head Attention Achieves Universal Sequence-to-Sequence Approximation

In this section, we present our main results:

A single layer of single-head self-attention preceded by one linear layer is a sequence-to-sequence universal approximator for continuous functions on any compact domain.

A single layer of single-head cross-attention preceded by one linear layer is likewise a sequence-to-sequence universal approximator for continuous functions on any compact domain.

Importantly, we achieve attention-only universal approximation for both the L p -norm and L ∞ -norm, whereas most existing results apply only to the L p -norm 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 our results for self-attention in Section 4.1 and for cross-attention in Section 4.2.

4.1 Single-Head Self-Attention as a Universal Seq-to-Seq Approximator

We now present our main result: a single-layer, single-head self-attention module, combined with a linear transformation, is sufficient to approximate any continuous map f : R d × n → R d × n on a compact domain U ⊆ [ -D,D ] d × n . We present the result first in terms of the L ∞ norm for continuous f and then extend it to L p integrable functions.

Theorem 4.1 ( L ∞ -Norm Universal Approximation) . Let f : R d × n → R d × n denote any continuous function on a compact domain U ⊂ R d × n and let /epsilon1 > 0 be any positive real number. There exists a self-attention Attn with a prepended Linear layer, such that

Theorem 4.1 indicates that a single-layer self-attention block, combined with a linear preprocessing layer Linear , approximates sequence-to-sequence f in the L ∞ -norm.

Overview of Proof Strategy. We adopt a proof strategy based on a key observation: self-attention is capable of approximating target functions via implicit MaxAff operations. Our proof consists of the following 4 steps:

Step 1: Partition Input Domain U via MaxAff . Construct a max-affine function MaxAff over U (i.e., input domain of target function f ) such that this MaxAff induces a partition of sizeN ma of U .

Step 2: Configure Linear and Attn to Imitate MaxAff over U . Use Linear and W K , W Q in Attn to map the input Z ∈ U to values of the affine components { y i ( Z ) = a /latticetop i ˜ Z + b i } i ∈ [ N ma ] of MaxAff . Here we flatten the input sequence Z ∈ R d × n to ˜ Z ∈ R dn 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 /latticetop Q so that Softmax( K /latticetop Q ) produces a near-one-hot vector as 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 f ( Z ) . Map each partition cell's indicator to the corresponding value of f . By continuity of f , refining the partitioned cell ensures ‖ f -Attn ◦ Linear ‖ ∞ ≤ /epsilon1 .

Proof Sketch. We elaborate above in detail. Consider a continuous function f : U ⊆ [ -D,D ] d × n → R d × n on a compact domain U . Let /epsilon1 > 0 . We aim to construct a single-layer, single-head selfattention mechanism Attn (prepended with a linear transformation Linear ) such that

Step 1: Partition Input Domain U via MaxAff .

Flattening Input. Each input Z ∈ R d × n is reshaped into a single vector ˜ Z ∈ R dn by stacking its rows or columns. This unifies the domain as ˜ Z ∈ [ -D,D ] dn .

Grid / Max-Affine Construction. Since f is uniformly continuous on the compact set U , choose δ > 0 such that

We subdivide [ -D,D ] dn into cubes of side ≤ δ , yielding G = P dn grid centers { v j } G -1 j =0 . We treat MaxAff as a piecewise (max-)affine or piecewise-constant partition: for each ˜ Z , there's a nearest v j within δ/ 2 .

Technical Highlight. This partition-based approach leverages uniform continuity to discretize U . The number of partitions can be large but finite, ensuring we only need a single-layer of attention to 'select' the correct grid cell.

Sum-of-Linear-Transformations Map Linear . Design Linear : R d × n → R M (for some dimension M ) to capture the dot products 〈 v j , ˜ Z 〉 . Essentially, Linear( Z ) arranges these

{ v /latticetop j ˜ Z } in a form accessible to attention. This ensures each grid center v j can be individually 'queried.'

Encoding Affine Components. Observe that max j {〈 v j , ˜ Z 〉 -1 2 ‖ v j ‖ 2 } is akin to a maxaffine function. We store terms v /latticetop j ˜ Z , plus -1 2 ‖ v j ‖ 2 , into K and Q for later use in Softmax( K /latticetop Q ) .

Technical Highlight. This step demonstrates how we embed {〈 v j , ˜ Z 〉} 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 K /latticetop Q . In the self-attention block, let K /latticetop Q ≈ R ( 〈 v j , ˜ Z 〉 -1 2 ‖ v j ‖ 2 ) , where R > 0 is large. This makes Softmax( K /latticetop Q ) favor the row j ∗ maximizing

Near-One-Hot Distribution. Hence the j ∗ -th row obtains probability close to 1, effectively identifying which grid center v j ∗ is nearest to ˜ Z . We interpret this as a near-one-hot 'indicator' vector for the correct partition cell.

Technical Highlight. This is the crux: attention's softmax can act as a continuous arg max by scaling the scores with R . As R →∞ , the distribution becomes more peaked, approximating a hard partition.

Step 4: Map the Indicator to the Target Value f ( Z ) .

Assigning Values. We place f ( v j ) in the 'value matrix' W V , so that once row j ∗ is selected, the attention output is ≈ f ( v j ∗ ) . Since Z is within δ/ 2 of v j ∗ , uniform continuity implies

Final Reshaping (If Needed). A small linear projection M can reshape the output back to R d × n . The essential logic is that the correct f ( v j ) is 'routed' to the final output via the near-one-hot attention distribution.

Technical Highlight. This reveals how a single-head attention layer, armed with linear preprocessing, suffices to replicate the entire function f . No feed-forward sub-layer or multiple heads are needed to achieve universal approximation.

In sum, combining these steps, we see that: (i) A finite grid subdivides U to handle uniform continuity. (ii) Linear encodes {〈 v j , ˜ Z 〉} . (iii) LargeR Softmax( K /latticetop Q ) selects the best anchor v j ∗ . (iv) A 'value matrix' translates that selection into f ( v j ∗ ) . We conclude that a single-layer, single-head self-attention block approximates f within /epsilon1 in the L ∞ norm.Please see Appendix E.1 for a proof.

Our result in L ∞ norm can be easily extended to L p norm, where it applies to not just the continuous functions but all Lebesgue integrable functions with compact support. Please see Corollary E.1.1 for more details.

4.2 Single-Head Cross-Attention as a Universal Seq-to-Seq Approximator

Here we extend self-attention universal approximation results from Section 4.1 to cross-attention. Importantly, we establish the first known universal approximation in cross-attention setting. First, we state our main result in L ∞ -norm.

Theorem 4.2 ( L ∞ -Norm Universal Approximation) . Let f : U K × U Q → R d × n denote any continuous function on a compact domain U K × U Q and let /epsilon1 be any positive real number. Here U K , U Q ∈ R d × n stands for the compact domain of the two input sequences of cross-attention. Then

there exists a cross-attention Attn prepended with a Linear layer such that

Theorem 4.2 indicates that a single-layer cross -attention block, prepended with a linear preprocessing layer Linear , approximates f : U K → U Q → R d × n in L ∞ -norm.

Proof Sketch. Our proof follows that of Theorem 4.1 except one additional step: use Attn to aggregate the max-affine functions on U K , U Q and merge into a MaxAff function on U K × U Q . The proof consists of the following steps:

Step 1: Partition the Input Domain U K and U Q with MaxAff K and MaxAff Q Respectively. Construct two max-affine function MaxAff K over U K and MaxAff Q over U Q such that this MaxAff K induces a partition of sizeN ma of U and MaxAff Q a same size partition on U Q .

Step 2: Configure Linear and Attn to Imitate MaxAff K , MaxAff Q over W K , U Q Respectively. Use Linear and W K , W Q in Attn to map the input Z K , Z Q ∈ U to values of the affine components { y i ( Z ) = a /latticetop i ˜ Z + b i } i ∈ [ N ma ] of MaxAff K and MaxAff Q respectively. Here we flatten the input sequence Z ∈ R d × n to ˜ Z ∈ R dn to express MaxAff concisely.

Step 3: Use Attn to Aggregate MaxAff K and MaxAff Q to Form a MaxAff : U K × U Q → R on Both Input Sequences. Use Attn to generate MaxAff( Z K , Z Q ) := MaxAff K ( Z K ) + MaxAff Q ( Z Q ) . This max-affine function merges the partition on U K and U Q to generate a unified partition on U K × U Q .

Step 4: Use Attn to Indicate the Position of the Both Input Sequence in the MaxAff -Generated Partition. Use Attn to generate an indicator to the max-affine partition generated by MaxAff (as defined in Definition 3.2). This indicator (approximately a one-hot vector) shows which part of the MaxAff -generated partition contains the Cartesian product of both input sequences Z K × Z Q .

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 a 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 a 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 L ∞ norm, or integrable functions under L p 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-defined functions.

Universal Approximation. With only a single linear layer and a single-head attention module, one can approximate arbitrary sequence-to-sequence maps in both the L ∞ and L p senses, for both self-attention (Theorem 4.1 and Corollary E.1.1) and cross-attention (Theorem 4.2 and Corollary E.2.1) 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-assumption 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 A 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 heterogeneous 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-head or deeper attention 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, our results establish a theoretical foundation for understanding attention-based architectures as universal function approximators. They illustrate how token-wise information is partitioned and reassigned to represent complex sequence-to-sequence 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 is 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 R01LM1372201, NSF AST-2421845, Simons Foundation MPS-AI-00010513, AbbVie , Dolby and Chan Zuckerberg Biohub Chicago Spoke Award. This research was supported in part through the computational resources and staff contributions provided for the Quest high performance computing facility at Northwestern 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.

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.08774 , 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-57211, 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 , 33:1877-1901, 2020.

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 preprint arXiv:2504.04702 , 2025.

Sitan Chen and Yuanzhi Li. Provably learning a multi-head attention layer. arXiv preprint arXiv:2402.04084 , 2024.

George Cybenko. Approximation by superpositions of a 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 arXiv:2310.11685 , 2023.

Jacob Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 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 at 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 ¸ Sahin. On the rate of convergence of a classifier based on a transformer encoder. IEEE Transactions on Information Theory , 68(12):8139-8155, 2022.

Kurt Hornik. Approximation capabilities of multilayer feedforward networks. Neural networks , 4(2): 251-257, 1991.

Jerry Yao-Chieh Hu, Hude Liu, Hong-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 fixed-weight 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 Issei Sato. Are transformers with one layer self-attention using low-rank weight matrices universal approximators? arXiv preprint arXiv:2307.14023 , 2023.

Tokio Kajitsuka and Issei Sato. Optimal memorization capacity of transformers. arXiv preprint arXiv:2409.17677 , 2024.

Feyza Duman Keles, Pruthuvi Mahesakya Wijewardena, and Chinmay Hegde. On the computational complexity of self-attention. In 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 on Neural Networks and Learning Systems , 35(2):2448-2459, 2022.

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.08633 , 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.10469 , 2024.

Shengjie Luo, Shanda Li, Shuxin Zheng, Tie-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 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 logarithmic depth. arXiv preprint arXiv:2402.09268 , 2024a.

Clayton Sanford, Daniel J 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-sequence 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.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:2302.13971 , 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, Łukasz Kaiser, and Illia Polosukhin. Attention is 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 J Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? arXiv preprint arXiv:1912.10077 , 2019.

Zhihan Zhou, Yanrong Ji, Weijian Li, Pratik Dutta, Ramana Davuluri, and Han Liu. Dnabert-2: Efficient foundation model and benchmark for multi-species genome. arXiv preprint arXiv:2306.15006 , 2023.

Appendix

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 on a high dimension input domain ( R d × n ), attention is often considered to approximate a function defined upon a small input domain X ⊂ R d × n .

To this end, we extend our method to the approximation rate of L -Lipschitz functions with a relatively small input domain. We state our result as the following theorem.

Theorem A.1. Let f : R d × n → R d × n denote an L -Lipschitz function (in terms of 2 -norm) whose input domain is X . For any /epsilon1 > 0 , assume X is contained in N x spheres by the radius of /epsilon1/ (3 L ) in 2 -norm. Then, there exists a Linear layer and a Attn layer such that:

Furthermore, Attn and Linear have a total number of O ( dnN x ) trainable parameters.

Proof Sketch. This proof only differs from the proof of Theorem 4.1 on the choice of partition. For universal approximation, we choose a partition that evenly partition the whole space. In this theorem, we change this partition to have each part centered on a different sphere described in the Theorem A.1. By characterizing our partition, we achieve a more precise approximation result.

Please see Appendix F.1 for a detailed proof.

Theorem A.1 states that when the input domain is contained in N x spheres of /epsilon1 -level radius, there exists a single-head self-attention layer that approximates the target function with a precision of /epsilon1 .

B Proof-of-Concept Experiments

MNIST

CIFAR10

FashionMNIST

120

350

225

200

000

100

Sum of Weights

Square Sum of Weights

175

250

150

80

200

125

60

150

100

100

75

40

50

50

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

Noise Level

Noise Level

Noise Level

In Proposition 3.2, we demonstrate domain-partition mechanism of attention. In this mechanism, the temperature of the Softmax function affects the precision of the max-affine partition generated by attention, which is crucial to the complex approximations accomplished in Theorem 4.1 and Corollary E.1.1.

Since the temperature of Softmax is equivalent to the scale of the matrix involved in computing the attention score matrix ( W K , W Q ), our theory suggests the scale of W K , W Q decreases when the input data contains more noise, as a result of the rise in difficulty to form a clear partition, and an approximation based on this partition.

To verify this conjecture, we test the correlation between the scale of W K , W Q and the noise level in the training data.

Objectives. Examine the relationship between scale of matrix involved in computing the attention score matrix in attention ( W K , W Q ) 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 to the complexity and different characteristics of 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 W K , W Q becomes evident in all settings. This aligns with our 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, we use attention to approximate a step function and observe the max-affine function generated by the weights in K and Q matrices in the attention score matrix. The result of this experiment is shown in Figure 3.

Plotoftargetfunction

PlotofMax-AffineFunction

10

targetfunction(x)

Max-AffineFunction

14000

5

12000

0

10000

8000

6000

-10

4000

15

2000

20

0

10.0

7.5

5.0

2.5

0.0

2.5

5.0

7.5

10.0

-10.0

7.5

5.0

2.5

0.0

2.5

5.0

7.5

10.0

X

X

The max-affine function generated in the attention score 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 our theory.

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, consider a max-affine function MaxAff( x ) = max i ∈ [ N ma ] { a /latticetop i x + b i } , and let X ⊂ R d x be its input domain. Then MaxAff generates a partition on X :

We call the partition P ma the max-affine partition of X induced by MaxAff .

Proof. If an x 0 is not grouped to any U i , i ∈ [ N MaxAff ] . Since MaxAff is define over X and thus defined on x 0 , we have:

/negationslash

This is contradictory to the definition of MaxAff .

Since in Section 3 we exclude the discussion on the overlapped regions of the affine components { y i = a /latticetop i x + b i } , { U i | i ∈ [ N MaxAff ] } 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 Partition) . Let X = [ X 1 , X 2 , · · · , X n ] ∈ R d × n denote any input sequence. We use X to denote the domain of all X i , i ∈ [ n ] . Let MaxAff be any max-affine function on X with N MaxAff components, and let /epsilon1 > 0 be any positive real number. We define P MaxAff = { U i | i ∈ [ N MaxAff ] } as the max-affine partition generated by MaxAff as in Proposition 3.1. Let E be the indicator of P MaxAff as defined in Definition 3.2. Under the above definitions, there exists a Linear layer and a self-attention Attn whose attention matrix satisfies

with exception of an arbitrarily small region. Here W K , W Q are the attention weights within Attn .

Proof. We first denote that according to the premise of Section 3, the intersection region of different affine components are omitted. This means for an arbitrarily small δ > 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:

Construct Linear and Attn .

Estimate the error between the attention score matrix of Attn ◦ Linear and the target indicator.

For the max-affine function MaxAff , we denote it as follows.

Definition D.1 (Max-Affine Function) . Let a i ∈ R d , b i ∈ R , i ∈ [ N MaxAff ] denote the coefficients of the affine components of MaxAff . In this definition, MaxAff writes out as

for any Z ∈ R d .

Remark D.1. For conciseness of presentation, we assume the top component of MaxAff exceeds the second-largest by a fixed ∆ > 0 , independent of the input and arbitrarily small.

Construction of Linear . Without loss of generality, assume N MaxAff ≥ n . We construct Linear (the layer of linear transformations) to be

The the output of Linear( X ) is

Construction of Attn . Since we only use the attention score matrix Softmax( K /latticetop Q ) , we only have to construct the W K and W Q matrices.

We construct them to be as follows

where R is a coefficient to control the precision of the approximation. Specifically, as R increases, Softmax is closer to maximum function, and the approximation is more precise.

In this construction, we now calculate the K and Q matrices of attention

and

/latticetop

Calculation of Softmax( K Q ) . We now calculate the attention score matrix as Softmax( K /latticetop Q )

This yields the entry on the k -th row of Softmax K /latticetop Q : ,i to be

When a /latticetop k x i + b k is the maximal affine component and a /latticetop k ′ x i + b k ′ is the second largest, we have

/negationslash

Thus when we have

which means

Moreover, since the sum of all entries in Softmax K /latticetop Q : ,i is 1 , we have

(D.3) and (D.3) are equivalent to

/negationslash

/negationslash

Thus

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 : R d → R d out be a piece-wise constant function which is separately constant on each U i , i ∈ [ N MaxAff ] . We show that for any /epsilon1 > 0 , there exists an self-attention Attn such that

for every X in X with exception of a region of arbitrarily small Lebesgue measure in R n .

Proof. Let Linear and the W K , W Q and W O matrices be the same as in Appendix D.2. Then by Appendix D.2, we have

for any /epsilon1 0 > 0 .

Let V i denote the value of F on U i .

Construction of W V . We construct W V to be

Thus V equals to

Thus we have

‖

V

Softmax

(

K

/latticetop

Q

)

W

O

-

[

F

(

X

1

)

, F

(

X

2

)

,

· · ·

, F

(

X

n

)]

‖

∞

≤‖

V

‖

∞

/epsilon1

0

.

Let ‖ V ‖ ∞ /epsilon1 0 ≤ /epsilon1 yields the final result. This completes the proof.

/negationslash

This yields

Thus, by the nature of ‖ · ‖ ∞ ,

We construct W O to discard Softmax K /latticetop Q n +1: N MaxAff ,i in Softmax K /latticetop Q :

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 L ∞ norm version whose target function are continuous. Then we combine this result with the well known Lusin's theorem and extend our result to Lebesgue integrable functions in terms of L p norm.

Theorem E.1 (Theorem 4.1 Restated: L ∞ -Norm Universal Approximation of Self-Attention) . Let f : R d × n → R d × n denote any continuous function on a compact domain U ⊂ R d × n and let /epsilon1 > 0 be any positive real number. Then, there exists a self-attention Attn with a prepended Linear layer, such that

Proof Sketch. Our proof consists of four conceptual steps.

Flattening Input. Each input Z ∈ R d × n is reshaped into a single vector ˜ Z ∈ R dn by stacking its rows or columns. This unifies the domain as ˜ Z ∈ [ -D,D ] dn .

Grid / Max-Affine Construction. Since f is uniformly continuous on the compact set U , choose δ > 0 such that

We subdivide [ -D,D ] dn into cubes of side ≤ δ , yielding G = P dn grid centers { v j } G -1 j =0 . We treat MaxAff as a piecewise (max-)affine or piecewise-constant partition: for each ˜ Z , there's a nearest v j within δ/ 2 .

Sum-of-Linear-Transformations Map Linear . Design Linear : R d × n → R M (for some dimension M ) to capture the dot products 〈 v j , ˜ Z 〉 . Essentially, Linear( Z ) arranges these { v /latticetop j ˜ Z } in a form accessible to attention. This ensures each grid center v j can be individually 'queried.'

Encoding Affine Components. Observe that max j {〈 v j , ˜ Z 〉 -1 2 ‖ v j ‖ 2 } is akin to a maxaffine function. We store terms v /latticetop j ˜ Z , plus -1 2 ‖ v j ‖ 2 , into K and Q for later use in Softmax( K /latticetop Q ) .

Step 3: Enginner Attn to Generate an Indicator of Which Partition Cell the Input Belongs To.

Construct K /latticetop Q . In the self-attention block, let K /latticetop Q ≈ R ( 〈 v j , ˜ Z 〉 -1 2 ‖ v j ‖ 2 ) , where R > 0 is large. This makes Softmax( K /latticetop Q ) favor the row j ∗ maximizing

Near-One-Hot Distribution. Hence the j ∗ -th row obtains probability close to 1, effectively identifying which grid center v j ∗ is nearest to ˜ Z . We interpret this as a near-one-hot 'indicator' vector for the correct partition cell.

Assigning Values. We place f ( ˜ v j ) in the 'value matrix' W V , so that once row j ∗ is selected, the attention output is ≈ f ( ˜ v j ∗ ) . Since Z is within δ/ 2 of v j ∗ , uniform continuity implies

Final Reshaping (If Needed). A small linear projection M can reshape the output back to R d × n . The essential logic is that the correct f ( ˜ v j ) is 'routed' to the final output via the near-one-hot attention distribution.

Thus, a single-head attention block with a minimal linear layer can approximate any continuous function on the domain. This completes the proof.

Proof. We divide our proof into two parts:

Part 1: Construction of Attn and Linear . We construct Attn and Linear in 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,D ] dn used in the construction of Linear and Attn .

These grid points are used to construct the max-affine partition. Specifically, the max-affine partition we use is a grid-partition and these points are the center points of these grids.

Construction of Grid Centers in [ -D,D ] dn . Let Z = [ z 1 , z 2 , · · · , z n ] ∈ R d × n denote the input to Linear . Define ˜ Z := [ z /latticetop 1 , z /latticetop 2 , · · · , z /latticetop n ] /latticetop . P ∈ N + is a parameter that controls the size of the attention block and the error of our approximation.

Definition E.1 (Grid Centers in [ -D,D ] dn ) . Define v k 1 ,k 2 , ··· ,k dn ∈ R dn as

for k i ∈ { 0 , 1 , 2 , · · · , P -1 } , i ∈ [ dn ] .

Remark E.1 (Scalar-Labeled Grid Centers) . For each multi-index ( k 1 , . . . , k dn ) with k i ∈ { 0 , . . . , P -1 } , we define

This baseP expansion gives a one-to-one map between the tuple and the scalar. This notation allows us to define another representation of the grid center:

We now construct functions E and T . They are linear functions of f : R d × n → R d × n playing crucial roles in the constructions of W K and W Q in Attn( · ) .

For every v ∈ V , we define

Construction of E and T . We first show that f is bounded. Because f is continuous within a closed region, its output value is bounded ∞ -norm. Let B 0 denote this bound

Wenowconstruct two functions E ( · ) , T ( · ) related to f . Their sum is a constant while their subtraction is scaled f . For any Z ∈ R d × n , we define

By the definition of E ( · ) and T ( · ) , we have

and for any Z ∈ R d × n .

Construction of the Layer of Sum of Linear Transformations. We now construct the Linear layer to be

where G = P dn .

This layer multiplies the flattened input with the grid centers in Definition E.1 and append a 2 dG -dimensional identity matrix below the matrix containing these multiplications.

We now express the output of Linear in a simpler form in the following discussion.

First, we show that

This yields

Explicitly, the last line is by

which implies

Here

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 an identity matrix ( I 2 dG ).

Construction of K and Q Matrices. We now construct the W k and W Q matrices in the selfattention block and calculate the output of Softmax ( K /latticetop Q ) .

We define W K as follows

The definition of W K yields that

where the last line follows from X 0 being multiplied by 1 and thus appearing in the first row of the output.

Next, we construct W Q to be

This yields that

We now calculate the attention matrix Softmax ( K /latticetop Q ) .

Calculation of Softmax( K /latticetop Q ) . First, K /latticetop Q writes out as

where the last line follows from the multiplication of block matrices. This multiplication between K /latticetop and Q is equivalent to first multiplying the first 2 columns in K /latticetop with R and then broadcasting their sum to the first n columns, and then adding the result with T and E related blocks. Columns are all filled with 0 except for the first n columns.

Remark E.2 (Interpretation of K /latticetop Q ) . The non-zero entries of K /latticetop Q is an aggregation of two matrices

In these two matrices, (E.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 so that the Softmax approximates a maximum function, and 'selects' the i of the maximal R ( v /latticetop i ˜ Z -‖ v i ‖ 2 2 2 )1 d × n for i ∈ { 0 , 1 , · · · , G -1 } . By "select" we mean only the entries with the selected label has a value not close to 0 in each column of Softmax( K /latticetop Q ) .

(E.8) does not include R related terms. Thus when R is 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 /latticetop Q ) , the output approximates a matrix whose columns are all-zero except for two sub-vector equal to 1 / 2 d · 1 d . This writes out as (here we only show the first n non-constant columns)

for any s ∈ [ G ] . The addition of (E.8) change the 1 d in (E.9) to

In later discussion, we use V to transform (E.10) to T ( ˜ v s -1 ) -E ( ˜ v s -1 ) = 2 f ( ˜ v s -1 ) / 2 dB 0 to obtain the final output.

Now, we divide the calculation of Softmax ( K /latticetop Q ) into two parts: the calculation of exp ( K /latticetop Q ) and the calculation of the denominator of every column of Softmax ( K /latticetop Q ) . This denominator explicitly writes out as ∑ 2 dG j =1 exp ( K /latticetop Q ) ij for each i ∈ [2 dG ] .

and

For exp ( K /latticetop Q ) , by (E.6), we have

For the denominator, we calculate it in columns. Let i denote the column which we calculate the denominator in Softmax . When i ∈ { n +1 , n +2 , · · · , 2 dG } , there are 1 · 2 dG = 2 dG columns. And when i ∈ [ n ] , we denote that

Observing from (E.12), ∑ 2 dG j =1 exp ( K /latticetop Q ) i,j is invariant of i for i ∈ [ n ] . In this case, we define

From (E.11) and (E.12), we have

Construction of W V and W O . We now construct the W V matrix and calculate the V matrix of the self-attention.

We define W V as:

where

is a matrix formed by stacking GI d matrix horizontally.

With this definition, we compute V matrix as follows

After the construction and calculation of V , we go on to construct W O as:

The sole purpose of W O 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 )

Let I d denote the d -dimensional identity matrix. We have

This yields

Part 2: Estimation of the Approximation Error between Attn ◦ Linear and f .

With above calculations of the output of Attn ◦ Linear , we now demonstrate how this output approximates our target function.

Essentially, we demonstrate that each term in the summation of (E.15), given by

approximates a max-affine indicator as R becomes sufficiently large. They are each multiplied with f ( ˜ v j ) , which is the value of the target function at the center point of the indicated region.

Definition E.2 (Max-Affine Function on ˜ Z ) . Let Aff j ∈ R dn → R with j ∈ { 0 , 1 , 2 , · · · , G -1 } denote a group of affine functions defined as:

Then let MaxAff ∈ R dn → R denote a max affine function whose affine components are { Aff j | j ∈ { 0 , 1 , 2 , · · · , G -1 }} . Explicitly defined as:

Because the target function f is a continuous function on a closed domain, the function f is uniformly continuous. Thus for /epsilon1 , there exists a δ > 0 such that for any Z 1 , Z 2 , as long as ‖ ˜ Z 1 -˜ Z 2 ‖ ∞ ≤ δ , we have ‖ f ( Z 1 ) -f ( Z 2 ) ‖ ∞ ≤ /epsilon1/ 3 .

According to this δ , we divide the affine components of MaxAff into three parts:

The maximal component, which has the smallest label j m .

( Equivalent to summing all blocks in S )

All affine components that match the maximal component or fall within δ of it ( J 0 as defined below).

The remaining Aff j for j ∈ { 0 , 1 , . . . , G -1 } ( J 1 as defined below).

We write out the labels of these groups of components as follows

For any pair of i, j ∈ { 0 , 1 , · · · , G -1 } , we have

Thus for j m , we have

This yields for all j ∈ { 0 , 1 , · · · , G -1 } .

This denotes j m is also the label of the closest v i to ˜ Z among all v i , i ∈ { 0 , 1 , · · · , G -1 } . Thus we have

Now, we prove v j m (the grid point nearest to ˜ Z ) has a distance to ˜ Z smaller than half of the grid width (e.g., D/g ) in infinite norm.

Let D := 2 D/g × {-1 , 0 , 1 } dn denote a set differences to v j m from the set of all v i ( i ∈ { 0 , 1 , · · · , G -1 } ) neighboring v j m . For any ∆ in D , from (E.16) we have

This means that, for any k ∈ [ dn ] , by selecting ∆ to be ± 2 D/ge ( dn ) k , we have:

Thus we have which implies

This yields

Set g to be larger than 2 D/δ ; we have thus

where the inequality holds by δ/ 2 < δ .

Calculation of ‖ Attn ◦ Linear -f ‖ ∞ . We now calculate the difference between the output in (E.15) and target function f

We now calculate each part in (E.18).

As previously stated, for any Z 1 , Z 2 , as long as ‖ ˜ Z 1 -˜ Z 2 ‖ ∞ ≤ δ , we have ‖ f ( Z 1 ) -f ( Z 2 ) ‖ ∞ ≤ /epsilon1/ 3 . Thus when we designate Z 1 = v j for any j ∈ J 0 and Z 2 = v j m , along with (E.17) we have

For j m , we have

When R is larger than 8 ( 3 )

Combining (E.19) and (E.20) yields

where the last line is by ∑ j ∈ J 0 ∪{ j m } 1 α ( Z ) exp ( R ( v /latticetop j ˜ Z -1 2 ‖ v j ‖ 2 2 ) ) ≤ 1 . We plug (E.21) and (E.22) to (E.18) and get

This completes the proof.

We also extend this L ∞ -Norm result we just proved to L p -Norm.

Corollary E.1.1 ( L p -Norm Universal Approximation) . Let f : R d × n → R d × n denote any Lebesgue integrable function on a compact domain U ∈ R d × n and let /epsilon1 > 0 be any positive real number. Then, there exists a self-attention Attn prepended with a Linear layer such that

Proof Sketch. The same partition-based construction applies almost everywhere; outside a negligible set, f is continuous (Lusin's theorem). Thus the L ∞ argument extends. Proof. Since f is Lebesgue integrable on a compact set, f is bounded almost every where. Let B p denote the bound of ‖ f ‖ p .

By Lusin's theorem, for f on a compact domain U , there exists a continuous function g which is equal to f in U except for a region D δ such that µ ( D δ ) ≤ ∆ . This can be written as

/negationslash

Here µ stands for the Lebesgue measure of a set.

By Theorem 4.1, there exists a net work Attn ◦ Linear , consists of a self-attention Attn and a layer of sum of linear transformation Linear such that

for any /epsilon1 0 > 0 .

This denote that for any Z ∈ U

Combine this with (E.23) and (E.24), we get

/negationslash since that f ( Z ) = g ( Z ) , ‖ Attn ◦ Linear( Z ) -g ( Z ) ‖ = ‖ Attn ◦ Linear( Z ) -f ( Z ) ‖ ≤ /epsilon1 0 . This yields

Set

We have

This completes our proof.

E.2 Proof of Theorem 4.2

Theorem E.2. Let U K ⊂ R d × n and U Q ⊂ R d × n be two compact domains, and let f : U K × U Q → R d × n be any continuous function that takes input from both domains. We use Z K , Z Q ∈ R d × n to denote the two inputs of f from U K and U Q respectively. Without loss of generality, suppose both input domains to be [ -D,D ] d × n , where D ∈ R + . Then for any /epsilon1 > 0 , there exists a single-head cross-attention Attn and two layers of sum of linear transformations, Linear K and Linear Q such that:

for any Z K , Z Q ∈ [ -D,D ] d × n .

Proof. Without loss of generality, assume U K = U Q = [ -D,D ] d × n for a D ∈ R + .

Construction of Grid Centers in U K , U Q . Same as in Appendix E.1, we define ˜ Z := [ z /latticetop 1 , z /latticetop 2 , · · · , z /latticetop n ] /latticetop . P ∈ N + is a parameter that controls the size of the attention block and the error of our approximation. Define v k 1 ,k 2 , ··· ,k dn ∈ R dn to be

Let V := { v k 1 ,k 2 , ··· ,k dn | k i ∈ { 0 , 1 , 2 , · · · , P -1 } , i ∈ [ dn ] } be the set of all v k 1 ,k 2 , ··· ,k dn . We also define another way to refer to a vector in V , denoted as

Please see Remark E.1 for the reason for the feasibility of such expression.

Following the notation in Appendix E.1, for every v ∈ V , we define

as a d × n matrix-form representation of v .

Construction of f Related Function E and T . The continuity of f within a closed region guarantees it to be bounded in ∞ -norm. Let B 0 denote this bound. For any a K , a Q ∈ R d × n , we define

We define ( E + T )( a K , a Q ) = E ( a K , a Q ) + T ( a K , a Q ) . By the definition of E and T , ( E + T )( a K , a Q ) ≡ 2 d × n for any a K , a Q ∈ R d × n .

Remark E.3 (Intuition behind E and T ) . E and T are constructed to satisfy 3 conditions:

T ( Z k , Z Q ) + E ( Z K , Z Q ) ≡ 2 .

T ( Z k , Z Q ) -E ( Z K , Z Q ) = 2 f ( Z K , Z Q ) /B 0 .

T, E > 0 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 E.1, define

We now construct the Linear K and Linear Q layers to be

Same as that in Theorem E.1, we have

We now calculate the output of Linear K and Linear Q .

For Linear K , we have

in which X K and X Q are defined as

We now construct the W k and W Q matrices in the self-attention block and calculate the output of Softmax ( K /latticetop Q ) .

In the following, we define W K in parts. First, we present it as a block matrix

We then define the submatrices in (E.30) as follows

in which

The definition of W K yields that

Next, we construct the W Q matrix as

In this definition, the Q matrix in attention can be calculated as follows

Now we calculate the attention matrix Softmax ( K /latticetop Q )

K /latticetop Q can be calculated as follows

The W /latticetop 1 X Q in the expression of K /latticetop Q matrix is further calculated as

We define Q 1 := W /latticetop 1 X Q , then W /latticetop 1 X Q can be denoted as stacking this block vertically for G times. In this definition, Q 1 matrix can be expressed as

.

The calculation of Softmax ( K /latticetop Q ) can be disassembled into two parts, the numerator exp ( Softmax ( K /latticetop Q )) in the expression of Softmax and the denominator of every column of Softmax ( K /latticetop Q ) , as in the expression of Softmax , explicitly written out as ∑ 2 dG j =1 exp ( K /latticetop Q ) ij for each i ∈ [2 dG ] .

We calculate exp ( K /latticetop Q ) as follows

For the denominator, we calculate it in columns. Let i denote the column which we calculate the denominator in Softmax . When i ∈ { n +1 , n +2 , · · · , 2 dG 2 } , the i -th column has 1 in every entry. Thus the sum of all entries in this column equals to 1 · 2 dG = 2 dG .

And when i ∈ [ n ] , we have

We observe from (E.33), that ∑ 2 dG 2 j =1 exp ( K /latticetop Q ) ij is invariant of i for i ∈ [ n ] . In this case, we define

to denote the 1 / 2 d of this value invariant of i for simplicity.

Because

from (E.32) and (E.33) we have

Now we've defined and calculated the attention score matrix Softmax( K /latticetop Q ) , we go on to construct the W V matrix and calculate the result of multiplying V = W V Linear( Z K ) to the attention score matrix.

We define W V as where

This yields the V matrix to be

With V , we compute the output of V Softmax( K /latticetop Q ) as follows

To further calculate V Softmax( K /latticetop Q ) , we now calculate the result of its non-trivial part (the part beside 0 dG 2 × (2 dG 2 -n ) )

We now calculate each part in (E.35)

In (E.36), we have

Thus (E.36) is equal to

We also calculate the other part exp ( ( RX /latticetop K + RW /latticetop 0 )1 1 × n + RW /latticetop 1 X Q ) in separate parts

and

Thus

Combing (E.37) and (E.38), we have

Thus we compute (E.35) as

We now put back the 1 / 2 dα ( Z K , Z Q ) in (E.34) and calculate the final output as

Next, we construct W O to be

This yields the final output of Attn ◦ Linear to be

Attn ◦ Linear( Z )

= V Softmax K /latticetop QW O

Estimation of Error between Attn ◦ 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 := [[ ˜ Z /latticetop K , ˜ Z /latticetop Q ] /latticetop ] to accommodate to the expression of affine functions.

Definition E.3 (Max-Affine Function on ˜ Z .) . Let Aff i,j ∈ R dn → R , j ∈ { 0 . 1 . 2 . · · · , G -1 } denote a group of affine functions defined as

Then let MaxAff ∈ R dn → R denote a max affine function whose affine components are { Aff i,j | i, j ∈ { 0 , 1 , 2 , · · · , G -1 }} . Explicitly defined as:

In the following discussion, we use η ∈ { 0 , 1 , · · · , G -1 } 2 to refer to a pair of coefficients ( i, j ) , and denote A i,j as A η for the corresponding η . Furthermore, we denote the two labels encapsulated in η as i η and j η

Because the target function f is a continuous function on a closed domain, the function f is uniformly continuous. Thus for /epsilon1 , there exists a δ > 0 such that for any Z (1) = [ Z (1) K , Z (1) Q ] , Z (2) = [ Z (2) K , Z (2) Q ] , as long as ‖ Z (1) -Z (2) ‖ ∞ ≤ δ , we have ‖ f ( Z (1) ) -f ( Z (1) ) ‖ ∞ ≤ /epsilon1/ 3 .

According to this δ , 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 η m , the group of affine components equal to the maximal component or smaller than it by no more than δ , and finally, the other Aff η . We write out the labels of these groups of components as follows

For any pair of η 1 , η 2 ∈ { 0 , 1 , · · · , G -1 } 2 , we denote that

Let v η := [ v /latticetop i η , v /latticetop j η ] /latticetop , denote a flatten stack of v i η and v j η . Same as v i , define ˜ v η := [ ˜ v i η , ˜ v j η ] . Then the above expression denotes η m is also the label of the v η closest to ˜ Z among all v η , η ∈ { 0 , 1 , · · · , G -1 } 2 . Thus we have

This means that v η m is the grid center closest to ˜ Z 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 := 2 D/g ×{-1 , 0 , 1 } dn denote a set of differences to v η m of all the v i ( i ∈ { 0 , 1 , · · · , G -1 } ) neighboring v η m . For any ∆ in D , from (E.42) we have

which means for any k ∈ [ dn ] , by selecting ∆ to be ± 2 D g · e ( dn ) k , we have

Thus we have

This is equivalent to

This yields

Set g to be larger than 2 D/δ , we have thus

Calculation of ‖ Attn ◦ Linear -f ‖ L ∞ . We now calculate the difference between the output in (E.39) and target function f

The last row is simply a separation of the summation in the row above.

We now calculate each part in (E.43).

As previously stated, for any Z 1 , Z 2 , as long as ‖ ˜ Z 1 -˜ Z 2 ‖ ∞ ≤ δ , we have ‖ f ( Z 1 ) -f ( Z 2 ) ‖ ∞ ≤ /epsilon1/ 3 . Thus when we designate Z 1 = v η for any η ∈ η 0 and Z 2 = v η m , along with (E.40) we have

For any η m , we have

When R is larger than 8 ln(3 / 2 · B 0 G/epsilon1 ) / (3 δ 2 ) , we have

Combing (E.44) and (E.45) yields

where the last line is by ∑ η ∈ E 0 ∪{ η m } exp ( R ( v /latticetop η ˜ Z -‖ vη ‖ 2 2 2 ) ) α ( Z ) ≤ 1 .

By (E.47) and (E.46), we have

This completes the proof.

Theorem 4.2 can be easily extended to Lebesgue integrable functions in L p norm in the following result.

Corollary E.2.1 ( L p -Norm Universal Approximation) . Let f : U K × U Q → R d × n denote any Lebesgue integrable function on a compact domain U K × U Q and let /epsilon1 be any positive real number. Here U K , U Q ∈ R d × n stands for the compact domain of the two input sequences of cross-attention. Then, there exists a cross-attention Attn prepended with a Linear layer such that

Proof. Without loss of generality, assume U K = U Q = [ -D,D ] d × n for a D ∈ R + .

Since f is Lebesgue integrable on a compact set, f is bounded almost every where. Let B p denote the bound of ‖ f ‖ p .

By Lusin's theorem, for f on a compact domain U , there exists a continuous function g which is equal to f in U except for a region D δ such that µ ( D δ ) ≤ ∆ . This can be written as

/negationslash

where µ stands for the Lebesgue measure of a set.

By Theorem 4.2, there exists a network Attn ◦ Linear , consists of a cross-attention Attn and a layer of sum of linear transformation Linear such that

for any /epsilon1 0 > 0 .

This denote that for any Z ∈ U × U

Combing this with (E.48) and (E.49), we get

/negationslash since if f ( Z ) = g ( Z ) , ‖ Attn ◦ Linear( Z ) -g ( Z ) ‖ = ‖ Attn ◦ Linear( Z ) -f ( Z ) ‖ ≤ /epsilon1 0 This yields

Set

We have

This completes the proof.

F Proof of Results in Appendix A

F.1 Proof of Theorem A.1

Theorem F.1 (Theorem A.1 Restated) . Let f : R d × n → R d × n denote an L -Lipschitz function (in terms of 2 -norm) whose input domain is X . For any /epsilon1 > 0 , assume X is contained in N x sphere by the radius of /epsilon1/ (3 L ) in 2 -norm. Then, there exists a Linear layer and a Attn layer such that:

Furthermore, Attn and Linear have a total number of O ( dnN x ) trainable parameters.

Proof sketch. This proof is identical with Theorem 4.1, except for an alteration on the set of v i .

Proof. We follow the proof of Theorem 4.1.

Notation of Sphere Centers. Let Z = [ z 1 , z 2 , · · · , z n ] ∈ R d × n denote the input to Linear . Define ˜ Z := [ z /latticetop 1 , z /latticetop 2 , · · · , z /latticetop n ] /latticetop . P ∈ N + is a parameter that controls the size of the attention block and the error of our approximation.

Let v i , i ∈ [ N x ] denote the centers of the N x spheres that covers X . Let V := { v i | i ∈ [ N x ] } denote the set of all v i .

Construction of f Related Functions. Because f is continuous within a closed region, its output value is bounded in ∞ -norm. Let B 0 denote this bound, we now construct two functions that. For any a ∈ R d × n , we define E ( a ) := 1 d × n -f ( a ) /B 0 and T ( a ) = 1 d × n + f ( a ) /B 0 . We define ( E + T )( a ) = E ( a ) + T ( a ) . By the definition of E and T , ( E + T )( a ) ≡ 2 d × n for any a ∈ R d × n .

Construction of the Layer of Sum of Linear Transformations. We now construct the Linear layer to be

where N x = P dn .

We now express the output of Linear in a simpler form in the following discussion. First, we show that

This yields

in which X 0 is defined as follows

Construction of K and Q Matrices. We now construct the W k and W Q matrices in the selfattention block and calculate the output of Softmax ( K /latticetop Q ) .

We define W K as follows:

The definition of W K yields

Next, we construct W Q to be

This yields that

We now calculate the attention matrix Softmax ( K /latticetop Q ) .

Calculation of Softmax( K /latticetop Q ) . First, K /latticetop Q can be expressed as follows

Now, we divide the calculation of Softmax ( K /latticetop Q ) into two counterparts, the calculation of exp ( K /latticetop Q ) and the calculation of the denominator of every column of Softmax ( K /latticetop Q ) , as in the expression of Softmax , explicitly written out as ∑ 2 dN x j =1 exp ( K /latticetop Q ) ij for each i ∈ [2 dN x ] .

For exp ( K /latticetop Q ) , we have

For the denominator, we calculate it in columns. Let i denote the column which we calculate the denominator in Softmax . When i ∈ { n +1 , n +2 , · · · , 2 dN x } , it obviously equals to 1 · 2 dN x = 2 dN x . And when i ∈ [ n ] , we denote that

We observe from (F.3), that ∑ 2 dN x j =1 exp ( K /latticetop Q ) ij is invariant of i for i ∈ [ n ] . In this case, we define

From (F.1) and (F.3) we have

Construction of W V and W O . We now construct the W V matrix and calculate the V matrix of the self-attention.

We define W V as where

is a matrix formed by stacking N x I d matrices horizontally.

In this definition, V matrix can be calculated as follows:

After the construction and calculation of V , we go on to construct W O as

The sole purpose of W O is to extract the non-zero entries of the final output.

Calculation of the Output of Attn ◦ Linear . We now calculate the final output of the self-attention block.

We have

This yields

Estimation of the Error between Attn ◦ 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 F.1 (Max-Affine Function on ˜ Z ) . Let Aff j ∈ R dn → R , j ∈ { 0 , 1 , 2 , · · · , N x -1 } denote a group of affine functions defined as

Then let MaxAff ∈ R dn → R denote a max affine function whose affine components are { Aff j | j ∈ { 0 , 1 , 2 , · · · , N x -1 }} . Explicitly defined as

Because the target function f is a continuous function on a closed domain, the function f is uniformly continuous. Thus for /epsilon1 , there exists a δ > 0 such that for any Z 1 , Z 2 , as long as ‖ ˜ Z 1 -˜ Z 2 ‖ ∞ ≤ δ , we have ‖ f ( Z 1 ) -f ( Z 2 ) ‖ ∞ ≤ /epsilon1/ 3 .

According to this δ , we divide the affine components of MaxAff into three parts, the maximal component(and also with the smallest label), whose label is denoted as j m , the group of affine components equal to the maximal component or smaller than it by no more than δ , and finally, the other Aff j , j ∈ { 0 , 1 , 2 , · · · , N x -1 } . We write out the labels of these groups of components as follows

For any pair of i, j ∈ { 0 , 1 , · · · , N x -1 } , we have

This denotes j m is also the label of the closest v i to ˜ Z among all v i , i ∈ { 0 , 1 , · · · , N x -1 } . Thus we have

Thus, when considering the Z in the input domain of f , which by definition is contained in N x spheres, the closest center to Z is the sphere containing Z . This gives

Then, with the L Lipschitzness of L we have

Difference between Attn ◦ Linear and f . We now calculate the difference between the output in (F.4) and target function f

We now calculate each part in (F.8).

For the L -Lipschitzness of f , for any Z 1 , Z 2 , as long as ‖ ˜ Z 1 -˜ Z 2 ‖ ∞ ≤ /epsilon1 3 L , we have ‖ f ( Z 1 ) -f ( Z 2 ) ‖ ∞ ≤ /epsilon1/ 3 . Thus when we designate Z 1 = v j for any j ∈ J 0 and Z 2 = v j m , along with (F.7) we have:

For j m , we have

When R is larger than 8 3 δ 2 ln ( 3 2 · B 0 N x /epsilon1 ) , we have:

Combing (F.10) and (F.11) yields

where the last line is by ∑ j ∈ J 0 ∪{ j m } exp ( R ( v /latticetop j ˜ Z -‖ v j ‖ 2 2 2 ) ) α ( Z ) ≤ 1 . We plug (F.12) and (F.13) to (F.8) and get

where

This concludes our result on the approximation error.

Estimation of the Number of Trainable Parameters. We now estimate the number of trainable parameter in the network we constructed to verify our claim on number of trainable parameters in the main text of 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 a sum over N x v /latticetop j ˜ Z, j ∈ [ N x ] , and thus contain dn · N x trainable parameters.

Then we do the estimation for W K and W Q . We restate the construction of W K and W Q :

From this, we observe they combined together have 2 d · N x +2 dn · N x trainable parameters. Finally, For W V and W O , we restate their definition:

W O contains n trainable parameters ( dB 0 ).

In conclusion, the whole network contains a total of

trainable parameters, which is of O ( dnN x ) level.

This completes the proof.

G Extended Related Work

Recent empirical 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. Our work builds 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-to-sequence mappings.

H Extension to Transformer Setting

In this section, we demonstrate how our proof technique in attention extends to the transformer setting. Specifically, we show it is possible to use a three-layer multi-head attention-only network to achieve universal approximation in L p norm.

Lemma H.1 (Attention Simulate Column-Wise Transformation) . For an input X ∈ R d × n . Let l ( X ) := AXB,A ∈ R d output × d , B ∈ R n × n be the linear operation (both token-wise and sequencewise included) we wish to get. Without loss of generality, suppose B has positive entries (since subtraction of positive matrices yields all matrix). We use an all-zero padding token. Let I n +1 be the positional encoding below X . Then for any /epsilon1 > 0 , there exists an single-head attention Attn such that

Proof. Let s i be the sum of all entries of B : ,i ( i -th column) and M := max i ∈ [ n ] s i . Define S = [ s 1 · · · s n ] . Construct W V := 3 M ∗ [ A 0 d output × ( n +1) ] , W K := [0 n × d ln ( B /latticetop ) ln (3 M · 1 n -S /latticetop )] ( ln is entry-wise), and W Q := [0 n × d I n T · 1 n ] (T will be turned to arbitrarily large), then the output of attention is 3 MA [ X 0 d ] Softmax ([ ln ( B /latticetop ) ln (3 M · 1 n -S /latticetop )] /latticetop [ I n T · 1 n ]) .

This equals

in which H 1 , H 2 are the sum of all columns in ln ( B ) , ln (3 M · 1 1 × n -S ) respectively.

This further equals to

in which C = 3 MAXa 0 , when T → + ∞ , a 0 → 0 d output and the padding token is preserved as well.

Lemma H.2 (Preservation of Identity Matrix) . For any /epsilon1 > 0 and n ∈ N + , there exists an attention head Attn that satisfies

Proof. Construct the W K , W Q , W V matrices of this attention head as

W V := I n

Softmax( RI n )

Hence when R is sufficiently large, the output is within an arbitrarily small error to I n .

Corollary H.0.1. For any /epsilon1 > 0 and input of form

where A ∈ R m × n denotes any matrix, there exists an attention head Attn such that

The output of this layer is

The proof is obvious. We only have to zero out A in W K , W Q , W V and do the rest as Lemma H.2. We now prove the universal approximation result in the transformer setting.

Theorem H.1. Let X ∈ R d × n denote the input sequence. Let f be a continuous function on a compact support. For any η > 0 , there exists a three-layer multi-head attention Attn ( i ) , i ∈ [3] such that

The First Layer

Let v i , i ∈ [ P dn ] be as defined in Remark E.1.

Let ˜ X represent the flattened input X and let ˜ v represents the R d × n sequence form of v ∈ R dn .

Let δ i , i ∈ [ d ] be defined as

The Second Layer

Label the heads in the second layer with ( v, g, s ) pair, v ∈ [ P dn ] , g ∈ [ n ] , s ∈ {± 1 } . Construct parameters in Attn (2) ( v,g,s ) to be

Let Attn (1) h , h ∈ [ n ] denote the h -th head of Attn (1) . According to Lemma H.1, for any /epsilon1 0 > 0 ,there exists such set of Attn h such that

in which

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 I n +1 in this case). So the output approximates

to an arbitrarily small error in the infinite norm.

The first attention layer is then defined as

in which D is

and R ∈ R + is a parameter that will be set in later process.

Then the attention score matrix of Attn (2) ( v,g,s ) is

When R is taken sufficiently large, for any given /epsilon1 1 > 0 and an arbitrary ∆ s > 0 invariant of /epsilon1 1 and R , if

for a specific i , then in the attention score matrix of Attn (2) ( v,g, 1) , the i -th row should be larger than the last row by ∆ s , and this difference after amplified by R makes the last row to be lesser than the given /epsilon1 1 .

Similarly, when

for a specific i , the ( n +1 , j ) entries in Attn (2) ( v,g, -1) ( j = 1 , · · · , n ) are smaller than /epsilon1 1 when setting R to be sufficiently large. And the sum the rest of the j -th column is larger than 1 -/epsilon1 1 .

Finally, when | δ /latticetop i ( x g -( ˜ v v ) : ,g ) | < 1 2 P 2 -∆ s , for all i ∈ [ n ] , it makes the last row largest in every entry (except the last column). Thus there exists a sufficiently large R , such that the last row of the attention score matrix is larger than 1 -/epsilon1 1 in every entry. This also makes the sum of the rest of the rows less than or equal to /epsilon1 1 in every column.

We show that the last situation is equivalent to the input token x g falling into the grid centered at ( ˜ v v ) : ,g in R d . (notice v v is a grid point in R dn and its column is a grid point in R d )

Therefore

is equivalent to x g being closer to ( ˜ v v ) : ,g than its neighboring grid points in R d (which are { ( ˜ v v ) : ,g ± δ i | i ∈ [ d ] } ), this concludes it to be within the grid centered at ( ˜ v v ) : ,g .

This means that when x g is within the grid centered at ˜ v v , both heads in Attn (2) ( v,g,s ) , s ∈ [ ± 1] has their last row in the attention score matrix larger than 1 -/epsilon1 1 (which makes their sum larger than 2 -2 /epsilon1 1 ) and if x g is not in the grid, their sum would be smaller than 1 + /epsilon1 1 (because at least one is smaller /epsilon1 and all entries should be smaller than 1 )

Let W V of head labeled ( v, g, s ) be in which

is the previously constructed attention score of the ( v, g, s ) head.

Finally, add a head Attn (2) I to preserve the identity matrix as that done in the first layer using Corollary H.0.1.

This head outputs (with an arbitrarily small error)

Define Attn (2) as the sum of previously defined heads

Then we have that when the input is

in which s v := vd +( g -1) , t v = P dn n -s v + n -1 . We note that here head ( v, g, 1) and ( v, g, -1) have the same head.

Then V is

The output of the ( v, g, s ) -th head in the second attention layers is

the output of this layer is within an arbitrary small error to

We note that only when input is within the grid centered at v , 2 1 × n -S ( v, g ) is smaller than 2 -(2 -2 /epsilon1 1 ) = 2 /epsilon1 1 for every g in all entries.

The Third Layer

The third layer consists of P dn blocks, the v -th block identifies if the input falls inside the v -th grid. For the v -th block, construct its W K , W Q as

in which R 1 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

This yields the attention score matrix to be

By n · 2 1 × n -∑ g ∈ [ n ] S ( v, g ) is smaller than 2 n/epsilon1 1 if input is in the grid centered at v and larger than 1 -/epsilon1 1 if not. Set /epsilon1 1 = 1 / 4 n , we get that for any /epsilon1 2 > 0 , when R 1 is sufficiently large,

when input is within the grid centered at v .

And when the input is not in the grid centered at v , we have

in which

Construct W V as

This makes the V matrix

when input has the form in (H.2).

Thus the third layer's v -th block outputs

Then for any /epsilon1 3 > 0 , there exists an R 1 such that the output of the third layer satisfies

when the input falls inside the grid centered at v v .

And

when the input does not fall inside the grid centered at v v . Here ∗ denotes the padded column that is 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 an error of P dn /epsilon1 3 to

with v v 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 µ (Ω) = V .

Let F ( x ) = (Attn (3) ◦ Attn (2) ◦ Attn (1) )( x ) denote the output of the network given ideal inputs Z i , and let F ∗ ( x ) denote the output given perturbed inputs Z ∗ i (which are the actual inputs).

We have formerly proven the following conditions for any /epsilon1 > 0 :

Ideal Approximation: || F ( x ) -f ( x ) || ∞ < /epsilon1 .

Perturbed Inputs: The input to the third layer, Z ∗ 3 ( x ) , satisfies || Z ∗ 3 ( x ) -Z 3 ( x ) || ∞ ≤ /epsilon1 on a set Ω \ B δ , where B δ is a region of arbitrarily small measure µ ( B δ ) < δ . On B δ , the error exceeds /epsilon1 but is still bounded.

Given the conditions above, for any precision η > 0 and any p ∈ [1 , ∞ ) , the actual output F ∗ approximates the target function f in the L p norm such that || F ∗ -f || p < η .

We aim to bound the L p error || F ∗ -f || p . By the Minkowski inequality (triangle inequality for L p norms), we can separate the error into an approximation component and a stability component:

We analyze these two terms separately.

1. Bounding the Approximation Error

From the problem statement, the ideal network approximates f uniformly with error /epsilon1 . Using the definition of the L p norm on the domain Ω :

Since /epsilon1 can be chosen arbitrarily small, this term converges to 0 as /epsilon1 → 0 .

2. Bounding the Stability Error

Consider the term || F ∗ -F || p , which represents the deviation caused by the perturbed intermediate inputs. Recall that F ( x ) = Attn (3) ( Z 3 ( x )) and F ∗ ( x ) = Attn (3) ( Z ∗ 3 ( x )) . To evaluate the integral over Ω , we partition the domain into the 'good' set S = Ω \ B δ and the 'bad' set B δ :

Analysis on S (Lipschitz Continuity): The Multi-Head Attention layer Attn (3) consists of linear projections and the Softmax function, which are smooth and differentiable operations. Consequently, Attn (3) is locally Lipschitz continuous on compact domains. Let K be the Lipschitz constant for Attn (3) . On the set S , we are given that || Z ∗ 3 -Z 3 || ≤ /epsilon1 . Therefore:

The integral over S is bounded by:

Analysis on B δ (Uniform Boundedness): On the region B δ , the Lipschitz bound fails. However, since Ω is compact and the functions comprising the network are continuous, the image of the network is bounded. There exists a constant M > 0 such that || F ( x ) || ≤ M and || F ∗ ( x ) || ≤ M for all x ∈ Ω . By the triangle inequality, the pairwise difference is bounded by 2 M . Thus:

Combining these results, the stability error satisfies:

3. Conclusion

Substituting the bounds back into the original inequality, we obtain:

To ensure || F ∗ -f || p < η , we select parameters as follows:

Choose /epsilon1 sufficiently small such that both /epsilon1V 1 /p and V 1 /p K/epsilon1 are smaller than η/ 3 .

Choose the region of perturbation B δ sufficiently small such that δ < ( η/ 3) p / (2 M ) p .

Since /epsilon1 and δ can be made arbitrarily small, for any η > 0 , the condition holds, and we get the global L p approximation error converges to zero.

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 a 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 it (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 to use the checklist as one of the factors in their evaluation. While "[Yes] " is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] " provided a proper justification is given (e.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 a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to 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 .

1. 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-attention universality - and no 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 in 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 it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Sec 5 'Concluding Remarks', paragraph 'Limitations' (lines 336-350) details assumptions on partition granularity, boundary effects, training difficulty, and distribution shift.

Guidelines:

The answer NA means that the paper has no 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 a 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 a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 fairness.

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 acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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 set 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-48).

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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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 or 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 experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

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 be 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 a 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 a 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 to 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 to replicate the results, access to 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 NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or 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. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: All datasets (MNIST, CIFAR-10, Fashion-MNIST) are public; an anonymized 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 NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) 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.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a 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) is recommended, but including URLs to data and code is permitted.

6. Experimental setting/details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to 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 of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment statistical significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

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, confidence intervals, or statistical significance tests, at 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 or the standard error of the mean.

It 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 a 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 of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments compute resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [No]

Justification: The proof-of-concept experiments run on a single commodity GPU, but exact hardware specifications and wall-clock times are not reported.

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 in the paper (e.g., preliminary or failed experiments that didn'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://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The work is purely theoretical/empirical on public data; no 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., if there is a special consideration 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 is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or 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., deployment of technologies 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 not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is 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, it 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 as 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 or unintentional) misuse of the technology.

If 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 a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

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 a 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or 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 use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a 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 are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset'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, or benchmarks.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an 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 to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: No human subjects or crowdsourcing 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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS 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 an equivalent approval/review based on 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 to 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 LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or 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 in standard writing support tools.

Guidelines:

The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.

Please refer to our LLM policy ( https://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

Peer reviews:
comments: 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.

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.

**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.
Overall, there is no major weakness in my opinion.

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?
- 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? 

Minor.
- I understand the definition of the softmax function is well-known but providing it in Sec. 2 will make the paper smoother. 
- 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
TITLE: Official Review
CLARITY: 4
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-06-16

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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..\n\nSTRENGTHS AND WEAKNESSES:\nStreangths:\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.\n\nQUESTIONS:\nAre 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.\n\nLIMITATIONS:\nNothing.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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).\n\nQUESTIONS:\nSee Weaknesses above.\n\nLIMITATIONS:\nAs 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\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\n\nQUESTIONS:\nOn 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\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 2, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 1, "ORIGINALITY": 2, "DATE": "2025-07-05"}