Scalable-Softmax Is Superior for Attention

Paper: arXiv
Author: Ken M. Nakanishi
Date of Publication: 31st January 2025

Overview

Softmax that is used in scaled dot product attention cause attention scores to flatten as context length increases. This reduces the model’s ability to prioritize to key information in context and also generalise to longer contexts not seen during training. To fix this author modifies softmax and introduces $SSMax$. $SSMax$ is defined as:

$$z_i\mapsto \frac{n^{s{z}_i}}{{\sum }_{j=1}^n{n}^{s{z}_j}}=\frac{e^{(s\log n)z_i}}{{\sum }_{j=1}^n{e}^{(s\log n)z_j}}$$

where, $n$ is context length and $s$ is a learnable parameter shared amongst all layers and heads.
Results show that $SSMax$ performs much, much better at information retrieval tasks and also generalises better to longer contexts even without being trained on longer contexts.

Problem with Softmax

Softmax transforms an input vector into a vector that can be interpreted as a probability distribution, where all elements are non-negative and sum up to one.

$$z_i\mapsto \frac{e^{z_i}}{{\sum }_{j=1}^n{e}^{z_j}}$$

In the attention layers of transformer, input vector size $n$ increases as the context length grows. Softmax plays a critical role in computing attention scores over all tokens in the context, determining how much ‘attention’ is allocated to each token. When $n$ grows, the denominator in softmax increases while the numerator remains independent of $n$. As a result the resultant distribution becomes increasingly flat (attention fading).
This reduces the model’s ability to focus on key tokens in context and also reduces ability to generalise to longer contexts.

Scalable-Softmax (SSMax)

$SSMax$ is defined as:

$$z_i\mapsto \frac{n^{s{z}_i}}{{\sum }_{j=1}^n{n}^{s{z}_j}}=\frac{e^{(s\log n)z_i}}{{\sum }_{j=1}^n{e}^{(s\log n)z_j}}$$

$SSMax$ also transforms the input vector into a probability distribution as it can clearly be seen from the definition. However the key difference is in the dependence of exponential base on input vector size $n$. This design helps in mitigating attention fading and the resulting attention scores remains focused on the key tokens.
Author provides very nice justifications for $SSMax$.

Rationale Behind the Design of SSMax

To investigate the optimal variant of softmax, author replaced softmax with the following function at all layers and heads:

$$z_i\mapsto \frac{e^{(s{p}_n+b)z_i}}{{\sum }_{j=1}^n{e}^{(s{p}_n+b)z_j}}$$

where, $s$ and $b$ are learnable parameters unique to each layer and head. And $p_n$ $(n=1,2,3,....N)$ represents learnable parameters shared across all layers and heads, depending solely on the input vector size $n$. $N$ denotes the size of context length used during training. Author trains a model with the above function replacing softmax in scaled dot product attention.

$p_n$ followed a logarthimic relationship of the form, $p_n\approx a_1\log n+a_2$
This finding suggested that softmax in attention mechanism could benefit from reformulation as:

$$z_i\mapsto \frac{e^{(s\log n+b)z_i}}{{\sum }_{j=1}^n{e}^{(s\log n+b)z_j}}=\frac{n^{s{z}_i}{e}^{b{z}_i}}{{\sum }_{j=1}^n{n}^{s{z}_j}{e}^{b{z}_j}}$$

where, $s$ and $b$ are layer and head specific learnable parameters. $b$ is referred as bias. Based on further evaluation (as we will see in Results), omitting $b$ turns out to be better and thus author arrived at $SSMax$.

Justification for the Design of SSMax

Let $z=(z_1,z_2,....,z_n)$ be an input vector of size $n$. Let $z_{max},z_{2nd},z_{min}$ denote its maximum, second maximum and minimum elements, respectively. Let $z_{max}>z_{2nd}$.

When $z$ is processed by softmax, $z_{max}$ is transformed as:

$$z_{max}\mapsto \frac{e^{z_{max}}}{{\sum }_{j=1}^n{e}^{z_j}}$$

We can replace denominator by $(n-1)e^{z_{min}}+e^{z_{max}}$ to obtain an upper bound.

$$\frac{e^{z_{max}}}{{\sum }_{j=1}^n{e}^{z_j}}\le \frac{e^{z_{max}}}{(n-1)e^{z_{min}}+e^{z_{max}}}$$

Multiplying and Dividing RHS by $e^{z_{max}}$ we get,

$$\frac{e^{z_{max}}}{{\sum }_{j=1}^n{e}^{z_j}}\le \frac{1}{\frac{n-1}{e^{z_{max}-z_{min}}}+1}.$$

As we can clearly see, as $\lim n\to \infty$, the maximum element of the output vector produced by Softmax approaches zero.

On the other hand, when $z$ is processed by $SSMax$, $z_{max}$ is transformed as

$$z_{max}\mapsto \frac{n^{s{z}_{max}}}{{\sum }_{j=1}^n{n}^{s{z}_j}}$$

Assuming $s>0$, we can obtain an upper bound for RHS.

$$\frac{n^{s{z}_{max}}}{{\sum }_{j=1}^n{n}^{s{z}_j}}\le \frac{n^{s{z}_{max}}}{(n-1)n^{s{z}_{min}}+n^{s{z}_{max}}}=\frac{1}{\frac{n-1}{n^{s(z_{max}-z_{min})}}+1}$$

Similarly we can obtain a lower bound.

$$\frac{n^{s{z}_{max}}}{{\sum }_{j=1}^n{n}^{s{z}_j}}\ge \frac{n^{s{z}_{max}}}{(n-1)n^{s{z}_{2nd}}+n^{s{z}_{max}}}=\frac{1}{\frac{n-1}{n^{s(z_{max}-z_{2nd})}}+1}.$$

Now we have,

$$\frac{1}{\frac{n-1}{n^{s(z_{max}-z_{2nd})}}+1}\le \frac{n^{s{z}_{max}}}{{\sum }_{j=1}^n{n}^{s{z}_j}}\le \frac{1}{\frac{n-1}{n^{s(z_{max}-z_{min})}}+1}$$

The maximum element output by $SSMax$ exhibits the following properties:

If $z_{max}-z_{2nd}>\frac{1}{s}$, the lower bound approaches $1$. So the output of $SSMax$ approaches $1$. Meaning attention is focused on the element with the highest value.
If $z_{max}-z_{min}<\frac{1}{s}$, the upper bound approaches $0$. So the output of $SSMax$ approaches $0$. Meaning attention is distributed across all the elements.

Thus, $SSMax$ ensures that attention is focused on elements whose values exceed others by approximately $\frac{1}{s}$, while distributing attention when all values are within a range of approximately $\frac{1}{s}$.

We can very easily convert softmax attention to $SSMax$ attention by

$$\text{SSMax}\left(\frac{{\mathbf{q}}_n{\mathbf{K}}_{1:n}^T}{\sqrt{d}}\right)=\text{Softmax}\left(\frac{(s\log n){\mathbf{q}}_n{\mathbf{K}}_{1:n}^T}{\sqrt{d}}\right)$$

Results

The loss during training. Author reports $SSMax$ is better during regular training too.

The grey dotted line is the context window at which models were trained. And models’ context length was extended by simply increasing the $\theta$ of RoPE with no additional training.

Results when evaluated on Needle-in-a-haystack.

Overall, results look very good with barely any extra learnable parameters. Big if true.

Thoughts

I really liked the paper. The intuition behind $SSMax$ makes sense. The results look pretty good too with almost zero additional compute cost. And very easy to implement. Big if true.
Combining this with differential attention may increase evaluation results on Needle-in-a-haystack even more. It would be nice to try that out.