Rethinking Attention: Polynomial Alternatives to Softmax in Transformers

TL;DR

This work questions the necessity of softmax-derived probabilistic properties in transformer attention and proposes that the success of softmax stems from implicit Frobenius-norm regularization of the self-attention matrix. By introducing polynomial activations with appropriate scaling, notably $\phi(x)=\frac{1}{\sqrt{N}}x^p$, the authors demonstrate comparable or superior performance across vision, object detection, long-range tasks, and physics-informed modeling, while violating non-negativity, normalization, and sparsity. Theoretical results bound the Frobenius norms and gradients for softmax and polynomial activations, providing a principled scaling regime, and extensive experiments validate the approach on ImageNet, COCO, Nyströmformer benchmarks, and PINNs. The findings offer a new lens on attention design, enabling alternative, potentially more stable mechanisms that preserve performance while loosening traditional probabilistic constraints. Overall, the paper broadens the design space for attention in transformers and suggests practical benefits for scalability and interpretability across diverse domains.

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.

Figures (7)

• Figure 1: Training ViT-Tiny with the activation $\phi(x) = x^3$ with different sequence lengths and different scales. As the sequence length gets larger, the $k$ scale (x-axis) needed to obtain good accuracy when using $\frac{1}{k}x^3$ as an activation increases validating the theory from section \ref{['subsec:poly_acts']}.
• Figure 2: Frobenius norm of the self-attention matrix with three different activations in layer 2, 7 and 12 of the ViT-Tiny architecture during training.
• Figure 3: Frobenius norm of the jacobian of the self-attention matrix with three different activations in layer 2, 7 and 12 of the ViT-Tiny architecture during training.
• Figure 4: Heat maps of the self-attention matrix in layer 2, head 8, of a ViT base architecture, comparing $x^3$ + dynamic (left) and softmax (right) activations after training. The stark difference in self-attention patterns between the two activations is evident, showing distinct distributions across input tokens.
• Figure 5: Heat maps of the self-attention matrix in layer 12, head 6, of a ViT base architecture, comparing $x^3$ + dynamic scale (left) and softmax (right) activations after training. The contrast in self-attention patterns between the two activations is clearly visible.

Theorems & Definitions (9)

• Theorem 4.1
• Theorem 4.2
• Corollary 4.3
• Theorem 4.4
• Corollary 4.5
• Remark 4.6
• proof : Proof of theorem \ref{['thm:softmax_regularity']}
• proof : Proof of theorem \ref{['thm:expectation']}
• proof : Proof of theorem \ref{['thm:grad_expectation']}