SAMConvex: Fast Discrete Optimization for CT Registration using Self-supervised Anatomical Embedding and Correlation Pyramid

Given a source image $I_{s}:\Omega_{s}\to{\mathbb{R}}$ and a target image $I_{t}:\Omega_{t}\to{\mathbb{R}}$ within a spatial domain $\Omega\subset{\mathbb{R}^{n}}$, our goal is to find the spatial transformation $\varphi^{-1}:\Omega_{t}\to\Omega_{s}$. We aim at minimizing the following energy:

where the transformation model is the displacement vector field (DVF) $\varphi^{-1}=Id+u$ with $Id$ being the identity transformation, the $E_{\mathcal{D}}$ is the similarity term measuring the similarity between $I_{t}$ and warped image $I_{s}\circ\varphi^{-1}$, $E_{\mathcal{R}}=||\nabla{u}||_{2}^{2}$ is the diffusion regularizer that advocates the regularity of $\varphi^{-1}$. $\lambda$ weights between the data matching term and the regularization term.

We conduct the optimization scheme proposed in  [4, 22, 29, 11] and we give a brief review here. By introducing an extra term into Eq. 1

the optimization then can be decomposed into two sub-optimization problems¹¹1For better illustration, we substitute ${\lambda}E_{\mathcal{R}}(u)$ and $E_{\mathcal{D}}(I_{s}\circ\varphi,I_{t})$ with $const$ in the two equations in Eq. 3, respectively.

with each can be solved via global optimizations. By alternatively optimizing the two subproblems and progressively reducing $\theta$ during the alternative optimization, we get $\hat{v}\approx\hat{u}$ and obtain the solution of Eq. 1 as $\varphi^{-1}=Id+\hat{v}$. To be noted, the first problem in Eq. 3 can be solved point-wise because the spatial derivatives of $v$ is not involved. We can search over the cost volume of each point to obtain the optimal solution of the first problem in Eq. 3. For the second problem, we are inspired by mean-field inference [14] and process $u$ using average pooling.

Presumably, handling complex registration tasks rely on: (1) distinct features that are robust to inter-subject variation, organ deformation, contrast injection, and pathological changes, and (2) global/contextual information that can benefit the registration accuracy in complex deformation. To achieve these goals, we adopt self-supervised anatomical embedding (SAM) [27] based feature descriptor that encodes both global and local embeddings. The global embeddings memorize the 3D contextual information of body parts on a coarse level, while the local embeddings differentiate adjacent structures with similar appearances, as shown on the left of Fig. 1. The former helps the latter to focus on small regions to distinguish with fine-level features.

To be specific, given an image $I\in{\mathbb{R}^{H\times W\times D}}$, we utilize a pre-trained SAM model that outputs a global embedding $f_{global}\in{\mathbb{R}^{128\times\frac{H}{8}\times\frac{W}{8}\times\frac{D}{2}}}$ and a local embedding $f_{local}\in{\mathbb{R}^{128\times\frac{H}{2}\times\frac{W}{2}\times\frac{D}{2}}}$. We resample the global embedding with bilinear interpolation to match the shape of $f_{global}$ to $f_{local}$ and concatenate them to get $f_{SAM}\in{\mathbb{R}^{256\times\frac{H}{2}\times\frac{W}{2}\times\frac{D}{2}}}$. We normalize $f_{global}$ and $f_{local}$ before concatenation and adopt the dot product $<\cdot,\cdot>$ as the similarity measure in the following part, with a higher value indicating better alignment.

With $f_{SAM}$, we construct $E_{\mathcal{D}}(\cdot,\cdot)$ as the cost volume in the SAM feature space within a neighborhood of $[-N,N]$, where $N$ represents searching radius. Given two images $I^{0}$ and $I^{1}$, the resulting $E_{\mathcal{D}}(\cdot,\cdot)$ can then be formulated as

where $f_{SAM}^{0}$ and $f_{SAM}^{1}$ are the SAM embedding of $I^{0}$ and $I^{1}$, respectively. $x$ is any voxel in the image domain and $d$ is the voxel displacement within the neighborhood. $E_{\mathcal{D}}(\cdot,\cdot)$ has the shape of $(2N+1)\times(2N+1)\times(2N+1)\times\frac{H}{2}\times\frac{W}{2}\times\frac{D}{2}$.

Since $E_{\mathcal{D}}(\cdot,\cdot)$ is computed within the neighborhood, the magnitude of deformation is bounded by the size of the neighborhood. Our approach to addressing this issue is based on a surprisingly simple but effective observation: performing registration based on a coarse-to-fine scheme with a small search range at each level instead of a large search range at one resolution benefits to 1) significantly improve the efficiency such as low computation burden and fast running time; 2) enlarge receptive field and improve registration accuracy.

We first build an image pyramid of $I_{s}$ and $I_{t}$. At each resolution, we warp the source image with the composed deformation computed from all the previous levels (starting from the coarsest resolution), transform the warped image and target image to the SAM space, and conduct the decoupling to convex optimizations strategy to obtain the deformation between the warped image and target at the current resolution. With such a coarse-to-fine strategy, we estimate a sequence of displacement fields from a starting identity transformation. The final field is computed via the composition of all the displacement fields [30].

To be noted, the cost volume at each level is computed by taking the inputs of $\{1,\frac{1}{2},\frac{1}{4}\}$ resolution. Adding one coarser level consumes less computation than doubling the size of the neighborhood at the fine resolution but yields the same search range.

We conduct the registration in 3 levels of resolutions. At each level, we solve Eq. 3 via five alternative optimizations with $\frac{1}{2\theta}=\{0.003,0.01,0.03,0.1,0.3,1\}$. To further improve the registration performance, we append a SAM-based instance optimization after the coarse-to-fine registration with a learning rate of $0.05$ for $50$ iterations. It solves an instance-specific optimization problem [3] in the SAM feature space. The optimization objective consists of similarity (dot product between SAM feature vectors on the highest resolution) and diffusion regularization terms. All the experiments are run on the CPU of Intel Xeon Platinum 8163 with 16 cores of 2.50GHz and the GPU of NVIDIA Tesla V100. Code will be available at https://github.com/Alison-brie/SAMConvex.

We split Abdomen and HeadNeck into a training set and test set to accommodate the requirement of the training dataset of comparing learning-based methods. Lung dataset contains 35 CT pairs, which is not sufficient for developing learning-based methods. Hence, it is only used as a testing set for optimization-based methods. All the methods are evaluated on the test set. The SAM is pre-trained on NIH Lymph Nodes dataset [27]. All 3 datasets in this paper are not used for pre-training.

We compare with five widely-used and top-performing deformable registration methods, including NiftyReg [24], Deeds [9], top-ranked ConvexAdam [20] in Learn2Reg Challenge [12], and SOTA learning-based methods LapIRN [18] and SAME [15]. All results are based on the SAM-affine [15] pre-alignment.

As shown in Table. 1, SAMConvex outperforms the widely-used NiftyReg [24] across the three datasets with an average of $10.0\%$ Dice score improvement. Compared with the best traditional optimization-based method Deeds [9], SAMConvex performs better or comparatively on the three datasets with approximately $20\sim 100$ times faster runtime. We also see a better performance of SAMConvex over ConvexAdam [20]. To be noted, the performance gap between SAMConvex and ConvexAdam is greater on Abdomen CT than the other two datasets. This may be because the deformation is more complex in Abdomen and the anatomical differences between inter-subjects make the registration more challenging. With a coarse-to-fine strategy, SAMConvex can better bear these issues. Although LapIRN [18] has the fastest inference time, it has slightly inferior registration accuracy as compared to SAMConvex. Moreover, it needs training when being applied to a new dataset. Equipped with SAM, SAME achieves the overall 2nd best performance in two inter-patient registration tasks but with a notably higher folding rate overall. Moreover, it also requires individual training for a new dataset and struggles with the intra-patient lung registration task (small dataset and no available training data). In summary, our SAMConvex achieves the overall best performance as compared to other methods with a fast inference time and comparable deformation field smoothness.

To better understand the performance of different deformable registration methods, we display organ-specific results of the inter-patient registration task of abdomen CT in Fig. 2 where large deformations and complex anatomical differences exist. As shown, SAMConvex is consistently better than the second-best and third-best registration methods on most examined abdominal organs, in 11 out of 13 organs.