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

Abstract: Self-attention has emerged as a core component of modern neural architectures, yet its theoretical underpinnings remain elusive. In this paper, we study self-attention through the lens of interacting entities, ranging from agents in multi-agent reinforcement learning to alleles in genetic sequences, and show that a single layer linear self-attention can efficiently represent, learn, and generalize functions capturing pairwise interactions, including out-of-distribution scenarios. Our analysis reveals that self-attention acts as a mutual interaction learner under minimal assumptions on the diversity of interaction patterns observed during training, thereby encompassing a wide variety of real-world domains. In addition, we validate our theoretical insights through experiments demonstrating that self-attention learns interaction functions and generalizes across both population distributions and out-of-distribution scenarios. Building on our theories, we introduce HyperFeatureAttention, a novel neural network module designed to learn couplings of different feature-level interactions between entities. Furthermore, we propose HyperAttention, a new module that extends beyond pairwise interactions to capture multi-entity dependencies, such as three-way, four-way, or general n-way interactions.

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

The paper, "A Theoretical Study of (Hyper) Self-Attention through the Lens of Interactions: Representation, Training, Generalization," delves deeply into the theoretical capabilities of self-attention mechanisms, situating them within the broader context of interaction learning. Recognizing the foundational role of self-attention in contemporary neural architectures, the authors address its applications across a diverse range of domains such as natural language processing, computer vision, and reinforcement learning. Crucially, the paper's contributions lie in offering a comprehensive theoretical framework that captures the essence of self-attention through mutual interactions and extends this understanding to novel constructions like HyperFeatureAttention.

The authors reframe self-attention as a mechanism for learning mutual interactions between entities, such as agents in a multi-agent system or alleles in genetic sequences. This approach simplifies many application scenarios and aligns with the core representational capabilities of linear self-attention. The paper shows that a single-layer linear self-attention is sufficient to represent a variety of pairwise interaction functions efficiently, with parameters scaling as $\mathcal{O}(k^2)$.

A critical strength of the paper lies in its theoretical proof that linear self-attention can achieve zero error on training data and generalizes effectively under conditions of data versatility. This is complemented by rigorous experimentation that confirms these theoretical insights, showcasing the practical applicability of theoretical predictions in controlled settings like a colliding agents environment.

A significant leap forward is the introduction of HyperFeatureAttention and HyperAttention, designed to capture complex interaction patterns beyond pairwise dependencies. HyperFeatureAttention enables the modeling of couplings between feature-level interactions, effectively reducing the computational overhead compared to standard self-attention. Algorithmically, HyperAttention further extends this framework to encompass higher-order interactions, such as three-way or $n$-way interactions. This is particularly valuable for domains requiring the modeling of intricate dependencies, such as higher-order language constructs or protein structures.

The implications of these findings are profound. HyperFeatureAttention and HyperAttention suggest new avenues for model design in tasks involving rich interaction dynamics, offering potential computational savings and enhanced representational capacity. The reduced parameterization and efficient computation of these modules suggest they could be integrated into existing Transformer-based architectures to enhance model performance without the corresponding increase in complexity.

While the theoretical developments are compelling, practical challenges remain, particularly regarding the computational efficiency of HyperAttention at high orders. The authors acknowledge these challenges and propose approximations that maintain performance while reducing compute requirements.

In summary, this paper provides substantial theoretical and empirical advancements in understanding self-attention and its derivatives. By framing self-attention as an interaction learning mechanism, the authors offer a novel perspective that enhances its applicability and efficiency in capturing complex dependencies in diverse applications. The introduction of HyperFeatureAttention and HyperAttention opens new doors for extending transformer capabilities, especially in environments where learning is underscored by multifaceted interactions.