Problema de desvanecimiento del gradiente en LSTM con PyTorch: ¿culpa de la inicialización?

Opera

Llevo semanas optimizando un modelo LSTM para series temporales financieras (ventana de 60 pasos, 3 capas ocultas de 256 unidades cada una). El problema: después de la época 5, el gradiente se vuelve prácticamente nulo en las capas inferiores (norma < 1e-7). He probado con nn.LSTM(batch_first=True, dropout=0.3), activación tanh por defecto, y optimizador Adam con lr=1e-3. ¿Alguien ha enfrentado esto? Sospecho que la inicialización ortogonal no es suficiente para secuencias tan largas. Adjunto snippet del entrenamiento:

for epoch in range(50):
    for x, y in loader:
        out, _ = model(x)
        loss = criterion(out, y)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        # grad norm aquí ya es casi cero

Surfing

El clipping excesivo puede estar matando los gradientes si el max_norm es muy bajo para tu arquitectura. Prueba aumentarlo a 5.0 o incluso 10.0. Además, verifica que no haya NaNs en las salidas ocultas.

Blues

También revisa la inicialización de los pesos de las puertas. Para LSTM, a veces conviene inicializar el sesgo de la puerta de olvido con valores altos (ej. 1.0 o 2.0) para evitar que el estado celular se sature a cero. Usa nn.init.constant_(lstm.bias_hh_l0[hidden_size:2*hidden_size], 1.0).

Opera

Gracias por las sugerencias. Lo del sesgo de forget gate tiene sentido, pero ¿cómo afecta eso al gradiente de las capas inferiores? Mi preocupación es que incluso en la primera época los gradientes son razonables, luego caen abruptamente. ¿Podría ser la función de pérdida? Uso MSELoss.

Polyglot

Con MSELoss y series temporales, es común que el gradiente se sature si los valores de salida están en rangos muy pequeños. Normaliza tus datos de entrada y salida a media 0 y desviación 1. Además, prueba con nn.L1Loss para ver si cambia la dinámica.

StreetArt

Otra posibilidad: el número de capas (3) puede estar causando desvanecimiento si no hay conexiones residuales. Aunque LSTM está diseñado para mitigarlo, en la práctica con 3 capas y 60 pasos aún sufre. Intenta reducir a 2 capas o añadir batch_norm entre capas LSTM (implementando un wrapper).

Opera

Efectivamente, al normalizar los datos (z-score) el problema mejoró, pero aún noto que la norma del gradiente baja después de la época 10. Implementaré el wrapper con batch_norm y probaré L1Loss. Gracias a todos por las ideas.

Standup

Si aún persiste, considera usar torch.nn.utils.clip_grad_value_ en lugar de clip_norm, o incluso desactivar el clipping temporalmente. A veces el clipping interfiere con la regularización implícita del optimizador.

Fitness

La mejora con z-score sugiere que el problema de escala era un factor, pero la caída de la norma del gradiente después de la época 10 apunta a saturación de activaciones o a un learning rate inadecuado. Prueba a reducir el learning rate gradualmente con un scheduler (StepLR o ReduceLROnPlateau) y verifica la distribución de las activaciones de las celdas LSTM; si muchas están cerca de -1 o 1 (con tanh), la inicialización ortogonal ayuda pero no es suficiente. También, si usas dropout, prueba a desactivarlo temporalmente para descartar que esté causando ruido excesivo. Por último, asegúrate de que la pérdida no se estancó en un mínimo local: visualiza la pérdida en escala logarítmica.

Coffee

Buen análisis. La caída de la norma del gradiente tras la época 10 es típica de saturación en tanh o sigmoid. Prueba inicialización ortogonal con torch.nn.init.orthogonal_ para las matrices de peso recurrentes y distribuir la inicialización de las bias a 0. Además, verifica si estás usando batch_first=True y si la secuencia es muy larga; considera truncar o usar pack_padded_sequence. También, si la pérdida se estanca, prueba con torch.nn.L1Loss combinado con torch.optim.AdamW y un scheduler ReduceLROnPlateau con paciencia 3. Otra opción es cambiar a GRU, que tiene menos parámetros y suele ser más estable.