RFR-WWANet: Weighted Window Attention-Based Recovery Feature Resolution Network for Unsupervised Image Registration

TL;DR

This work tackles the challenge of unsupervised deformable image registration in abdominal CT, where preserving fine-grained spatial detail and modeling long-range voxel relationships are difficult for existing transformer-based methods. The authors integrate a Swin-transformer–based encoder with a Recovery Feature Resolution (RFR) mechanism and a novel Weighted Window Attention (WWA) to recover high-resolution features and establish global interactions across windows. The resulting RFR-WWANet achieves state-of-the-art performance on atlas-based abdominal registration benchmarks (WORD and BTCV) and demonstrates statistical significance over strong baselines, with ablations confirming the contributions of both the recovery branch and WWA. The approach also generalizes reasonably to brain MRI (OASIS), suggesting robustness across modalities and an impact on improving precise anatomical correspondences for downstream clinical analysis.

Abstract

The Swin transformer has recently attracted attention in medical image analysis due to its computational efficiency and long-range modeling capability. Owing to these properties, the Swin Transformer is suitable for establishing more distant relationships between corresponding voxels in different positions in complex abdominal image registration tasks. However, the registration models based on transformers combine multiple voxels into a single semantic token. This merging process limits the transformers to model and generate coarse-grained spatial information. To address this issue, we propose Recovery Feature Resolution Network (RFRNet), which allows the transformer to contribute fine-grained spatial information and rich semantic correspondences to higher resolution levels. Furthermore, shifted window partitioning operations are inflexible, indicating that they cannot perceive the semantic information over uncertain distances and automatically bridge the global connections between windows. Therefore, we present a Weighted Window Attention (WWA) to build global interactions between windows automatically. It is implemented after the regular and cyclic shift window partitioning operations within the Swin transformer block. The proposed unsupervised deformable image registration model, named RFR-WWANet, detects the long-range correlations, and facilitates meaningful semantic relevance of anatomical structures. Qualitative and quantitative results show that RFR-WWANet achieves significant improvements over the current state-of-the-art methods. Ablation experiments demonstrate the effectiveness of the RFRNet and WWA designs. Our code is available at \url{https://github.com/MingR-Ma/RFR-WWANet}.

Figures (8)

• Figure 1: Overview of the RFR-WWANet. The parameter on each block indicates the number of channels output by that block. The input yellow and green cuboid images represent the moving and fixed image, and $\phi$ is the deformation field output from RFR-WWANet.
• Figure 2: The details of SCPE and the details of two successive Swin transformer blocks with the proposed WWA blocks. The SCPE is utilized at the full resolution stage and extracts the feature representations to the 1/4 resolution stage. WWA blocks are exploited after the regular window partitioning operation and shifted window partitioning operation of two successive Swin transformer blocks.
• Figure 3: Overview of the proposed WWA block. $N$, $K$, and $C$ represent the number of windows, the number of elements in each window, and the number of channels in each window, respectively.
• Figure 4: Two example slices of the window sequence input into and output from WWA block. Each image denotes the window with 96 channels (horizontal axis) and 48 elements (vertical axis) at the $1/4$ resolution stage. Column (a) represents a window before inputting into the WWA block, and column (b) represents the output from the WWA block.
• Figure 5: Images showing an example of a registration image pair. Four organs in each slice are the stomach, left and right kidneys, and liver. They are marked in light blue, dark blue, yellow, and red, respectively. In these two parts, the first row shows the input image pair, and the second row shows the warped image by different methods. Warped grids are used to observe deformations roughly. Each channel in the RGB image corresponds to a direction in the deformation field, where each pixel represents the displacement of the voxel at that position in three directions..