Swin-PSAxialNet: An Efficient Multi-Organ Segmentation Technique

Contemporary clinical intervention, including the diagnosis of diseases, the formulation of treatment plans, and the tracking of treatment outcomes, relies on the accurate segmentation of medical images¹. However, the complex structural relationships among abdominal organs²make it a challenging task to achieve accurate segmentation of multiple abdominal organs³. Over the past few decades, the flourishing developments in medical imaging and computer vision have presented both new opportunities and challenges in the field of abdominal multi-organ segmentation. Advanced Magnetic Resonance Imaging (MRI)⁴ and Computed Tomography (CT) technology⁵ enable us to acquire high-resolution abdominal images. The precise segmentation of multiple organs from CT images holds significant clinical value for the assessment and treatment of vital organs such as the liver, kidneys, spleen, pancreas, etc.^6,7,8,9,10 However, manual annotation of these anatomical structures, especially those requiring intervention from radiologists or radiation oncologists, is both time-consuming and susceptible to subjective influences¹¹. Therefore, there is an urgent need to develop automated and accurate methods for abdominal multi-organ segmentation.

Previous research on image segmentation predominantly relied on Convolutional Neural Networks (CNNs), which improve segmentation efficiency by stacking layers and introducing ResNet¹². In 2020, the Google research team introduced the Vision Transformer (VIT) model¹³, marking a pioneering instance of incorporating Transformer architecture into the traditional visual domain for a range of visual tasks¹⁴. While convolutional operations can only contemplate local feature information, the attention mechanism in Transformers enables the comprehensive consideration of global feature information.

Considering the superiority of Transformer-based architectures over traditional convolutional networks¹⁵, numerous research teams have undertaken extensive exploration in optimizing the synergy between the strengths of Transformers and convolutional networks^16,17,18,19. Chen et al. introduced the TransUNet for medical image segmentation tasks¹⁶, which leverages Transformers to extract global features from images. Due to the high cost of network training and the failure to utilize the concept of feature extraction hierarchy, the advantages of Transformer have not been fully realized.

To address these issues, many researchers have started experimenting with incorporating Transformers as the backbone for training segmentation networks. Liu et al.¹⁷ introduced the Swin Transformer, which employed a hierarchical construction method for layered feature extraction. The concept of Windows Multi-Head Self-Attention (W-MSA) was proposed, significantly reducing computational cost, particularly in the presence of larger shallow-level feature maps. While this approach reduced computational requirements, it also isolated information transmission between different windows. To address this problem, the authors further introduced the concept of Shifted Windows Multi-Head Self-Attention (SW-MSA), enabling information propagation among adjacent windows. Building upon this methodology, Cao et al. formulated the Swin-UNet¹⁸, replacing the 2D convolutions in U-Net with Swin modules and incorporating W-MSA and SW-MSA into the encoding and decoding processes, achieving commendable segmentation outcomes.

Conversely, Zhou et al. highlighted that the advantage of conv operation could not be ignored when processing high-resolution images¹⁹. Their proposed nnFormer employs a self-attention computation method based on local three-dimensional image blocks, constituting a Transformer model characterized by a cross-shaped structure. The utilization of attention based on local three-dimensional blocks significantly reduced the training load on the network.

Given the problems with the above study, an efficient hybrid hierarchical structure for 3D medical image segmentation, termed Swin-PSAxialNet, is proposed. This method incorporates a downsampling block, Space-to-depth (SPD)²⁰ block, capable of extracting global information²¹. Additionally, it adds a parameter shared axial attention (PSAA) module, which reduces the learning parameter count from quadratic to linear and will have a good effect on the accuracy of network training and the complexity of training models²².

Swin-PSAxialNet network
The overall architecture of the network adopts the U-shaped structure of nnU-Net²³, consisting of encoder and decoder structures. These structures engage in local feature extraction and the concatenation of features from large and small-scale images, as illustrated in Figure 1.

Figure 1: Swin-PSAxialNet schematic diagram of network architecture. Please click here to view a larger version of this figure.

In the encoder structure, the traditional Conv block is combined with the SPD block²⁰ to form a downsampling volume. The first layer of the encoder incorporates Patch Embedding, a module that partitions the 3D data into 3D patches , (P₁, P₂, P₃) represents non-overlapping patches in this context,  signifies the sequence length of 3D patches. Following the embedding layer, the next step involves a non-overlapping convolutional downsampling unit comprising both a convolutional block and an SPD block. In this setup, the convolutional block has a stride set to 1, and the SPD block is employed for image scaling, leading to a fourfold reduction in resolution and a twofold increase in channels.

In the decoder structure, each upsample block after the Bottleneck Feature layer consists of a combination of an upsampling block and a PSAA block. The resolution of the feature map is increased twofold, and the channel count is halved between each pair of decoder stages. To restore spatial information and enhance feature representation, feature fusion between large and small-scale images is performed between the upsampling blocks. Ultimately, the upsampling results are fed into the Head layer to restore the original image size, with an output size of (H × W × D × C, C = 3).

SPD block architecture
In traditional methods, the downsampling section employs a single stride with a step size of 2. This involves convolutional pooling at local positions in the image, limiting the receptive field, and confining the model to the extraction of features from small image patches. This method utilizes the SPD block, which finely divides the original image into three dimensions. The original 3D image is evenly segmented along the x, y, and z axes, resulting in four sub-volume bodies. (Figure 2) Subsequently, the four volumes are concatenated through "cat" operation, and the resulting image undergoes a 1 × 1 × 1 convolution to obtain the down-sampled image²⁰.

Figure 2: SPD block diagram. Please click here to view a larger version of this figure.

PSAA block architecture
In contrast to traditional CNN networks, the proposed PSAA block is more effective in conducting global information focus and more efficient in network learning and training. This enables the capture of richer images and spatial features. The PSAA block includes axial attention learning based on params sharing in three dimensions: Height, Width, and Depth. In comparison to the conventional attention mechanism that performs attention learning for each pixel in the image, this method independently conducts attention learning for each of the three dimensions, reducing the complexity of self-attention from quadratic to linear. Moreover, a learnable keys-queries parameter sharing mechanism is employed, enabling the network to perform attention mechanism operations in parallel across the three dimensions, resulting in faster, superior, and more effective feature representation.