Paper: A Plug-and-Play Image Registration Network (ICLR-2024)

Reviewer comment (Evaluation):
The incremental improvements in results, especially in the second decimal place, raise concerns about the practical implications of such minor improvements, particularly for downstream tasks, especially that the proposed method is an iterative optimization-based approach that significantly increases inference time compared to deep learning based methods.

Author response:
We would like to highlight that PIRATE+ can gain noticeable improvement (more than 3% improvement in Dice score) than DL base methods (VoxelMorph, SYMNet, and GraDIRN). Compared to another iterative method, NODEO, PIRATE+ achieved slightly better results with  50% less inference time. Prompted by your comments, we also presented the best and the worst cases of registration results in Figure 2 and Figure 12, respectively. Note how PIRATE+ can gain visual improvement than the baseline, especially for small anatomical structures in the segmentation mask, in both the cases.

Reviewer comment (Writing):
The paper makes heavy use of acronyms, which affects readability. Notably, the acronym "DU" is mentioned without a clear definition.

Author response:
Prompted by your comment, the revised manuscript listed all acronyms and corresponding full names in Table 7.

Reviewer comment (Experiments):
While the paper's focus on brain MRI datasets is appreciated, it would have been beneficial to see the adaptability of PIRATE and PIRATE+ to other anatomies and imaging modalities.

Author response:
Prompted by your comment, we conducted additional experiments on the suggested EMPIRE10 lung CT dataset. The revised manuscript reports the numerical result in Table 4 and visual results in Figure 6 and 7. These results showed that PIRATE and PIRATE+ also had better performance than other baselines on the new lung CT dataset.

Reviewer comment (Evaluation):
How do the minor improvements in quantitative results translate to real-world applications, especially considering the potentially longer inference time of the iterative optimization-based approach?

Author response:
We would like to highlight that PIRATE+ can gain noticeable improvement (more than 3% improvement in Dice score) than DL base methods (VoxelMorph, SYMNet, and GraDIRN). Compared to another iterative method, NODEO, PIRATE+ achieved slightly better results with  50% less inference time. Prompted by your comments, we also presented the best and the worst cases of registration results in Figure 2 and Figure 12, respectively. Note how PIRATE+ can gain visual improvement than the baseline, especially for small anatomical structures in the segmentation mask, in both the cases.

Reviewer comment (Experiments):
Is registration performed on the full 3D scan, or are 2D slices of roughly pre-aligned images used?

Author response:
We performed 3D registration: Both PIRATE and PIRATE+ performed on full 3D scans. All figures in the manuscript show a 2D slice of the registered 3D volumes.

Reviewer comment (Evaluation):
How does PIRATE and PIRATE+ compare in terms of computational efficiency and scalability?

Author response:
Prompted by your comment, we tested the computational efficiency and scalability of PIRATE and PIRATE+ in Table 6 and Figure 11. Table 6 shows that PIRATE+ has higher computational and memory complexity than PIRATE. Figure 11 shows that PIRATE+ and PIRATE has similar scalability in terms of time and memory cost for inference.

Reviewer comment (Evaluation):
The methods compared in the paper use very different loss or cost functions as well as different models for the parameterization of the deformation field. This makes the papers' comparison of the registration accuracy in terms of voxels with negative Jacobian very difficult to come across methods. Registration accuracy measured in terms of Dice is more meaningful. At the same time, it would have been good if the authors had used some additional non-brain datasets which have landmarks and thus allow the calculation of quantities such as the target registration error. One such dataset is from the EMPIRE10 challenge...

Author response:
Prompted by your comment, we conducted additional experiments on the suggested EMPIRE10 lung CT dataset. The revised manuscript reports the numerical result in Table 4 and visual results in Figure 6 and 7. These results show that PIRATE and PIRATE+ can maintain better performance than other baselines on the new lung CT dataset.

Reviewer comment (Evaluation):
The run-time of the proposed framework is significantly higher than those of other DL methods. This is a significant disadvantage for clinical applications.

Author response:
The inference time of PIRATE and PIRATE+ are longer than other DL based methods (VoxelMorph, SYMNet, and GraDIRN) due to their iterative scheme. However, our results show that PIRATE+ can gain noticeable improvement (more than 3% improvement in Dice score) than those DL methods. Moreover, the testing time of PIRATE and PRIATE+ (around 40 seconds per 3D volume) is still negligible in clinical applications when real-time registration is not needed. One such application is to generate aligned image pairs required for training image translation models. Note also that our approach does not focus on real-time clinical application at the current stage.

Reviewer comment (Evaluation):
The paper proposes two methods, PIRATE and PIRATE+. I am a bit unclear on what are the conclusions: Is PIRATE+ better than PIRATE? When should PIRATE be used? When should PIRATE+ be used?

Author response:
Our results show that PIRATE+ outperforms PIRATE with a designated iterative algorithm when trained to optimize the task-specific algorithmic loss. The choice between PIRATE and PIRATE+ depends on the availability of training datasets and trade-off between performance and training complexity. In PIRATE+, one needs pairs of MRI data for training and to optimize a sophisticated DEQ loss related to a specific registration task. On the other hand, PIRATE only requires a set of registration fields, and the training loss corresponds to a simple regression task. The revised manuscript also reported the training complexity of PIRATE and PIRATE+ in Table 6.

Reviewer comment (Experiments):
How are the parameters in the other registration methods chosen, especially in the context of trading off registration accuracy and regularisation of the deformation field?

Author response:
We used the official codes of other registration baselines with the suggested parameters (parameters on brain datasets). We retrained all models on our datasets to make fair comparisons.

Reviewer comment (Evaluation):
The run-time of the proposed framework is significantly higher than those of other DL methods. What are the reasons for this?

Author response:
The inference time of PIRATE and PIRATE+ are longer than other DL based methods (VoxelMorph, SYMNet, and GraDIRN) due to their iterative scheme. However, our results show that PIRATE+ can gain noticeable improvement (more than 3% improvement in Dice score) than those DL methods. Moreover, the testing time of PIRATE and PRIATE+ (around 40 seconds per 3D volume) is still negligible in clinical applications when real-time registration is not needed. One such application is to generate aligned image pairs required for training image translation models. Note also that our approach does not focus on real-time clinical application at the current stage.

Reviewer comment (Experiments):
I didn't understand the motivation of using pre-trained denoisor for regularizations. Why isn't the trainable denoiser used for this task?

Author response:
Motivation of using pre-trained denoiser: The recent literature has shown the potential of training neural network denoisers to model the statistical distribution of high-dimensional data (e.g., denoising diffusion models and plug-and-play priors). Our work is the first to adopt this idea for registration by training a **denoiser prior** over the registration fields. The revised manuscript has clarified the choice of using denoiser prior for image registration.

Why not trainable denoiser: We would like to highlight that we did use trainable denoiser in PIRATE+. In PIRATE+, we train the denoiser on a dataset by using the loss specific to the registration task computed using the registration field, the warped image and the fixed image.

Reviewer comment (Experiments):
No ablation studies about the necessity of the usage of a pre-trained model.

Author response:
Motivation of using pre-trained denoiser: The recent literature has shown the potential of training neural network denoisers to model the statistical distribution of high-dimensional data (e.g., denoising diffusion models and plug-and-play priors). Our work is the first to adopt this idea for registration by training a **denoiser prior** over the registration fields. The revised manuscript has clarified the choice of using denoiser prior for image registration.

Why not trainable denoiser: We would like to highlight that we did use trainable denoiser in PIRATE+. In PIRATE+, we train the denoiser on a dataset by using the loss specific to the registration task computed using the registration field, the warped image and the fixed image.

Reviewer comment (Evaluation):
Generalization needs to be validated with group analysis.

Author response:
We would like to highlight that **we did conduct group analysis in Table 1, 2, and 3 in our original manuscript**. In those tables, we reported the mean and standard deviation values of the entire testing dataset, which consists of 100 pairs of 3D MRI volumes.

Reviewer comment (Evaluation):
What is the corner case for the proposed method?

Author response:
Prompted by your comment, the revised manuscript shows the best (with the most dice improvement, **not** absolute Dice value) and the worst case (with the least dice improvement) in Figure 2 and Figure 12, respectively. In both cases, PIRATE+ can gain better performance than other baseline models.

Reviewer comment (Reproducibility):
Memory consumption needs to be reported according to fixed size of datasets.

Author response:
Prompted by your comment, the revised manuscript shows the memory consumption on both training and inference stages of PIRATE, PIRATE+ and other baselines in Table 6. Note that PIRATE and PIRATE+ has lower memory complexity compared to other DL baselines.

Reviewer comment (Writing):
There was virtually no intuitive explanation of what plug-and-play does, why it is useful compared to other approaches, etc. Given that one of the reported contributions of the paper is to show that denoising priors can be used for registration, the lack of any explanation of why denoising is appropriate, how it works, etc. is quite conspicuous.

Author response:
The recent literature has shown the potential of training neural network denoisers to model the statistical distribution of high-dimensional data (e.g., denoising diffusion models and plug-and-play priors). Our work is the first to adopt this idea for registration by training a **denoiser prior** over the registration fields. The revised manuscript has clarified the choice of using denoiser prior for image registration. Please also refer to our analysis in section 3.1.

Reviewer comment (Novelty):
What is the main drawback with current deformable image registration models that this current approach is addressing? Why might plug and play be better than other forms of priors for registration? The introduction does not address this.

Author response:
PIRATE and PIRATE+ do not focus on addressing issues in existing DL methods. Our approach instead is a novel optimization method for image registration that leverages a CNN-based prior over registration fields. Our paper shows superior performance using the learned prior compared to handcrafted priors (such as smoothness), while also achieving the SOTA performance compared to the DL methods.

Reviewer comment (Writing):
The methods section is poorly written, without high-level insights presented before details. For example, the methods starts immediately with the PIRATE iteration updates instead of presenting what the objective being optimized is, what the prior is capturing, etc.

Author response:
Prompted by your comment, we added high-level insights of PnP/Denoiser prior in the introduction. We also included the objective function in the method section of the revised manuscript as below.
	
The object function that PIRATE optimizes is formulated as a variation of (1)

$\pmb{\hat{\phi}} = \arg \mathop{\min}\limits_{\pmb{\phi}} \lbrace g(\pmb{f},\pmb{\phi} \circ \pmb{m}) + h(\pmb{\phi}) \rbrace$

where $h$ represents the regularizer consisting of an explicit smoothness constraint and the implicit denoising regularizer. The iteration of PIRATE can be expressed as……..

Reviewer comment (Theory):
It is unclear what issue PIRATE+ is trying to improve upon. The methods section simply says "PIRATE+ uses DEQ to fine-tune the regularizer D in the PIRATE iteration by minimizing...", but there was no explanation of why the regularizer might not be accurate in the first place.

Author response:
The regularizer in PIRATE+ is learned by minimizing **the loss specific to the registration task computed using the registration field, the warped image and the fixed image**. On the other hand, the regularizer in PIRATE is learned by minimizing MMSE loss between denoised registration fields and noisy registration fields, without incorporating any specific information about registration tasks. For the specific task PIRATE+ is trained on, PIRATE+ can gain improved performance compared to PIRATE. The revised manuscript has highlighted this difference.

Reviewer comment (Presentation):
The method section would benefit by using pointers back to Fig. 1 and explaining the various steps of that figure.

Author response:
We added a pointer back to Fig. 1 in section 3 paragraph 1 of the revised manuscript.

Reviewer comment (Novelty):
What advantages do you think plug and play priors offer for registration?

Author response:
The recent literature has shown the potential of training neural network denoisers to model the statistical distribution of high-dimensional data (e.g., denoising diffusion models and plug-and-play priors). Our work is the first to adopt this idea for registration by training a **denoiser prior** over the registration fields. The revised manuscript has clarified the choice of using denoiser prior for image registration. Please also refer to our analysis in section 3.1.

Reviewer comment (Theory):
Why might denoising be appropriate as a prior for registration?

Author response:
The recent literature has shown the potential of training neural network denoisers to model the statistical distribution of high-dimensional data (e.g., denoising diffusion models and plug-and-play priors). Our work is the first to adopt this idea for registration by training a **denoiser prior** over the registration fields. The revised manuscript has clarified the choice of using denoiser prior for image registration. Please also refer to our analysis in section 3.1.

Reviewer comment (Writing):
I am not able to understand what PIRATE+ is addressing. The methods says "PIRATE+ fine-tunes the AWGN denoiser into a task-specific regularizer using DEQ" but this is quite opaque to me. What does "task-specific regularization" mean in this case?

Author response:
The regularizer in PIRATE+ is learned by minimizing **the loss specific to the registration task computed using the registration field, the warped image and the fixed image**. On the other hand, the regularizer in PIRATE is learned by minimizing MMSE loss between denoised registration fields and noisy registration fields, without incorporating any specific information about registration tasks. For the specific task PIRATE+ is trained on, PIRATE+ can gain improved performance compared to PIRATE. The revised manuscript has highlighted this difference.

Reviewer comment (Experiments):
Are there any limitations of the method compared to traditional CNN-based deformation models like VoxelMorph in terms of usability, training time, etc.?

Author response:
The training time of PIRATE+ is longer than other traditional CNN-based methods. The additional training time comes from the fixed point iterations of PIRATE. However, our results show that PIRATE+ can gain noticeable improvement (more than 3% improvement in Dice score) than traditional CNN-based methods. The revised manuscript has clarified this limitation in the conclusion section.