Puntuaciones escaladas, normalizadas sobre las claves

logits  = np.einsum('bhsd,bhtd->bhst', q, k) / np.sqrt(D_k)
weights = softmax(logits, axis=-1)
values  = np.einsum('bhst,bhtd->bhsd', weights, v)

En el einsum, s es el token consulta, t el token clave y d la característica de cabeza. Solo d desaparece en la contracción de las puntuaciones.

L contiene puntuaciones (logits); A contiene pesos normalizados. Dividir por √D_k contrarresta el crecimiento de la varianza del producto punto cuando las varianzas de los componentes son cercanas a uno. El softmax corre sobre las claves, el último eje, y la versión del cuaderno no tiene máscara ni dropout.