fast.ai · Parte 2 · Unidad 2 de 5

Por qué las redes profundas no explotan (cuando están bien hechas)

Entrenar una red de más de unas pocas capas sin cuidado da NaN o ceros por todos lados. Esta unidad es la explicación matemática precisa de por qué pasa eso, y las tres técnicas (inicialización, normalización, skip connections) que lo resuelven.

Antes de leerU1 §3, backprop · lección 8 §3, batchnorm en una CNN

🎯 Al terminar vas a poder

Explicar con tus palabras por qué multiplicar 50 veces por una matriz aleatoria explota o colapsa, y derivar en tres líneas de dónde sale el 1/√n_in de Xavier y el √2 extra de Kaiming.
Inicializar un modelo en PyTorch con nn.init.kaiming_normal_ vía model.apply, y saber cuándo no hacerlo (pesos preentrenados).
Colgar un register_forward_hook para medir media, desvío e histograma de las activaciones de cualquier capa, y leer esos números como diagnóstico.
Distinguir LayerNorm de BatchNorm por sobre qué dimensiones promedian, y por qué BatchNorm necesita model.eval().
Implementar LSUV (medir, corregir, repetir) y explicar qué problema resuelve un ResBlock que la init y la norm no resuelven.

1. El problema: 50 capas sin cuidado dan NaN o cero

La demostración es brutal y directa:

x = torch.randn(200, 100)
for i in range(50): x = x @ torch.randn(100,100)
x[0:5,0:5]  # resultado: NaN en todos lados

Multiplicando 50 veces por pesos aleatorios estándar, los valores explotan a infinito (y de ahí a NaN). Reduciendo la escala de los pesos (* 0.01) pasa lo opuesto: todo colapsa a cero. Ninguno de los dos extremos entrena nada — hace falta encontrar la escala exacta que mantenga la varianza de las activaciones estable capa tras capa.

Para ver que es un problema de escala y no de suerte, alcanza con imprimir el desvío estándar después de cada capa: con pesos randn sin escalar, cada matmul multiplica el desvío por aproximadamente √100 = 10. Después de 50 capas eso es 10⁵⁰, muy por encima del máximo representable en float32 (del orden de 3·10³⁸) — de ahí el inf, y de inf − inf sale el NaN. Con * 0.01 el factor por capa es 0.1, y 0.1⁵⁰ es cero a todo efecto práctico. El punto fijo está exactamente en el medio: un factor por capa de 1.

El experimento de las 50 capas en una sola imagen. Es un problema de escala, no de arquitectura: cada matmul multiplica la varianza por n_in·σ_w², así que cualquier factor distinto de 1 se eleva a la 50. Solo la corrección de Kaiming, que compensa la mitad que ReLU descarta, deja la curva plana.

2. Inicialización: Xavier/Glorot y Kaiming/He

Xavier Glorot y Yoshua Bengio (2010) derivaron la escala exacta: si cada capa tiene n_in entradas, escalar los pesos por 1/√n_in mantiene la desviación estándar de las activaciones en 1 a través de capas sucesivas — matemáticamente, porque sumar n productos de variables independientes de media 0 y varianza 1 da una varianza total de n, y dividir por √n la normaliza de vuelta a 1.

La derivación completa entra en tres líneas. Cada salida de una capa lineal es y = Σᵢ wᵢ·xᵢ, una suma de n_in productos:

Si w y x son independientes y de media 0, cada producto tiene media 0 y varianza Var(wᵢ)·Var(xᵢ) = σ_w² · 1.
La varianza de una suma de términos independientes es la suma de varianzas: Var(y) = n_in · σ_w². Por eso la varianza se multiplica por n_in en cada matmul — no es una propiedad de las redes, es aritmética de sumas de variables aleatorias.
Para que Var(y) = 1 hace falta σ_w² = 1/n_in, es decir σ_w = 1/√n_in. Ese es Xavier.

Pero Xavier asume que no hay no-linealidad entre medio. Con ReLU (que descarta la mitad de los valores, los negativos), la varianza se reduce a la mitad en cada capa — por eso Kaiming He et al. (2015, "Delving Deep into Rectifiers") derivaron una escala corregida: √(2/n_in), el factor 2 compensando exactamente lo que ReLU descarta. (Con una distribución simétrica alrededor de 0, E[relu(y)²] = ½·E[y²]: los negativos aportan cero, los positivos aportan la mitad del segundo momento.)

📐 Demostración matemática rigurosa: Xavier vs. Kaiming He

Sea una neurona cuya pre-activación es $y = \sum_{i=1}^{n_{\text{in}}} w_i x_i$ , donde los pesos $w_i$ y las entradas $x_i$ son variables aleatorias independientes e idénticamente distribuidas (i.i.d.), con $\mathbb{E}[w_i] = 0$ .

Esperanza y Varianza de la suma: $\mathbb{E}[y] = \sum_{i=1}^{n_{\text{in}}} \mathbb{E}[w_i]\mathbb{E}[x_i] = 0$ La varianza total se expande como: $\text{Var}(y) = \mathbb{E}[y^2] = \sum_{i=1}^{n_{\text{in}}} \mathbb{E}[w_i^2 x_i^2] + \sum_{i \neq j} \mathbb{E}[w_i w_j x_i x_j]$ Por independencia y al tener $\mathbb{E}[w_i]=0$ , los términos cruzados se anulan idénticamente: $\mathbb{E}[w_i w_j x_i x_j] = \mathbb{E}[w_i]\mathbb{E}[w_j]\mathbb{E}[x_i x_j] = 0$ . Por ende: $\text{Var}(y) = n_{\text{in}} \mathbb{E}[w_i^2]\mathbb{E}[x_i^2] = n_{\text{in}} \text{Var}(w) \mathbb{E}[x^2]$
El régimen lineal / Xavier (Glorot & Bengio, 2010): Si no hay no-linealidad o es simétrica centrada en cero (como Tanh), $\mathbb{E}[x] = 0 \implies \mathbb{E}[x^2] = \text{Var}(x)$ . Para conservar la varianza capa a capa ( $\text{Var}(y) = \text{Var}(x)$ ): $n_{\text{in}} \text{Var}(w) = 1 \implies \text{Var}(w) = \frac{1}{n_{\text{in}}} \implies \sigma_w = \frac{1}{\sqrt{n_{\text{in}}}}$
El colapso bajo ReLU y el factor 2 de Kaiming He (He et al., 2015): Con activación ReLU, la entrada de la capa siguiente es $x = \max(0, y_{\text{prev}})$ . Si la pre-activación $y_{\text{prev}}$ sigue una distribución simétrica alrededor de 0 con densidad $p(y)$ , el segundo momento de $x$ es: $\mathbb{E}[x^2] = \int_{-\infty}^{\infty} (\max(0, y))^2 p(y) \, dy = \int_0^{\infty} y^2 p(y) \, dy$ Como $p(y) = p(-y)$ , la integral sobre el semieje positivo es exactamente la mitad de la integral total: $\mathbb{E}[x^2] = \frac{1}{2} \int_{-\infty}^{\infty} y^2 p(y) \, dy = \frac{1}{2} \mathbb{E}[y_{\text{prev}}^2] = \frac{1}{2} \text{Var}(y_{\text{prev}})$ Sustituyendo $\mathbb{E}[x^2]$ en la ecuación de varianza: $\text{Var}(y) = n_{\text{in}} \text{Var}(w) \cdot \left( \frac{1}{2} \text{Var}(y_{\text{prev}}) \right) = \frac{n_{\text{in}}}{2} \text{Var}(w) \text{Var}(y_{\text{prev}})$ Para que la varianza no varíe ( $\text{Var}(y) = \text{Var}(y_{\text{prev}})$ ), requerimos obligatoriamente: $\frac{n_{\text{in}}}{2} \text{Var}(w) = 1 \iff \text{Var}(w) = \frac{2}{n_{\text{in}}} \iff \sigma_w = \sqrt{\frac{2}{n_{\text{in}}}}$ Si se usara la inicialización de Xavier en una red con ReLU, tras 50 capas la varianza se multiplicaría por $(1/2)^{50} \approx 8.88 \times 10^{-16}$ , aniquilando la señal por completo (underflow).

# Sin corrección: colapsa a cero tras 50 capas
x = torch.randn(200, 100)
for i in range(50): x = relu(x @ (torch.randn(100,100) * 0.1))

# Con Kaiming/He: estable
for i in range(50): x = relu(x @ (torch.randn(100,100) * sqrt(2/100)))

En PyTorch esto es nn.init.kaiming_normal_, y la forma idiomática de aplicarlo a todas las capas es model.apply, que recorre recursivamente cada submódulo:

import torch.nn as nn
from functools import partial

def init_weights(m, leaky=0.):
    if isinstance(m, (nn.Conv1d, nn.Conv2d, nn.Conv3d, nn.Linear)):
        nn.init.kaiming_normal_(m.weight, a=leaky)   # a = pendiente del leak; 0 para ReLU puro

model.apply(partial(init_weights, leaky=0.1))   # visita cada submódulo, in-place

El parámetro a le dice a Kaiming cuánta señal deja pasar la no-linealidad por el lado negativo, porque el factor de corrección deja de ser exactamente 2 cuando hay leak (es 2/(1+a²)). El notebook también introduce GeneralRelu: un ReLU con leak (para no matar completamente los negativos), sub (restar una constante, recentrando la distribución) y un máximo — cada parámetro corrigiendo una asimetría distinta que el ReLU estándar introduce en la distribución de activaciones:

import torch.nn.functional as F

class GeneralRelu(nn.Module):
    def __init__(self, leak=None, sub=None, maxv=None):
        super().__init__()
        self.leak, self.sub, self.maxv = leak, sub, maxv

    def forward(self, x):
        x = F.leaky_relu(x, self.leak) if self.leak is not None else F.relu(x)
        if self.sub  is not None: x = x - self.sub        # recentrar: la media de relu(x) es > 0
        if self.maxv is not None: x = x.clamp_max(self.maxv)  # cortar colas largas
        return x

act = partial(GeneralRelu, leak=0.1, sub=0.4, maxv=6.)   # los valores que usa el notebook

El sub tiene una justificación simple: si la entrada tiene media 0, relu(x) tiene media estrictamente positiva (solo sobreviven los positivos), y esa media positiva se acumula capa tras capa. Restar una constante chica la devuelve cerca de cero. El maxv evita que unos pocos valores enormes dominen la varianza medida.

3. Mirar las activaciones: hooks y ActivationStats

Todo lo anterior es teoría sobre lo que debería pasar con la varianza. Lo que la Parte 2 enseña de verdad es a mirarla: un forward hook es una función que PyTorch llama cada vez que un módulo termina su forward, con acceso a la entrada y a la salida. Con eso podés registrar media, desvío y un histograma de cada capa sin tocar el modelo:

stats = {}   # nombre de capa → lista de (mean, std, histograma)

def make_hook(name):
    def hook(mod, inp, outp):            # firma fija: módulo, tupla de entradas, salida
        a = outp.detach().float()
        stats.setdefault(name, []).append(
            (a.mean().item(), a.std().item(), a.abs().histc(40, 0, 10))   # 40 bins entre 0 y 10
        )
    return hook

handles = [m.register_forward_hook(make_hook(f"capa{i}"))
           for i, m in enumerate(model) if isinstance(m, GeneralRelu)]

model(xb)                                # un forward: los hooks llenan `stats`
for h in handles: h.remove()            # siempre sacarlos, o quedan colgados para siempre

ActivationStats en miniai es exactamente esto envuelto en un callback: cuelga un hook por capa de activación al empezar el entrenamiento, acumula las tuplas en cada batch y al final dibuja tres cosas. El gráfico de medias y desvíos por batch te dice si la escala se mantiene (idealmente media ≈ 0, std ≈ 1 en todas las capas, desde el primer batch). El histograma apilado en el tiempo (el "color_dim" de fast.ai) muestra si la distribución se va concentrando en cero — en cuyo caso la red está muriendo — o se dispersa. Y el porcentaje de activaciones en el primer bin (dead chart) es la métrica más cruda: con ReLU sin corrección, ese porcentaje trepa al 90 % y más en las últimas capas, que es lo que "colapsa a cero" significa concretamente.

Es la herramienta con la que se justifica cada técnica de esta unidad: la init de Kaiming se ve como un primer batch con std ≈ 1 en vez de 0.01; el sub de GeneralRelu se ve como la media bajando a 0; BatchNorm se ve como las curvas quedándose planas durante todo el entrenamiento en vez de solo al principio.

4. Normalización: LayerNorm y BatchNorm

Inicializar bien resuelve el problema al principio del entrenamiento — pero a medida que los pesos cambian, la distribución de activaciones de cada capa vuelve a moverse ("internal covariate shift", el término del paper original de Ioffe & Szegedy 2015). La solución es normalizar activamente, no solo una vez al inicio.

LayerNorm normaliza sobre todos los canales/alto/ancho de cada ejemplo individual:

class LayerNorm(nn.Module):
    def __init__(self, dummy, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.mult = nn.Parameter(torch.tensor(1.))
        self.add  = nn.Parameter(torch.tensor(0.))

    def forward(self, x):
        m = x.mean((1,2,3), keepdim=True)   # un valor por ejemplo del batch
        v = x.var ((1,2,3), keepdim=True)
        x = (x-m) / ((v+self.eps).sqrt())
        return x*self.mult + self.add   # parámetros aprendidos, no fijos

BatchNorm normaliza sobre el batch completo (por canal), y mantiene un promedio corrido de media/varianza para usar en inferencia (cuando no hay un batch completo de donde calcular estadísticas):

def update_stats(self, x):
    m = x.mean((0,2,3), keepdim=True)   # un valor por canal, promediando batch y espacio
    v = x.var ((0,2,3), keepdim=True)
    self.means.lerp_(m, self.mom)  # promedio móvil exponencial
    self.vars.lerp_ (v, self.mom)
    return m,v

def forward(self, x):
    if self.training:
        with torch.no_grad(): m,v = self.update_stats(x)
    else: m,v = self.means, self.vars     # en eval: las estadísticas acumuladas
    x = (x-m) / (v+self.eps).sqrt()
    return x*self.mults + self.adds

La diferencia entre ambas cabe en una tabla — y la columna "dimensiones sobre las que promedia" es la que importa, porque es donde se cometen los errores:

	LayerNorm	BatchNorm
Promedia sobre	C, H, W de cada ejemplo (dims 1,2,3)	N, H, W de cada canal (dims 0,2,3)
Depende del batch	No — cada ejemplo se normaliza solo	Sí — la estadística es del batch entero
Estado entre train y eval	Ninguno, se comporta igual	Promedio corrido; `model.eval()` cambia el comportamiento
Batch chico	Indiferente	Estadísticas ruidosas; con batch de 1 no hay varianza que calcular
Dónde domina hoy	Transformers (sobre la dimensión de embedding)	CNNs de visión (ResNets)

Con BatchNorm el learning rate estable sube notablemente (de 0.2 a 0.4 en el experimento del notebook, con el resto de hiperparámetros iguales) — normalizar activamente permite pasos de gradiente más agresivos sin que el entrenamiento se desestabilice.

La única diferencia que importa entre las dos normalizaciones es la dirección en la que promedian. Una fila (ámbar) es un ejemplo con todos sus canales; una columna (teal) es un canal a través de todo el batch. De ahí salen todas las demás: por qué BatchNorm necesita model.eval(), por qué falla con batch de 1, y por qué los transformers usan la otra.

🔧 La LayerNorm que ya tenés adentro

En un transformer como Gemma la "LayerNorm" (en su variante RMSNorm, que omite la resta de la media) promedia sobre la última dimensión — el vector de embedding de cada token — no sobre las tres últimas como en la versión de imágenes de arriba. Es la misma operación con otro dim: el principio es "normalizar cada unidad de datos por sí sola", y qué es una unidad depende de la arquitectura (una imagen acá, un token allá). Cuando cuantizás a NVFP4 esas capas de normalización son justamente las que se suelen dejar en precisión alta: son baratas y su salida fija la escala de todo lo que sigue.

5. LSUV: inicializar mirando las activaciones reales

Layer-wise Sequential Unit-Variance (Mishkin & Matas, "All You Need is a Good Init") es un enfoque distinto: en vez de calcular la escala correcta analíticamente, se ejecuta un forward pass real y se ajusta cada capa hasta que su output tenga media 0 y varianza 1 medida, no derivada por fórmula. Usa el mismo mecanismo de hooks de §3 — un hook que guarda media y desvío de la salida de la capa de activación — y un loop que corrige la conv que la alimenta hasta converger:

class Hook:
    def __init__(self, m, f): self.hook = m.register_forward_hook(partial(f, self))
    def remove(self): self.hook.remove()

def _lsuv_stats(hook, mod, inp, outp):
    acts = outp.detach().float()
    hook.mean, hook.std = acts.mean().item(), acts.std().item()

def lsuv_init(model, m, m_in, xb, tol=1e-3):
    # m: la capa cuya salida medimos (la activación); m_in: la conv/linear que la alimenta
    h = Hook(m, _lsuv_stats)
    with torch.no_grad():
        model(xb)                                   # primer forward: llena h.mean y h.std
        while abs(h.std - 1) > tol or abs(h.mean) > tol:
            m_in.bias -= h.mean                     # corregir la media con el bias
            m_in.weight.data /= h.std               # corregir el desvío con la escala de los pesos
            model(xb)                               # volver a medir con los pesos ajustados
    h.remove()

# en el notebook: una (activación, conv) por bloque, en orden, con un solo batch real
for act, conv in zip(relus, convs): lsuv_init(model, act, conv, xb)

Dos detalles que importan. El orden es secuencial y de la primera capa a la última (de ahí el "layer-wise sequential"): cuando ajustás la capa 3, las capas 1 y 2 ya están corregidas, así que la entrada que ve es la que va a ver en el entrenamiento real. Y se corrige la conv anterior a la activación, no la activación misma: restar la media al bias y dividir los pesos por el desvío son las dos únicas perillas lineales que hay, y mueven la salida de la ReLU exactamente en la dirección deseada.

Es un enfoque empírico y directo: medí, ajustá, repetí hasta converger — útil sobre todo cuando la arquitectura es lo suficientemente rara como para que la derivación analítica de Kaiming/He no aplique limpiamente (activaciones nuevas, bloques con normalización adentro, capas con escalas heredadas).

6. ResNets: skip connections

Incluso con init y norm perfectos, apilar demasiadas capas puede empeorar el resultado — no por overfitting, sino porque el gradiente tiene que atravesar todas las capas para llegar a las primeras, y se degrada en el camino. La solución de Kaiming He et al. (2015, "Deep Residual Learning"): cada bloque no aprende la transformación completa, aprende solo la diferencia respecto a la identidad, sumada al input original:

class ResBlock(nn.Module):
    def __init__(self, ni, nf, stride=1, act=act_gr, norm=nn.BatchNorm2d):
        super().__init__()
        self.convs = _conv_block(ni, nf, stride, act=act, norm=norm)   # dos convoluciones
        self.idconv = fc.noop if ni==nf else conv(ni, nf, ks=1, stride=1, act=None)
        self.pool = fc.noop if stride==1 else nn.AvgPool2d(2, ceil_mode=True)
        self.act = act()

    def forward(self, x): return self.act(self.convs(x) + self.idconv(self.pool(x)))

La suma self.convs(x) + self.idconv(self.pool(x)) es el "skip connection": si el bloque completo no aporta nada útil, puede converger a aproximar la función identidad (salida ≈ entrada) en vez de tener que "recordar" activamente cómo pasar la información sin distorsión — mucho más fácil de aprender por gradiente que la identidad exacta desde cero. Los dos noop existen para que la suma cierre: si cambia el número de canales hace falta una conv 1×1 en el atajo, y si cambia la resolución hace falta un pooling; cuando nada cambia, el atajo es literalmente x.

Vista desde el gradiente, la razón por la que esto ayuda es todavía más directa: la derivada de f(x) + x respecto de x es f'(x) + 1. Ese +1 es un camino por el que el gradiente llega a las primeras capas sin atravesar ninguna multiplicación — aunque f' sea diminuto en 50 bloques seguidos, el término identidad no se degrada. Con ResBlocks + BatchNorm + OneCycleLR, el notebook llega a resultados que igualan modelos de timm (la librería de arquitecturas de visión preentrenadas de referencia) entrenados desde cero.

📐 Demostración matemática: Por qué la identidad garantiza el flujo de gradiente

En una red secuencial profunda estándar sin atajos, la salida de la capa $L$ es el producto sucesivo de transformaciones: $\mathbf{x}_L = \mathcal{F}_{L-1}(\mathcal{F}_{L-2}(\dots \mathcal{F}_l(\mathbf{x}_l)))$ . Por la regla de la cadena, el gradiente hacia una capa intermedia $l$ es:

\frac{\partial \mathcal{L}}{\partial \mathbf{x}_l} = \frac{\partial \mathcal{L}}{\partial \mathbf{x}_L} \prod_{i=l}^{L-1} \frac{\partial \mathbf{x}_{i+1}}{\partial \mathbf{x}_i} = \frac{\partial \mathcal{L}}{\partial \mathbf{x}_L} \prod_{i=l}^{L-1} \mathbf{W}_i^\top

Si el radio espectral de los pesos $\rho(\mathbf{W}_i) < 1$ , el gradiente decae exponencialmente como $\rho^{L-l} \to 0$ (gradiente desvanecido). Si $\rho(\mathbf{W}_i) > 1$ , explota hacia $\infty$ .

En una ResNet, la ecuación recurrente de cada bloque residual es $\mathbf{x}_{i+1} = \mathbf{x}_i + \mathcal{F}(\mathbf{x}_i, \mathcal{W}_i)$ . Desarrollando la recursión hasta la capa final $L$ :

\mathbf{x}_L = \mathbf{x}_l + \sum_{i=l}^{L-1} \mathcal{F}(\mathbf{x}_i, \mathcal{W}_i)

Al derivar el costo $\mathcal{L}$ respecto a la activación de la capa temprana $\mathbf{x}_l$ :

\frac{\partial \mathcal{L}}{\partial \mathbf{x}_l} = \frac{\partial \mathcal{L}}{\partial \mathbf{x}_L} \frac{\partial \mathbf{x}_L}{\partial \mathbf{x}_l} = \frac{\partial \mathcal{L}}{\partial \mathbf{x}_L} \left( \mathbf{I} + \frac{\partial}{\partial \mathbf{x}_l} \sum_{i=l}^{L-1} \mathcal{F}(\mathbf{x}_i, \mathcal{W}_i) \right)

Esta ecuación formal revela el secreto de la estabilidad:

El término $\frac{\partial \mathcal{L}}{\partial \mathbf{x}_L} \cdot \mathbf{I}$ fluye directamente sin pasar por ninguna multiplicación de matrices de peso intermedio.
Incluso si todas las capas intermedias colapsan o se saturan y $\frac{\partial}{\partial \mathbf{x}_l} \sum_{i=l}^{L-1} \mathcal{F}_i \to \mathbf{0}$ , el gradiente total nunca se anula: $\frac{\partial \mathcal{L}}{\partial \mathbf{x}_l} \to \frac{\partial \mathcal{L}}{\partial \mathbf{x}_L} \neq \mathbf{0}$ .
Esta misma autopista aditiva de gradientes es la razón por la cual los Transformers modernos (con bloques residuales $x + \text{Attention}(x)$ y $x + \text{MLP}(x)$ ) pueden entrenarse con 80, 96 o más capas de profundidad sin sufrir desvanecimiento de gradiente.

Por qué una ResNet de 50 bloques entrena y una red plana de 50 capas no. Cada bloque aprende una diferencia respecto de la identidad (ámbar), y la identidad (teal) viaja intacta en el forward y en el backward: ese +1 en la derivada es el camino limpio del gradiente, y es la misma construcción que sostiene a un transformer.

🟢 Es la misma idea que sostiene a un transformer

Cada bloque de un transformer es x = x + attn(norm(x)) seguido de x = x + mlp(norm(x)): dos residuales por bloque, exactamente la construcción de ResBlock. Sin ese "residual stream" no se podrían apilar los 30 a 80 bloques que tienen los modelos de lenguaje actuales. Lo vas a ver en código en la Unidad 4, donde el residual es además el sitio de un bug clásico.

7. Errores comunes

Casi todos tienen el mismo diagnóstico: colgá un hook (§3) y mirá media y desvío antes de tocar nada más.

Síntoma	Causa probable	Qué hacer
Cargaste un modelo preentrenado y "no sabe nada": la loss arranca como si fuera aleatorio.	Llamaste a `model.apply(init_weights)` después de `load_state_dict`: la init pisó los pesos aprendidos.	Inicializar solo capas nuevas (la head que agregaste), o inicializar antes de cargar. Verificá con un hash de un peso antes y después.
Error o NaN en la capa de BatchNorm al entrenar; la loss salta entre batches.	Batch de tamaño 1 (o muy chico): la varianza sobre un solo ejemplo no existe, y con 2-4 ejemplos es puro ruido.	Subir el batch, usar `drop_last=True` para evitar el último batch de 1, o cambiar a GroupNorm/LayerNorm, que no dependen del batch.
El modelo con BatchNorm da resultados distintos (y peores) en inferencia que en validación durante el entrenamiento.	Olvidaste `model.eval()`: BN sigue usando estadísticas del batch actual — y con un batch de 1 en inferencia, eso es basura.	`model.eval()` antes de inferir (y `model.train()` al volver). Lo mismo aplica a Dropout.
Con ReLU y muchas capas, el histograma de activaciones se aplasta contra 0 y la loss no baja.	Init de Xavier (o `randn * 0.01`) sin la corrección de Kaiming: cada ReLU tira la mitad de la varianza y a las 10 capas no queda nada.	`kaiming_normal_` con el `a` correcto, y/o BatchNorm. Confirmá con el dead chart de §3.
LayerNorm "funciona" pero el modelo entrena peor que sin ella.	Normalizaste sobre la dimensión equivocada: por ejemplo, sobre el batch (dim 0) en vez de sobre cada ejemplo, o sobre los canales cuando la unidad natural era el token.	Escribí explícitamente sobre qué dims promediás (`x.mean((1,2,3))` para imágenes, `x.mean(-1)` para tokens) y chequeá que la salida tenga media 0 y std 1 por ejemplo.
Los hooks acumulan memoria hasta reventar la GPU.	Guardaste `outp` entero (con su grafo de autograd) en vez de un escalar `.detach()`eado, o nunca llamaste a `h.remove()`.	Guardar solo estadísticas (`.item()`, histogramas chicos) y sacar los hooks al terminar.

8. Resumen

Sin cuidado, apilar capas hace que las activaciones exploten a NaN o colapsen a cero — un problema de escala, no de arquitectura: cada matmul multiplica la varianza por n_in·σ_w².
Xavier/Glorot deriva la escala correcta sin no-linealidad (1/√n_in); Kaiming/He la corrige para ReLU (factor √2 extra). En PyTorch: kaiming_normal_ vía model.apply.
Los forward hooks permiten medir media, desvío e histograma de cada capa sin tocar el modelo — es el instrumento con el que se justifica todo lo demás.
LayerNorm (por ejemplo) y BatchNorm (por canal, sobre el batch) normalizan activamente durante el entrenamiento, no solo al inicio — permiten learning rates más altos. BatchNorm tiene estado y necesita model.eval().
LSUV inicializa empíricamente, midiendo activaciones reales capa por capa en vez de derivar una fórmula.
ResNets resuelven la degradación de redes muy profundas haciendo que cada bloque aprenda una diferencia respecto a la identidad — el +1 en la derivada es el camino limpio del gradiente, y es la misma construcción que sostiene a los transformers.

9. Ponete a prueba

Respondé antes de abrir cada respuesta. Si fallás una, volvé a la sección indicada.

1. Una capa lineal tiene 512 entradas y va seguida de ReLU. ¿Con qué desvío estándar inicializás sus pesos, y por qué no con 1/√512?

Con √(2/512) ≈ 0.0625. 1/√512 (Xavier) mantiene la varianza de la salida lineal en 1, pero la ReLU descarta los negativos y deja la mitad del segundo momento; sin el factor 2 la varianza se va reduciendo a la mitad en cada capa. (§2)

2. ¿Por qué la varianza de y = Σ wᵢxᵢ es proporcional a n_in y no a √n_in?

Porque las varianzas de términos independientes se suman: n_in términos de varianza σ_w² dan n_in·σ_w². La raíz aparece recién al pasar a desvío estándar, y por eso la corrección de los pesos es 1/√n_in. (§2)

3. Entrenás una CNN con BatchNorm, validás bien, y en producción con imágenes de a una da resultados absurdos. ¿Qué falta?

model.eval(). En modo train BatchNorm normaliza con la media y la varianza del batch actual; con un solo ejemplo esa varianza es ruido o directamente no existe. En eval usa el promedio corrido acumulado durante el entrenamiento. (§4)

4. En LSUV, ¿por qué se modifica el bias y los pesos de la conv anterior a la activación, y no la activación misma?

La activación (ReLU) no tiene parámetros. La conv sí, y sus dos perillas son lineales: restar la media al bias desplaza la salida, dividir los pesos por el desvío la reescala. Se mide después de la activación porque esa es la distribución que ve la capa siguiente. (§5)

5. ¿Qué mide el "dead chart" de ActivationStats y qué valor indica problema?

El porcentaje de activaciones que caen en el primer bin del histograma, es decir, prácticamente cero. Un porcentaje alto (del orden del 90 %) en las últimas capas significa que casi todas las unidades están muertas: la red colapsó, típicamente por ReLU sin corrección de Kaiming. (§3)

6. ¿Qué pasa con el gradiente que llega a la primera capa en 50 bloques residuales, comparado con 50 capas comunes?

En cada bloque residual la derivada es f'(x) + 1; el término identidad se propaga sin multiplicarse por nada, así que aunque cada f' sea diminuto el gradiente no se degrada. Sin residual, es un producto de 50 jacobianos y cualquier factor menor que 1 lo aplasta. (§6)

10. Para practicar

✏️ Ejercicio 1 — Medir la escala capa por capa (CPU, sin datos)

Reproducí el experimento de las 50 capas con tres inicializaciones y mirá el desvío estándar de la activación después de cada una. Tendrías que ver que solo Kaiming se mantiene cerca de 1.

import torch, math
torch.manual_seed(0)
n = 100

def run(scale, capas=50):
    x = torch.randn(200, n)
    stds = []
    for i in range(capas):
        x = torch.relu(x @ (torch.randn(n, n) * scale))
        stds.append(x.std().item())
    return stds

for nombre, scale in [("randn*0.01", 0.01), ("xavier", 1/math.sqrt(n)), ("kaiming", math.sqrt(2/n))]:
    s = run(scale)
    print(f"{nombre:>10}: capa 1 = {s[0]:.3g}   capa 10 = {s[9]:.3g}   capa 50 = {s[-1]:.3g}")

Pista / qué esperar

Xavier cae aproximadamente a la mitad de varianza por capa (el desvío por 1/√2 ≈ 0.71), así que a la capa 50 está en el orden de 1e-8. Kaiming queda en un rango de 0.5 a 1 todo el camino (no es exactamente 1 porque la ReLU también corre la media). Probá con capas=200 para ver que Kaiming sigue viva.

✏️ Ejercicio 2 — Hooks + LSUV sobre un MLP

Armá un MLP de 6 capas con ReLU, colgale hooks a cada ReLU para imprimir media y desvío, y después aplicale LSUV capa por capa. Compará las estadísticas antes y después.

import torch, torch.nn as nn
from functools import partial
torch.manual_seed(0)

capas = []
for _ in range(6): capas += [nn.Linear(256, 256), nn.ReLU()]
model = nn.Sequential(*capas)
xb = torch.randn(64, 256)

def stats(name, mod, inp, outp):
    print(f"{name}: mean={outp.mean():.3f} std={outp.std():.3f}")

def mostrar(titulo):
    print("---", titulo)
    hs = [m.register_forward_hook(partial(stats, f"relu{i}")) for i, m in enumerate(model) if isinstance(m, nn.ReLU)]
    with torch.no_grad(): model(xb)
    for h in hs: h.remove()

class Hook:
    def __init__(self, m, f): self.hook = m.register_forward_hook(partial(f, self))
    def remove(self): self.hook.remove()

def _lsuv_stats(hook, mod, inp, outp):
    hook.mean, hook.std = outp.mean().item(), outp.std().item()

def lsuv_init(model, m, m_in, xb, tol=1e-3):
    h = Hook(m, _lsuv_stats)
    with torch.no_grad():
        model(xb)
        while abs(h.std - 1) > tol or abs(h.mean) > tol:
            m_in.bias -= h.mean
            m_in.weight.data /= h.std
            model(xb)
    h.remove()

mostrar("init por defecto de PyTorch")
lins  = [m for m in model if isinstance(m, nn.Linear)]
relus = [m for m in model if isinstance(m, nn.ReLU)]
for r, l in zip(relus, lins): lsuv_init(model, r, l, xb)
mostrar("después de LSUV")

Pista / qué esperar

Con la init por defecto de nn.Linear (uniforme, escala 1/√n_in, sin factor 2) el desvío va bajando capa a capa. Después de LSUV cada ReLU imprime mean≈0.000 std≈1.000. Ojo: un mean de 0 con std de 1 después de una ReLU implica que el bias quedó negativo — LSUV mueve la distribución para que la mitad de las salidas queden en cero; es exactamente lo que sub hace en GeneralRelu de forma fija.

11. Fuentes y para seguir

Practical Deep Learning for Coders — Part 2 — las lecciones 17 y 18 en video; la 17 es donde se construyen los hooks y ActivationStats en vivo.
fastai/course22p2 — los notebooks 11_initializing.ipynb y 13_resnet.ipynb con el código completo de esta unidad; corré el primero de punta a punta al menos una vez.
Delving Deep into Rectifiers (He et al., 2015) — la derivación de Kaiming; la sección 2.2 son dos páginas y vale la pena leerla después de §2.
Batch Normalization (Ioffe & Szegedy, 2015) — el paper original; el algoritmo 1 es exactamente update_stats de §4.
All You Need is a Good Init (Mishkin & Matas, 2015) — LSUV; corto y directo.
Deep Residual Learning for Image Recognition (He et al., 2015) — ResNet; la figura 1 (la red de 56 capas peor que la de 20) es el problema que §6 resuelve.