LNCC loss instability for large kernel size

[ebrahimebrahim]

I experienced some numerically unstable training while using LNCC loss to train a registration network. Exploding gradients, clipping them not helping, training loss increasing even with a low learning rate, etc.

This is in the context of synthetic data, where the images contained large homogeneous blobs, and so the variances used to compute LNCC were probably quite low.

I see that small variance has been a problem in the past and that it has been addressed. However, I seem to still have the issue when I set the kernel_size argument for LNCC loss to be large. Setting it back to a small value like 3 fixes the problem completely. Setting it to a large value (say, half the image size) leads to problems.

I think that giving a large kernel_size to LNCC is a natural thing to do, because NCC (normalized cross correlation, without the "local") is the case of having kernel_size equal image size, if I'm understanding correctly.

Is there any reason to expect that large kernel size would worsen the instability problem? From the code, it's not clear to me why the LNCC kernel size would matter. Thoughts, anyone?

[wyli]

Hi @kate-sann5100 would you be able to look into this one?
cc @YipengHu @mathpluscode

[YipengHu]

@ebrahimebrahim Can you elabrate a bit on the exploded gradeitns? i gathered it only happens when kernel_size is large, correct? how large you think we can re-produce? did you clip the output or the input of the LNCC? where did you observe the high gradient?

[ebrahimebrahim]

Here is a demo of the issue: https://github.com/ebrahimebrahim/deep-atlas/blob/main/lncc_with_large_kernel.ipynb
This short notebook creates synthetic data consisting of solid balls that have varying size and position,
then tries to train a registration network using lncc loss alone.

In the two training runs, you can look at the gradient size that is printed out. Also the training loss plot. The only difference in the two runs is the lncc loss kernel_size.

Lowering the learning rate in the large kernel_size run does not seem to help much.

it only happens when kernel_size is large, correct?

correct-- as far as I have seen

how large you think we can re-produce?

In one trial that gave me issues, the kernel size was 17, while the image size was 64x64x64
In the example I just linked, kernel_size=25 is shown

did you clip the output or the input of the LNCC?

no I did not try that. I just tried using torch.nn.utils.clip_grad_norm_ on the model itself, but the problem was worse than just having spikes in the gradient every now and then. It was consistently unstable. So this approach did not help.

I also did not try messing with smooth_nr and smooth_dr parameters of the lncc.
These might help-- idk-- but when I look at the code I am puzzled as to why there's a problem with kernel_size in the first place.

[kate-sann5100]

Hi, I am so happy to see the use of registration components in the DeepAtlas tutorial.
For the LNCC loss instability issue, the demo is really helpful for me to understand the problem. I will look into that and hopefully get back to you soon.

[ebrahimebrahim]

The registration components are very well put together 😀

Thanks for looking into this.

[ebrahimebrahim]

Should this be converted to an issue?

[mikami520]

@ebrahimebrahim I am working on Deep Atlas and using the local normalized cross correlation loss function (return -LNCC instead of 1-LNCC) to compute the similarity loss between warped moving image and fixed image. However, the loss is too big (too negative shown in the screenshot) and makes me confused.

similarity_loss_reg.txt

Did you face same problem when you implemented?

[ebrahimebrahim]

@mikami520 The problem that I faced was large gradients resulting in training instability.

I think the problem is worse when the images have more homogeneous regions, because then the variances in the denominator of LNCC become zero.

See here that "smoothing" is applied to prevent division by zero. Depending on the choices of smooth_dr and smooth_nr, it can be possible for this -LNCC function to return values below -1.

If you are having training stability issues I suggest playing with the values of smooth_dr and smooth_nr in the loss constructor, or consider using a different image similarity measure. I found that it's good to choose them such that smooth_nr << smooth_dr. Having smooth_nr $\approx$ smooth_dr can be bad because it causes the loss function in the zero-variance case to be "good" and therefore encourages the network to discover warps that decrease pixel variation as much as possible. I've gotten some pretty funny behavior out of this, like learning to shrink an image of existence, or learning to translate it out of frame.

[mikami520]

@mikami520 The problem that I faced was large gradients resulting in training instability.

I think the problem is worse when the images have more homogeneous regions, because then the variances in the denominator of LNCC become zero.

See here that "smoothing" is applied to prevent division by zero. Depending on the choices of smooth_dr and smooth_nr, it can be possible for this -LNCC function to return values below -1.

If you are having training stability issues I suggest playing with the values of smooth_dr and smooth_nr in the loss constructor, or consider using a different image similarity measure. I found that it's good to choose them such that smooth_nr << smooth_dr. Having smooth_nr ≈ smooth_dr can be bad because it causes the loss function in the zero-variance case to be "good" and therefore encourages the network to discover warps that decrease pixel variation as much as possible. I've gotten some pretty funny behavior out of this, like learning to shrink an image of existence, or learning to translate it out of frame.

@ebrahimebrahim Sounds good. I am appreciated your help. I'll have a try and come back to share my results with you. Thank you so much !!!