Rethinking Attention: Polynomial Alternatives to Softmax in Transformers (2410.18613v2)

Abstract: This paper questions whether the strong performance of softmax attention in transformers stems from producing a probability distribution over inputs. Instead, we argue that softmax's effectiveness lies in its implicit regularization of the Frobenius norm of the attention matrix, which stabilizes training. Motivated by this, we explore alternative activations, specifically polynomials, that achieve a similar regularization effect. Our theoretical analysis shows that certain polynomials can serve as effective substitutes for softmax, achieving strong performance across transformer applications despite violating softmax's typical properties of positivity, normalization, and sparsity. Extensive experiments support these findings, offering a new perspective on attention mechanisms.

• The paper reveals that softmax’s key contribution lies in its implicit regularization of the attention matrix’s Frobenius norm, promoting stable training.
• Empirical results show that scaled polynomial activations, such as cubic forms, rival or exceed softmax performance on image classification and object detection tasks.
• These insights challenge the traditional probability-based view, opening avenues for norm-centric design of advanced attention mechanisms in transformers.

Rethinking Softmax: Self-Attention with Polynomial Activations

The paper entitled "Rethinking Softmax: Self-Attention with Polynomial Activations" presents a theoretical and empirical examination of the canonical role of softmax in transformer architectures, contesting the prevailing assumption that its primary success stems from generating a probability distribution for attention allocation. The authors argue that the effectiveness of softmax is more deeply rooted in its implicit regularization of the Frobenius norm of the attention matrix during training.

Theoretical Insights

The foundation of the paper lies in uncovering the theoretical underpinnings of softmax's regularization properties. The authors present a theorem elucidating that the Frobenius norm of the self-attention matrix under softmax grows sub-linearly with respect to the input sequence length. This regulation is further mirrored in the gradient norms, indicating stable training dynamics which are vital for gradient descent methods. The paper extends this insight by deriving polynomial activations that similarly regularize the Frobenius norm, showcasing that polynomial mappings such as those of the form $\frac{1}{\sqrt{N}}x^p$ achieve comparable performance constraints.

Empirical Evaluation

The empirical investigations focused on validating the theoretical claims by benchmarking these polynomial activations against softmax across several transformer tasks. The experiments highlighted that polynomial forms, particularly $\frac{1}{16}x^3$ and $\frac{1}{16}x$, not only rival softmax in terms of accuracy but also surpass it under certain conditions. These findings were consistent across multiple architectures including ViT, DeiT, Swin Transformer, and XCiT, evaluated on standard datasets like ImageNet-1k and COCO.

In tasks such as image classification and object detection, these polynomial activations were shown to be highly effective when appropriately scaled. The paper offers visual insights into how these activations alter the attention landscapes, demonstrating distinct patterns that challenge our traditional understanding shaped by softmax.

Implications and Future Directions

The implications of this paper are pivotal: it redefines the rationale behind attention mechanisms' success, urging a shift from probability-centric interpretations towards norm-based regularization perspectives. The results suggest that transformer architectures might not be inherently dependent on probabilistic attention distributions, opening pathways to design novel activation functions that can further enhance model robustness and efficiency.

Future research could explore extending these theoretical constructs beyond self-attention to other types of attention mechanisms. Additionally, investigating the implications of these findings in other domains like NLP, and further refinement of polynomial activations for optimized performance, would form a compelling continuation of this work.

In conclusion, this paper provides a rigorous theoretical examination and an insightful empirical paper challenging established norms of softmax-based self-attention, paving the way for innovative attention mechanisms in AI systems.