A Naive Trick to Accelerate Training of LNCC-Based Deep Image Registration Models

1. Introduction

Local Normalized Cross-Correlation Loss (LNCC)[1], measuring the negative correlation of two input images in a sub-window manner, has been widely adopted for deformable image registration tasks [2,3]. Specifically, for a pair of source image I and reference image J, it is defined as

$$L_{\mathrm{lncc}}(I,J)=1-{\displaystyle \frac{1}{\left|\mathcal{N}\right|}}\sum _{k\in \mathcal{N}}{\displaystyle \frac{{\sum }_{i\in {\mathcal{W}}_k}(I_i-{\overline{I}}_k)(J_i-{\overline{J}}_k)}{\sqrt{{\sum }_{i\in {\mathcal{W}}_k}{(I_i-{\overline{I}}_k)}^2}\sqrt{{\sum }_{i\in {\mathcal{W}}_k}{(J_i-{\overline{J}}_k)}^2}}},$$

where:

$$\begin{array}{cc}\hfill \mathcal{N}& =\mathrm{set}\mathrm{of}\mathrm{all}\mathrm{voxels},\hfill \\ \hfill {\mathcal{W}}_k& =\mathrm{local}\mathrm{window}\mathrm{centered}\mathrm{at}\mathrm{voxel}k,\hfill \\ \hfill {\overline{I}}_k& ={\displaystyle \frac{1}{|{\mathcal{W}}_k|}}\sum _{i\in {\mathcal{W}}_k}{I}_i(\mathrm{mean}\mathrm{of}I\mathrm{in}\mathrm{window}{\mathcal{W}}_k),\hfill \\ \hfill {\overline{J}}_k& ={\displaystyle \frac{1}{|{\mathcal{W}}_k|}}\sum _{i\in {\mathcal{W}}_k}{J}_i(\mathrm{mean}\mathrm{of}J\mathrm{in}\mathrm{window}{\mathcal{W}}_k),\hfill \\ \hfill I_i,J_i& =\mathrm{intensity}\mathrm{values}\mathrm{of}\mathrm{images}I\mathrm{and}J\mathrm{at}\mathrm{voxel}i.\hfill \end{array}$$

Following the traditional convolution-based implementation [1], many PyTorch-based registration models, such as VoxelMorph [4,5]1, DeepReg [6]2, SYMNet[7] 3, MultiPropReg [8]4, TransMorph [9]5, LKU-Net[10]6, MONAI [11]7, RAN [12]8, WiNet [13]9 and Rethink-Reg [14]10, just to name a few, implement LNCC using five sequential convolution operations, as outlined in Algorithm 1. We, however, found that the computations may be accelerated.

Algorithm 1: LNCC loss calculation with five sequential convolutions

Require:I (input image), J (reference image), w (window size) 1: Define a convolutional filter W of size $w\times w\times w$ with all elements equal to ${\displaystyle \frac{1}{w^3}}$ 2: Compute ${\mu }_I=W*I$▹ Local mean of I 3: Compute ${\mu }_J=W*J$▹ Local mean of J 4: Compute ${\sigma }_I^2=W*\left(I^2\right)-{\mu }_I^2$▹ Local variance of I 5: Compute ${\sigma }_J^2=W*\left(J^2\right)-{\mu }_J^2$▹ Local variance of J 6: Compute ${\sigma }_{IJ}=W*(I·J)-{\mu }_I·{\mu }_J$▹ Local covariance of I and J 7: Compute $\mathrm{LNCC}={\displaystyle \frac{{\sigma }_{IJ}}{\sqrt{{\sigma }_I^2·{\sigma }_J^2+ϵ}}}$▹ Add small $ϵ$ to avoid division by zero 8: Compute loss: $\mathrm{Loss}=1-\mathrm{mean}\left(\mathrm{LNCC}\right)$ 9: return Loss

2. Implementation

2.1. Pseudo Code

The approach in Algorithm 1 uses five consecutive convolutions to compute the LNCC loss, which fails to fully leverage the parallel computing capabilities of modern hardware. We propose to replace the five consecutive convolutions with one single group convolution, as illustrated in Algorithm 2. This simple change can rapidly reduce the computational time on modern hardware (as illustrated in Figure 1), while the numerical results are not affected. (The computed loss values and gradients will be identical.)

Algorithm 2: LNCC loss calculation with one group convolution

Require:I (input image), J (reference image), w (window size) 1: Define a convolutional filter W of size $w\times w\times w$ with all elements equal to ${\displaystyle \frac{1}{w^3}}$ 2:  Compute $Set=W*(I,J,I^2,J^2,I·J)$▹ Group convolution 3: Get ${\mu }_I=Set\left[0\right]$▹ Local mean of I 4: Get ${\mu }_J=Set\left[1\right]$▹ Local mean of J 5: Compute ${\sigma }_I^2=Set\left[2\right]-{\mu }_I^2$▹ Local variance of I 6: Compute ${\sigma }_J^2=Set\left[3\right]-{\mu }_J^2$▹ Local variance of J 7: Compute ${\sigma }_{IJ}=Set\left[4\right]-{\mu }_I·{\mu }_J$▹ Local covariance of I and J 8: Compute $\mathrm{LNCC}={\displaystyle \frac{{\sigma }_{IJ}}{\sqrt{{\sigma }_I^2·{\sigma }_J^2+ϵ}}}$▹ Add small $ϵ$ to avoid division by zero 9: Compute loss: $\mathrm{Loss}=1-\mathrm{mean}\left(\mathrm{LNCC}\right)$ 10: return Loss

Additionally, this approach may be applicable to other similarity measures that require repeated sliding window-based computations, such as the Structural Similarity Index Measure (SSIM).

2.2. Runtime Analysis

In Figure 1, we compare the training time of a deep registration model taken using the previous Algorithm 1 and the new Algorithm 2. Our experimental setup is as follows:

• Data: 403 $160\times 192\times 224$ training pairs from the IXI dataset (https://brain-development.org/ixi-dataset/) pre-processed by TransMorph [9] are selected as training data.
• Architectures: 5 benchmarking deep registration models, i.e., VoxelMorph-1, VoxelMorph-2 [4,5], ${\mathrm{LKU-Net}}_{8,5}$, ${\mathrm{LKU-Net}}_{16,5}$ [10], and TransMorph [9] are selected as comparing methods.
• Training: Each one of the five models is trained in five epochs using the Algorithm 1 and the new Algorithm 2, respectively.
• Results: The training time required for each epoch is reported to compare the difference between the two algorithms.
• Others: The batch size of all networks is set to 1. The sub-window of LNCC is fixated at $9\times 9\times 9$. All models are trained on the same Nvidia A100-40G GPU with Pytorch. (More experiments are in the Appendix A.)

We can observe that the runtime of the new LNCC implementation (Algorithm 2) is consistently and significantly less than the Algorithm 1 on all compared architectures. Specifically, the average training time on 5 epochs is reduced 52.49%, 51.72%, 47.54%, 43.69%, and 37.61% for VoxelMorph-1, VoxelMorph-2, ${\mathrm{LKU-Net}}_{8,5}$, ${\mathrm{LKU-Net}}_{16,5}$, and TransMorph. (Detailed numbers are included in Table 1.) The trend is easy to interpret as the feed-forward and back-propagation will account for more computations in larger models, making the computations of LNCC less critical.

3. Conclusion

We are gaining access to millions of scans from various imaging modalities to develop AI-driven healthcare applications, of which, image registration is often a fundamental step [15]11. Researchers in the registration community contrive to accelerate the registration process for such large-scale cohorts.

Many deep image registration models have proven that their inference time can be orders of magnitude faster than conventional models, while the lengthy training time for deep models is the elephant in the room. Using a single group convolution, the training time for LNCC-based deep registration models can be nearly halved on modern hardware such as Nvidia A100 GPU, significantly accelerating the development of more advanced registration and downstream medical analysis models. We hope this approach will demonstrate practical value to the community.