Support Smooth Softmax in GroupQueryAttention

[tianleiwu]

Description

This implements a modified softmax that adding one to the denominator. Similar idea was implemented as add_zero_attn in torch.nn.MultiheadAttention about 6 years ago, and later in softmax_with_extra_logit in flaxformer about 3 years ago. Also recently the formula was analyzed in Attention Is Off By One by Evan Miller.

Background

Softmax (formula 1) is like the following:

$$y_{i} = \frac{exp(x_{i})}{\sum_{j} exp(x_{j})}$$

After applying softmax, each element will be in the range of $[0, 1]$, and the elements will add up to 1, so that they can be interpreted as probabilities.

In transformers model like BERT or LLAMA, softmax is used to normalize batch_size x num_heads x query_sequence_length rows, each row has size of key_sequence_length (total tokens in key).

For such application, softmax has two potential issues:

• When all elements are -inf (for example, a whole row is masked when a query token is padding), the result is not a number (NaN) since exp(-inf)=0 and divided-by-zero is encountered in the above formula.
• Why do we need normalize in a way that each row are treated as equal important (each row has sum equals to1)? In language model, some words in a query may carry more weight because they provide key information that determines the context or meaning of the query.

For example, consider the query: "Find the best Italian restaurant near me."
In this sentence:

• "Italian" and "restaurant" are the most important words because they define what the user is specifically looking for.
• "Best" is also significant as it indicates the user is seeking a high-quality option.
• "Near me" adds location-based context to the search.
• Words like "find" and "the" are less important in this context because they don't add specific content to the query; they are more functional in nature.

Thus, in processing this query, a language model might give more attention or weight to the words "Italian," "restaurant," and "best" to understand and respond accurately.

Smooth Softmax (formula 2) is a modified version that introduces a smooth factor like the following:

$$s_{i} = \frac{exp(x_{i})}{1+ \sum_{j} exp(x_{j})}$$

A straight-forward explanation of such formula is that there is an additional virtual element equal to zero, then exp(0) becomes the extra 1 in denominator.

This formula could tackle the above two issues:

• It could handle the special case that all elements are -inf: the result $s_{i}$ is 0 for every element in such case.
• Sum of all elements $\sum_{i}{s_{i}} = \frac{\sum_{i}{exp(x_{i})}}{1+ \sum_{j} exp(x_{j})}$ is in the range of [0, 1), so that we can train the model to assign different importance to different rows. For a query word or a head that is not important, the whole row might be zeros.

Since exponential is prone to overflow or underflow, to get stable result, formula 3 can be used:

$$s_{i} = \frac{exp(x_{i} + c)}{exp(c)+ \sum_{j} exp(x_{j} +c)}$$

c can be any value in theory. In practical, choice of constant c shall avoid $exp(c)$ and $exp(x_{i} +c)$ overflow (or underflow) at the same time. A reasonable choice is like formula 4:

$$c=-\max_{j} \{ x_j \}$$

or apply a constraint that c <=0 like the following formula 5:

$$c=-\max(0, \max_{j} \{ x_j \})$$

The latter one (formula 5) ensures that $s_{i}$ will fallback to formula 2 when all elements are negative.

For CPU provider, smooth softmax is implemented in MLAS. CPU implementation uses formula 5.

@wangyems implemented the smooth softmax in flash attention for CUDA, which requires Ampere or newer GPU. The implementation of smooth softmax in flash attention uses formula 4.

[yufenglee]

For this "When $x_{i}$ is much smaller than 0, $exp(x_{i})$ will be close to 1. When all elements are much smaller than 0, the result will be close to 1 / (1 + N).", is it a typo "When $x_{i}$ is much smaller than 0"? It should be $x_{i}$ is much close to 0.

[tianleiwu]

For this "When x i is much smaller than 0, e x p ( x i ) will be close to 1. When all elements are much smaller than 0, the result will be close to 1 / (1 + N).", is it a typo "When x i is much smaller than 0"? It should be x i is much close to 0.

When x_i is much smaller than 0, exp(x_i) will be close to 0. When all elements are much smaller than 0, the result will be close to 0.
When x_i is close to 0, exp(x_i) will be close to 1. When all elements are close to 0, the result will be close to 1 / (1+N).