# Is Softmax Necessary as the Activation Function for Self-Attention Mechanisms?

I’m curious about the mathematical reasoning behind the use of the softmax function as the activation function in self-attention mechanisms within neural networks. Specifically, I’m interested in understanding if there is a theoretical basis that necessitates the use of softmax over other activation functions.

Softmax is commonly employed to convert raw attention scores into a probability distribution, ensuring that the sum of attention weights equals 1. This normalization allows the model to effectively focus on certain parts of the input sequence. However, I wonder if there are alternative activation functions that could be less constraining and still allow the optimization process to determine the best way to allocate attention, similar to how tanh or other activations work in different layers of a neural network.

1. Is there a mathematical justification for the necessity of softmax in self-attention mechanisms?
2. Could other activation functions, perhaps with fewer constraints, be used effectively in place of softmax, allowing the optimization process more flexibility?

Any insights or references to relevant literature would be greatly appreciated.

In the case of attention, many of such properties do not hold, and it's used merely for its ability to emulate a categorical distribution. However, there is not specific reason to use it instead of some other formulation, for example also the following activation will be a valid one: $$\frac{p_i^2}{\sum_j p_j^2}$$ where the $$p_i$$ are the row logits from $$QK^T$$. However, even though this formulation is valid, as the sum adds up to 1 and all probabilities are positive, it's much more unstable, as the denominator might get to 0, or close to 0, thus introducing instability during learning
At the end of the day, any non parametric activation that maps a vector $$v\in R^n$$ to a vector $$u$$ such that $$||u||_1=1$$ and $$u_i \ge 0\forall i$$ is indeed a valid "attention" activation
• @Kasia yes, if you relax that assumption (debatable calling it attention then though) you can even solve the $O(n^2)$ problem of transformers, which is what indeed linear transformer do, and it's even used in the paper of "RWKV: Reinventing RNNs for the Transformer Era", since they exploit that using linear attention, you don't need to know the future tokens attention, and thus you can decay to a sort of RNN with additive recurrent layer Dec 8, 2023 at 10:39