Transpuesta, producto exterior y rango

Álgebra lineal para LLMs. Episodio tres: transpuesta, producto exterior y rango. Con matrices ya sabemos multiplicar. Hoy, tres ideas que explican por qué LoRA, low-rank adaptation, entrena tan pocos números y cómo se puede editar un dato de un modelo tocando una sola matriz.

De dónde venimos

Alineación, ep. 2: LoRA

Alineación, ep. 7: ROME edita un dato

En Alineación, episodio dos, LoRA: a una matriz congelada W se le suma B por A, escalado por alfa sobre r. En el dibujo, B es alta y flaca, de d por r; A es baja y ancha, de r por k; y su producto es una matriz grande de d por k. Y en el episodio siete, ROME, rank-one model editing, cambia un dato del modelo sumándole a una matriz una columna por una fila transpuesta. En el dibujo, una columna por una fila da una matriz entera. Al final de hoy vamos a saber por qué las dos son matrices «de rango bajo», y qué quiere decir eso.

Parte 1 · La transpuesta

La operación más simple de todas: dar vuelta una matriz. Aparece en casi todas las fórmulas.

Dar vuelta filas y columnas

La transpuesta de A, A punto T: las filas pasan a ser columnas. Una de dos por tres queda de tres por dos. El elemento cero, dos de A es el elemento dos, cero de la transpuesta. Ahora escribimos los vectores como columnas, de tres por uno. Una fila por una columna: x transpuesta por y es un número, el producto punto del episodio uno. Al revés, columna por fila, sale una matriz de tres por tres: el producto exterior, que vemos enseguida. Y una regla: la transpuesta de un producto da vuelta el orden. x transpuesta y da siete: cuatro, más cero, más tres. x por y transpuesta, de tres por tres. Y la regla se cumple.

A = np.array([[1, 2, 3],
              [4, 5, 6]])
print("A.T =\n", A.T)                  # filas ↔ columnas
print("A[0, 2] =", A[0, 2], " A.T[2, 0] =", A.T[2, 0])

x = np.array([[1], [2], [3]])          # vectores columna: 3 × 1
y = np.array([[4], [0], [1]])
print("xᵀ y =", (x.T @ y).item(), " (1×3 por 3×1)")
print("x yᵀ forma:", (x @ y.T).shape, " (3×1 por 1×3)")

B = np.array([[1, 0], [0, 1], [1, 1]])
print("(A B)ᵀ == Bᵀ Aᵀ:", np.array_equal((A @ B).T, B.T @ A.T))

A.T =
 [[1 4]
 [2 5]
 [3 6]]
A[0, 2] = 3  A.T[2, 0] = 3
xᵀ y = 7  (1×3 por 3×1)
x yᵀ forma: (3, 3)  (3×1 por 1×3)
(A B)ᵀ == Bᵀ Aᵀ: True

La transpuesta, en fórmula

espejo sobre la diagonal

el producto punto, como matrices

el orden se invierte, para que las formas encajen

La definición: el elemento i j de A transpuesta es el elemento j i de A. Es un espejo sobre la diagonal. Con eso, un vector columna transpuesto es una fila. En el dibujo, una fila de uno por tres por una columna de tres por uno: una sola celda. Esa celda es el producto punto. Uno por cuatro, más dos por cero, más tres por uno: siete. Por eso los papers escriben x transpuesta y donde nosotros escribíamos x punto y. Y la transpuesta de A por B es B transpuesta por A transpuesta. Si A es de cuatro por tres y B de tres por dos, A B es de cuatro por dos, y su transpuesta, de dos por cuatro: solo encaja en ese orden.

Parte 2 · El producto exterior

Fila por columna da un número. Columna por fila da una matriz entera, muy particular.

u vᵀ: una columna por una fila

Dos vectores: u es uno, dos, tres; v es uno, cero, menos uno. np punto outer arma la matriz u v transpuesta: el elemento i j es u i por v j. Ahora la aplicamos a un x: dos, uno, uno. Primero, v punto x: cuánto de v hay en x. Y u multiplicado por ese número. Por último, el rango, que definimos en un rato. La matriz: la primera fila es v, la segunda es dos v, la tercera es tres v. M por x da uno, dos, tres. v punto x da uno, y u por uno: lo mismo. Rango, uno.

u = np.array([1, 2, 3])
v = np.array([1, 0, -1])

M = np.outer(u, v)                 # M[i, j] = u[i] * v[j]
print("u vᵀ =\n", M)

x = np.array([2, 1, 1])
print("M @ x      =", M @ x)
print("v·x        =", v @ x)        # cuánto de v tiene x
print("u * (v·x)  =", u * (v @ x))  # escribe u, con ese peso
print("rango(M)   =", np.linalg.matrix_rank(M))

u vᵀ =
 [[ 1  0 -1]
 [ 2  0 -2]
 [ 3  0 -3]]
M @ x      = [1 2 3]
v·x        = 1
u * (v·x)  = [1 2 3]
rango(M)   = 1

u vᵀ, celda por celda

La misma matriz, dibujada. La primera fila es u sub uno, uno, por v: uno, cero, menos uno. La segunda, dos veces v. La tercera, tres veces v. Mirala por columnas: cada columna es u por un número. Nueve números en la tabla, pero alcanzan seis para describirla: los tres de u y los tres de v. Esa es la gracia del producto exterior.

Leer con v, escribir con u

columna × fila

se reagrupa: primero el número $v^\top x$

escribe u, pesado por lo que leyó

La celda i j de u v transpuesta es u i por v j. En el dibujo, la columna u por la fila v transpuesta: cada celda es un solo producto, no una suma. Y lo importante: aplicada a un vector x, la matriz se puede reagrupar. u por v transpuesta por x es u por el número v transpuesta x. En el dibujo: dos, más cero, menos uno. Da uno. Así que el resultado es uno por u. Una matriz de rango uno hace exactamente esto: lee una sola dirección de la entrada, la v, y escribe una sola dirección de salida, la u. Ese es el truco de ROME.

Parte 3 · El rango

Cuánta información tiene realmente una matriz, más allá de su tamaño.

Contar direcciones independientes

matrix rank de NumPy cuenta el rango. La identidad de tres por tres, unos en la diagonal: rango tres. Una de dos por dos donde la segunda fila es el doble de la primera: rango uno. Un producto exterior de seis por seis, rango uno; la suma de dos productos exteriores, rango dos. Una B de seis por dos por una A de dos por seis: la matriz es de seis por seis, pero... Y una matriz de seis por seis al azar. Tres, uno, uno, dos. B por A: rango dos, aunque tenga treinta y seis celdas. Y la al azar, rango seis, el máximo.

rango = np.linalg.matrix_rank
print("identidad 3×3     :", rango(np.eye(3)))
print("[[1,2],[2,4]]     :", rango(np.array([[1, 2], [2, 4]])))

u1, v1 = rng.normal(size=6), rng.normal(size=6)
u2, v2 = rng.normal(size=6), rng.normal(size=6)
print("u1 v1ᵀ            :", rango(np.outer(u1, v1)))
print("u1 v1ᵀ + u2 v2ᵀ   :", rango(np.outer(u1, v1) + np.outer(u2, v2)))

B = rng.normal(size=(6, 2))          # 6 × 2
A = rng.normal(size=(2, 6))          # 2 × 6
print("B @ A (6×6)       :", rango(B @ A))
print("al azar 6×6       :", rango(rng.normal(size=(6, 6))))

identidad 3×3     : 3
[[1,2],[2,4]]     : 1
u1 v1ᵀ            : 1
u1 v1ᵀ + u2 v2ᵀ   : 2
B @ A (6×6)       : 2
al azar 6×6       : 6

Qué es el rango

• Cuántas columnas independientes tiene: las demás se arman con esas
• Rango de $u v^\top$ = 1: todas las columnas son múltiplos de $u$
• Suma de $r$ productos exteriores: rango $\le r$
• $BA$ con $B$ de $d \times r$: rango $\le r$, aunque sea $d \times k$

El rango es la cantidad de columnas independientes: las que no se pueden armar sumando múltiplos de las otras. Da igual contar filas o columnas: el número es el mismo. En u v transpuesta, todas las columnas son múltiplos de u: rango uno. Cada producto exterior agrega, como mucho, una dirección nueva. Sumar r de ellos da rango r o menos. Y B por A es justamente eso, como vemos ahora: una suma de r productos exteriores. Tiene la forma de una matriz grande, pero solo r direcciones.

B A es una suma de productos exteriores

columna $j$ de $B$ por fila $j$ de $A$

cada término aporta una dirección

números para guardar la matriz

Multiplicar B por A también se puede leer como suma: la columna uno de B por la fila uno de A, más la columna dos por la fila dos, y así hasta r. En el dibujo, dos productos exteriores que se suman en B A. Por eso el rango de B A es, como mucho, r. En el dibujo, una de seis por dos por una de dos por seis: treinta y seis celdas, y dos direcciones. Y la cuenta de números. Guardar la matriz entera son d por k. Guardar B y A son r por d más r por k. En el dibujo, con las formas de q proj de Gemma y r igual a ocho: dieciséis mil trescientos ochenta y cuatro más dieciocho mil cuatrocientos treinta y dos.

Cuántos números hay que entrenar (q_proj, r = 8)

En barras. Entrenar la matriz q proj completa son cuatro millones setecientos dieciocho mil quinientos noventa y dos números. LoRA con rango ocho entrena treinta y cuatro mil ochocientos dieciséis: ciento treinta y seis veces menos. La barra casi no se ve. Y la apuesta de LoRA es que el cambio que necesita un fine-tuning cabe en esas pocas direcciones.

• W completa: 4.718.592
• B y A (LoRA): 34.816

Parte 4 · LoRA, con las formas de Gemma

Ahora LoRA de verdad, con el tamaño de una matriz de Gemma.

B arranca en cero; BA tiene rango r

Las formas de q proj de Gemma, dos mil cuarenta y ocho por dos mil trescientos cuatro, y rango ocho. W acá es al azar: lo que importa es la forma. LoRA inicializa B en cero, así que al empezar W más B A es exactamente W: el modelo arranca igual que antes. Hacemos de cuenta que ya entrenó: B con números. El cambio, B por A, tiene la forma completa. Y medimos su rango. Al inicio, igual a W. B A es de dos mil cuarenta y ocho por dos mil trescientos cuatro, de rango ocho. Cuatro millones setecientos mil números contra treinta y cuatro mil ochocientos: ciento treinta y seis veces menos. En el episodio seis, con la SVD, singular value decomposition, vamos a ver que los pesos de Gemma no tienen rango bajo; el cambio de LoRA, sí.

Cada línea, su término:

• d, k, r = 2048, 2304, 8 → $W \in \mathbb{R}^{2048 \times 2304},\; r = 8$ — la forma de q_proj
• B = torch.zeros(d, r) → $B_0 = 0 \;\Rightarrow\; W + B_0 A = W$ — al inicio, el modelo no cambia
• delta = B @ A → $\Delta W = BA \in \mathbb{R}^{2048 \times 2304}$ — grande por fuera...
• matrix_rank(delta) → $\mathrm{rango}(\Delta W) = 8$ — ...pero de rango 8

d, k, r = 2048, 2304, 8            # q_proj de Gemma 2 2B, rango 8
W = torch.randn(d, k)              # congelada (acá, al azar)
B = torch.zeros(d, r)              # LoRA arranca B en cero
A = torch.randn(r, k) * 0.01

print("W + BA == W al inicio:", torch.equal(W + B @ A, W))

B = torch.randn(d, r)              # como si ya hubiera entrenado
delta = B @ A                      # 2048 × 2304
print("forma de BA:", tuple(delta.shape))
print("rango de BA:", torch.linalg.matrix_rank(delta).item())

completa, lora = d * k, r * (d + k)
print(f"números: {completa:,} vs {lora:,}  ({completa / lora:.0f}×)")

W + BA == W al inicio: True
forma de BA: (2048, 2304)
rango de BA: 8
números: 4,718,592 vs 34,816  (136×)

Hu et al., 2021 — LoRA: Low-Rank Adaptation of Large Language Models · arXiv:2106.09685

10.000× · «Compared to GPT-3 175B fine-tuned with Adam, LoRA can reduce the number of trainable parameters by 10,000 times and the GPU memory requirement by 3 times.»

En el paper de LoRA, de Microsoft, la cuenta a escala: comparado con hacer fine-tuning de GPT tres, de ciento setenta y cinco mil millones de parámetros, con Adam, adaptive moment estimation, el optimizador de siempre, LoRA reduce los parámetros entrenables diez mil veces, y la memoria de la GPU, graphics processing unit, la placa que hace las cuentas, tres veces. Es la misma cuenta de la barra, con matrices más grandes y aplicada solo a algunas.

Parte 5 · ROME: editar con rango uno

El producto exterior también sirve para lo contrario: cambiar una sola cosa de una matriz sin tocar el resto.

Cambiar lo que W responde a una clave

Una versión de juguete de ROME. Pensá a W como una memoria: entra una clave, sale un valor. La clave k es uno, uno, cero; el valor nuevo que queremos, cero, cero, cinco. El cambio es un producto exterior: la columna es lo que falta, v menos W k; la fila es k transpuesta, dividida por k transpuesta k, el largo al cuadrado de k. Comprobamos que la clave da el valor nuevo. Y probamos con un x perpendicular a k. Antes, W k daba tres, uno, uno. Después, cero, cero, cinco. Para el x perpendicular, la salida no cambió: menos uno, menos uno, uno, antes y después. Rango del cambio: uno.

Cada línea, su término:

• k = np.array → $k_*,\; v_*$ — la clave y el valor nuevo
• delta = np.outer(v - W @ k, k) / (k @ k) → $\Delta W = \frac{(v_* - W k_*)\, k_*^\top}{k_*^\top k_*}$ — rango uno: columna × fila
• print("W' k → $(W + \Delta W)\, k_* = v_*$ — la clave da el valor nuevo
• x = np.array([1., -1., 0.]) → $x \perp k_* \;\Rightarrow\; \Delta W x = 0$ — lo perpendicular, intacto

W = np.array([[1., 2., 0.],
              [0., 1., 1.],
              [1., 0., 1.]])
k = np.array([1., 1., 0.])           # la clave: «el sujeto»
v = np.array([0., 0., 5.])           # el valor nuevo que queremos
print("W k antes:", W @ k)

delta = np.outer(v - W @ k, k) / (k @ k)    # rango uno
W2 = W + delta
print("W' k     :", W2 @ k)

x = np.array([1., -1., 0.])          # perpendicular a k
print("x ⟂ k:  W x =", W @ x, "  W' x =", W2 @ x)
print("rango(delta) =", np.linalg.matrix_rank(delta))

W k antes: [3. 1. 1.]
W' k     : [0. 0. 5.]
x ⟂ k:  W x = [-1. -1.  1.]   W' x = [-1. -1.  1.]
rango(delta) = 1

Por qué funciona

la fila $k_*^\top$ lee la clave entera

perpendicular: no lee nada

ROME: lee con $C^{-1}k_*$ en vez de $k_*$

Aplicado a la clave: el cambio es la columna por k transpuesta k, sobre k transpuesta k. Eso vale uno, así que suma exactamente lo que faltaba. En el dibujo, k transpuesta k: uno más uno, dos, que se cancela con el divisor. Aplicado a un x perpendicular, k transpuesta x es cero, y el cambio desaparece. En el dibujo: uno, menos uno, cero. Suma cero. La fórmula real de ROME es la del principio. Lambda es la columna, lo que falta escalado; y la fila no es k sino C a la menos uno por k. C es la covarianza de las claves que el modelo ve en texto común, y C a la menos uno es su inversa. Así el cambio molesta lo menos posible a las claves habituales. Sigue siendo una columna por una fila: rango uno.

Meng et al., 2022 — Locating and Editing Factual Associations in GPT · arXiv:2202.05262

rango 1 · «minimize ‖ŴK − V‖ such that Ŵk∗ = v∗ by setting Ŵ = W + Λ(C⁻¹k∗)ᵀ», con C = KKᵀ estimada sobre texto de Wikipedia.

Es el paper que vimos en Alineación, episodio siete. Plantean la matriz de proyección de un MLP, multilayer perceptron, como una memoria de claves y valores, y buscan el cambio que haga que la clave del sujeto dé el valor nuevo, alterando lo menos posible las demás. La solución cerrada es esa suma de rango uno, con C estimada sobre texto de Wikipedia. Un dato editado, una columna por una fila.

Resumiendo

• $(A^\top)_{ij} = A_{ji}$; $x^\top y$ = producto punto; $(AB)^\top = B^\top A^\top$
• $u v^\top$: columna × fila; lee con $v$, escribe con $u$
• Rango = columnas independientes; $u v^\top$ tiene rango 1
• LoRA: $BA = \sum_j b_j a_j^\top$, rango $\le r$; q_proj con r = 8: 136× menos números
• ROME: $\Delta W = \Lambda (C^{-1}k_*)^\top$, un cambio de rango 1
• Próximo: normas, LayerNorm y RMSNorm

Papers del episodio

• Hu et al., 2021 — LoRA: Low-Rank Adaptation of Large Language Models
• Meng et al., 2022 — Locating and Editing Factual Associations in GPT