Lección 2 de 10 · Vectores y matrices · 17 min

Matrices: fila por columna

Formas, una matriz como función, y buscar una fila

Notebook del episodio

El código de esta lección, para correrlo y modificarlo en Google Colab, gratis y con tu cuenta. Para guardar tus cambios: Archivo → Guardar una copia en Drive.

Abrir en Colab Ver en GitHub

Álgebra lineal para LLMs. Episodio dos: matrices. En el episodio uno, el producto punto: dos vectores, un número. Hoy, una matriz: muchos productos punto a la vez. Es la operación de la que está hecho casi todo un modelo de lenguaje.

De dónde venimos

x_0^{(t)} = W_E[\,\text{id}_t\,]

Alineación, ep. 1: buscar una fila

Q = X\, W_Q

Alineación, ep. 1: una proyección

\ell = W_U\, x_L

Alineación, ep. 1: un logit por token

En Alineación, episodio uno, aparecieron tres matrices. W sub E, la tabla de embeddings: en el dibujo, el id de un token elige su fila y la copia. Q igual a X por W sub Q: X tiene una fila por token, y en el dibujo una fila de X recorre una columna de W sub Q y llena una celda. Y los logits: W sub U por el último vector; en el dibujo, cada fila de W sub U por x da el logit de un token.

Parte 1 · Una tabla de números

Una matriz es una tabla rectangular. Lo primero que hay que mirar es su forma.

Filas, columnas y forma

Una matriz A de dos filas y tres columnas: una lista de filas. Su forma, el shape, es dos, tres: siempre primero filas, después columnas. A de uno es la fila uno: un vector de tres números. NumPy cuenta desde cero. Los dos puntos quieren decir «todas las filas»: A de dos puntos coma dos es la columna dos, un vector de dos. Y con dos índices, un solo número. Forma dos, tres. Fila uno: cuatro, cinco, seis. Columna dos: tres, seis. El elemento uno, dos: seis.

python

A = np.array([[1, 2, 3],
              [4, 5, 6]])         # 2 filas, 3 columnas

print("forma    :", A.shape)       # (filas, columnas)
print("fila 1   :", A[1])          # una fila: un vector de 3
print("columna 2:", A[:, 2])       # una columna: un vector de 2
print("A[1, 2]  :", A[1, 2])       # fila 1, columna 2
print("Aᵀ forma :", A.T.shape)     # transpuesta: (3, 2)

text

forma    : (2, 3)
fila 1   : [4 5 6]
columna 2: [3 6]
A[1, 2]  : 6
Aᵀ forma : (3, 2)

A, de 2 × 3

La misma A, dibujada. La fila cero: uno, dos, tres. La fila uno: cuatro, cinco, seis. En matemática se escribe A sub i j: el elemento de la fila i y la columna j. En Gemma, la matriz de embeddings tiene esta misma pinta, con doscientas cincuenta y seis mil filas y dos mil trescientas cuatro columnas.

Parte 2 · Matriz por vector

Multiplicar una matriz por un vector tiene dos lecturas. Las dos dan lo mismo, y las dos sirven.

Dos formas de leer A @ x

Una matriz de dos por dos, y un vector x: uno, dos. Primera lectura: cada fila de A hace un producto punto con x, y cada resultado es una coordenada de la salida. Segunda lectura: x dice cuánto de cada columna de A hay que sumar. Y el arroba hace todo junto. Fila cero por x: dos menos dos, cero. Fila uno por x: uno más dos, tres. Por columnas: también cero, tres. Y A arroba x: cero, tres.

python

A = np.array([[2, -1],
              [1,  1]])
x = np.array([1, 2])

# vista 1: cada fila de A, producto punto con x
for i in range(2):
    print(f"fila {i} · x = {A[i]} · {x} = {A[i] @ x}")

# vista 2: columnas de A, pesadas por las coordenadas de x
y = x[0] * A[:, 0] + x[1] * A[:, 1]
print("1·col0 + 2·col1 =", y)

print("A @ x =", A @ x)

text

fila 0 · x = [ 2 -1] · [1 2] = 0
fila 1 · x = [1 1] · [1 2] = 3
1·col0 + 2·col1 = [0 3]
A @ x = [0 3]

A x, dibujado

y_i = \textstyle\sum_{j} A_{ij}\, x_j

fila $i$ de $A$ , producto punto con $x$

y_1 = 2 \cdot 1 + (-1) \cdot 2 = 0

los números del snippet

y = x_1\, a_{:1} + x_2\, a_{:2}

o: columnas de $A$ pesadas por $x$

La primera lectura en fórmula: la coordenada i de y es la suma, sobre j, de A sub i j por x sub j. Es el producto punto de la fila i con x. En el dibujo, la fila recorre x término a término y llena una celda. Con los números: dos por uno, más menos uno por dos. Cero. La segunda lectura: y es x sub uno por la primera columna de A, más x sub dos por la segunda. A sub dos puntos uno es la columna uno. En el dibujo, cada columna pesada por su coeficiente, y la suma. Es la lectura que usamos en Alineación, episodio uno, para el MLP, multilayer perceptron: una suma de columnas pesadas por neuronas.

Una matriz mueve el plano

La segunda lectura tiene un dibujo. En la grilla, e sub uno es el vector uno, cero; e sub dos, el cero, uno; y entre los dos, el cuadrado unidad. En verde, nuestro x: una vez e sub uno más dos veces e sub dos. Ahora le aplicamos A. e sub uno va a dos, uno: la primera columna de A. e sub dos va a menos uno, uno: la segunda. La grilla entera los sigue: las rectas siguen rectas y paralelas, y el cuadrado se vuelve un paralelogramo. Y x va a una vez la primera columna más dos veces la segunda: cero, tres. Una matriz mueve el plano entero, y queda definida por adónde manda los ejes.

Matrices como funciones

A por e sub uno da la columna cero. Y A por x es la combinación de columnas. Una matriz conocida: rotar noventa grados. Otra: estirar al doble el eje horizontal y achicar a la mitad el vertical. Aplicar primero S y después R es lo mismo que aplicar la matriz R arroba S de una. Rotado, x pasa de uno, dos a menos dos, uno. Estirado, a dos, uno. Y las dos formas de componer dan menos uno, dos.

python

A = np.array([[2, -1],
              [1,  1]])
e1, e2 = np.array([1, 0]), np.array([0, 1])
print("A @ e1 =", A @ e1, " (columna 0)")
print("A @ e2 =", A @ e2, " (columna 1)")

x = np.array([1, 2])                 # x = 1·e1 + 2·e2
print("A @ x  =", A @ x, " = 1·(A e1) + 2·(A e2)")

R = np.array([[0, -1], [1, 0]])      # rotar 90°
S = np.array([[2, 0], [0, 0.5]])     # estirar x, achicar y
print("R @ x  =", R @ x)
print("S @ x  =", S @ x)
print("R @ (S @ x) =", R @ (S @ x), " = (R @ S) @ x =", (R @ S) @ x)

text

A @ e1 = [2 1]  (columna 0)
A @ e2 = [-1  1]  (columna 1)
A @ x  = [0 3]  = 1·(A e1) + 2·(A e2)
R @ x  = [-2  1]
S @ x  = [2. 1.]
R @ (S @ x) = [-1.  2.]  = (R @ S) @ x = [-1.  2.]

Estirar, y después rotar

Las dos del código, en la grilla. S estira al doble el eje horizontal y achica a la mitad el vertical: el cuadrado se vuelve un rectángulo acostado, y x va a dos, uno. R por S: primero estirar, después rotar noventa grados. El rectángulo gira entero y x llega a menos uno, dos, como en la salida.

Parte 3 · Matriz por matriz

Si una matriz por un vector son varios productos punto, una matriz por otra matriz es una grilla entera de productos punto.

Fila por columna, a mano

A, la de dos por tres. B, de tres por dos. Para cada fila i de A y cada columna j de B, el producto punto va a la celda i j de C. El arroba hace lo mismo. B por A, en el otro orden, ni siquiera tiene la misma forma. Y A por A falla: las columnas de la primera tienen que coincidir con las filas de la segunda. Las dos dan cuatro, cinco; diez, once. B por A es de tres por tres. Y A por A, error: tres no es dos.

python

A = np.array([[1, 2, 3], [4, 5, 6]])     # 2 × 3
B = np.array([[1, 0], [0, 1], [1, 1]])   # 3 × 2

C = np.zeros((2, 2), dtype=int)
for i in range(2):                       # fila i de A
    for j in range(2):                   # columna j de B
        C[i, j] = A[i] @ B[:, j]
print("a mano:\n", C)
print("A @ B :\n", A @ B)
print("(B @ A).shape =", (B @ A).shape)  # otro orden, otra cosa

try:
    A @ A                                # 2×3 por 2×3
except ValueError as e:
    print("A @ A: error, 3 ≠ 2")

text

a mano:
 [[ 4  5]
 [10 11]]
A @ B :
 [[ 4  5]
 [10 11]]
(B @ A).shape = (3, 3)
A @ A: error, 3 ≠ 2

El producto de matrices

C_{ij} = \textstyle\sum_{k} A_{ik}\, B_{kj}

fila $i$ de $A$ · columna $j$ de $B$

C_{10} = 4 \cdot 1 + 5 \cdot 0 + 6 \cdot 1 = 10

la celda del dibujo, con los números del snippet

(m \times n)\,(n \times p) = (m \times p)

las del medio tienen que coincidir, y se van

La celda i j de C es la suma, sobre k, de A sub i k por B sub k j: la fila i de A por la columna j de B. En el dibujo, la fila uno de A recorre la columna cero de B. Con números: cuatro por uno, más cinco por cero, más seis por uno. Diez. Y la regla de las formas: m por n, por n por p, da m por p. La n del medio es el largo de cada producto punto, así que tiene que coincidir, y desaparece.

nn.Linear: un batch de tokens

En PyTorch, una capa lineal es una matriz. nn punto Linear de cuatro a tres guarda los pesos con forma tres por cuatro: primero la salida, después la entrada. R a la tres por cuatro: una matriz de tres por cuatro números reales. X apila cinco tokens, cada uno de cuatro números: T por d. A varios ejemplos juntos se les dice batch. La capa calcula X por W transpuesta, más el sesgo b. Y cada fila de la salida depende solo de su fila de entrada. Pesos de tres por cuatro. Cinco por cuatro entra, cinco por tres sale. Igual a la cuenta a mano, y el token dos, solo, da lo mismo que dentro del batch.

Cada línea, su término:

nn.Linear(4, 3) → $W \in \mathbb{R}^{3 \times 4}$ — PyTorch guarda (salida, entrada)
X = torch.randn(5, 4) → $X \in \mathbb{R}^{T \times d},\; T = 5$ — un token por fila
a_mano = X @ lin.weight.T + lin.bias → $Y = X\,W^\top + b$ — lo que calcula la capa
fila = lin(X[2]) → $y_t = x_t\,W^\top + b$ — cada fila, por su cuenta

python

lin = nn.Linear(4, 3)            # entra d=4, sale 3
print("weight:", tuple(lin.weight.shape), " (salida, entrada)")

X = torch.randn(5, 4)            # 5 tokens, uno por fila
Y = lin(X)
print("X:", tuple(X.shape), "→ Y:", tuple(Y.shape))

a_mano = X @ lin.weight.T + lin.bias
print("igual a X @ W.T + b:", torch.allclose(Y, a_mano))

fila = lin(X[2])                  # un solo token, aparte
print("fila 2 suelta = fila 2 del lote:", torch.allclose(fila, Y[2]))

text

weight: (3, 4)  (salida, entrada)
X: (5, 4) → Y: (5, 3)
igual a X @ W.T + b: True
fila 2 suelta = fila 2 del lote: True

Parte 4 · Las matrices de Gemma

Ahora las formas de verdad: las matrices de la primera capa de Gemma dos, dos B.

Las formas, sin cargar nada

Seis matrices: los embeddings y cinco de la capa cero. safetensors nos deja leer la forma de cada tensor sin traerlo a memoria. Y la imprimimos, filas por columnas. Embeddings, doscientos cincuenta y seis mil por dos mil trescientos cuatro. q proj, dos mil cuarenta y ocho por dos mil trescientos cuatro. k proj, la mitad: mil veinticuatro. up proj sube a nueve mil doscientos dieciséis, y down proj baja de nuevo. Todas guardadas como salida por entrada.

python

nombres = ["model.embed_tokens.weight",
           "model.layers.0.self_attn.q_proj.weight",
           "model.layers.0.self_attn.k_proj.weight",
           "model.layers.0.self_attn.o_proj.weight",
           "model.layers.0.mlp.up_proj.weight",
           "model.layers.0.mlp.down_proj.weight"]
for n in nombres:
    archivo = hf_hub_download(REPO, mapa[n])
    with safe_open(archivo, framework="pt") as f:
        forma = f.get_slice(n).get_shape()   # sin cargar el tensor
    corto = n.split(".")[-2]              # q_proj, up_proj, ...
    print(f"{corto:13} {forma[0]:>6} × {forma[1]}")

text

embed_tokens  256000 × 2304
q_proj          2048 × 2304
k_proj          1024 × 2304
o_proj          2304 × 2048
up_proj         9216 × 2304
down_proj       2304 × 9216

Por dónde pasa un token, en formas

El viaje de un token, contado en formas. Un entero, el id, elige una fila: dos mil trescientos cuatro números. La atención lo proyecta a consultas de dos mil cuarenta y ocho, que son ocho attention heads de doscientos cincuenta y seis. Las claves y los valores tienen solo cuatro heads, mil veinticuatro: es GQA, grouped-query attention, dos heads de consulta comparten cada clave. Y o proj lo devuelve a dos mil trescientos cuatro. El MLP sube a nueve mil doscientos dieciséis y baja otra vez. Y al final, doscientos cincuenta y seis mil logits, uno por token.

id: 1 número → x₀: 2304 → q: 2048 = 8×256 → k, v: 1024 = 4×256 → atención → 2304 → up: 9216 → down: 2304 → ℓ: 256.000

Buscar una fila = multiplicar por one-hot

Buscar una fila también es multiplicar. Armamos un vector one-hot: doscientos cincuenta y seis mil ceros, y un uno en la posición del token perro. El vector por W sub E es una suma de filas pesadas por sus coordenadas: todas pesan cero salvo la fila i. Así que da exactamente la fila, que es lo que hace el indexado, sin ninguna cuenta. Un detalle de Gemma: en el forward, el embedding se multiplica por raíz de dos mil trescientos cuatro, cuarenta y ocho. Idénticas, bit a bit. La fila mide uno coma cincuenta y tres; escalada, setenta y tres coma seis.

Cada línea, su término:

uno[i] = 1.0 → $e_i = (0, \dots, 0, 1, 0, \dots, 0)$ — one-hot: un 1 en la posición $i$
por_matriz = uno @ W_E → $e_i^\top W_E = \textstyle\sum_k (e_i)_k\, W_E[k] = W_E[i]$ — solo sobrevive la fila $i$
buscar = W_E[i] → $W_E[i]$ — lo mismo, sin cuentas
x0 = buscar * 2304 ** 0.5 → $x_0 = \sqrt{2304}\; W_E[i] = 48\, W_E[i]$ — Gemma escala el embedding

python

(i,) = tok(" perro", add_special_tokens=False).input_ids
print("id de ' perro':", i)

uno = torch.zeros(256_000)
uno[i] = 1.0                         # one-hot: un 1 y 255.999 ceros
por_matriz = uno @ W_E               # (256000) @ (256000 × 2304)
buscar = W_E[i]                      # la fila i, sin cuentas
print("forma:", tuple(por_matriz.shape))
print("iguales:", torch.equal(por_matriz, buscar))

x0 = buscar * 2304 ** 0.5            # Gemma la escala por √d = 48
print("‖fila‖ =", round(buscar.norm().item(), 2),
      " ‖x0‖ =", round(x0.norm().item(), 1))

text

id de ' perro': 40472
forma: (2304,)
iguales: True
‖fila‖ = 1.53  ‖x0‖ = 73.6

Puntuar: logits = W_U x

En Gemma, la matriz de salida es la misma que la de entrada: tied embeddings. Un vector de prueba: gato más perro. Es de prueba, no lo produce el modelo. W sub U por x: doscientos cincuenta y seis mil productos punto, uno por fila. El logit de mascota es la fila de mascota por x. Y los cuatro más altos. Mascota, uno coma setenta y tres, igual de las dos formas. Arriba de todo, tokens raros de textos viejos, con la ese larga. Gato recién cuarto. Mirá las normas: esas filas miden más de cuatro y medio; la de gato, uno coma seis. El producto punto premia el largo, no solo la dirección.

python

W_U = W_E                        # Gemma: la misma matriz entra y sale
x = vec(" gato") + vec(" perro") # un vector de 2304, de prueba

logits = W_U @ x                 # (256000 × 2304) @ (2304) → 256000
print("logits:", tuple(logits.shape))

k = tok(" mascota", add_special_tokens=False).input_ids[0]
print("ℓ[mascota] =", round(logits[k].item(), 2),
      "= fila·x =", round((W_U[k] @ x).item(), 2))

for t in logits.topk(4).indices:     # los 4 logits más altos
    print(f"{tok.decode(int(t)):10} ℓ = {logits[t]:.2f}"
          f"   ‖fila‖ = {W_U[t].norm():.2f}")
print("‖fila‖ promedio:", round(W_U.norm(dim=1).mean().item(), 2))

text

logits: (256000,)
ℓ[mascota] = 1.73 = fila·x = 1.73
 Monfieur  ℓ = 3.59   ‖fila‖ = 4.60
 myſelf    ℓ = 3.51   ‖fila‖ = 4.87
 Theſe     ℓ = 3.38   ‖fila‖ = 4.51
 gato      ℓ = 3.36   ‖fila‖ = 1.63
‖fila‖ promedio: 1.78

Gemma Team, 2024 — Gemma 2: Improving Open Language Models at a Practical Size · arXiv:2408.00118

d_model = 2304 · Tabla 1, modelo 2B: 26 capas, 8 heads y 4 KV heads (GQA), head size 256, feedforward dim 18432 y «tied embedding: yes».

La tabla uno del reporte técnico de Gemma dos dice, para el modelo de dos B: d model, dos mil trescientos cuatro; veintiséis capas; ocho heads de consulta y cuatro de clave y valor; head size, doscientos cincuenta y seis; y embeddings atados, la misma matriz para entrar y salir. El feedforward figura con dieciocho mil cuatrocientos treinta y dos: las dos matrices de subida juntas.

Las tres matrices, leídas enteras

x_0^{(t)} = 48\; e_{\text{id}_t}^\top W_E

buscar = one-hot por la tabla, y escalar

Q = X\, W_Q^\top:\; (T \times 2304)(2304 \times 2048)

q_proj se guarda como 2048 × 2304

\ell_k = W_U[k] \cdot x_L

un producto punto por token del vocabulario

De vuelta a Alineación. El vector inicial de cada token: el one-hot por la tabla, que es buscar su fila, y en Gemma, por cuarenta y ocho. En el dibujo, tres ids eligen tres filas y se apilan: esa es X, de T por dos mil trescientos cuatro. Q igual a X por W sub Q: en el código, X por q proj transpuesta, porque PyTorch guarda salida por entrada. T por dos mil trescientos cuatro, por dos mil trescientos cuatro por dos mil cuarenta y ocho: da T por dos mil cuarenta y ocho. Y cada logit, la fila k de W sub U por el vector. Es un producto punto: las filas largas tienen ventaja.

Resumiendo

Forma $m \times n$ : primero filas, después columnas
$A x$ : cada fila · $x$ , o columnas de $A$ pesadas por $x$
Las columnas de $A$ dicen adónde van los ejes; componer = multiplicar
$C_{ij}$ = fila $i$ de $A$ · columna $j$ de $B$ ; $(m{\times}n)(n{\times}p) = m{\times}p$
Gemma: $W_E$ 256.000 × 2304, q_proj 2048 × 2304; buscar = one-hot · $W_E$
Próximo: transpuesta, producto exterior y rango (LoRA)

Papers del episodio

Gemma Team, 2024 — Gemma 2: Improving Open Language Models at a Practical Size

¿Qué tal esta lección?

Sé el primero en calificarla

Entra para calificar

¿Terminaste la lección?

Márcala para sumar +50 XP y mantener tu racha.

Comentarios

Entra para comentar en el minuto exacto del video y responder dudas.Entrar