The Effect of Attention Head Count on Transformer Approximation
Authors: Penghao Yu
Haotian Jiang
Zeyu Bao
Ruoxi Yu
Qianxiao Li
Venue: ICLR 2026 Poster (2026)
Keywords: Transformer
Approximation Theory
Lower Bound
Primary area: learning theory

Abstract: Transformer has become the dominant architecture for sequence modeling, yet a detailed understanding of how its structural parameters influence expressive power remains limited. In this work, we study the approximation properties of transformers, with particular emphasis on the role of the number of attention heads. Our analysis begins with the introduction of a generalized $D$-retrieval task, which we prove to be dense in the space of continuous functions, thereby providing the basis for our theoretical framework. We then establish both upper and lower bounds on the parameter complexity required for $\epsilon$-approximation. Specifically, we show that transformers with sufficiently many heads admit efficient approximation, whereas with too few heads, the number of parameters must scale at least as $O(1/\epsilon^{cT})$, for some constant $c$ and sequence length $T$. To the best of our knowledge, this constitutes the first rigorous lower bound of this type in a nonlinear and practically relevant setting. We further examine the single-head case and demonstrate that an embedding dimension of order $O(T)$ allows complete memorization of the input, resulting in the approximation entirely achieved by the feed-forward block. Finally, we validate our theoretical findings with experiments on both synthetic data and real-world tasks, illustrating the practical relevance of our results.

Published as a conference paper at ICLR 2026
THE EFFECT OF ATTENTION HEAD COUNT ON TRANS- FORMER APPROXIMATION
Penghao Yu
Haotian Jiang
Department of Mathematics National University of Singapore penghaoyuEu. nus
Institute for Functional Intelligent Materials
National University of Singapore haot anotian@ee an@nus edu.ss sg
wus..uuuuuuuuuuuuueue.uee... .du. edu
Zeyu Bao
Ruoxi Yu
Department of Mathematics National University of Singapore zeyußu. nus edu
Center for Data Science
Peking University
yuruoxi @estu.pku
www.:/p.uueuu....o. edu. cn
Qianxiao Li
Department of Mathematics Institute for Functional Intelligent Materials National University of Singapore qi gianxiao@nnu .ww.aaeeuuueuuuuuu..... sg
ABSTRACT
Transformer has become the dominant architecture for sequence modeling, yet detailed understanding of how its structural parameters influence expressive power remains limited. In this work, we study the approximation properties of transform- ers, with particular emphasis on the role of the number of attention heads. Our analysis begins with the introduction of a generalized D-retrieval task, which we prove to be dense in the space of continuous functions, thereby providing the basis for our theoretical framework. We then establish both upper and lower bounds on the parameter complexity required for e-approximation Specifically, we show that transformers with sufficiently many heads admit efficient approximation, whereas with too few heads, the number of parameters must scale at least ( for some constant e and sequence length T. To the best of our knowledge, this con- stitutes the first rigorous lower bound of this type in nonlinear and practically relevant setting. We further examine the single-head case and demonstrate that embedding dimension of order (T)) allows complete memorization of the in- put, where approximation is entirely achieved by the feed-forward block. Finally, we validate our theoretical findings with experiments on both synthetic data and real-world tasks, illustrating the practical relevance of our results.
INTRODUCTION
The transformer architecture (Vawwani et al. 2017) has become the foundation of modern sequence modeling, driving progress in natural language processing (Devlin et al. 2019 Brown et al. 2020). computer vision (Dosovitskiy ability to scale has enabled breakthroughs such as BERT, GPT, and ViT, making it the dominant paradigm across domains. Despite this remarkable empirical success, the theoretical principles un- derlying transformer expressivity remain incomplete. In particular, while universal approximation results establish that transformers can approximate arbitrary eequence-to-sequene mappings (Yun) et al. 2020a Pérez et al. 2021), much less is known about how their structural hyperparameters influence approximation efficiency.
202T), and multi-modal learning (Radford al.) 2021). Its
Among transformer hyperparameters, the number of attention heads plays central role. In practice,
large models often adopt head counts such as 32, 64, or 128 (c.g.. [Delin et al. (2019); Dosovitskiy
al. (2021): Touvron et al. (2023); Jiang et al. 2023) Grattafiori al. (2024) see Table 10 for

---

Published as a conference paper at ICLR 2026
more) ). yet this choice is largely heuristic: there is no principled understanding of how many heads are needed for a given task, nor of the costs incurred when the head count is insufficient. Theoreti- this question has so far been limited. Most existing results focus upper bounds, showing that transformers with sufficiently many heads or with extremely large embedding dimen- sion in the single-head case can achieve universal approximation or good approximation rate, but offering little insight into the limitations that arise when the head count is insufficient. Moreover, many analyses rely on strong simplifications-sueee as restricting linear embeddings, isolating the attention block, or linearizing the architecture. While these assumptions make the problem more tractable, they severely restrict the model's 's expressive power and prevent the derivation of rigorous lower bounds in realistic nonlinear settings.
cal progress
In this work, we address this gap by analyzing single-layer transformers on sequence-to- vector tasks. To this end, we introduce a new function class, the generalized D-retrieval tasks, which we design as a structured but expressive family motivated by retrieval problems. Each coordinate is defindd by (Xr) min.ss, fi( I..... D for subsets S; S [T', and the overall target takes the form (X,) = (2 (X))..... ...).... ED(X))). By construction, this class extends retrieval-style problems while being dense in the space of continuous seguence-to-vectto mappings, ensuring that results obtained in this setting reflect general approximation behavior.
The central challenge arises when the number of heads h is smaller than the intrinsic dimension D of the target. In this case, multiple coordinates 2; (Xr) must be represented by the same head, creating information bottleneck: the attention layer maps distinct sequences nearly indistin- guishable representations, forcing the feed-forward network to perform the separation. We show that overcoming this bottleneck requires parameter growth exponential in the sequence length T', namely O(1/E()) parameters for e-accuracy, thus establishing the first rigorous lower bounds for transform- ers nonlinear settings. In contrast, when h 2 D, heads can specialize distinct coordinates 24, eliminating the bottleneck and enabling efficient approximation.
Our results advance the theoretical understanding of attention by showing, that insufficient head count provably limits expressivity in realistic regimes. Experiments on both synthetic tasks and real-world retrieval data confirm that the predicted scaling laws persist in practice.
Contributions. Our main contributions are as follows:
First, we establish the first rigorous lower bounds for transformers in nonlinear settings, showing that when h < D. parameter complexity grows exponentially with sequence length.
Second, we provide constructive upper bounds, proving that h 2 D enables efficient approximation with parameter growth independent sequence length
Third, in the memorization regime, single-head transformers with embedding dimension > Td approximate by memorizing sequences, with the complexity residing in the feed-forward block.
2 RELATED WORK
Several works have studied the approximation and expressivity properties of transformers. The uni- versal approximation property first established in Yun al. (2020a), and later extended to trans- formers with sparse attention matrices in Yun et al. (2020b). The approximation rate of single-layer transformers with one head was analyzed in Jiang & Li (2024). Amsel et al. (2024) investigated how the rank of the attention matrix influences expressivity for a specific nearest-neighhor target can be constructed. They showed that when the rank is insufficient, the number of heads required for approximation grows exponentially, independent of sequence length. In a related direction, (Bhojanapalii et al. 2020) argued that setting the rank of the attention matrix equal to the sequence length enhances ex expressivity. Beyond dimensiimeliional settings. Takakura & Suzuki (2023) consid- (Tkkang
ered sequences of infinite dimension, characterizing approximation rates in terms of target function smoothness. Similarly, Wang et al. (2024) studied special classes of target functions and demon- strated that approximation error scales polynomially with the number of heads. In addition to these approximation-theorceee results, several works have investigated broader notions of expressivity. Dehghani et al.] (2019); Pérez et al.] (202T) established the Turing completeness of transformers, and Giannou al. (2023) showed that transformers represent arbitrary computer programs in
2

---

Published as a conference paper at ICLR 2026
looped setting. Finally, Mahdavi al. (2023) examined memorization capacity, proving that the number samples that can be stored scales linearly with the number heads.
3 PRELIMINARIES
Input and Output We consider the input space
X7 = { [0, [T]]
where [T] = (I....,T), We call T the length of input sequence. The output a single vector y E R², where I is independent of T and specified by the task.
For example, T the number of candidates, and the size of the output representation.
text retrieval task one may take to be the max number of tokens per candidate,
Input Representation. Each token is mapped to an F-dimensional vector by trainable encoder
Po [0,10 [T] R R
(2r, + (a( 8),
which jointly incorporates the content æ and its position s. Given Xx = {-(s)]]_ E Xr, the embedded sequence is
X = ( = Po(x(s), ER :s [T]}.
This formulation subsumes common designs where Pa combines a content embedding with either learned or deterministic positional encoding.
For example, if Emb(z) is a content embedding map and p(t) a positional code, then common schemes correspond
Pa(x((), Emb(x(1)) p(t) (additive PE),
()) = (Emb(x(()), p(t)
(concatenated PE).
Following common practice, we append a trainable classification token èo E RE to the sequence. The final input the transformer is
( = ( (()), E EE((++)
(3)
and the output ý is taken from the (7+ 1)-th position corresponding to âg-
Transformer Hypothesis Class
With the input space and embedding defined, we then formulate
the transformer hypothesis space.
We consider single-layer transformer based on the standard architecture (Vaswani et al. 2017 with two modifications introduced for analytical simplicity. Firstly, we omit layer normalization to simplify the analysis, while acknowledging its practical importance, and we conjecture that our key lower bound (Theorem 2 (2)) still holds when layer normalization is present. Secondly, exclude residual connections outside the feed-forward network. In the single-layer sequence-to- vector setting, where the output is read from designated classification token, the residual branch can be merged into the foed-forward transformation by reparameterization, thus these likewise do alter the expressive power of the architecture.
For an h-head, single-layer transformer, let denote the embedding dimension per head and E nh the total embedding dimension. The output
و = ( P(co + Wo Concat
Σ a " Wvrs(())).
(4)
where for each head Wo,,, WK, E R the query/key projection matrices, Wy. E R the value projection, Wo RREX is the output projection applied the concatenated heads, and
3

---

Published as a conference paper at ICLR 2026
F:R - R is a feed- forward network which we call it the feed-forward block. The softmax with scaling factor 3 defined by
(xp(3p(t))
ol0(() Er= (x exp()
3>
(5)
p(t)))
and 8 > O may be chosen arbitrarily large in order to make the softmax attention mechanism approximate hardmax.
We denote this family by
H(h, n, d, T, M),
(6)
the class of single-layer transformers with h heads, per-head embedding dimension input dimen- sion d, sequence length T, and parameter count M. Each E H(h, n, d, .)) is a mapping
H I2 +R'
implemented by the encoder Pe [0, 1]"h x [T] + IR"A, concatenation of the classification token ćo. a multi-head attention layer with projections (Wq.s, WK,,, Wv, Wo, and a feed-forward network F. Raa - R.. Thus H has the form equation with parameter count k referring to the weights and biases in FENs (Po. F).
Approximation Problem With the hypothesis class specified, now formalize the approxima- tion problem, which provides the framework for analyzing the expressive power of transformers. Definition (e-approximation) Let XT p be compact domain, and let F Nx - R be target function. We say that the hypothesis class H(h, n, d, M) e-approximates F on Xy if there exists H E H(h, n, d, T, M) such that
[(((X))) F(Xr) <c.
sup
XrENr
GENERALIZED D-RETRIEVAL TASKS
Target functions.
To motivate our construction, consider a simple one-dimensional example: Xr (Xr (x.).......(TT)) xxt) }]
with target
H(xT) = max r(t) + min x(t). ISrST
ISTT
This task requires the model to extract two distinct features from the sequence- the maximum and the minimum-heffor combining them. It can thus be viewed retrieval problem with two independent features being aggregated.
This example illustrates the broader idea behind our target class: retrieval-style problems where mul- tiple salient features must be identified and combined. We now formalize this intuition by defining the family of generalized D-retrieval tasks.
Mathematical Formulation Formally, for each i L,..., D, let fa define
[0, 1] + [0, be C2 and
1 T,
2((T) = min f((x(t),, iES,
S. S.
|S| 2
(8)
4
so that E(X,)= =
(x())) E 10, The target is then H(x)) F((x(XT))
where Fo : [0, + is C2, For vector-valued targets H : J0, 1]axT R2' defined with the same functions fi. subsets S, and an outer map Fu 1] + R., the extension is applied coordinate- wise, since each coordinate function of Fu can be approximated independently. Therefore, it suffices to consider the scalar-valued case. We denote by Fo."" the class all such functions H.
Related sparse sequence-to-ssqueeee retrieval tasks, such the q-sparse token regression (qSTR) model of Hoossavi- et al., 2025), where each output position depends on at most relevant input tokens, can be viewed as souuence-t--eqqeeee analogues of our formulation. Their results the sample complexity of single-layer Transformers with at least heads are complementary to our aoproximation--theneeii guarantees in the generalized D-retrieval setting.
4

---

Published as a conference paper at ICLR 2026
Assumptions on the target class For the theoretical analysis to be tractable impose the fol- lowing conditions:
Assumption (Model constraints). The model constraints are follows:
(1.1) The embedding Pa satisfies
(1)( 11,
E Xx E &r,
ensuring embedded inputs remain uniformly bounded.
(1.2) The post-attention mapping F is a two-layer forward network with 1-Lipschitz activation, hence universal approximator on compact domains.
(1.3) All weights in F and entries of the attention matrices (Wa.s WK,5, Wy,), Wo are bounded in magnitude by 1, ensuring stability of the model.
Assumption (Target class constraint). The target functions defined in equation satisfy the fol-
lowing:
(2.1) Each fi 10, + [0. 1] attains its unique global minimum 2; at a point x0)) E 10,
(2.2) The minimizers {-:(0) D are pairwise distinct.
(2.3) The Hessian V (zx((()))) is positive definite for all i = L....,..
(2.4) The gradient F ( has all coordinates strictly nonzero.
Remark. Assumption 2 excludes only degenerate cases while preserving broad generality for both the functions and the outer map Fo- More specifically: Assumptions (2.1) and (2.4) ensure that each ft behaves regularly around its minimizer. A degenerate example excluded by these assump- tions is fx( co for constant co. which totally independent of the input; Assumption (2.2) requires distinct minimizers, which allows partitioning the space into D disjoint basins around each minimizer _.(), A degenerate example excluded by this assumption is fi ʃ2 = .... fD: As- sumption (2.3) enforces sensitivity of the target to small perturbations near the minimizers, ruling out trivial flat cases (such as Fo for constant co) where no meaningful separation can be made.
Having introduced the generalized D-retrieval tasks, it remains to ask whether this family is suffi- ciently expressive. To address this, we now establish the universality of the target class: the family dense in the space of continuous sequence-to-vecto mappings.
Theorem (Density of the target class). For fixed d, T, the family {FD] is dense in C(XT).
That for every F E C(Xr) and every O, there exists D and f E FD such that
max [(() f(x) XEXT
The proof is deferred to Appendix A.I This density property highlights that our specially designed target family is not overly restrictive; rather, it forms a sufficiently general class to capture arbitrary continuous sequence-to-vecte dimension of this target, which means that it is the unique D < T for which the target H can be represented in the generalized D-retrieval task form.
mappings. Beyond density, we show that D is indeed the intrinsic
Corollary (Uniqueness of intrinsic dimension). 1f H can be represented by (ffi, 2 Fo) and (ff., 5, Fo), satisfying Assumption 1 and with DJ + D2 < 50 then D1 D2
This corollary justifies that comes from the intrinsic property of the target and is invariant across its form of representation. The proof is deferred to Appendix A.2 A.2
5 APPROXIMATION RATE OF GENERALIZED D-RETRIEVAL TASKS
Theorem [ establishes that the generalized D-retrieval tasks form a dense family in the space of continuous sequennce-to-vector functions. The next step is to analyze the efficiency with which trans- formers approximate these functions. To this end, we begin by stating two standard approximation assumptions regarding how well the fundamental building blocks of the target can be approximated. Assumption 3 (Approximation of components). We assume the following approximation properties hold.
(A1) There exist constants C1 > 0 and 7 > 0 such that for every 5 > 0, the function F6 [0, 1] +
E can be 5-approximated by two-layer feew-forward network 4 of width most C1/67, i.c.
sup (F((x) $s(x)) < 6.
=E(o.11P
5

---

Published as a conference paper at ICLR 2026
(A2) There exist constants C2 > and 7 0 (possibly different from (A1)) such that for each ? = 1...., .,. and every 8 > O, the function fü [0. 1 + [0, 1]
can be 6-approximated by a
two-layer wred-forward network Vi, of width at most C2/87. i...
sup (f((z) 12 [x)(()) & ==(00,1
These assumptions are reasonable: by the classical result of (Cybenko 1989), two-layer networks can approximate continuous functions on compact domains. In particular, if the Barron norm finite, one may take 2 (Barron [993); even in the worst case, setting
max(d, D) yields approximation rates comparable to uniform grid discretizations, which still suffices for our analysis.
We now present our main theoretical result. It establishes upper and lower bounds on the approxi- mation rates of transformers within the generalized D-retrieval framework. In particular, the lower bound in part (2) provides the first rigorous evidence that insufficient head count h < D leads to exponential parameter complexity, revealing a fundamental expressivity bottleneck.
Theorem 2 (Approximation rates of transformers). Fix d, T. Under Assumption hold for any target H E Fb?:
3 the following
(1) Sufficient expressivity with heads. For = D and embedding dimension n = 2 per head, there exists constant CA.D.T > O such that VM > Ca.o.T the hypothesis class H(,, n, d, T, M) e-approximates H.
(2) Lower bound with DD heads. For h < define
l: T 8 D 1)
-,
(n+ 1)s +1
then
min M H(h, n, d. T, M) c-approximates H 2(2).
(3) Single-head large embedding dimension. For h and per-head embedding dimension n Z Td, if the feed-forward block is 5-layer ReLU neural network, then there exists constant Ca.D,r 0 such that for all M Ca.p.r, the hypothesis class H(h, n, d T M) can e-approximate H.
Remark. We clarify the precise role of the assumptions and constants in Theorem 2
(1)Theorems 2 and 2 (3) require only Assumption 3 while Theorem 2 (2) relies only on As-
sumptions]] T and
(2) The constant in Theorem can be made explicit as Ca.DII = Ca,D (rT) ooo where > 0
is determined by the local behavior of the functions fi around ((()) and of Fo, and a.D depends only on d and D. This form is valid the regime d,
< T < l/e.
(3) The exponent coefficients in Theorems 2 (1) and 2 (3) differ because, in Theorem (3) the
network F also needs Q(T/E) parameters approximate the "max-like" operation. This yields bound of the form M < maxttl.) and for notational simplicity we write M I 1/c
We provide the detailed proof in Appendix A.2 We also justify the tightness of Theorem Appendix A.2.2
Theorem 2 highlights how approximation efficiency depends on head count: enough heads allow specialization, too few force inefficient compression, and single large head can rely on memoriza- tion. To illustrate these cases concretely, we now revisit the toy example from equation and discuss how each part of the theorem works in that setting.
Case (1): 2 D heads. Theorem (1) shows that when the number of heads matches the intrinsic dimension D of the target, the transformer allocate one head per component feature, allowing each head to specialize and leaving the feed-forward block to aggregate their outputs. This yields efficient approximation with ((M-1/) error for parameter count M, independent of sequence length T.
In the toy example with D = 2, one head naturally tracks the maximum and the other the minimum, so the task is solved directly without incurring inefficiency. This illustrates how having "enough heads" removes the unfavorable scaling in and explains the practical advantage of multiple heads beyond universal approximation results (e.g., Kajitsuka & Sato (003))

---

Published as a conference paper at ICLR 2026
establishes that when the number of heads is smaller than the intrinsic dimension D, the parameter count required to achieve a given accuracy can grow expo- nentially in the sequence length T. This lower bound highlights why insufficient heads lead to severe inefficieney. Intuitively, each head can be viewed as specializing in one coordinate of the minima structure in equation | When h D, single head must encode multiple roles simultaneously.
Case (2): h D heads. Theorem 2 (2)
In the toy example with = 2, one head is forced to capture both the maximum and the minimum across all T positions. Since softmax attention only produces weighted averages, the head must effectively encode information from multiple sequence elements simultaneously. As T increases, the number of relevant elements to distinguish grows linearly with T, yet they are compressed into an n-dimensional vector whose size does not scale with T. The feed-forwer block must then disentangle these increasingly entangled representations, which requires parameters exponential in T. This explains why the parameter requirement scales as ((1/eeT) and why the scaling improves dramatically once h 2/ D.
Theorem 2 is) proved with the following idea: (i) each head selects its most responsive locations (g), ty) in D disjoint minima basins around 20);; (ii) because < D, there is at least one segment G. C ((x r) that no head focuses on. We then consider the segment G; in it; (iii) along this segment (suppose it is G1), we construct many candidate subsequences and, by a pigeonhole ar- gument, obtain two subsequences Z1, Z2 whose post-attention representations are almost identical but whose contribution to S1 differs; (iv) these subsequences are then embedded into full sequences W1. W2, which the target function separates by at least 3c, while the attention block maps them within O(e*+1), forcing a large sced-forward network.
The intuition by which we derive the large network is different from geometric arguments in existing works such as (Yehudai & Shamir 2019). We directly made use of the fact that the network must be either large or have large parameters approximate function with large Lipschitz norm.
Case (3): single head with large embedding. Theorem 2 (3) shows that when the embedding dimension scales with the sequence length, E = n 2 Td, the model can encode the entire sequence into the classification token ég- Concretely, each input can be embedded e @ xx() TTd where e, is the f-th standard basis vector, so that trivial attention aggregates to
T (a(1)..... x (T)) [0.
which preserves the full sequence. The ferw-forward block can then recover the target relation efficiently. Unlike memorization of training data (Mahdavi et al.) 2023), this mechanism can gener- alize since it captures the relation itself. Moreover, approximating extrema functions such as max and min with shallow ReLU network is straightforward (see Lemma 9), requiring width ((T/e) for accuracy. However, this regime is impractical, it demands embedding dimensions that grow linearly with T.
As deeper transformer are more commonly used in practice, here we conjecture the extension of Theorem 2 (2) onto the L-layer case.
Conjecture (Multilayer transformer case). necessary condition for efficient approximation is L h 2 D. When the head number is insufficient across layers, we conjecture the lower-bound scaling for some constants az, by > 0 depending only on depth L to be
log PgpPaamCooun) = (1llog
lo
2x T nh
Experiments on the synthetic dataset in Section 6. with 2-layer transformer with no positional encoding and no layer norm has also verified this transition at D h. (See Table Appendix)
EXPERIMENTS
Theorem provides theoretical insights into how the approximation ability of transformers depends on the number of heads. In this section, we illustrate these insights empirically. We begin with syn- thetic tasks that mirror the structure of the generalized D-retrieval tasks, and then turn to real datasets (MS MARCO and CIFAR-10) to examine whether similar scaling behaviors arise in practice.
7

---

Published as a conference paper at ICLR 2026
6.1 NUMERICAL VERIFICATION OF THEOREM 2 WITH SYNTHETIC DATASET
We design a synthetic task aligned with the target class analyzed in Theorem Given a sequence
X (=(1),..,, of length T with x(1) E R the output is
max af xt))..
y
ISES
where at,..... a4 E I are fixed. Inputs are sampled i.i.d. from =(t) (4) For T {8, 16, 32, 64, 128} we generate 8000 training and 2000 validation examples.
On this task, evaluate single-layer transformers with head numbers {1, 11,2,3, 4,5) and fixed per-head embedding dimension. Each x=t) is embedded via two-layer ReLU MLP and concate- nated with a trainable classification token co. after which single-layer multi-head attention block (without residuals or normalization) processes the sequence. A two-layer GeLU MLP then outputs the scalar prediction. Both MLPs have the same hidden dimension N.
Then for each (h,T), T), models are trained under multiple random seeds. We report the minimal normalized mean squared error (NMSE) across seeds reduce optimization noise and highlight expressivity. NMSE, equivalent to R², corrects for the variance shrinkage of maxima as T grows, thus enabling fair comparison across lengths. Further training details and explanations are given in Appendix B.1.1
0.05
Validation NMSE (Min)
0.04
Sequence Length (T)
0.03
T-16
T-32
Log Hidden Dimension
lilden Dimentioon
4.0
0.01
T-128
2.5
Twa
T-10
T1:32
3.0
2c
0.00
Number of Heads
-388
Log Min Valiiattoo NMSE
-0.8
Log Min Validation NMSE
-b.5
-0.4
-6.2
-6.0
-5.8
(a) NMSE vs. Number of Heads
(b) Log vs. Log Accuracy (NMSE)
Figure 1: Results the synthetic example. (a) NMSE vs. number of heads h for sequence lengths
T E {8, 16, 32, 64, 128]. hidden dimension fixed at = 32. Note that there is transition at h = 4. (A table of mean and variance values corresponding these eurves is provided in Table (b) Log
Hidden Dimension N vs. Log Accuracy for different sequence lengths T. The parameter count k for the MLPs change linearly with N. (Plots for h 1 and h 2 are Figure
Figure Ta shows minimal validation NMSE versus head number h across sequence lengths T. Per- formance improves monotonically with h and exhibits clear transition near the intrinsic dimension D = 4. For h < D, NMSE grows with T, as limited heads must encode multiple extrema and the FFN becomes inefficient. Once h > D. curves flatten across T, indicating that heads special- ize to different coordinates and the FFN aggregates them very effectively, Normalization by NMSE ensures comparability across T, despite the increasing concentration of the max-oofGGussian target.
Figure bb highlights a phase transition between h = 3 and h = 4, with h = D = 4 equal to the intrinsic dimension of the target. When < 3, the negative log NMSE scales approximately linearly with the log parameter count (proportional to the MLP hidden dimension N). in agreement with Theorem 2 (2) Moreover, for fixed parameter count, larger T yields higher NMSE (worse approximation). Equivalently, as indicated by the fitted scaling lines, achieving the same error re- quires larger parameter counts when T increases, in line with Theorem (2) In contrast, for these trends change qualitatively. Validation error reaches the order of 10 indicating near- perfect generalization, yet the slope with respect to parameter count reverses: larger MLPs yield slightly higher validation NMSE, signature of tiny overfitting. The dependence on T also changes in this regime; see Remark I. in Appendix for details.
We also conducted experiments 32 constant(For h = 3, 5, we choose per- head embedding dimension to be (82/h], and total embed- ding dimension becomes F. 33, 35. See Table 3 in Appendix for details), well experiments
synthetic data with the widely used scheme of fixing E = nh
8

---

Published as a conference paper at ICLR 2026
on synthetic datasets with D 3 (See Table 4 in Appendix for details). Both of the above periments demonstrate similar trends to the experiments, implying the robustness of our results.
6.2 EXPERIMENTS ON REAL DATASETS
We conduct two additional experiments on real datasets assess the practical relevance of our the- oretical findings. The first is text retrieval task based on MS MARCO, and the second is an image classification task based on CIFAR-10. As there is no natural NMSE-like metric on retrieval tasks and accuracy is most widely used, we focus on training accuracy to isolate architectural expressivity from issues related to optimization or data scarcity. For completeness, for both experiments in Table 7 and 9 in Appendix. The experiments examine whether the phase also manifests in practice.
also report test accuracy
transition around the intrinsic dimenston D, predicted by Theorem 2
MS MARCO (text retrieval). We construct retrieval-style datasets from the MS MARCO passage
ranking collection (Bajaj et al. 2016), where each query is paired with one positive passage and 7 mined hard negatives (T {8, 16, 32, 64)). We train two-layer transformer encoder with per-head embedding dimension fixed at 32, varying the number of heads across {1, 2, 4, 6, 8, 10, 12, 14, 16). Input text is tokenized using the BERT tokenizer, and word, positional, and segment embeddings from pretrained BERT (Devlin et al. 2019) are kept frozen. These 768-dimensional embeddings are linearly projected to the embedding size E heads 32, after which only the projection and trans- former layers are trained. Full dataset construction and training details are given in Appendix B.2
CIFAR-10 (image classification). We further evaluate on the CIFAR-10 dataset (Krizhevsky, 2009) using four-layer Vision transformer (ViT) (Dosovitskiy et al. 2021). Each image of size 32 x 32 is divided into non-overlapping 8 x 8 patches (patch size = 8), which are linearly em- bedded. The transformer encoder has per-head embedding dimension fixed 16, and we vary the number of heads across h E {1, 8, 10, 11, 12, 13, 14, 16, 20, 24). To vary the sequence length, we extend the border with interpolation around each image to enlarge its side length, after which the
sequence length is given by some of the examples. Full dataset preprocessing and training details are provided in Appendix B.3
[mage longth
+ 1, including the classification token. Figure 6 shows
patch size
Result analysis. Both experiments exhibit the same qualitative trend in the synthetic setting. Figure 2a shows that in the text retrieval experiment, when < 12, accuracy declines as the se- quence length T increases, consistent with Theorem [ (2) Once h > 12, this dependence on T' disappears, and performance remains stable. Taking Err(h, T") = 1 - Accuracy(A (, ?") as er- ror, by using e exp(ah/T")
approximate (Err(h, 7") in log scale under MAE and drop-outs (h - 1, 12, 14, 16 are dropped out as outliers, 8 = 0.25 > 0, a = -1.40 < 0). figure 2b illustrates that when h 12, log( Err(h, T)) x h/T", highly consistent with the order in Theorem 2 (2) under fixed parameter count M. The flattening of curves after h > 12 is also consistent with theory.
Figure shows similar trend in image classification, with intrinsic dimension at 10. Figures 2d 1 and 2e illustatts weighted reversal score, calculated by R(h) wy Σr, <T; max max((err(T))
err (2)). 0) with normalization factor WA maxy err(T) - miny err(T), detects the existence of longer T yielding smaller error for this head number h. Such phenomenon leads to positive R(h), and it also indicates phase transition as explained in remark B.I.I Figure 2e further verified this.
CONCLUSION
In this work we investigated the approximation properties of single-layer transformers. We first introduced structured target family, the generalized D-retrieval task, that is broad enough to capture general sequence-to-vector mappings (Theorem []. Within this setting, we analyzed how the approximation efficiency of transformers depends architectural choices, especially the number of head. Our results indicate that having a sufficient number of heads leads to efficient approximation, while an insufficient number of heads forces the parameter count to grow exponentially with sequence length T. We also examined the single-head case, where large embedding dimension allows sequence memorization but shifts the complexity the feed-forward

---

Published as a conference paper at ICLR 2026
100%
10a
Train Accuracy (%)
80%
10-1
Error
60%
Seq Len (T)
10-3
Seq Len (T)
40%
T-8
10
T=8
T-16
T-16
20%
T-32
10-5
T-32
T=64
T=64
0%
10
1
10
12
14
16
Number of Heads
Number of Heads
(a) Accuracy vs. Number of Heads for different T (Text Retrieval).
(b) Log(I-Accuracy) and its prediction (Text
Retrieval)
Train Accuracy (%)
100
Reversal Score
Reversal Score
Seq Len T-65
T=145
60
T=257
2-
T-1025
24 810121416 20
24
4
10121416 20 24
2
4
5
Number Heads
Number Heads
Number Heads
(c) Train accuracy (%) vs. Number (d) of Heads (Image).
Weighted Reversal score vs. (c) Weighted Reversal score vs. Head Number (Image)
Head Number (Synthetic)
Figure 2: Experiments on real datasets. Training performance with different numbers of heads h across different sequence lengths T. (a) Accuracy number of heads for different T in text retrieval; phase transition near h 12. Mean and standard deviation see Table 66. MRR shows a similar trend, Fig. 5 in the appendix. (b) Phase transition for text retrieval. (c) Accuracy vs. num- ber of heads for different T in image classification; phase transition near h = 10. Mean and standard deviation see Table 8 (d) Weighted Reversal Score for Image Classification, err
Accur acy. The plot becomes positive when h 2 10, indicating phase transition. (e) Weighted Reversal Score for Synthetic Experiment, becomes positive h exactly the intrinsic dimension of the task.
block (Theorem 2. These findings clarify the roles played by head count in transformer expressivity.
Our experiments both synthetic and real datasets reveal non-trivial phase transition around the intrinsic dimension D, consistent with theoretical analysis. When the number heads below D, models exhibit higher error for the parameter count sequence length T increases. Once the head count reaches or exceeds D, approximation rate becomes independent of sequence lengths T. This transition is also observed in real-world datasets with deeper architectures, indicating that the notion of intrinsic dimension is not only theoretical but also practically relevant. In particular, beyond fully training models, analyzing head contributions early in training to estimate how many heads meaningfully affect the output, or training multiple models with varying depths and head counts while tracking how error scales with 7 are potential ways to probe the task's intrinsic dimen- sion. These experiments might demonstrate whether the inferred intrinsic dimension is stable across architectures, thereby informing head-count selection and the head number to retain under pruning.
Limitations. We conclude by noting several limitations of this study. Firstly, although the lyzed target class is dense, the phenomena interest are most naturally manifested retrieval-style tasks aligned with our setting. Secondly, our analysis restricted while experiments on real datasets supports Conjecture in deeper architectures, a rigorous multi- layer theory remains open. Finally, the tradeoff between sequence memorization and pattern learn ing observed for shorter sequences (cf. Remark B.I.I has not yet been established rigorously and warrants further investigation.
single-layer transformers;
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
This research is supported by the National Research Foundation, Singapore, under the NRF fel- lowship (project No. NRF-NRFFI3-22011-00 The computational work for this article was fully performed on resources of the National Supercomputing Centre, Singapore (htps://www.nsecc.s)
REFERENCES
Noah Amsel, Gilad Yehudai, and Joan Bruna. Quality over Quantity in Attention Layers: When Adding More Heads Hurts. In The Thirteenth International Conference on Learning Representa- tions, October 2024.
Payal Bajaj, Daniel Campos, Nick Craswell, Li Deng, Jianfeng Gao, Xiaodong Liu, Rangan Ma- jumder, Andrew McNamara, Bhaskar Mitra, Tri Nguyen, al. Ms marco: A human generated machine reading comprehension dataset. arXiv preprint arXiv: 1 1.09268, 2016.
A.R. Barron. Universal approximation bounds for superpositions of sigmoidal function. 1EEE Transactions on Information Theory, 39(3):3309945, May 1993, ISSN 1557-9654, doi: 10.1 109/ 18.256500.
Srinadh Bhojanapalli, Chulhee Yun, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Low- rank bottleneck in multi-head attention models. In International conference on machine learning, pp. PMLR, 2020.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
G. Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of Con- trol, Signals and Systems, 2(4):303-314, December 1989. ISSN 1435-568X. doi: 10.1007/ BF02551274.
Mostafa Dehghami, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal Transformers. In arXiv: 1807.03819 /Cs, Stat], March 2019.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. 4rXiv:1810.08805 [cs]. May 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An Image Worth 16x16 Words: Transformers for Image Recognition at Scale. In International Conference on Learning Representations, 2021.
Angeliki Giannou, Shashank Rajput, Jy- Yony Sohn, Kangwook Lee, Jason D. Lee, and Dimitris Papailiopoulos. Looped Transformers as Programmable Computers. In Proceedings of the 40th International Conference Machine Learning, pp. 11398-11442. PMLR, July 2023.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The Ilama 3 herd of models. arXiv preprint arXiv:2407.21788 2024.
Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chap- lot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anee Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mistral 7b. arXiv preprint arXiv:2310.06825. 2023.
Haotian Jiang and Qianxiao Li. Approximation Rate of the Transformer Architecture for Sequence Modeling. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024.
Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion scale similarity search with gpus. JEEE Transactions Big Data, 733::535-547, 2019.

---

Published as a conference paper at ICLR 2026
Tokio Kajitsuka and Issei Sato. Are transformers with one layer self- using low-rank weight matrices universal approximators? arXiv preprint 1XXi:23377.10233 2023.
A. Krizhevsky. Learning Multiple Layers of Features from Tiny Images. In Technical Report O, 2009.
Sadegh Mahdavi, Renjie Liao, and Christos Thrampoulidis. Memorization capacity of multi-head attention in transformers. arXiv preprint arXiv:2306. 220000200000, 2023.
Alireza Mousavi-Hosseiie Clayton Sanford, Denny Wu, and Murat A Erdogdu. When do trans-
formers outperform feedforward and recurrent networks? a statistical perspective. arXiv preprint arXiv:2503. 12272, 2025.
Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is Turing-Complete. Journal of Machine Learning Research, 22(75):1-35, 2021. ISSN 1533-7928.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, al. Learning transferable visual models from natural language supervision. In International conference machine learning, pp- 8748-8763. PmLR, 2021.
Stephen Robertson and Hugo Zaragoza. The Probabilistic Relevance Framework: BM25 and Be- yond. Found. Trends Inf. Retr., 3(4):333-33) April 2009. ISSN 1554-0669. doi: 10.1561/ 1500000019.
Shokichi Takakura and Taiji Suzuki. Approximation and estimation ability of transformers for seuuence-to-eqqueeee functions with infinite dimensional input. In International Conference on Machine Learning. pp. 14446633447. PMLR, 2023.
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 arXi::2302. 1397,,
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- fion processing systems, 30, 2017.
Mingze Wang et al. Understanding the expressive power and mechanisms of transformer for se- quence modeling. Advances in Neural Information Processing Systems, 37:25781-2555 2024.
Gilad Yehudai and Ohad Shamir. On the power and limitations of random features for understanding neural networks. Advances in neural information processing systems, 32, 2019.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J. Reddi, and Sanjiv Kumar. Are Transformers universal approximators of soguence-to-sequeeee functions? In arXiv: 1XXX::91/2.1007 [Cs, February 2020a.
Chulhee Yun, Vin-Wen Chang, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and San- jiv Kumar. O(n) Connections are Expressive Enough: Universal Approximability of Sparse Trans- formers. In Advances in Neural Information Processing Systems, volume 33, pp. 13783-13794. Curran Associates, Inc., 2020b.
12

---

Published as a conference paper at ICLR 2026
A PROOFS OF MAIN THEOREMS
A.1 PROOF OF THEOREM 1
Proof Sketch. The proof proceeds in three steps. First, by Lemma I we approximate a broader function class that relaxes the smoothness requirements and the assumptions in Assumption 2 Sec- ond, Lemmas 2 and [ show that by constructing appropriate S,, we can faithfully recover all in- formation from the original input sequence with simple ft. Finally, the outer function Fo can be applied to approximate an arbitrary sequence-to-vector target within this class. Together, these steps establish the result.
Proof of Theorem
7 To prove Theorem we first establish a few auxiliary lemmas.
Lemma (Relaxed target class and closure equivalence), Let Fb be defined as in Ibut with only 1]0, [0, 11). i.e. we drop "unique minimizer", "pairwise distinct" and "PD Hessian" .Set fi E 1[",
Fd..T = -UD21 FD² and FulT = UD21 Jo.
Then
11 = FdTT
Fd,T
on Nr.
Proof. Fix H(X7)= ((.,(X),...,E211..
FAT with (XT) = mings, [((((() and E D
(0,, Let > 0. We will construct H E Fd. with || l S E. Firstly, by F E C([O, 1I.
Stone- weieessssssssss, choose p: E - (10, 1ja) so that |pp; - filloo I where > 0 will be fixed later. Because the uniform approximation can slightly leave [0, 1J. compose with a smooth strictly increasing squashing [-c, 1 + c| + [0, 1] with s(u)
on [0, and ||s pp pilloc S (for small enough > 0), and replace p. by o Pt. We still write p. and retain |p: - fill S 2n.
Secondly, let & E arg jo, (xxx) (nonempty by compactness). Pick r E (0, 1 small and C bump o, supported in (f,,, 2r) m j0, 1J% with (E,) = 1. Vol(E.) = 0, and with ('). negative definite Define, for parameters òr,1, 2r,2 > 0 to be fixed,
(x(x) = p=(x) 8,,1 dx(x) + 81,2 (x) || E.1l
(i) Since gi(E,) = (C.) - 6,1 while (x) 2 pz(x) whenever (x(x) = 0 and gi(x)) > pr(x) for BEi 2r) \ (Ei), we get that & is the unique global minimizer of gi-
(ii) At E,, because Vol(() = O.
V'iitE)) = vPpc(E)) 5,1 Vori()) + 28,
Here - (VDoi(()) > 0, so choosing (5r,1, )x,2) suitably makes (Es) > 0 (PD). Because ||0:||>0 < 1 and || - 1 I d on [0, 1]d
11.. polloo S 6,,1 + 81,2 (2r)
Hence, by taking small and then ,,,,, و1,2 small (using the - scaling in (i to keep the Hessian PD), we can ensure both PD at & and lig: - -olloo S 17.
Thirdly, we use tiny translation remove distinctiveness. It may happen that & &vr for some of i'. Choose pairwise distinct small vectors v E IRd and fix a smooth cutoff E C 1J [o, 1)) that equals on [r, 1 r]d and vanishes near the boundary. Define a C diffeomorphism of the cube,
(x)x) == x E, x(x) Dss
with E, > 0 small.
Then 4 is arbitrarily close to the identity in C² for small e,, maps [o, 1]d to itself, and h; = gr o has (unique) minimizer at 10( क- (E.). For different i, these points are distinet if the v, 's distinct and e,'s are small but nonzero. Moreover, because Vg:(E.)) 0, the Hessian at x() satisfies
())(((()) = ( V (E,
D DT.(x()))
> 0,
For instance take (x) d1[ - //1///2) with (0)) = 1, v 0 near 0, = 0 on [1, no); then 2 (fi) x 0 and its scales like
13

---

Published as a conference paper at ICLR 2026
PD is preserved. Since 9i is Lipschitz the compact cube, ||,, gilloo I E for some Li, hence by taking E, small we get ||,; 9sllas $ 17.
Finally, if needed, compose with the same strictly increasing squashing in Step and set
fo so h4 ( 1] [])
Because is strictly increasing, it preserves the minimizer location and, at the minimizer 2(1) (s0 (z)()())) s(())((())) 2 (x(9)) 0. Also Il hilloo S 1j by construction.
Collecting the bounds from previous deduction:
Ilfe filloo I lpp fillo + 9. Poloc + ||4 gillee + llfs helloo S 5m.
S2n
For each i. the map + minges, ue is 1-Lipschitz in || Iloo- Hence the corresponding features 2(xr) = mines, fi (=(()) and 2 (Xr) = mintes, fi())t) satisfy || - /o/oc S 5n.. Let wi be a modulus of continuity of Fo on [0, _IP. Choose so small that wF, (5n) S s/2. Then
(=(T)) - Fo(E(X)))ll 2e/2.
Finally, approximate Fo uniformly on [0, lll by some Fo C² 119) within e/2 (Stone Weierstrass), Setting
H(XT) = (XT ((x(()))).,,,((x)) E FTTT
we obtain
H H1 < [Fo(2) Fo()) + Fo(2) - (o( S 8.
Se/2
Se/2
This shows Ful C _dd 1o
The reverse inclusion is simple, hence we have the lemma.
Remark. Thus, we now focus on the relaxed class F.d.T and Lemma lifts the result to the original class Fd...
Lemma (Order-statistic in the relaxed class). Without loss of generation, suppose 4/T. Let m
T For each j [d] and Xy = U,(XT)= max min a()j
(T)),, define
4.
BC uEB
and each fixed 1 E [T).
(X)( =
max ACITI | 1 m tEA
min xuin)
Ze5(XT) := max AC
min(1 z(u);). aEA
uEd
|
TEA
(T), and set U, = Um.j- For the Let U1.j ... 2 vry be the sorted values of (nonincreasing) and wis -.. S WT. multi-set (a(u)} E [T]}, let Vıjj 2 ... 2 Vr. (nondecreusing), Then we have
YeJ = min(x(t),, Umil.
1 Ze, = max(x(()s, Wm.sh.
In particular:
a(1); 5 Un.j 3 Yr, = =z)=<(())-
x(1), 2 Wmj + 1 - Z = x(1))-
Proof. Among all m-subsets B, the maximum of minwe a I(u), is attained by picking the m largest coordinates, U, = Um.j. Forcing é E A. the choice of the other minimum is the m - 1 largest among {a(u), : u * t}, hence Y,, = min (x(t); Um.5). For Zıj, note that min x(u),) 1 max 2(u);- as minimizing m z(u),. which picks t plus the (m - 1) smallest leading to Zeg. The particular statements follow immediately.
indices to maximize the
maximizing it over A is the same
14

---

Published as a conference paper at ICLR 2026
Lemma 3 (Smooth selector). For q define
(t) == ea( 1- 71. um. m )² (1 Z1,5) cal Mis-umm l Ye.j
ea( 1-1-22 11, Wm,, s V1. Vem
Then (1), - z(t) uniformly on Xy as + 0c.
mm Proof. From Lemma 2 if (t); 2 Uj. we have Yıj f and (Yr, Um 2
0 while 1 Z
:(t) and (1 Zt.j Wm. ) 0. Thus, as q 0o, the weight concentrates on (1 Zt,)) x());
If x=tt; I Uj and Wm.jj 2 x(t),, we have (1 Ze, Wm,)² )² 0 while Yi, x(t), and
(Yo,) Um.j) )² > concentrating on Yı,- xx())- If Wm.j I x(1, S Uj. we have Zz.j
Yi, =((),. either of three settings leads to x(1),- j. The compactness of Xr gives us the uniform property.
Now, we begin formal proof for Theorem
By Lemma 1 for any F C(Ar) and 0, suffices construct an H Fd.T with || HI I e/2, since the lemma could lift it to fdd, with another e/2.
Fix m
F. For each coordinate j and each S [T] with [S] = m, include the relaxed primitives
( = min at)s, rES
2018) = min x())))
ES
To form the thresholds Um.jj and Wm.j needed in Lemma
additionally include:
min
(t); for all S S (T) \ (t) with |S| rES
and
min(1 xx())) for all S S (7) (1) with |S| T m, EES
Using smooth log-sum-ext (softmax) in the outer function Fo. we can recover the subset-wise max- imum required compute U, Yi,; Zt,j. Um.j and um.y from these primitives.
By Lemma for any >> there exists q such that
(t,) x((ss] s 6.
max max X (E(T),
By uniform continuity of F on the compact Xr, choose 6 so that this implies 1F(X7) F(Xx(T)) S =/4 for all Xr. where Xy(T)) stacks the coordinates (t) We approximate the con- tinuous map + Fu) on [0, 1J uniformly by a polynomial P within e/4 (Stone-Weierstrass) Define
H(Xr) = (P vee) (8,)({)(T));
which is a C² function of the inner features. Hence 11 E Frd and
F 11 S F F & + | 8 -
Po Po 8 & 8/2.
Se/4
Se/4
Apply Lemma 1 to replace each relaxed primitives by admissible C functions with unique mini- mizers and replace Fo by function Fo so that the final error increases by most e/2. This leads to f e FAT with || - ~ S.
A.2 PROOF OF THEOREM 2
A.2.1 PROOF OF THEOREM 2 (1)
Here prove Theorem (1)
15

---

Published as a conference paper at ICLR 2026
Proof Sketch. The idea is straightforward. Each attention head is assigned to approximate one term minees, ft(). Once these components are extracted, the outer function Fo can be approxi- mated by a suitable F, completing the construction.
Proof of Theorem (1) Sufficient expressivity with D heads. Fix d, T and a E (0, 1), and let >> 0 be given. Throughout the proof, constants depending only on (d, D, T) are absorbed into Ca.D.T O which may change from line to line.
In the display equation 4 each head produces an n-dimensional vector and Concat 4 gives a vector in Th before F. For the construction, we realize the usual block-by-head parameterization, which means that the encoder outputs a block decompeed embedding
x(1)
(1) RDDD
(1(( E ER
and the i-th head only reads the i-th block via block- diackna Wo.+. Wv,, (entries bounded by 1). This keeps the parameter counts within the same order. We therefore set the per-head embedding dimension to n 2.
Firstly, from Assumption (A2), for any 5 > 0 there exist two-layer FFNs Vi, 10, 1] - [0, 1] such that
C2
max ==[0,1]4
[f(()) - -W; (z)s S width(W...) S
(10)
8%1"
where Y/ 0 is the exponent from (A2). Define for each head the position gate
0,
s SE
r((s)
S [T.
-1,
s $ Sr
(Recall S; C [T] with |St] 2 aT by equation 8
We implement the encoder Pg. that its i-th block
is
(t) (L,((ttt)), ri(t)) E [0, 1] x {-1,0} c 111.
(11)
This choice follows ||á(t)]22 S V2D. After a fixed rescaling (absorbed into 3), this meets the norm constraint.
Secondly, would like head-wise attention to isolate the minimum on S. For each head we take a single attention logit (m, = by choosing
Wo.: = T,
WK.- [-1 1),
WQ = Wv, 1z.
All entries are within the allowed bound 1. With the block equation token t head i
I1 the (pre-softmax) score of
pt(t) (W (t)) (Wo, = -v.. + rtt).
(12)
Let [[p;] be the softmax equation 5 with O. Define the head-i value readout (first coordinate of the head output)
(XT) Σ dlpl(t)) (x(())))).
(13)
Here, the second coordinate is unused. F could ignore it via a fixed linear projection, counted in the constant Ca,DTT
Now, we give a uniform bound on S,. Take a; == (x(((((()) E [o, 1] and split the sum into S and ST. From rt())= 0 on S; and ri(t)= -1 Sy, we have
e Dar
t Si
3 ((-aetr((1)
Σ Eues, e-
alp:](1) Σ Eves, e Bau + Evese e- 3(a++1) S
-Ba
e 8
E Sp.
EnEs,e e-
16

---

Published as a conference paper at ICLR 2026
Hence, we have
Enes, are ape Enes, 3au
te-8 Enese are 8a
= E clai(()at +
Σ oloo)(()) BEST
(14)
Eues, -Ga
LES.
Gibbs mean St
To simplify, we denote a, == minkes, and b; == E [0, for E S.. Then, we have
Eies, arE Sa,
+ Eies, b
a+ Σ bie a, S -
Enes, Ban.
Enes, Bbu
eB
LES,
The inequality comes from su be- 9b denominator in the middle fraction is 2 1. For S term in 14 we use S 1 and Eues, e Sau >
c-118 for 3 1 and one of the by 0, so the
1 to get
- Eiesy a.e Sa,
ΣvEs, E-Bo. S Se S P..
Combining the two bounds, we have the uniform estimate:
|S 1
min 1. (z(t)) < 3 (Xr) S min (x)() rES;
T
(82 1).
(15)
teS,
eB
In particular, since |S|| - T, there a constant CT with
0 S (Xr) min v (x(()) s Cr ( + (
Cr max(Te, T}.
(16)
tES:
Thirdly, we need to lift bounds from & to Zi. From equation and the definition of 2,.
min
[x(((() min Va. (x(()) rES;
s 6.
rES:
Together with equation 16
(xxr) (XT)) 8 + CT (3 + for all Xy Hr, i i= I....,.. (17)
Let Lg = sup=c(o.10
| FVF((>)ll < so (compactness and C.). Choose
B B. == = 1, 4CTLoD log( (CCLoo))]
8 = D
E
Then by equation 17
[E(XT) =(x
2L for all XT E Xr.
(18)
Finally, constreut the approximation for Fo and count the number of parameter. By Assump- tion (A1), there exists a two-layer FFN 4% [0, 1] - with width < C1/58 (for some "o > 0) such that
(b(2)) 4 (2) < 80. D
max
Set 6n :- 8/2. Define the model's final ceed-forward fe to project R2D RD by keeping the first coordinate of each head (a fixed linear map with entries in {0, 1}) and apply Ps..
Then for all Xr. we have [Cooncat, Fo(2(XX)) < 10, (ă(XT)) - 25 (=(Xr)) 2. (=(XT)) Fo(2(XT))
< Lo ([x((T)) - =XX))l1 + 50
< Lo D ||ā(X)) (00)==(Xr)11)) =(Xr) +e/2
< 8/2 + e/2 = s,
where we used equation 18 in the last inequality.
Here, the trainable components composed of three parts:
17

---

Published as a conference paper at ICLR 2026
the D subnetworks 14,5 inside the encoder blocks equation
the fixed-size projections WQ WK,,, Wy, (size O(D) and independent of E);
the two-layer FFN 2, used inside F.
Thus
1 87S
M
C
C"
C."
for constants C², , " Ca.D.T-
,
With 8 = (e) and 6n (s) chosen above,
Ca.D.T M
7 max 70),
E7
and the construction uses h D heads with per-head dimension = 2 and achieves C- approximation on Xy. This proves Theorem (1)
A.2.2 PROOF OF THEOREM 2 (2)
Proof Sketch. The argument proceeds in two parts. The core idea is construct two sequences whose representations after the attention layer are indistinguishably close, on the order of O(e*+1), yet whose target outputs differ by least 3c. Lemma then implies the lower bound on the param- eter count required for approximation.
Using Lemmas D > h, there exists least one neighborhood selected by the heads. Within this gion, the pigeonhole principle guarantees the existence of two distinct subsequences. By carefully designing these subsequences, we ensure that their outputs after the attention layer are nearly distinguishable, while their target values differ by at least 3c. Extending them to full sequences completes the construction.
6 and we obtain D disjoint neighborhoods around the minima x(0) Since
We now turn to the full proof. To establish Theorem 2 (2) we begin by introducing several auxiliary lemmas that will serve as building blocks for argument. Lemma and are only to set up the approximation problem into more tractable form.
Lemma 4. Let U1, E R Suppose
|| 02112 S and [1²(v))- F(22) 2 B,
where F. RR Rm is a two-layer feed feed-forwar network satisfying the constraints stated above.
Then F must use at least
B Avn
Ω
parameters.
Proof of Lemma 4 with width p be
Let Ar 01 U2 and AF = F(v1) ((2). Suppose the two-layer network
F(x) = (U + b) +c,
where Rexn V Rmx b E R E Fm is 1-Lipschitz acting coordinate- wise and every entry of U, V,b, has magnitude at most 1.
For the j-th output coordinate, have
AF
Σ uo + -o bb))
where u. is the r-th row of U. Using the 1-Lipschitz property of and Schwarz inequality, we have
AF <
Vyrl |u Ar
<
Virl ||u ||A2||2-
18

---

Published as a conference paper at ICLR 2026
$ 1. Therefore, for all we have By the entrywise weight bound, ||url22 vn and [Vrr
| I Pvn |
(19)
be the norm used in the lemma statement. By equivalence in finite dimensions, there exists Cm E (0, 1] depending only on the chosen norm and m such that
Let ||-
lllll
=
blvloo for all E R
Cm
|AF] 2 B implies | 2 Cm so there is some jr with
J Z Cm B.
Combining this with equation 19 and || S A, we have
Cm B
Cm B S PVnA
p
Avn
Finally, let Peff S p be the number of hidden units that actually affect the output, i.e., those with nonzero row in U and a nonzero entry in the j"-th row of V. The above bound holds with Poff in place of p. hence Peft Z Cm B/ (AA/t) Each such unit uses least one nonzero parameter in and one in V, so the parameter counts satisfy k 2 Pet. Therefore
Cm B Avn
B AVn
l 2
Ω
which proves the lemma.
Lemma 5. There exists R O such that for every E (1...., D} and every v < R. there exist constants > O and Lg with the following property: there exists a segment G, C B(x of length & such that
[f(x)) - S(()) 2 Lg 1| -1l2.
Vz,ye G.
and moreover
f1(x) > 2i.
Vx G5:
Proof of Lemma 5 Fix i E {l...., D} and denote I² := x(0,) f ;= fi and H. == V=f(x*). By positive definiteness, let Amin (... By continuity of V2 there exists R such that vaf(2) = +
for all x ((x", R
Set μ4 X//2 0. Because the domain is [0, 1]d and x² E [0, 1] could choose unit vector 2/ pointing strictly into the cube at x² (if x² is interior, take any unit vector). Define
7 == supst t
x² sv; [0, 1]dt for all e [0,t]} >
and set R; = min R n... Take R = minggig R > O. Fix any r E (0, R) and consider the restriction
g(t) = f(* + tvi),
t 10
Then
(() = v. f(* + tv.) 0z 2 μs for all t E [0,r].
Since x* minimizes on 1] and i is feasible inward, we have g(t) > g(0) for small t 2 leading to the one-sided derivative g' (0+) 2 0 (if xt is interior then Vf(**) 0 so g'(0) 0). Because g" 2 púx the derivative g' is increasing and thus
g((1) 2 g0(++) + +ual 2 pal,
[0,r]]
Let : r/4 and b r/2 and define the segment
G; {x* + to; te b]} C B(x*.r).
whose length 6, = h - a = r/4. For any x = x² + bu; and x² + s0; in G3 with t. > 8, the
mean value theorem gives some E E (s, [a, b] such that
[f(z) - f(v) = -g(()) g(s) = Ig' It - s| 2 n4a - = (>rr)
-(axr) 2 8112.
19

---

Published as a conference paper at ICLR 2026
Therefore the choice
L
0
8
works uniformly for all x; E Gj- The segment G; does not contain x², for all its points are at distance at least r/4 from x². By uniqueness of the minimizer, f(x) f(x*) 2, for all x E Gi.
The lemma holds with 6; r/4 and L, (Amin (2
Lemma 6. Let (21...., denote the minima defined above. Then there exist constants ro and Lo > 0 such that the following holds: for any E {l...., D) and any perturbation Pg with 1ö0l ro.
Fo(z . do...., ,D) Fo =p) z Lo Jool.
Proof of Lemma Denote e; for the i-th standard basis vector of RD. By assumption, m, == (0) > 0 for each i. Since Fo E C',, the map u + 0,Fu(u) is continuous at 2. Hence for each i, there exists pfant > 0 such that
[.Fo(u)] 2 2 m, whenever ||u
=1100 < rront.
If necessary, shrink goout so that the line segment { te; |t| goott lies in 1j.. Define uniform constants
Lo
min m, 0, To Isis
min Isis
r
cont 0.
Fix and dg with 160| < ru- Consider the one-dimensional slice gt()) == Fulz + te:) for [t] < To- Then gi O and (t) 0.Fo(z (z + te;). By the mean value theorem, there exists (0, 1) such that
Fo + 6)t,) Fo(2) (066) 6, 0,Fo( 2 + 05gez) d0-
Taking absolute values and using the lower bound on [-) inside the loo radius rn around have
Fo(= + 50es) - Fo(z)] 2 To lool
which is the desired inequality.
Lemma 7. Let 2; min f((x) and let x!0) denote the unique minimizer of fc (as assumed above). Then there exist constants Ro > O and Eu 0 such that:
1. The open balls (B(x()) (B) are pairwise disjoint.
2. For each i (I...., and every x E [0, 11dh ((x(0) Ro).
fi(x) > 2; +
Proof of Lemmal)
7
Since the minimizers
(a((0) li=l are pairwise distinct and finite in number, we
have
A min (z(() 20) ll, 0.
ifj
Set Ro = L. If If + and æ E ((x Ro). by the triangle inequality, have
|| - 1ا 2 [| - x 1| || z 2|2 A Ro = Ro.
so x & ((x(),Ro)) Hence the balls are pairwise disjoint, proving the first part.
For the second part, fix i and define the compact set K; : [0, 1]d B(x()) Ro). The continuity of fi implies that the minimum
m, = min fx K
is attained on K;. Because (1() & K; and (00 is the unique global minimizer on [0, 1] we have m; 25. Let E; = m. - 2; > and set
= min
> 0.
20

---

Published as a conference paper at ICLR 2026
Notation flow (dependency structure)
Meaning
(x(
Point where /a achieves minimum
B(x(0)r)
Basin region for retrieval coordinate
(Basin around 2!(i)
GG,, K;
Monotone local segment near 2x()
(In ((x(),,))
Po
The set of all candidate points.
(We only choose 2; E Po) Index partition for retrieval coordinate (i = 1,,.,,, D)
S
Attention head j
Defines response at position
+ (z(2xt) t) (1jst))
Attention score
Maximum-attention point selected by head in Po x Sj Chosen maximizers attention score (one per head) Value embedding
+ Y = (V11...,,9''
+ (x,
WLOG, suppose K₁ n Y = 0.
+ To
To C Si. indices not in (yj 42), .,....8 Coordinate system on G₁ Rewriting filG, into the coordinate system.
+ [0, 1] Gi
79 = f
+ U/
Discrete grid on [0, 1] at index t
(1)
Candidate point for subsequence l, zt E (U,))
Adversarial subsequences Z, = (I). (To))
Two subsequences almost indistinguishable by attention head. Full sequence embedding Ze Token of Wz of index t
We we
J, Iz, I3
Partition of indices: differ agree / remaining
Per-head analysis
Q A 1
Attention mass on I, (j E 11,,,3), (1..... Weighted value average on Ij
Table Flow-style dependency map of notation introduced in the proof of Theorem 2.2.
Then, for every x E K,, we have
f(x)) 2 m; zi t 2 z; + 2e0 > z1 + 60,
Thus, we have proved this lemma.
Before the proof, also provide a notation table to help with understanding in Table
Proof of Theorem (2) Given the target function under the assumptions. For any given single-layer transformer defined in the main context, our goal is to find two different sequences such that their output in the part
T a (() Wo Wv,a(())
Concat
(20)
are very close (differs by only O 1)), but their output from the target function differs by least 3e, then according to lemma we have the required parameter count for the FFN f to be at least
2(1 (t
21

---

Published as a conference paper at ICLR 2026
Notations For each head = i,...... S, define the attention weight function
Ax(x,t) (WI, (W):())))))))))
and the value mapping
(((,,) = Wv,,Po () E RR"
where y >> 0 is the softmax scaling factor, WQ,,, WK, R=xE are the query and key projections, and Wv. RnxE is the value projection for head i.
Notation of sets Without loss of generality, assume x:(() belongs to the interior of [0, 1] and the other case can be treated with the same method below. From lemma [ lemma 6 and lemma 7
have that there exists R >> and segments G; B(x (0, RR) I,..., D and L, ō0, >
satisfying the following:
Vi, Yz, y E Go we have [f((x) f(())) L|| yl12.
Vj # and Vo E B(al)) R),y E B(x(), R), we have fi(y) fi(x)) > do-
The length of G, rs vi L....., D.
For any i (1,.... D} and any perturbation 8, with |ó1| < m RY (f((())))) 2;).
Folz,..... 24 sI...., =D)
ED) 2 [
We denote by K; = (, 1,... D. and Po UP_ K2- Recall that T-s- D+1
1.
(a+1)s+1
We assume without loss of generality that k > 0 and T D + > 0, otherwise the result would be trivial.
Max weight for each head For j 1,..., s, define recursively the pairs (yi, ty) as follows:
For the first head,
(y1,t1) arg max (1 t). PEPo rES1
For j 1,
(u5, (4) = arg
max PC TES,
A5 1).
(f2.... 1)
If maximum can be obtained at multiple (y, t), then choose one of them.
Let Y (vi,.... Va). Since the Ki,..., Kn are pairwise disjoint and least one index i E (I....,)) D} such that
D, there exists
K, nnY= = Ø. Without loss of generality, we assume that we have that there exists set of (Lz...., ,,,) such that
we have |SA| 2 T s D 1
1,111,11110 = I,..., D.
$ (t te), for j = ...... D.
t are pairwise distinct.
t E Sy.j = ........ D.
Let To 1T
D 1 > and assume that To is integer. Then have [S]
(t....., , t, L. ,b]l 2 To > Without loss of generality, suppose {1, (1.2....; Tol C
Si (h,..., 12 ........ (D)
22

---

Published as a conference paper at ICLR 2026
Sequences to be considered As 11 is segment of length then it is natural to assign coordinate system 17 : [0, 1] - G1 on G1, with == fi 1 being a monotonically increasing function on [0, 1]. The monotone property is result of [f fi (v) 2 2 yl12-
We denote by Tol 3T As MM Construct the following To sets:
j-1 + To
TB M
Uj
1...,..... To
(21)
M
We have |U,|= | 0
Claim 2.1. Existence of two distinct sub- sequence There exists two subsequences 21 (1)...., 21 (Tu) and 22(1)....22(() with 2x()) E (UU) satisfying the following conditions.
To Adeet0).0eetas00.0
EE ((ee())))))(((011)))
S 37p
for i ,,...,8.
Ee ((())())))
Ee (1()2(1)))
For each = I,....,,, s, either of the following holds:
1. Ee (((2(()))) E/e (e(2)(t).)
E 1/(11
c+1 ), 1+ A+1 1278
1213
max Ees (t), *+1 Ew=r (e
Proof. We compare the orders of 1/e appearing both sides of the conditions.
First, since U.| = O(1/c) for each i, the total number of possible choices of subsequences (=(1),...,,<)) =TT)) is most O(1/cTn).
Next, to satisfy condition (1), note that both vectors involved are n--dimensional with norms bounded by 1. Thus, the discretization required achieve accuracy *+1 (33%) in the &g leads to most
distinct possibilities, since there s heads.
For condition (2), observe that
Σ (Bwrtwe 2
To max
(t),
Hence, for each the relevant interval be partitioned into most O logs
sub-intervals.
Taken across heads, this contributes at most O((- log e)* /e possibilities.
Combining the two conditions, the total number of distinct admissible cases is bounded above by
of ((++)))+((++1))
eo
Since To 2 (k 1)ns + (k + 1)s 1, we have
(
o
of lo)
(c(tt))ns((+11)))
To
Therefore, by the pigeonhole principle, there must exist two distinct subsequences 21 and (2) (TT)) satisfying all the conditions of Claim 2.1
Construction of Distinct sequences From Claim 2.1
have constructed two sub-sequences Z,, Zz satisfying the given conditions. We now consider the construction of two full input sequence Wi, W2:
For t = 1,..., To. if 21( = 22(t), then (t) = wz(() = (D), Otherwise, w1(t) =
(t),, (2(() 22))
wyy(,) Vi. = 1,...,8; 8; j 1,2.
wy(()) = ,,0) 1= 2...., j j= 1,2.
For all other w)(() (D).
23

---

Published as a conference paper at ICLR 2026
Difference of W1, W2 applied to target function Denote by I the of all indices t with
21 # (t), and 1 = [To] J, 13 = [T] I1. clear from the difference Zi. Zz that 1 Ø.
We then define the following notations for the simplicity calculation (defined for each head = 8))
Qi, Exen Ar(won ((), ().
Q2 = EtEn, (=(wz(()) (4),
V1,+ = (EEE, Am(wn (1), (w) (1,44))/1.
V221 EIEN (1) t) (w2 (t).
Q3 ExEL, (2 (t), t), which is also the same if defined on Zz.
V3,1 = EiEn (= (1), () (2) 0)/2,,, which the same if defined on Z2-
Q4.² = ExErs (w)) (4), 1), which is the same if defined on W2-
V4 (EeS,, (wn (1), ( (1),))/Q4,., which is the same if defined on W2.
As X²() maps to positive values, Vis are convex combinations of u;(). whose norm is bounded by according to the constraint section Therefore ||Va 1, 3,2,3,44
As fi 1) is monotone
[0, 1). let f maxter, t, then have
maxiss, Ji (w) (1)) fi (wi (w). maxres, Ji (w2 fi (w2(
And by construction we know that
[1( (f) (w2(1) r
(22)
M
which is the minimal distance for any two points in Uz. Then have
rL M
(i Si (w2(0))
(23)
As we have for = 2..... D
maxies; S:(wi (1)) 20))
maxtes, (w) (0)
Then following the perturbation property of Fu defined above have that the difference of output between the two sequence be least rt" that [W)))((W.) Model(W,, 2
Mr which is greater than 3c. Then e-approximation requires
W1 and We are close after attention layer For any given head i, we consider the the two cases given 2.1
Case Case can be rewritten as follows:
211
Qu.++Qa.1
Q2.+
375
QuitRaa Q2.+Raa.
11(11 +
k+1
)
1+
A-1
1277
1277
24

---

Published as a conference paper at ICLR 2026
Without loss of generality, we assume Q1 2 Q2 By calculation, we have
Q1,iVi + V3,4Va.
Q2. + Q3,4Va.
(24)
Q1, Q3
Q2,1 Q3
Q3 (V,,) V221)
Q1
+ Q1 Q3,,
(Vi, 122z).
(25)
(Q1, + (22, + Q3,i)
We have already known that Q4,: 2
Q2. (As Q4,; has the max weight of each head in it). Then
To
(Q2. (V4, V221)
(26)
(Q1, + Q4,i) + Q4,i)
(Q2, - Qv,)) V22)) -V2s)
(27)
(Q1,1 + Q1,5)
T2(Qu, - (Va, V2.i)
(28)
(Q1, + Q3,1)
The k+1 IVa.s VV2.))
(29)
1273
k+1
(30)
6To
Similarly, we also have
- -2,1)
kt1
(31)
(Q1,, + (Q2, + Q3,1)
6To
From inequality 26 and substituting equation 24 we have
Q1.i (V1, 12,i)l
k-1
e k+1
h+1
(32)
Q1. + Q3,5
6Tu
3To
2To
Therefore
Q1 (V1 V2.i)l1
(33)
Q Q4
ToQ1,
<1
(V1, - (V,
(34)
21 Q3
e k+1
(35)
2
Thus
V1, + QuiVa Q2,4 + Q1.Va.
12
(36)
Q1 + Qu,
Q2,i + Q4,1
Q2, QVa.) - -V2,)) || + Q1 + Q
Q1,
(V1, 1
(37)
(Q1, + +(,i)) + Q4,1))
k+1
h+1
+
(38)
6To
2
<e k+1
(39)
Case 2 Case can be rewritten as follows:
Qa 19..a.4
1|2 < Atl
Qh.tt0a,s
Q2,+44s,
370
Q1, + Q3.i <
Al1 Ei Ae(Uhes bw) < x1 4 Q1.
4
Q2, + Qa,i d+1 Ee
dely bo)
5 - Q4,
25

---

Published as a conference paper at ICLR 2026
Thus
11. V1. Qu,Vis
22 Q4.
112
(40)
Q1, + Q4,i
22 + Q4,5
Q0.1 (Va.) -V,
+ Q1
(V1, - 12,1)
(41)
(Q1, + (4,;) + +4,,))
Q1 + Q1
Q2.: Q1,1)(V4,1 V221)
01
++]
(Vis - V221)l
(42)
(Q2, + Q4,i)
Q1, + Q4,i
(V4, - V22a)
A+1
+
(V1, 122i
(43)
4
8+1
(44)
And it can be seen from definition that
Qu.. is the output of the i-th head of the attention layer with input sequence Q1i+Q4,i
W1. (which means Vt.VV..t4..Va.
[ (t)) (q. (t))
Q1,+Q4.:
to. Q2,+ W2. (which means Qu. Va.Va.tQuuaa
is the output of the i-th head of the attention layer with input sequence
Et=
(t)) :q (t).
Q4.+
Then for each i I,..., s, we have that
Co Wv, (W (Q
Σ a (WW) Wo Wyg s k+1
(45)
Therefore, as Wo have entries bounded by we have
|| c + Wo Concat
[(W WQ (v.nū (1)
(46)
(00 + Wo Concat &
a [(Kx,)u2(()) Wo, ((o W (v),
(47)
k+1
(48)
However, has been proven above that we need (Mddel(W)) - Model(W2)) 2 e to achieve e- approximation of the target function. According to lemma 4 the required parameter count of the FFN F is of order ((t/ Thus the parameter count required to achieve c-approximation is ((1/c)).
Remark. Tightness of Theorem 2 (2) The lower bound in Theorem 2 (2) remains essentially tight under several relaxations of the of 1-Lipschitz activations, matching upper bounds can be constructed, but this case is impractical since Heaviside activations are rarely used in practice. If parameter norms are permitted to scale as O(T1/s), the parameter count can be reduced to O(1 (11/7++)), though this scenario is likewise unrealistic in practical settings. Finally, if F is allowed up to five layers, the lower bound changes to 2(1 k/1)), which does not alter the qualitative conclusion.
block F. If F uses Heaviside activations instead
A.3 PROOF OF THBOREM (3)
Proof Sketch. The argument is based on an explicit construction. We begin with trivial attention, that the post-attention output is simply the averaged concatenation (x) x(T)) RTd The feed-forward block can then be used to compute the transformations f((())), perform the necessary comparisons, and approximate Fg. ensured by Lemmas and
and
26

---

Published as a conference paper at ICLR 2026
Having outlined the main idea, proceed to the detailed proof. As a first step, introduce several auxiliary lemmas that will be used the argument.
Lemma 8. Fix a pointwise activation a (e.g., ReLU or any activation used in this paper). Let be a 2-layer fully connected network, F2 = Im 3-layer fully connected
Fi :m + m
network, and F3 R + R a 2-layer fully connected network. Let W1, W2, W3 denote their respective (maximum) hidden widths, and W = max W2, W3}- Then there exists 5- fully connected network G m + with activation a and hidden width at most W such that
G(z) (()
for all x E ER.
Proof. Proof of Lemma activation o):
Write the three networks in affine-nonlinearity form (with a pointwise
(x) A (Bir + b2) a1
E R Fi(x) E R"
E Rm F2(u) R (u) C2 o D Equ 62) d2) C2.
Fs( Ps aQ3 9s) r3:
E In Fs(e) R.
Define 5-layer fully connected network G mh IR by stacking the hidden layers of Fi (one), F2 (two), and F3 (one), keeping their original widths:
(x) == oB bi);
u(x) =- (1h1(x) +ai
h2(x) == (Ezu(z)
ha(x) (1) +2))
v(x) == C2ha(v) +
14(x) == (Q3u(z)) +
G(x) == psh1(1) +
By construction,
G(z) oo.fealGoddDdddEdddddaaaaaaaaaad
(())(F (or
Thus G realizes the composition exactly, has hidden layers (hence 5 layers total), and its hidden widths are precisely those of the constituent hidden layers of Fi, F2, and F3.
Lemma (Approximating max with shallow ReLU network). Let f f(11.... .T) = max(21, max(xi....,x7 For any c E (0, 1), there exists fully connected ReLU net- work f with three layers (i.e., two hidden layers and output layer), whose hidden-layer layer widths
[0, [0.00
be
each most 27 [1/c], such that f c-approximates f.
Proof. Proof of Lemma Let
n n= [1/e].
For each coordinate t E [I'] and each grid index = 0, I 1, define the first hidden layer neurons by
h1 (t,
ReLU 4)
For each j = 0, ......, 1, define the second hidden layer neurons by
h2(s) = ReLU Σ
T
mi(t.))) - ReLU
Σ
T
(t.tt) -
Finally, the output of the network is given by
n-1
f(z .T)
Σ
h2(f).
27

---

Published as a conference paper at ICLR 2026
Clain. Fix (O....,, 1} and set
T
T
5 Σ
(1(1))
Σ
ReLU (x
-2) n
1=1
1=1
By definition,
h2(i) = ReLU((S)) ReLU (s) 1 - n
1) If h2(s) O, then necessarily S > 0 (since ReLU(z) > 0 iff 2 > 0), hence there exists some with
ReLU 2, - 1 0
Thus h2(i) > 0 only if 3t with 22 > j/n.
2) If there exists with H,
n,, then
5 2 ReLU at - 2)) > +.
Therefore S₃ 1, and we get
hz(s) = Ss - -sp 11 = #
Fix [0, 1)Td and let j be such that max, de E (j/n, (j + +1/m]. By construction,
(k)) = 0 for 2 1, h2(k) for S 1,
and for = j we have
T Sy 2
0 - hz(1) = ReLU(S,) - ReLU (5) - s 1 n
ReLU 81 21)
Hence
f(x) Σ hk(k) = + 2.
4 1
1
i+1
k=
k=
Since max; E (j/n, (j + 1)/n), it follows that
1 n S 1f(x) max n.. S
S e.
Therefore f e-approximates f(x) max, He on [0, 1|"
Proof. Theorem 2 (3)
We begin by fixing the embedding with positional information. Let Pø : 1j x [7] + Ref be defined by
(a(xtt) = (O...., z(t), ,,...,
where the vector xtf) occupies the t-th block of dimension d, and all other blocks zero. With the classification token Eg = 0. the attention layer reduces to a trivial aggregation, and the output (prior the feed-forward network)
+
(T)) [0,
Given target accuracy 0, construct three feed-forward networks F1, F2, F3 follows.
Step 1: Approximating the component functions.
Define
F :::400 1jdxT + DxT, Fi(+ ................................ Fo(T)) (u(1),...., u(T)),
where each u(!) E D satisfies
- - (+((1)) S for all i 1 D.
By Assumption 3 such an approximation can be implemented by a two-layer FFN with parameter count O(1/e)).
28

---

Published as a conference paper at ICLR 2026
Step 2: Approximating the minimization. Let F2g R uD),
RR be defined by
(,. u(T)
U/ = min ut);
sES:
By Lemma (which works the same for taking minimum), there exists a layer ReLU network with ((1/c) parameters that e-approximates F'z. We denote this approximation by F2.
Step 3: Approximating the outer function. Finally, F3 : RRD + R be a two-layer FFN that e-approximates Fo. with parameter count O(1/e7).
Composition. Since Fn C on a compact domain, it is Lipschitz with constant L, and the min operator is I-Lipschitz. Therefore, the composed network
P3 o P2 Fi
provides an Le-approximation of the target function, with total parameter count
0(1/c-1))
According to lemma F3 o F2 Fj can be written equivalently as a five-layer FFN.
A.4 PROOF OF COROLLARY
Proof Sketch. It is direct corollary of Theorem 2 (1) and 2 (2)
Proof. Suppose D1 < D2. Let Mg be the minimal positive integer such that H(D1, 2, d, T, Mo) c-approximates H. Then with representation ((), Fo), Theorem 2 (1) suggests that there exists positive Ca,D,.T such that
Mo 1 Cd.D,.T
(49)
E7
With representation (3 such that
Fu), Theorem (2) suggests that there exists a positive Ca.D..r
CAPATT for ( Di D2 + Mo 2
(50)
ce
3D1 +
As fi and Fo are at least C smooth, we have y S max( D2). Thus with D3 D3 I 30 T. we
have k > Then there exist c > 0 such that
Cd.D
Ca.D..T
(51)
ek
67
This leads to contradiction. Thus D D2.
29

---

Published as a conference paper at ICLR 2026
B EXPERIMENT DETAILS
B.1 DETAILS FOR EXPERIMENT 1
B.1.1 EXPERIMENTAL DETAILS FOR SECTION 6.1
Data generation. The intrinsic dimension of the synthetic task is 4. For each sequence length T E {8, 16, 32, 64, ]28] we generate 8000 training and 2000 validation examples. The inputs i..d. Gaussian samples xtt) (N 4).
Model architecture. Each input vector xxt) first mapped to R by a two-layer forwafo network with hidden dimension N and ReLU activations, ensuring a per-head embedding dimen- sion of 8. A trainable classification token co is appended, and no positional encoding is used since the task permutation invariant. The sequence is processed by single-layer multi-head attention block without residual connections or layer normalization, consistent with the theoretical setting. The output is concatenated and passed through two-layer GeLU-activated feed-forward network with hidden dimension N, yielding the final scalar prediction. The fixed hidden size ensures com- parability of parameter counts across different h.
Training protocol. Each configuration (h, T) is trained separately under multiple random seeds. To reduce the effect of optimization variance, we report the minimal validation error achieved across seeds. This choice isolates expressivity limitations the architecture from randomness in training dynamics.
Evaluation metric. We adopt the normalized mean squared error (NMSE), defined mean squared error divided by the variance of the targets. As T increases, maxima of Gaussian sam- ples concentrate, shrinking target variance and making trivial predictors appear competitive under raw MSE. (An intuition is that suppose the target output be Yr = maxi Te with input tokens T, NO, , 1) independently, then Var(Yr) decreases as T increases, because Yr concentrates more tightly around its growing mean. Normalization by variance corrects this effect and ensures compa- rability across lengths. NMSE also equivalent to R², where R is the standard coefficient of determination.
Variance across seeds. While mean performance across seeds is also informative, reporting the minimal validation NMSE highlights the best achievable accuracy for a given architecture. This em- phasizes limitations due model capacity rather than training noise. Tables showing seed variance are included for completeness (Table 2}.
Remark. When h > D = 4, also observe that the validation NMSE first decreases rapidly and then increases slowly as T grows. For shorter sequences, the model with enough heads can either capture the pattern through attention (Theorem (1) or rely
memorization-on based strategy with the feed-forward network (Theorem (3) Both approaches generalize reasonably well, but the memorization-basse
does less effectively. For longer sequences, memorization becomes infeasible and the model relies on attention, which generalizes better; however, longer sequences may also be more sensitive to parameterization, and the observed curve likely reflects tradeoff between these effects. See Figure 3Ji Appendix.
B.1.2 FIGURES AND TABLES FOR SYNTHETIC EXPERIMENT 6.
Heads
1=16
T=M4
7.01
10
+ 500
L00
10
+
+000
1.00
+0.30
1.4
+8.3
1.4
120
7.31 10
+ 4.7
8.41 x
7.97
10
9.42 x
10
+ 6.42
1.31
+116
1.47 x
1.21
2290
1.5
10-144
10-4
10-6
5.84
2.6
10
4.19
8.30
3.99 10
Table 2: Error bar for synthetic dataset. NMSE(Mcan Standard Deviation) for different sequence lengths T and number of heads.
30

---

Published as a conference paper at ICLR 2026
1e -
8
Validation NMSE (Min)
6
4
2
Number of Heads
-
0
20 40
60
80
100
120
Sequence Length
Figure 3: zoom in plot of Figurel which shows that when the number of head is enough, the loss first decreases and then increases, as explained in the remark B.I.I
H == 1
H = 2
Hidden Dimension
4.0
Log Hidden Dimension
3.5
3.5
T-8
T=16
T-32
3.0
T-128
2.5
2.5
-1,6
-1.4
-1.2
-2.2
-2.1
-2.0
Log Mín Validation NMSE
Log Min Validation NMSE
Figure 4: Additional plot of lb for H = and H = 2.
Heads
T=8
T=16
T=32
T=64
T=128
1.75 X 10-2
1.98
10
2.06
10-2
2.54 x 10-2
3.03
10-2
7.17 x 10-3
7.39 x 10 3
7.82
10-3
8.57 x 10-3
1.02
10 10-2
2.11 x 10-4
2.17 x 10-4
2.73
10-4
3.71 x 10-4
4.77 x 10
4
1.32 x 10 10-6 6
5.59 X 10 7
3.40 x 10-7
3.46 x 10-7
5.70 X 10
5
2.19 x 10-6
4.33 X 10-7
3.22 X 10-7
2.73 x 10-7
2.66 x 10-7
Table 3: Validation NMSE under fixed total embedding dimension E nh 32.
Heads
T=8
T=16
T=32
T=64
T=128
1.38 x 10
1.63 x 10
1.84 x 10
2.17 x 10-2
2.31 X 10
2
1.09 x 10 3
7.08
10
7.24
10
7.76 x 10
1.11
10
3
4.18
10-7
1.72
10-7
1.17
10-7
3.58 x 10-7
2.11
10-7
5.56 x 10 10-7
1.22 x 10
6.89
1.85 x 10-7
3.48
x 10
Table 4: Approximation error for the D = 3 retrieval task under fixed total embedding dimension E = nh.
31

---

Published as a conference paper at ICLR 2026
Heads
T=8
T=16
T=32
T=64
T=128
2.12 x 10-1
1.85 x 10
2.22 x 10
3.13 x 10
4.28 x 10
2
7.22 x 10-6
2.69 x 10-6
3.50 x 10-6
3.07 x 10-6
3.83 x 10-6
3
8.16 10-6
1.83 10-6
1.73 10
3.86 x 10-6
3.50
10-6
4
3.68 x 10-6
1.87
10-6
2.60
10-6
4.32 x 10-6
5.98
10-6
6.15 x 10-6 3.02 x 10-6
3.31 x 10-6
3.78 X 10-6 5.34 x 10-6
Table 5: Two-layer transformer on the synthetic task (D 4, NoPE, NoLN, fixed total embedding dimension E nh 32).
B.2 MS MARCO TEXT RETRIEVAL
B.2.1 EXPERIMENT DETAILS FOR MS MARCO (TEXT RETRIEVAL) EXPERIMENT
Dataset construction. We construct retrieval-style datasets from the MS MARCO passage rank- ing collection (Bajaj et al. 2016). Since the original dataset associates each query with only a few candidate passages, we enlarge the candidate set by mining hard negatives. Specifically, BM25 [Robertson & Zaragoza 2009) is used to mine local negatives and FAISS (Johnson et al. 2019) sim- ilarity search to retrieve global negatives, reducing redundancy across queries. For each query. the sequence length 7 defined as the total number of candidates (one positive and T 1 negatives), with T {8, 16, 32, 64). We build datasets containing 28,000 training queries and 2,000 validation queries for each T.
Model and training setup. We evaluate a two-layer Transformer encoder with per-head embed- ding dimension fixed at 32, while varying the number of heads across {1, 2, 4, 6, 8, 10, 12, 14, 16). Tokenization and input embeddings follow the BERT tokenizer and frozen BERT word, position, and segment embeddings (Devlin al. 2019), projected the model hidden size h heads 32. Only the projection and Transformer layers are trained. We report training top-1 accuracy, focusing on training performance since MS MARCO with BM25-mined negatives is particularly challenging for validation, and the difference can be seen in training metrics. Training MRR is also reported in Fig S
with similar trend as training accuracy.
B.2.2 FIGURES AND TABLES FOR EXPERIMENT
Heads
T=8
T=16
T=32
T=64
0.597 0.003
0.450 + 0.005
0.303
0.003
0.154 0.002
0.771 0.003
0.647
0.003
0.486
0.286 0.0002
4
0.956 0.002
0.900
+ 0.002
0.793
0.580 0.004
0.992 + 0.000
0.001
0.937=00000
0.814
0.998 0.000
0.995 0.000
0.983 0.001
0.932 - 0.002
12
1.000 + 0.000
0.999
0.000
0.998 + 0.000
0.991 0.001
16
1.000 + 0.000
1.000 + 0.000
0.999 + 0.000
0.996 + 0.000
Table 6: Error bar for MS Marco dataset. Accuracy (Mean = Standard Deviation) for different sequence lengths T and number of heads.
B.3 CIFAR-10 IMAGE CLASSIFICATION
B.3.1 DATASET CONSTRUCTION
We create image classification datasets from the CIFAR-10 dataset using a padded preprocessing
approach. The original CIFAR-10 images have dimensions of 32 32 pixels. To generate datasets with larger image sizes, we apply padding to achieve sizes the set {32, 48, 64, 96, 128]. The original image randomly positioned within the enlarged frame, with padding filled using the colors of the border pixels. An illustration is provided in Figure 6 By apply this padding method we are creating tasks with increasing difficulty. The background is enlarged, making models need more
32

---

Published as a conference paper at ICLR 2026
Heads
T=8
T=16
T=32
T=64
1
0.5107 + 0.0069
0.3917
0.2542 + 0.0049
0.1257 + 0.0057
2
0.5221 + 0.0102
0.4205
0.2712
0.0067
0.1369
+ 0.0038
4
0.5076 0.0112
0.4048 0.0070
0.2547
0.0093
0.1139 + 0.0061
6
0.5153 0.0112
0.3865
+ 0.0103
0.2397 + 0.0098
0.1018 + 0.0057
8
0.5058 + 0.0082
0.3801 = 0.0084
0.2308
0.0068
0.0983 + 0.0050
12
0.5184 + 0.0073
0.3721 = 0.0107
0.2219 + 0.0091
0.0902 + 0.0054
16
0.5021 + 0.0123
0.2816 0.0418
0.2170 0.0075
0.0878 + 0.0058
Table MS Marco Validation Accuracy (Mean Standard Deviation) for different sequence lengths 7' and number of heads.
1.0
Train MRR (Max)
0.8
0.6
Sequence Length (T)
T=8
- T=16
0.4
- T=32
T=64
2
4
6
8
12
16
Number of Heads
Figure 5: Plot of training mrr for MS MARCO dataset.
effort to learn how to extract useful information. The random placement make sure the padded outside aera cannot be simply ignored by position encodings.
Each image divided into non-overlapping patches of size 8 pixels, resulting sequence of patches for each image. For each image size, the sequence length T is defined as the total number of patches plus one additional class token, with = {17,38 37, 65, 145, 257). We adopt the standard CIFAR-10 data splits, which include 50, 000 training images and 10, 000 test images across 10 classes
B.3.2 MODEL TRAINING SETUP
We evaluate a Vision Transformer (ViT) with four layers and a per-head embedding dimension of 16, while varying the number of attention heads in different configurations. Fach image patch embedded through lincar projection, and positional embeddings are added along with a learnable class token. No global convolutional embedding used.
Input processing follows the standard ViT procedure, including patch embedding of size 8 x 8, po- sitional encoding, and aggregation of the class token for final classification. The model is trained using the AdamW optimizer with cosine annealing learning rate scheduling. Standard architectural techniques, such as layer normalization, residual connections, and dropout, are employed for regu- larization.
33

---

Published as a conference paper at ICLR 2026
Figure 6: Examples of the padded images from the dataset.
Seq=145
Seq=257
Seq=577
Seg
4.78
#4
1.
10
4.37
2 4.70
10
4.30 x
5.00
4.08 x 10
6.0
10
4.00
10
7.20
10
5.97
10
532
33,0
10
5.34
1.10
10
5.15 x
4.70
10
5.08
47.00
10
7.03
6.8
6.58
7.20
9.51
x 10a
1.50
10
9.26
1
10
9.14
6.20
10
9.07
10
1.02
10"
9.02
10
1.00
10P
9.81
10
10
9.73
+ 2.40
10
9.87
9.65
3.20
9.67
47.30
10-1
9.92
10
33.00
1
9.89
1
10
9.86
+
10
9.36
1.00
10
9.56 x
x
13
9.94
16
+-6.0
10
933
10 +6600
10
9.11
10
1.10
10
9.00 10
2,00
10
9.01 10 11000-222000 +2..0
10
-0
10-2
10
+2.00
10-1
9.97
9.96
2.000
9.85
9.96
1.40
0.96
20
9.98
16a
1.00
10
9.97 x 10
5.500
10-2
9.97
x 10
8.00
10
9.98 x
8.00
10-3
9.99 x 10
7.00
10
10
10
9,88
2.00
10-2
9.88
10
4.00
10
9.99
5.00 10-2
Table 8: Error bar for Image task. Accuracy (Mean Standard Deviation) for different sequence
lengths and number of heads.
C LARGE LANGUAGE MODEL USAGE
Large language models were used only for linguistic refinement (e.g., polishing sentences and check- ing grammar). The ideas, theoretical results, experimental design, and analyses presented in this paper were entirely developed by the authors without assistance from large language models.
34

---

Published as a conference paper at ICLR 2026
Heads
T=65
T=145
T=257
T=577
T=1025
50.50 + 0.44
45.85
43.53 + 0.58
41.43 0.97
0.75
60.01 + 0.57
53.12 0.20
4.80
0.83
4
67.98 + 0.43
61.42 0.70
$7.19
69.63 = 0.55
0.53
0.95
56.18 57.59
10
69.70 = 0.24
6
62.45 0.49
57.21
5.
11
69.84 + 0.36
65.26 0.43
61.97 0.29
56.95 0.63
53.77 + 2.91
12
69.66 + 0.39
6
62.63 0.36
56.57 = 0.66
53.17 1.06
13
69.72 + 0.18
61.66 + 0.58
0.58
54.81 2.23
52.90 + 1.45
14
69.49 + 0.48
0.425
6
=
50.42 + 2.48
16
69.69 = 0.33
6
59.29
49.68 = 1.39
48.51 + 2.59
20
67.99
0.35
6
55.25 2.12
48.65 = 0.85
46.89 + 1.07
24
65.12 = 0.68
56.80 = 1.09
52.14 + 0.56
48.74 0.60
48.39 + 0.76
Table 9: Error bar for Image task. Validation Accuracy (Mean = Standard Deviation) for different sequence lengths and number of heads.
Table 10: Hyperparameter settings of popular transformer models. Only d (embedding dimension),
(number of layers), and (number attention heads) are shown for brevity.
Model
d
Year
8
Attention is all you need
512
6
2017
8
Gemma 2B
2,048
18
2024
12
GPT
768
12
2018
16
BERT-Large
1,024
24
2019
16
ViT-Huge
1,280
32
2021
16
Gemma 7B
3,072
28
2024
28
Turing-NI.G
4,256
2020
32
LLaMA-7B
4,096
32
2023
32
Baichuan 2-7B
4,096
32
2023
32
Mistral 7B
4,096
32
2023
32
Yi-6B
4,096
32
2023
32
LLaMA 3-8B
4,096
32
2024
32
Mixtral 8x7B
4,096
32
2024
40
LLaMA-13B
5,120
40
2023
40
Baichuan 2-13B
5,120
40
2023
56
Yi-34B
7.168
60
2023
64
LLaMA-65B
8,192
80
2023
64
1lama-2-70B
8.192
80
2023
64
LLaMA 3-70B
8,192
80
2024
96
GPT-3
12,288
96
2020
96
Jurassic-1
13,824
76
2021
128
MT-NLG
20,480
105
2021
128
LaMDA
8,192
64
2022
128
LLaMA 3.1-B05B
16,384
126
2024
128
DeepSeek- V2
5,120
60
2024
35


Reviews:
[{"review_id": "92cmpYLRBj", "reviewer": "Reviewer_FaEJ", "summary": "This paper provides a theoretical analysis of the expressivity and approximation efficiency of single-layer Transformer networks, focusing specifically on the role of the number of attention heads, $h$. It introduces a function class, the generalized $D$-retrieval task, which models problems that combine multiple sequence elements determined by minimizing different component functions. The authors establish the generality of this task in Theorem 1 by proving that the class is dense in all continuous sequence-to-vector mappings. \n\nThe core of the paper is then captured in Theorem 2, which provides a set of sharp approximation bounds:\n* Positive Result (Thm 2.1): A transformer with $h=D$ heads can efficiently ($\\epsilon$-approximation in $L_\\infty$ norm) approximate $D$-retrieval tasks with constant per-head embedding dimension ($n=2$) and MLP size $M$ scaling independently of sequence length $T$.\n* Negative Result (Thm 2.2): If the number of heads $h$ is insufficient ($h<D$), and the sequence length $T$ is much larger than the per-head embedding dimension $n$, the MLP size $M$ must grow exponentially in $T$.\n* Positive Result (Thm 2.3): If the embedding dimension $n$ is scales linearly with the $T$, the full sequence can be effectively \"memorized\" in the attention output, and the task can be completed efficiently without an exponentially large FFN.\n\nThe theoretical findings are validated with experiments. The synthetic task clearly demonstrates a sharp phase transition in learning difficulty at $h=D=4$. Experiments on simple image and text retrieval tasks show that performance increases as the number of heads (with fixed parameters per head) increases.", "questions": "Why did the synthetic experiments focus exclusively on the case where $D=4$ (4-retrieval task)? It would be valuable to know how the critical performance threshold changes as a function of $D$. Did preliminary experiments suggest the sharpness of the transition is invariant to $D$?\n\nThe restriction to single-layer transformers is a noted limitation. While a rigorous multi-layer theory is beyond the scope of this paper, it would be highly insightful to share a conjecture or intuition about what is expected in that regime. Is the intrinsic dimension D of a task expected to be distributed across layers, or does the bottleneck effect remain a strong limiting factor at every layer? Insights based on the real-world experiments, which use deeper architectures, would be helpful to include in the discussion.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The work provides a fundamental theoretical basis that cleanly captures the trade-offs on the expressive power of attention heads. The model is deliberately structured to reveal that when $h<D$, heads are forced to encode multiple retrieval roles simultaneously, creating an information bottleneck that MLP can only overcome with exponential complexity in $T$. Taken together, the bounds demonstrate sharp trade-offs between difference scaling regimes depending on head and embedding dimension capacity. The use of a simple, illustrative toy example greatly aids in explaining the bounds.\n\nA key achievement is how the work uses the concept of intrinsic dimension $D$ to distinguish performance at $h$ vs. $h+1$ heads. This analysis moves beyond earlier approximation results that either focused on universal capacity (where one large head is enough) or analyzed rank restrictions in isolation. The demonstration that $h<D$ leads to a $T$-dependent performance scaling is a useful distinction.\n\nThe synthetic experiments offer strong validation. They demonstrate a real threshold in model performance exactly at the theoretical intrinsic dimension D=4. This empirical finding strongly suggests that the theoretical representational gap translates directly into a learning difficulty gap in practice.", "weaknesses": "The theoretical section lists model assumptions (Assumption 1) and target function constraints (Assumption 2) that appear to apply to the entire set of theorems (Thm 2.1, 2.2, 2.3). While it is appreciated that the bounds are in a consistent setting, it would improve the presentation and clarity of the paper if the authors would specify the minimal subset of assumptions required for each individual bound. For example, does Theorem 2.1 strictly require the Hessian constraints of Assumption 2.3? This small clarification would make the theoretical scope of each result easier to grasp.\n\nThe derivation of the lower bound in Theorem 2.2 relies on Lemma 4, which bounds all weights in the MLP by 1 (Assumption 1.3). This weight constraint simplifies the analysis, but is (to my knowledge) nonstandard in theoretical work or practical models. It would be valuable to discuss whether the lower bound can be maintained by instead considering a joint bound on the network's width and its weight norm in the lower bound, instead of a strict magnitude bound on every entry. This would align the result more closely with well-known complexity bounds for feed-forward networks (e.g., those from the line of work including [Yehudai-Shamir '19](https://arxiv.org/abs/1904.00687)). \n\nWhile the proofs appear largely correct upon detailed inspection, the proof of Theorem 2.2 is dense and hard to follow, involving numerous variables and recursive definitions that must be kept in the reader's head. I recommend the authors consider adding visual aids or diagrams that illustrate the construction of the adversarial sequences $W_1$ and $W_2$ and the partitioning of weights across the sequence indices. More descriptive notation would also help.\n\nThe experiments on real tasks show that increasing the number of heads $h$ (with a constant embedding dimension $n$ per head) helps performance. However, this is a key confounder, as the total embedding dimension $E=nh$ increases linearly with $h$, and thus the raw parameter count of the attention layers increases. Given the theoretical claim is about the specialization enabled by multiple heads, is it possible to design an experiment to disentangle the effect of the number of heads $h$ from the total parameter count?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "b6Dbv4cwdt", "reviewer": "Reviewer_SJw1", "summary": "This paper studies the task of approximating sequence-to-vector mappings with a particular structure, where the output is a function of $D$ summary statistics of the input sequence with length $T$. The authors prove that with a number of heads $h \\geq D$, a Transformer with a single attention layer can efficiently approximate this class of mappings. On the other hand, with $h < D$, and an embedding dimension $n \\ll T$, approximating up to $\\epsilon$ error requires $(1/\\epsilon)^{\\Omega(T)}$ parameters. When the embedding dimension is $n \\geq Td$ where $d$ is the input token dimension, the Transformer can efficiently approximate these mappings by memorizing the input sequence.", "questions": "* There seems to be a sudden transition in between cases 1 and 2 of Theorem 2, i.e. even with $h = D - 1$ one has a lower bound $M \\geq 1/\\epsilon^{\\Omega(T/(nD))}$, while with $h = D$ we suddenly get $M \\leq 1/\\epsilon^{\\gamma}$. Is there an intuition on this sudden change?\n\n* Similarly, what is the reason for the difference between $M \\leq 1/\\epsilon^{\\gamma}$ for Case 1 and $M \\leq 1/\\epsilon^{1 + \\gamma}$ for Case 3 in Theorem 2?\n\n* A relevant work that could be discussed is [1]. The authors there consider a sparse sequence-to-sequence task, and in one of their results they separate the approximation power of a single-head and multi-head attention in Transformers. The task studied in [1] could also be of interest here, as it allows to move from sequence-to-vector mappings to sequence-to-sequence ones where the length of the output sequence can grow with the input.\n\n* It would be nice if $C_{d,D,T}$ could be made explicit, at least in terms of $T$, since long inputs are of particular interest here.\n\n\n[1] A. Mousavi-Hosseini et al. \"When Do Transformers Outperform Feedforward and Recurrent Networks? A Statistical Perspective\". NeurIPS 2025.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper presents an interesting setting where one can clearly prove the benefits of multiple heads for certain long-context tasks. It shows a nice tradeoff where one either needs to use as many heads as the intrinsic dimension of the mapping, or needs to use an embedding dimension that grows with the sequence length; in practical settings the former is always preferred for long-context problems. The authors also demonstrate the relevance of their introduced task by showing that it is a dense subset of the class of continuous sequence-to-vector mappings.", "weaknesses": "* A key weakness of the results is that they only apply to a single-layer Transformer, i.e. a Transformer using only one layer of attention. It is unclear how the lower bounds would change if one were to use multiple layers with fewer heads. \n\n* It would be nice if the authors could provide some intuition on the necessity of Assumption 2. Some of the items in this assumption seem a bit specific, e.g. requiring non-zero coordinates in the gradient of $F\\_0$.\n\n* It would also be nice if the authors could include a brief intuition of the proof strategy, especially the strategy for proving the lower bound, in the main text.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "R33NaH6rP4", "reviewer": "Reviewer_YCdB", "summary": "This paper studies how the approximation properties of single-layer transformers on sequence-to-vector tasks vary with the number of attention heads. To investigate this question, they first introduce a new 'generalized D-retrieval' task, which they prove to be dense in the space of continuous sequence-to-vector functions. They then theoretically find that the number of parameters necessary to approximate their D-retrieval task is small when the number of attention heads $\\geq D$ and scales at least with $O\\left(1 / \\epsilon^{c T}\\right)$ where $T$ is the sequence length if the number of attention heads $< D$.\nThey validate their theoretical findings empirically on toy as well as more realistic datasets.", "questions": "1. In line 190, why do you make the assumption $\\left|S_i\\right| \\geq \\frac{1}{4} T$? Where is this assumption needed?\n2. Should the input space of $F$ in line 195 depend on $T$? If yes, could you quickly explain this.\n3. Could you expand a bit on Assumption 2? Why or why not do you think they are reasonable assumptions?\n4. Could you explain the following sentence in slightly more detail:\n> As T increases, the number of relevant elements to distinguish grows linearly with T , yet they are compressed into an n-dimensional vector whose size does not scale with T.\n5. Could you explain the following sentence in more detail:\n> NMSE, equivalent to $1−R^2$, corrects for the variance shrinkage of maxima as T grows, thus enabling fair comparison across lengths.\n6. Could you expand on why you used the minimal validation NMSE achieved across multiple random seeds on the synthetic tasks, while using the training accuracy on the real datasets?\n7. Do you have further thoughts on how one might be able to verify that the experiments on real datasets indeed identify the 'correct' intrinsic dimension of the task? Might it be useful to train multiple slightly different transformers (e.g., different numbers of layers) on the same task to check whether the discovered intrinsic dimension is constant?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The story of the paper was clear and the paper was well-organized. The introduction  of the generalized D-retrieval task and the connection to the necessary number of heads of a single-layer transformer to efficiently approximate it is original. In particular, the lower bound on the number of parameters necessary to achieve a certain accuracy is interesting and useful. The experiments were well designed and presented.", "weaknesses": "The presentation of the experimental results could be improved. In particular, it is not clear to me why they used the minimal validation NMSE achieved across multiple random seeds on the synthetic tasks, while using the training accuracy on the real datasets and the analysis of the results on the real datasets was kept very short, making it hard to evaluate the applicability of these results to real-world problems.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XoJ0rAAkYB", "reviewer": "Reviewer_vXWe", "summary": "The paper shows that target functions that can be concisely expressed in a given form can be efficiently approximated with transformers with a sufficient amount of heads, and inefficiently so when the number of heads is lower than a given threshold.", "questions": "Minor comments:\n\n1. On line 182 you use the term “intrinsic dimension” for the first time before it is defined, it feels a bit handwavy at this point. Perhaps you can explain what it means.\n2. There seems to be a mix in the notation in Figure 1 and the experiments section: both $h$ and $H$ seem to refer to the number of heads in the model.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "1. The paper is generally well-written and explained. Intuition is provided after most mathematical definitions and results which makes it easy to follow the reasoning and the motivation.\n2. The introduction of the class of generalized D-retrieval tasks and proving that they are dense in the space of continuous functions is quite neat. It does remind me of the Kolmogorov-Arnold representation theorem, although there are a few details that might be of concern (see Weaknesses).\n3. The experiments demonstrate the predictions from theory indeed do appear in practice.", "weaknesses": "1. Theorem 1 is purely a density result. It says nothing as to how the intrinsic dimensionality $D$ and the smoothness of $f$ depend on the target function $F$.\n2. This is a problem because the authors then go on to study how $D$ affects the required number of heads for a transformer to approximate $F$ well. However, if the link between the $F$ and $D$ is not established, it is difficult to be convinced whether the conclusions of the paper actually say anything of practical value. For instance, there could be trivial constructions that require a very large $D$ and $f$ with very high Lipschitz constant. Then all that the inner functions $f_i$ would do is to embed the sequence into a $D$-long vector and then let the MLP do the heavy lifting. But then that MLP might have to be extremely large. Overall, this interplay does not seem to be taken good care of.\n3. The experiments on real datasets do not study the interplay between the number of layers and the number of heads. How were the depths used (two and four layers) chosen? Additional experiments with varying depth possibly depth vs number of heads plot could be nice to have.\n4. On line 131 you say that “we omit layer normalization, noting that its removal does not harm the approximation capacity of the model”. However, layer norm (depending on where it is applied) might limit the ability of softmax to be “picky” and to approximate hardmax. Similarly, on line 151 you say that $\\beta$ can be chosen “arbitrarily large in order to make the softmax attention mechanism approximate a hardmax.” However, for both numerical stability and learnability, we do not want arbitrarily large activations in models. I understand that these are necessary assumptions for the theory to work in the limit but I think you should be more upfront with that they are significant and an important departure from how models work in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)