Geometry of Lightning Self-Attention: Identifiability and Dimension

We thank the reviewer for the constructive feedback.

We acknowledge that our paper, in the submitted form, is lacking motivations and discussion, especially from a practical perspective. To this end, we have re-uploaded a new version of the manuscript, where we have elaborated on several points including the importance of dimension and its relation to sample complexity, the role of fibers and singularities in the learning dynamics, a more accessible explanation of our main result (Theorem 3.17), and specific ideas for future work. All these points are also summarized in the collective comment we posted above.

We additionally wish to comment on the raised weakness raised about the relation between number of parameters and dimension of the neuromanifold. We argue that the difference between the two can potentially be quite large, especially for attention mechanisms, since, differently from other models, their parametrization exhibits plenty of symmetries – i.e., symmetries between queries and keys of the same layer, and between attention weights and values of successive layers. This applies to both the lightning case and, assuming Conjecture 3.19, to the traditional one. Therefore, the dimension of the neuromanifold can differ significantly from the number of parameters. To illustrate this, there are instances where the dimension of queries/keys equal to the embedding dimension (see the original paper [1]), i.e., $\alpha_i = \delta$ in Corollary 3.8. Assuming $d_0 = d_l = \delta$ as well, the number of parameters is $\frac{3}{2}d^2 l$, while according to our formula the dimension is $d^2( l + 1) - l$. Asymptotically, their ratio is $\frac{3}{2}$, i.e., the number of parameters is 50 % more than the dimension. We have added this discussion in the re-uploaded version of the manuscript in lines 369-377.

We hope that these discussions help to clarify the motivation and significance of our results.

[1] Vaswani et al., Attention is all you need, 2017.

We thank the reviewer for the feedback and the appreciation. Below, we address the questions raised.

Q1 and Q2: We acknowledge that it is important to expand on the motivation and the significance of computing the dimension of neuromanifolds and, in particular, its relation to sample complexity, which is a primary motivation behind our main result. To this end, we have re-uploaded a new version of the manuscript with a discussion on the relation between dimension and sample complexity – see the section Introduction and Related Work. This discussion is also reported in the collective comment we posted above. Moreover, after Corollary 3.8, we have discussed the implications of our formula for the dimension and, in particular, how the latter compares to the number of parameters, which is commonly used as a measure of expressivity and/or sample complexity.

Q3: As mentioned by the reviewer, we perform multiple trials when estimating the Jacobian, and then take the maximum. We acknowledge that this might sound superfluous, since taking the rank of the Jacobian at an arbitrary parameter yields the correct result almost everywhere w.r.t. such parameter, as explained in lines 436-438 in our paper. However, in practice this does not hold due to the discrete nature of computations – i.e., the floating-point approximation of real numbers – and due to numerical errors when calculating the rank, leading to an underestimate of the desired generic rank. In order to amend for this, we perform a few estimates at random parameters, and take the maximum over the ranks obtained. The resulting estimation is more robust to numerical errors. We did not mention this in the text since we believe that it constitutes a small implementation detail, that also depends on the programming framework – more specifically, on its numerical precision – and that is often unnecessary.

Q4: We wish to additionally comment on the relevance of singularities of neuromanifolds and, in particular, on their relation to the training dynamics. Singularities of neuromanifolds are a central focus in Information Geometry, and specifically in Singular Learning Theory [1]. According to the latter, singularities of neuromanifolds play a central role in deep learning, since the function learned by a neural network (via gradient descent) often corresponds to a singular point of the neuromanifolds [2]. In other words, singularities often attract the learning dynamics, resulting in a form of `implicit bias’ associated to the neural architecture. According to our result for a single-layer lightning attention network (Corollary 3.8), singularities of the neuromanifold arise exactly when both the attention matrix $A$ or the value matrix $V$ have rank $\leq 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 [3], and our result might suggest a mathematical explanation to this phenomenon, at least in the single-layer and lightning case. Since we believe that this discussion is important to include, we have re-uploaded a new version of the manuscript that includes it in the section Single-Layer Geometry.

[1] Watanabe, Algebraic geometry and statistical learning theory, 2009.

[2] Amari et al., Singularities affect dynamics of learning in neuromanifolds, 2006.

[3] Huh et al., The low-rank simplicity bias in deep networks, 2021.

We thank the reviewer for the several comments and suggestions, resulting in constructive feedback. We have re-uploaded an updated version of the manuscript, where we address several weaknesses raised. As suggested, we have moved several technical parts (and a figure) to the appendix, and used the space to motivate and explain the results. In particular, we have elaborated on the importance of dimension and of the fibers in the section Introduction and Related Work, and on the role of singularities in section Single-Layer Identifiability. We have discussed practical implications of our main result (see Corollary 3.8), and extended the section Conclusions and Future Work. Lastly, as suggested in weakness d, we have mentioned that the reparametrization is crucial for stating our main result. Below, we address the questions raised.

Question 1: We agree with the reviewer on the importance of elaborating on possible future extensions. While in our paper we considered a `vanilla’ self-attention mechanism, we believe that it might be possible to extend the theory to incorporate the variations that are used in practice. Two such variations, which are ubiquitous in contemporary Transformers, are skip connections and multiple attention heads. With both these additions, the lightning self-attention mechanism is still polynomial, which enables to potentially apply similar techniques as in our paper to their analysis. Note that skip connections make the model non-homogeneous, which breaks the scaling symmetry in the parameterization. We conjecture that in this case the parameterization via $(**M**, L)$ is generically one-to-one, similarly to the traditional case (Conjecture 3.16), 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 [1]. Therefore, these two variations give rise to interesting phenomena in terms of symmetries of the parameterization, and we believe that these directions are worthy of exploration. We have incorporated this discussion in the section Conclusions and Future Work of the re-uploaded manuscript.

Question 2: Our intent in Figures 2 and 3 is to illustrate via a diagram the statement of Lemma 3.12 and the proof of Theorem 3.13, respectively. The statement of the Lemma involves a subtle recursion, which, in our opinion, is difficult to interpret from equations. The diagram in Figure 2 visualizes this recursion by unrolling it as a tree. The main step in the proof of Theorem 3.13 exploits a symmetry of this tree, which is illustrated intuitively in Figure 3. While we believe it is important to convey an intuition for Theorem 3.13, we agree that the space in the main body can be used to better build motivation and convey significance; we have moved Figure 3 to the appendix, as suggested.

Question 3: We agree with the reviewer that it is crucial to elaborate on the implications of our main theorem – i.e., the computation of the fibers and, as a consequence, of the dimension of the neuromanifold. To this end, we have re-uploaded a new version of the manuscript with a discussion on the importance of dimension, especially its relation to sample complexity in machine learning, which has practical implications – see the section Introduction and Related Work. This discussion is also reported in the collective comment we posted above. Moreover, after Corollary 3.8, we have discussed the implications of our formula for the dimension and, in particular, how the latter compares to the number of parameters, which is commonly used as a measure of expressivity and/or sample complexity.

[1] Kileel et al., On the expressive power of deep polynomial neural networks, 2019.

We thank the reviewer for the comments and the suggestions. Below, we address the questions raised.

We acknowledge that it is important to elaborate on the importance of understanding the dimension of the neuromanifold. To this end, we have re-uploaded a new version of the manuscript with a discussion on the importance of dimension, especially its relation to sample complexity in machine learning, which has practical implications – see the section Introduction and Related Work. This discussion is also reported in the collective comment we posted above.

Regarding the question raised on "approximate fibers", we believe that this is an interesting point. As mentioned by the reviewer, this might provide a relaxed identifiability condition, that can potentially inspire improvements in terms of architecture design. One way to formalize this is by considering pre-images $\varphi^{-1}(U)$ via the parameterization map of neighbourhoods $U \subset \mathcal{M}$ in the neuromanifold $\mathcal{M}$ – for example, balls $U = B_{\varepsilon} \cap \mathcal{M}$, $\varepsilon \in \mathbb{R}_{>0}$. Such pre-images contain parameters that are mapped to functions that are close in the neuromanifold, according to some prescribed metric. One issue is that, different from fibers, pre-images of balls do not have a regular behavior: it is not true, for example, that they have the same dimension, for a generic ball in the neuromanifold. The diameter of these preimages is closely-related to – more precisely, it is upper bounded by – the Lipschitz constant of the parameterization map. Studying this Lipschitz constant is interesting by itself, since it is related to the regularity of the parameterization, and, as a consequence, to the "flatness" of the loss landscape for a given dataset/task. The Lipschitz constant of a polynomial map depends on the degree of the polynomial, which for deep lightning self attention networks equals to $3^l$, where $l$ is the number of layers (see lines 151-152 in the paper). This suggests that the degree – which is a fundamental invariant in algebraic geometry – might play an important role when considering `approximate fibers’, and might be connected to more general questions on optimization. Therefore, we believe that this is an interesting direction to explore.

We wish to comment on the extension of the result to include skip connections and/or MLPs. First, including skip connections into lightning self-attention layers still results in a polynomial model, but makes it non-homogeneous. This breaks the scaling symmetry in the parameterization. We conjecture that in this case the parameterization via $(**M**, L)$ is generically one-to-one, similarly to the traditional case (Conjecture 3.16), which is also non-homogeneous. We believe that this might be possible to prove with techniques similar to ours, and we have some ideas for approaching this. However, we do not have a complete argument at this stage, and we plan to work on this as a next step. Regarding MLPs, it is possible to include them in the model and potentially apply algebraic techniques similar to ours, as long as their activation function is polynomial – see [1] for an overview of polynomial MLPs. However, the fibers of the parameterization and the dimension of the neuromanifold are not well-understood for polynomial MLPs, and characterizing them is a major open problem [2], although recent progress has been made [3]. Even though we agree that this is an important future direction since it reflects the Transformer architectures that are used in practice, we believe that the complete understanding is far away, and this would require significant theoretical advances in order to be addressed. Still, our work could provide a basis for reasonable next steps along this path. In summary, both these directions constitute important lines for future work, and we have incorporated the content of this discussion in the new version of the manuscript (see Conclusions and Future Work), which we have re-uploaded.

[1] Kileel et al., On the expressive power of deep polynomial neural networks, 2019.

[2] Kubjas et al., Geometry of polynomial neural networks, 2024.

[3] Finkel et al., Activation thresholds and expressiveness of polynomial neural networks, 2024.