brudfors / UniRes

Model-based super-resolution of medical images in PyTorch.
MIT License
75 stars 12 forks source link

Pre-conditioning CG #12

Open balbasty opened 3 years ago

balbasty commented 3 years ago

Hi Mikael,

I think you can get CG to converge dramatically faster by using a relatively simple preconditioner. From what I understand, CG is really just gradient descent, except that you don't descend along directions that you have already visited (you project the current gradient on the subspace orthogonal to all previous gradients). This means that without preconditioning, we can take very tiny steps, and preconditioning can be seen as using an approximation of the Hessian to take larger (better) steps. I found that using a diagonal preconditioner that is more positive definite than the true Hessian makes a very good conditioner. In your case, your Hessian (the system that you are solving) is t*A'A + lam*D'D, so your preconditioner can be inv(t*diag(A'A*1) + lam*diag(D'D)). Since D'D is just a convolution, its diagonal is constant and equal to 2/sum(vx ** 2) (see the Dartel paper: that is the central weight in the membrane energy).

So at each step, preconditioning just means dividing voxel-wise by t*A'A*1 + 2/sum(vx ** 2)

Happy Christmas!