Loss suddenly increases

[YukinoshitaLove]

I came across a problem when training resfusion_restore on the LOLDataset. At the start of the training, the loss decreases as expected. But after 17.6K~18.1K iterations, the loss suddenly explodes to a very large number(~1e9). I tried several times but the same problem occurs, using command: python train_resfusion_restore.py.

[YukinoshitaLove]

I printed all the inputs of the denoising module and the predict object of it. They were all fine, only the pred_resnoise went wrong. This indicates that the problem lies with the neural network or optimizer.

pred_resnoise = self.denoising_module(x=x_t, time=t, input_cond=X_0_hat)
loss = F.mse_loss(input=pred_resnoise, target=resnoise)
self.log('X_0_hat_max', X_0_hat.max(), on_epoch=True, prog_bar=False, logger=True, sync_dist=True)
self.log('X_0_hat_min', X_0_hat.min(), on_epoch=True, prog_bar=False, logger=True, sync_dist=True)
self.log('resnoise_max', resnoise.max(), on_epoch=True, prog_bar=False, logger=True, sync_dist=True)
self.log('resnoise_min', resnoise.min(), on_epoch=True, prog_bar=False, logger=True, sync_dist=True)
self.log('x_t_max', x_t.max(), on_epoch=True, prog_bar=False, logger=True, sync_dist=True)
self.log('x_t_min', x_t.min(), on_epoch=True, prog_bar=False, logger=True, sync_dist=True)
self.log('t_max', t.float().max(), on_epoch=True, prog_bar=False, logger=True, sync_dist=True)
self.log('t_min', t.float().min(), on_epoch=True, prog_bar=False, logger=True, sync_dist=True)
self.log('pred_resnoise_max', pred_resnoise.max(), on_epoch=True, prog_bar=False, logger=True, sync_dist=True)
self.log('pred_resnoise_min', pred_resnoise.min(), on_epoch=True, prog_bar=False, logger=True, sync_dist=True)

[YukinoshitaLove]

I tried gradient clip. There was still a loss explosion, but the training process didn't fail.

parser.add_argument('--gradient_clip', default=1, type=float)
trainer = Trainer(
  ...,
  gradient_clip_val=args.gradient_clip,
)

[YukinoshitaLove]

It seems that the problem lies in the Adam optimizer. ( ᐛ )
Accordding to:

https://discuss.pytorch.org/t/loss-suddenly-increases-using-adam-optimizer/11338

In anycase, the problem with Adam is that it uses moving average in the denominator term. So if the gradients get really small and the whole of denominator will be small. Since the gradients are already small, the denominator results in blowup thus pushing you very far away hence huge loss. You may have a look at https://openreview.net/forum?id=ryQu7f-RZ 560 . I think there are many recent methods which avert this problem including AMSGrad (in the earlier mentioned paper), Hyper-gradient descent (for Adam) etc. Also look for comments in the openreview forum, there seems to be further discussion on this issue. Unfortunately i do not know of any pytorch implementation of these algorithms.

I replaced the Adam optimizer with AMSGrad (and still used gradient_clip_val=1). The problem is solved ヽ( ຶ▮ ຶ)ﾉ! It indeed comes from the stabilization issue of the Adam itself. Now the train_loss is normal and we can happily train the resfusion_restore. (◍•ᴗ•◍)

optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.parameters()), lr=lr,
                               weight_decay=self.hparams.weight_decay, amsgrad=True)