Rethinking Attention: Polynomial Alternatives to Softmax in Transformers

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.

💡 Research Summary

The paper challenges the widely held belief that the success of softmax‑based attention in transformers stems from its three canonical properties: non‑negativity, row‑wise normalization (interpretable as a probability distribution), and sparsity. Instead, the authors argue that the crucial factor is an implicit regularization effect: softmax constrains the Frobenius norm of the attention matrix to grow at most on the order of √N, where N is the sequence length. This bound limits the magnitude of gradients during back‑propagation, thereby stabilizing training.

Theoretical contributions begin with Theorem 4.1, which proves that for any N×N matrix A, the softmax‑transformed matrix satisfies ‖softmax(A)‖_F ≤ √N and its Jacobian satisfies ‖∇softmax(A)‖_F ≤ 2√N. The authors then propose polynomial activation functions ϕ(x)=x^p (p≥1) scaled by a factor α=1/√N, i.e., ϕ_α(x)=α·x^p. Theorem 4.2 and Corollary 4.3 show that, under i.i.d. Gaussian assumptions on inputs X and projection matrices Q, K, the expected Frobenius norm of (XQKᵀXᵀ)^p is O(N); after scaling by 1/√N it becomes O(√N), matching the softmax bound. Theorem 4.4 and Corollary 4.5 extend the analysis to gradients with respect to Q and K, demonstrating analogous O(√N) bounds after scaling.

Empirical validation focuses on Vision Transformers (ViT‑Tiny) trained on Tiny‑ImageNet. The authors first sweep the scale k in the activation k·x³ across four sequence lengths (N=256, 64, 16, 8). Consistent with the theory, the optimal k grows proportionally to √N, confirming that the regularization effect is scale‑dependent. In a second experiment they compare three configurations: standard softmax, unscaled cubic activation (x³), and scaled cubic activation (1/√N·x³). Measurements of the attention matrix Frobenius norm and the Jacobian norm across layers 2, 7, 12 show that the scaled polynomial aligns closely with softmax, while the unscaled version yields substantially larger norms and unstable gradients. Accuracy results reveal that the scaled cubic (1/16·x³ for N=256) attains a Top‑1% accuracy of 50.5 %, marginally surpassing softmax’s 50.26 %, whereas the unscaled cubic drops to 45.3 %.

Further experiments extend the polynomial attention to a broad set of tasks: image classification, object detection, instance segmentation, text classification, and physics‑based modeling. Across these domains, the scaled polynomial consistently matches or slightly exceeds softmax‑based baselines, while offering a simpler computational graph and potential hardware advantages.

The key insight is that the probabilistic interpretation of attention is not essential for performance; what matters is controlling the magnitude of the attention matrix via Frobenius‑norm regularization. By abandoning the traditional constraints of positivity, normalization, and sparsity, and using appropriately scaled polynomial activations, one can design attention mechanisms that are both theoretically grounded and empirically competitive. This work opens a new design space for transformer architectures, especially for scenarios demanding lightweight, hardware‑friendly, or analytically tractable attention modules.