1
$\begingroup$

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.

Improve this question
$\endgroup$

1 Answer 1

Reset to default
0
$\begingroup$

Softmax was first used for its properties, as it is differentiable, has not domain problem even though it has a division, and it's gradient is well behaved (and in conjunction with Categorical Cross Entropy it gets simplified a lot becoming linear)

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

Improve this answer
$\endgroup$
3
  • $\begingroup$ Part of my question is whether the attention mechanism has to generate outputs that sum to 1, creating a valid probability distribution. Would lifting this condition change anything other than the ease of interpretation from our point of view? $\endgroup$
    – Kasia
    Commented Dec 7, 2023 at 23:24
  • 1
    $\begingroup$ @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 $\endgroup$
    – Alberto
    Commented Dec 8, 2023 at 10:39
  • $\begingroup$ Just wanted to add a quite relevant and recent paper on the topic: arxiv.org/abs/2410.18613 $\endgroup$
    – cjferes
    Commented Oct 28 at 0:12

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .