Geometry of Lightning Self-Attention: Identifiability and Dimension
Authors: Nathan W. Henry
Giovanni Luca Marchetti
Kathlén Kohn
Venue: ICLR 2025 Poster (2025)
Keywords: Lightning Self-Attention
Neuromanifolds
Algebraic Geometry
Primary area: learning theory
TL;DR: We study neuromanifolds of lightning self-attention networks, discussing, in particular, dimension and identifiability.

Abstract: We consider function spaces defined by self-attention networks without normalization, and theoretically analyze their geometry. Since these networks are polynomial, we rely on tools from algebraic geometry. In particular, we study the identifiability of deep attention by providing a description of the generic fibers of the parametrization for an arbitrary number of layers and, as a consequence, compute the dimension of the function space. Additionally, for a single-layer model, we characterize the singular and boundary points. Finally, we formulate a conjectural extension of our results to normalized self-attention networks, prove it for a single layer, and numerically verify it in the deep case.

Published as a conference paper at ICLR 2025
GEOMETRY OF LIGHTNING SELF-ATTENTION: IDENTIFIABILITY AND DIMENSION
Nathan W. Henry
Giovanni Luca Marchetti
University of Toronto
Royal Institute of Technology (KTH)
nathan. hanra@@aaaammmma utoronto.ca
g glma@kth. meaaktt.ss.
Kathlén Kohn
Royal Institute of Technology (KTH)
kathlenekth.se
ABSTRACT
We consider function spaces defined by self-attention networks without normal- ization, and theoretically analyze their geometry. Since these networks are poly- nomial, we rely tools from algebraic geometry. In particular, study the identifiability of deep attention by providing a description of the generic fibers the parametrization for arbitrary number of layers and, consequence, com- pute the dimension of the function space. Additionally, for a single-layer model, we characterize the singular and boundary points. Finally, we formulate conjec- tural extension of our results to normalized self-attention networks, prove it for a single layer, and numerically verify in the deep case.
Figure 1: A slice of the space of lightning self-attention mechanisms.
INTRODUCTION AND RELATED WORK
The self-attention mechanism is the cornerstone the Transformer - a modern machine learning architecture that is nowadays popular in vast variety of domains, ranging from natural language processing (Vaswani et al., 2017), to vision (Dosovitskiy et al. 2020). to sound (Huang al, 2018). In all of these domains, self-attention mechanisms have showcased outstanding performance due to their ability to model long-rang dependencies within data sequences. Lightning self-attention mechanisms (Schlag et al.] 202T) standard variants where, differently from the original pro- posal, the attention weights are left un-normalized. As a result, the computational complexity of a forward pass is linear with respect to the sequence length, substantially improving on the quadratic complexity of the original model.
Despite their effectiveness, the theoretical understanding of self-attention mechanisms is superficial, and many aspects have yet to be clarified. In particular, understanding the geometry of function
"Equal contribution.

---

Published as a conference paper at ICLR 2025
spaces defined by neural networks - typicalll referred to neuromanifolds (Marchetti al. 2025
Kohn 2024 Calin 2020) - a fundamental challenge due to its intimate connection to several ma- chine learning aspects, such as sample complexity and expressivity. Morcover, since neural networks learn by following a gradient flow over the neuromanifold, the geometry of the latter controls several aspects of the training dynamies (Trager al. 2019). While neuromanifolds
well-understood
for several architectures, such as fully-connected (Kileel et al. 2019 Kubjas et al., (2024) and con-
volutional networks (Kohn al. 2022 2023 Shahverdi et al. 2024). they have not been considered
for self- attention mechanisms,
In this work, we study the geometry of neuromanifolds associated to lightning self attention mech- anisms. These models are of algebraic nature, since they are tri-linear in their weights and cubical in the input. This enables us to analyze neuromanifolds via ideas and tools from algebraic ge- ometry a rich field concerned with spaces defined by polynomial equations. In particular, it is possible to compute geometric quantities such as the dimension of the neuromanifold. The latter measure of expressivity of the underlying model. More concretely, it is intimately linked with sample complexity. According to the Fundamental Theorem of Learning, the dimension controls, linearly, the sample complexity of learnability (Shalev-Shward is typically formulated for (binary) classifier:] and the notion of dimension is a discrete one the Sopnik-Chervonenie (VC) dimension, specifically. In the continuous setting. the dimension of the neuromanifold is the natural analogue of the combinatorial VC dimension, and controls the sample complexity of learnability. An expression for sample complexity can be used both to select the ap- propriate crcht//architectuee given an available dataset, and to collect appropriate amounts of data train a given model. This is especially important for attention-based models, that are nowadays popular in several domains, and are trained at extremely scales.
& Ben-David 2014). This theory
The dimension of the neuromanifold is related the dual question of identifiability (Grigsby et al. 2023 -Sona-Pellissii et al. 20233] Fefferman et al. 1994) a problem concerned with characterizing the parameters corresponding to the same function. Geometrically, such parameters define fibers of the parametrization of the neuromanifold, and their (generic) dimension measures the difference be- tween the dimension of the neuromanifold and the number of parameters. Therefore, characterizing fibers leads to estimate of sample complexity which is more precise than the common practice of counting parameters. Moreover, understanding the fibers can be interesting beyond their relation to the dimension, since they control aspects of the training dynamics. Indeed, fibers induce invariances of the loss function which are iata--ndependeet, meaning that for any dataset, the loss will be con- stant for parameters within the same fiber. This gives rise to the phenomenon of flatness of the loss landscape (Zhao et al., 20225), where minima are not isolated but instead belong to a continuous set. Even further, it understood that the symmetries of the loss landscape control training dynamics as gradient directions must be orthogonal to fibers of the loss, which also induces a constraint on the Hessian (Kunin et al., 2020). This has recently been exploited to design optimizers that 'teleport' along fibers (Zhao et al. 2022a), improving learning efficiency.
1.1 SUMMARY OF CONTRIBUTIONS
Our core contribution is description of the (generic) fibers of the parametrization of lightning self-attention networks. As a consequence, the expression for the dimension of the neuromanifold follows immediately. More specifically, results are summarized as follows.
For a single layer of lightning self-attention (Section 3.1), we describe all the fibers, and addition- ally study various aspects of the geometry of the neuromanifold. Specifically, we prove that it is Euclidean closed and compute its singular and boundary points.
For deep lightning self-attention network (Section 3.3 we compute the generic fibers. Our proof involves a reparametrization which can be interpreted as introducing 'virtual weights' - and a subtle induction argument based a closed-form algebraic expression for the network w.r.t. the new parameters. Assuming the network has a bottleneck architecture, we derive a formula for the dimension of the neuromanifold. In particular, the dimension is strictly lower than the number of parameters, with redundancies arising from a scaling symmetry. from inter-layer symmetries, and from the rank constraint on the attention weights.
In this context, neuromanifolds are referred to "hypothesis spaces', and are usually considered in combinatorial version.
2

---

Published as a conference paper at ICLR 2025
Lastly, study traditional self-attention by re-introduuing the softmax normalization (Section 3.4 We prove that the parametrization single layer is generically one-to-one, and state a conjecture verified via numerical experiments for the generic fibers of deep self-attention networks.
2 LIGHTNING SELF-ATTENTION
Fix positive integers d, d', a, E N and matrices Q, K Raxd V E Rd'' xd' The latter are deemed query weights, key weights, and value weights respectively self-attention mechanism a map parametrized by (Q, K, V) sending sequences of length of vectors in Rd to sequences of vec- tors in Rd Here consider the variant of self-attention mechanisms deemed lightning. which is computationally efficient and fully algebraic.
Definition 1. The lightning self-attention mechanism associated to the weights W (Q, K, V) is the map:
PW: RRdxt
Re'xt xt
(xr (xi)
Σ ISisc
2) (K (KTQ)ar 2; Voj
(1)
isist
Intuitively, every component 24 of the input corresponds token and 'attends' bilinearly every other Zj. producing a scalar weight zj K Qu, E R. These weights are used to aggregate the values Væ, E Rdd,, obtaining the corresponding component of the output. The map defined by Equation is tri-linear in (Q, K, V) and homogeneous cubical in It is often convenient to write Equation in matrix form: if = (#t)isist is interpreted as d x matrix, then:
Fw (X) VXX K QX..
Moreover, will often simplify the parametrization by introducing the attention matrix:
A -=K Q.
The latter will always be interpreted as a bilinear form x Ay.
Lightning attention mechanisms variants of the traditional ones (Vaswani et al. 2017). where the attention weights x Ax; are normalized to a probability distribution across j see Section 3.4 for further details. The major practical advantage of the lightning variant is its computational efhciency with respect to the sequence length. Specifically, Equation can be computed O(t) time, while traditional self-attention mechanisms require (1²²) time due normalization. The improvement in efficiency motivates the term 'lightning'.
Self-attention mechanisms be stacked in order to obtain deep network architecture. To this end, fix positive integers i, d = (do...., d)) (K1,... Ki), V (Vi...,, . Vi), with Qu K; Ra, V44 E gd xd:-1, Definition 2. A deep self-attention network associated to the weights W = (Q, K, V) is the map:
(oo....d)).. a = (a1,.... a1), and weights Q = (Q1,..., Q.).
pw: : do gdixt
given by the composition yw Pww Pw.
Again, deep self-attention network is homogeneous of degree 3 in æ. Based on this, denote by Symy: dRod xt) the vector space of homogeneous polynomial functions from Rd to Rd of degree 3' in all the output co-ordinates.
Definition 3. The neuromanifold of deep self-attention network the image of the parametrization map W + PWw
Ma. = { fow W E ((0.ttd))) Symg: (Rdaxe, Rdixt)
(5)
The neuromanifold is semi-algebraic set by the Tarski-Seidenherg Theorem, meaning that it can be defined by finite number of polynomial equalities and inequalities in Sym, (Rdo xf Rdxx).
An alternative standard notation for Q, K, V is Wo, WK, Wy. Our choice is motivated by better read- ability.
3

---

Published as a conference paper at ICLR 2025
3 RESULTS
In this section, we study the neuromanifold of lightning attention networks, focusing on its parametrization and its dimension. Our core focus will be the description of the fibers of the parametrization map W + sw, meaning that we will describe the sets of weights that define the same function. More precisely, the fiber of Fw E Ma,a is the
(W' <ww = (w)
(6)
Once the fibers are understood, the dimension of the neuromanifold can be computed. To this end, is actually sufficient to describe the generic fibers, i.e., the ones corresponding to 'almost all W or, more precisely, to W lying outside of the common zeros of a polynomial system. The co-dimension of such fibers is constant and coincides with the dimension of the neuromanifold.
In order to study the parametrization map and its fibers, is convenient to split the problem by considering self-attention mechanisms as parametrized via the attention matrix. More precisely, will think of self-attention mechanisms as parametrized, by abuse of notation, via weights W = (A, V), where A E gaxd an arbitrary matrix, and will study the matrix multiplication map (Q, A = K Q independently. We begin by considering the latter. When a < d, the matrix constrained to have rank S a. In other words, the
multiplication map is not surjective, since image of this map is the determinantal variety - defined as the set of matrices in Rdd of rank at most a. On the other hand, the fibers of the matrix multiplication map are subtle, since they are closely related to the problem of matrix factorization. Yet, it is still possible to describe the generic fibers. To this end, note that the map exhibits the following invariance: KT FTQ K' Q', where K" = CK and Q' = C- Q for an arbitrary invertible matrix C' E GLa (RR. Conversely, the following elementary result shows that this is the only symmetry of generic fiber.
Lemma 3.1. Suppose that A = K Q K" Q' A' and that rk(A) rk(A') = < d. Then there exists a unique invertible matrix C E GL (R) such that K' = CK and Q' = C C-QQ.
Proof. See Appendix A.I
If follows from the above result that, for < d, the generic fibers of the matrix multiplication map are isomorphic to GL, (RR, and therefore have dimension a². This recovers the well-known formula for the dimension of the determinantal variety, which coincides with 2ad a² a(2d a).
3.1 LINGLE-LAYER IDENTIFIABILITY
We now describe completely the fibers of the parametrization of a lightning self-attention mechanism in terms of the attention matrix. By abuse of notation, we will write fw for W (A, V). Firstly, note that it is always possible to rescale the weights without changing the function. That is, (A, V)
and (AA, (V) belong to the same fiber for all X E (0). Therefore, we will focus on the fibers up rescaling.
Theorem 3.2. Suppose > 2. The fiber of fpw E Ma.d' a for given W (A,V) is follows:
If rk(A) = rk(V) = 1, given tensor decompositions = k @ q and = @ u for some g, k. E Rd h E Rd {0), the fiber consists, up to rescaling, of W and
W = (u h.hh
If pww = 0, the fiber consists of W = (A', V') such that or V =
Otherwise, the fiber consists only of rescalings of W.
Proof. See Appendix A.2
Note that the second condition of Theorem 3.2 is negligible, i.e., it does not hold for almost all weights W, even under the constraint rk I a. Moreover, if d, 2 or d, a 2 2, the first condition is negligible as well. Therefore, the generic fibers are one-dimensional. As consequence, it is possible to compute the dimension (in the sense of algebraic geometry) of the neuromanifold, even when parametrized via queries and keys, or equivalently, when A is restricted have rank < a.
4

---

Published as a conference paper at ICLR 2025
Corollary 3.3. Suppose that t 2 2 and that d, d' 22 or d, a 2 2. The dimension of the neuroman- ifold is:
{ 2ad + dd' a² -1 if S d, dim (Madd' = d8 + dd' -1
otherwise.
Proof. The formula follows from the fact that the generic fibers of the parametrization are one- dimensional and that the dimension of the determinantal variety is 2nd + dd' a², where min(a, d} (see discussion after Lemma 3.1
3.2 SINGLE-LAYER GEOMETRY
We will now describe the geometry of the neuromanifold of a single layer in more detail. We will return to the question of identifiability for deep networks in Section 3.3 Throughout this section, we assume that 2 2 and that either d, d' 2 2 or d, a 2 2.
Theorem 3.4. The newromanifold Ma.8 is closed in the Euclidean topology. Its (relative) bound- ary points are those P(AV)) of the form == k&&q and V look for some q. E TR,, h E TRdd More- over, Ma,d' is not a smooth manifold: its singular points are the (P(A,V) satisfying rk( rk(A)rk(V) S 1.
Proof. This is an amalgamation of Corollaries A.1
A.4 and A.6 in Appendix A.3
Figure I provides a visualization of M2, for f 2. The latter has dimension 5 and is embedded in
the 40-dimensional space Sym3 [R= ,R²).. The illustration shows a rendering of the neuromanifold in 3-dimensional affine slice of the ambient space. This slice cuts 2-dimensional section of the neuromanifold. The yellow line denotes the singular locus, whose dotted segment emerges by taking
the closure in the Zariski topology.
The above result has several consequences, from a machine learning perspective and, in particular, in terms of learning dynamics. The fact that the neuromanifold is closed in the Euclidean topology implies that it contains its limit points. In particular, when the model is trained via a dynamical system - which the case for gradient descent any training trajectory that converges to an equi- librium in the ambient space will converge within the neuromanifold. Moreover, singularities of neuromanifolds are a central focus Information Geometry, and specifically in Singular Learning Theory (Watanabe| 2009). According to the latter, singularities of neuromanifolds play central role in deep learning, since the function learned by a neural network (via gradient descent) often corre- sponds singular point of the neuromanifolds [8]. In other words, singularities often attract the learning dynamics, resulting in a form of 'implicit bias' associated to the neural architecture. Ac- cording Theorem (.4 singularities of the neuromanifold arise exactly when both A and V have rank < 1 (or vanish). Therefore, this result suggests an implicit bias in attention mechanism towards inferring (extremely) low-rank functions. Such bias has been empirically observed in a variety of neural architectures (Amari et al., 2006), and our result might suggest mathematical explanation this phenomenon, least the single-layer and lightning case.
3.3
DEEP NETWORKS
In this section, we completely describe the symmetries in the parameters of a generic function arising from a deep network of lightning self- olf-attention layers. We show that there are only three symmetries: 1) each layer can be scaled by constant, 2) the keys and queries within each layer can be scaled by an invertible matrix as Lemma 3.1 and 3) the output of one layer can be scaled by an invertible matrix if the next layer cancels out this scaling. We now describe the latter type of parameter symmetry, before formally stating our main result in Theorem 3.7
To this end, consider dimension vectors d, and weights W (A, V) of network with layers. For < i & 1, define:
L = II Vi-
M2 = LT ALL-1
L-g
Moreover, set Mj = Ai. It follows immediately from Definition 2 that a deep self- network can be written in terms of (M, T.), where M (M,...., Mi) and I. L,, Therefore, introduce
5

---

Published as a conference paper at ICLR 2025
further re- pa-aaamttriiztionn and write, with abuse of notation, Pw for W (M,L), We call the parameters in this parametrization virtual weights.
This parametrization has symmetries. Namely, invertible matrices C2 E GLd, (R) for S < consider:
= -CCV,
A/+1 = ATTAIOO,
In the above, we set Co and C₁ to the identity. Replacing Vi. V/+1, Ai+1 with V,, Vi4 A/+1 does not alter the virtual weights M and L. Intuitively, this operation consists of transforming the output of the i-th layer and cancelling back the transformation in the next layer. The following result states that, for generie V and with assumption the dimensions d,, the above procedure completely characterizes the generic fibers of the map (A,V) + (M, L.).
Lemma 3.5. Suppose that for some N holds that d, = 6 for all < < and do di 6. Let (M, L) be virtual weights (Equation 8y obtained from both (A, V) and (A', V'). If
then for every there exists a unique C₁ GLa, (R) such that V and 1/411 are obtained from Vi and A4t1 via Equation 9
Proof. See Appendix A.4
The assumption on the dimensions d; in the above result can be practically interpreted as 'bot- tleneck' architecture. Specifically, this architecture equivalent including a low-dimensional embedding and high-dimennional unembedding layer, which is common in the literature, espe- cially an interpretability context (Elhage et al. 202I).
Before computing the fibers of the ce-parametrization, it is convenient to rephrase the latter. The following result provides a recursive expression of a deep attention network in terms of (M, L).
Lemma 3.6. For X Rdaxt, define inductively:
[ Do =1 D; = D; M.D, X-
(10)
Then:
(w(X) LX
II 1iit
D M
(11)
Proof. See Appendix A.5
The above result can be rephrased without recursion. Given X = (xi)isiso since XX
Eisig T,t Equation 11 states that (o((X) can be written for every k as:
Σ
Li
II 1535i-1
k M
x7 A1
(12)
ki....
where i - (3' 1)/2 and My is defined follows. Let ay the index of the first non-zero digit
from the right of j in base 3 (i.e., its 3-adic valuation plus one). Then M, equals to Ma, if this digit is 1, and to M.T. if it is 2. We provide diagram illustrating Equation 12 in Figure 2 where the matrices in the equation correspond, in order, to the regions between the lines.
We now discuss the fibers, as anticipated. Firstly, similarly to the one-layer case, W (M, L) can be rescaled without altering the corresponding function. More precisely, if M/ = X,M, and L' = pL for some Xi P E R (0), then yw = Pw if:
II (1) 1 18151
(13)
P
Conversely, we now show that rescaling characterizes the generic fibers. The following is the main result of this work.

---

Published as a conference paper at ICLR 2025
M
Ma
M2
M
M2
Mi
Mi
M²
M,
Mi
Mi
Mi
Figure 2: Diagrammatic illustration of Equation
12
Theorem 3.7. Let W (M, L). Suppose that 2 and:
For 1 1si5 1, rk(M)) 2,
For Forl<is M, is not skew-symmetric
Then the fiber of fow consists of rescalings of W.
Proof. The proof is highly technical. Here, provide short summary; for the full proof, Appendix [.6] Using unique factorization of polynomials, we first show that, up to rescalings, the fiber consists of W (M', L),, where Mf M,+ and E, is skew-symmetric matrix. We then proceed to show that Σ; = 0 via an induction argument over The argument involves analyzing specific monomials where 2; appears, and proving that most of them vanish due the symmetries
of the unrolled recursion tree in Figure
The conditions of the above theorem are generic if d,, 4; 2 for all i. To summarize, the theorem states that generic function in the neuromanifold has precisely three types of symmetries in its parameters: 1) scaling each layer by a constant as in Equation [3 2) scaling the keys and queries of each layer by invertible matrices
in Lemma 3.1 and 3) scaling the output of one layer by an invertible matrix and reverting this scaling in the next layer as in Equation 9 Similarly to the one-layer case, this results leads to the computation of the dimension of the neuromanifold.
Corollary 3.8. Suppose that 2 3, a; 2 for all and for some d > 2 holds that dN 6 for all 0 < i < and do, di 2 8. The dimension of the neuromanifold is:
20,do - aî + 8(do + d)) - -1 Σ (20,0 a.)) 1si5
(14)
where a; = min(a;;d441)
Proof. See Appendix A.7
As discussed in Section] I an exact expression for the dimension of the neuromanifold enables one to estimate the sample complexity of the model. While the latter is commonly measured as the number of parameters, the dimension which constitutes the theoretically-conee
estimate can, sometimes, significantly differ from it. To illustrate this, in some instances the dimension of queries/keys is set to be equal to the embedding dimension (Vaswani et al. 2017). which translates to setting a; for all i. Assuming do di = ò, the number of parameters is Ldl,, while according to Equation
14 the dimension is dl + 1) - 1. Asymptotically, their ratio is 3 i.e., in this case the number of parameters 50% larger than the dimension.
A square matrix M is skew-symmetric if MT -MM.

---

Published as a conference paper at ICLR 2025
3.4 TRADITIONAL SELF-ATTENTION
In this section, we briefly consider the case of traditional self-attention mechanisms, where the at- tention weights normalized. We compute the fibers in the single-layer and state conjecture around the deep case.
In order to introduce the normalization, consider map 8 I Rso that injective and such that 8(0) 1. A typical choice is 8(z) e*/T, where E Rzo is the temperature hyperparameter. The traditional self-attention mechanism is then defined for W = (A, V) and X E Raxe as:
1 Si
fFw (X)
Σ 1S15I
s (x, Ax;) Vaj
Si = Σ ISKSI
(15)
1gi5t
Note that, for simplicity, adhere the convention of parametrizing via the attention matrix. Intuitively, in Equation 15 the attention weights (xr] Ax;) are normalized to sum to 1, forcing the model "distribute" its attention across the input components. The following result deseribes the fibers of the parametrization, analogously to Theorem 3.2
Theorem 3.9. Suppose that 2 If yow = 0, then V = 0. Otherwise, the fiber of fw is singleton {W}-
Proof. See Appendix A.8
Therefore, differently from lightning self-attention, in this case, the parametrization is generically one-to-one.
We now consider the deep reparametrized via M and L, as defined in Section 3.3 Therefore, the fibers of the parametrization will be unaffected by the symmetries inside the attention matrices from Lemma 3.J] and the transfor- mations C; from Equation However, this time no rescaling is possible. Therefore, we conjecture that the parametrization via (M, L) generically one-to-one; other words, that normalization only breaks the layer-wise scaling symmetry of the parametrization.
Note that with normalization, deep attention networks be
Conjecture 3.10. For normalized deep self-altention networks, the generic fibers of the parametrization via (M, L) are singletons.
In particular, suppose that for some it holds that d, = 6 for all 0 < i < I and do, di 2 8. Similarly to Corollary 3.8 the above conjecture implies that the dimension of the neuromanifold equals to:
20, do of + 5(d + do) 82
Σ
(2a,5 az),
(16)
141
which coincides with the dimension in the lightning case (Equation 14p. summed with the number of layers / due to the removal scaling symmetries. This is an inconsequential difference for large models, where is significantly smaller than the number of parameters per layer.
X
Espected Estimated
Wesromenifad Stoeersion
120
200
40
20
X
Dimerrsion
Figure 3: Plot of the estimated and expected dimensions of the neuromanifold as & varies.
8

---

Published as a conference paper at ICLR 2025
To this end, implement deep attention network with softmax normalization (i.e., 8(x) = E²) ) and estimate the dimension its neuroman- ifold. The latter is subtle problem since, differently from the lightning case, the neuromanifold is not a priori embedded in a finite-dimensional vector space. Therefore, we rely on a stochastic finite element approach by randomly sampling N 250 input points in Rdoxt from normal distribution and restricting fw to this finite space. As result, the neuromanifold is embedded in a (N x t x di)- dimensional vector space. Its dimension can then be computed the rank of the Jacobian of the parametrization at random parameter W. This provides the correct result with probability w.rt. the sampling of W. with the only possible error coming from the discretization, which can be corrected for by increasing the number of input samples. Our Python code is available at public repository [ The results are visualized
We provide empirical evidence for Conjecture 3.10
Figure 3 for a deep attention network with 1 = 2 layers, f == 3, ag 2 for all and de ō varying from 3 to 10. The plot shows both the dimension esti- mated via the numerical approach ( ("Estimated") and the one computed via Equation I6 ('Expected') The two values coincide for all 6, confirming Conjecture 3.00
CONCLUSIONS AND FUTURE WORK
In this work, we have analyzed the geometry of neuromanifolds of lightning self-attention networks. In particular, we have provided description of the fibers of the parametrization, and consequently computed the dimension of the neuromanifold for an arbitrary number of layers. Finally, we have formulated an analogous conjecture for traditional self-attention networks.
Our work represents a first step towards the mathematical understanding of neuromanifolds de- fined by attention networks. As such, it is subject to limitations and leaves several questions open. Specifically, the attention networks we consider are simplified version of the ones deployed in practice, since we omit popular architectural variations (Brauwers Frasincar] (02T). Two such variations, which are ubiquitous in contemporary Transformers, are skip connections and multiple heads. With both these additions, the lightning self-attention mechanism still polynomial. Note that skip connections make the model non non-homogeneous, which breaks the scaling symmetry in the parameterization. We believe that in this case the paramcterization via (M, L) is generically one-to-one, similarly to the traditional case (Conjecture 3.) which is also non-homogeneous, In contrast, including multiple attention heads introduces new symmetries due to permutation of heads, similarly to the permutation symmetries of traditional Multi-Layer Perceptrons (Kileel et al.] 2019). Therefore, these two variations give rise to interesting phenomena in terms of symmetries of the parameterization. and define future directions that worthy of exploration.
From a wider perspective, the research program of applying the tools offered by algebraic geometry the field of deep learning remains open and worthy of exploration. Even further, going beyond the polynomial setting- c.g., addressing problems such as Conjecture 3.10 is an even more general challenge that lies the foundations of the theoretical understanding deep learning.
ACKNOWLEDGEMENTS
This work was partially supported by the Wallenberg AI, Autonomous Systems and Software Pro- gram (WASP) funded by the Knut and Alice Wallenherg Foundation.
REFERENCES
Shun-ichi Amari, Hyeyoung Park, and Tomoko Ozeki. Singularities affect dynamics of learning in neuromanifolds. Neural computation, 1007)160, 2006.
Joachim Mona-Pellissier François Bachoc, and François Malgouyres. Parameter identifiability of deep feedforward relu neural network. Machine Learning, 1 11(11):443-4492 2023.
Gianni Brauwers and Flavius Frasincar. A general survey on attention mechanisms in deep learning. IEEE Transactions on Knowledge and Data Engineering, 35(4):3279-329 2021.
https://github..co -archettt: /NeuroDim

---

Published as a conference paper at ICLR 2025
Ovidiu Calin. Neuromanifolds. Deep Learning Architectures: A Mathematical Approach, pp. 465- 504, 2020.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 1929, 2020.
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Gan- guli, Zac HatfieldDodds, Danny Hernandez, Andy Jones. Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. https::
wtw.:/ts...p//222/E/eeeeeeeeee fra.eworrrawwok.....mm index. html
Charles Fefferman et al. Reconstructing a neural net from its output. Revista Matemática Iberoamer- icana, 10(3):077-556, 1994.
Elisenda Grigsby, Kathryn Lindsey, and David Rolnick. Hidden symmetries of relu networks. In International Conference Machine Learning, pp. 11734-11760. PMLR, 2023.
Robin Hartshorne. Algebruic geomelry, volume 52. Springer Science & Business Media, 2013.
Cheng-Zhi Anna Huang, Ashish Vaswani, Jakob Uszkoreit, Noam M. Shazeer, Ian Simon, Curtis Hawthorne, Andrew M. Dai, Matthew D. Hoffman, Monica Dinculescu, and Douglas Eck. Music transfomer: Generating music with long-term structure. In International Conference on Learn- ing Representations, 2018. URL https://apt
54477714
Joe Kileel, Matthew Trager, and Joan Bruna. On the expressive power of deep polynomial neural networks. Advances in neural information processing systems, 32, 2019.
Kathlén Kohn. The geometry of the neuromanifold. Collections, 57(06), 2024.
Kathlén Kohn, Bernt Ivar Utstøl Nødland, and Paolo Tripoli. Secants, bitangents, and their con- gruences. Combinatorial Algebraic Geometry: Selected Papers From the 2016 Apprenticeship Progrum, pp. 87-112, 2017.
Kathlén Kohn, Thomas Merkh, Guido Montúfar, and Matthew Trager. Geometry of linear convolu- tional networks. SIAM Journal on Applied Algebra and Geometry, 6(3):368-406, 2022.
Kathlén Kohn, Guido Montúfar, Vahid Shahverdi, and Matthew Trager. Function space and critical points of linear convolutional networks. arXiv preprint arX1v:2304.05752, 2023.
Kaie Kubjas, Jiayi Li, and Maximilian Wiesmann. Geometry of polynomial neural networks. arXiv preprint arXiv:2402.000 2024.
Daniel Kunin, Javier Sagastuy-Brena Surya Ganguli, Daniel LK Yamins, and Hidenori Tanaka. Neural mechanics: Symmetry and broken conservation laws in deep learning dynamics. arXiv preprint arXiv:2012.04728, 2020.
Giovanni Luca Marchetti, Vahid Shahverdi, Stefano Mereta, Matthew Trager, and Kathlén Kohn. An invitation to neuroalgebraic geometry, 2025. URL http:/tta///ax....rr/////////// http://aaxxvv..o/bbbssbs.ss//// artp.///50///1111111
18915
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight
programmers. In International Conference on Machine Learning, pp. 9355-9366. PMLR, 2021.
Vahid Shahverdi, Giovanni Luca Marchetti, and Kathlén Kohn. On the geometry and optimization of polynomial convolutional networks. 2024. URL htttttts:///att
ottpga////.//aaa//////////// 00722
Shai Shwaev-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algo-
rithms. Cambridge university press, 2014.
10

---

Published as a conference paper at ICLR 2025
Matthew Trager, Kathlén Kohn, and Joan Bruna. Pure and spurious critical points: geometric study of linear networks. arXiv preprint arXiv:1910001661.00000 2019.
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- tion processing systems, 30, 2017.
Sumio Watanabe. Algebraic geometry and statistical learning theory, volume 25, Cambridge uni- versity press, 2009.
Bo Zhao, Nima Dehmamy, Robin Walters, and Rose Yu. Symmetry teleportation for accelerated optimization. Advances in neural information processing systems, 11116666-16999 2022a.
Bo Zhao, Tordan Ganev, Robin Walters, Rose Yu, and Nima Dehmamy. Symmetries, flat minima, and the conserved quantities of gradient flow. arXiv preprint arXiv:2210.12216, 2022b.

---

Published as a conference paper at ICLR 2025
A PROOFS OF THEORETICAL RESULTS
In this appendix, we include the proofs of the theoretical results in the main body of the paper.
A.1 PROOF OF LEMMA 3.1
Proof. Since a 2 Qk( 2 rk(A) a, we deduce rk(K) a, and similarly for Q K" Q' Therefore, without loss of generality, we assume that K Kı K2, where Ki E GL, and K2 E Rax(d-o), and similarly for K", ",, Q'. If follows immediately that exists unique C' E (R) such that Ki = CK1. Since by hypothesis K2 Q1 = Ki Q1, we deduce Q1 = C 9j But then K. Q1 K2 Qi Ky C-*Q1, implying that K2 K2 C-T and, similarly,
Q2 = C' Q2 as desired.
A.2 PROOF OF THEOREM 3.2
Proof. Suppose that fou (X) Pw (X) for all X = (s,) R where W (A', V*). By comparing the terms containing only in the polynomial identity (x(()) VW (X),, we obtain:
Vzi x A2 V'a1 A'z1.
(17)
The above equation is product of a quadratic form and a linear multivariate form. If AA)) 1, the quadratic form irreducible, and from the unique factorization property of polynomials it follows that, up to multiplicative scalars, V coincides with V', and the quadratic form associated to A coincides with the one associated to A'. The same holds if rk(V) > 1, since the sides of the above equation represent cubical polynomials that share quadratic factor, but whose remaining linear factors are independent. order to prove that A and A' coincide (up to scaling), consider the terms linear in æ1 and quadratic in 22 in the identity yo (w(()) vw (X)1. This leads to:
Vxa x. Az = V'ra x22 A'x,
(18)
follows that the bilinear forms associated to A and A' coincide (up to scaling), as desired for the third case of the claim.
If instead rk rk(A) rk(V) 1, then it is possible to factor A k @ q and V h & v for some vectors q. k 2 E Rd \\ {0). h e IRd \. {0}. Exchanging the role of and v does not alter Equation D which provides the first case of the claim. Finally, since the ring of polynomials is an integral domain, if yow (X) = 0 then from Equation 17 follows that either = 0 or V = 0, as desired for the second case of the claim.
A.3 PROOF OF THBOREM 3.4
To prove Theorem 3.4 we begin with the following key insight: consequence of Theorem 3.2 that it is possible to projectify the neuromanifold and its parametrization (with respect to the attention matrix). To this end, denote by IPV = V/(R) \ {0)) the projectification of vector space V. The second condition of Theorem 3.2 implies that the parametrization descends to a well-defined morphism:
c PRpdxd x PRdxd + PSym, dddxt Rdd Rdxx)).
(19)
We denote by IMM. a the image of the map T referred to as projective neuromanifold. The fact that Tp is well-defined implies the following topological result.
Corollary A.1. The neuromanifold Ma,d' is closed in the Euclidean topology.
Proof. A continuous map between a compact and a Hausdorff space is closed, and in particular has a closed image. These properties are satisfied by 5, where both the domain and the codomain are equipped with the Euclidean topology over projective spaces. It follows that PM4. .o is Euclidean closed in IPSym, dxd Rdd xt). Since the parametrization is homogeneous, the neuromanifold coincides with the affine cone of its projectification. Since affine cones of closed subspaces are closed, the claim follows.
[
12

---

Published as a conference paper at ICLR 2025
Next, we characterize the lightning self-attention mechanisms that singular in the neuromanifold, i.e. whose tangent space has a dimension greater than Equation 7 For that, we combine Theorem
3.2 with computation of the critical points of the network parametrization map. Firstly, remark
that the last point in Theorem 3.2 implies that yp is generically one-to-one onto its image MPM..,d'
i.e,, is birational.
Corollary A.2. The map v is birational onto its image PMM as with special fibers of cardinality 2.
Proof. This follows from the first and the last point of Theorem 3.2
Next, we compute the critical points of the parametrization. To this end, recall that a point critical for differentiable map if the differential at that point does not have maximal rank.
Lemma A.3. A point (A, V) = W is critical for the parametrization map W + &ww if and only if A = k @ and V = h @ k for some 9. k E R h E R
Proof. Given W (A, V), by computing the partial derivatives of yo. we see that the differential sends a tangent vector (A, V) to VXXT AX + VXX ÁX. The latter is interpreted as an element of the vector space Sym;
[dxt Rdd xtdd'x)) which is identified with the tangent spaces of all of its elements. Since the dimension of the neuromanifold is one less than the dimension of the parameter space, the point W is critical if and only if the kernel of the differential has dimension > 1. The kernel consists of those (A, V) that, for all X, satisfy
VXXT AX = -VXX ÁX.
(20)
V 0. the differential vanishes whenever V O; and similarly for A. This provides the first case of the claim. If rk( AV)) 2 2, Theorem 3.2 implies that the kernel consists of (V -VV, Á XA) for À E R, and W is therefore not critical. Finally, suppose that rk(A) rk(V) = 1, and write A h @ q and V h @ for some 9. k. R h R In order to obtain larger kernel, by Theorem 3.2 it is necessary that = (up to scaling). In that case, for every Á of the form Á p @ 9, have that V = -h @ is a solution to Equation 20 implying that W V) is critical.
By exploiting the fact that & birational with finite fibers, the above result leads us to the charac- terization of the singular points of Ma a.
Corollary A.4. P(AAV) is singular in the neuromanifold if and only if (((rrk(()) S
Proof. The morphism & is finite and birational by Corollary [A.2] For such map. a standard fact from algebraic geometry (Kohn et al. 2017 Lemma 3.2) says that point yw is singular in the projective neuromanifold IPM..a'.. if and only if either its fiber under & has cardinality 2 W is critical for Φ.. By Lemma A.3 and Theorem 3.2 a point with rk( rk(A)rk(V)) S 1 is either critical or has a fiber of cardinality 2, and P(A.V) is therefore singular in PM Since Ma,ad' is the affine cone of PMa.d'... as point is singular in the projective neuromanifold if and only if the corresponding line is singular in the neuromanifold, from which the claim follows.
Finally, we compute the boundary points of Ma,d' a We will see that they are precisely the critical values of the parametrization map. For that, we interpret the neuromanifold as Segre variety on linear forms, as follows. Given vector space V. we consider its dual space 1 and the outer product of vectors of linear forms:
ã: (Vr)d (Vr)d' + (V) (a.v) + (am Vn, P SmSSdisnn.
We can projectify this map analogously to 7, obtaining birational morphism (if d > 1 or d' > 1), but not an embedding, If < d, we need restrict the first factor of the map to tuples of linear forms that span a subspace of rank at most a. We denote this space by (V*)d Sa.
Proposition A.5. The neuromanifold Maddr is linearly isomorphic to the image of o restricted to (V*)g x (Vryd', where V = Rdd
13

---

Published as a conference paper at ICLR 2025
Proof. We denote by am the linear form that takes the inner product with the m-th column of A and by Pn the inner product with the n-th row of V. We have seen in the proof of Theorem 3.2 that every PWw Ma is uniquely determined by its term that is linear in 21 and quadratic in 22, which is
22 A Væ² =
Σ
2 1,m O'm (12) Vn (x22
Isngd'
Since the 21.m are formal variables, this expression is the collection of all products linear forms am V/ for 1 i S d and 1 S $ d', which shows the claim.
In what follows, relative boundary refers the set of points (in the Euclidean topology) in Ma,d' o that limit points of sequences in Ma,d', \ Ma,d', where Ma.d' denotes the closure of the neuromanifold in the Zariski topology of its ambient space Sym3 (Rext, Rddxxt
Corollary A.6. fo is on the relative boundary of the neuromanifold if and only if W it is critical for the parametrization map.
Proof. For simplicity, we write M. for Ma.d' and start by computing its Zariski closure M. The Zariski closure a subset S of R consists precisely the real points the complex Zariski closure of S viewed subset of C We compute the complex Zariski closure of M by considering the parametrization map Ø in (9) over C instead of over R. That map Pe is also well-defined morphism between projective spaces. The main theorem projective varieties (Hartshorne 2013 II, $4, Theorem 4.9) states that the image PMC of Pc is Zariski closed. Hence, the real Zariski closure M consists precisely of the real points ((A.V) with A E dd V Cdded'.
Next, we compute the points in the complement M M. Using the identification in Proposition A.S points that complement look like real matrices with entries Om Vns where Om, Un. are complex linear forms, and real parametrization exists. So one of the a'm is non-real. Since each of the am Pn is real, every Vn has to be the complex conjugate am times a real scalar. Analogously, all a, have be Vy times real scalar, i.e., they need coincide with am up to scaling. This shows that MV . consists of ((AVV) with A 1 E q and V @ T for some q E Fd {0}, h E R \ {0), & Cd \ Rd. In the limit, sequences of W = (A, V) in this form can become real if either one among and vanish, or the complex conjugated pair k, becomes the same real vector k.
A.4 PROOF OF LEMMA 3.5
Proof. Note that V, is invertible for every < by the rank hypothesis. Moreover, since L Nosssii Vi- = IIos Vi'z. it is possible to apply Lemma 3.1 by splitting the latter factorization
of L at the first and last layer. This way, obtain matrices C GL,(RR) for < i < such that V = C.V.C But then we have:
AL,A, =
= -L
(21)
Since L; is surjective by the rank hypothesis, we conclude that CT, = AA.
A.5 PROOF OF LEMMA 3.6
Proof. For < i < denote by X; the output of the i-th layer of the network. In other words, for i > O. We wish to prove that for all i:
Xo = X and X; = Pw, (X,_1) =
x; L II
(22)
15J5i
14

---

Published as a conference paper at ICLR 2025
Denote by Y, the right- side of Equation
22 We have:
Yi = Li XD
M
X II
D-1--XTM--,X
(23)
Vih-11
15jsi
V
(I=-1
A, L
II D 15j<i
(24)
Therefore, in order show that Y; satisfies the same recurrence relation X,, need to prove that L,DD, (L,X)* Y,Y," for all i, which in turn reduces to show that D; Z where:
Z, II (Di-,X Mi-+t1X).
(25)
IS/Si
To this end, since Zi = Di-1 X M,X 7 we have:
2,Z, = D M,X
X
(26)
If we assume inductively that Z = Di-1, and using the fact that D, is a symmetric matrix (i.e., D_ Dı-1), the above recurrence relation coincides with the one defining D, in Equa-
tion 11 This implies that Z D;, desired.
A.6 PROOF OF THEOREM
Proof. Pick weights W. W' such that (x) = Pw (X) for all X = (I;)ıgise E Rd Our proof strategy will involve considering the polynomial identity
Fw (X) ww (X)
(27)
and comparing monomial terms arising from Equation 12 with specific degrees in the wy's.
To this end, Equation 12 implies that all such terms contain T1. Even further, the unique term in (p(()) that is linear in 2ı and of degree 3' in 22 can be written as:
Li22
II (o) Mx2) 25i5i
31 (o) (oT)
xy A1Z
(28)
Put simply, the above expression is product of a linear form in 22, several quadratic forms in 22, and a bilinear form in 21 and 22. Moreover, the quadratic forms are coprime by hypothesis since rk(M,) > 1 for i > 1. By comparing this term with the corresponding term on the right-hand side of Equation 27 from the unique factorization property of polynomials, it follows that up to multiplicative scalar, the following hold:
L coincides with L/,
Ai coincides with A₁.
The quadratic forms associated to M, and M; coincide for all 2 < iis.
The last condition above implies that for every 2 i < there exists a skew-symmetric matrix Σ; such that M₁¹ coincides with M,+ + E, up to a multiplicative sealar. Since Equation 12 is multi-linear in each occurrence of M,. by substituting Mf M, + Σ, in (X): for all j, we obtain of expressions in the form of Equation 12 but where an arbitrary number of occurrences of My has been replaced by Σy. The only expression with no replacement coincides with pw Therefore, Equation [7 reduces to a vanishing sum of copies Equation 12 where at least one M, has been replaced with Σ.z.
In order to conclude, we wish to show that Σ, 0 for i > 2. We will proceed by induction on Specifically, given and assuming that Σ, for j i, we will show that Σ, 0. To this
end, consider the monomial terms in sw (X) of degree 1 in Tı, of degree 3--1 - 1 in 33. and
15

---

Published as a conference paper at ICLR 2025
of remaining degree in Z2. This is possible since assume t 2 3. From the discussion above, it
follows that Equation 27 reduces to a vanishing sum of copies of Equation 12. where an arbitrary
(non-zero) number of occurrences of Mj Maj, with aj 2 i, has been replaced with Ea, (and
similarly with transposes), and where k, 3) for all S j S i (3' 1)/2.
We will now argue that several terms cancel due to the symmetries of Equation 12 illustrated in figure 4 Namely, consider monomial term with some multi-index h. containing factor of the form 2k 2 Bkjzk for some j. with a, > Since 2y cannot appear more than 3- times in the monomial, there exists _ I (3 - 1 such that kj-j kjtj+1 = 2. Moreover, due to the inductive hypothesis, My coincides with May or with its transpose for all
since ay' < i. Consider the multi-index Kg such that ký-j" ky+5'+1 for all -j S 3 S and Kj, ky for all the other j'. Intuitively, K, 'reflects' hi locally around j see Figure 4 for illustration. By construction, the monomial corresponding to ki is identical to the one corresponding to hi, except for the term
k Therefore, these two monomials cancel each other out.
1ks Ca, Thy - -ky Log 32ay12
(29)
Σ
Mi
M
M1
M
Figure 4: Diagrammatic illustration of the symmetry involved in the cancellation argument.
After the above cancellation argument, we left with monomials containing occurrences of Eay only for 0j = i. In fact, the cancellation argument still applies verbatim to these cases, except for when j = (3' 1/21)/2 and ky-j' * ky+j+1 for all j' S (3i-1 1)/2. In this case, due to the k. and K,, are not opposite, but contain different factors given by x. A13221 dy and - -x3 Aιoı x22 A122, respectively. At the end of the day,
presence of 3ı, the monomials corresponding
Equation 27 reduces to an expression of the form:
z3 Σ (I3 A A 2J A1T1 A122) Σ (((En,)))) =0
(30)
where h. is an opportune multi-index with values in (2,3), and p is a polynomial consisting of products of bilinear/quadratic forms associated to the M,'s 's and a multi-variate linear factor asso- ciated to T.. The right-most factor in Equation 30 is not the zero polynomial (in the variables 22 and æ²) since, by imposing the condition Ø2 = Ø3, it becomes (up to a multiplicative scalar) a product of quadratic forms associated to the M,'s. s. The latter are non-vanishing since M, is not skew-symmetric for all j by hypothesis.
We wish to prove that the second factor in Equation 30 is non non vanishing well. To this end, note that the coefficients of that factor, seen as a polynomial in I T2, 23, are of the form
(A1)a, (A1) + (A)),, (A1) (A1)..) - (
(31)
for < a, 0,8,7,8 < dg. Suppose by contradiction that the above expression vanishes for all o, 8, 7, When B 7, Equation 31 coincides with (twice) the determinant of arbitrary 2
16

---

Published as a conference paper at ICLR 2025
minor of A1 intersecting the diagonal. If (A1)8.9 O for some 3. then from the Kronecker's bordered matrix theorem it follows that A1 has rank 1, contradicting the hypothesis rk(Ar) 2 2. We conclude that the diagonal entries of A₁ are zero, and, when / y Equation [3 reduces to (A1)a,e(A) s., = O. Therefore, if (A1),.,8 o 0 for some a * S. then the 35thh row of A must vanish. But then Equation 31 reduces (up sign) to the determinant of arbitrary 2 x 2 minor intersecting (a, B), and we again obtain a contradiction from Kronecker's bordered matrix theorem.
In conclusion, the left-most factor in Equation 30 must vanish, meaning that Σ, 0, desired.
A.7 PROOF OF COROLLARY 3.8
Proof. Recall that the dimension of the determinantal variety of d:-1 x d;-1 matrices of rank at most a; is 20,di-1 - a² and therefore the space of parameters (A, V) has dimension
20, do af Ō(do + di) + (1 2)6²
Σ (2a, af))
(32)
KiS
Moreover, by Lemma 3.5 the generic fibers of the re-parametrization (A, V) + (M, L) have di-
mension (Z-1D 1)82, while by Theorem 3.7 the generic fibers with respect to (M, L) have dimension I (recall the constraint on rescaling given by Equation [3). Since the dimension of the image of map coincides with the co-dimension of the generic fibers, the result follows.
A.8 PROOF OF THEOREM 3.9
Proof. Suppose that Pw (X) Pw (X) for all X = (I:) E Rdxt, where W V'), We assume that V j 0, since the case FW 0 follows immediately. Due to normalization, if x; Ij for all j,j, then 4ow (X)1 Va1 V (w:(X),, implying V V'. If instead x; 0 for
i > 1. since S(0) 1, we obtain:
8 Ax)
8 (2T A'x1)
8 (0] Ar1) + Vxg
(33)
8 (xx A' +t
Since 8 is injective, we deduce that the quadratic form associated to A coincides with the one associated to A'. In order prove that A A', suppose that 24 2i for i, 2 2, obtaining:
8 ( Az)) Ven 8 8 [xT Ar2) Vaa [1] Am (t 1)8 [5) Ara)
8 (TTTA'r) V'a (2] A w2)) Vaaa
(34)
( A'r + (1 1)8 (z) A'r2)
After elementary algebraic manipulations using the fact that x. Ar = x o and V = V the above equation reduces to:
(t - (8 A'a2) 8 (z Ax2) Vzi = (z) (2) - (x. Arz)) Vx2.
(35)
Since V * 0, Vz1 # Vr2 for generic æ1, 22. Therefore, 8 (z A'r2) - 8 (o Az2) must vanish generically, implying A = A' due to the injectivity of S, as desired.
17


Reviews:
[{"review_id": "dCLTNhtlpk", "reviewer": "Reviewer_fk8M", "summary": "This paper provides a theoretical analysis of the geometry of the function spaces corresponding to self-attention networks. The core of the work concentrates on the “lightning” variant of self-attention (for which forward computation computation is linear in the sequence length) since it is amenable to an algebraic analysis. The authors calculate the dimension of the hypothesis space (“neuromanifold”) through an analysis of the (generic) fibers of the model’s parametrization. The submission provides an extension to normalized attention modules, provably for the case of a single layer, and conjecturally for deep networks.", "questions": "Q1 & Q2. What is the authors’ position on W1 and W2?\n\nQ3. In the (much appreciated!) anonymized code for the verification of Conjecture 3.16, L19 instantiates `n_iter_array = [5, 5, 5, 5, 5, 5, 5, 8]`. These iterations are used as “trials” out of which the maximum rank Jacobian is extracted in L41 `dim_vec.append(max(ranks))`. Could the authors explain why such multiple “trials” are needed? And also why is this component left out of the experimental description in L472-484?\n\nQ4. Could the authors elaborate on the implications of the non-smoothness of the $\\mathcal{M}_{d, d’, a}$ manifold from an optimization standpoint?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "S1. The paper exposition is quite clear, striking a balance between the lingo of algebraic geometry and the context of the ICLR audience.\n\nS2. The empirical popularity of attention-based networks, and the relative lack of understanding on the properties of said architectures (with respect to fully connected and convolutional ones, as the authors point out in L51-53) sufficiently justifies the theoretical exploration presented in this submission.\n\nS3. The authors identify lightning self-attention as a fruitful compromise that is both (1) amenable to their algebraic-geometric analysis and (2) a reasonably realistic neural network module. I believe the presented implicit hierarchy of modules (single-layer lightning self-attention, deep lightning self-attention, single-layer normalized self-attention, deep normalized self-attention) can be of use in further theoretical analyses of attention networks.\n\nS4. The authors introduce a “virtual weights” parametrization which enables an elegant analysis of deep lightning self-attention networks.\n\nS5. The submission poses a conjecture on the fibers for deep normalized self-attention networks (under their virtual weights parametrization). This conjecture is theoretically validated in the case of a single layer and empirically tested for deep networks. The communication of this conjecture to the relevant communities within the context of the ICLR conference can be beneficial towards advancing on its resolution.", "weaknesses": "While S2 can warrant sufficient interest from the ICLR community, I think the submission falls short of adequately examining the implications of its findings. I think the authors could have expanded more on the “so what?” question after reaching their (interesting!) core theoretical goals. \n\nW1. For example, in L416-419, the authors are quite descriptive about the dimension of the neuromanifold for a contemporary language model being ~8.2B vs the crude parameter dimension of ~8.6B. Yet, there is not even a high-level attempt to prescribe how to exploit this information.\n\nW2. The authors partially justify the importance of the analysis of the dimension of the neuromanifold via an appeal to learning theory (L58-61) and its connection with sample complexity. (Although this may be considered outside of the scope!) The submission does not present a tangible demonstration of the implications of the knowledge of the neuromanifold’s dimension on the sample complexity, not even in a very simple toy-ish experimental setting.", "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": "v37xUrSABr", "reviewer": "Reviewer_sSei", "summary": "The paper studies the geometry of lighting attention models by characterizing their fibers (redundant parameterizations of a neural network) and neuromanifolds (the function spaces induced by all possible parameterizations). Critically, they assume for most of the paper that (1) the attention matrix has no normalization (i.e. no softmax); (2) there are no MLP layers of the transformer; and (3) there are no residual connections. They prove the following:\n* From Theorem 3.2, in most cases, the fibers are 1-dimensional (identifiable functions up to rescaling), except under certain rank-1 cases, where the fiber has higher-dimensional symmetries. The dimensionality of these fibers can be used to exactly capture the neuromanifold dimension in Corollary 3.3.\n* Theorem 3.4 shows that the neuromanifold is closed in Euclidean space and that its boundaries are defined by networks with rank-1 matrices.\n* Theorem 3.15 shows that self-attention layers have singleton fibers, and Conjecture 3.16 predicts that deeper self-attention models also have singleton fibers. The conjecture is backed up with experiments that estimate the dimensionality of the neuromanifold dimension and shows a tight alignment to the conjecture.", "questions": "While the results are certainly mathematically interesting, it's unclear to me what practical guidance can be gleaned from this line of work. What are the consequences of understanding the dimensionality of the neuromanifold for the model's expressivity and generalization? \n\nI think parameter redundancy is an interesting angle, but this work seems to suggest that existing models have little room for improvement on that front. Are there modifications to this work (or questions to be pursued by future work) that might look at notions of fibers that pertain to parameters that map to _approximately_ the same output, rather than exactly the same output?  \n\nI appreciate the section that extends this analysis to self-attention layers. Can this analysis be modified to account for the inclusion of residual connections or MLPs?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The contributions reveal interesting properties of the geometry of lightning attention models and provide a rigorous characterization of the resulting functional space. The results are well-described and relatively straight-forward to follow.\n\nAs far as I can tell (although I am no expert in algebraic geometry), the claims are clearly stated and the mathematics are correct.", "weaknesses": "There are several key gaps between the models studied in this paper and practical models:\n* The lack of non-linearities in the attention unit (in particular, the lack of softmax).\n* The lack of residual connections between layers.\n* The lack of element-wise multi-layer perceptrons between layers.\nAs far as I am aware, all three components are important for practical transformers, and that a lack of residual connections or MLPs tend to lead to degeneracy in training. While the manifold analysis in the paper is interesting, it's unclear how it would extend to cases where the model does not have such clear polynomial behavior.\n\nWhile the theory is interesting in its own right, it's somewhat unclear to me what can be inferred about realistic models based on this theory, especially because the softmax results point to no parameter redundancy.", "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": "fXVsWMoRJT", "reviewer": "Reviewer_qe2M", "summary": "This paper studies the geometry of self-attention networks, focusing on \"lightning\" self-attention mechanisms. The paper offers a theoretical analysis of lightening attention using algebraic geometry, focusing on dimension of the neurmoanifold and the identifiability of deep attention. Key contributions include a description of the fibers of the parametrization map and the computation of neuromanifold dimensions. Additionally, the authors extend their findings to traditional, normalized self-attention networks (essentially with softmax activation) and present a conjecture, supported by empirical evidence, regarding their fibers of parametrization.", "questions": "1. The authors say in the conclusion that “Extending our algebra-geometric analysis to include these variations represents an interesting line for future investigation” (line 496). How? Can you be more specific? What do you think are the future direction of the research? Can you specify 2-3 specific directions?\n\n2. What purpose do figures 2 and 3 serve?\n\n3. What do you think are the main implications of the main theorem?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 4, "presentation": 2, "contribution": 2, "strengths": "The paper presents a novel and rigorous theoretical analysis of undoubtedly important topics, i.e., attention mechanisms. The authors use algebraic geometry to develop a solid theoretical understanding of the neuromanifold of neural networks using attention mechanisms. Although the analysis is offered in a simplified context, that is standard and inevitable for a rigorous theoretical analysis and does not reduce the importance of their contribution. I believe the introduction motivates the study quite well (though it could be improved). Understanding the dimension and the topology of the neuromanifold is a consequential result, as it relates to the learning dynamics of transformers in the overparameterized regime (as authors point out, e.g., lines 246-248), and it could enable a series of future studies. The paper could potentially provide foundational insights with a potentially wide range of applications in understanding the mathematical properties of attention mechanisms in deep learning (though, whether this is achieved or not is not fully clear, as I discuss under weaknesses).", "weaknesses": "**Key weakness:** The main weakness of the paper which is present throughout the paper is that the ultimate goal of the theoretical results as well as their implications tend to seem lost and poorly explained. I believe many among the ML audience of this paper might raise the objection that “this paper belongs to a math journal, not ICLR”. I take the liberty to clarify that that is not my objection, and I believe even a paper that solely focuses on deriving results and developing mathematical tools from a purely theoretical perspective could be valuable and could very well belong to ICLR. That being said, I expect such a paper to clarify the implications of their results for the wider community, or at least for future theoretical research on foundations of deep learning. I regret to see that this paper does not meet this criteria. As I mentioned among the strengths, I understand the results are foundational with potentially important consequences, but they are highly (perhaps overly) mathematicized, without proper explanations that clarify what these mathematicized observations entail for the ML applications, e.g., how transformers learn their representation, or how these result allow us to develop a better understand of the convergence of GD-based optimization on transformers (which I think they could do). \n\n\n\n**A list of some of the main weaknesses:**\n\n1. The implications of the theoretical statements are often not discussed (with a small number of exceptions). For instance, the main result of the paper is stated in Theorem 3.13, but no discussion of its implication is provided. In an ML paper, all such theoretical results should follow with a discussion of the implications that are more accessible to ML researchers. \n\n2. Even when implications are discussed, the ultimate goal or the consequences of the implications are not discussed. e.g., the authors discuss the implications of corollary 3.5, but they can clarify what this implies for the success of the attention mechanisms, or at least remark how it could enable future studies on that.\n\n3. The experiments lead to a conjecture, but they are from empirical relevance. I would like to clarify again that this is fine with me, as long as at least the theoretical relevance for future work that could eventually lead to an empirically relevant application is discussed.\n\n4. There are too many technical details leading to the proofs or sketch of the proofs of the results, which I do not think belong to the main text of a machine learning paper (e.g., lines 165-173 or 227-240). These details could be move to the appendix and replaced with a more high-level discussion of the intuition behind the results and the implications of the results.\n\n\n\n**More specific points about the weaknesses:**\n\na. There is almost no review of the literature. There are many interesting (and some quite relevant) recent works on developing a theoretical understanding of transformers or relevant topics (e.g. learning dynamics in overparameterized regimes). It is important to discuss this and explain the gap in the literature. This helps motivate the paper as well.\n\nb. While the introduction motivates the paper, the abstract does not. The abstract could be improved to reflect the motivation explained in the second paragraph of the introduction. Moreover, the second paragraph of the introduction itself could mention more specific examples of the applications and the potential future direction the paper opens a door to, in order to better motivate the paper.\n\nc. The purpose of the visualizations in figures 2 and 3 is unclear.\n\nd. Line 297 says “we introduce yet another re-parameterization”. While I understand you later use this, I suggest you explain closer to where you introduce the reparameterization and why you are introducing it.\n\ne. The summary of the proof of Theorem 3.13 could also be moved to the appendix. It is crucial to discuss the implications of the theorem, instead of explaining the proof in the main paper.\n\nf. The presentation can be improved to be more accessible to ML researchers.", "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": "kTZ1sm6z65", "reviewer": "Reviewer_SY47", "summary": "This paper studies the neuromanifold, the image of the parametrization map that maps network weights to the corresponding input-output map of the network, of single-layer and multi-layer linear-self-attention networks. Specifically, the dimensions of these neuromanifolds are computed, and the fibers of the parametrization of these neuromanifolds are discussed.", "questions": "See \"Weaknesses\". The authors do not need to spend too much time responding to this review, rather, they should revise the paper to be more accessible to a broader audience and motivate the results from a more practical perspective.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "Analysis of linear self-attention networks with tools from algebraic geometry", "weaknesses": "The reviewer is incapable of evaluating the technical contributions of this paper. However, the reviewer has concerns about the significance of the results to ML community and the paper's lack of accessibility to a broader ML audience.\n\n1. This paper computes the intrinsic dimensions of the neuromanifolds for single- and multi-layer self-attention networks, and they are of the same scale as the number of parameters. (single-layer: $2ad+dd'$ v.s. $2ad+dd'-a^2-1$, and multi-layer 8.2b v.s. 8.6b, in an LLM example). The reviewer understands the research effort behind computing the intrinsic dimensions and this review does not mean to downplay its importance; but still, the results do not provide a better characterization of the complexity of the function space than simple parameter counting. \n\n2. This paper also studies the fiber of the parametrization, the boundary of the neuromanifold, critical points of the parametrization map, etc. They are presumably important characteristics to be studied in algebraic geometry, but are they important to the deep learning community? Most of the results are stated without discussing their relevance to the training or generalization of these linear self-attention networks, except for Corollary 3.5: \"Therefore, any training trajectory that converges to an equilibrium in the ambient space will converge within the neuromanifold\", which is, in reviewer's opinion, too vague to imply anything useful. \n\nTherefore, while the results presented in this paper may be of interest to the algebraic geometry community, the reviewer does not see it getting well appreciated by the ML community in its current form.", "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)