Tetrahedron-Net for Medical Image Registration

TL;DR

Medical image registration remains challenging due to large deformations and complex anatomy. Tetrahedron‑Net introduces a two‑level decoder that cooperates with a shared encoder to refine deformation fields in a coarse‑to‑fine manner, and demonstrates consistent gains across multiple U‑Net backbones and benchmarks. The approach yields higher Dice scores and better diffeomorphic regularity, indicating strong practical impact for rapid, accurate MIR, with broad applicability to existing architectures. The results suggest that the proposed two‑level decoding strategy offers a simple yet effective plug‑in to boost dense prediction tasks in medical imaging and beyond.

Abstract

Medical image registration plays a vital role in medical image processing. Extracting expressive representations for medical images is crucial for improving the registration quality. One common practice for this end is constructing a convolutional backbone to enable interactions with skip connections among feature extraction layers. The de facto structure, U-Net-like networks, has attempted to design skip connections such as nested or full-scale ones to connect one single encoder and one single decoder to improve its representation capacity. Despite being effective, it still does not fully explore interactions with a single encoder and decoder architectures. In this paper, we embrace this observation and introduce a simple yet effective alternative strategy to enhance the representations for registrations by appending one additional decoder. The new decoder is designed to interact with both the original encoder and decoder. In this way, it not only reuses feature presentation from corresponding layers in the encoder but also interacts with the original decoder to corporately give more accurate registration results. The new architecture is concise yet generalized, with only one encoder and two decoders forming a ``Tetrahedron'' structure, thereby dubbed Tetrahedron-Net. Three instantiations of Tetrahedron-Net are further constructed regarding the different structures of the appended decoder. Our extensive experiments prove that superior performance can be obtained on several representative benchmarks of medical image registration. Finally, such a ``Tetrahedron'' design can also be easily integrated into popular U-Net-like architectures including VoxelMorph, ViT-V-Net, and TransMorph, leading to consistent performance gains.

Figures (9)

• Figure 1: Visualization of unregistered image pairs and the overlapping images. Overlapping the two images with different colors suggests that there are significant mismatches in the image pairs.
• Figure 2: The architecture of the proposed Tetrahedron-Net registration network. The registration network used in the figure is the U-UNet network. Firstly, the fixed image $f$ and the moving image $m$ are concated in the channel dimension as the input. After the encoder extracts the features and then two decoders generate the deformation field $\phi$. Then the spatial transformation network (STN) uses the generated deformation field $\phi$ to deform the moving image $m$ to obtain the Deformed image ($m\circ \phi$), the loss of smoothness (loss_smooth) is calculated for the generated deformation field, and the loss of similarity (loss_sim) is calculated for the generated deformed and fixed images. The structure of the encoder(Enc) and Decoder1(Dec1) as same as UNet, and the Decoder2(Dec2) with the same structure as Dec1. The circles in the figure represent the concat operation and the squares represent two consecutive $3 \times 3 \times 3$ Convolution and ReLU layers.
• Figure 3: U-UNet++ is adopted as the second-level decoder. Each node is composed of convolution and ReLU layers.
• Figure 4: The second-level decoder is designed by U-UNet3+.
• Figure 5: The second-level decoder is designed by U-DenseUNet.