LNCC Training: Unable to learn (very unstable)

[JannikLa]

Task (what are you trying to do/register?)

I am registering two T1 3D images from the OASIS data set (L2Reg challenge). The scans are preprocessed (skull-stripped, aligned, normalized). I am taking a fixed image from one subject and randomly select another subject as moving image (examples below).

I invested a lot if time in figuring out what is wrong. Maybe someone can give hints on what I could try or what might be wrong in my setup. Any help is appreciated!

What have you tried

I am using the pytroch implemention of voxelmorph. A dataloader loads the inputs to the network.

• The network works as expected with MSE
• With LNCC the loss is very unstable an does not decrease

Details of experiments

• My "basline" setup is: lr: 1e-4, bs: 4, Reg weight: 0.1, loss: LNCC (win size 9)
  • I changed hyper-parameters separately in both directions (e.g. increasing and decreasing the learning rate)
  • I also tried setting the Regularization weight to zero
  • If I set the initialization of the flow layer a lot higher (to 1 or 10) the network is able to overfit
  • self.flow.weight = nn.Parameter(Normal(0, 1e-5).sample(self.flow.weight.shape)) (networks.py, l 214)
  • I also tried debugging the LNCC loss. It works as expected (e.g. tried computing the loss between identical images)

Image of training loss with maxed out smoothing in visualization ("baseline" setup from above)

Image of training loss (LNCC) when overfitting on a pair of images with large initialization weights (Normal(0, 10))

Image of input data

[adalca]

I haven't worked with the pytorch version that much especially recently, but just so I know -- when you say the MSE version worked, what were the hyperparameters that you used?

What do you mean by this line -- what does 'overfit' mean here?

If I set the initialization of the flow layer a lot higher (to 1 or 10) the network is able to overfit