La pérdida salta abruptamente cuando disminuyo la tasa de aprendizaje con el optimizador Adam en PyTorch

Estoy entrenando una auto-encoderred con Adamoptimizador (con amsgrad=True) y MSE losspara la tarea de separación de fuente de audio de un solo canal. Cada vez que disminuyo la tasa de aprendizaje por un factor, la pérdida de la red salta abruptamente y luego disminuye hasta la próxima disminución en la tasa de aprendizaje.

Estoy usando Pytorch para la implementación y capacitación de redes.

Following are my experimental setups:

 Setup-1: NO learning rate decay, and 
          Using the same Adam optimizer for all epochs

 Setup-2: NO learning rate decay, and 
          Creating a new Adam optimizer with same initial values every epoch

 Setup-3: 0.25 decay in learning rate every 25 epochs, and
          Creating a new Adam optimizer every epoch

 Setup-4: 0.25 decay in learning rate every 25 epochs, and
          NOT creating a new Adam optimizer every time rather
          using PyTorch's "multiStepLR" and "ExponentialLR" decay scheduler 
          every 25 epochs

Estoy obteniendo resultados muy sorprendentes para las configuraciones # 2, # 3, # 4 y no puedo razonar ninguna explicación al respecto. Los siguientes son mis resultados:

Setup-1 Results:

Here I'm NOT decaying the learning rate and 
I'm using the same Adam optimizer. So my results are as expected.
My loss decreases with more epochs.
Below is the loss plot this setup.

Parcela-1:

optimizer = torch.optim.Adam(lr=m_lr,amsgrad=True, ...........)

for epoch in range(num_epochs):
    running_loss = 0.0
    for i in range(num_train):
        train_input_tensor = ..........                    
        train_label_tensor = ..........
        optimizer.zero_grad()
        pred_label_tensor = model(train_input_tensor)
        loss = criterion(pred_label_tensor, train_label_tensor)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    loss_history[m_lr].append(running_loss/num_train)

Setup-2 Results:  

Here I'm NOT decaying the learning rate but every epoch I'm creating a new
Adam optimizer with the same initial parameters.
Here also results show similar behavior as Setup-1.

Because at every epoch a new Adam optimizer is created, so the calculated gradients
for each parameter should be lost, but it seems that this doesnot affect the 
network learning. Can anyone please help on this?

Parcela-2:

for epoch in range(num_epochs):
    optimizer = torch.optim.Adam(lr=m_lr,amsgrad=True, ...........)

    running_loss = 0.0
    for i in range(num_train):
        train_input_tensor = ..........                    
        train_label_tensor = ..........
        optimizer.zero_grad()
        pred_label_tensor = model(train_input_tensor)
        loss = criterion(pred_label_tensor, train_label_tensor)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    loss_history[m_lr].append(running_loss/num_train)

Setup-3 Results: 

As can be seen from the results in below plot, 
my loss jumps every time I decay the learning rate. This is a weird behavior.

If it was happening due to the fact that I'm creating a new Adam 
optimizer every epoch then, it should have happened in Setup #1, #2 as well.
And if it is happening due to the creation of a new Adam optimizer with a new 
learning rate (alpha) every 25 epochs, then the results of Setup #4 below also 
denies such correlation.

Parcela-3:

decay_rate = 0.25
for epoch in range(num_epochs):
    optimizer = torch.optim.Adam(lr=m_lr,amsgrad=True, ...........)

    if epoch % 25 == 0  and epoch != 0:
        lr *= decay_rate   # decay the learning rate

    running_loss = 0.0
    for i in range(num_train):
        train_input_tensor = ..........                    
        train_label_tensor = ..........
        optimizer.zero_grad()
        pred_label_tensor = model(train_input_tensor)
        loss = criterion(pred_label_tensor, train_label_tensor)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    loss_history[m_lr].append(running_loss/num_train)

Setup-4 Results:  

In this setup, I'm using Pytorch's learning-rate-decay scheduler (multiStepLR)
which decays the learning rate every 25 epochs by 0.25.
Here also, the loss jumps everytime the learning rate is decayed.

Como lo sugirió @Dennis en los comentarios a continuación, probé con ambas ReLUy sin 1e-02 leakyReLUlinealidades. Pero, los resultados parecen comportarse de manera similar y la pérdida primero disminuye, luego aumenta y luego se satura a un valor más alto de lo que alcanzaría sin disminuir la tasa de aprendizaje.

La gráfica 4 muestra los resultados.

Parcela-4:

scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer=optimizer, milestones=[25,50,75], gamma=0.25)

scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer=optimizer, gamma=0.95)

scheduler = ......... # defined above
optimizer = torch.optim.Adam(lr=m_lr,amsgrad=True, ...........)

for epoch in range(num_epochs):

    scheduler.step()

    running_loss = 0.0
    for i in range(num_train):
        train_input_tensor = ..........                    
        train_label_tensor = ..........
        optimizer.zero_grad()
        pred_label_tensor = model(train_input_tensor)
        loss = criterion(pred_label_tensor, train_label_tensor)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    loss_history[m_lr].append(running_loss/num_train)

EDICIONES:

• Como se sugiere en los comentarios y la respuesta a continuación, he realizado cambios en mi código y he entrenado el modelo. He agregado el código y las parcelas para lo mismo.
• Probé con varios lr_scheduleren PyTorch (multiStepLR, ExponentialLR)y parcelas de la misma se enumeran en Setup-4lo sugerido por @Dennis en los comentarios a continuación.
• Intentando con leakyReLU como lo sugiere @Dennis en los comentarios.

Alguna ayuda. Gracias

No veo ninguna razón por la cual las tasas de aprendizaje en decadencia deberían crear el tipo de saltos en las pérdidas que estás observando. Debería "ralentizar" la rapidez con la que "se mueve", lo que en el caso de una pérdida que de otra manera se reduce de manera consistente, en el peor de los casos, solo debería conducir a una meseta en sus pérdidas (en lugar de esos saltos).

Lo primero que observo en su código es que vuelve a crear el optimizador desde cero en cada época. Todavía no he trabajado lo suficiente con PyTorch para asegurarlo, pero ¿esto no solo destruye el estado interno / memoria del optimizador cada vez? Creo que debería crear el optimizador una vez, antes de recorrer las épocas. Si esto es realmente un error en su código, también debería ser un error en el caso de que no use la disminución de la tasa de aprendizaje ... pero tal vez simplemente tenga suerte allí y no experimente los mismos efectos negativos del insecto.

Para la disminución de la tasa de aprendizaje, recomendaría usar la API oficial para eso , en lugar de una solución manual. En su caso particular, querrá crear una instancia de un planificador StepLR , con:

• optimizer = el optimizador ADAM, que probablemente solo debería crear una instancia una vez.
• step_size = 25
• gamma = 0.25

Luego, puede simplemente llamar scheduler.step()al comienzo de cada época (o tal vez al final; el ejemplo en el enlace de la API lo llama al comienzo de cada época).

Si, después de los cambios anteriores, aún experimenta el problema, también sería útil ejecutar cada uno de sus experimentos varias veces y trazar resultados promedio (o líneas de trama para todos los experimentos). Teóricamente, tus experimentos deberían ser idénticos durante las primeras 25 épocas, pero aún vemos grandes diferencias entre las dos figuras, incluso durante las primeras 25 épocas en las que no se produce una disminución de la tasa de aprendizaje (por ejemplo, una cifra comienza con una pérdida de ~ 28K, la otra comienza con una pérdida de ~ 40K). Esto puede deberse simplemente a diferentes inicializaciones aleatorias, por lo que sería bueno promediar esa falta de determinación de sus gráficos.