In scaled dot product attention, we scale our outputs by dividing the dot product by the square root of the dimensionality of the matrix:

enter image description here

The reason why is stated that this constrains the distribution of the weights of the output to have a standard deviation of 1.

Quoted from https://www.tensorflow.org/tutorials/text/transformer:

For example, consider that $Q$ and $K$ have a mean of 0 and variance of 1. Their matrix multiplication will have a mean of 0 and variance of $d_k$. Hence, square root of $d_k$ is used for scaling (and not any other number) because the matmul of $Q$ and $K$ should have a mean of 0 and variance of 1, and you get a gentler softmax.

Why does this multiplication have a variance of $d_k$?

If I understand this, I will then understand why dividing by $\sqrt({d_k})$ would normalize to 1.

Trying this experiment on 2x2 arrays I get an output of 1.6 variance:

enter image description here

  • $\begingroup$ Even though here programming questions are generally off-topic, in case the code is necessary, it's preferable to provide it rather than an actual screenshot. So, please, if you can still provide the actual code, edit your post to include it. $\endgroup$ – nbro Jan 16 at 23:20

In statistics, if $X$ and $Y$ are independent and randomly distributed variables:

$\mathbb{E}[X + Y] = \mathbb{E}[X] + \mathbb{E}[Y] \\ Var(X + Y) = Var(X) + Var(Y) \\ \mathbb{E}[XY] = \mathbb{E}[X]\mathbb{E}[Y] \\ Var(XY) = (Var(X) + \mathbb{E}[X]^2)(Var(Y) + \mathbb{E}[Y]^2) - \mathbb{E}[X]^2\mathbb{E}[Y]^2$

Let $Q$ and $K$ be random $d_k$ x $d_k$ matrices, where each entry is a some random distribution with $0$ mean and $1$ variance. Every entry is independent from each other.

Since each entry of $Q$ and $K$ have identical distribution, we can focus only on the top-left-most element of $QK$ without loss of generality. The same applies to every other element.

The top-left-most element of $QK$ is $\displaystyle \sum_{i=0}^{d_k} Q_{1,i} K_{i, 1}$.

Since $Q$ and $K$ are independent:

$\mathbb{E}[Q_{1, i} K_{i, 1}] = \mathbb{E}[Q_{1, i}] \mathbb{E}[K_{i, 1}] = 0 \\ Var(Q_{1, i} K_{i, 1}) = (Var(Q_{1, i}) + \mathbb{E}[Q_{1, i}]^2)(Var(K_{i, 1}) + \mathbb{E}[K_{i, 1}]^2) - \mathbb{E}[Q_{1, i}]^2\mathbb{E}[K_{i, 1}]^2 = 1$

And so summing up $d_k$ of them:

$\mathbb{E} \left[\displaystyle \sum_{i=0}^{d_k} Q_{1,i} K_{i, 1} \right] = \displaystyle \sum_{i=0}^{d_k} \mathbb{E} \left[ Q_{1,i} K_{i, 1} \right] = 0 \\ Var\left(\displaystyle \sum_{i=0}^{d_k} Q_{1,i} K_{i, 1} \right) = \displaystyle \sum_{i=0}^{d_k} Var\left( Q_{1,i} K_{i, 1} \right) = d_k$

For your code block, you are computing the dot product of matrices $a$ and $b$, when you should be doing a matrix multiplication (the attention function multiplies $Q$ by $K$ after all, which is the vectorized form of dot-product -- it doesn't actually do dot product). It should work out to unit variance.

edit: the last paragraph is incorrect, as dot is the same as matrix multiply in the above case

  • $\begingroup$ Awesome, thank you this is a great answer. I am still confused on the code part though looking at the docs for numpy the np.dot function should be equivalent to np.matmul in 2 dimensions. $\endgroup$ – Jacob B Dec 8 '20 at 20:40
  • 1
    $\begingroup$ Ah, you are right! dot is the same as matmul. I'm not sure how to explain the code part in that case. Maybe dimension=2 is too small? The sample variance will itself have a variance that reduces with larger dimension size. $\endgroup$ – user3667125 Dec 8 '20 at 20:53
  • 1
    $\begingroup$ Hmmm I see I will try with larger shapes, in any case your answer provided great insight to my initial question so thanks for taking the time! $\endgroup$ – Jacob B Dec 8 '20 at 22:23

It might help to take two small matrices that match the assumptions (mean of zero and variance of one) and just do the matrix multiplication. The dimensionality of K scales Q in the multiplication, scaling the variance simultaneously.

  • $\begingroup$ Hey, I did do this but it does not come out as expected. I have edited the post $\endgroup$ – Jacob B May 21 '20 at 1:09

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

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