Attention Mechanism, Max-Affine Partition, and Universal Approximation

Thank you for your detailed review. We have addressed all your comments and questions in the following responses.

Q1. Why does the theoretical result use Attn∘Linear instead of Linear∘Attn as in practice?

Response: The paper deliberately reverses the order for theoretical convenience. It applies a linear (feed-forward) layer first and then attention (Attn∘Linear), whereas practical transformers do attention then linear. This reversal simplifies the universality proof by giving the attention mechanism a richer, “position-mixed” input to work with. In effect, the initial linear layer acts like an embedding or positional encoding step. It injects positional information and mixes the sequence tokens before self-attention is applied. This is crucial because vanilla self-attention is permutation-invariant (it can’t distinguish token order without position info). By encoding position and content via the linear layer first, the attention module can attend over a sequence where each token representation carries positional context, granting the attention full access to order-specific patterns.

In short, placing the linear layer first is a theoretical choice that ensures the attention sees an input enriched with positional cues (much like adding positional encodings), which is essential for expressive power. It also clarifies the roles in the proof: the pre-attention linear layer provides necessary position-dependent mixing, and then the attention layer performs contextual weighting over the entire sequence. This makes the one-layer transformer mathematically easier to analyze while retaining universality. Namely, the attention can leverage the position-encoded sequence to approximate any function (rather than being limited to order-invariant functions), and the final feed-forward can then map the attended context to the desired output.

Q2. Does the parameter count scale as $\mathcal{O}(\varepsilon^{-d})$?

Response: Yes, you are correct. To approximate a general continuous function in $d$ dimensions within error $\varepsilon$, the construction essentially partitions the input space into fine grid cells of side length about $\varepsilon$. The number of such small regions grows on the order of $(1/\varepsilon)^d$ (since in $d$ dimensions an $\varepsilon$-grid yields roughly $\varepsilon^{-d}$ cells).

The proof indeed leverages this kind of partitioning (approximating the target function by piecewise-constant outputs on a very fine grid), so the total parameters needed scale as $\mathcal{O}(\varepsilon^{-d})$ for a worst-case $d$-dimensional function. This is in line with standard uniform approximation results. It reflects the “curse of dimensionality.”

In other words, without additional structure or smoothness assumptions, the number of parameters (or regions) required to achieve error $\varepsilon$ typically grows exponentially in $d$. This mirrors other universal approximation schemes (e.g. using ReLU networks or Fourier series), which also require on the order of $(1/\varepsilon)^d$ basis functions or neurons to approximate an arbitrary continuous function on a $d$-dimensional compact domain to accuracy $\varepsilon$. Therefore, the parameter count in the transformer’s universality proof does scale as $\mathcal{O}(\varepsilon^{-d})$, consistent with expectations from uniform approximation theory and existing UAP proofs.

Q3 (minor). Consider including the definition of the softmax function in Section 2, even if it's well known, to improve flow.

Response: We have added a formal definition of the softmax function in Section 2, as suggested. This inclusion ensures clarity and improves the flow of the presentation. Specifically, we now define the softmax function as $\mathrm{softmax}(z)_i = \frac{\exp(z_i)}{\sum_j \exp(z_j)}$ when it is first introduced in Section 2.

Q4 (minor). Typo in Proposition 3.3: should the output space be written as $\mathbb{R}^{d_{\text{out}}}$ rather than $\mathbb{R}^{d}_{\text{out}}$?

Response: Yes are absolutely correct. This was indeed a typo. We have corrected the notation in Proposition 3.3. The output space is now written as $\mathbb{R}^{d_{\text{out}}}$, as suggested by the reviewer.

Thanks again for your kind words and detailed review! We have taken all comments into careful consideration and have made corresponding revisions to address the concerns raised.

We look forward to further feedback and discussion!

Thank you for your detailed review. We have addressed all your comments and questions in the following responses.

Q1. Are you able to revise the paper carefully, write it much more clearly and delete unnecessary and faulty parts in the rebuttal phase?

Response: Yes. We promise and have already conducted, and will continue to refine a line‑by‑line revision:

• Main proof: fully re‑structured into concise conceptual steps. All lengthy algebra moved to the appendix.
• Appendix: redundant material deleted, English polished, typos corrected, and every formula resized to fit the margins.
• Section 3: rewritten as a self‑contained intuition module with the tie‑breaking Assumption 3.1 made explicit. It can stay in the main text, be moved to the appendix, or be removed entirely.
• Ongoing updates: the draft has been under continuous revision since submitted because the results serves as the basis of several active projects in our group.

Although policy prevents us from uploading the new PDF at rebuttal time, these revisions are complete and will appear in the final version.

Since no PDF sharing is allowed, all revisions will only appear in the final version. Thus, our responses below are all in future tense (even if the revisions had been made).

W1. Section 3 is unnecessary. Appendix proofs have missing assumptions on affine coefficients.

Response: We respectfully clarify that Section 3 is included to provide intuition for our main results. In this section, we illustrate how attention can act as a max-affine partitioning mechanism and approximate an indicator function over those partitions. We wrote section 3 in consideration of the diverse potential readers of this conference. We hope people with relatively weak theoretical background are still capable of understanding the intuition of our work. That said, we acknowledge that Section 3 is not strictly required for the proofs. We are willing to streamline or relocate Section 3 (for example, moving some of its content to an appendix) to maintain focus on the core results, if the final publication format demands it.

For the assumption that regards $a_j $ and $b_j $, we do have noticed the assumption in the appendix isn't brought up in a standard and straightforward way. We do not think it is missing (as mentioned in remark D.1 in the submitted version), but we do acknowledge that it is not presented in a rigorous manner (that is, explicitly stated with mathematical expression and addressed as "assumption" instead of "remark"). We are sincerely sorry for any confusion caused. We have refined it's form and fix the above mentioned presentation problems.

Finally, we are willing to trim or relocate the proofs from Section 3 to the appendix (or further condense them in the appendix). So that Section 3 takes a smaller proportion and the main text remains focused.

W2. Main proof is unstructured. Derivation to $ L_p$ is trivial and redundant.

Response: We thank the reviewer for this valuable feedback. We agree that the proof of Theorem 4.1 can be presented more clearly. In the revised version, we will reorganize the proof of Theorem 4.1 to highlight its key ideas in a step-by-step, modular fashion. Specifically, we will break the proof into clearly-labeled steps focusing on:

• (1) constructing the max-affine partition of the domain,
• (2) approximating indicator functions on each partition region, and
• (3) reassigning function values on these regions to complete the approximation.

This structured approach will make the logical flow transparent and ensure that important ideas do not get buried in technical details.

For the $L_p$ derivation, we do admit it is obvious to people with a basic knowledge of the measure theory. The proofs of these derivations are simple and short so we don't think they cause very serious redundancy problems. The proofs are meant for people who aren't familiar with the measure theory, but we are open to cutting them off and make the $L_p$ derivations self-evident corollaries given in remarks.

W3. The core proof ideas are strong but obscured by long symbolic computations. They should be reorganized around conceptual steps.

Response: We agree that the proof of our main theorem is currently obscured by excessive symbolic computations. In the revision, we will reorganize this proof around its conceptual mechanism and key ideas, rather than lengthy algebra. Detailed computations will be either summarized or moved to the appendix, allowing the main text to focus on the clear, structured development of the core argument. This change will highlight the underlying concepts and make the proof more accessible to readers.

W4. Key steps are described in a convoluted way and assume readers have already parsed the full appendix. The exposition is circular and hard to follow.

Response: We acknowledge that the original exposition of some key steps was convoluted and hard to follow without the appendix. We will rewrite these parts of the proof to be self-contained and easily understandable, so that readers do not need to consult the appendix to grasp the logic. Every crucial step will be explained clearly in the main text, and we will remove any circular references. This revision will make the argument flow in a straightforward, reader-friendly manner.

W5. The appendix layout is messy: formulas overflow margins, formatting is inconsistent, etc.

Response: We will address the formatting issues in the appendix.

All formulas will be adjusted to ensure they fit within the margins, and we will enforce consistent formatting throughout the appendix. This includes fixing any overflow, aligning equations properly, and using uniform styles for text and math. The updated appendix will have a clean layout with no margin overflow and improved consistency.

W6. The English is sometimes faulty and needs proofreading.

Response: We have conducted 3 extra rounds of proofreading the entire manuscript for grammar and clarity. All grammatical errors and awkward phrasings will be corrected in the next iteration. By the final version, the paper will read smoothly and meet high standards of English usage.

Corrections. Section 3 lacked rigor, Proposition 3.2 was ill‑specified, norms were re‑defined, notation ($X_i,\mathcal{P},|\cdot|_\infty$) inconsistent or undefined, Theorem 4.1 and corollaries were repeated, trivial “technical highlights” and proof steps cluttered the text, matrix dimensions in (D.2) unclear, exceptional sets handled informally, validation logic for Lusin’s theorem and experiments opaque.

Response: To reiterate, we have made the following modifications:

• Simplify the proof sketch part of Section 3 to avoid over-specifying.
• Redundant norm definitions have been removed, and notation is unified: $X_i\in\mathcal{X}$ and a single partition symbol $P_{\text{ma}}$ replace ambiguous uses of $\mathcal{P}$.
• Repetitive material around Theorem 4.1 has been consolidated
• corollaries that follow immediately are now mentioned only in prose.
• Superfluous technical highlights and step‑by‑step derivations have been deleted or moved to an appendix.
• Appendix D labels every block matrix with its shape, resolving dimension ambiguities in $(D.2)$.
• Exceptional sets are further checked rigorously and revamped several obscure descriptions.
• Assumptions on Proposition 3.2 in the appendix is now stated in a more rigorous fashion.
• Extra rounds of proofreadings
• Finally, the use of Lusin’s theorem and the experimental validation are clarified, showing how partition error converges as predicted.

These edits address all reviewer concerns while making our draft shorter and clearer.

Thank you again for the kind words, and detailed review and proofreading (wow!). We believe the above changes and clarifications address all the points raised. The draft has been updated to reflect these corrections. Your efforts have helped us improve the clarity and quality of our paper. We are eager for further discussions!

Thank you for your detailed review. We have addressed all your comments and questions in the following responses.

W1. Missing citation of Lai et al. (2024) for ReLU-based attention UAP. Discuss differences vs. softmax attention.

Response: Thank you for pointing this out. We will add a citation to Lai et al. (2024) and discuss the distinction in our related work section. Lai et al. prove that with ReLU-based attention, the transformer’s attention module can be viewed as a piecewise-polynomial (cubic spline) function, and under certain assumptions this implies a universal approximation property. In essence, their analysis leverages the max-affine spline structure arising from ReLU activations, which partition the input space into linear regions.

In contrast, softmax attention introduces a normalization constraint that yields a nonlinear convex weighting of values. For example, a single-head softmax attention produces output $y$ as a convex combination of value vectors $v_j$:

and thus $\sum_{j=1}^n \alpha_j = 1$. This smooth softmax mechanism does not create the hard, piecewise-linear partitions that ReLU-based attention (or a “hardmax” selection) would. Consequently, our proof required new techniques: we approximate a max-affine partition using softmax weights despite the coupling caused by normalization. We agree that clarifying these differences will highlight the novelty of our approach, which tackles the softmax-specific challenges (normalization and smoothness) absent in the ReLU-based spline analysis.

W2. Contribution underemphasized: first UAP result for softmax attention (distinct from prior ReLU-based results).

Response: We appreciate the reviewer emphasizing this point. To the best of our knowledge, our work is indeed the first to establish a universal approximation theorem for a single-head softmax attention layer (with minimal linear pre/post-processing). We will make this significance more explicit in the revised paper. Prior universal approximation results in the literature either focused on attention variants with ReLU/hardmax-style activations (piecewise-linear splines as in Lai et al.) or required adding extra network components (e.g. Kajitsuka & Sato (2023) needed one attention layer plus two feed-forward networks to achieve universality). In contrast, our result shows that the standard softmax attention mechanism alone (with a single head and one layer) has sufficient expressive power. This novel finding matters because it theoretically validates the expressiveness of the vanilla softmax attention used in practice, without resorting to simplifications or additional architectures.

W3. Multi-layer extension: discuss if/how the result extends to stacked softmax attention heads (deep Transformers).

Response: We agree that extending the result to multi-layer (stacked) softmax attention is an important consideration. In principle, if a single-layer, single-head attention is a universal approximator, then a multi-layer architecture (with multiple attention heads per layer, as in real Transformers) should be at least as expressive. Additional layers cannot reduce expressive power, and could potentially approximate complex functions with fewer units per layer. Indeed, one could always let one layer carry out the construction from our proof while the other layers perform identity mappings. Thus, we expect that deep transformers also retain universal approximation capability.

That said, formally extending the proof to multiple layers is non-trivial. The challenge is that the composition of softmax attentions may not preserve the same neat max-affine partition structure we exploited in the single-layer case. Each softmax layer produces a smooth convex combination of its inputs rather than a strict partition. Stacking them can yield highly entangled representations where the “partitioning” of the input space from one layer is distorted by the next. This means a direct induction on layers would require careful new analysis. Additionally, while depth can theoretically be traded for width, a multi-layer constructive proof would need to coordinate attention heads across layers (and ensure the errors don’t compound). This is a significant technical leap beyond our current scope.

In summary, we believe multi-head, multi-layer Transformers are also universal approximators (since one layer already suffices), but proving it rigorously would demand new ideas to handle the composition of softmax layers. We will add a brief discussion to acknowledge this extension. We also note that analyzing deeper architectures remains an interesting open challenge (e.g., how successive softmax layers might refine or interfere with the learned partitions). This point will be clarified in the paper, and we thank the reviewer for the insightful suggestion.

Limitations. Result is constructive but not learnable in practice. Generalization bounds for softmax Transformers are still challenging.

Response: Thanks for pointing these out. To clarify, our goal is expressivity, not learnability. The theorem shows that a single‑head softmax Transformer can approximate any compactly supported function. It does not claim gradient descent will find those weights or provide sample‑complexity guarantees.

Recent progress begins to address these gaps:

• Learnability. Hu et al. (NeurIPS ’24) give a covering‑number analysis of softmax attention. Li et al. (’24) prove GD learns a 1‑NN rule with one softmax head.
• Generalization. Capacity and covering‑number bounds for softmax Transformers appear in 2110.10090, 2404.03828, 2407.01079, and 2411.17522

We will cite these works and note that extending universal approximation to provably learnable and generalizing algorithms is an open, promising direction.

Thank you again for the detailed review! We're open to any further questions or clarifications you might have about our work.

Thank you for your detailed review. We have addressed all your comments and questions in the following responses.

W1, “Linear layer” part of Q1 & Q2 The reviewer is concerned with the linear layer that is applied before the attention layer. The main reasons are as follows:

1. (In Weaknesses) The linear layer wouldn’t hold in classical frameworks where attention is applied.
2. (In Questions) The linear operator is not orthodox, which does not align with the structure of the linear layer in a transformer. Permitting such linear operations means allowing the model to duplicate the input sequence. Also, a bias term is added to the input sequences that varies along the sequence length. The Reviewer also raised concerns about the target function class not reflecting the permutation-invariant nature of attention.
3. (In Questions) Based on the second point, the reviewer thinks that our result on attention might not transfer to the setting of transformers. The reviewer also argues that allowing such a sequence-wise linear layer to precede the attention layer makes the overall architecture more akin to a generic FFN, since it skips much of the constraints of a transformer.

Response: Thank you for your thoughtful review. Our work focuses on the approximation ability of attention and thus uses a sequence-wise linear transformation mainly to keep the structure concise and the method easy to follow, which is also the style we strive for in section 3. However, we don’t think it is critically harmful to the genuineness and transferability of our work.

1. Resemblance to FFN We wish to clarify that the linear layer in our network is completely affine, that is, it doesn’t incorporate any non-linear activations. This also means our approach of theoretical analysis is vastly different from ones used in UAP of FFN. In FFN, the non-linear activations are separate and hence the whole model can be dealt with in parts. In the structure considered in our work, the softmax function is the only non-linear activation and it cannot be broken into smaller counterparts.

2. The Necessity of the Linear Layer In our UAT for single-head attention, we do need to duplicate the input sequence. However, that does not harm the transferability of our technique. In fact, duplicating the sequence is only necessary if we wish to use one head to approximate the whole target function. In practical settings with multiple heads, we can configure the network into transformer-style setting. For example, if we divide the target function into cutoff functions on separate supports, when these supports are small enough, each cutoff function can be approximated by a head that preserves the sequence length under a given precision. We admit that in the single head setting sequence-wise operations are necessary, but it does not defy our work’s purpose of demonstrating a method to analyze attention. Also, it does not harm this method’s transferability into transformer settings as discussed above.

3. Permutation Invariance and the Position‑Dependent Bias We believe you mean permutation equivariance. Because we added a position-dependent bias, the permutation equivariance is thus removed. The position-dependent bias can be seen as positional encoding and is widely used in theoretical works like

Questions (minor). K/Q projections expand input dimension ($d' > d$), unlike low-rank projections in practice.

Response: We acknowledge that our construction uses key/query projection matrices that expand the input dimension ($d' > d$), which deviates from the typical low-rank ($d' < d$) projections used in practice. However, this choice is a deliberate and mathematically valid design for the purpose of our universality proof. By expanding the feature space, the attention softmax can partition the input domain into a larger number of distinct regions. This enables more fine-grained max-affine partitions of the input. This finer partitioning is crucial for accurately representing complex target functions in our theoretical construction. Importantly, using a higher $d'$ does not compromise the theoretical integrity of our results. It remains within the standard definition of attention and is essential for achieving the universality guarantee.

Corrections. Fix scattered grammar/word‑choice issues, sharpen notation, clarify novelty vs. Kajitsuka & Sato [2023], tighten Assumption 3.1 wording and boundary cases, correct minor typos in Propositions 3.1–3.2.

Response: Thanks for pointing these out! We have made the following modifications in our latest draft accordingly.

• Grammatical edits include consistent tense, correct plurals (“dimensions”, “results”), and article use (“the attention module”)
• Ambiguous phrases have been replaced (e.g., “sum‑of‑linear transformations” → “a linear transformation”; “series of research” → “research works”)
• Notation is unified: $Attn_s$, $Attn_c$ are explicitly labeled as layers
• $Z_Q, Z_K$ are defined as query and key/value sequences
• $\lVert\cdot\rVert_\infty$ is “largest absolute entry”
• We also emphasize that our one‑layer, attention‑only network achieves universal approximation without skip‑connections or permutation‑equivariance, unlike Kajitsuka & Sato 2024
• Assumption 3.1 now rules out measure‑zero ties, so “the highest” is used throughout
• Boundary comments are added
• Proposition 3.2 references Lebesgue measure in $\mathbb{R}^{d\times n}$ and now reads “with the exception of a region of arbitrarily small measure.”
• All remaining typos (e.g., “a one‑hot”) are fixed.

Thank you once again for your constructive feedback, attention to detail, and encouraging kind words. Your review, particularly your suggestions regarding presentation and organization, has made this work more robust and accessible to a broader audience.

Please feel free to reach out if you have any further questions or need additional clarification about our work. Thank you!