Andrés Quintero

// writing /

The impact of normalization in the Transformer architecture

Over the last few months I have been studying the Deep Learning book by Christopher Bishop and Hugh Bishop. While studying the Transformer architecture (which I find fascinating for the many mathematical principles behind it, but I’ll leave that for other posts) I realized that when using the softmax function in the attention layer, it is extremely important for the inputs to be normalized. That sparked my curiosity about why, and it is the origin of this post.

When we talk about attention, we are talking about a concept where (in the world of words) we capture the semantic nature of each word with respect to all the others, within the context of a sentence.

12.1.5 Scaled Self-Attention

There is a small refinement we need to make to the self-attention layer. Recall that the gradients of the softmax function become exponentially smaller for inputs of large magnitude.

Proof

For an input:

z=[z1,…,zk]z=[z_1,\ldots,z_k]

We have:

P=Softmax(z)=ezi∑j=1kezjP=Softmax(z)=\frac{e^{z_i}}{\sum_{j=1}^{k}e^{z_j}}

If we take:

Pi=Softmax(zi)=ezi∑j=1kezjP_i=Softmax(z_i)=\frac{e^{z_i}}{\sum_{j=1}^{k}e^{z_j}}

Then, for P1P_1:

P1=ez1ez1+⋯+ezkP_1=\frac{e^{z_1}}{e^{z_1}+\cdots+e^{z_k}}

Now we want to know how much P1P_1 changes with respect to z1z_1:

∂P1∂z1\frac{\partial P_1}{\partial z_1}

In the previous equation we can replace the denominator with:

S=ez1+⋯+ezkS=e^{z_1}+\cdots+e^{z_k}

Then:

P1=ez1SP_1=\frac{e^{z_1}}{S}

or:

P1=ez1S−1P_1=e^{z_1}S^{-1}

Now we apply the product rule:

∂P1∂z1=∂ez1∂z1S−1+ez1∂S−1∂z1\frac{\partial P_1}{\partial z_1} = \frac{\partial e^{z_1}}{\partial z_1}S^{-1} + e^{z_1}\frac{\partial S^{-1}}{\partial z_1}

Taking into account:

∂ex∂x=ex\frac{\partial e^x}{\partial x}=e^x

and:

∂x−1∂x=−x−2\frac{\partial x^{-1}}{\partial x}=-x^{-2}

We get:

∂P1∂z1=ez1S−1−ez1S−2∂S∂z1\frac{\partial P_1}{\partial z_1} = e^{z_1}S^{-1} - e^{z_1}S^{-2}\frac{\partial S}{\partial z_1}

But:

S=ez1+⋯+ezkS=e^{z_1}+\cdots+e^{z_k}

therefore:

∂S∂z1=ez1\frac{\partial S}{\partial z_1}=e^{z_1}

Substituting:

∂P1∂z1=ez1S−1−ez1S−2ez1\frac{\partial P_1}{\partial z_1} = e^{z_1}S^{-1} - e^{z_1}S^{-2}e^{z_1}

Then:

∂P1∂z1=ez1S−e2z1S2\frac{\partial P_1}{\partial z_1} = \frac{e^{z_1}}{S} - \frac{e^{2z_1}}{S^2}

But:

P1=ez1SP_1=\frac{e^{z_1}}{S}

and:

P12=e2z1S2P_1^2=\frac{e^{2z_1}}{S^2}

Therefore:

∂P1∂z1=P1−P12\frac{\partial P_1}{\partial z_1} = P_1-P_1^2

Finally:

∂P1∂z1=P1(1−P1)\boxed{ \frac{\partial P_1}{\partial z_1}=P_1(1-P_1) }

Now let’s see what happens.

If:

P1=1P_1=1

then:

∂P1∂z1=1(1−1)≈0\frac{\partial P_1}{\partial z_1} = 1(1-1) \approx 0

If:

P1=0.99P_1=0.99

then:

∂P1∂z1=0.99(1−0.99)\frac{\partial P_1}{\partial z_1} = 0.99(1-0.99) =0.0099=0.0099

That is, when P1P_1 approaches 1, the gradient approaches 0.

Now let’s bring this to self-attention.

In self-attention we compute the scores:

QKTQK^T

If the vectors have a large magnitude, these scores can also have a large magnitude.

When these values go through softmax, we can end up with probabilities very close to 0 and 1.

For example:

P=[0.001,0.998,0.001]P=[0.001,0.998,0.001]

At this point the function is saturated and its gradients are small.

That is why we make a small refinement:

Softmax(QKTdk)\boxed{ Softmax\left(\frac{QK^T}{\sqrt{d_k}}\right) }

We divide the scores by dk\sqrt{d_k}.

This reduces their magnitude before applying softmax.

With this we keep softmax from saturating so easily.

And that is where the name comes from:

Scaled Self-Attention.

The idea is quite simple:

large scores→saturated softmax→small gradients\text{large scores} \rightarrow \text{saturated softmax} \rightarrow \text{small gradients}

That is why we need to scale the scores before applying softmax.