Lección 10 de 10 · Entrenar bien · 16 min

Profundidad

Inicialización, gradientes que se apagan o explotan, y residuales

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

Deep learning a secas. Episodio diez, el último: profundidad. En el episodio nueve, una red de seis capas sin normalizar no aprendía nada. Hoy la estiramos a veinte, medimos capa por capa qué se rompe, y vemos los dos arreglos que hicieron posibles las redes profundas: inicializar bien y sumar en vez de reemplazar.

Por qué más capas es más difícil

Hacia adelante: la señal se apaga o explota
Hacia atrás: el gradiente es un producto de 20 factores
Arreglo 1: inicializar con la varianza justa
Arreglo 2: conexiones residuales

En el forward, cada capa multiplica por una matriz. Si agranda un poco los números, veinte capas los hacen enormes; si los achica, los hacen cero. En el backward pasa lo mismo con los gradientes: por la regla de la cadena, el gradiente de la primera capa es un producto con un factor por cada capa. El primer arreglo es elegir la escala de los pesos al azar para que cada capa no agrande ni achique. El segundo es cambiar la arquitectura: cada capa suma su resultado a lo que ya había. Ese es el residual stream de Alineación.

Parte 1 · La señal, hacia adelante

Antes de entrenar nada: ¿qué pasa con un vector que cruza veinte capas con pesos al azar?

Una capa multiplica la varianza

Primero una capa, sin activación. Vectores de doscientos cincuenta y seis números al azar, con varianza uno: qué tan desparramados están alrededor de cero. Pesos también con varianza uno. Y otra vez, con los pesos divididos por raíz de n. Comparamos con raíz de n. Con pesos de varianza uno, la salida tiene un desvío de dieciséis coma cero tres: dieciséis veces más grande que la entrada, que es justo raíz de doscientos cincuenta y seis. Con pesos de varianza uno sobre n, el desvío queda en uno. Una capa no agranda ni achica.

python

n = 256
x = torch.randn(1000, n)               # 1000 entradas, varianza 1

w = torch.randn(n, n)                  # pesos con varianza 1
print("Var(w) = 1   → std de y:", round((x @ w.T).std().item(), 2))

w = torch.randn(n, n) / n ** 0.5       # pesos con varianza 1/n
print("Var(w) = 1/n → std de y:", round((x @ w.T).std().item(), 2))
print("raíz de n =", n ** 0.5)

text

Var(w) = 1   → std de y: 16.03
Var(w) = 1/n → std de y: 1.0
raíz de n = 16.0

De dónde sale el 1/n

y = \textstyle\sum_{i=1}^{n} w_i x_i \;\Rightarrow\; \mathrm{Var}(y) = \sum_{i=1}^{n} \mathrm{Var}(w_i)\,\mathrm{Var}(x_i)

términos independientes, de media 0: las varianzas se suman

\mathrm{Var}(y) = n\,\mathrm{Var}(w)\,\mathrm{Var}(x) \;\Rightarrow\; \mathrm{Var}(w) = \tfrac{1}{n}

Glorot / Xavier: que la capa no cambie la escala

\mathrm{ReLU}\text{ anula la mitad} \;\Rightarrow\; \mathrm{Var}(w) = \tfrac{2}{n}

Kaiming / He: compensar lo que la ReLU tira

Una neurona suma n productos de peso por entrada. Si son independientes y de media cero, las varianzas se suman, y la de cada producto es la del peso por la de la entrada. En el dibujo, seis términos de varianza uno: la salida tiene varianza seis. Como los n términos son iguales, queda n por varianza de w por varianza de x. Para que la salida tenga la misma varianza que la entrada, la varianza de w tiene que ser uno sobre n. En el dibujo, cada barra aporta un sexto, y suman uno. La ReLU, rectified linear unit, pone en cero la mitad negativa, y con eso se va la mitad de la varianza. Se compensa con dos sobre n.

Veinte capas, sin entrenar

Ahora veinte capas seguidas. Multiplicamos por una matriz al azar y aplicamos la activación, veinte veces. Después de cada capa medimos el desvío: cuánta señal queda. Tres casos. Tangente hiperbólica con pesos chiquitos, de desvío cero coma cero uno. ReLU con varianza uno sobre n, y ReLU con dos sobre n. Capas uno, cinco, diez y veinte. Con pesos chiquitos, en la cinco ya no queda nada. Con uno sobre n, la ReLU pierde de a poco: en la veinte queda cero coma cero cero cero siete. Con dos sobre n, cero coma siete.

python

n, capas = 256, 20
x0 = torch.randn(512, n)                    # un batch de 512

def propagar(std_w, act):
    torch.manual_seed(1)                    # mismos pesos al azar
    x, stds = x0, []
    for _ in range(capas):                  # 20 capas sin entrenar
        W = torch.randn(n, n) * std_w
        x = act(x @ W.T)
        stds.append(x.std().item())         # ¿cuánta señal queda?
    return stds

casos = {"tanh 0,01": (0.01, torch.tanh),           # std de los pesos
         "ReLU 1/n": (1 / math.sqrt(n), torch.relu),  # varianza 1/n
         "ReLU 2/n": (math.sqrt(2 / n), torch.relu)}  # varianza 2/n
for nombre, (s, act) in casos.items():
    st = propagar(s, act)
    print(f"{nombre:10}", " ".join(f"{st[i]:.4f}" for i in (0, 4, 9, 19)))

text

tanh 0,01  0.1567 0.0001 0.0000 0.0000
ReLU 1/n   0.5872 0.1509 0.0257 0.0007
ReLU 2/n   0.8305 0.8535 0.8208 0.7068

¿Cuánta señal queda en cada capa?

Las veinte capas de las dos ReLU, en un gráfico. La tangente hiperbólica con pesos chiquitos ni aparece: se muere en tres capas, y desde ahí todas las neuronas dan cero, para cualquier entrada. Con varianza uno sobre n, cada capa pierde una parte, y después de veinte capas casi no queda nada. Con dos sobre n, la línea se mantiene, con altibajos, cerca de cero coma siete. Es lo que ve la primera pasada de entrenamiento.

El init de PyTorch y el de Kaiming

¿Y qué hace PyTorch por defecto? Una capa lineal de sesenta y cuatro a sesenta y cuatro. Sus pesos salen de una distribución uniforme entre menos y más uno sobre raíz de sesenta y cuatro. Medimos el desvío. Y la comparamos con la inicialización de Kaiming He, que pide varianza dos sobre n para ReLU. El peso más grande es cero coma uno dos cinco, la cota. El desvío es cero coma cero siete: la cota dividida por raíz de tres, porque es uniforme. Varianza de un tercio sobre n: menos de lo que pide la ReLU. Por eso la red del episodio nueve se apagaba. Kaiming da cero coma uno siete ocho.

python

capa = nn.Linear(64, 64)                     # init por defecto de PyTorch
cota = 1 / math.sqrt(64)
print("defecto: máx |w| =", round(capa.weight.abs().max().item(), 4))
print("         1/√64   =", cota)
print("         std(w)  =", round(capa.weight.std().item(), 4))
print("         cota/√3 =", round(cota / math.sqrt(3), 4))

nn.init.kaiming_normal_(capa.weight, nonlinearity="relu")
print("Kaiming: std(w)  =", round(capa.weight.std().item(), 4))
print("         √(2/64) =", round(math.sqrt(2 / 64), 4))

text

defecto: máx |w| = 0.125
         1/√64   = 0.125
         std(w)  = 0.0718
         cota/√3 = 0.0722
Kaiming: std(w)  = 0.1785
         √(2/64) = 0.1768

Glorot y Bengio, 2010 — Understanding the difficulty of training deep feedforward neural networks · fuente

1/n · «We find that the logistic sigmoid activation is unsuited for deep networks with random initialization»

La cuenta de la varianza tiene nombre. Xavier Glorot y Yoshua Bengio, en dos mil diez, estudiaron cómo cambian las activaciones y los gradientes capa por capa. Encontraron que la sigmoide no sirve para redes profundas inicializadas al azar, y propusieron una inicialización que converge bastante más rápido. Por eso a veces se la llama Xavier. Cinco años después, Kaiming He y su equipo derivaron la versión que tiene en cuenta la ReLU: el dos sobre n.

Parte 2 · El gradiente, hacia atrás

La red de veinte capas, con el init por defecto. Un solo backward, y medimos el gradiente de cada capa.

El backward cruza todas las capas

Una red profunda, dibujada con seis columnas. La entrada. La primera capa. Las del medio. Más capas. La última. Y la salida. Ahora el backward: el pulso vuelve de la salida hacia la entrada, y en cada capa el gradiente se multiplica por la derivada de esa capa. Si cada factor es menor que uno, llega casi apagado.

La norma del gradiente, capa por capa

Veinte capas lineales de sesenta y cuatro con ReLU, init por defecto. Un backward sobre un batch de dígitos, sin entrenar. Recorremos de la última capa a la primera. Y medimos la norma del gradiente de los pesos de cada una La capa veinte recibe uno coma seis por diez a la menos dos. La diez, diez a la menos seis. Y la primera, cinco coma siete por diez a la menos nueve: seis órdenes de magnitud menos que la última. Es el vanishing gradient: la primera capa prácticamente no aprende.

python

modelo = Profunda("plana")        # 20 × (Linear 64→64 + ReLU), init por defecto
loss = F.cross_entropy(modelo(X[train[:256]]), y[train[:256]])
loss.backward()                   # un solo backward, sin entrenar

for i in (19, 14, 9, 4, 0):       # de la salida hacia la entrada
    g = modelo.capas[i].weight.grad.norm().item()
    print(f"capa {i + 1:2}: norma del gradiente {g:.1e}")

text

capa 20: norma del gradiente 1.6e-02
capa 15: norma del gradiente 1.1e-04
capa 10: norma del gradiente 1.5e-06
capa  5: norma del gradiente 1.6e-08
capa  1: norma del gradiente 5.7e-09

Un producto de veinte factores

\frac{\partial L}{\partial h_0} = \frac{\partial L}{\partial h_{20}} \cdot \frac{\partial h_{20}}{\partial h_{19}} \cdots \frac{\partial h_1}{\partial h_0}

regla de la cadena, una vez por capa

0{,}5^{20} \approx 9{,}5 \cdot 10^{-7}, \qquad 0{,}25^{20} \approx 9{,}1 \cdot 10^{-13}

factores menores que 1: vanishing

1{,}5^{20} \approx 3325

factores mayores que 1: exploding

h sub l es la salida de la capa l. Por la regla de la cadena, el gradiente respecto de la entrada es el gradiente de la salida multiplicado por la derivada de cada capa respecto de la anterior: veinte factores. En el dibujo, la sigmoide: su pendiente nunca pasa de cero coma veinticinco. Si cada factor vale un medio, veinte capas dan nueve coma cinco por diez a la menos siete. Con cero coma veinticinco, por diez a la menos trece. En el dibujo, la exponencial que cae. Eso es el vanishing gradient. Y si cada factor vale uno coma cinco, el producto da tres mil trescientos veinticinco: el exploding gradient. Hace falta que cada factor valga casi uno.

El gradiente de cada capa, en escala logarítmica

Las veinte capas, en escala logarítmica: cada raya del eje vertical es un factor diez. Con el init por defecto, la curva es una rampa: de menos dos en la capa veinte a menos ocho en la primera. Con Kaiming, la línea queda plana, entre cero coma dos y cero coma tres en todas las capas: el factor por capa vale casi uno. Y con bloques residuales con LayerNorm, los de la parte tres, también queda plana, un poco más arriba. Las dos últimas curvas son redes donde todas las capas aprenden a la vez.

Parte 3 · Residuales

La idea que permitió pasar de veinte capas a ciento cincuenta: que cada capa sume, en vez de reemplazar.

El residual stream

Este dibujo es el de Alineación, episodio uno, donde cada capa de un transformer tiene dos bloques: la atención y el MLP, multilayer perceptron, una red densa como las de esta serie. Una cinta por la que viaja el vector, y bloques al costado. Cada bloque lee de la cinta, calcula algo, y lo suma de vuelta: h igual a h más f de h. Por eso la cinta engorda: al final, el vector es la suma de todo lo que escribieron todos los bloques. Eso es el residual stream. Y para el gradiente, la cinta es una autopista: un camino de la salida a la entrada sin ninguna multiplicación.

Por qué el gradiente no se apaga

h_l = h_{l-1} + f_l(h_{l-1})

el bloque suma su resultado

\frac{\partial h_l}{\partial h_{l-1}} = 1 + f_l'(h_{l-1})

el 1 es el camino directo

h_{20} = h_0 + f_1 + f_2 + \cdots + f_{20}

la salida: la entrada más 20 aportes

f sub l es lo que calcula el bloque l: en nuestro código, una LayerNorm, una capa lineal y una ReLU. El bloque no reemplaza a h: le suma f. En el dibujo, dos columnas que se suman. Al derivar h sub l respecto de h sub l menos uno, la h da uno y el bloque da su derivada. Cada factor del producto es uno más algo, no algo solo. En el dibujo: aunque la derivada del bloque sea chica, el factor queda cerca de uno, y veinte factores cerca de uno no se apagan. Y desarmando la suma, la salida es la entrada más los veinte aportes. Es la cinta, como fórmula.

Un bloque residual

El bloque es una clase de PyTorch. Tiene una LayerNorm y una capa lineal. Y el forward devuelve h más ReLU de la capa aplicada a h normalizada. La normalización va adentro del bloque, no sobre la cinta: el llamado pre-norm de los transformers. Apilamos veinte bloques y medimos los gradientes. El bloque no cambia la forma: entra de cuatro por sesenta y cuatro y sale igual. Y los gradientes quedan todos del mismo orden: cero coma nueve en la capa veinte, uno coma cinco en la diez, tres coma ocho en la primera.

python

class BloqueResidual(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.norma = nn.LayerNorm(d)
        self.capa = nn.Linear(d, d)

    def forward(self, h):
        return h + torch.relu(self.capa(self.norma(h)))   # h + f(h)

torch.manual_seed(0)
bloque = BloqueResidual(64)
h = torch.randn(4, 64)
print("entra:", tuple(h.shape), " sale:", tuple(bloque(h).shape))

modelo = Profunda("residual_ln")   # 20 bloques como este
g = grads_por_capa(modelo)
for i in (19, 9, 0):
    print(f"capa {i + 1:2}: norma del gradiente {g[i]:.1e}")

text

entra: (4, 64)  sale: (4, 64)
capa 20: norma del gradiente 9.3e-01
capa 10: norma del gradiente 1.5e+00
capa  1: norma del gradiente 3.8e+00

Las cuatro redes, entrenadas

La prueba final, con dígitos. Veinte hidden layers, cinco epochs, SGD, stochastic gradient descent, con learning rate cero coma cero cinco. Cuatro versiones: plana, con Kaiming, residual sin normalizar y residual con LayerNorm. Entrenamos cada una y medimos en validación. La plana no aprende: dos coma tres cero cinco de pérdida, nueve coma siete por ciento. Con Kaiming arranca, pero despacio: dieciséis coma nueve por ciento en cinco epochs. La residual sin normalizar explota: la pérdida da nan, not a number. Y la residual con LayerNorm llega a noventa y seis por ciento.

python

# 20 hidden layers, 5 epochs, SGD con learning rate 0,05
for modo in ("plana", "kaiming", "residual", "residual_ln"):
    lv, acc = entrenar(Profunda(modo))
    print(f"{modo:12} val {lv:.3f}  acc {acc:.3f}")

text

plana        val 2.305  acc 0.097
kaiming      val 2.050  acc 0.169
residual     val nan  acc 0.089
residual_ln  val 0.188  acc 0.960

Veinte capas, cinco epochs

En barras. La plana, nueve coma siete por ciento: adivinar. Kaiming, dieciséis coma nueve: la señal llega, pero veinte capas siguen siendo difíciles de entrenar con SGD a secas. La residual sin normalización, ocho coma nueve, con la pérdida en nan: sumar veinte veces sin control hace crecer el vector hasta desbordar. Es el exploding del otro lado. Y la residual con LayerNorm, noventa y seis por ciento. La suma evita que el gradiente se apague, y la normalización, que explote.

plana: 9,7 %
Kaiming: 16,9 %
residual: nan · 8,9 %
residual + LayerNorm: 96,0 %

He, Zhang, Ren y Sun, 2015 — Deep Residual Learning for Image Recognition · arXiv:1512.03385

152 capas · «We explicitly reformulate the layers as learning residual functions»

Las conexiones residuales son de dos mil quince, de Kaiming He y su equipo, el mismo He de la inicialización. Cada capa aprende una función residual: lo que hay que sumarle a la entrada. Con eso entrenaron redes de hasta ciento cincuenta y dos capas, ocho veces más profundas que las anteriores, y un conjunto de esas redes logró tres coma cincuenta y siete por ciento de error en el test de ImageNet. Los transformers adoptaron la misma cinta.

Resumiendo

Una capa multiplica la varianza por $n\,\mathrm{Var}(w)$ : init con $1/n$ o $2/n$
El gradiente es un producto de un factor por capa: se apaga o explota
Residual: $h_l = h_{l-1} + f_l(h_{l-1})$ , cada factor es $1 + f'$
20 capas: 9,7 % plana, 96 % residual con LayerNorm
La serie: neurona, capa, pérdida, autograd, loop, regularizar, normalizar
Próximo: la serie «Álgebra lineal para LLMs»

Papers del episodio

¿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