A Theoretical Study of (Hyper) Self-Attention through the Lens of Interactions: Representation, Training, Generalization

RESPONSE 1 (About the references): In the introduction we referenced several works on "theoretically understanding attention" but we agree that attention modifications could be discussed more. In the final version of the paper we will mention these approaches, i.e, the papers such as (DEBERTA) https://arxiv.org/pdf/2006.03654; (Improving sampling by modifying the effective diffusion) https://arxiv.org/abs/2410.05258; (Selective Attention) https://arxiv.org/pdf/2410.02703.

In addition, we will expand the discussion with the "mechanistic interpretability" literature such as (Tranformer Circuits) https://transformer-circuits.pub and previous studies attempted to understand and modify attention and modify attention.

Also, We will briefly mention NA2Q (ICML, 2023) and other relevant methods in the revised version of our paper.

RESPONSE 2 (Use of HyperFeatureAttention and HyperAttention in MARL Settings):Our work views tokens as interacting entities—a perspective that naturally encompasses agents in multi-agent reinforcement learning (MARL). In MARL, interactions are often more complex than simple pairwise relationships, and our proposed extensions are well suited to address such challenges. In particular, HyperAttention is designed to capture higher-order (e.g., three-way or n-way) interactions, which are common in collaborative or competitive multi-agent environments where an agent’s outcome may depend on the joint behavior of multiple agents. Similarly, HyperFeatureAttention enables the coupling of different feature-level interactions (such as spatial positions, velocities, or task-specific attributes), providing a more expressive representation of agent dynamics. Although our current experiments primarily focus on controlled settings, the theoretical framework and empirical evidence presented in Sections 5 and 6 of our paper indicate that these modules can be naturally extended to MARL scenarios. In future work, we plan to investigate and validate these extensions in MARL benchmarks, which we believe will further demonstrate the versatility and potential impact of our approach in such applications.

RESPONSE 3 (Potential Drawbacks or Challenges for Larger Datasets and More Complex Tasks): We strongly suggest you to refer RESPONSE 6 of our reply to reviewer xG3k where we summarize the potential drawbacks in terms of complexity of models and how to overcome them.

LINK (please open in a private window): https://drive.google.com/file/d/1lJ3HYR6i02jpm3CAxg4cu4J6icQXubJ8/view?usp=share_link

RESPONSE 1 (On the experiments for "further validation in more complex, real-world scenarios", "more complex or empirical scenarios", "well-established benchmark datasets"): For self-attention, our main goal is to provide rigorous theories for the understanding of the existing self attention mechanism. Consequently, our current experiments primarily validate the theoretical framework of linear self-attention, offering strong empirical support as opposed to conducting large-scale experiments on self-attention (which has already done on many literature as we cited). That being said, we have done large-scale experiments such as famous many body problem where we predict trajectories of bodies (please refer to LINK)

RESPONSE 2 ("clarity regarding symbolic definitions and intuitive explanations of mathematical formulations could be further improved" and "presentation of theorems and formulations ... with clearer definitions and intuitive explanations"): We made the symbolic definitions more distinguishable (for instance we will replace $s^{(n)}(i)$ "index to corresponding element mapping" function with $\mathcal{X}^{(n)}(i)$ since the former version may be confused with $\mathbf{s}^{(n)}$ "count vector") and we further improve the mathematical formulations. Further, in the final version of the paper we will provide clearer definitions and intuitive explanations

RESPONSE 3 (on the references on "Transformer interpretability"): Please refer to RESPONSE 1 of our reply to reviewer Ct39.

RESPONSE 4 (on the concerns and questions on complexity and practicality of Hyper(Feature)Attention): please refer to the "RESPONSE 6 (Comparison Between Attention Models)" part of our part of our reply to Reviewer xG3k.

RESPONSE 5 (Regarding empirical evidence for Hyper(Feature) Attention): Please refer to RESPONSE 3 in our reply to reviewer N655.

RESPONSE 6 (to Question 2): When we only remove the $d<N$ assumption, we start to get nonzero errors i.e. approximations to the exact solutions to the task (which is the practical real world situation). This is explained in both at the end of Section 3 "Why d=N?" part and in Theorem B.4. Shortly, "... Our goal is to understand how self attention captures mutual interactions, and d = N ensures orthogonal domain embeddings, yielding an exact and transparent representation. Compressing to $d < N$ is perpendicular to this focus and can be addressed separately using standard techniques... Starting with $d = N$ allows us to establish theorems that elucidate how self-attention captures pairwise interactions, while reducing to $d < N$ merely introduces a small approximation gap without altering the core theory. For completeness, we provide.. (Theorem~B.4)". That is being said, we runned experiments to quantify this sensitivity on the same colliding agents environment with $d<N$. As a result, we had nonzero errors (train, test, out of distribution). As the embedding dimension approached to the dictionary size the errors approached to zero. Although it is an important insight, that discussion is orthogonal to the theories for this paper,so we left those results for a future work. However, we may add them to the final version of the paper.

RESPONSE 7: Lastly, we are thankful for pointing out the typos, formatting issues and symbol explanation issues. Considering the length of the paper, unfortunately we made those small errors which are easy to fix and will be totally fixed in the final version of the paper.

RESPONSE 1 (the concerns on practicality of proposed models and what is their use while self-attention works in NLP): We agree that traditional self-attention architectures can also represent these complex interactions "in surprising ways". However, the main distinction is that the proposed models can represent those with fewer parameters. For instance, it is discussed in Appendix E.1 and Appendix E.2's "a short note on approximate embedding" paragraph. We agree that we need to carry some of those discussions to the main text and make it clear in the main text, which will be done in the revised version. Another point is we do not advocate replacing conventional self-attention with the proposed models. For more details we strongly suggest referring to "RESPONSE 6" part of our reply to reviewer xG3k.

RESPONSE 2 (insightful experiments): As for the small-scale experiments, we have promising preliminary results, please refer to "RESPONSE 3" part of our reply to reviewer N655.

RESPONSE 3 (the concerns on proof of Thm B.4): We agree with your concerns and we think that old version was not clear. Therefore, we completely changed the theorem statement and the proof, keeping the essence the same as old version. In the new version, we use low rank approximations to the matrices that would solve the task with zero error if embedding dimension were large enough. We bound the infinity norm between the output of the linear self attention having $d<N$ and the correct labels. The bound is in terms of singular values of the matrices that would solve the task with zero error if embedding dimension were large enough. Unfortunately, it is not allowed to provide whole theorem and proof with anonymous links.

RESPONSE 4 ("earlier works often do not delve into interpreting the roles of the learned parameters beyond attention patterns") While making that statement we were mostly thinking about the papers not only theoretical but also mathematical, with rigorous theories expectations and validations of those experiments. We apologize for the confusion and in final version we will clarify. In more detail, our approach is different from conventional mechanistic interpretability approach that our approach is rigorous/mathematical and our experiments are for validating our predictions. Their approach is like "phenomenology", and our approach is "fundamental".

RESPONSE 5 (About the scepticism on assumption 4.2): In the revised version of the paper we will add a section to the appendix that justifies Assumption 4.2. In short, assuming $\mathbf{s}^{(n)} \in \mathbb{R}^{|\mathcal{S}|}$ are sampled from a distrubution such that covariance of $\mathbf{s}^{(n)}$ is positive definite, we will use matrix Bernstein inequality and Weyl's inequality and prove that the statement in the Assumption 4.2 hold with high probability: for some $\gamma >0$, B is the number of data samples, $\mathbb{P}\left(\mathrm{rank}(\mathbf{S}_{{B}\mu}) \le |\mathcal{S}|\right)\le e^{-\gamma B}$

RESPONSE 6 (real LLM's not falling to skip trigram bug): We agree that LLMs as a whole do not have such bugs. As you mentioned, the skip connections from earlier layers help to approximately get rid of those problems and by adding more layers, the approximations gets better. In our analysis, we focus on self-attention itself (not a large neural network) and compared it with our new modules (small modules not large networks). We strongly suggest looking at the RESPONSE 6 part of our reply to reviewer xG3k (especially the last paragraph of it), which we will add to the paper.

RESPONSE 7 (On the question about the usefulness of HyperFeatureAttention (HFA)): Both HFA and self-attention have same complexity, so we compared them by parameter count,since even MLPs can represent any function with enough parameters. Our intentially unfair comparison used two-layer multihead self-attention against a single-layer single-head HFA. A fairer comparison (single-layer single-head vs. single-layer single-head) shows HFA captures self-attention. We acknowledge that some interactions require two-layer multihead self-attention, but for more on HFA’s utility, please see “RESPONSE 3” in our reply to reviewer N655.

RESPONSE 8(connection between theoretical results and the proposed modifications): Various colliding agent example scenarios in Appendix B.1.3 and "Non-Identical Agents Revisited" in Appendix E.1, we introduce how HFA is a natural construction from the theoretical insights on self-attention. Though seemingly unrelated, they were stepping stones. Thank you; we’ll clarify.

RESPONSE 9: For essential references please refer to "RESPONSE 1" to reviewer Ct39.

LINK (OPEN IN PRIVATE WINDOW): https://drive.google.com/file/d/1lJ3HYR6i02jpm3CAxg4cu4J6icQXubJ8/view?usp=share_link

RESPONSE 1 (the results of one-hot embedding): The results were explained as "negligible error $\Theta(10^{-7})$ on test and out of distribution data". However, we agree and for complenetess we will add the learned parameters to the paper (see LINK).

RESPONSE 2 (additional experiments for theory validation): We experimented the theories on the genotype phenotype mapping task (explained Sec. 3 Example 2 and Appendix B.2) and got same results $\Theta(10^{-7})$ error on test and OOD data.We may add them to paper.

RESPONSE 3 (the figures in the appendix "which are neither discussed nor referenced in the main text"): discussed in Sec.7 as "shown in Figures 3 and 4, these matrices are indistinguishable from the theoretical counterparts"

RESPONSE 5: Sanford et al. shares our $\mathcal{O}(L^3)$ complexity. The footnote 5 is unnecessary and will be removed.

RESPONSE 6 (Comparison Between Attention Models): We will add the following discussion and the attached table (LINK) to the paper. Letting embedding dim to be $d$, sequence length to be $L$, we compare those models in terms of number of parameters, computational complexity, and abilities.

• For self-attention, it captures mutual interactions between entities (cf.Thms 3.1,4.4,4.6,4.8).If it has multiple heads, it can capture summation over mutual interactions between features of the entities. It has ${\Theta}(d^2)$ parameters, and its computational complexity is ${\Theta}(L^2)$. In order to reduce the computational complexity to ${\Theta}(L)$, people came up with approximations called "Linear Attention" (Katharopolus et al 2020, Wang et al 2020). However, despite the name "linear", the method is used to approximate nonlinear selfattention.

• HyperFeatureAttention captures couplings of mutual interactions between features (cf. Sec.5, AppendixE). If it has multiple heads, it can capture summation over couplings of mutual interactions between features. Same as self-attention, HyperFeatureAttention has ${\Theta}(d^2)$ parameters, and its computational complexity is ${\Theta}(L^2)$. The same Linear Attention approximations can be applied to HyperFeatureAttention, reducing its computational complexity to ${\Theta}(L)$. Seeing that the main goal of the paper is not this approximation for HyperFeatureAttention, we did not show it explicitly.

• As for HyperAttention of order $n$, it captures up to $n^{th}$ order interactions (Refer to Sec.6 and AppendixF). If it has multiple heads, it can capture summation over up to $n^{th}$ order interactions between features of the tokens. It has ${\Theta}(d^2)$ parameters, and its computational complexity is ${\Theta}(L^n)$. Using similar Linear Attention approximations, in Appendix~F.3, we explained how to reduce computational complexity of HyperAttention to ${\Theta}(L)$. We will add the requirements on the query and key matrices that if their entries are less than $o(\sqrt[3]{\log L})$ then the infinite norm between the actual outputs and approximate outputs is less than $1/\mathrm{poly}(L)$.

• While we do not advocate replacing conventional self-attention with HyperAttention or HyperFeatureAttention, we propose these mechanisms as complementary enhancements. In a typical multi-head architecture, certain heads may employ standard self-attention while others utilize HyperFeatureAttention (of varying orders) or HyperAttention to capture richer, higher-order interactions. Depending on the computational constraints, the HyperAttentions may leverage the linear approximations described in Appendix~F.3.

RESPONSE 7 ("the violations that make expected guarantees not generally true in the real world" and "why dont we have zero test errors and what happens when $d<N$"): When we only have $d<N$ assumption, we start to get nonzero errors i.e. approximations to the exact solutions to the task (which is the practical real world). This is explained in both at the end of Sec. 3 "Why d=N?" part and in Thm B.4. Another key assumption in our analysis is the (weak, strong, universal) realizability. Although these conditions may not hold in practical scenarios, implying that a single layer, single-head self-attention mechanism may not fully capture the complexities of real-world tasks, it is important to note that SOTA models typically employ multiple layers of self-attention blocks. Such architectures can approximate the desired functions even when the strict realizability assumptions are violated.Actually, the realizability and d=N assumptions are very similar in essence because when letting d large enough, we simply showed in Appendix B that realizability is satisfied (with several example scenarios such as colliding agents, genotype phenotype mapping, vision task, time series prediction). A thorough investigation of these approximations, lies beyond the scope of this work.

LINK (please open in a private window): https://drive.google.com/file/d/1lJ3HYR6i02jpm3CAxg4cu4J6icQXubJ8/view?usp=share_link

RESPONSE 1 (On Computational Complexity of Models and Response to Question 1): In our appendix, we showed that by leveraging linear self-attention approximations, the computational complexity of HyperAttention with even softmax activation function can be reduced to $\mathcal{O}(L)$. We suggest looking at "RESPONSE 6" part of our reply to reviewer xG3k, where we detailly answered this question and what kind of revisions we will do.

RESPONSE 2 (to Question2): For self attention, our main goal is to provide rigorous theories for the understanding of the existing self attention mechanism, as opposed to conducting large-scale experiments (which has already been done on many literature as we cited). Our current experiments validate the theoretical framework of linear self-attention.

RESPONSE 3 (Regarding emprical evidence for Hyper(Feature)Attention): We agree experimental verification of HFA and HA are important. As this work is primarily theoretical, we now present small-scale experiments while working on large-scale versions in a following paper. In our small-scale experiments, we evaluate next-token prediction under two setting. 1) 1-layer 1-head self-attention, HyperFeatureAttention with context window of 1024 but all of the rest of the hyperparemeters are the same as GPT-2 paper (e.g. embedding dimension is 768 etc.)2) 3 layer 3 head self-attention, HyperFeatureAttention experiments with all of the hyperparameters the same as GPT-2 (context window of 1024, embedding dimension of 768 and so on). 3) 1-layer 1-head self-attention, HyperAttention, HyperFeatureAttention with context window of 256 but all of the rest of the hyperparemeters are the same as GPT-2 paper (e.g. embedding dimension is 768 etc.) In both experiments the networks we compare have the same number of parameters. Self-attention and HyperFeatureAttention networks have the same computational complexities but HyperAttention has $\mathcal{O}(L^3)$ complexity (we did not apply the approximations reducing complexity yet). We deliberately adopted certain design choices that do not favor our proposed models which implicitly favor self attention to ensure a rigorous evaluation. Although these results are preliminary, they consistently indicate that networks incorporating HyperAttention and HyperFeatureAttention achieve better perplexities and cross entropy looses compared to those based solely on standard self-attention. For the comparison over the results, please refer to the LINK. As for large-scale experiments with comparisons to SOTA, we have deferred extensive, large-scale experimens for the new models to a forthcoming follow-up study because 1) this paper focuses on fundamental theoretical insights and 2) we have possible variations that preserve the essence of Hyper(Feature)Attention models yet may enhance empirical performance, thoroughly exploring these refinements and different possible architectures involves extensive research.

In addition, in the various colliding agent example scenarios in Appendix B.1.3 and "Non-Identical Agents Revisited" in Appendix E.1, with several examples, we step by step induce how HyperFeatureAttention is a natural construction from the theoretical insights on self-attention. Similarly, we explain in Section 6 how HyperAttention is natural construction from our self-attention theories. It is also worth noting that HyperFeatureAttention and HyperAttention are designed to supplement rather than replace existing self-attention mechanisms. For instance, in a multi-head architecture, some heads may employ standard self-attention, while others integrate HyperFeatureAttention (to capture cross-coupled feature interactions) some employ HyperAttention (e.g., order-3 or order-4) to capture higher order interactions -even combination of HyperFeature and Hyper Attention is possible.

RESPONSE 4 (the interpretability of parameters with sinusoidal embedding and response to question 3): We believe that the confusion was merely due to a narration problem. It was stated "..., the learned parameters do not overlap with the parameters we originally devised, especially for the sinusoidal embedding, seen in Figure 2 -these learned parameters lack easy interpretation ... as discussed in Corollary D.5, this outcome is a natural consequence of the generalization theories ... Therefore we focus on" transformed versions of the parameters ... Further our generalization theories are on those transformed versions. "As shown in figures 3 and 4, these matrices are indistinguishable from the theoretical counterparts. With only ${\Theta}(10^{-4})$ mean square distance between their entries. Thus, in this simple setting, we can fully interpret the meaning of the parameters". We will fix narration.

RESPONSE 5: For references look at RESPONSE 1 to reviewer Ct39.