2
$\begingroup$

We know that in transformer architecture, we first compute $Q,K$ and $V$ matrices from the token embeddings. Then we compute the attention weights $\alpha$ as $$\alpha = Softmax(QK^T)$$

In the case of masked attention, we can use additive mask as follows

$$\alpha = Softmax(QK^T+UTMask)$$

where $UTMask$ is a strict upper triangular matrix with elements value $-\infty$ . This serves the purpose. However, we could achieve the same with a multiplicative mask with the output of softmax as follows,

$$\alpha = Softmax(QK^T)*LTMask$$ where $LTMask$ is the Lower triangular matrix with elements value $1$ (including the diagonal).

I am just wondering why Masked Multi-head Attention layer in transformer architecture prefers additive mask over multiplicative (Boolean) mask. Is it due to the behaviour of how back-prop for work additive and multiplicative operators? Are there any other reasons?

$\endgroup$

1 Answer 1

1
$\begingroup$

What you propose actually won’t have the same behavior. $Softmax(QK^T) * LTMask$ will do the following operations:

  1. Compute inner products between all pairs of the input sequence.
  2. Compute softmax scores based off of (1).
  3. Zero out softmax scores for time steps larger than the time step index in (2).

This will result in rows in the self-attention to not add up to 1, which is expected. Furthermore, this can cause major issues! Say for example, token 3 and token 8 has the highest inner product (when token 3 is the query vector), in fact it dominates 99%. Then when we mask it out using $LTMask$, we are left with alphas that are $< 1\%$, remaining unscaled.

In the original formulation, this issue is removed, because softmax has a built-in guarantee in most implementations that -inf will be mapped to 0 for numerical stability.

Alternatively, you might have meant to include the $LTMask$ inside the softmax; however, this also doesn’t work, because inner products can be positive or negative.

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.