Skip to content

English · Español

05 — Cálculo completo de attention: secuencia de longitud 4 con valores numéricos

Hacemos toda la atención de cabeza única para una secuencia de longitud 4 con números reales. Q, K, V salen de un embedding §A13. Cada paso — QK^T, escalado, máscara, softmax, multiplicación por V — lo escribimos con valores concretos. Después calculamos ∂L/∂Q por la regla de la cadena. Esto convierte la sección §02 (derivación general) en algo numérico que cabe en una hoja.

Anclas: LYNX_CORTEX.md §4 / PHASE 15; theory §02 scaled dot-product; lab §00 attention by hand; Fase 27 (Flash attention reutilizará esta misma aritmética).

Planteamiento — tensores concretos

Secuencia: ["I", "will", "work", "."]. Embebida a d_model = 4. Usaremos attention single-head, single-layer con d_k = d_v = 4.

import numpy as np
np.set_printoptions(precision=3, suppress=True)

X = np.array([
    [ 0.50,  0.30, -0.20,  0.10],   # "I"
    [-0.10,  0.40,  0.20, -0.30],   # "will"
    [ 0.20, -0.10,  0.50,  0.10],   # "work"
    [ 0.00,  0.00,  0.00,  0.20],   # "."
])

# Tiny W_Q, W_K, W_V (identity-ish for clarity)
W_Q = np.eye(4) * 0.5
W_K = np.eye(4) * 0.5
W_V = np.eye(4) * 1.0

Calcula Q, K, V:

Q = X @ W_Q = 0.5 · X
K = X @ W_K = 0.5 · X
V = X @ W_V = X

Así que Q = K = 0.5 · X, V = X. (Elegir pesos identidad hace la aritmética legible.)

Paso 1 — Q K^T

Q @ K.T = (0.5 · X)(0.5 · X).T = 0.25 · X X.T

Calcula X X.T (una matriz 4×4 de productos escalares):

       I       will    work     .
I   |  0.39   -0.04   -0.01    0.02  |
will| -0.04    0.30    0.04   -0.06  |
work| -0.01    0.04    0.31    0.02  |
 .  |  0.02   -0.06    0.02    0.04  |

(Comprobación de sanidad: X[0] · X[0] = 0.25 + 0.09 + 0.04 + 0.01 = 0.39 ✓.)

Q K^T = 0.25 · X X.T:

       I       will    work     .
I   |  0.098  -0.010  -0.003   0.005  |
will| -0.010   0.075   0.010  -0.015  |
work| -0.003   0.010   0.078   0.005  |
 .  |  0.005  -0.015   0.005   0.010  |

Paso 2 — Escalar por sqrt(d_k) = 2

S = (Q K^T) / sqrt(d_k) = (Q K^T) / 2

       I       will    work     .
I   |  0.049  -0.005  -0.002   0.003  |
will| -0.005   0.038   0.005  -0.008  |
work| -0.002   0.005   0.039   0.003  |
 .  |  0.003  -0.008   0.003   0.005  |

Estos son los logits pre-softmax. Como usamos pesos identidad y embeddings pequeños, están todos cerca de 0. (El modelo real de la Fase 17 tendrá una dispersión mayor.)

Paso 3 — Máscara causal

Pon el triángulo superior a -∞:

       I       will    work     .
I   |  0.049   -inf    -inf    -inf   |
will| -0.005   0.038   -inf    -inf   |
work| -0.002   0.005   0.039   -inf   |
 .  |  0.003  -0.008   0.003   0.005  |

Esto hace que el softmax en el paso 4 dé 0 para las posiciones enmascaradas.

Paso 4 — Softmax por filas

Para cada fila, resta el máximo de la fila (por estabilidad), exponencia, normaliza.

Fila 0 ("I"): Solo una entrada sin enmascarar, 0.049. Softmax = [1, 0, 0, 0].

Fila 1 ("will"):

logits = [-0.005, 0.038, -inf, -inf]
shifted = logits - 0.038 = [-0.043, 0, -inf, -inf]
exp     = [exp(-0.043), 1, 0, 0] = [0.958, 1.0, 0, 0]
sum     = 1.958
softmax = [0.489, 0.511, 0, 0]

Fila 2 ("work"):

logits = [-0.002, 0.005, 0.039, -inf]
shifted = logits - 0.039 = [-0.041, -0.034, 0, -inf]
exp     = [0.960, 0.967, 1.0, 0]
sum     = 2.927
softmax = [0.328, 0.330, 0.342, 0]

Fila 3 ("."):

logits = [0.003, -0.008, 0.003, 0.005]
shifted = logits - 0.005 = [-0.002, -0.013, -0.002, 0]
exp     = [0.998, 0.987, 0.998, 1.0]
sum     = 3.983
softmax = [0.250, 0.248, 0.250, 0.251]

(Aproximadamente uniforme — las diferencias en los logits pre-softmax son demasiado pequeñas para hacer la distribución afilada. El §02 de la Fase 15 deriva por qué: en init, attention es aproximadamente uniforme; el aprendizaje la afila.)

Matriz de atención A:

A = | 1.000  0      0      0     |
    | 0.489  0.511  0      0     |
    | 0.328  0.330  0.342  0     |
    | 0.250  0.248  0.250  0.251 |

Paso 5 — A @ V

V = X. Calcula Y = A V:

Fila 0 ("I"): Y[0] = 1.0 · X[0] = X[0] = [0.50, 0.30, -0.20, 0.10]. Ella misma.

Fila 1 ("will"): suma ponderada de X[0] y X[1]:

Y[1] = 0.489·X[0] + 0.511·X[1]
     = 0.489·[0.5, 0.3, -0.2, 0.1] + 0.511·[-0.1, 0.4, 0.2, -0.3]
     = [0.2445 - 0.0511, 0.1467 + 0.2044, -0.0978 + 0.1022, 0.0489 - 0.1533]
     ≈ [0.194, 0.351, 0.004, -0.104]

Fila 2 ("work"):

Y[2] ≈ 0.328·X[0] + 0.330·X[1] + 0.342·X[2]
     ≈ [0.164, 0.099, -0.066, 0.033] + [-0.033, 0.132, 0.066, -0.099] + [0.068, -0.034, 0.171, 0.034]
     ≈ [0.199, 0.197, 0.171, -0.032]

Fila 3 ("."): promedio aproximado de las cuatro filas de X.

Y[3] ≈ 0.250·X[0] + 0.248·X[1] + 0.250·X[2] + 0.251·X[3]
     ≈ [0.150, 0.155, 0.123, 0.020]

Paso 6 — Proyección de salida y pérdida

Supón que el objetivo es predecir el siguiente token (LM). La cabeza de LM da logits sobre el vocab; la pérdida es entropía cruzada. No la expandiremos aquí — la Fase 17 lo hace — pero asume ∂L/∂Y[3] = g_3 (el gradiente en la última posición).

Gradiente backward — a través del softmax

Queremos ∂L/∂Q. La cadena completa es:

  1. Y = A V, así que ∂L/∂A = (∂L/∂Y) V^T y ∂L/∂V = A^T (∂L/∂Y).
  2. A = softmax(S) (por filas). El jacobiano del softmax por fila i es diag(A_i) - A_i A_i^T. Así que ∂L/∂S = A ⊙ (∂L/∂A - (∂L/∂A · A^T) · ⊙_rows) — la fórmula estándar del backward del softmax.
  3. S = (Q K^T) / sqrt(d_k), así que ∂L/∂Q = (∂L/∂S) K / sqrt(d_k).

La aritmética de formas para nuestro caso T = 4, d_k = 4:

∂L/∂Y : (4, 4)        [from cross-entropy]
∂L/∂A : (4, 4) = (∂L/∂Y) V.T = (4,4) @ (4,4)
∂L/∂V : (4, 4) = A.T @ (∂L/∂Y) = (4,4) @ (4,4)
∂L/∂S : (4, 4) = softmax_backward(A, ∂L/∂A)
∂L/∂Q : (4, 4) = (∂L/∂S) @ K / sqrt(d_k)
∂L/∂K : (4, 4) = (∂L/∂S).T @ Q / sqrt(d_k)

Crítico: el /sqrt(d_k) sobrevive al backward — aparece en ∂L/∂Q y ∂L/∂K. Si te olvidas de escalar en el forward, pierdes este factor también en el backward — el /break de la Fase 15 lo expone.

Lo que muestra el ejemplo numérico

  1. En init, attention es aproximadamente uniforme. Cada fila de A tiene entradas en [0.24, 0.51]. El patrón "mírame" no existe todavía.
  2. La máscara causal pone a cero las contribuciones futuras. La fila 0 solo atiende a sí misma; la fila 1 a las posiciones 0–1; etc.
  3. El modelo puede hacer cualquier cosa ajustando W_Q, W_K. Con init aleatorio, attention es uniforme; el entrenamiento la afilará en los patrones que vemos en modelos pre-entrenados (el lab de la Fase 17 los visualiza).
  4. El softmax preserva la forma (T, T). Las dimensiones no cambian dentro de attention — solo en el paso A V donde contraemos sobre el eje T (de keys).

Presupuesto de cómputo para attention de longitud 4

Q K^T            : 4 · 4 · 4 = 64 muls
scale            : 16 ops
softmax (per row): 4 exp + 4 sum + 4 div = ~32 ops per row × 4 = 128
mask             : 6 sets to -inf
A V              : 4 · 4 · 4 = 64 muls

Total           : ~270 muls + ~150 small ops + 4 exponentials per row

Para longitud-32 (una frase realista del §A13): T^2 d_k = 1024 · 4 = 4096 muls por matriz — manejable en CPU. Para longitud-2048 (una historia pequeña): 2048^2 · 64 = 268M muls — necesita la GPU de la Fase 23. Aquí es donde la KV cache de la Fase 22 y Flash attention de la Fase 27 se vuelven críticos.

Citas

  • Vaswani, A. et al. 2017. "Attention is All You Need." arXiv:1706.03762. La sección 3.2 deriva attention; la §3.2.1 explica el escalado sqrt(d_k).
  • El ejemplo numérico trabajado sigue el patrón en "The Annotated Transformer" de Sasha Rush (https://nlp.seas.harvard.edu/2018/04/03/attention.html) pero se calcula de forma independiente para el §A13.

Recapitulación en un párrafo

Un pase de attention single-head de longitud 4 sobre embeddings §A13 calcula Q K^T / sqrt(d_k), aplica una máscara causal, toma softmax por filas para obtener A, luego Y = A V. Con init identidad la matriz de atención es aproximadamente uniforme a lo largo de las posiciones permitidas: fila 0 = self; fila 1 = 50/50; fila 3 = ~uniforme sobre 4. El gradiente backward ∂L/∂Q = (∂L/∂S) K / sqrt(d_k) preserva el factor de escalado — olvidarlo en el forward también rompe el backward. El coste de cómputo es O(T² d) para el bloque de attention; en T = 4 son unos pocos cientos de ops, en T = 2048 es el cuello de botella que abordan la Fase 22 y la Fase 27.

Anterior: 04-masking.md Siguiente: Fase 16 (codificaciones posicionales).