Improving Instance Optimization in Deformable Image Registration with Gradient Projection

Yi Zhang,Yidong Zhao,Qian Tao

2024-10-23

Abstract:Deformable image registration is inherently a multi-objective optimization (MOO) problem, requiring a delicate balance between image similarity and deformation regularity. These conflicting objectives often lead to poor optimization outcomes, such as being trapped in unsatisfactory local minima or experiencing slow convergence. Deep learning methods have recently gained popularity in this domain due to their efficiency in processing large datasets and achieving high accuracy. However, they often underperform during test time compared to traditional optimization techniques, which further explore iterative, instance-specific gradient-based optimization. This performance gap is more pronounced when a distribution shift between training and test data exists. To address this issue, we focus on the instance optimization (IO) paradigm, which involves additional optimization for test-time instances based on a pre-trained model. IO effectively combines the generalization capabilities of deep learning with the fine-tuning advantages of instance-specific optimization. Within this framework, we emphasize the use of gradient projection to mitigate conflicting updates in MOO. This technique projects conflicting gradients into a common space, better aligning the dual objectives and enhancing optimization stability. We validate our method using a state-of-the-art foundation model on the 3D Brain inter-subject registration task (LUMIR) from the Learn2Reg 2024 Challenge. Our results show significant improvements over standard gradient descent, leading to more accurate and reliable registration results.

Computer Vision and Pattern Recognition,Image and Video Processing

What problem does this paper attempt to address?

This paper attempts to solve the multi - objective optimization (MOO) problem in deformable image registration (DIR). Specifically, the paper focuses on how to improve the accuracy of registration through instance optimization (IO) at test time.

### Main problems:
1. **Multi - objective optimization conflict**: The two objectives of image similarity and deformation regularity often conflict with each other, resulting in optimization results falling into unsatisfactory local minima or slow convergence speed.
2. **Poor performance of deep learning methods at test time**: Although deep learning methods perform well in handling large - scale data sets and achieving high precision, their performance at test time is usually not as good as that of traditional optimization techniques, especially when there are distribution differences between training data and test data.
3. **Limitations of instance optimization**: Although traditional IO methods can combine the generalization ability of deep learning and the advantages of instance - specific optimization, the optimal balance point for each test pair may be different, resulting in limited actual performance improvement.

### Solutions:
To solve the above problems, the author introduced a new instance - optimization algorithm based on gradient projection. This method projects conflicting gradients into a common space, thereby better aligning dual objectives and enhancing the stability of optimization. The specific steps are as follows:

1. **Define the optimization problem**: Given a pair of 3D images \(I_A\in\mathbb{R}^{D\times H\times W}\) and \(I_B\in\mathbb{R}^{D\times H\times W}\), find a dense transformation \(\phi\in\mathbb{R}^{3\times D\times H\times W}\) such that the deformed moving image \(I_A\circ\phi\) is anatomically similar to the fixed image \(I_B\). The optimization problem can be expressed as:
\[
\hat{\phi}=\arg\min_{\phi}L_{\text{sim}}(I_A\circ\phi, I_B)+\lambda L_{\text{reg}}(\phi)
\]
where \(L_{\text{sim}}\) is the similarity term, \(L_{\text{reg}}\) is the regularization term, and \(\lambda\) is a trade - off parameter.

2. **Gradient projection technique**: To alleviate the conflicting updates in MOO, the author used the gradient projection technique. When the gradients \(\nabla L_{\text{sim}}\) and \(\nabla L_{\text{reg}}\) of the two loss functions conflict, one of the gradients will be randomly projected onto the normal space of the other gradient. For example, if \(\nabla L_{\text{sim}}\) is selected for correction, the formula for calculating the corrected gradient is:
\[
\nabla L_{\text{sim}}=\nabla L_{\text{sim}}-\frac{\langle\nabla L_{\text{sim}},\nabla L_{\text{reg}}\rangle}{\|\nabla L_{\text{sim}}\|^2}\nabla L_{\text{reg}}
\]
where \(\langle\cdot,\cdot\rangle\) represents the inner product of two vectors.

### Experimental verification:
The author verified the method using the LUMIR task in the Learn2Reg 2024 challenge and showed results significantly superior to standard gradient descent, especially with obvious improvements in accuracy and reliability.

In summary, this paper aims to improve instance optimization through gradient projection techniques.