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

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.

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.

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.

# 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 init.kaiming_normal_(m.weight, a=leaky) aplicado a cada capa convolucional. 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.

3. 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 forward(self, x):
        m = x.mean((1,2,3), keepdim=True)
        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)
    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

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.

4. 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:

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

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.

5. 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):
        self.convs = _conv_block(ni, nf, stride)   # dos convoluciones
        self.idconv = fc.noop if ni==nf else conv(ni, nf, ks=1, act=None)
        self.pool = fc.noop if stride==1 else nn.AvgPool2d(2, ceil_mode=True)

    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. 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.

6. Resumen

1. Sin cuidado, apilar capas hace que las activaciones exploten a NaN o colapsen a cero — un problema de escala, no de arquitectura.
2. Xavier/Glorot deriva la escala correcta sin no-linealidad; Kaiming/He la corrige para ReLU (factor √2 extra).
3. LayerNorm y BatchNorm normalizan activamente durante el entrenamiento, no solo al inicio — permiten learning rates más altos.
4. LSUV inicializa empíricamente, midiendo activaciones reales en vez de derivar una fórmula.
5. ResNets resuelven la degradación de redes muy profundas haciendo que cada bloque aprenda una diferencia respecto a la identidad, no la transformación completa.