A Plug-and-Play Image Registration Network
Authors: Junhao Hu
Weijie Gan
Zhixin Sun
Hongyu An
Ulugbek Kamilov
Venue: ICLR 2024 poster (2024)
Keywords: deformable image registration
plug-and-play priors
deep equilibrium models
iterative algorithms
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We propose the first plug-and-play methods for deformable image registration: (a) PIRATE that uses deep denoiser trained on registration fields as prior, and (b) PIRATE+ that improves PIRATE by using deep equilibrium models to fine-tune the prior.

Abstract: Deformable image registration (DIR) is an active research topic in biomedical imaging. There is a growing interest in developing DIR methods based on deep learning (DL). A traditional DL approach to DIR is based on training a convolutional neural network (CNN) to estimate the registration field between two input images. While conceptually simple, this approach comes with a limitation that it exclusively relies on a pre-trained CNN without explicitly enforcing fidelity between the registered image and the reference. We present plug-and-play image registration network (PIRATE) as a new DIR method that addresses this issue by integrating an explicit data-fidelity penalty and a CNN prior. PIRATE pre-trains a CNN denoiser on the registration field and "plugs" it into an iterative method as a regularizer. We additionally present PIRATE+ that fine-tunes the CNN prior in PIRATE using deep equilibrium models (DEQ). PIRATE+ interprets the fixed-point iteration of PIRATE as a network with effectively infinite layers and then trains the resulting network end-to-end, enabling it to learn more task-specific information and boosting its performance. Our numerical results on OASIS and CANDI datasets show that our methods achieve state-of-the-art performance on DIR.

Published as a conference paper at ICLR 2024
A PLUG-AND-PLAY IMAGE REGISTRATION NETWORK
Junhao Hu; Weijie Gan; Zhixin Sun, Hongyu An, Ulugbek S. Kamilov Washington University in St. Louis, St. Louis, MO, USA {hjunhao, wei gan, zhixin. sun, hongyuan, kamilovjewustl.cd
ABSTRACT
Deformable image registration (DIR) is active research topic in biomedical imaging. There is a growing interest in developing DIR methods based on deep learning (DL). A traditional DL approach to DIR is based on training a convolu- tional neural network (CNN) to estimate the registration field between two input images. While conceptually simple, this approach comes with limitation that it exclusively relies on pre-trained CNN without explicitly enforcing fidelity be- tween the registered image and the reference. We present plug-and-play image registration network (PIRATE) as a new DIR method that addresses this issue by integrating an explicit data-fidelie penalty and CNN prior. PIRATE pre-traim a CNN denoiser on the registration field and "plugs" it into an iterative method as a regularizer. We additionally present PIRATE+ that fine-tunes the CNN prior in PI- RATE using deep equilibrium models (DEQ). PIRATE+ interprets the fixed-point iteration of PIRATE as a network with effectively infinite layers and then trains the resulting network end-to-end, enabling it to learn more task-specific informa- tion and boosting its performance. Our numerical results on OASIS and CANDI datasets show that our methods achieve kate-of-thee art performance on DIR.
INTRODUCTION
Deformable image registration (DIR) is important component in modern medical imaging (Soti- ras et al., 2013; Haskins et al.. 2020; Maintz & Viergever, 1998). DIR seeks to estimate a dense registration field align voxels between a moving image and fixed image. DIR has become an ac- tive research area with many applications, including radiation therapy planning (Brock et al., 2010), disease progression tracking (Ashburner & Friston, 2000), and image-enhanced endoscopy (Uneri et al., 2013). DIR is often formulated as an optimization problem that minimizes an energy function composed of two terms: a penalty function measuring the similarity between the aligned image and the fixed image, and regularizer imposing prior constraints on the registration field (e.g., smooth- ness via the gradient loss penalty). This optimization problem is usually solved using iterative methods (Thirion, 1998; Bajesy& & Kovačič, 1989; Rueckert et al., 1999; Glocker et al., 2008).
Deep learning (DL) has recently gained importance in DIR due to its promising performance (Hask- ins al., 2020: Fu et al., 2020; Boveiri et al., 2020). A common approach in this context is to use a convolutional neural network (CNN) to directly estimate the registration fields between the moving and fixed images. The training process uses energy function the loss function. How- ever, a potential limitation of this approach that its results exclusively depend on the pre-trained CNN, without incorporating an explicit penalty function for imposing data consistency during in- ference. Several studies have explored the integration of explicit penalty and CNN priors in the inference phase by training a recursive network comprising multiple blocks that incorporate both explicit penalty and CNNs (Qiu et al., 2022; Jia et al., 2023; Wang et al., 2023b). In the context of inverse problems imaging, such methods are often referred to as model-based deep learning (MBDL) (Ongie et al., 2020; Wen et al., 2023; Kamilov et al., 2023).
Plug-and-play (PnP) priors are widely used MBDI. frameworks for solving inverse problems in imaging (Venkatakrishnan et al., 2013). The central idea of PnP is that one can use pretrained neural network as a Gaussian denoiser as a regularizer for solving any inverse problem. The motiva- tion behind using denoiser prior is that training neural network denoisers can model the statistical
Denotes co-first author.

---

Published as a conference paper at ICLR 2024
PIRATE
PIRATE+
DEQ Update
Converged
Iteration
Iteration 2
Iteration oc
PIRATE
Fixed Point
Loss
Fixed Image
Field "
Warped Image f
Iteration
Field d
Loss
Loss Jocoiin
STN
Energy Function g w.
Moving Image
00
Field db
Warped Image f"
Denoiser
Update d
Lostaretlrrty
Figure I: Illustration the PIRATE and PIRATE+ pipelines. PIRATE updates the registration field & using the penalty function that measures the similarity between the warped image and the fixed image, as well as pre-trained CNN denoiser used as a regularizer. The DEQ update PIRATE+ enables fine-tune the CNN by calculating the gradients using implicitly differentiation through the fixed point of the forward iteration. As described in this paper, the DEQ update PIRATE+ is computed using the weighted loss consisting of similarity loss, smoothness loss, and Jacobian loss.
distribution of high- dinensiiien data. PnP has been successfully used in many imaging applications such as super-resolution, phase retrieval, microscopy, and medical imaging (Nachaoui et al., 2021; Wu al., 2019; Zhang al., 2017b; Sun al., 2019; Wei et al., 2020; Liu al., 2020; Zhang et al., 2019) (see also reviews Ahmad et al. (2020); Kamilov et al. (2023)). Practical success of PnP has also motivated novel extensions, theoretical analyses, statistical interpretations, as well as connections to related approaches such as score matching and diffusion models (Chan et al., 2016; Romano al., 2017; Buzzard al., 2018; Reehorst Schniter, 2018; Ryu et al., 2019; Kadkhodaie & Simoncelli, 2021; Cohen et al., 2021; Hurault et al., 2022a;b; Laumont et al., 2022; Gan et al., 2023a).
Despite the rich literature on PnP, the existing work on the topic has primarily focused using denoisers to specify priors on images. To the best our knowledge, the potential denoisers to specifyy priors over registration fields in DIR has never been investigated before. We address this gap by proposing a new DIR method slug-and-play Image RegistrATion nEtwork (PIRATE). PIRATE is the first PnP approach that trains CNN-based denoiser on the registration field and integrates this denoiser as a regularizer within iterative methods. We additionally present PIRATE+ as an ex- tension of PIRATE based on deep equilibrium learning (DEQ) (Bai al., 2019) that can fine-tune the denoiser to learn more task-specific information from the training dataset. PIRATE+ interprets network with effectively infinite layers and trains the
the fixed-point iteration of ATR sulting network end-to-end by calculating implicit gradients based on fixed point convergence. We propose three-term loss function in DEQ training: normalized cross-correlation (NCC), gradient loss, and Jacobian loss. We present an extensive validation of PIRATE and PIRATE+ on two widely used datasets: OASIS (Marcus et al., 2007) and CANDI (Kennedy et al., 2012). We used Dice similarity coefficient (DSC) and ratio of negative Jacobian determinant (JD) to evaluate the quality of registration. On both two datasets and metrics, the experimental results show that both PIRATE and PIRATE+ achieves state-of-the-arr performance. This work thus addresses a gap in the current literature by providing a new application to PnP and introducing a new principled framework for infusing prior information registration fields in DIR.
2 RELATED WORK
Deformable image registration. DIR refers to the process of obtaining registration field that maps the coordinates of the moving image m to those of the fixed image f by comparing the content of the corresponding image (Sotiras et al., 2013; Haskins et al., 2020; Maintz & Viergever, 1998). DIR is typically formulated as an optimization problem
b arg min (g(f, Φo m) ) + r()))
(1)
2

---

Published as a conference paper at ICLR 2024
where Φ represents aligned (warped) image obtained by using Φ warp m, 9 is the penalty function measuring the dissimilarity between the aligned and fixed images, and r is a regularizer imposing prior constraints the registration field. Examples of g in DIR include mean squared error (MSE), global cross-correlation (GCC), and normalized cross-corelltiion (NCC). A commonly used r is the smoothness regularizer implemented by the gradient loss | (Balakrishnan al., 2019; Mok & Chung, 2020; Wu et al., 2022), where V denotes to the gradient of Φ.
Iterative algorithms. Iterative algorithms are traditional approaches to solve the optimization problem in (1) on each image pair with satisfactory performance. Widely used examples include Demons (Thirion, 1998; Vercauteren al., 2009), elastic type methods (Bajesy & Kovačič, 1989; Davatzikos. 1997), b-splines based methods (Rueckert et al., 1999; Xie & Farin, 2004), and discrete methods (Glocker et al., 2008; 2011). The key idea of these methods to update the registration field with the gradient of an energy function consisting of their unique design of the penalty func- tion and regularizers. For example, Demons (Thirion, 1998) uses the first derivative of intensity difference and Gaussian smoothing regularizer. Log-Demon (Vercauteren et al., 2009) improves Demons by using the second derivative. Its regularizer consists of a Gaussian fluid-like regularizer and Gaussian diffusion-like like regularizer.
Deep learning. DL in DIR has gained widespread attention over the last few years due to its ex- cellent performance (Haskins et al., 2020; Fu et al., 2020; Boveiri al., 2020). A traditional ap- proach of DL is to train a CNN no that models & following & = no(f. m)) where O are network parameters updated by training loss (Balakrishnan et al., 2019; Mok & Chung, 2020; De Vos et al., 2019; 2017; Eppenhof et al., 2018). This method usually involves spatial transform net- work (STN) (Jaderberg et al., 2015) ensure a differentiable warping operator o. For example, Voxelmorph (Balakrishnan et al., 2019) uses the U-net as ne to exploit features from the image pairs. The moving image warped by using the output registration field in STN. Another approach, SYMNet (Mok & Chung, 2020), involves fully convolutional network (FCN) as pap. It performs stronger topology preservation and invertibility by learning the symmetric deformation fields that also formulates the inverse transformation. Moreover, recent DI. approaches such as reinforcement learning (Luo et al., 2022; Hu et al., 2021; Lotfi et al., 2013; Liu & Liu, 2022) and generative mod- els (Kim et al., 2022; Tanner et al., 2018; Mahapatra et al.. 2018; Zhang et al., 2020) have shown their potentials in DIR. However, the output of DL methods in the inference phase exclusively de- pends on the pre-traine traitrained CNN without incorporating the penalty function that can promote data consistency.
Deep Unfolding. DU is DL paradigm with roots in sparse coding (Gregor & LeCun, 2010; Chen et al., 2018h) that has recently been used in DIR (Qiu et al., 2022; Jia et al., 2023; Wang et al., 2023b), due to its ability to provide systematic connection between iterative algorithms and CNN architectures. For example, GraDIRN (Qiu et al., 2022) updates & in iterations of form t o' >Vg'(V) + Vno (f, m, Φ')). A practical limitation of DU is that the number of iterations is usually limited due to the the high computational and memory complexity of end-to-end training. In GraDIRN, this limitation
reflected in its network structure that contains only iterations.
Plug-and-play PnP is a MBDL approach that has been extensively used for solving inverse prob- lems imaging (Venkatakrishhan al., 2013; Romano et al., 2017; Kamilov et al., 2023). The key idea behind PnP is that one can use pre-trained deep denoisers to specify regularizers for various inverse problems. The key difference between PnP and DU is that PnP does not seek to train a MBDL architecture by instead focusing on an easier task of traininig a Gaussian. For example, widely used variant of PnP is the steepest descent regularization by denoising (SD-RED) (Romano et al., 2017)
x (+1 = TTox'')= = =z' -''->'''a- 1 T(=' (aa''
(2)
where 9 denotes the dita-fidelity term that measures data consistency. The data-fidelity term is conceptually similar to the use of g in (1) for DIR. The operator D denotes an additive white Gaussian noise (AWGN) denoiser (Zhang et al., 2017a; 2021; Rudin et al., 1992; Dabov et al., 2007) on æ. It is worth noting that PnP based methods typically train the denoiser on images. Recent work has shown the potential of training the denoiser on other modalities. For example, recent work (Gan et al., 2023a) has explored the effectiveness of training the denoiser the measurement operators, such as blur kernels. To the best of our knowledge, no prior work has explored using denosers within PnP to regularize registration fields in DIR.
3

---

Published as a conference paper at ICLR 2024
Deep equilibrium learning. DEQ (Bai et al., 2019; Winston & Kolter, 2020; Gilton et al., 2021) and neural ordinary differential equations (NODE) (Chen et al., 2018a; Dupont et al., 2019; Kelly et al., 2020) are two related methods for training infinite-depth, weight-tied feedforward networks. Both of them have shown their effectiveness in inverse problems (Pramanik & Jacob, 2022; Pra- manik et al., 2023; Liu al., 2022; Zhao et al., 2022). The key idea of NODE is represent the network with an infinite continuous-time perspective, while DEQ updates the network by analyt- ically backpropagating through the fixed points using implicit differentiation. The DEQ output specified implicitly as fixed point of the operator To parameterized by weights O
T To
In forward pass, DEQ runs fixed-point iteration to estimate ã. The DEQ backward pass pro- duces the gradients with respect to 0 by implicitly differentiating through the fixed point =00) For example, for the least-squares loss function, get analytical gradient
1 {(0) マ. (1( - x21 + VE(O) = (VoTo())) (1 (+ x²) (4)
where xx is the training label, and is the identity mapping. Two recent papers have investi- gated DEQ and ODE in DIR. REQT (Bai et al., 2022) extends the recurrent flow estimator RAFT (Teed & Deng, 2020) into infinite layers and updates the resulting network using implicit DEQ differentiation. NODEO (Wu et al., 2022) models each voxel in moving images as a moving particle and considers the set of all voxels as high-dimensional dynamical system whose trajectory determines the registration field. NODEO uses a neural network to represent the registration field and optimizes that dynamical system using adjoint sensitivity method (ASM) (Chen al., 201 8a) in NODE. The proposed PIRATÉ/PIRATE+ methods are different from these prior approaches since they combine an explicit penalty function g and CNN prior. DEQ-RAFT does not have any penalty function, while NODEO does not involve learning from training data.
Our contribution. (1) Our first contribution the use of learned denoisers for regularizing the reg- istration fields. The proposed PIRATE method can thus be seen as an extension of PnP framework to DIR. (2) Our second contribution a new DEQ method for fine-tuning the regularizer within PnP iterations. The proposed PIRATE+ method thus extends prior work on DEQ for DIR by inte- grating both penalty function measuring the dissimilarity between aligned and fixed images, and a CNN regularizer for infusing prior knowledge. (3) Our third contribution is the extensive validation of the proposed methods on OASIS and CANDI datasets. Our numerical results show the SOTA performance of PIRATE and PIRATE+, thus highlighting the potential of PnP and DEQ for DIR.
A PLUG-AND-PLAY IMAGE REGISTRATION NETWORK
In this section, we present two new methods: PIRATE and PIRATE+. Fig. 1 illustrates the pipelines of PIRATE and PIRATE+. PIRATE is new PnP algorithm that integrates pre-trained denoiser for estimating Φ. PIRATE+ fine-tunes the AWGN denoiser into a task-specific regularizer using DEQ. For each image pair during DEQ training, the forward iterations of PIRATE run until fixed point. The backward DEQ iterations update the regularizer by backpropagating through the fixed point using implicit differentiation of the loss.
3.1 PIRATE
The objective function that PIRATE optimizes is formulated as a variation of (1)
e arg min {g(f, &b m) + h(Φ)}
(5)
where h represents the regularizer consisting of an explicit smoothness constraint and the implicit
denoising regularizer. The iterations of PIRATE can be expressed as
b'+1 7 (Vg(f. m) + vr((") + T(0' D(o')))
(6)
where Ø' is the estimated registration field at iteration t, 7 > 0 is a step size, and a > and 7 >
regularization parameters.
4

---

Published as a conference paper at ICLR 2024
The PIRATE updates consist of three terms. The first term g is the penalty function measuring similarity between the aligned and fixed images, which in our implementation corresponds to the global cross correlation (GCC) function
gff Φ m) =
fmme Σ(- (Φ m (goom)
(7)
where j and a denote the mean and the standard deviation of the warped image o m and the fixed image f. The second term the classical smoothness promoting regularizer r()) = V
The third term is the PnP regularizer consisting of the residual term D(z), where D is pre- trained AWGN denoiser. We will consider denoisers trained to perform minimum mean squared error (MMSE) estimation of Φ from
with Φ Po. n ~ N(). MMSE denoisers can be analytically expressed as
2 =
(8)
[[ = /
- /
(G(=-))pole dΦp
((2)
Φ Pol ( z)
P=(=)
where we used the probability density function of the noisy registration field
pz(z)= = / Go - ())
(10)
The function Go in (9) denotes the Gaussian density with the standard deviation >> 0. The
remarkable property of the MMSE denoiser is that its residual satisfies (Robbins & Monro, 1951)
Φ (()) = - log P= (Φ).
(11)
which implies that PIRATE seeking to compute high probability registration fields from ps. which is a Gaussian smoothed version of the prior p
3.2 PIRATE+
PIRATE+ uses DEQ to fine-tune the regularizer D in the PIRATE iteration by minimizing the fol- lowing loss
fl f. m) = smismi(f (f. o m) + W1 lmme + (2f)sc(6)
(12)
where Wo, W1, and t2 are non-negative weights. The term fymt in (12) corresponds to the same smoothness regularizer in (6). The term smi is the normalized cross correlation (NCC) function
L 1- Σ o ((())
1
E(((r- us() @ ( m(r) H4po(rr)) (13)
where fr) and m(r) refer to the 3D sliding windows centered on the pixel (æx, y, 2) in the fixed image and the aligned image, respectively. The quantities and a denote the mean and the standard deviation in those windows. The term eaa in (12) represents Jacobian loss
()ac( Σ (
(14)
where [J(i denotes the Jacobian determinant (JD) each in Φ. Negative JD in registration field indicates non- non-physical transformations such as flipping. A value greater than one indicates local expansion. value less than but greater than zero indicates local compression, and a negative value indicates a physically implausible inversion. The (2-norm promotes the sparsity on negative JD by filtering positive values using ReLU activation function, leading physically plausible trans- formations in DIR. Note that regularizer in PIRATE+ is learned by minimizing the loss of a specific registration task. On the other hand, regularizer in PIRATE learned by minimizing a MSE loss of purely denoising noisy registration fields.
We adopt Jacobian-Free DEQ (JFB) (Fung al., 2022) to efficiently minimize the PIRATE+ loss function
Vere ()) = (VoTo(T)) (( (()(()) f, f))
(15)
where VF (0) is been descent direction for the loss function E with respect to 0 (Fung al., 2022). The effectiveness of JFB has been shows in prior work (Wang et al., 2023a; Gan et al., 2023b; Li al., 2022).
5

---

Published as a conference paper at ICLR 2024
Table 1: Numerical results of DSC, ratio of negative JD, and inference time for PIRATE, PIRATE+, and benchmarks on OASIS and CANDI datasets. The variances are shown in parentheses. Note that the best and the second-best result are highlighted in red and blue. The result shows that PI- RATE+ performs state-of-the arr performance comparing with other baselines. Moreover, PIRATE can achieve competitive performance without DEQ
OASIS
CANDI
Method
Time (s)
Avg. DSC f
Neg. JD % |
Avg. DSC t
Neg. JD % 4
SyN
0.6986 (0.028)
0.1471 (0.002)
0.7274 (0.049)
0.1211 (0.001)
42.52
VoxeIMorph
0.7655 (0.033)
0.0940 (0.001)
0.7459 (0.014)
0.1149 (0.020)
0.17
SYMNet
0.7664 (0.032)
0.0796 (0.001)
0.7489 (0.016)
0.0912 (0.001)
0.32
GraDIRN
0,7311 (0.054)
0.1418 (0.002)
0.7362 (0.019)
0.1532 (0.028)
0.49
Log- Bog-Demons
0.7931 (0.040)
0.0912 (0.007)
0,7553 (0.036)
0.0997 (0.006)
46.72
NODEO
0.7947 (0.027)
0.0531 (0.001)
0.7629 (0.019)
0.0726 (0.001)
89.01
PIRATE
0.7948 (0.033)
0.0567 (0.002)
0.7624 (0.020)
0.0833 (0.002)
32.01
PIRATE+
0.7952 (0.029)
0.0495 (0.001)
0.7633 (0.027)
0.0660 (0.001)
49.47
4 NUMERICAL EXPERIMENTS
4.1 SETUP
Datasets. We validated PIRATE and PIRATE+ on two widely used datasets: OASIS-1 (Marcus et al., 2007) and CANDI (Kennedy et al., 2012). The OASIS-1 dataset includes cross-sectional collection of 416 subjects with T1-weighted brain MRI scans and anatomical segmentations. We followed widely used preprocessing method', including affine pre-registration, skull stripping, and rescaling to 0-1. All images were padded to the size [192. 128, 224]. We used anatomical seg- mentations of 35 different brain structures provided by OASIS-I I for evaluation. The CANDI dataset comprises 54 l'--heettted brain MRI scans and anatomical segmentations. The preprocessing does not include affine pre-registration and skull stripping. since images in CANDI pre-registered and do not contain skulls. Besides the preprocessing methods used on OASIS, we excluded structures smaller than 1000 pixels in anatomical segmentations and used the remaining 28 structures for eval- uation in CANDI. For both datasets, we randomly shuffled the images and allocated 100 unique image pairs for training and another 100 unique image pairs for evaluation.
Evaluation metrics. We adopted two evaluation metries: Dice similarity coefficient (DSC) and ratio of negative JD, to evaluate the accuracy and quality of the registration. The DSC quantifies the spatial overlap between the anatomical segmentations of the fixed image and the warped image. We obtained the segmentation mask of the warped image by applying the associated registration field to the segmentation mask of the moving image. JD evaluates physical plausibility of the transformation by quantifying the expansion and compression of each voxel's neighbors. We employed the ratio of negative JD of each registration field to evaluate the quality of the transformation.
4.2 IMPLEMENTATION DETAILS
PIRATE setup. The penalty function in PIRATE is based on GCC, which, unlike NCC, relies on the whole image instead of sliding windows. Due lower computational complexity of GCC, we used it instead of NCC in the forward iterations. In our implementation, we used dynamic cosine-based step size in (6)
1 7
(1 + cos(
(16)
fmax
which we observed to give better results compared to the fixed y.
We chose DnCNN (Rudin et al., 1992) as a deep learning based AWGN denoiser. We obtained the ground truth for DrCNN training by using registration fields estimated by NODEO. In training

---

Published as a conference paper at ICLR 2024
VoxelMorph
GraDIRN
NODEO
Log-Demons
PIRATE+
Fixed Image
C66660
0.636
0.655
0.717
DICE SCORE
Figure 2: The visual results warped images (top) and correlated warped segmentation masks (bottom) from PIRATE+ and selected benchmarks on the OASIS dataset. The regions of interest are highlighted within red box. The DSC for each image displayed the bottom row. Note that the result of PIRATE+ is more consistent with the fixed image with fewer artifacts compared with other baselines.
Table Numerical results DSC, Jacobian determinant, and inference time for PIRATE+ and its ablated variants. The variances are shown in parentheses. We denote P as the penalty function, R as the denoiser regularizer, S as the smoothness regularizer, and as the DEQ refined denoiser. this case, PIRATE is formulated as P+R+S, and PIRATE+ is formulated as P+D+S. Note that the best and the second-best result are highlighted in red and blue. The result shows that the dual-regularizer structure and DEQ significantly reduce the negative Jacobian.
OASIS
CANDI
Method
Time (s)
Avg. DSC
JD
Avg. DSC
JD
0.7281 (0.044) 6.2931 (0.002)
0.7156 (0.047) 7.5804 (0.003)
20.20
P+R
0.7989 (0.023) 1.3319 (0.002)
0.7644 (0.001) 2.1713 (0.031)
28.37
P+S
0.7782 (0.012) 0.2091 (0.001)
0.7492 (0.019) 1.5847 (0.001)
25.64
P+D
0.7897 (0.008) 0.1 (0.001)
0.7585 (0.011) 1.1844 (0.001)
27.64
PIRATE: P+R+S
0.7948 (0.033) 0.0567 (0.002)
0.7624 (0.020) 0.0833 (0.002)
32.01
PIRATE+:
P+D+S 0.7952 (0.029) 0.0495 (0.001)
0.7633 (0.027) 0.0660 (0.001)
49.47
phase of the denoiser, used Adam (Kingma & Ba, 2014) optimizer with learning rate 1e-4 400 epochs. We selected the best persforeroninn denoiser from 10 denoisers trained on different noise levels (Gaussian noise with standard deviations starting from 1 to 10). Moreover, we downsampled the registration field to half the size of the image. The downsampled registration fields were interpolated back the original size before being applied to the moving image. Note that the denoiser was also trained on the downsampled data. We tested different values of of a, and in (6). PIRATE achieved the best performance by assigning 70 x 10a, a to 5 x 10-1 and 7 to 1 x 10-7 for OASIS- 1. a was assigned to 5 x 10 for CANDI while 70 and stayed the same.
DEQ training in PIRATE+. We used the Anderson solver (Anderson, 1965) to accelerate the fixed-pone iteration in PIRATE. For DEQ in PIRATE+, used Adam optimizer with learning rate x 10 for 50 epochs. We assigned wy to 1, w/y to 5, W2 to for both datasets. Note that only deep learning based denoiser (DnCNN in experiment) can be fine-tuned by DEQ.

---

Published as a conference paper at ICLR 2024
P+R
P+5
P+D
PIRATE: P+R+5
PIRATE+: P+D+S
J0X 99849
2.067
0.859 0.206
0.104
0.077
The visual results for the warped grid (top) and negative JD (yellow points bottom) of PIRATE+ and its five ablated variants on CANDI dataset. We denote P as the penalty function, R as the denoiser regularizer, P+R+S, and PIRATE+ is formulated as P+D+S. The gradient loss is shown in the top row, and the ratio of negative JD is shown the bottom row. Note that PIRATE's architecture is optimal, since it significantly reduces the negative JD and provides smoother registration field.
Figure 3:
the smoothness regularizer, and D as DEQ. PIRATE formulated as
Table 3: Numerical results of DSC, Jacobian determinant, and inference time for PIRATE+ with fixed and dynamic step size in its forward iteration on OASIS and CANDI datasets. The variances are shown parentheses. Note that the optimal results are highlighted in red. The result shows that the dynamic step size in the forward iteration improves the overall performance.
OASIS
CANDI
Method
Time (s)
Avg. DSC 1
JD 4
Avg. DSC f
JD L
Fixed
0.7902 (0.044)
0,0742 (0.001)
0.7556 (0.071)
0.0988 (0.002)
31.42
Dynamic
0.7952 (0.029)
0.0495 (0.001)
0.7633 (0.027)
0.0660 (0.001)
32.01
4.3 RESULTS
Comparison with benchmarks. We compared PIRATE and PIRATE+ against three iterative meth- ods (SyN (Avants et al., 2008), bog-Demons (Vercauteren et al., 2009), and NODEO (Wu et al., 2022)), two CNN based methods (VoxeIMorph (Balakrishnan Chung, 2020)), and one DU based method (GraDIRN (Qiu et al., 2022)). We retrained all DL based methods and DU based method using the optimal hyperparameters mentioned in their papers.
al., 2019) and SYMNet (Mok &
Table 1 shows the numerical results of PIRATE, PIRATE+ and benchmarks on OASIS and CANDI datasets. On both two datasets and metrics, PIRATE+ shows superior performance compared with other benchmarks. It is worth noting that PIRATE: outperforms NODEO, although we used NODEO to generate the ground truth for the denoiser. The results of PIRATE show that the pre- trained AWGN denoiser can provide competitive results without refinement of the DEQ. Moreover, PIRATE+ performs competitive inference time while maintaining state-of-the-ae performance com- pared with iterative baselines.
Fig. 2 shows the visual results of warped images (top) and correlated warped segmentation masks (bottom) from PIRATE+ and benchmarks the OASIS dataset. In the top row, the result of PI- RATE+ the region of interest (ROI) is more consistent with the fixed image with fewer artifacts than other benchmarks. In the bottom row, the DSC shows that PIRATE+ achieves better structure
8

---

Published as a conference paper at ICLR 2024
Image
Itrraion
Iteration
Iteration
Image
I10 -0000 1,"""1""
10'
20
10-
100
200
900
400
Figure 4: The visualization and plot for the warped image from PIRATE+ with different iterations. The error map of ROI is shown at the upper right corner. The results show a convergence behaviour of PIRATE+.
matching. The ROI in the output of PIRATE+ shows better alignments to the fixed image while maintaining the original anatomical structure.
Ablation study. We conducted an ablation study to demonstrate the effectiveness of each part in PIRATE+. We denoted P as the penalty function, as the denoiser regularizer, as the smoothness regularizer, and D DEQ refined denoiser regularizer. Note that PIRATE and PIRATE+ can be expressed as + + R and + S + D, respectively.
Table 2 shows the numerical results of PIRATE+ and its five ablated variants. The results show that the dual-regularizer structure and DEQ of PIRATE+ significantly reduce the ratio of negative JD with acceptable decrement in DSC.
Fig. visualizes the warped grid (top) and negative JD (yellow points in bottom) from PIRATE+ and its five different variants. It is evident that the two regularizers and DEQ contribute to smooth registration fields while maintaining low ratio of negative JD, which underlines the effectiveness of the PIRATE+'S architecture.
Convergence. We designed an experiment to validate the convergence of PIRATE+. Fig. 4 illus- trates the warped images from PIRATE+ with different iterations on CANDI dataset. The result shows a convergence behavior.
Step size. We tested the effectiveness of the dynamic step size compared with the fixed step size. Table shows the numerical results of using fixed step size and dynamic step size in PIRATE+ on OASIS and CANDI datasets. The result shows that the dynamic step size improves the overall performance of PIRATE+.
5 CONCLUSION
This paper presents PIRATE as the first PnP-based method for DIR and PIRATE+ as its further DEQ extension. The key advantages of PIRATE and PIRATE+ are follows: (a) PIRATE achieves competitive performance in comparison to other baselines. This superior performance is attributed to its ability seamlessly integrate the penalty function with learned CNN priors; (b) PIRATE treats the CNN prior as an implicit regularizer, which renders it compatible with other regularizers. For instance, in experiments, combined it with the gradient loss further enhance its perfor- mance; (c) ATRATE+ achieves state-of-the-an performance in comparison to other baselines. DEQ PIRATE+ fine-tuned the denoiser PIRATE, enabling it learn more task-specific informa- tion. The training time of PIRATE+ is longer than traditional DL methods. However, our results show that PIRATE+ can gain noticeable improvement than traditional CNN-based methods. It's also worth mentioning that our validation of PIRATE and PIRATE+ in this work was based on brain MRI datasets. Potential future endeavors might explore their adaptability to different anatomies and imaging modalities, such as whole-body registration and cross-mode modsl registration.
ACKNOWLEDGEMENT
This work was supported by the following grants: NIH RO1HLI29241, RFINS116565, R21NS127425, Ro1 EB032713, and the NSF CAREER award under grants CCF-2043134.

---

Published as a conference paper at ICLR 2024
REFERENCES
R. Ahmad, C. A. Bouman, G. T. Buzzard, S. Chan, S. Liu, E. T. Reehorst, and P. Schniter. Plug-and- play methods for magnetic resonance imaging: Using denoisers for image recovery. IEEE Signal Process. Mag., 37(1): 2020.
D. G. Anderson. 1965.
Iterative procedures for nonlinear integral equations. ACM, 12(4))547-560,
Ashburner and K. J. Friston. Voxel-based morphometry-thhe the methods. Neurolmage, 11(6))805- 821, 2000.
B. B. Avants, C. L. Epstein, M. Grossman, and J. C. Gee. Symmetric diffeomorphic image registra- tion with cross-correlation:: evaluating automated labeling of elderly and neurodegenerative brain. Med. Image Anal., 12(1):26-44 2008.
Bai, Z. Kolter, and V. Koltun. Deep equilibrium models. Advances Neural Information Processing Systems (NeurIPS), 32, 2019.
S. Bai, Z. Geng, Y. Savani, and J. Z. Kolter. Deep equilibrium optical flow estimation. Proc. IEEE Conf. Computer Vision and Pattern Recognition (CVPR), pp. 620-630, 2022.
R. Bajesy and S. Kovačič. Multiresolution elastic matching. Comput. graph. image process., 46(1): 1-21, 1989.
G. Balakrishnan, A. Zhao, M. R. Sabuncu, J. Guttag, and A. V. Dalca. Voxelmorph: a learning framework for deformable medical image registration. IEEE Trans. Med. Imag., 38(8):1788- 1800, 2019.
H. R. Boveiri, R. Khayami, R. Javidan, and Mehdizadeh. Medical image registration using deep neural networks: comprehensive review. Comput. Electr. Eng., 87:106767, 2020.
K. K. Brock, Deformable Registration Accuracy Consortium, et al. Results of a multi multi-institution deformable registration accuracy study (midras). Int. J. Radiat. Oncol. Biol. Phys., 76(2):583- 596, 2010.
G. T. Buzzard, S. H. Chan, S. Sreehari, and C. A. Bouman. pld--and unplugged: Optimization- free reconstruction using consensus equilibrium. SIAM J. Imaging Sci., 11(3):2001-2020, 2018.
Chan, X. Wang, and O. A. Elgendy. Plug-and-play admm for image restoration: Fixed-point convergence and applications. IEEE Trans. Comput. Imag., 3(1):84 98, 2016.
R. T. Q. Chen, Y. Rubanova, Bettencourt, and D. K. Duvenaud. Neural ordinary differential equations. Advances in Neural Information Processing Systems (NeurIPS), 31, 2018a.
X. Chen, Liu, Z. Wang, and W. Yin. Theoretical linear convergence of unfolded ISTA and its prac- tical weights and thresholds. In Advances in Neural Information Processing Systems (NeurIPS), 2018b.
pp.
R. Cohen, Y. Blau, D. Freedman, and Rivlin. It has potential: Gradient-driven denoisers for convergent solutions to inverse problems. In Advances in Neural Information Processing Systems (NeurIPS), pp. 1152-18164, December 6-14, 2021.
K. Dabov, A. Foi, V. Katkovnik, and K. Egiazarian. Image denoising by sparse 3-d transform- domain collaborative filtering. IEEE Trans. Image Prucess., 16(8):0800-0955, 2007.
C. Davatzikos. Spatial transformation and registration of brain images using elastically deformable models. Comput. Vis. Image Underst., 66(2):207-222, 1997.
B. D. De Vos. F. F. Berendsen, M. A. Viergever, M. Staring, and I. Iggum. End-to-end unsupervised deformable image registration with a convolutional neural network. Medical Image Computing and Computer-Assister Intervention Workshop (MICCA/W), pp. 204-212, 2017.
10

---

Published as a conference paper at ICLR 2024
B. D. De Vos, F. F. Berendsen, M. A. Viergever, H. Sokooti, M. Staring, and I. 1šgum. A deep learning framework for unsupervised affine and deformable image registration. Med. Image Anal., 52:128- 103,-1432222 2019.
E. Dupont, A. Doucet, and Y. W. Teh. Augmented neural odes. Advances in Neural Information Processing Systems (NeurIPS), 32, 2019.
K. A. J. Eppenhof, M. W. Lafarge, P. Moeskops, M. Veta, and J. P. W. Pluim. Deformable image registration using convolutional neural networks. Med. Image Anal., 10574:192-197, 2018.
Y. Fu, Y. Lei. T. Wang, J. Curran, T. Liu, and X. Yang. Deep learning in medical image registra- tion: review. Phys. Med. Biol., 65(20):20TRO, 2020.
S. W. Fung, H. Heaton, Li, D. McKenzie, S. Osher, and W. Yin. Jfb: Jacobian-free backpropaga- tion for implicit networks. Proc. AAAI Conf. Artif. Intell., 36(6) 6656, 2022.
W. Gan, S. Shoushtari, Y. Hu, J. Liu, H. An, and U. S. Kamilov. Block coordinate plug-and- play methodsfor blind inverse problems. Advances in Neural Information Processing Systems (NeurIPS), 2023a.
W. Gan, C. Ying. E. Boroojeni, T. Wang, C. Eldeniz, Y. Hu, J. Liu, Y. Chen, An, and U. S. Kamilov. Seff-supervised deep equilibrium models with theoretical guarantees and applications to mri reconstruction. JEEE Trans. Comput. Imag., 2023b.
D. Gilton, G. Ongie, and R. Willett. Deep equilibrium architectures for inverse problems in imaging. IEEE Trans. Comput. Imag. 1:112331133, 2021. 2021.
B. Glocker, N. Komodakis, G. Tziritas, N. Navab, and N. Paragios. Dense image registration through mrfs and efficient linear programming, Med. Image Anal., 12(6):731-741, 2008.
B. Glocker, A. Sotiras, N. Komodakis, and N. Paragios. Deformable medical image registration: setting the state of the art with discrete methods. Annu. Rev. Biomed. Eng., 13:219-244, 2011.
K. Gregor and Y. LeCun. Learning fast approximation of sparse coding. In Proc. 27th Int. Conf. Machine Learning (ICML), pp. 399-406, Haifa, Israel, June 21-24, 2010.
G. Haskins, U. Kruger, and P. Yan. Deep learning in medical image registration: a survey. Mach. Vis. Appl., 31:1-18, 2020.
J. Hu, Z. Luo, X. Wang. S. Sun, Y. Yin, K. Cao, Q. Song, S. Lyu, and X. Wu. End-to-end multimodal image registration via reinforcement learning. Med. Image Anal., 68:101878 2021.
S. Hurault, A. Leclaire, and N. Papadakis. Gradient step denoiser for convergent hlu-annd-lay. In Int. Conf. Learming Representations (ICLR), Kigali, Rwanda, May 1-5, 2022a.
S. Hurault, A. Leclaire, and N. Papadakis. Proximal denoiser for convergent plug-and-play opti- mization with nonconvex regularization. In Proc. 39th Int. Conf. Machine Learning (ICML), pp. 9483-9505, Baltimore, MD, July 17-23, 2022b.
M. Jaderberg, K. Simonyan, A. Zisserman, and Kavukcuoglu. Spatial transformer networks. Advances in Neural Information Processing Systems (NeurIPS), 28, 2015.
X. Jia, Thorley, A. Gomez, W. Lu, D. Kotecha, and J. Duan. Fourier-nett: Leveraging band- limited representation for efficient 3d medical image registration. arXiv:2307.02997 2023.
Z. Kadkhodaie and E. Simoncelli. Stochastic solutions for linear inverse problems using the prior implicit in denoiser. In Advances in Neural Information Processing Systems (NeurIPS), pp. 13242-13254, December 6- 64, 2021.
U. S. Kamilov, C. A. Bouman, G. T. Buzzard, and B. Wohlberg. Plug-and-pll and-play methods for integrat- ing physical and learned models in computational imaging: Theory, algorithms, and applications. IEEE Signal Process. Mag,, 40(1):85-97, 2023.
11

---

Published as a conference paper at ICLR 2024
J. Kelly, J. Bettencourt, M. J. Johnson, and D. K. Duvenaud. Learning differential equations that are easy to solve. Advances in Neural Information Processing Systems (NeurlPS). 33:4370- 2020.
D. N. Kennedy, C. Haselgrove, S. M. Hodge, P. S. Rane, N. Makris, and J. A. Frazier. Candishare: a resource for pediatric neuroimaging data. Neuroinformatiis, 10:319-322, 2012.
B. Kim, I. Han, and C. Ye. Diffusemorph: unsupervised deformable image registration using diffusion model. European Conf. on Comput. Vis., pp- 347-364, 2022.
D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. Int. Conf. Learning Representations (ICLR), 2014.
R. Laumont, V. De Bortoli, A. Almansa, J. Delon, A. Durmus, and M. Pereyra. Bayesian imaging using plug & play priors: When Langevin meets Tweedie. SIAM J. Imaging Sciences, 15(2): 701-737, 2022.
M. Li, Y. Wang. Y. Wang, and Z. Lin. Unbiased stochastic proximal solver for graph neural networks with equilibrium states. Int. Conf. Learning Representations (ICLR), 2022.
J. Liu, Y. Sun, C. Eldeniz, W. Gan, H. An, and U. S. Kamilov. Rare: Image reconstruction using deep priors learned without groundtruth. IEEE J. Sel. Top. Signal Process., 14(6):1088-1099, 2020.
J. Liu, X. Xu, W. Gan, S. Shoushtari, and U. S. Kamilov. Online deep equilibrium learning for regularization by denoising. Advances in Neural Information Processing Systems (NeurIPS), 35: 25363-28 0663-25376,, 2022.
T. Liu and J. Liu. Application of reinforcement learning in multimodal non-rigid image registration. Proc. CISC, pp. 7766384,777 2022.
Lotfi, L. Tang, S. Andrews, and G. Hamarneh. Improving probabilistic image registration via reinforcement learning and uncertainty evaluation. Medical Image Computing and Computer- Assisted Intervention Workshop (MICCA/W), pp. 187-194, 2013.
Z. Luo, J. Hu, X. Wang, Hu, B. Kong. Y. Yin, Q. Song. X. Wu, and S. Lyu. Stochastic planner- actor-critic for unsupervised deformable image registration. Proc. AAAI Conf. Artif. Intell., 36(2): 1917-1925, 2022.
D. Mahapatra, B. Antony, S. Sedai, and R. Gamavi. Deformable medical image registration using generative adversarial networks. Proc. Int. Symp. Biomedical Imaging, pp. 1449-1453, 2018.
J. B. A. Maintz and M. A. Viergever. A survey of medical image registration. Med. Image Anal., (1):1-36, 1998.
D. Marcus, T. H. Wang, J. Parker, G. Csernansky. C. Morris, and R. L. Buckner. Open access series of imaging studies (oasis): cross-sectional mri data in young. middle aged, nondemented, and demented older adults. J. Cogn. Neurosci.,
2007.
C. W. Mok and C. S. Chung. Fast symmetric diffeomorphic image registration with convolutional neural networks. Proc. IEEE Conf. Computer Vision and Pattern Recognition (CVPR), pp. 4644- 4653, 2020.
M. Nachaoui, L. Afraites, and A. Laghrib. A regularization by denoising super-resolution method based genetic algorithms. Signal Process.: Image Commun. 99:1 16505, 2021.
G. Ongie, A. Jalal, C. A. Metzler, R. G. Baraniuk, A. G. Dimakis, and R. Willett. Deep learning techniques for inverse problems in imaging. IEEE J. Sel. Areas Inf. Theory, 1(1):39-56, May 2020.
A. Pramanik and M. Jacob. Improved model based deep learning using monotone operator learning (mol). Proc. Int. Symp. Biomedical Imaging, pp. 1-4, 2022.
A. Pramanik, M. B. Zimmerman, and M. Jacob. Memory-efficient model-based deep learning with convergence and robustness guarantees. IEEE Trans. Comput. Imag., 2023.
12

---

Published as a conference paper at ICLR 2024
H. Qiu, K. Hammernik, C. Qin, C. Chen, and D. Rueckert. Embedding gradient- based optimization in image registration networks. Medical Image Computing and Computer- Intervention (MICCAI), pp. 56-65, 2022.
E. T. Reehorst and P. Schniter. Regularization by denoising: Clarifications and new interpretations.
IEEE Trans. Compul. Imag. 5(1):52-67, 2018.
H. Robbins and S. Monro. A stochastic approximation method. Ann. Math. Stat., 22(3):400-407, September 1951.
Y. Romano, M. Elad, and P. Milanfar. The little engine that could: Regularization by denoising (RED). SIAM J. Imaging Sci., 10(4):1804-1844, 2017.
L. I. Rudin, S. O, and E. Fatemi. Nonlinear total variation based noise removal algorithms. Phys. D: Nonlinear Phenom., 60(1100)))))))) (1):1449999922 2688, 1992.
D. Rueckert, L. I. Sonoda, Hayes, D. L. G. Hill, M. O. Leach, and D. J. Hawkes. Nonrigid registration using free-form deformations: application to breast mr images. IEEE Trans. Med. Imag., 1888):122-721, 1999.
E. K. Ryu, Liu, Wnag, Chen, Wang, and Yin. Plug-and-play methods provably converge with properly trained denoisers. In Proc. 36th Int. Conf. Machine Learning (ICML), pp- 5546-5557, Long Beach, CA, USA, June 2019.
A. Sotiras, C. Davatzikos, and N. Paragios. Deformable medical image registration: survey. IEEE Trans. Med. Imag., 32(7): 153-1190, 2013.
Y. Sun, S. Xu, Y. Li, L. Tian, B. Wohlberg, and U. S. Kamilov. Regularized fourier ptychography using an online plug-and-yg play algorithm. Proc. IEEE Int. Conf. Acoustics, Speech and Signal Process. (ICASSP)., pp. 7665-7669, 2019.
C. Tanner, F. Ozdemir, R. Profanter, V. Vishnevsky, E. Konukoglu, and O. Goksel. Generative adversarial networks for mr-et deformable image registration. arXiv:1807.07344 2018.
Z. Teed and J. Deng. Raft: Recurrent all-pairs field transforms for optical flow. European Conf. on Comput. Vis., pp. 402-419, 2020.
J. P. Thirion. Image matching diffusion process: analogy with maxwell's demons. Med. Image Anal., 2(3):243-260, 1998.
A. Uneri, S. Nithiananthan, S. Schafer, Y. Otake, J. W. Stayman, G. Kleinszig, M. S. Sussman, J. L. Prince, and J. H. Siewerdsen. Deformable registration of the inflated and deflated lung in cone- beam ct-guided thoracic surgery: Initial investigation of combined model-and image-driven approach. Med. Phys., 40(1):017501, 2013.
S. Venkatakrishnan, C. A. Bouman, and Wohlberg. Plug-and and-ala priors for model based reconstruction. Proc. IEEE Global Conf. Signal Process. and Inf. Process., pp. 945-948, 2013.
T. Vercauteren, X. Pennec, A. Perchant, and N. Ayache. Diffeomorphic demons: Efficient non- parametric image registration. Neurolmage, 45(1):561-77,, 2009.
Wang, Y. Teng, and I.. Wang. Deep equilibrium object detection. arXiv:2308.09564, 2023a.
Y. Wang, Qiu, and Qin. Conditional deformable image registration with spatially-variant and adaptive regularization. Proc. Inf. Symp. Biomedical Imaging, 2023b.
K. Wei, A. Aviles-Rives J. Liang, Y. Fu, C. B. Schönlieb, and H. Huang, Tuning-free plug-and- play proximal algorithm for inverse imaging problems. Proc. 37th Int. Conf. Machine Learning (ICML), pp. 10158 10169, 2020.
B. Wen, S. Ravishankar, Z. Zhao, R. Giryes, and J. C. Ye. Physics-driven machine learning for computational imaging. IEEE Signal Process. Mag., 40(1):28-30, 2023.
E. Winston and Z. Kolter. Monotone operator equilibrium networks. Advances in Neural Infor- mation Processing Systems (3: 13:10118-770 2020.
13

---

Published as a conference paper at ICLR 2024
Y. Wu, T. Z. Jiahao, J. Wang, P. A. Yushkevich, M. A. Hsich, and J. C. Gee. Nodeo: A neural ordinary differential equation based optimization framework for deformable image registration. Proc. IEEE Conf. Computer Vision and Pattern Recognition (CVPR), pp. 01804-208133 2022.
Z. Wu, Y. Sun, J. Liu, and U. S. Kamilov, Online regularization by denoising with applications phase retrieval. Proc. IEEE Int. Conf. Comp. Vis. (ICCV), pp. 0-0, 2019.
Z. Xie and G. E. Farin. Image registration using hierarchical b-splines. IEEE Trans. Vis. Comput. Graph., 10(1):85 94, 2004.
K. Zhang, W. Zuo, Y. Chen, D. Meng, and L. Zhang. Beyond a gaussian denoiser: Residual learning of deep for image denoising. IEEE Trans. Image Process., 26(7):314233155, 2017a.
K. Zhang, W. Zuoo S. Gu, and L.. Zhang. Learning deep enn denoiser prior for image restoration. Proc. IEEE Conf. Computer Vision and Pattern Recognition (CVPR), pp. 3929-3938, 2017b.
K. Zhang, W. Zuo, and L. Zhang. Deep plug-and-plas super-resolution for arbitrary blur kernels. Proc. IEEE Conf. Computer Vision and Puttern Recognition (CVPR). pp. 1671-1681, 2019.
K. Zhang, Y. Li, W. Zuo, L. Zhang. L. Gool, and R. Timofte. Plug-and-play image restoration with deep denoiser prior. IEEE Trans. Pattern Anal. Mach. Intell., (4(10))6360063777 2021.
X. Zhang, W. Jian, Y. Chen, and S. Yang. Deform-gan: unsupervised learning model for de- formable registration. arXiv:2002.11450, 2020.
Y. Zhao, S. Zheng, and X. Yuan. Deep equilibrium models for snapshot compressive imaging, Proc. AAAI Conf. Artif. Intell., 2022.
14

---

Published as a conference paper at ICLR 2024
A APPENDIX
Our main manuscript presents PIRATE and PIRATE+ as the first PnP based method DIR. Our ex- perimental results show the effectiveness of the architectures PIRATE and PIRATE+ and demon- strate their state-of-the-art performance. The additional materials in this section supplement our experimental results and further support the conclusions mentioned above. appendix, we provide the visualization of warped images and segmentation masks on CANDI dataset as supplement to the results on OASIS dataset in main manuscript. We also provide additional results on the EM- PIREIO lung CT dataset. In addition to the image similarity and DSC used in main manuscript, we illustrate the warped grids on OASIS and CANDI datasets to further evaluate the performance using the smoothness of registration fields. We design experiment by using total variation (TV) denoiser (Rudin et al., 1992) in PIRATE to demonstrate the compatibility of PIRATE with non- learning based denoisers. Moreover, we show the visualization of warped images and segmentation masks from PIRATE using dynamic and fixed step size on OASIS dataset, which matches the nu- merical results in main manuscript. To further evaluate our methods, computational efficiency and scalability. To clarify the acronyms, we provide a table of acronyms and corresponding full names in this paper. We also provide the corner case of method and the number of function evaluation(NFF)) along with the training iterations of PIRATE+.
include the analysis of
Fig. 5 shows the visual results for warped images and correlated warped segmentation masks from PIRATE+ and benchmarks on the CANDI dataset. In the top and bottom row, the results of PIRATE+ in ROI are more consistent with the fixed image with fewer artifacts than other benchmarks.
Table. 4 shows the numerical results of PIRATE, PIRATE+ and benchmarks on the EMPIREI0 lung CT dataset. PIRATE+ shows superior performance compared with other benchmarks. The results of PIRATE show that the pre-trained AWGN denoiser can provide competitive results without refinement of the DEQ.
Fig. 6 and Fig. 7 show the illustration of warped images and correlated warped segmentation masks from PIRATE+ and benchmarks on the EMPIREIO lung CT dataset. both warped images and seg- mentation masks, the results of PIRATE+ are more consistent with the reference than other bench- marks.
Fig. 8 shows the visual results of warped grid from OASIS and CANDI datasets of PIRATE and different baselines. The results show that PIRATE maintains the smoothness of the registration fields.
Fig. shows the visual results of using dynamic step size and fixed step size from OASIS and CANDI datasets. It shows that dynamic step size can provide results more similar to the fixed images and segmentation masks with less artifact.
Fig. 10 shows the visual results of using TV denoiser and DnCNN denoiser from OASIS dataset. Table shows the numerical result of using TV denoiser and DrCNN denoiser from OASIS and CANDI datasets. The results show that TV denoiser can still provide reasonable result, although with some performance degradation.
Table. 6 shows the numerical results of memory cost, running time in training and inference for PIRATE, PIRATE+, and benchmarks on OASIS and CANDI datasets. Table. shows that PIRATE and PIRATE+ has lower memory complexity compared to other DL baselines.
Fig. 1l shows the visual results of memory cost and inference time of PIRATE and PIRATE+ with different size of inputs. Fig. 11 shows that PIRATE and PIRATE+ have similar scalability.
Table. 7 shows acronyms and corresponding full names in this paper.
Fig. 12 shows the visual results of the corner case of PIRATE+ (with the least improvement in Dice) compared with results from other benchmarks on this case. Note that the result of PIRATE+ still maintains state-of-the-art performance in the corner case.
Fig. 13 shows the visual results of NFE along with the training iterations of PIRATE+. The figure shows that the NFE converges in the training process.
15

---

Published as a conference paper at ICLR 2024
VoxelMorph
GraDIRN
NODEO
Log-Demons
PIRATE+
Fixed Image
0.736
0.731
0.743
0.244 0.749
DICE SCORE
Figure 5: The visual results of warped images (top) and correlated warped segmentation masks (bottom) from PIRATE+ and selected benchmarks on the CANDI dataset. The regions of interest are highlighted within red box. The DSC for each image is displayed in the bottom row. Note that result PIRATE+ more consistent with the fixed image with fewer artifacts comparing with other basclines.
Table Numerical results of DSC, ratio of negative JD, and inference time for PIRATE, PIRATE+ and benchmarks on the EMPIRE1O lung CT dataset. The variances are shown in parentheses. Note that the best and the second-best result are highlighted red and blue. The result shows that PI- RATE+ performs state-of-the the- performance comparing with other baselines. Moreover, PIRATE can achieve competitive performance without DEQ
Method
Avg. DSC
Neg. JD % |
Time (s)
SyN
0.9246 (0.045)
0.0854 (0.008)
37.28
VoxelMorph
0.9681 (0.022)
0.0521 (0.012)
0.16
SYMNet
0.9701 (0.012)
0.0542 (0.009)
0.28
GraDIRN
0.9644 (0.012)
0.0542 (0.009)
0.44
Log- -og-Demons
0.9772 (0.019)
0.0571 (0.011)
43.29
NODEO
0.9802 (0.028)
0.0441 (0.008)
82.17
PIRATE
0.9797 (0.032)
0.0546 (0.009)
29.77
PIRATE+
0.9811 (0.027)
0.0467 (0.007)
45.32
16

---

Published as a conference paper at ICLR 2024
Fiked
PIRATE
NODEO
Demons
VM
GfaDIRN
SYMnet
SyN
Figure 6: The visual results of warped images from PIRATE+ and selected benchmarks the EMPIRE10 lung CT dataset. The regions of interest are highlighted within red box. Note that the result of PIRATE+ is more consistent with the fixed image with fewer artifacts comparing with other baselines.
Input
PIRATE+
NODEO
Demons
DICE: 0.9211
DICE: 0.9840
DICE: 0.9805
DICE 0.9531
VM
GraDIRN
SYMnet
SyN
DICE: 0.9796
DICE: 0,9764
DICE: 0.9742
DICE: 0.9601
Figure 7: The visual results warped segmentation masks from PIRATE+ and selected benchmarks on the EMPIRE 10 lung CT dataset. The red region is the mask from fixed image. and the gray region the misalignment of the warped mask. The DSC for each image is displayed in the bottom. Note that the result of PIRATE+ is more consistent with the fixed mask comparing with other baselines.
17

---

Published as a conference paper at ICLR 2024
VoxelMorph
GraDIRN
NODEO
Log-Demons
PIRATE+
OASIS
0.082
CANDI
0.07
Figure 8: The visual results of warped grid from OASIS (top) and CANDI (bottom) of PIRATE+ and different baselines. The gradient loss is shown the bottom of each image.
PIRATE(Fixed y)
PIRATE(Dynamicy) y
Fixed Image
0.701
0.729
DICE SCORE
Figure 9: The visual results of warped images (top) and correlated warped segmentation masks (bottom) from PIRATE(fixed 7) and PIRATE<dynamic 7) on the OASIS dataset. The regions of interest highlighted within red box. The DSC for each image displayed in the bottom row.
18

---

Published as a conference paper at ICLR 2024
Table 5: Numerical results of DSC, Jacobian determinant, and inference time for PIRATE with total variation (TV) denoiser and DrCNN denoiser on OASIS and CANDI datasets. The variances are shown in parentheses. Note that the optimal results are highlighted in red.
OASIS
CANDI
Method
Time (s)
Avg. DSC
JD
Avg. DSC
JD
PIRATE(TV)
0.7617 (0.054)
0.1281 (0.002)
0.7481 (0.049)
0.0992 (0.001)
118.20
PIRATE(DHCNN)
0.7948 (0.033) 0.0567 (0.002)
0.7624 (0.020)
0.0833 (0.002)
32.01
PIRATE(TV)
PIRATE(DACNN)
Fixed Image
0.673
0.729
DICE SCORE
Figure 10: The visual results of warped images (top) and correlated warped segmentation masks (bottom) from PIRATE(TV) and PIRATE(DiCNN) on the OASIS dataset. The regions of interest are highlighted within red box. The DSC for each image is displayed in the bottom row.
Scalability: Memory
Scalability: Time
14000
250
(spaar)
30000
8000
200
4000
50
2000
(80.
(160,
Input
inpur
Figure 11: The illustrations of memory cost and inference time of PIRATE and PIRATE+ with different size of inputs. This figure shows that PIRATE and PIRATE+ have similar scalability.
19

---

Published as a conference paper at ICLR 2024
Table 6: Numerical results of memory cost, rumning time in training and inference for PIRATE, PIRATE+, and benchmarks on OASIS and CANDI datasets. This table shows that PIRATE and PIRATE+ has lower memory complexity compared other DL. baselines.
Inference
Training
Method
Memory (MB) +
Running time (s) 4
Memory (MB)
Running time (s/epoch)
SyN
2252
42.52
VoxeIMorph
7642
0.17
12846
78
SYMNet
5976
0.32
16448
185
GraDIRN
4386
0.49
11642
112
Log-Demons
2709
46.72
NODEO
5534
89.01
PIRATE
3090
32.01
5408
30.02
PIRATE+
3314
49.47
8866
1821.90
Table 7: Acronyms and corresponding full names in this paper.
Acronyms
Full name
Acronyms
Full name
DIR
deformable image registration
deep learning
CNN
convolutional neural network
PIRATE
plug- auu--add-pln play image registration network
DEQ
deep equilibrium models
MBDL
model-based deep learning
PnP
blug-and-plas
NCC
normalized cros-coorclation
GCC
global cross-cores coss-correlation
DSC
Dice similarity coefficient
JD
Jacobian determinant
MSE
mean squared error
STN
spatial transform network
FCN
fully convolutional network
DU
deep unfolding
SD-RED
steepest descent regularization by denoising neural ordinary differential equations minimum mean squared error region of interest
AWGN
additive white Gaussian noise
NODE
ASM
adjoint sensitivity method
MMSE
JFB
Jacobian-Free deep equilibrium models
ROI
TV
total variation
20

---

Published as a conference paper at ICLR 2024
VoxelMorph
GraDIRN
NODEO
Log-Demons
PIRATE+
Image
DICE SCORE
The visual results of warped images (top) and correlated warped segmentation masks (bottom) from the corner case of PIRATE+ (with the least improvement (nor absolute value) in Dice) compared with results from other benchmarks on this case. The regions of interest are highlighted within a red box. The DSC for each image is displayed in the bottom row. Note that the result of PIRATE+ is still more consistent with the fixed image on this corner case.
Figure 12:
NFE's progression across training iterations
413
N
412
Forward iteration
411
410
409
408
407
10
20
30
40
50
60
Training Epoch
Figure 13: The plot for NFE along with the training iterations of PIRATE+. The results show convergence behaviour.
21


Reviews:
[{"review_id": "k5hzZG4o6j", "reviewer": "Reviewer_kX3w", "summary": "This paper introduces a new deformable registration framework for medical imaging. Their main contribution is the inclusion of a \"plug and play\" prior into the registration framework. A novelty of the work is using a denoiser to specify priors over registration fields. They also propose an additional model (PIRATE+) that fine-tunes the CNN prior in PIRATE using deep equilibrium models (DEQ). The authors then evaluate their work on standard brain MRI datasets used in the registration literature.", "questions": "1. What advantages do you think plug and play priors offer for registration ?\n\n2. Why might denoising be appropriate as a prior for registration?\n\n3. I am not able to understand what PIRATE+ 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?\n\n4. Are there any limitations of the method compared to traditional CNN-based deformation models like VoxelMorph in terms of usability, training time, etc.?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper seems to be the first at using plug-and-play priors for registration, which could be a useful contribution in this already rich space. \n\nThe work obtains good quantitative results on standard benchmarks for deformable medical image registration.", "weaknesses": "My main issue is with the writing and presentation of this paper:\n\n- 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. \n\n- 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.\n\n- 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. \n\n- 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.\n\n- The method section would benefit by using pointers back to Fig. 1 and explaining the various steps of that figure.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VY0yL2TihC", "reviewer": "Reviewer_yfCW", "summary": "They proposed a new DIR method which is the plug-and-play approach that trains a CNN-based denoiser on the registration field. With\nthis denoiser as a regularizer within iterative methods, deep equilibrium learning is used as a fixed point iterator. The authors use a pre-trained denoiser for such regularization problems. The proposed method achieved the best performances on OASIS & CANDI datasets with reasonable qualitative results corresponding to quantitative results.", "questions": "+ What is the corner case for the proposed method?\n+ Memory consumption needs to be reported according to fixed size of datasets.", "limitations": "", "rating": 6, "confidence": 1, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "+ The proposed method used the DEQ approach for iterative registrations. The adapting approach is unique and reasonable. Their training loss looks going down well on the datasets.\n+ DEQ approach successfully addressed registration problems with PnP(plug-and-play) method with appealing gain in Tables 1 & 2.", "weaknesses": "+ I didn't understand the motivation of using pre-trained denoisor for regularizations. Why isn't the trainable denoiser used for this task?\n+ No ablation studies about the necessity of the usage of a pre-trained model\n+ Less analysis on the DEQ model : I would like to see the convergence of the number of function evaluation(NFE) along with the training iterations.\n+ In general, the result section is short of significant experiments to support their claim. For example, Figure 3 is a single instance analysis rather group analysis.\n+ Generalization needs to be validated with group analysis.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ULznLZV5sl", "reviewer": "Reviewer_3727", "summary": "The paper addresses the problem of deformable (non-rigid) image registration in the context of biomedical image analysis. In particular, it proposes two methods (PIRATE and PIRATE+) for the regularisation of the deformation field. In contrast to existing deep learning approaches to image registration, the approach presented here explicitly integrates a data-fidelity penalty as well as a CNN prior (this is pre-trained and acts as regularizer for the deformation field). The authors argue that this improves the fidelity between the registered image and the reference image. The second approach, PIRATE+, is similar to this but uses a CNN prior that is trained using deep equilibrium models.", "questions": "- The paper proposes two methods, PIRATE and PIRATE+. I am a bit unclear on what are the conclusions: Is PIRATE+ is better than PIRATE? When should PIRATE be used? When should PIRATE+ be used?\n\n- 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?\n\n- The run-time of the proposed framework is significantly higher than those of other DL methods. What are the reasons for this?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The proposed methods (PIRATE and PIRATE+) are novel methodological contributions which is positive. Additionally, the proposed framework is compared to a number of different methods, including DL and non-DL methods, which is very good. This also includes the best-performing method from recent comparative studies (Mok and Chung, CVPR 2020, MICCAI 2020). The results reported outperform this method in terms of registration accuracy. Another strength of the paper is the careful review of the state-of-the-art in the field. This is well done and comprehensive, allowing the reader to place the proposed work in the context of the SOTA.", "weaknesses": "The weaknesses are mainly related to the evaluation of the proposed framework:\n\n- 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...\n\n- The run-time of the proposed framework is significantly higher than those of other DL methods. This is a significant disadvantage for clinical applications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Scb8Xg1MY9", "reviewer": "Reviewer_UpRR", "summary": "The paper introduces a new deformable image registration (DIR) method called Plug-and-play Image Registration Network (PIRATE). PIRATE offers a new approach to DIR by integrating explicit data fidelity and a CNN prior. The paper also presents an extended version, PIRATE+, that fine-tunes the CNN prior using deep equilibrium models (DEQ). Both methods are validated on the OASIS and CANDI datasets, with results indicating state-of-the-art performance in DIR.", "questions": "- 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?\n\n- Is registration performed on the full 3D scan, or are 2D slices of roughly pre-aligned images used?\n\n- How does PIRATE and PIRATE+ compare in terms of computational efficiency and scalability?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- Deformable image registration is a challenging task in medical image analysis, and the authors' approach of using a plug-and-play method to address this challenge is commendable.\n\n- There is a comprehensive coverage of related work, which provides a solid foundation and context for the proposed method.\n\n- The introduction of learned denoisers for regularizing the registration fields and the use of DEQ to fine-tune the regularizer within PnP iterations are innovative contributions.\n\n- The extensive validation on two widely used datasets, OASIS and CANDI, highlights the robustness and general applicability of the proposed methods.\n\n- The qualitative visual results presented in the paper convincingly demonstrate the superiority of PIRATE and PIRATE+ compared to existing deep learning and iterative methods.", "weaknesses": "- 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.\n\n- The paper makes heavy use of acronyms, which affects readability. Notably, the acronym \"DU\" is mentioned without a clear definition.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)