Training instability with Dice Loss/Tversky Loss

[martaranzini]

I am training a 2D UNet to segment fetal MR images using MONAI and I have been observing some instability in the training when using MONAI Dice loss formulation. After some iteration, the loss jumps up and the network stops learning, as the gradients drop to zero. Here is an example (orange is loss on training set computed over 2D slices, blue is loss on validation computed over 3D volume).

After investigating several aspects (using the same deterministic seed), I've narrowed down the issue to the presence of the smooth term in both the numerator and denominator of the Dice Loss:
f = 1.0 - (2.0 * intersection + smooth) / (denominator + smooth)

When using the formulation:
f = 1.0 - (2.0 * intersection) / (denominator + smooth)
without the smooth term in the numerator, the training was stable and no longer showed unexpected behaviour:

[Note: this experiment was trained for much longer to make sure the jump would not appear later in the training]

The same pattern was observed also for the Tversky Loss, so it could be worth investigating the stability of the losses to identify the best default option.

Software version
MONAI version: 0.1.0+84.ga683c4e.dirty
Python version: 3.7.4 (default, Jul 9 2019, 03:52:42) [GCC 5.4.0 20160609]
Numpy version: 1.18.2
Pytorch version: 1.4.0
Ignite version: 0.3.0

Training information
Using MONAI PersistentCache
2D UNet (as default in MONAI)
Adam optimiser, LR = 1e-3, no LR decay
Batch size: 10

Other tests
The following aspects were investigated but did not solve the instability issue:

• Gradient clipping
• Different optimisers (SGD, SGD + Momentum)
• Transforming the binary segmentations to a two-channel approach ([background segmentation, foreground segmentation])
• Choosing smooth = 1.0 as default here (https://github.com/MIC-DKFZ/nnUNet/blob/master/nnunet/training/loss_functions/dice_loss.py). However, this made the behaviour even more severe and the jump would happen sooner in the training.

The following losses were also investigated

• Binary Cross Entropy --> stable
• Dice Loss + Binary Cross Entropy --> unstable
• Dice Loss (no smooth at numerator) + Binary Cross Entropy --> stable
• Tversky Loss --> Unstable
• Tversky Loss (no smooth at numerator) --> stable

[wyli]

thanks for this nice report, will look into this (cc @FabianIsensee @ericspod @holgerroth)

[FabianIsensee]

Interesting observation. I must admit that I have never investigated this in depth in nnU-Net. I should also note that while the loss function defaults to smooth=1, nnU-Net actually uses smooth=1e-5 (see here https://github.com/MIC-DKFZ/nnUNet/blob/9524dc33425627d1e4b21336a5202b1bcc8157e5/nnunet/training/network_training/nnUNetTrainer.py#L112).

Have you tried training the 2D configuration of nnU-Net with this dataset? It would be quite interesting in knowing whether the instability happens there as well

[martaranzini]

No, I actually haven't tried and trained nnUnet with my data, but I can look into it in the next days and feedback whether I observe the same behaviour :) (@LucasFidon may find this conversation useful as well)

[FabianIsensee]

When you do, please beware that there is a memory leak on Ubuntu 18 and above with Turing cards. Unfortunately there is nothing I can do about it because this is related to mixed precision training with apex/amp (and also occurs with pytorch autocast). I already opened an issue about this and I hope someone will look into this. If you observe this on your system as well, please let me know. You can get rid of it by training in fp32 with the -fp32 option (you will need more GPU memory).

Let me know if you need any assistance getting 2D data to run. This is not straightforward right now. You basically need to create dummy 3D niftis that have shape 1 in the first axis.

Best,
Fabian

[wyli]

I'm trying to replicate this issue with some synthetic data, looks like using sigmoid=True instead of softmax=True is much more stable. @martaranzini are you using the softmax=True option, and could you please try sigmoid=True if possible?

i guess softmax is causing some over/underflow I'll look further into this

[martaranzini]

Hi @wyli
I was actually using sigmoid=True already in the experiments I reported in the issue. When using a single-channel output, I put sigmoid=True and softmax=False. Conversely, in the two-channel test I did, I used the softmax instead of the sigmoid. In both cases I would observe the reported behaviour.

We are also working on training the nnU-net on the same data to see if it shows the same behaviour. We are going to report back about it as soon as possible.

[wyli]

ok, thanks for the information, I haven't tested a single-channel case yet, I'll do that as well. for now two-channel + sigmoid seems to be more stable compared with two-channel + softmax. looks like it's independent of the choice of network afaik

[martaranzini]

Ok, I see, thanks! I can try and test two-channel + sigmoid instead of softmax on my data as well and see if it gets more stable. I will let you know.

[martaranzini]

Hi all,

First, apologies for the delay in getting back to you about this issue – we have been running a few experiments to put together a MONAI and nnU-Net comparison and this required quite some time. We hope you will find our results interesting and informative.

@LucasFidon kindly ran a few experiments with nnU-Net and this allowed us to identify some implementation differences. Both MONAI and nnU-Net use the Dice formulation:

However, as a default MONAI computes the Dice per element in the batch and then the loss is averaged across the batch (we will refer to this simply as “Dice”). We noticed that in nnU-Net the Dice is instead computed directly as a single value across the whole batch (i.e. not per image). The average is computed only across the channels, not across the batch elements. We refer to this approach as “Batch Dice”.
We experimented with both formulations in both frameworks, and also tested for Dice loss and Dice + Cross entropy loss at training, as well as for the use of a single- or 2-channel approach.

The performance of the different trained models on the validation set is reported below (note: they are separated in groups A-E for our own clinical interpretation, but the group separation is not particularly relevant for this issue). The sampling strategy at training is also reported.

A few specifications:

• MONAI – Dice no smooth at numerator used the formulation:
  
• nnU-Net – Batch Dice + Xent, 2-channel, ensemble indicates ensemble performance from 5-fold cross validation at training
• NeuroImage indicates a published two-step approach on our dataset, and it is reported just for reference.

We gather two main observations from these experiments:

1. Training instability: On our dataset, nnU-Net did not present the training instability observed with the default MONAI implementation of Dice. This is also confirmed when known implementation differences were ruled out (nnU-Net – Dice, 2-channel, uniform sampling). Also, nnU-Net generally provides better performance. @FabianIsensee, are there any other implementation differences that could justify our results?
2. Dice vs Batch Dice: In both frameworks, the Batch Dice implementation clearly outperforms the “normal” Dice computation. This could be an interesting feature to be added in MONAI – happy to open another issue/PR about this.

@wyli: I did retrain also the two-channel with sigmoid instead of softmax. With respect to the single-channel, the gradients do not drop to zero (with either sigmoid and softmax), but I still observe some instability in the loss which I cannot fully explain:

Looking forward to hearing your comments, and happy to run more experiments to investigate this further!

[FabianIsensee]

Hi @martaranzini , thank you so much for the detailed report! It's a very interesting read and I am happy to see that nnU-Net did not disappoint.

Just so that I think about the results in the correct context (@LucasFidon ) : Did you provide separate 2D slices as training examples? Or did you just run the 2D nnU-net configuration on the 3D images? My guess is the latter, but it would be important to be sure.

The batch dice (this is what I call it as well) is implemented on purpose in nnU-Net, but it is not always active. nnU-Net sets that automatically:

• if the patch size does not cover the entire image nnU-Net uses batch dice. Since small patches do not represent the true class distribution anyways we might as well benefit from the increased stability of batch dice. Batch dice is used for all 2D U-Nets or 3d_fullres U-Nets if no 3d_lowres is present
• nnU-Net uses 'sample dice' (that is what I call the regular Dice) for all cases where the patch size covers almost entire training cases, so that would be the 3d_lowres U-Net and the 3d_fullres U-Net if no lowres is present.

Note that there are cases where sample Dice is better than batch dice, so replacing it categorically is not a good idea.

I don't know about the MONAI implementation, but using the sample Dice in 2D segmentation tasks in nnU-Net is a bad idea. That is because nnU-Net never optimizes the background class with the Dice (just the CE). What this translates to is a loss function that will not yield any useful gradients on slices where no foreground voxels are present - essentially ignoring these slices in optimization. If you want to run sample Dice with nnU-Net, make sure to set do_bg=True in the SoftDiceLoss :-) (this problem does not exist when pairing the Dice with CE because the CE term saves it ;-) )

In my experiments, using uniform sampling gives pretty much the same results as oversampling foreground. Right now we see uniform sampling in nnU-Net only for a scenario which is also unstable in non-uniform sampling. It would be very interesting to see the nnU-Net performance for uniform sampling with the default loss function as well to ensure that oversampling it not skewing the results.

I am not quite sure I understand what you mean by 'ensemble indicates ensemble performance from 5-fold cross validation at training'. Is that simply the Dice score averaged over the five folds? Or do you have a heldout validation set where you applied the ensemble to?

Since your segmentation problem seems to be 3D - what was your motivation to go with a 2D U-Net? I believe you could get better results by running the 3D U-Net. If you have compute resources to spare you could just run the 3d_fullres config. If your dataset is public I can also do that for you :-)

Let me know what you think!

Best,

Fabian

[LucasFidon]

Hi @FabianIsensee, thank you for the quick reply!

Regarding 3D vs 2D, in fetal MRI the 3D MRI are motion-corrupted stacks of 2D slices. In addition the in-plane resolution is typically several times higher than the resolution across planes (please see [1] for more details).
In [1] it has already been shown empirically that 2D CNNs were superior to 3D CNNs for the task of fetal brain extraction. As a result, I have trained nnU-Net only in '2d' mode. For training, the 2D slices were used as input (with an extra first dimension of size 1 as you suggested above).

Regarding the validation dataset used in the figure of @martaranzini, it is a heldout validation set.
The nnU-Net ensemble is the ensemble of the five 2D U-Nets trained for the five folds created by nnUNet.

The dataset is not publicly available... But I will train nnU-Net for the default configuration (batch Dice + CE) except I will deactivate the oversampling :)

Do you think that it is worth trying do_bg=True also with batch Dice + CE or does do_bg=False always work better in your experience?

Thanks,
Lucas

[1] Ebner, Michael, et al. "An automated framework for localization, segmentation and super-resolution reconstruction of fetal brain MRI." NeuroImage 206 (2020): 116324.

[martaranzini]

Hi @FabianIsensee ,

Thanks for your comments, very interesting, especially in the analysis of when batch Dice is more suitable than sample Dice. I totally agree that it should not replace it categorically, but it would be useful to add an option in MONAI for the user to choose one over the other.

I will leave to @LucasFidon the answers to the nnU-Net experiments as he actually did all the hard part on that :)

I can however reply to the details about our segmentation problem. Our task is brain segmentation of fetal MR images for subsequent super-resolution-reconstruction. Our data is 3D but because of heavy (inter-slice) motion artifacts (and heavy anisotropic resolution), we need to approach it as a 2D problem, where we segment each slice independently. Only at inference we perform a 3D assessment, by computing the Dice for the whole 3D image - and this 3D assessment is reported for the validation set in the figure above. We did run some tests with 3D training but it was clearly underperforming the 2D approach.

In addition, the FOV of our images are quite large and we do have a lot of background-only slices - and we want to be able to correctly predict them as empty. So the comment you raised about the background is spot-on and we will look into greater detail about it. I know for sure we did include the background in MONAI (in the two-channel approaches), but we will check the single-channel and nnU-Net better for this.
Re: the single-channel, I think that removing the smooth term at the numerator in the Dice will have the same effect of ignoring the background slices. On the other hand, we have this instability in the Dice loss and in the Tversky loss in the MONAI framework when keeping the smooth term at the numerator that we cannot explain.

Many thanks!
Marta

[martaranzini]

Oups, I think we posted about at the same time, thanks for answering the nnU-Net questions, @LucasFidon!