Rethinking Attention: Polynomial Alternatives to Softmax in Transformers

TL;DR

This paper proposes polynomial functions as alternatives to softmax in transformer attention, demonstrating they can provide similar regularization benefits and stable training without requiring probability distribution properties.

Contribution

It introduces polynomial-based attention mechanisms that serve as effective substitutes for softmax, challenging the necessity of probability normalization in transformers.

Findings

Polynomials can replace softmax in attention mechanisms.

Polynomial attention achieves comparable performance across tasks.

Regularization via polynomials stabilizes training.

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.