Gradientes y la regla de la cadena

Álgebra lineal para LLMs. Episodio nueve: gradientes y la regla de la cadena. Cómo se calcula cuánto cambia una pérdida si movés cada uno de los miles de millones de pesos de un modelo. La respuesta es una sola regla, aplicada con matrices.

Las fórmulas que nos traen

Alineación, ep. 2: el gradiente de la entropía cruzada

Alineación, ep. 9: attribution patching

Dos fórmulas de Alineación que usan gradientes. La del episodio dos: la derivada de la pérdida respecto de los logits es la probabilidad predicha, p, menos el one-hot del token correcto, y. La de rulo es la derivada parcial: cuánto cambia la pérdida si movés un solo logit. En el dibujo, con los logits del episodio siete: el token que vino tiene gradiente negativo, los otros positivo, y suman cero. Y la del episodio nueve, attribution patching: el cambio de una métrica se aproxima con el cambio de la activación por la derivada. En el dibujo, la recta que toca la curva. Hoy vemos de dónde salen las dos.

Parte 1 · La derivada es una pendiente

Si movés la entrada un poquito, ¿cuánto se mueve la salida? Ese cociente es la derivada.

La derivada, a mano y con autograd

Una función: f de x igual a x al cuadrado. Movemos x, que vale tres, un paso h, y dividimos lo que cambió f por h, con pasos cada vez más chicos. Con PyTorch: requires grad le pide que registre las operaciones. backward calcula la derivada y la deja en x punto grad. Y usamos la derivada para predecir f en tres coma dos. Con h igual a uno, el cociente da siete; con cero coma uno, seis coma uno; con un milésimo, seis coma cero cero uno. Se acerca a seis, y autograd da seis: dos por x. La predicción da diez coma veinte; el valor real, diez coma veinticuatro.

def f(x):
    return x ** 2

x0 = 3.0
for h in [1.0, 0.1, 0.001]:              # cociente de diferencias
    print(f"h={h:<6} (f(x+h) - f(x)) / h = {(f(x0 + h) - f(x0)) / h:.4f}")

x = torch.tensor(3.0, requires_grad=True)
y = f(x)
y.backward()                              # autograd: dy/dx
print("autograd: f'(3) =", x.grad.item())

d = 0.2                                   # aproximación lineal
print(f"f(3,2) = {f(3.2):.2f}   f(3) + f'(3)*0,2 = {f(3) + 6 * d:.2f}")

h=1.0    (f(x+h) - f(x)) / h = 7.0000
h=0.1    (f(x+h) - f(x)) / h = 6.1000
h=0.001  (f(x+h) - f(x)) / h = 6.0010
autograd: f'(3) = 6.0
f(3,2) = 10.24   f(3) + f'(3)*0,2 = 10.20

Derivada, aproximación lineal y gradiente

x² en x = 3 → 6

9 + 6 · 0,2 = 10,20 · real: 10,24

una derivada parcial por cada entrada

La derivada de f en x, f prima, es el límite del cociente cuando h tiende a cero: en el dibujo, la inclinación de la recta que toca la curva. Para x al cuadrado es dos x: menos dos, dos y seis. Cerca de x, la función se parece a su recta tangente: f de x más delta es casi f de x más la derivada por delta. Con la sigmoide: pega cerca, se despega lejos. Con muchas entradas, cada una tiene su derivada parcial. Juntas forman un vector, el gradiente, que se escribe con nabla, el triángulo invertido. En el dibujo, el de x al cuadrado más y al cuadrado en dos, uno: cuatro y dos.

El gradiente apunta cuesta arriba

f igual a x al cuadrado más y al cuadrado: un cuenco. Las líneas son curvas de nivel, puntos con el mismo valor de f; la estrella es el mínimo. El gradiente es dos x, dos y. En el punto dos, uno, el gradiente es cuatro, dos; dibujamos menos el gradiente, achicado: apunta al centro. En menos dos, uno coma cinco. Y en uno, menos uno coma cinco. Cada flecha cruza las curvas de nivel en ángulo recto: menos el gradiente es la dirección de bajada más empinada. Es lo que usa el episodio diez.

Parte 2 · La regla de la cadena

Una red es una función adentro de otra, adentro de otra. Para derivar el total, se multiplican las derivadas de cada pieza.

Una neurona, derivada a mano

Una neurona del episodio uno de Deep learning, con una pérdida. La entrada x vale dos; el objetivo, t, uno. El forward: z es w por x más b, y la salida y es la sigmoide de z. La pérdida: la distancia al objetivo, al cuadrado. Las tres derivadas locales: L respecto de y; y respecto de z, que para la sigmoide es y por uno menos y; y z respecto de w, que es x. Las multiplicamos. Y comparamos con autograd. y da cero coma seis dos dos cinco; la pérdida, cero coma uno cuatro dos cinco. Las locales: menos cero coma siete cinco cinco, cero coma dos tres cinco y dos. Su producto, menos cero coma tres cinco cinco, es lo que da autograd.

x, t = 2.0, 1.0                          # entrada y objetivo
w = torch.tensor(0.5, requires_grad=True)
b = torch.tensor(-0.5, requires_grad=True)

z = w * x + b                            # forward, paso a paso
y = torch.sigmoid(z)
L = (y - t) ** 2
print(f"z={z.item():.2f}  y={y.item():.4f}  L={L.item():.4f}")

dL_dy = 2 * (y - t)                      # derivadas locales
dy_dz = y * (1 - y)
dz_dw = x
print(f"locales: {dL_dy.item():.4f} {dy_dz.item():.4f} {dz_dw}")
print(f"a mano: dL/dw = {(dL_dy * dy_dz * dz_dw).item():.4f}")

L.backward()                             # autograd
print(f"autograd: dL/dw = {w.grad.item():.4f}  dL/db = {b.grad.item():.4f}")

z=0.50  y=0.6225  L=0.1425
locales: -0.7551 0.2350 2.0
a mano: dL/dw = -0.3549
autograd: dL/dw = -0.3549  dL/db = -0.1774

El grafo de cómputo

Lo mismo como grafo: cajas para valores, círculos para operaciones. Las flechas van de las entradas a la pérdida. El forward, de izquierda a derecha: w por x da uno, más b da z igual a cero coma cinco, la sigmoide da cero coma sesenta y dos, la pérdida, cero coma catorce. El backward, de derecha a izquierda, en rojo. La pérdida respecto de sí misma vale uno. Cada nodo multiplica lo que le llega por su derivada local. La suma reparte igual a sus dos entradas; el producto le pasa a w el gradiente por x, y a x el gradiente por w.

La regla de la cadena, con los números del grafo

x = 2, w = 0,5, b = −0,5, t = 1

una derivada local por eslabón, multiplicadas

0,6225 · 0,3775 = 0,235

La pérdida, una función dentro de otra. En el dibujo, la neurona del snippet: da cero coma seis dos. La regla de la cadena: la derivada del total es el producto de las derivadas de cada eslabón. Dos por y menos t, menos cero coma siete cinco cinco; la derivada de la sigmoide en z, cero coma dos tres cinco; y x, dos. En el dibujo, esa pendiente de la sigmoide en cero coma cinco. La sigmoide tiene una derivada cómoda: ella misma por uno menos ella. En el dibujo, cero coma seis dos por cero coma tres ocho da cero coma dos tres cinco. Por eso el código la calcula con y.

Parte 3 · Jacobianos: la cadena con vectores

Con vectores de entrada y de salida, cada derivada es una matriz. La regla de la cadena pasa a ser un producto de matrices.

Jacobianos en PyTorch

f multiplica por una matriz W de dos por dos. g aplica la tangente hiperbólica, tanh, a cada elemento. jacobian, de torch punto autograd, arma la matriz de todas las derivadas parciales. Comprobamos que el Jacobiano de g después de f es el producto de los dos. Y backward con un vector v calcula v transpuesta por el Jacobiano, sin armarlo. El Jacobiano de f es W. El de g es diagonal: cero coma noventa y nueve y cero coma ocho cinco seis. El producto coincide. Y v transpuesta por J, con v igual a uno, cero, da la primera fila de J: cero coma noventa y nueve y uno coma noventa y ocho.

W = torch.tensor([[1.0, 2.0], [-1.0, 0.5]])
f = lambda x: W @ x                       # R² → R²: lineal
g = lambda u: torch.tanh(u)               # R² → R²: por elemento
x = torch.tensor([0.3, -0.2])

Jf = jacobian(f, x)                       # es W
Jg = jacobian(g, f(x))                    # diagonal: 1 - tanh²
Jgf = jacobian(lambda x: g(f(x)), x)
print("J_f =", Jf.tolist())
print("J_g =", r(Jg))
print("J_g @ J_f == J_(g∘f):", torch.allclose(Jg @ Jf, Jgf))

v = torch.tensor([1.0, 0.0])              # backward = v^T J
xr = x.clone().requires_grad_()
g(f(xr)).backward(v)
print("v^T J      =", r(xr.grad))
print("fila 0 de J =", r(Jgf[0]))

J_f = [[1.0, 2.0], [-1.0, 0.5]]
J_g = [[0.99, 0.0], [0.0, 0.856]]
J_g @ J_f == J_(g∘f): True
v^T J      = [0.99, 1.98]
fila 0 de J = [0.99, 1.98]

El Jacobiano y por qué backprop no lo arma

$J$: $m \times n$, una fila por salida, una columna por entrada

regla de la cadena = producto de Jacobianos

de izquierda a derecha: siempre un vector por una matriz

El Jacobiano, J, tiene una fila por salida y una columna por entrada; la celda i j es cuánto cambia la salida i si movés la entrada j. Es la aproximación lineal con vectores: en el dibujo, J por un cambio chico da el cambio de la salida. La regla de la cadena: el Jacobiano de la composición es el producto de los Jacobianos. El truco de backprop: la pérdida es un solo número, así que se arranca con un vector. Multiplicando de izquierda a derecha, cada paso es un vector por una matriz, nunca dos matrices gigantes.

Forward y backward en una red

La misma idea en una red chica. Tres entradas. Una hidden layer de cuatro neuronas con ReLU, rectified linear unit. Y una salida, de donde sale la pérdida. El backward vuelve por el mismo camino, capa por capa. En cada capa, el vector que llega se multiplica por el Jacobiano local, y de paso se calcula el gradiente de los pesos de esa capa. Cuesta más o menos lo mismo que el forward, para todos los pesos a la vez.

Parte 4 · Backprop con matrices

Ahora la cuenta completa para una red de dos capas, a mano, y comparada con autograd.

Backprop a mano en un MLP

Un MLP, multilayer perceptron, de tres a cuatro a uno. El forward: W uno por x más b uno, ReLU, y W dos por h. La pérdida, la mitad del error al cuadrado. El backward arranca en la salida: la derivada respecto de y es y menos t. La llamamos delta. El gradiente de W dos: el producto exterior de delta con h. Para seguir hacia atrás, delta se multiplica por W dos transpuesta. La ReLU deja pasar el gradiente donde z uno era positivo. Y el gradiente de W uno, otro producto exterior. Los dos coinciden con autograd. Tres de las cuatro ReLU estaban activas. Y cada gradiente tiene la forma de su matriz.

z1 = W1 @ x + b1                  # forward: 3 → 4 → 1
h = torch.relu(z1)
y = W2 @ h + b2
L = 0.5 * ((y - t) ** 2).sum()

d_y = y - t                       # backward a mano
d_W2 = torch.outer(d_y, h)        # δ hᵀ
d_h = W2.T @ d_y                  # W2ᵀ δ
d_z1 = d_h * (z1 > 0)             # la ReLU deja pasar o corta
d_W1 = torch.outer(d_z1, x)       # δ xᵀ

L.backward()                      # autograd, para comparar
print("dW2 iguales:", torch.allclose(d_W2, W2.grad))
print("dW1 iguales:", torch.allclose(d_W1, W1.grad))
print("ReLU activas:", (z1 > 0).int().tolist())
print("formas:", tuple(d_W1.shape), tuple(d_W2.shape))

dW2 iguales: True
dW1 iguales: True
ReLU activas: [1, 1, 0, 1]
formas: (4, 3) (1, 4)

Las cuatro líneas de backprop

producto exterior: 1 × 4

el error vuelve por los mismos pesos, transpuestos

$\odot$: producto elemento a elemento

otro producto exterior: 4 × 3

Delta es el error de la salida. El gradiente de W dos es delta por h transpuesta: columna por fila, el producto exterior del episodio tres. Cada peso recibe el error de su salida por su entrada. Para llevar el error a la hidden layer, se multiplica por W dos transpuesta: los mismos pesos, al revés. El círculo con punto es el producto elemento a elemento, y el uno con corchetes vale uno si z es positivo. En el dibujo, la neurona tres estaba apagada: su error no pasa. Y el gradiente de W uno, otra vez error por entrada transpuesta. Con un batch, cada producto exterior pasa a ser un producto de matrices.

Baydin, Pearlmutter, Radul y Siskind — Automatic differentiation in machine learning: a survey · arXiv:1502.05767

autodiff · «Automatic differentiation (AD), also called algorithmic differentiation or simply “autodiff”, is a family of techniques similar to but more general than backpropagation…»

Backprop es un caso de algo más general. Este survey, publicado en el Journal of Machine Learning Research en dos mil dieciocho, lo dice en el abstract: la diferenciación automática, o autodiff, es una familia de técnicas parecidas a backpropagation pero más generales. Backprop es el modo hacia atrás, el de v transpuesta por J, aplicado a redes. El paper original de backprop lo vimos en Deep learning, episodio cinco.

Attribution patching: cuándo alcanza la recta

Volvemos a attribution patching. Una métrica no lineal de una activación a de cuatro números. Un solo backward da el gradiente. Una dirección de largo uno para mover la activación. El cambio real, corriendo la métrica de nuevo. Y la predicción lineal: el cambio por el gradiente, sin correr nada. Con un salto de cero coma cero cinco, los dos dan cero coma cero uno dos. Con cero coma tres, cero coma cero siete nueve contra cero coma cero siete uno. Con uno coma cinco, cero coma cuatro uno siete contra cero coma tres cinco siete: la recta se queda corta.

def metrica(a):                           # una métrica no lineal de a
    return torch.tanh(a) @ r

a_corr = torch.tensor([0.2, -0.1, 0.4, 0.0], requires_grad=True)
m = metrica(a_corr)
m.backward()                              # un solo backward
grad = a_corr.grad
d = torch.tensor([0.5, -0.5, 0.5, 0.5])   # dirección, largo 1

for escala in [0.05, 0.3, 1.5]:
    delta = escala * d                    # a_limpio - a_corr
    real = metrica(a_corr.detach() + delta) - m.detach()
    lineal = delta @ grad                 # Δ · ∂m/∂a
    print(f"|Δ|={escala}: real {real.item():+.3f}  lineal {lineal.item():+.3f}")

|Δ|=0.05: real +0.012  lineal +0.012
|Δ|=0.3: real +0.079  lineal +0.071
|Δ|=1.5: real +0.417  lineal +0.357

Volvemos a las fórmulas de Alineación

la derivada de LSE es el softmax (ep. 7)

la aproximación lineal: un backward en vez de N forwards

Alineación, ep. 4: la misma cadena, a través del log-softmax

Ahora se leen enteras. La pérdida de un token es menos su logit más el log-sum-exp del episodio siete. La derivada de log-sum-exp es el softmax; la del menos logit, menos uno en el token correcto: p menos y. En el dibujo: cero coma dos cuatro dos menos uno, menos cero coma siete cinco ocho. Attribution patching es la aproximación lineal: buena para saltos chicos, no para grandes. Y el gradiente de política del episodio cuatro de Alineación es el gradiente del log de una probabilidad: la misma cadena, a través del log-softmax, hasta cada peso. En el dibujo, su primer eslabón: one-hot menos p,.

Resumiendo

• Derivada: la pendiente; $f(x+\Delta) \approx f(x) + f'(x)\Delta$
• Gradiente $\nabla f$: una parcial por entrada; $-\nabla f$ es la bajada más empinada
• Regla de la cadena: multiplicar las derivadas locales de cada eslabón
• Jacobiano $J_{ij} = \partial f_i / \partial x_j$; backprop calcula $v^\top J$ sin armar $J$
• MLP: $\nabla W = \delta\, x^\top$ y el error vuelve por $W^\top$
• Próximo: optimización, del descenso por gradiente a Lagrange