Geometry of Lightning Self-Attention: Identifiability and Dimension

## Abstract

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.

## 1 Introduction and Related Work

The *self-attention mechanism* is the cornerstone of the transformer – a modern machine learning architecture that is nowadays popular in a vast variety of domains, ranging from natural language processing (Vaswani et al., ), to vision (Dosovitskiy et al., ), to sound (Huang et al., ). In all of these domains, self-attention mechanisms have showcased outstanding performance, due to their ability to model long-range dependencies within data sequences. *Lightning* self-attention mechanisms (Schlag et al., ) are standard variants where, differently from the original proposal, 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 spaces defined by neural networks – typically referred to as *neuromanifolds* (Calin, ) – is a fundamental challenge due to its intimate connection to the training dynamics. Since neural networks learn by following a gradient flow over the neuromanifold, the geometry of the latter controls several aspects of the training dynamics (Trager et al., ). While neuromanifolds are well-understood for several architectures, such as fully-connected (Kileel et al., ; Kubjas et al., ) and convolutional networks (Kohn et al., ; ), they have not been considered for self-attention mechanisms.

In this work, we study the geometry of neuromanifolds associated to lightning self-attention mechanisms. These models are of algebraic nature, since they are tri-linear in their weights and cubical in the input. This enables to analyze neuromanifolds via ideas and tools from *algebraic geometry* – 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 is a fundamental measure of expressivity of the given model, and it plays a fundamental role in statistical learning theory(^1^11In this context, neuromanifolds are referred to as ‘hypothesis spaces’, and are usually considered in a combinatorial version.). In particular, according to the Fundamental Theorem of Learning (Shalev-Shwartz & Ben-David, ), the dimension controls the sample complexity of learnability. Thus, it is useful in practice for estimating the dataset size appropriate for the given model, and vice versa. The dimension of the neuromanifold is in turn related to the dual question of *identifiability* (Grigsby et al., ; Bona-Pellissier et al., ; Fefferman et al., ) – a problem concerned with characterizing the parameters corresponding to the same function. Geometrically, such parameters define *fibers* of the parametrization of the neuromanifold.

Recently, the concept of *transformer circuits* (Elhage et al., ; Bhaskar et al., ) has been proposed to study self-attention mechanisms in a compositional fashion. This enables to break down a transformer into atomic components and to analyze them individually. The method is motivated by the challenge of *mechanistic interpretability* (Bereska & Gavves, ) – a field aimed at ‘de-compiling’ the function defined by a machine learning model. The identifiability of parameters of a transformer circuit plays a central role in determining the degree of ambiguity of the resulting interpretation. Our work is related to this line of research, since we consider a deep (lightning) self-attention network – which can be interpreted as a rectilinear transformer circuit – and study the identifiability of its parameters.

### 1.1 Summary of Contributions

Our core contribution is a 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, our results are summarized as follows.

For a single layer of lightning self-attention (Section 3.1), we additionally 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 a 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 on 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.

Lastly, study traditional self-attention by re-introducing the softmax normalization and arguing how the lightining self-attention geometry should change (Section 3.4). We prove that the parametrization of a 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^{\prime},a,t\in\mathbb{N}$ and matrices $Q,K\in\mathbb{R}^{a\times d}$, $V\in\mathbb{R}^{d^{\prime}\times d}$. The latter are deemed *query weights*, *key weights*, and *value weights* respectively(^2^22An alternative standard notatiopn for $Q,K,V$ is $W_{Q},W_{K},W_{V}$. Our choice is motivated by better readability.). A self-attention mechanism is a map parametrized by $(Q,K,V)$ sending sequences of length $t$ of vectors in $\mathbb{R}^{d}$ to sequences of vectors in $\mathbb{R}^{d^{\prime}}$. Here we 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:

$$ $\begin{array}[]{ccl}\varphi_{W}\colon\mathbb{R}^{d\times t}&\rightarrow& \mathbb{R}^{d^{\prime}\times t}\\ \left(x_{i}\right)_{1\leq i\leq t}&\rightarrow&\displaystyle\left(\sum_{1\leq j \leq t}x_{j}^{\top}\left(K^{\top}Q\right)x_{i}\ Vx_{j}\right)_{1\leq i\leq t} \end{array}$ (1) $$

Intuitively, every component $x_{i}$ of the input corresponds to a *token* and ‘attends’ bilinearly to every other $x_{j}$, producing a scalar weight $x_{j}^{\top}K^{\top}Qx_{i}\in\mathbb{R}$. These weights are used to aggregate the values $Vx_{j}\in\mathbb{R}^{d^{\prime}}$, obtaining the corresponding component of the output. The map defined by Equation 1 is tri-linear in $(Q,K,V)$ and homogeneous cubical in $x$. It is often convenient to write Equation 1 in matrix form: if $X=(x_{i})_{1\leq i\leq t}$ is interpreted as a $d\times t$ matrix, then:

$$ $\varphi_{W}(X)=VXX^{\top}K^{\top}QX.$ (2) $$

Moreover, we will often simplify the parametrization by introducing the *attention matrix*:

$$ $A=K^{\top}Q.$ (3) $$

The latter will always be interpreted as a bilinear form $x^{\top}Ay$.

Lightning attention mechanisms are variants of the traditional ones (Vaswani et al., ), where the attention weights $x_{j}^{\top}Ax_{i}$ 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 efficiency with respect to the sequence length. Specifically, Equation 1 can be computed in $\mathcal{O}(t)$ time, while traditional self-attention mechanisms require $\mathcal{O}(t^{2})$ time due to normalization. The improvement in efficiency motivates the term ‘lightning’.

Self-attention mechanisms can be stacked in order to obtain a deep network architecture. To this end, fix positive integers $t,l,\mathbf{d}=(d_{0},\ldots,d_{l}),\mathbf{a}=(a_{1},\ldots,a_{l})$, and weights $\mathbf{Q}=(Q_{1},\ldots,Q_{l}),\mathbf{K}=(K_{1},\ldots,K_{l}),\mathbf{V}=(V_
{1},\ldots,V_{l})$, with $Q_{i},K_{i}\in\mathbb{R}^{a_{i}\times d_{i}},V_{i}\in\mathbb{R}^{d_{i}\times d
_{i-1}}$.

###### Definition 2 .

A *deep self-attention network* associated to the weights $\mathbf{W}=(\mathbf{Q},\mathbf{K},\mathbf{V})$ is the map:

$$ $\varphi_{\mathbf{W}}\colon\mathbb{R}^{d_{0}\times t}\rightarrow\mathbb{R}^{d_{ l}\times t}$ (4) $$

given by the composition $\varphi_{\mathbf{W}}=\varphi_{W_{l}}\circ\cdots\circ\varphi_{W_{1}}$.

Again, a deep self-attention network is homogeneous of degree $3^{l}$ in $x$. Based on this, we denote by $\textnormal{Sym}_{3^{l}}\left(\mathbb{R}^{d_{0}\times t},\mathbb{R}^{d_{l}
\times t}\right)$ the vector space of homogeneous polynomial functions from $\mathbb{R}^{d_{0}\times t}$ to $\mathbb{R}^{d_{l}\times t}$ of degree $3^{l}$ in all the output co-ordinates.

###### Definition 3 .

The *neuromanifold* of a deep self-attention network is the image of the parametrization map $\mathbf{W}\mapsto\varphi_{\mathbf{W}}$:

$$ $\mathcal{M}_{\mathbf{d},\mathbf{a}}=\left\{\varphi_{\mathbf{W}}\ |\ \mathbf{W} \in\mathbb{R}^{\sum_{i}d_{i}(2a_{i}+d_{i-1})}\right\}\ \subset\textnormal{Sym} _{3^{l}}\left(\mathbb{R}^{d_{0}\times t},\mathbb{R}^{d_{l}\times t}\right).$ (5) $$

The neuromanifold is a semi-algebraic set by the Tarski-Seidenberg Theorem, meaning that it can be defined by a finite number of polynomial equalities and inequalities in $\textnormal{Sym}_{3^{l}}\left(\mathbb{R}^{d_{0}\times t},\mathbb{R}^{d_{l}
\times t}\right)$.

###### Definition 1 .

###### Definition 2 .

###### Definition 3 .

## 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 $\mathbf{W}\mapsto\varphi_{\mathbf{W}}$, meaning that we will describe the sets of weights that define the same function. More precisely, the fiber of $\varphi_{\mathbf{W}}\in\mathcal{M}_{\mathbf{d},\mathbf{a}}$ is the set

$$ $\{\mathbf{W}^{\prime}\ |\ \varphi_{\mathbf{W}^{\prime}}=\varphi_{\mathbf{W}}\}.$ (6) $$

Once the fibers are understood, the dimension of the neuromanifold can be computed. To this end, it is actually sufficient to describe the *generic* fibers, i.e., the ones corresponding to ‘almost all $\mathbf{W}$’ or, more precisely, to $\mathbf{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 neuromanfiold.

In order to study the parametrization map and its fibers, it is convenient to split the problem by considering self-attention mechanisms as parametrized via the attention matrix. More precisely, we will think of self-attention mechanisms as parametrized, by abuse of notation, via weights $W=(A,V)$, where $A\in\mathbb{R}^{d\times d}$ is an arbitrary matrix, and will study the matrix multiplication map $(Q,K)\mapsto A=K^{\top}Q$ independently. We begin by considering the latter. When $a<d$, the matrix multiplication map is not surjective, since $A$ is constrained to have rank $\leq a$. In other words, the image of this map is the *determinantal variety* – defined as the set of matrices in $\mathbb{R}^{d\times d}$ 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: $K^{\top}Q={K^{\prime}}^{\top}Q^{\prime}$, where $K^{\prime}=CK$ and $Q^{\prime}=C^{-\top}Q$ for an arbitrary invertible matrix $C\in\textnormal{GL}_{a}(\mathbb{R})$. Conversely, the following elementary result shows that this is the only symmetry of a generic fiber.

###### Lemma 3.1 .

Suppose that $A=K^{\top}Q=K^{\prime\top}Q^{\prime}=A^{\prime}$ and that $\textnormal{rk}(A)=\textnormal{rk}(A^{\prime})=a\leq d$. Then there exists a unique invertible matrix $C\in\textnormal{GL}_{a}(\mathbb{R})$ such that $K^{\prime}=CK$ and $Q^{\prime}=C^{-\top}Q$.

###### Proof.

See Appendix A.1
∎

If follows from the above result that, for $a<d$, the generic fibers of the matrix multiplication map are isomorphic to $\textnormal{GL}_{a}(\mathbb{R})$, and therefore have dimension $a^{2}$. This recovers the well-known formula for the dimension of the determinantal variety, which coincides with $2ad-a^{2}=2a(d-a)$.

###### Lemma 3.1 .

###### Proof.

### 3.1 Single-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 $\varphi_{W}$ for $W=(A,V)$. Firstly, note that it is always possible to rescale the weights without changing the function. That is, $(A,V)$ and $\left(\lambda A,\frac{1}{\lambda}V\right)$ belong to the same fiber for all $\lambda\in\mathbb{R}\setminus\{0\}$. Therefore, we will focus on the fibers up to rescaling.

###### Theorem 3.2 .

Suppose $t\geq 2$. The fiber of $\varphi_{W}\in\mathcal{M}_{d,d^{\prime},a}$ for a given $W=(A,V)$ is as follows:

- •
If $\textnormal{rk}(A)=\textnormal{rk}(V)=1$, given tensor decompositions $A=k\otimes q$ and $V=h\otimes v$ for some $q,k,v\in\mathbb{R}^{d}\setminus\{0\},h\in\mathbb{R}^{d^{\prime}}\setminus\{0\}$, the fiber consists, up to rescaling, of $W$ and $W^{\prime}=(v\otimes q,h\otimes k)$.
- •
If $\varphi_{W}=0$, the fiber consists of $W^{\prime}=(A^{\prime},V^{\prime})$ such that $A^{\prime}=0$ or $V^{\prime}=0$.
- •
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 $\textnormal{rk}(A)\leq a$. Therefore, the generic fibers are one-dimensional. Moreover, if $d,d^{\prime}\geq 2$ or $d,a\geq 2$, the first condition is negligible as well. As a consequence, it is possible to compute the dimension (in the sense of algebraic geometry) of the neuromanifold. Since the two even when parametrized via queries and keys i.e., when $A$ is restricted to have rank $\leq a$.

###### Corollary 3.3 .

Suppose that $t\geq 2$ and that $d,d^{\prime}\geq 2$ or $d,a\geq 2$. The dimension of the neuromanifold is:

$$ $\textnormal{dim}\left(\mathcal{M}_{d,d^{\prime},a}\right)=\begin{cases}2ad+dd^ {\prime}-a^{2}-1&\textnormal{if}\ a\leq d,\\ d^{2}+dd^{\prime}-1&\textnormal{otherwise}.\end{cases}$ (7) $$

###### 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 $2\alpha d+dd^{\prime}-\alpha^{2}$, where $\alpha=\min\{a,d\}$ (see discussion after Lemma 3.1).
∎

###### Theorem 3.2 .

###### Proof.

###### Corollary 3.3 .

###### Proof.

### 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 $t\geq 2$ and that either $d,d^{\prime}\geq 2$ or $d,a\geq 2$.

###### Theorem 3.4 .

The neuromanifold $\mathcal{M}_{d,d^{\prime},a}$ is closed in the Euclidean topology.
Its (relative) boundary points are those $\varphi_{(A,V)}$ of the form
$A=k\otimes q$ and $V=h\otimes k$ for some $q,k\in\mathbb{R}^{d},h\in\mathbb{R}^{d^{\prime}}$.
Moreover, $\mathcal{M}_{d,d^{\prime},a}$ is *not* a smooth manifold:
its singular points are the $\varphi_{(A,V)}$ satisfying $\textnormal{rk}(A)\textnormal{rk}(V)\leq 1$.

###### Proof.

This is an amalgamation of Corollaries 3.5, 3.8 and 3.10 below.
∎

Figure 1 provides a visualization of $\mathcal{M}_{2,1,2}$ for $t=2$. The latter has dimension $5$ and is embedded in the $40$-dimensional space $\textnormal{Sym}_{3}\left(\mathbb{R}^{4},\mathbb{R}^{2}\right)$. The illustration shows a rendering of the neuromanifold in a $3$-dimensional affine slice of the ambient space. This slice cuts a $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.

To prove Theorem 3.4, we begin with the following key insight:
a consequence of Theorem 3.2 is that it is possible to *projectify* the neuromanifold and its parametrization (with respect to the attention matrix). To this end, denote by $\mathbb{P}\mathcal{V}=\mathcal{V}/(\mathbb{R}\setminus\{0\})$ the projectification of a vector space $\mathcal{V}$. The second condition of Theorem 3.2 implies that the parametrization descends to a well-defined morphism:

$$ $\overline{\varphi}\colon\mathbb{P}\mathbb{R}^{d\times d}\times\mathbb{P} \mathbb{R}^{d\times d^{\prime}}\rightarrow\mathbb{P}\textnormal{Sym}_{3}\left( \mathbb{R}^{d\times t},\mathbb{R}^{d^{\prime}\times t}\right).$ (8) $$

We denote by $\mathbb{P}\mathcal{M}_{d,d^{\prime},a}$ the image of the map $\overline{\varphi}$ – referred to as *projective neuromanifold*.
The fact that $\overline{\varphi}$ is well-defined implies the following topological result.

###### Corollary 3.5 .

The neuromanifold $\mathcal{M}_{d,d^{\prime},a}$ is closed in the Euclidean topology.

###### Proof.

See Appendix A.3.
∎

The above result has an interpretation in terms of the learning dynamics. Namely, the neuromanifold contains its limit points and, therefore, any training dynamics that converges to an equilibrium in the ambient space will converge within the neuromanifold.

Next, we characterize the lightning self-attention mechanisms that are singular in the neuromanifold, i.e., whose tangent space has dimension greater than Equation 7.
For that, we combine Theorem 3.2 with a computation of the critical points of the network parametrization map.
Firstly, we remark that the last point in Theorem 3.2 implies that $\overline{\varphi}$ is generically one-to-one onto its image $\mathbb{P}\mathcal{M}_{d,d^{\prime},a}$, i.e., it is *birational*.

###### Corollary 3.6 .

The map $\overline{\varphi}$ is birational onto its image $\mathbb{P}\mathcal{M}_{d,d^{\prime},a}$, with special fibers of cardinality $2$.

###### Proof.

This follows from the first and the last point of Theorem 3.2.
∎

Now, recall that a point is critical for a differentiable map if the differential at that point does not have maximal rank.

###### Lemma 3.7 .

A point $(A,V)=W$ is critical for the parametrization map $W\mapsto\varphi_{W}$ if and only if $A=k\otimes q$ and $V=h\otimes k$ for some $q,k\in\mathbb{R}^{d}$, $h\in\mathbb{R}^{d^{\prime}}$.

###### Proof.

See Appendix A.4.
∎

###### Corollary 3.8 .

$\varphi_{(A,V)}$ is singular in the neuromanifold if and only if $\textnormal{rk}(A)\textnormal{rk}(V)\leq 1$.

###### Proof.

See Appendix A.5.
∎

Finally, we compute the boundary points of $\mathcal{M}_{d,d^{\prime},a}$.
We will see that they are precisely the critical values of the parametrization map $W\mapsto\varphi_{W}$.
For that, we interpret the neuromanifold $\mathcal{M}_{d,d^{\prime},a}$ as a *Segre variety on linear forms*, as follows.
Given a vector space $\mathcal{V}$, we consider its dual space $\mathcal{V}^{\ast}$ and the outer product of vectors of linear forms:

$$ $\displaystyle\tilde{\sigma}\colon(\mathcal{V}^{\ast})^{d}\times(\mathcal{V}^{ \ast})^{d^{\prime}}$ $\displaystyle\to\mathrm{Sym}_{2}(\mathcal{V})^{d\times d^{\prime}},$ $\displaystyle(\alpha,\nu)$ $\displaystyle\mapsto(\alpha_{m}\cdot\nu_{n})_{1\leq m\leq d,1\leq n\leq d^{ \prime}}.$ $$

We can projectify this map analogously to $\overline{\varphi}$, obtaining birational morphism (if $d>1$ or $d^{\prime}>1$), but not an embedding.
If $a<d$, we need to restrict the first factor of the map $\tilde{\sigma}$ to tuples of linear forms that span a subspace of rank at most $a$. We denote this space by $(\mathcal{V}^{\ast})^{d}_{\leq a}$.

###### Proposition 3.9 .

The neuromanifold $\mathcal{M}_{d,d^{\prime},a}$ is linearly isomorphic
to the image of $\tilde{\sigma}$ restricted to $(\mathcal{V}^{\ast})^{d}_{\leq a}\times(\mathcal{V}^{\ast})^{d^{\prime}}$, where $\mathcal{V}=\mathbb{R}^{d}$.

###### Proof.

See Appendix A.6.
∎

In what follows, *relative boundary* refers to the set of points (in the Euclidean topology) in $\mathcal{M}_{d,d^{\prime},a}$ that are limit points of sequences in $\overline{\mathcal{M}}_{d,d^{\prime},a}\setminus\mathcal{M}_{d,d^{\prime},a}$, where $\overline{\mathcal{M}}_{d,d^{\prime},a}$ denotes the closure of the neuromanifold in the Zariski topology of its ambient space $\textnormal{Sym}_{3}\left(\mathbb{R}^{d\times t},\mathbb{R}^{d^{\prime}\times t
}\right)$.

###### Corollary 3.10 .

$\varphi_{W}$ is on the relative boundary of the neuromanifold if and only if $W$ it is critical for the parametrization map.

###### Proof.

See Appendix A.7.
∎

###### Theorem 3.4 .

###### Proof.

###### Corollary 3.5 .

###### Proof.

###### Corollary 3.6 .

###### Proof.

###### Lemma 3.7 .

###### Proof.

###### Corollary 3.8 .

###### Proof.

###### Proposition 3.9 .

###### Proof.

###### Corollary 3.10 .

###### Proof.

### 3.3 Deep Networks

Before discussing the fibers of the parametrization of a deep lightning self-attention network, we introduce yet another re-parametrization. To this end, consider dimension vectors $\mathbf{d}$, $\mathbf{a}$ and weights $\mathbf{W}=(\mathbf{A},\mathbf{V})$ of a network with $l$ layers. For $1\leq i\leq l$, define:

$$ $L_{i}=\prod_{l-i\leq j<l}V_{l-j}\hskip 50.00008ptM_{i}=L_{i-1}^{\top}A_{i}L_{i -1}.$ (9) $$

Moreover, set $M_{1}=A_{1}$. It follows immediately from Definition 2 that a deep self-attention network can be written in terms of $(\mathbf{M},L)$, where $\mathbf{M}=(M_{1},\ldots,M_{l})$ and $L=L_{l}$. Therefore, we introduce a further re-parametrization and write, with abuse of notation, $\varphi_{\mathbf{W}}$ for $\mathbf{W}=(\mathbf{M},L)$. We call the parameters in this parametrization *virtual weights*.

This parametrization has symmetries. Namely, invertible matrices $C_{i}\in\textnormal{GL}_{d_{i}}(\mathbb{R})$ for $1\leq i<l$, consider:

$$ $V^{\prime}_{i}=C_{i}V_{i}C_{i-1}^{-1}\hskip 50.00008ptA^{\prime}_{i+1}=C_{i}^{ -\top}A_{i+1}C_{i}^{-1}.$ (10) $$

In the above, we set $C_{0}$ and $C_{l}$ to the identity. Replacing $V_{i}$, $V_{i+1}$, $A_{i+1}$ with $V^{\prime}_{i}$, $V^{\prime}_{i+1}$, $A^{\prime}_{i+1}$ does not alter the virtual weights $\mathbf{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 generic $\mathbf{V}$ and with an assumption on the dimensions $d_{i}$, the above procedure completely characterizes the generic fibers of the map $(\mathbf{A},\mathbf{V})\mapsto(\mathbf{M},L)$.

###### Lemma 3.11 .

Suppose that for some $\delta\in\mathbb{N}$ it holds that $d_{i}=\delta$ for all $0<i<l$ and $d_{0},d_{l}\geq\delta$. Let $(\mathbf{M},L)$ be virtual weights (Equation 9) obtained obtained from both $(\mathbf{A},\mathbf{V})$ and $(\mathbf{A}^{\prime},\mathbf{V}^{\prime})$. If $\textnormal{rk}(L)=\delta$ then for every $i$ there exists a unique $C_{i}\in\textnormal{GL}_{d_{i}}(\mathbb{R})$ such that $V_{i}^{\prime}$ and $A_{i+1}^{\prime}$ are obtained from $V_{i}$ and $A_{i+1}$ via Equation 10.

###### Proof.

See Appendix A.8.
∎

The assumption on the dimensions $d_{i}$ in the above result can be practically interpreted as a ‘bottleneck’ architecture. Specifically, this architecture is equivalent to including a low-dimensional embedding and a high-dimensional unembedding layer, which is common in the literature, especially in an interpretability context (Elhage et al., ).

Before computing the fibers of the re-parametrization, it is convenient to rephrase the latter. The following result provides a recursive expression of a deep attention network in terms of $(\mathbf{M},L)$.

###### Lemma 3.12 .

For $X\in\mathbb{R}^{d_{0}\times t}$, define inductively:

$$ $\begin{cases}D_{0}=I\\ D_{i}=D_{i-1}X^{\top}M_{i}XD_{i-1}X^{\top}M_{i}^{\top}XD_{i-1}.\end{cases}$ (11) $$

Then:

$$ $\varphi_{\mathbf{W}}(X)=L_{l}X\left(\prod_{1\leq i\leq l}D_{l-i}X^{\top}M_{l-i +1}X\right).$ (12) $$

###### Proof.

See Appendix A.9.
∎

The above result can be rephrased without recursion. Given $X=(x_{i})_{1\leq i\leq t}$, since $XX^{\top}=\sum_{1\leq i\leq t}x_{i}x_{i}^{\top}$, Equation 12 states that $\varphi_{\mathbf{W}}(X)_{k}$ can be written for every $k$ as:

$$ $\sum_{k_{1},\ldots,k_{\tilde{l}}}L_{l}x_{k_{1}}\left(\prod_{1\leq j\leq\tilde{ l}-1}x_{k_{j}}^{\top}\widetilde{M}_{j}x_{k_{j+1}}\right)x_{\tilde{l}}^{\top}A_ {1}x_{k},$ (13) $$

where $\tilde{l}=(3^{l}-1)/2$ and $\widetilde{M}_{j}$ is defined as follows. Let $\alpha_{j}$ be 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 $\widetilde{M}_{j}$ equals to $M_{\alpha_{j}}$ if this digit is $1$, and to $M_{\alpha_{j}}^{\top}$ if it is $2$. We provide a diagram illustrating Equation 13 in Figure 2.

Figure: Figure 2: Diagrammatic illustration of Equation 13.

We now discuss the fibers, as anticipated. Firstly, similarly to the one-layer case, $\mathbf{W}=(\mathbf{M},L)$ can be rescaled without altering the corresponding function. More precisely, if $M^{\prime}_{i}=\lambda_{i}M_{i}$ and $L^{\prime}=\rho L$ for some $\lambda_{i},\rho\in\mathbb{R}\setminus\{0\}$, then $\varphi_{\mathbf{W}^{\prime}}=\varphi_{\mathbf{W}}$ if:

$$ $\prod_{1\leq i\leq l}(\lambda_{i})^{3^{l-i}}=\frac{1}{\rho}.$ (14) $$

Conversely, we now shot that rescaling characterizes the generic fibers. The following is the main result of this work.

###### Theorem 3.13 .

Let $\mathbf{W}=(\mathbf{M},L)$. Suppose that $t\geq 3$ and:

- •
For $1\leq i\leq l$, $\textnormal{rk}(M_{i})\geq 2$,
- •
For $1<i\leq l$, $M_{i}$ is not skew-symmetric(^3^33A square matrix $M$ is skew-symmetric if $M^{\top}=-M$.).

Then the fiber of $\varphi_{\mathbf{W}}$ consists of rescalings of $\mathbf{W}$.

###### Proof.

The proof is highly technical. Here, we provide a short summary; for the full proof, see Appendix A.10. Using unique factorization of polynomials, we first show that, up to rescalings, the fiber consists of $\mathbf{W}^{\prime}=(\mathbf{M}^{\prime},L)$, where $M_{i}^{\prime}=M_{i}+\Sigma_{i}$ and $\Sigma_{i}$ is a skew-symmetric matrix. We then proceed to show by that $\Sigma_{i}=0$ via an induction argument over $i$. The argument involves analyzing specific monomials where $\Sigma_{i}$ appears, and proving that most of them vanish due to the following symmetry of the diagram in Figure 2:

Figure: Figure 3: Diagrammatic illustration of the symmetry involved in the cancellation argument.

∎

The conditions of the above theorem are generic if $d_{i},a_{i}\geq 2$ for all $i$. Similarly to the one-layer case, this leads to the computation of the dimension of the neuromanifold.

###### Corollary 3.14 .

Suppose that $t\geq 3$, $a_{i}\geq 2$ for all $i$, and for some $\delta\geq 2$ it holds that $d_{i}=\delta$ for all $0<i<l$ and $d_{0},d_{l}\geq\delta$. The dimension of the neuromanifold is:

$$ $2\alpha_{1}d_{0}-\alpha_{1}^{2}+\delta(d_{0}+d_{l})-\delta^{2}-l+\sum_{1<i\leq l }(2\alpha_{i}\delta-\alpha_{i}^{2}),$ (15) $$

where $\alpha_{i}=\min\{a_{i},d_{i-1}\}$.

###### Proof.

See Appendix A.11.
∎

###### Lemma 3.11 .

###### Proof.

###### Lemma 3.12 .

###### Proof.

###### Theorem 3.13 .

###### Proof.

###### Corollary 3.14 .

###### Proof.

### 3.4 Traditional Self-Attention

In this section, we briefly consider the case of traditional self-attention mechanisms, where the attention weights are normalized. We compute the fibers in the single-layer case, and state a conjecture around the deep case.

In order to introduce the normalization, consider a map $\mathcal{S}\colon\mathbb{R}\rightarrow\mathbb{R}_{>0}$ that is injective and such that $\mathcal{S}(0)=1$. A typical choice is $\mathcal{S}(x)=e^{x/\tau}$, where $\tau\in\mathbb{R}_{>0}$ is the temperature hyperparameter. The traditional self-attention mechanism is then defined for $W=(A,V)$ and $X\in\mathbb{R}^{d\times t}$ as:

$$ $\varphi_{W}(X)=\displaystyle\left(\frac{1}{\zeta_{i}}\sum_{1\leq j\leq t} \mathcal{S}\left(x_{i}^{\top}Ax_{j}\right)\ Vx_{j}\right)_{1\leq i\leq t} \hskip 30.00005pt\zeta_{i}=\sum_{1\leq k\leq t}\mathcal{S}\left(x_{i}^{\top}Ax _{k}\right).$ (16) $$

Note that, for simplicity, we adhere to the convention of parametrizing via the attention matrix. Intuitively, in Equation 16 the attention weights $\mathcal{S}\left(x_{i}^{\top}Ax_{j}\right)$ are normalized to sum to $1$, forcing the model to ‘distribute’ its attention across the input components. The following result describes the fibers of the parametrization, analogously to Theorem 3.2.

###### Theorem 3.15 .

Suppose that $t\geq 2$. If $\varphi_{W}=0$, then $V=0$. Otherwise, the fiber of $\varphi_{W}$ is a singleton $\{W\}$.

###### Proof.

See Appendix A.12.
∎

Therefore, differently from lightning self-attention, in this case the parametrization is generically one-to-one.

We now consider the deep case. Note that even with normalization, deep attention networks can be reparametrized via $\mathbf{M}$ and $L$, as defined in Section 3.3. Therefore, the fibers of the parametrization will be unaffected by the transformations $C_{i}$ from Equation 10. However, this time no rescaling is possible. Therefore, we conjecture that the parametrization via $(\mathbf{M},L)$ is generically one-to-one.

###### Conjecture 3.16 .

For normalized deep self-attention networks, the generic fibers of the parametrization via $(\mathbf{M},L)$ are singletons.

In particular, suppose that for some $\delta$ it holds that $d_{i}=\delta$ for all $0<i<l$ and $d_{0},d_{l}\geq\delta$. Similarly to Corollary 3.14, the above conjecture implies that the dimension of the neuromanifold equals to:

$$ $2\alpha_{1}d_{0}-\alpha_{1}^{2}+\delta(d_{l}+d_{0})-\delta^{2}+\sum_{1<i\leq l }(2\alpha_{i}\delta-\alpha_{i}^{2}).$ (17) $$

Figure: Figure 4: Plot of estimated and expected dimension of the neuromanifold as $a$ varies.
Refer to caption: x1.png

We provide empirical evidence for Conjecture 3.16. To this end, we implement a deep attention network with softmax normalization (i.e., $\mathcal{S}(x)=e^{x}$), and estimate the dimension of its neuromanifold. The latter is a subtle problem since, differently from the linear 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=100$ input points in $\mathbb{R}^{d_{0}}$ from a normal distribution, and restricting $\varphi_{\mathbf{W}}$ to this finite space. As a result, the neuromanifold is embedded in a $(N\times t\times d_{l})$-dimensional vector space. Its dimension can then be computed as the rank of the Jacobian of the parametrization at a random parameter $\mathbf{W}$. For almost all parameters, this provides the correct result, with the only error coming from the discretization. Our Python code is available at a public repository(^4^44[https://github.com/giovanni-marchetti/NeuroDim](https://github.com/giovanni-marchetti/NeuroDim)). The results are visualized in Figure 4 for a deep attention network with $l=3$ layers, $t=3$, $d_{i}=4$ for all $i$, and $a_{i}=a$, with varying $a=1,\ldots,d$. The plot shows both the dimension estimated via the numerical approach (‘Estimated’), and the one computed via Equation 17 (‘Expected’). The two values coincide for all $a$, confirming Conjecture 3.16 empirically.

###### Theorem 3.15 .

###### Proof.

###### Conjecture 3.16 .

## 4 Conclusions and Future Work

In this work, we have analyzed the geometry of neuromanifolds of lightning self-attention networks. In particular, we have provided a 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 defined by attention networks. As such, it is subject to limitations and leaves several questions open. Specifically, the attention networks we consider are a simplified version of the ones that are deployed in practice, since we omit popular architectural variations (Brauwers & Frasincar, ), e.g., residual connections, positional encoding, masking, and multiple heads. Extending our algebro-geometric analysis to include these variations represents an interesting line for future investigation.

From a wider perspective, the research program of applying the tools offered by algebraic geometry to the field of deep learning remains open and worthy of exploration. Even further, going beyond the polynomial setting – e.g., addressing problems such as Conjecture 3.16 – is an even more general challenge that lies at the foundations of the theoretical understanding of deep learning.

## Acknowledgements

This work was partially supported by the Wallenberg AI, Autonomous Systems and Software Program (WASP) funded by the Knut and Alice Wallenberg Foundation.

## Appendix 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\geq\textnormal{rk}(Q)\geq\textnormal{rk}(A)=a$, we deduce $\textnormal{rk}(K)=a$, and similarly for $Q,K^{\prime},Q^{\prime}$. Therefore, without loss of generality, we can assume that $K=K_{1}\oplus K_{2}$, where $K_{1}\in\textnormal{GL}_{a}(\mathbb{R})$ and $K_{2}\in\mathbb{R}^{a\times(d-a)}$, and similarly for $K^{\prime},Q,Q^{\prime}$. If follows immediately that exists a unique $C\in\textnormal{GL}_{a}(\mathbb{R})$ such that $K^{\prime}_{1}=CK_{1}$. Since by hypothesis $K_{1}^{\top}Q_{1}=K_{1}^{\prime\top}Q_{1}^{\prime}$, we deduce $Q_{1}=C^{\top}Q_{1}^{\prime}$. But then $K_{2}^{\top}Q_{1}=K_{2}^{\prime\top}Q_{1}^{\prime}=K_{2}^{\prime\top}C^{-\top}
Q_{1}$, implying that $K_{2}^{\top}=K_{2}^{\prime\top}C^{-\top}$ and, similarly, $Q_{2}=C^{\top}Q_{2}^{\prime}$, as desired.
∎

###### Proof.

### A.2 Proof of Theorem 3.2

###### Proof.

Suppose that $\varphi_{W}(X)=\varphi_{W^{\prime}}(X)$ for all $X=(x_{i})_{1\leq i\leq t}\in\mathbb{R}^{d\times t}$, where $W^{\prime}=(A^{\prime},V^{\prime})$. By comparing the terms containing only $x_{1}$ in the polynomial identity $\varphi_{W}(X)_{1}=\varphi_{W^{\prime}}(X)_{1}$, we obtain:

$$ $Vx_{1}\ x_{1}^{\top}Ax_{1}=V^{\prime}x_{1}\ x_{1}^{\top}A^{\prime}x_{1}.$ (18) $$

The above equation is a product of a quadratic form and a linear multivariate form. If $\textnormal{rk}(A)>1$, the quadratic form is irreducible, and from the unique factorization property of polynomials it follows that, up to multiplicative scalars, $V$ coincides with $V^{\prime}$, and the quadratic form associated to $A$ coincides with the one associated to $A^{\prime}$. The same holds if $\textnormal{rk}(V)>1$, since the sides of the above equation represent cubical polynomials that share a quadratic factor, but whose remaining linear factors are independent. In order to prove that $A$ and $A^{\prime}$ coincide (up to scaling), consider the terms linear in $x_{1}$ and quadratic in $x_{2}$ in the identity $\varphi_{W}(X)_{1}=\varphi_{W^{\prime}}(X)_{1}$. This leads to:

$$ $\bcancel{Vx_{2}}\ x_{2}^{\top}Ax_{1}=\bcancel{V^{\prime}x_{2}}\ x_{2}^{\top}A^ {\prime}x_{1}.$ (19) $$

It follows that the bilinear forms associated to $A$ and $A^{\prime}$ coincide (up to scaling), as desired for the third case of the claim.

If instead $\textnormal{rk}(A)=\textnormal{rk}(V)=1$, then it is possible to factor $A=k\otimes q$ and $V=h\otimes v$ for some vectors $q,k,v\in\mathbb{R}^{d}\setminus\{0\}$, $h\in\mathbb{R}^{d^{\prime}}\setminus\{0\}$. Exchanging the role of $k$ and $v$ does not alter Equation 1, which provides the first case of the claim. Finally, since the ring of polynomials is an integral domain, if $\varphi_{W}(X)=0$ then from Equation 18 it follows that either $A=0$ or $V=0$, as desired for the second case of the claim.
∎

###### Proof.

### A.3 Proof of Corollary 3.5

###### Proof.

A continuous map between a compact and a Hausdorff space is closed, and in particular has closed image. These properties are satisfied by $\overline{\varphi}$, where both the domain and the codomain are equipped with the Euclidean topology over projective spaces. It follows that $\mathbb{P}\mathcal{M}_{d,d^{\prime},a}$ is Euclidean closed in $\mathbb{P}\textnormal{Sym}_{3}\left(\mathbb{R}^{d\times t},\mathbb{R}^{d^{
\prime}\times t}\right)$. 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.
∎

###### Proof.

### A.4 Proof of Lemma 3.7

###### Proof.

Given $W=(A,V)$, by computing the partial derivatives of $\varphi_{\bullet}$, we see that the differential sends a tangent vector $(\dot{A},\dot{V})$ to $\dot{V}XX^{\top}AX+VXX^{\top}\dot{A}X$. The latter is interpreted as an element of the vector space $\textnormal{Sym}_{3}\left(\mathbb{R}^{d\times t},\mathbb{R}^{d^{\prime}\times t
}\right)$, 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 $(\dot{A},\dot{V})$ that, for all $X$, satisfy

$$ $\dot{V}XX^{\top}AX=-VXX^{\top}\dot{A}X.$ (20) $$

If $V=0$, the differential vanishes whenever $\dot{V}=0$; and similarly for $A$. This provides the first case of the claim.
If $\textnormal{rk}(A)\textnormal{rk}(V)\geq 2$, Theorem 3.2 implies that the kernel consists of $(\dot{V}=-\lambda V,\dot{A}=\lambda A)$ for $\lambda\in\mathbb{R}$, and $W$ is therefore not critical.
Finally, suppose that $\textnormal{rk}(A)=\textnormal{rk}(V)=1$, and write $A=k\otimes q$ and $V=h\otimes v$ for some $q,k,v\in\mathbb{R}^{d},h\in\mathbb{R}^{d^{\prime}}$.
In order to obtain a larger kernel, by Theorem 3.2, it is necessary that $v=k$ (up to scaling). In that case, for every $\dot{A}$ of the form $\dot{A}=p\otimes q$, we have that $\dot{V}=-h\otimes p$ is a solution to Equation 20, implying that $W=(A,V)$ is critical.
∎

###### Proof.

### A.5 Proof of Corollary 3.8

###### Proof.

The morphism $\overline{\varphi}$ is finite and birational by Corollary 3.6.
For such a map, a standard fact from algebraic geometry (Kohn et al., , Lemma 3.2) says that a point $\varphi_{W}$ is singular in the projective neuromanifold $\mathbb{P}\mathcal{M}_{d,d^{\prime},a}$ if and only if either its fiber under $\overline{\varphi}$ has cardinality $\geq 2$ or $W$ is critical for $\overline{\varphi}$. By Lemma 3.7 and Theorem 3.2, a point with $\textnormal{rk}(A)\textnormal{rk}(V)\leq 1$ is either critical or has a fiber of cardinality $2$, and $\varphi_{(A,V)}$ is therefore singular in $\mathbb{P}\mathcal{M}_{d,d^{\prime},a}$. Since $\mathcal{M}_{d,d^{\prime},a}$ is the affine cone of $\mathbb{P}\mathcal{M}_{d,d^{\prime},a}$, a point is singular in the projective neuromanifold if and only if the corresponding line is singular in the neuromanifold, from which the claim follows.
∎

###### Proof.

### A.6 Proof of Proposition 3.9

###### Proof.

We denote by $\alpha_{m}$ the linear form that takes the inner product with the $m$-th column of $A$ and by $\nu_{n}$ the inner product with the $n$-th row of $V$.
We have seen in the proof of Theorem 3.2 that every $\varphi_{W}\in\mathcal{M}_{d,d^{\prime},a}$ is uniquely determined by its term that is linear in $x_{1}$ and quadratic in $x_{2}$,
which is

$$ $\displaystyle x_{2}^{\top}Ax_{1}Vx_{2}=\left(\sum_{m=1}^{d}x_{1,m}\,\alpha_{m} (x_{2})\,\nu_{n}(x_{2})\right)_{1\leq n\leq d^{\prime}}.$ $$

Since the $x_{1,_{m}}$ are formal variables, this expression is the collection of all products of linear forms $\alpha_{m}\cdot\nu_{n}$ for $1\leq m\leq d$ and $1\leq n\leq d^{\prime}$, which shows the claim.
∎

###### Proof.

### A.7 Proof of Corollary 3.10

###### Proof.

For simplicity, we write $\mathcal{M}$ for $\mathcal{M}_{d,d^{\prime},a}$, and start by computing its Zariski closure $\overline{\mathcal{M}}$.
The Zariski closure of a subset $S$ of $\mathbb{R}^{N}$ consists precisely of the real points in the complex Zariski closure of $S$ viewed as a subset of $\mathbb{C}^{N}$.
We compute the complex Zariski closure of $\mathcal{M}$ by considering the parametrization map $\overline{\varphi}$ in (8) over $\mathbb{C}$ instead of over $\mathbb{R}$.
That map $\overline{\varphi}_{\mathbb{C}}$ is also a well-defined morphism between projective spaces.
The main theorem on projective varieties (Hartshorne, , II, §4, Theorem 4.9) states that the image $\mathbb{P}\mathcal{M}^{\mathbb{C}}$ of $\overline{\varphi}_{\mathbb{C}}$ is Zariski closed.
Hence, the real Zariski closure $\overline{\mathcal{M}}$ consists precisely of the real points $\varphi_{(A,V)}$ with $A\in\mathbb{C}^{d\times d},V\in\mathbb{C}^{d\times d^{\prime}}$.

Next, we compute the points in the complement $\overline{\mathcal{M}}\setminus\mathcal{M}$.
Using the identification in Proposition 3.9, points in that complement look like real matrices with entries $\alpha_{m}\cdot\nu_{n}$, where $\alpha_{m},\nu_{n}$ are complex linear forms, and no real parametrization exists. So one of the $\alpha_{m}$ is non-real. Since each of the $\alpha_{m}\cdot\nu_{n}$ is real, every $\nu_{n}$ has to be the complex conjugate $\overline{\alpha_{m}}$ times a real scalar. Analogously, all $\alpha_{r}$ have to be $\overline{\nu_{n}}$ times a real scalar, i.e., they need to coincide with $\alpha_{m}$ up to scaling. This shows that
$\overline{\mathcal{M}}\setminus\mathcal{M}$ consists of $\varphi_{(A,V)}$ with $A=k\otimes q$ and $V=h\otimes\overline{k}$ for some $q\in\mathbb{R}^{d}\setminus\{0\},h\in\mathbb{R}^{d^{\prime}}\setminus\{0\},k
\in\mathbb{C}^{d}\setminus\mathbb{R}^{d}$. In the limit, sequences of $W=(A,V)$ in this form can become real if either one among $A$ and $V$ vanish, or the complex conjugated pair $k,\overline{k}$ becomes the same real vector $k=\overline{k}$.
∎

###### Proof.

### A.8 Proof of Lemma 3.11

###### Proof.

Note that $V_{i}$ is invertible for every $1<i<l$ by the rank hypothesis. Moreover, since $L=\prod_{0\leq j<l}V_{l-j}=\prod_{0\leq j<l}V^{\prime}_{l-j}$, it is possible to apply Lemma 3.1 by splitting the latter factorization of $L$ at the first and last layer. This way, we obtain matrices $C_{i}\in\textnormal{GL}_{\delta}(\mathbb{R})$ for $1\leq i\leq l$ such that $V^{\prime}_{i}=C_{i}V_{i}C_{i-1}^{-1}$. But then we have:

$$ $L_{i-1}^{\prime\top}A_{i}^{\prime}L_{i-1}^{\prime}=L_{i-1}^{\top}C_{i-1}^{\top }A_{i}^{\prime}C_{i-1}L_{i-1}=L_{i-1}^{\top}A_{i}L_{i-1}.$ (21) $$

Since $L_{i}$ is surjective by the rank hypothesis, we conclude that $C_{i-1}^{\top}A_{i}^{\prime}C_{i-1}=A_{i}$.
∎

###### Proof.

### A.9 Proof of Lemma 3.12

###### Proof.

For $0\leq i\leq l$, denote by $X_{i}$ the output of the $i$-th layer of the network. In other words, $X_{0}=X$ and $X_{i}=\varphi_{W_{i}}(X_{i-1})=V_{i}X_{i-1}X_{i-1}^{\top}A_{i}X_{i-1}$ for $i>0$. We wish to prove that for all $i$:

$$ $X_{i}=L_{i}X\prod_{1\leq j\leq i}\left(D_{i-j}X^{\top}M_{i-j+1}X\right).$ (22) $$

Denote by $Y_{i}$ the right-hand side of Equation 22. We have:

$$ $\displaystyle Y_{i}=$ $\displaystyle\underbrace{L_{i}}_{V_{i}L_{i-1}}XD_{i-1}X^{\top}\underbrace{M_{i }}_{L_{i-1}^{\top}A_{i}L_{i-1}}X\prod_{1\leq j<i}D_{i-1-j}X^{\top}M_{i-j}X$ (23) $\displaystyle=$ $\displaystyle V_{i}\ L_{i-1}XD_{i-1}\left(L_{i-1}X\right)^{\top}A_{i} \underbrace{L_{i-1}X\prod_{1\leq j<i}D_{i-1-j}X^{\top}M_{i-j}X}_{Y_{i-1}}.$ (24) $$

Therefore, in order to show that $Y_{i}$ satisfies the same recurrence relation as $X_{i}$, we need to prove that $L_{i}XD_{i}\left(L_{i}X\right)^{\top}=Y_{i}Y_{i}^{\top}$ for all $i$, which in turn reduces to show that $D_{i}=Z_{i}Z_{i}^{\top}$, where:

$$ $Z_{i}=\prod_{1\leq j\leq i}\left(D_{i-j}X^{\top}M_{i-j+1}X\right).$ (25) $$

To this end, since $Z_{i}=D_{i-1}X^{\top}M_{i}XZ_{i-1}$, we have:

$$ $\displaystyle Z_{i}Z_{i}^{\top}=D_{i-1}X^{\top}M_{i}XZ_{i-1}Z_{i-1}^{\top}X^{ \top}M_{i}^{\top}XD_{i-1}^{\top}.$ (26) $$

If we assume inductively that $Z_{i-1}Z_{i-1}^{\top}=D_{i-1}$, and using the fact that $D_{i}$ is a symmetric matrix (i.e., $D_{i-1}^{\top}=D_{i-1}$), the above recurrence relation coincides with the one defining $D_{i}$ in Equation 12. This implies that $Z_{i}Z_{i}^{\top}=D_{i}$, as desired.
∎

###### Proof.

### A.10 Proof of Theorem 3.13

###### Proof.

Pick weights $\mathbf{W}$, $\mathbf{W}^{\prime}$ such that $\varphi_{\mathbf{W}}(X)=\varphi_{\mathbf{W}^{\prime}}(X)$ for all $X=(x_{i})_{1\leq i\leq t}\in\mathbb{R}^{d_{0}\times t}$. Our proof strategy will consist in considering the polynomial identity

$$ $\varphi_{\mathbf{W}}(X)_{1}=\varphi_{\mathbf{W}^{\prime}}(X)_{1}$ (27) $$

and comparing monomial terms arising from Equation 13 with specific degrees in the $x_{i}$’s.

To this end, Equation 13 implies that all such terms contain $x_{1}$. Even further, the unique term in $\varphi_{\mathbf{W}}(X)_{1}$ that is linear in $x_{1}$ and of degree $3^{l}-1$ in $x_{2}$
can be written as:

$$ $L_{l}x_{2}\left(\prod_{2\leq i\leq l}\left(x_{2}^{\top}M_{i}x_{2}\right)^{3^{l -i}}\right)\left(x_{2}^{\top}A_{1}x_{2}\right)^{3^{l-1}-1}x_{2}^{\top}A_{1}x_{ 1}.$ (28) $$

Put simply, the above expression is a product of a linear form in $x_{2}$, several quadratic forms in $x_{2}$, and a bilinear form in $x_{1}$ and $x_{2}$. Moreover, the quadratic forms are coprime by hypothesis since $\textnormal{rk}(M_{i})>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 a multiplicative scalar, the following hold:

- •
$L_{l}$ coincides with $L_{l}^{\prime}$,
- •
$A_{1}$ coincides with $A_{1}^{\prime}$,
- •
The quadratic forms associated to $M_{i}$ and $M_{i}^{\prime}$ coincide for all $2\leq i\leq l$.

The last condition above implies that for every $2\leq i\leq l$ there exists a skew-symmetric matrix $\Sigma_{i}$ such that $M_{i}^{\prime}$ coincides with $M_{i}+\Sigma_{i}$ up to a multiplicative scalar. Since Equation 13 is multi-linear in each occurrence of $M_{j}$, by substituting $M^{\prime}_{j}=M_{j}+\Sigma_{j}$ in $\varphi_{\mathbf{W}^{\prime}}(X)_{1}$ for all $j$, we obtain a sum of expressions in the form of Equation 13, but where an arbitrary number of occurrences of $M_{j}$ has been replaced by $\Sigma_{j}$. The only expression with no replacement coincides with $\varphi_{\mathbf{W}}(X)_{1}$. Therefore, Equation 27 reduces to a vanishing sum of copies of Equation 13, where at least one $M_{j}$ has been replaced with $\Sigma_{j}$.

In order to conclude, we wish to show that $\Sigma_{i}=0$ for $i\geq 2$. We will proceed by induction on $i$. Specifically, given $i$ and assuming that $\Sigma_{j}=0$ for $j<i$, we will show that $\Sigma_{i}=0$. To this end, consider the monomial terms in $\varphi_{\mathbf{W}^{\prime}}(X)_{1}$ of degree $1$ in $x_{1}$, of degree $3^{i-1}-1$ in $x_{3}$, and of remaining degree in $x_{2}$. This is possible since we assume $t\geq 3$. From the discussion above, it follows that Equation 27 reduces to a vanishing sum of copies of Equation 13, where an arbitrary (non-zero) number of occurrences of $\widetilde{M}_{j}=M_{\alpha_{j}}$, with $\alpha_{j}\geq i$, has been replaced with $\Sigma_{\alpha_{j}}$ (and similarly with transposes), and where $k_{j}\in\{2,3\}$ for all $1\leq j\leq\tilde{l}=(3^{l}-1)/2$.

We will now argue that several terms cancel due to the symmetries of Equation 13. Namely, consider a monomial term with some multi-index $k_{\bullet}$ containing a factor of the form $x_{k_{j}}^{\top}\Sigma_{\alpha_{j}}x_{k_{j+1}}$ for some $j$, with $\alpha_{j}>i$. Since $x_{3}$ cannot appear more than $3^{i-1}-1$ times in the monomial, there exists a $\underline{j}\leq(3^{i-1}-1)/2$ such that $k_{j-\underline{j}}=k_{j+\underline{j}+1}=2$. Moreover, due to the inductive hypothesis, $\widetilde{M}_{j^{\prime}}$ coincides with $M_{\alpha_{j^{\prime}}}$ or with its transpose for all $j-\underline{j}\leq j^{\prime}\leq j+\underline{j}$ since $\alpha_{j^{\prime}}<i$. Consider the multi-index $k^{\prime}_{\bullet}$ such that $k^{\prime}_{j-j^{\prime}}=k_{j+j^{\prime}+1}$ for all $-\underline{j}\leq j^{\prime}\leq\underline{j}$, and $k^{\prime}_{j^{\prime}}=k_{j^{\prime}}$ for all the other $j^{\prime}$. Intuitively, $k^{\prime}_{\bullet}$ ‘reflects’ $k_{\bullet}$ locally around $j$ – see Figure 3 for an illustration. By construction, the monomial corresponding to $k^{\prime}_{\bullet}$ is identical to the one corresponding to $k_{\bullet}$, except for the term

$$ $x_{k^{\prime}_{j}}^{\top}\Sigma_{\alpha_{j}}x_{k^{\prime}_{j+1}}=x_{k_{j+1}}^{ \top}\Sigma_{\alpha_{j}}x_{k_{j}}=-x_{k_{j}}^{\top}\Sigma_{\alpha_{j}}x_{k_{j+ 1}}.$ (29) $$

Therefore, these two monomials cancel each other out.

After the above cancellation argument, we are left with monomials containing occurrences of $\Sigma_{\alpha_{j}}$ only for $\alpha_{j}=i$. Actually, the cancellation argument still applies verbatim to these cases, except for when $j=(3^{l}-3^{i-1})/2$ and $k_{j-j^{\prime}}\not=k_{j+j^{\prime}+1}$ for all $j^{\prime}\leq(3^{i-1}-1)/2$. In this case, due to the presence of $x_{1}$, the monomials corresponding to $k_{\bullet}$ and $k^{\prime}_{\bullet}$ are not opposite, but contain different factors given by $x_{3}^{\top}A_{1}x_{2}x_{2}^{\top}A_{1}x_{1}$ and $-x_{3}^{\top}A_{1}x_{1}\ x_{2}^{\top}A_{1}x_{2}$, respectively. At the end of the day, Equation 27 reduces to an expression of the form:

$$ $x_{3}^{\top}\Sigma_{i}x_{2}\ \left(x_{3}^{\top}A_{1}x_{2}x_{2}^{\top}A_{1}x_{1 }-x_{3}^{\top}A_{1}x_{1}\ x_{2}^{\top}A_{1}x_{2}\right)\ \sum_{h_{\bullet}}p(( x_{h_{j}})_{j})=0,$ (30) $$

where $h_{\bullet}$ 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_{j}$’s, and a multi-variate linear factor associated to $L_{l}$. The right-most factor in Equation 30 is not the zero polynomial (in the variables $x_{2}$ and $x_{3}$) since, by imposing the condition $x_{2}=x_{3}$, it becomes (up to a multiplicative scalar) a product of quadratic forms associated to the $M_{j}$’s. The latter are non-vanishing since $M_{j}$ is not skew-symmetric for all $j$ by hypothesis.

We wish to prove that the second factor in Equation 30 is non-vanishing as well. To this end, note that the coefficients of that factor, seen as a polynomial in $x_{1},x_{2},x_{3}$, are of the form

$$ $(A_{1})_{\alpha,\beta}(A_{1})_{\gamma,\delta}+(A_{1})_{\alpha,\gamma}(A_{1})_{ \beta,\delta}-(A_{1})_{\alpha,\delta}(A_{1})_{\beta,\gamma}-(A_{1})_{\alpha, \delta}(A_{1})_{\gamma,\beta}$ (31) $$

for $1\leq\alpha,\beta,\gamma,\delta\leq d_{0}$. Suppose by contradiction that the above expression vanishes for all $\alpha,\beta,\gamma,\delta$. When $\beta=\gamma$, Equation 31 coincides with (twice) the determinant of an arbitrary $2\times 2$ minor of $A_{1}$ intersecting the diagonal. If $(A_{1})_{\beta,\beta}\not=0$ for some $\beta$, then from the Kronecker’s bordered matrix theorem it follows that $A_{1}$ has rank $1$, contradicting the hypothesis $\textnormal{rk}(A_{1})\geq 2$. We conclude that the diagonal entries of $A_{1}$ are zero, and, when $\beta=\gamma$, Equation 31 reduces to $(A_{1})_{\alpha,\beta}(A_{1})_{\beta,\delta}=0$. Therefore, if $(A_{1})_{\alpha,\beta}\not=0$ for some $\alpha\not=\beta$, then the $\beta$-th row of $A_{1}$ must vanish. But then Equation 31 reduces (up to sign) to the determinant of an arbitrary $2\times 2$ minor intersecting $(\alpha,\beta)$, 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 $\Sigma_{i}=0$, as desired.
∎

###### Proof.

### A.11 Proof of Corollary 3.14

###### Proof.

Recall that the dimension of the determinantal variety of $d_{i-1}\times d_{i-1}$ matrices of rank at most $a_{i}$ is $2\alpha_{i}d_{i-1}-\alpha_{i}^{2}$, and therefore the space of parameters $(\mathbf{A},\mathbf{V})$ has dimension

$$ $2\alpha_{1}d_{0}-\alpha_{1}^{2}+\delta(d_{0}+d_{l})+(l-2)\delta^{2}+\sum_{1<i \leq l}(2\alpha_{i}\delta-\alpha_{i}^{2}).$ (32) $$

Moreover, by Lemma 3.11 the generic fibers of the re-parametrization $(\mathbf{A},\mathbf{V})\mapsto(\mathbf{M},L)$ have dimension $(l-1)\delta^{2}$, while by Theorem 3.13 the generic fibers with respect to $(\mathbf{M},L)$ have dimension $l$ (recall the constraint on rescaling given by Equation 14). Since the dimension of the image of a map coincides with the co-dimension of the generic fibers, the result follows.
∎

###### Proof.

### A.12 Proof of Theorem 3.15

###### Proof.

Suppose that $\varphi_{W}(X)=\varphi_{W^{\prime}}(X)$ for all $X=(x_{i})_{1\leq i\leq t}\in\mathbb{R}^{d\times t}$, where $W^{\prime}=(A^{\prime},V^{\prime})$. We assume that $V\not=0$, since the case $\varphi_{W}=0$ follows immediately. Due to normalization, if $x_{i}=x_{j}$ for all $i,j$, then $\varphi_{W}(X)_{1}=Vx_{1}=V^{\prime}x_{1}=\varphi_{W^{\prime}}(X)_{1}$, implying $V=V^{\prime}$. If instead $x_{i}=0$ for $i>1$, since $\mathcal{S}(0)=1$, we obtain:

$$ $\frac{\mathcal{S}\left(x_{1}^{\top}Ax_{1}\right)}{\mathcal{S}\left(x_{1}^{\top }Ax_{1}\right)+t-1}\bcancel{Vx_{1}}=\frac{\mathcal{S}\left(x_{1}^{\top}A^{ \prime}x_{1}\right)}{\mathcal{S}\left(x_{1}^{\top}A^{\prime}x_{1}\right)+t-1} \bcancel{V^{\prime}x_{1}}$ (33) $$

Since $\mathcal{S}$ is injective, we deduce that the quadratic form associated to $A$ coincides with the one associated to $A^{\prime}$. In order to prove that $A=A^{\prime}$, suppose that $x_{i}=x_{j}$ for $i,j\geq 2$, obtaining:

$$ $\frac{\mathcal{S}\left(x_{1}^{\top}Ax_{1}\right)Vx_{1}+\mathcal{S}\left(x_{1}^ {\top}Ax_{2}\right)Vx_{2}}{\mathcal{S}\left(x_{1}^{\top}Ax_{1}\right)+(t-1) \mathcal{S}\left(x_{1}^{\top}Ax_{2}\right)}=\frac{\mathcal{S}\left(x_{1}^{\top }A^{\prime}x_{1}\right)V^{\prime}x_{1}+\mathcal{S}\left(x_{1}^{\top}A^{\prime} x_{2}\right)V^{\prime}x_{2}}{\mathcal{S}\left(x_{1}^{\top}A^{\prime}x_{1} \right)+(t-1)\mathcal{S}\left(x_{1}^{\top}A^{\prime}x_{2}\right)}.$ (34) $$

After elementary algebraic manipulations using the fact that $x_{1}^{\top}Ax_{1}=x_{1}^{\top}A^{\prime}x_{1}$ and $V=V^{\prime}$, the above equation reduces to:

$$ $(t-1)\left(\mathcal{S}\left(x_{1}^{\top}A^{\prime}x_{2}\right)-\mathcal{S} \left(x_{1}^{\top}Ax_{2}\right)\right)Vx_{1}=\left(\mathcal{S}\left(x_{1}^{ \top}A^{\prime}x_{2}\right)-\mathcal{S}\left(x_{1}^{\top}Ax_{2}\right)\right) Vx_{2}.$ (35) $$

Since $V\not=0$, $Vx_{1}\not=Vx_{2}$ for generic $x_{1},x_{2}$. Therefore, $\mathcal{S}\left(x_{1}^{\top}A^{\prime}x_{2}\right)-\mathcal{S}\left(x_{1}^{
\top}Ax_{2}\right)$ must vanish generically, implying $A=A^{\prime}$ due to the injectivity of $\mathcal{S}$, as desired.
∎

###### Proof.

## Peer Reviews

confidence: 2
contribution: 2
initial_rating: 6
presentation: 3
questions: Q1 & Q2. What is the authors’ position on W1 and W2?

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

Q4. Could the authors elaborate on the implications of the non-smoothness of the $\mathcal{M}_{d, d’, a}$ manifold from an optimization standpoint?
soundness: 3
strengths: S1. The paper exposition is quite clear, striking a balance between the lingo of algebraic geometry and the context of the ICLR audience.

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

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

S4. The authors introduce a “virtual weights” parametrization which enables an elegant analysis of deep lightning self-attention networks.

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

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

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

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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? 

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

I 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?
soundness: 4
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.

As far as I can tell (although I am no expert in algebraic geometry), the claims are clearly stated and the mathematics are correct.
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:
* 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.
* Theorem 3.4 shows that the neuromanifold is closed in Euclidean space and that its boundaries are defined by networks with rank-1 matrices.
* 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.
weaknesses: There are several key gaps between the models studied in this paper and practical models:
* The lack of non-linearities in the attention unit (in particular, the lack of softmax).
* The lack of residual connections between layers.
* The lack of element-wise multi-layer perceptrons between layers.
As 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.

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

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 2
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?

2. What purpose do figures 2 and 3 serve?

3. What do you think are the main implications of the main theorem?
soundness: 4
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).
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.
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). 



**A list of some of the main weaknesses:**

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

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

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

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



**More specific points about the weaknesses:**

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

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

c. The purpose of the visualizations in figures 2 and 3 is unclear.

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

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

f. The presentation can be improved to be more accessible to ML researchers.

---

confidence: 2
contribution: 2
initial_rating: 3
presentation: 2
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.
soundness: 3
strengths: Analysis of linear self-attention networks with tools from algebraic geometry
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.
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.

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

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

Therefore, 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.

Meta-review: This paper studies the geometry of single-layer and multi-layer linear-self-attention networks via the neuromanifold. It reveals interesting properties of the geometry of lightning attention models with a rigorous characterization of the resulting functional space. Besides, the paper exposition is clear, striking a balance between the lingo of algebraic geometry and the context of the ICLR audience. I recommend to accept.

Decision: Accept (Poster)