Correlation-aware Coarse-to-fine MLPs for Deformable Medical Image Registration (2406.00123v2)

Abstract: Deformable image registration is a fundamental step for medical image analysis. Recently, transformers have been used for registration and outperformed Convolutional Neural Networks (CNNs). Transformers can capture long-range dependence among image features, which have been shown beneficial for registration. However, due to the high computation/memory loads of self-attention, transformers are typically used at downsampled feature resolutions and cannot capture fine-grained long-range dependence at the full image resolution. This limits deformable registration as it necessitates precise dense correspondence between each image pixel. Multi-layer Perceptrons (MLPs) without self-attention are efficient in computation/memory usage, enabling the feasibility of capturing fine-grained long-range dependence at full resolution. Nevertheless, MLPs have not been extensively explored for image registration and are lacking the consideration of inductive bias crucial for medical registration tasks. In this study, we propose the first correlation-aware MLP-based registration network (CorrMLP) for deformable medical image registration. Our CorrMLP introduces a correlation-aware multi-window MLP block in a novel coarse-to-fine registration architecture, which captures fine-grained multi-range dependence to perform correlation-aware coarse-to-fine registration. Extensive experiments with seven public medical datasets show that our CorrMLP outperforms state-of-the-art deformable registration methods.

• The paper presents CorrMLP, a novel correlation-aware coarse-to-fine MLP framework that enhances deformable medical image registration precision.
• It introduces the Correlation-aware Multi-window MLP block to capture both local and multi-range dependencies, outperforming conventional CNN and transformer methods.
• Extensive tests on brain and cardiac datasets demonstrate improved Dice scores and transformation smoothness, paving the way for real-time clinical applications.

Summary of Correlation-Aware Coarse-to-Fine MLPs for Deformable Medical Image Registration

The paper addresses the critical challenge of deformable medical image registration, proposing a novel solution that leverages multi-layer perceptrons (MLPs) in a coarse-to-fine registration framework. Historically, convolutional neural networks (CNNs) and transformers have dominated the medical image registration landscape. While transformers have demonstrated superiority over CNNs by capturing long-range dependencies, they are constrained by significant computational costs and memory demands, particularly when applied to full-resolution images. This limitation stymies their ability to register images with precise pixel-level accuracy, which is crucial in medical contexts where subtle anatomical details matter.

Introducing the Correlation-aware MLP (CorrMLP), the paper pioneers the use of MLPs for deformable registration tasks, combining computational efficiency with the ability to capture fine-grained dependencies at full resolution. The central innovation in the CorrMLP is the Correlation-aware Multi-window MLP (CMW-MLP), which integrates correlation awareness into the registration process. By employing CMW-MLP blocks in a hierarchical coarse-to-fine architecture, the method enhances registration accuracy by calculating local correlations and capturing multi-range dependencies through multiple window-based MLP branches.

Technical Contributions and Results

1. Novel Architecture: The authors designed a correlation-aware coarse-to-fine registration architecture. This framework utilizes hierarchical feature extraction and processes feature maps through CMW-MLP blocks to achieve spatial alignment between moving and fixed images. The architecture is meticulous in ensuring that both image-level and step-level correlations are considered, enriching contextual information at each registration step.
2. Correlation-aware Multi-window MLP: This block calculates local correlations and uses a set of MLP branches with various window sizes to capture intricate dependencies, enhancing the model's capability to address both large and subtle deformations. Importantly, this block is tailored for deformable registration, a domain where MLPs were previously under-utilized.
3. Extensive Validation: CorrMLP has been rigorously tested on several well-known datasets, demonstrating superior performance in both brain and cardiac image registration tasks. In comparative studies, CorrMLP outperformed existing state-of-the-art methods, including traditional CNN and transformer-based approaches, in terms of Dice Similarity Coefficient (DSC) without sacrificing transformation smoothness.

Implications and Future Work

The paper's contributions underscore the potential of MLPs as viable replacements for computationally intensive transformers in precise medical image registration tasks. It presents a compelling case for incorporating correlation-aware methodologies in coarse-to-fine registration pipelines, potentially setting a new standard for the field.

The proposed CorrMLP framework could have profound implications for clinical applications where precise image registration is crucial, such as in monitoring tumor progression or planning surgical interventions. The efficiency gains also make real-time applications feasible, opening new avenues for on-the-fly analysis and decision-making in medical settings.

Future avenues for research could involve refining the CMW-MLP to handle multi-modal image registration tasks or extending the methodology to other domains of medical imaging with varying modalities, such as PET-CT or MRI-PET fusion. Moreover, adaptations to include unsupervised learning paradigms could further bolster the model's application breadth and relevance.

In conclusion, the paper offers a robust foundation for future explorations and advances in the domain of medical image registration, providing a novel perspective on leveraging MLPs in addressing complex, real-world challenges in medical imaging.