Andrés Quintero

// blog /

El impacto de la normalización en la arquitectura Transformer

Durante los últimos meses me he puesto a estudiar el libro de Deep Learning de Christopher Bishop y Hugh Bishop. Estudiando la arquitectura Transformer (que me parece muy interesante por diversos principios matemáticos que dan mucho de qué hablar, pero lo dejaré para otros posts) me di cuenta de que al usar la función softmax en la capa de atención, es supremamente importante que los inputs sean normalizados, lo cual despertó mi curiosidad de por qué, y es el origen de este blog.

Cuando estamos hablando de atención, estamos hablando de un concepto donde (en el mundo de las palabras) capturamos la naturaleza de la semántica de cada una de las palabras para con las demás, en el contexto de una frase.

12.1.5 Scaled Self-Attention

Hay que hacer un pequeño refinamiento que debemos realizar a la capa self-attention. Recordemos que los gradientes de la función softmax se vuelven exponencialmente más pequeños para inputs de altas magnitudes.

Demostración

Para una entrada:

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

Tenemos:

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

Si tomamos:

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

Entonces, para P1P_1:

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

Ahora queremos saber qué tanto cambia P1P_1 con respecto a z1z_1:

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

En la ecuación anterior podemos reemplazar el denominador por:

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

Entonces:

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

o:

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

Ahora aplicamos la regla del producto:

∂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}

Teniendo en cuenta:

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

y:

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

Tenemos:

∂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}

Pero:

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

por lo tanto:

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

Reemplazando:

∂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}

Entonces:

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

Pero:

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

y:

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

Por lo tanto:

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

Finalmente:

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

Ahora revisemos qué pasa.

Si:

P1=1P_1=1

entonces:

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

Si:

P1=0.99P_1=0.99

entonces:

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

Es decir, cuando P1P_1 se acerca a 1, el gradiente se acerca a 0.

Ahora llevemos esto a self-attention.

En self-attention calculamos los scores:

QKTQK^T

Si los vectores tienen una magnitud grande, estos scores también pueden tener una magnitud grande.

Al pasar estos valores por softmax, podemos terminar con probabilidades muy cercanas a 0 y 1.

Por ejemplo:

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

En este punto la función está saturada y sus gradientes son pequeños.

Por esto hacemos un pequeño refinamiento:

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

Dividimos los scores por dk\sqrt{d_k}.

Así reducimos su magnitud antes de aplicar softmax.

Con esto evitamos que softmax se sature tan fácilmente.

Y de aquí viene el nombre:

Scaled Self-Attention.

La idea es bastante simple:

scores grandes→softmax saturada→gradientes pequen˜os\text{scores grandes} \rightarrow \text{softmax saturada} \rightarrow \text{gradientes pequeños}

Por eso necesitamos escalar los scores antes de aplicar softmax.