## SegStitch: Multidimensional Transformer for Robust and Efficient Medical Imaging Segmentation

Shengbo Tan, Zeyu Zhang, Ying Cai, Daji Ergu, Lin Wu, Binbin Hu, Pengzhang Yu, Yang Zhao

Abstract -Medical imaging segmentation plays a significant role in the automatic recognition and analysis of lesions. State-ofthe-art methods, particularly those utilizing transformers, have been prominently adopted in 3D semantic segmentation due to their superior performance in scalability and generalizability. However, plain vision transformers encounter challenges due to their neglect of local features and their high computational complexity. To address these challenges, we introduce three key contributions: Firstly, we proposed SegStitch, an innovative architecture that integrates transformers with denoising ODE blocks. Instead of taking whole 3D volumes as inputs, we adapt axial patches and customize patch-wise queries to ensure semantic consistency. Additionally, we conducted extensive experiments on the BTCV and ACDC datasets, achieving improvements up to 11.48% and 6.71% respectively in mDSC, compared to state-ofthe-art methods. Lastly, our proposed method demonstrates outstanding efficiency, reducing the number of parameters by 36.7% and the number of FLOPS by 10.7% compared to UNETR. This advancement holds promising potential for adapting our method to real-world clinical practice. The code will be available at https://github.com/goblin327/SegStitch

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00496v1-figure-9deb9eb3-71f2-447f-b985-695f902e3938.jpg)

Index Terms -Medical Imaging Segmentation, Semantic Segmentation, Transformers, Ordinary Differential Neural Networks

## I. INTRODUCTION

3D segmentation stands as a pivotal challenge within the medical imaging domain, with broad implications for clinical practices. Notably, delineating the intricate internal structures of the human body, including organs, within medical imagery is essential for key clinical applications like Computer-Aided Diagnosis (CAD), Computer-Assisted Surgery (CAS), and Radiotherapy (RT) [1].

In recent years, the mainstream approaches to medical image segmentation have predominantly relied on Convolutional Neural Networks (CNN) [2], such as U-Net [3] and its related architectures [4], [5], [6], [7], [8]. However, these methods often struggle with convolutional localization, limiting their ability to effectively capture contextual information[9].To address this limitation, there is an urgent need for efficient medical segmentation networks. However, ViT [10] has demonstrated state-of-the-art performance in various computer vision [11], [12] and medical image analysis tasks [13]. Firstly, ViT

Ying Cai is the corresponding author: caiying34@yeah.net.

Shengbo Tan, Ying Cai, Daji Ergu, Binbin Hu, and Pengzhang Yu are with Key Laboratory of Electronic Information Engineering, College of Electronic and Information, Southwest Minzu University, Chengdu 610225, China.

Zeyu Zhang is with College of Engineering, Computing and Cybernetics, the Australian National University, Canberra ACT 2601, Australia.

Lin Wu is with Hunan University of Technology, Zhuzhou 410072, China. Yang Zhao is with Department of Computer Science and Information Technology, La Trobe University, Bundoora VIC 3086, Australia.

Fig. 1. Relationship between model parameter count, computational complexity, and Dice similarity coefficient. The size of the spheres indicates the model parameter count. Compared to other models, our SegStitch achieves the highest mDSC while maintaining smaller model size and lower computational complexity.

[14] leverages its outstanding capability in capturing longrange interactions, demonstrating excellent generalization and robust features [15]. Secondly, this approach exhibits high parallelism, facilitating the training and inference processes of large-scale models [16].

In terms of algorithm improvements, a focal point of attention has been the enhancement of core building blocks, particularly focusing on attention mechanisms. However, such performance comes at a cost, due to attention requiring the computation of feature information across all spatial positions, it incurs high computational complexity and significant memory overhead. Nevertheless, queries in different semantic regions actually involve entirely different key-value pairs, so forcing all queries to attend to the same token set may be suboptimal. Therefore, we extended the self-attention module by dividing the attention mechanism into fine-grained and coarsegrained modules, and enabling them to work in tandem. This approach further enhances the model's ability to understand and represent complex features.

According to research findings, in neural ODE-based networks (ODENets), even for data subjected to slight perturbations, the difference between the corresponding perturbed version of the integral curve and the original curve is not significant [17]. Inspired by Neural Memory Ordinary Differential Equations (nmODE) [18], we have integrated the SegStitch module with nmODE to facilitate the dynamic and continuous updating of feature maps, leading to further

optimized segmentation results.

To further assess our proposed method, We validate our SegStitch method using two datasets: Synapse [19] and ACDC [20]. Both qualitative and quantitative results demonstrate the effectiveness and superior segmentation capability of SegStitch in terms of segmentation accuracy and model efficiency. On the Synapse dataset, SegStitch achieves high-quality segmentation of long-range targets (see Figure 7), with an absolute gain of 0.97% in Dice Score (as shown in Table V), while significantly reducing model complexity. Compared to the baseline UNETR, SegStitch reduces parameters by 36.7% and FLOPs by 10.7% [21],As shown in Figure 1.

Our contributions can be summarized as follows:

- (1) First, we propose SegStitch , an innovative architecture that combines transformers with denoising ODE blocks. In this transformer, we customized shared queries between different patches to enhance contextual relationships using shared query keys, ensuring semantic consistency. This achieves highquality segmentation of long-range targets through local-toglobal self-attention.
- (2) Additionally, we conducted extensive experiments on the BTCV and ACDC datasets. Compared to state-of-the-art methods, we achieved improvements of 11.48% and 6.71% in the mDSC metric, respectively.
- (3) Finally, our proposed method demonstrates outstanding efficiency. Compared to UNETR, it achieves a 36.7% reduction in the number of parameters and a 10.7% reduction in FLOPS. These advancements suggest promising prospects for applying our method in practical clinical settings.

## II. RELATED WORK

## A. 3D Medical Image Segmentation

Three-dimensional medical image segmentation is a crucial technique that distinguishes organs and pathological areas from medical images on a pixel-by-pixel basis. The development of this technology has greatly benefited from the remarkable achievements of Convolutional Neural Networks (CNNs) [2]. Typical work, such as 3D U-Net and nnU-Net [22], has been widely adopted for medical imaging segmentation tasks [23], [24], [25]. Researchers have been actively exploring the integration of Transformers alongside differential neural networks to overcome the limitations of CNNs and enhance segmentation accuracy. In recent years, significant efforts have been made to integrate these technologies into medical image segmentation, aiming to maximize accuracy. In the following sections, we will introduce and analyze the advancements made in the domain of three-dimensional medical image segmentation through the combination of CNNs, Transformers, and differential neural networks.

To address the limitation of local receptive fields in convolutional neural networks, researchers have explored the use of dilated convolutions and large kernel convolutions as alternatives to standard convolutions [26]. Inspired by the significant success of Transformers in natural language processing, methods like TransUNET [27] and similar approaches [28] [29] [30] combine Transformers with convolutions as building blocks for encoders. Transclaw Unet [31] integrates convolutional operations with Transformers on the encoder side, achieving detailed segmentation and long-range relationship learning. Meanwhile, Azad et al. [32] introduced TransDeepLab, which enhances DeepLab with different window strategies for skin lesion segmentation. They also introduced TransCeption [33], which captures multi-scale representations in a single stage through an improved patch merging module. Additionally, Huang et al. [34] enhanced segmentation performance by integrating global information across different scales within a U-shaped architecture.Wu et al. [35] achieved more efficient feature extraction and representation through patch embedding processing. Ali et al. [36] introduced an encoder based on 3D Swin Transformers and a decoder that combines CNN and Transformer-based techniques, effectively reconstructing high-quality segmentation images. Yun et al. [37] proposed an enhanced Transformer module to connect the encoder and decoder, effectively reducing information loss between hierarchical features. Li et al. [38] proposed an innovative upsampling method that simulates long-range dependencies and integrates global information, capturing and preserving spatial edge details more accurately. Shaker et al. [39] introduced the Efficient Pairwise Attention (EPA) module, which effectively learns features through spatial and channel-based attention.

Currently, methods based on Vision Transformers (ViT) have achieved significant success in medical image segmentation, primarily due to their strong representation capabilities. However, Transformers face challenges in handling long-range dependencies, as the attention mechanism requires calculating attention between all positions, leading to increased computational complexity with longer sequences. Furthermore, the sparsity of medical image data makes models prone to overfitting during training. An underexplored area is the application of multimodal attention mechanisms to reveal longrange edge information in three-dimensional medical images. Due to the inherent complexity of medical images, extracting long-range edge information remains a relatively unexplored and challenging topic.

## B. Ordinary Differential Neural Network

In recent years, only a few segmentation methods have combined neural ODEs with network design [40] [41] [42]. In a study [43], the structure of U-Net was enhanced by introducing elements of neural ODEs. In this design, the repetitive residual blocks in each branch of the U-Net were replaced by neural ODEs, which operated only around a convolutional block. Despite the widespread recognition of U-Net, recent research indicates that new network designs can achieve superior results compared to U-Net. A recent study [44] proposed an innovative approach that combines neural ODEs with the level set method. This method utilizes parameterized derivative contours as neural ODEs, implicitly learning the governing functions describing contour evolution. Additionally, several studies [45],[46],[47] have applied neural ODEs to defense against adversarial attacks, achieving notably significant results.

Building upon existing research, our work incorporates the use of neural memory ordinary differential equation blocks

Fig. 2. SegStitch Network Encoder Structure Diagram. The entire network comprises four layers, each containing a downsampling module and stages. The downsampling module consists of convolution and Group Normalization, while the stages consist of three Transformer Blocks. Each Transformer Block computes different patches, together forming the complete downsampling module.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00496v1-figure-9b4cc9cc-0778-48fe-b7c9-0a830540b61c.jpg)

Fig. 3. Overall Architecture. The SegStitch method utilizes a hierarchical encoder-decoder structure. The output of the downsampling modules is passed to the decoder through skip connections, and each decoder module generates the final segmentation mask using ODE blocks.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00496v1-figure-3d5b1242-e91b-413f-a55d-97764bdee64b.jpg)

into the last fully connected layer of the Vision Transformer (VIT) architecture. This innovative design has achieved stateof-the-art performance on multiple datasets, demonstrating superior stability and reliability compared to traditional Convolutional Neural Networks (CNNs). By integrating neural memory ordinary differential equation with VIT, we have achieved significant improvements in medical segmentation tasks, providing a new paradigm for the application of neural memory ordinary differential equation in network design.

## III. METHOD

## A. Overall Architecture

We propose SegStitch, a novel 3D image segmentation network. The overall architecture of SegStitch is illustrated

Fig. 4. Position-weighted token. we generate a set of learnable parameters Token wt with a size of 1×d (where d represents the total length after flattening the 3D patch). These parameters are used to learn the spatial positional information of the 3D patch.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00496v1-figure-e8b7f77d-1c84-4dd0-95c5-f122f5e375b6.jpg)

in Figure 3. Our approach enhances traditional methods by improving the positional information weighting of each pixel in the 3D image. We then introduce a dedicated 3D image feature extraction module (as shown in Figure 2) to effectively extract features from 3D data. Finally, the outputs of the selfattention mechanism are fed into a Neural Memory Ordinary Differential Equation network, aimed at improving the model's generalization capability and stability.Specifically, given a multi-modal MRI medical image input X ∈ R M C H × × × W D × (where (H, W, D) denotes the input resolution, M represents the number of modalities, and C denotes the number of channels), the encoder layers process the input starting from the first layer. The output of each encoder layer is processed through the SegStitch module, generating planar feature maps with half the dimensions of the input. Consequently, at the final layer, we obtain a global feature vector of size [1×64×156]. The decoder layers then convert the dense features back into segmentation image planes, with upsampled features concatenated with the outputs from the corresponding encoder layers to better distinguish between foreground and background pixels. This results in a more robust segmentation mask.A key aspect of our design is the introduction of the effective feature extraction block, SegStitch. In the SegStitch module, self-attention employs parallel attention modules to generate diverse features through shared queries, thus effectively learning rich feature representations. Finally, the outputs from the two attention modules are fused and then fed into the ODE block to generate high-quality 3D medical image masks.

## B. Position-Weighted Tokens

To address the loss of spatial positional information when directly flattening 3D patches before input tokenization, we generate a set of learnable parameters Token wt with a size of 1×d (where d represents the total length after flattening the 3D patch). we sum them with Token org to obtain Token final , which possesses positional information (as illustrated in Figure 4). The definition of Token final is as follows:

$$\text{Token} _ { f i n a l } = \text{ADD} ( \text{Token} _ { o r g }, \text{Token} _ { w t } ) \quad \ \ ( 1 )$$

Where Token final , ADD, Token org , and Token wt respectively represent the token with spatial positional information,

matrix summation function, original tokens, and the learned spatial information weights of the 3D patch.

## C. Network Encoder

Applying the Transformer architecture to medical imaging tasks involves a critical consideration: converting the entire image into a sequence for the self-attention mechanism. This approach can lead to issues such as excessive computational costs, significant storage requirements, and the inability to fully retrieve spatial information. To address these drawbacks, the self-attention module in the SegStitch method utilizes shared query techniques. This not only demonstrates enhanced feature learning capabilities but also effectively addresses the challenges of long-distance edge segmentation through the attention mechanism, achieving outstanding performance across multiple datasets.

1) Shared Queries: When constructing the 3D feature extraction module SegStitch, we adopted a novel self-attention mechanism that introduces both fine-grained and coarsegrained attention modules through shared queries. The entire image is divided into non-overlapping patches for parallel processing, reducing the computational complexity of the selfattention mechanism from quadratic to linear levels, thus achieving more efficient computation. The fine-grained attention module effectively focuses on small target features within the input sequence, while the coarse-grained attention module effectively establishes dependencies between different positions in the input sequence. As shown in Figure 5, The computations for the two attention modules are as follows:

$$\text{Attn} _ { f } = \text{FA} ( Q _ { c o m m o n }, K _ { f i n e }, V _ { f i n e } ) \quad \ \ ( 2 ) \ \text{chann}$$

$$K _ { R }, V _ { R C } = K _ { f i n e } \cdot \text{RP} ( M ), V _ { c o a r s e } \cdot \text{RP} ( M ) \quad ( 3 ) \quad \mathfrak { w h e }$$

$$\ A t t { n } _ { c } = C A ( Q _ { c o m m o n }, K _ { R }, V _ { R C } ) \quad \quad ( 4 ) \quad r e p r e s$$

Where Attn f and Attn c respectively represent the finegrained and coarse-grained attention maps. In this system or module, FA denotes the fine-grained attention module, while CA denotes the coarse-grained attention module, RP represents the random learnable parameter module, and M is a twodimensional matrix initialized to all zeros. Q common, K fine , K , R V fine , and V RC are matrices used for shared queries, finegrained keys, coarse-grained keys, fine-grained value layers, and coarse-grained value layers respectively.

Fine-grained Attention: Given a normalized tensor X of shape HWD × C, we use three linear layers to compute projections for Q common , K fine , and V fine , obtaining Q common = W Q X , K fine = W K X , and V fine = W V X , where X is the input feature map, and W Q , W , and W K V are the projection weights for Q common , K fine , and V fine respectively, with sizes H×W×D×C. The fine-grained attention is defined as follows:

$$\text{Attn} _ { f } = \text{Softmax} \left ( \frac { Q _ { \text{common} } K _ { f i n e } ^ { T } } { \sqrt { d } } \right ) \cdot \tilde { V } _ { \text{fine} } \quad ( 5 ) \quad \text{Our}$$

TABLE I DETAILS OF THE MEDICAL IMAGE DATASETS USED FOR EVALUATION.

| Dataset   | Image resolution range (mm 3 )            | Format     |   Classes |
|-----------|-------------------------------------------|------------|-----------|
| ACDC      | ([0.7 ∼ 1.95]×[0.7 ∼ 1.95]×[5.0 ∼ 10.0])  | 16x160x160 |         3 |
| Synapse   | ([0.54 ∼ 0.54]×[0.98 ∼ 0.98]×[2.5 ∼ 5.0]) | 64x128x128 |         8 |

Where Q common , K fine , and V fine respectively represent the shared query, the projected shared key, and the projected fine-grained feature value layer, and d is the size of each vector.

Coarse-grained Attention: In this module, we inherit Q common from the fine-grained attention module and combine it with K R and V RC generated by Formula 3 to form a new attention mechanism. This approach enhances the ability to capture dependencies between different positions within the input sequence. Coarse-grained attention is defined as follows:

$$\text{Attn} _ { c } = \text{Softmax} \left ( \frac { \mathbb { Q } _ { c o m m o n } ^ { T } K _ { R } } { \sqrt { d } } \right ) \cdot \mathbb { V } _ { R C } \quad \ \ ( 6 )$$

Where V RC , Q common , and K R respectively represent the coarse-grained value layer, the shared query, and the coarsegrained key, and d is the size of each input token.

2) nmODE: Yi [18] proposed a novel type of neural memory ordinary differential equation, the nmODE, which overcomes the limitations of traditional differential equations while fully leveraging the storage capacity provided by dynamical systems. Therefore, the nmODE module can simultaneously learn and memorize dependencies between characteristics during the mapping process. For individual samples, the high nonlinearity and memory capacity of the nmODE more effectively refine and integrate feature information from different channels, thereby obtaining more global feature information.

$$\dot { y } = \text{-} y + \sin ^ { 2 } ( y ( t ) + W ^ { ( 1 ) } x + b ) \text{ \quad \ \ } ( 7 )$$

Where y ∈ R n , x ∈ R m , b ∈ R n × n , and W (1) = W (1) ij × m ∈ R n × m . From the perspective of neural networks, x represents the external input of the feature map, y ( ) t represents the state of the network at time t. W and b respectively represent the learned weights and biases.

In this study, an nmODE module is embedded within the decoder to further optimize the model's upsampling and skip connections, fully integrating feature information from different channels. This leverages the powerful non-linear expressive capabilities of nmODE, enhancing the model's ability to represent and learn in complex scenarios. This integration is illustrated in Figure 6.In summary, the final output ˆ X is expressed as:

$$X _ { o u t } = n m O D E ( \text{Attn} _ { f } + \text{Attn} _ { c } ) \text{ \quad \ \ } ( 8 )$$

where Attn f and Attn c represent the coarse-grained and fine-grained attention maps, respectively, and nmODE is the ordinary differential neural network.

## IV. EXPERIMENTS AND RESULTS

## A. Dataset

Our network underwent rigorous evaluation on two formidable medical image datasets: the Synapse Multi-organ

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00496v1-figure-a92c0f8a-546e-4b74-a8d2-3d2d55624e81.jpg)

Fig. 5. The shared query structure of the internal self-attention mechanism in SegStitch.the input feature map x is fed into the fine-grained and coarse-grained attention modules of the SegStitch. The weights of the linear layers for Q are shared between the two attention modules. Additionally, using the global learnable parameter matrix M , new parameters K R and V R F are obtained by multiplying M with K fine and V fine respectively, resulting in a new set of QKV combinations. Thereby, enabling different attention tasks for each self-attention module.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00496v1-figure-de9e6bb7-5405-4487-99a2-a7f75af660ab.jpg)

with the standard procedure.

Fig. 6. Structure diagram of nmODE. By introducing the ODE block in SegStitch, we have achieved a more stable feature representation without sacrificing recognition ability.

TABLE II ACDC AND SYNAPSE EXPERIMENTAL DATASET DISTRIBUTION.

|         | Image count   | Image count   | Image count   | Image count   |
|---------|---------------|---------------|---------------|---------------|
| dataset | Training      | Validation    | Testing       | Total         |
| ACDC    | 70            | 10            | 20            | 100           |
| Synapse | 18            | -             | 12            | 30            |

Segmentation Dataset (Synapse) [19] and the Automatic Cardiac Diagnosis Challenge (ACDC) [20]. A comprehensive summary of the datasets' characteristics is delineated in Table II and Table I. These datasets are publicly accessible, offering a robust framework for the assessment of image segmentation methodologies against ground-truth (GT) masks.

(1)The Automatic Cardiac Diagnosis Challenge (ACDC) dataset comprises 100 three-dimensional cardiac magnetic resonance imaging (MRI) cases collected from diverse patients, encompassing the cardiac organ from the base of the left ventricle to the apex. The series offers a wide range of slice thicknesses (from 5.0mm to 10.0mm) and in-plane resolutions (from 0.7x0.95mm 2 to 0.7x1.95mm ). For the experiment at 2 hand, the dataset has been partitioned into training, validation, and testing sets in a 7:1:2 ratio. To ensure fair comparison with previous research methodologies, images have been randomly cropped to a model input size of 160x160x16, in accordance

(2)In our research, we utilized the Synapse Multi-Organ Segmentation Dataset, a publicly available resource that originated from the 2015 MICCAI Multi-Atlas Abdominal Labeling Challenge. The dataset encompasses 30 abdominal computed tomography (CT) scans, demonstrating a range of slice thicknesses from 2.5mm to 5.0mm, as well as varying inplane resolutions from 0.54x0.54mm² to 0.98x0.98mm². With meticulous attention to detail, we meticulously divided the dataset into a training set comprising 18 cases and a testing set comprising 12 cases. To ensure a fair and comparable evaluation against previous research methodologies, all images were uniformly cropped to a size of 128x128x64 voxels during the inference phase, tailored to the model's input specifications.

## B. Evaluation criteria

HD95 is a metric used to assess the difference between two point sets. It focuses on the closest 95% of points between the predicted and the ground truth segmentation, excluding the most extreme 5% to reduce the impact of outliers. This method provides a reliable measure of segmentation accuracy and reflects the model's performance more realistically. The Dice Similarity Coefficient (mDSC) is used to quantify the similarity between the predicted and ground truth segmentations in medical images. By using both HD95 and mDSC metrics simultaneously, we can present a more comprehensive view of the model's performance and accuracy.

## C. Training details

Pretraining is crucial for Transformer-based models, such as Vision Transformers (ViT) [10], where the model's performance is widely acknowledged to be heavily reliant on pretraining, a notion corroborated by empirical evaluations. Established medical image segmentation techniques also commonly leverage pretrained weights to initialize their models [48],[49],[50],[51]. However, pretraining Transformer-based

TABLE III

ABLATION STUDIES ON SELF-ATTENTION AND ODE BLOCK IN SEGSTITCH SEGMENTATION.THE BEST SCORES ARE INDICATED IN BOLD .

TABLE IV

| Methods          | self-attention √   | ODE block   |   mDSC(%) |
|------------------|--------------------|-------------|-----------|
| SegStitch (Ours) | 1                  | - √         |      88.3 |
| SegStitch (Ours) | 2 -                |             |      87.8 |
| SegStitch (Ours) | 3 - √              | - √         |      87.6 |
| SegStitch (Ours) | 4                  |             |      89.8 |

SEGMENTATION PERFORMANCE OF DIFFERENT METHODS ON THE ACDC DATASET.THE BEST SCORES ARE INDICATED IN BOLD .

| Method             |   RV ↑ |   Myo |   LV ↑ |   AVG ↑ |
|--------------------|--------|-------|--------|---------|
| R50 U-Net[54]      |  87.1  | 80.63 |  94.92 |   87.55 |
| R50 Att-UNet[55]   |  87.58 | 79.2  |  93.47 |   86.75 |
| VIT[10]            |  81.46 | 70.71 |  92.18 |   81.45 |
| R50 VIT[10]        |  86.07 | 81.88 |  94.75 |   87.57 |
| TransUNet[56]      |  88.86 | 84.54 |  95.73 |   89.71 |
| Swin-UNet[57]      |  88.55 | 85.62 |  95.83 |   90    |
| LeVit-Unet-384[49] |  89.55 | 87.64 |  93.76 |   90.32 |
| nnFormer[50]       |  90.22 | 89.53 |  95.59 |   91.78 |
| MISSFormer[58]     |  86.36 | 85.75 |  91.59 |   87.9  |
| D-Former[35]       |  91.33 | 89.6  |  95.93 |   92.29 |
| UNETR[21]          |  85.29 | 86.52 |  94.02 |   86.61 |
| UNETR++[39]        |  91.89 | 90.61 |  96    |   92.83 |
| SegStitch (Ours)   |  93.05 | 90.46 |  96.46 |   93.32 |

models presents two challenges. Firstly, it is computationally intensive, often demanding substantial resources. Secondly, unlike the natural scene images in ImageNet [52], the medical imaging field lacks large-scale datasets for pretraining. Consequently, our SegStitch model opted for training from scratch, meaning we randomly initialized the model weights. Despite this, our experiments on two major public datasets, ACDC and Synapse, yielded promising performance comparable to state-of-the-art pretrained methods.

We implemented our models in PyTorch 1.8.0 and trained them on an NVIDIA GeForce RTX 3090 GPU with 24 GB memory. The batch size was 2, and training was capped at 1000 epochs. We used the SGD optimizer with a momentum of 0.99 [53].

## D. Results and comparisons with recent methods

In this section, the performance of the proposed SegStitch is compared with state-of-the-art methods on the ACDC and Synapse datasets.

1) Comparisons with state-of-the-art methods on ACDC datase: We evaluated the performance of our proposed SegStitch model on the ACDC dataset and compared it with various state-of-the-art models. The segmentation results on the ACDC dataset are presented in Table IV, from which relevant conclusions can be drawn. SegStitch achieved the best average Dice Similarity Coefficient (mDSC) of 93.32% without pre-training. Compared to the R50 U-Net and R50 Att-UNet methods, our approach resulted in average mDSC improvements of 5.77% and 6.57%, respectively. When compared to methods based on concurrent Transformers, our approach still obtained performance gains of 0.45% and 1.54% in average mDSC over UNetr++ and nnFormer, respectively. Specifically, SegStitch demonstrated superior segmentation

Fig. 7. Qualitative validation and comparison of segmentation results from different methods on the Synapse multi-organ dataset and the ACDC cardiac segmentation dataset. From left to right: UNETR++, SegStitch (Ours), Ground Truth. Comprehensive analysis shows that we achieve finer and more accurate segmentation results.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00496v1-figure-46d4cf22-bf81-4e5a-aa00-906b7e7b0c9c.jpg)

accuracy across all crucial cardiac components, including the Right Ventricle (RV), Myocardium (Myo), and Left Ventricle (LV).

2) Comparisons with state-of-the-art methods on Synapse datase: We conducted experiments comparing the Synapse multi-organ segmentation dataset with state-of-the-art methods, as shown in Table V. Firstly, it can be observed that traditional CNN methods still exhibit good performance. However, the addition of Transformers or the use of pure Transformer architectures has been shown to have significant effectiveness, surpassing CNNs to some extent. Among these frameworks, R50 U-Net, R50 Att-Unet, and R50 ViT achieved mDSC scores of over 70% compared to V-Net, DARR, and ViT, indicating the overall superiority of Transformer-based models over traditional CNN models.

In our comparison with the UNETR++ model, which integrates a Transformer into a multi-attention mechanism, our experimental results demonstrate that our approach has achieved superior segmentation performance. Specifically, it has achieved an average accuracy of 89.80% on the Synapse dataset, along with an HD95 metric of 8.48. It is evident that our algorithm achieves a 2.58% improvement in average accuracy. Our model outperforms a range of CNN-based approaches (e.g., R50 U-Net, R50 Att-Unet, U-Net, Att-Unet) and Transformer-based methods (e.g., R50 ViT, TransUnet, Swin-Unet), including the Transformer-integrated UNETR++. It not only delivers superior overall segmentation results but also achieves the highest performance in segmenting specific organs such as the Aorta, Kidney(L), Kidney(R), and Stomach, surpassing the next best by margins of 1.09%, 1.57%, 0.95%, and 1.62%, respectively.

The experimental results further validate the feasibility of our proposed Transformer-based multi-attention approach. Figure 7 also demonstrates some segmentation results. From these images, it can be observed that our framework can improve segmentation accuracy to a certain extent. The SegStitch model can learn high-level semantic features and low-level texture features, achieving precise localization and segmentation.

TABLE V

SEGMENTATION PERFORMANCE OF DIFFERENT METHODS ON THE SYNAPSE DATASET.THE BEST SCORES ARE INDICATED IN BOLD .

| Method           |   Aor ↑ |   Gal |   LKid ↑ |   RKid ↑ |   Liv |   Pan |   Spleen |   Sto ↑ |   AVG ↑ | HD95 ↓   |
|------------------|---------|-------|----------|----------|-------|-------|----------|---------|---------|----------|
| V-Net[59]        |   75.34 | 51.87 |    77.1  |    80.75 | 87.84 | 40.04 |    80.56 |   56.98 |   68.81 | -        |
| DARR[3]          |   74.74 | 53.77 |    72.31 |    73.24 | 94.08 | 54.18 |    89.9  |   45.96 |   69.77 | -        |
| R50 U-Net[3]     |   87.74 | 63.66 |    80.6  |    78.19 | 93.74 | 56.9  |    85.87 |   74.16 |   74.68 | -        |
| R50 Att-UNet[55] |   55.92 | 63.91 |    79.2  |    72.71 | 93.56 | 49.37 |    87.19 |   74.95 |   75.57 | -        |
| U-Net[3]         |   89.07 | 69.72 |    77.77 |    68.6  | 93.43 | 53.98 |    86.67 |   75.58 |   76.85 | 39.7     |
| Att-UNet[60]     |   89.55 | 68.88 |    77.98 |    71.11 | 93.57 | 58.04 |    87.3  |   75.75 |   77.77 | 36.02    |
| VIT[10]          |   70.19 | 45.1  |    74.7  |    67.4  | 91.32 | 42    |    81.75 |   70.44 |   67.86 | 36.11    |
| R50 VIT[10]      |   73.73 | 55.13 |    75.8  |    72.2  | 91.51 | 45.99 |    81.99 |   73.95 |   71.29 | 32.87    |
| TransUNet[56]    |   87.23 | 63.13 |    81.87 |    77.02 | 94.08 | 55.86 |    85.08 |   75.62 |   77.48 | 31.84    |
| Swin-UNet[57]    |   85.47 | 66.53 |    83.28 |    79.61 | 94.29 | 56.58 |    90.66 |   76.6  |   79.13 | 22.96    |
| nnFormer[50]     |   92.04 | 71.09 |    87.64 |    87.34 | 96.53 | 82.49 |    92.91 |   89.17 |   87.4  | 10.63    |
| MISSFormer[58]   |   86.99 | 68.65 |    85.21 |    82    | 94.41 | 65.67 |    91.92 |   80.81 |   81.96 | 18.20    |
| D-Former[35]     |   92.12 | 80.09 |    92.6  |    91.91 | 96.99 | 76.67 |    93.78 |   86.44 |   88.83 | -        |
| UNETR++[39]      |   92.52 | 71.25 |    87.54 |    87.18 | 96.42 | 81.1  |    95.77 |   86.01 |   87.22 | 7.53     |
| SegStitch (Ours) |   93.61 | 77.11 |    94.17 |    92.86 | 96.92 | 80.39 |    92.61 |   90.79 |   89.8  | 8.48     |

TABLE VI COMPUTATIONAL COMPLEXITY ANALYSIS OF THE PROPOSED SEGSTITCH. ALL COMPARISONS WERE CONDUCTED ON THE SYNAPSE DATASET.THE BEST SCORES ARE INDICATED IN BOLD .

| Method           |   Params(M) |   GFLOPs |   mDSC(%) |
|------------------|-------------|----------|-----------|
| TransUNet        |       96.07 |    88.91 |     77.48 |
| UNETR            |       92.49 |    75.46 |     78.35 |
| Swin-UNETR       |       62.83 |   384.2  |     83.48 |
| nnFormer         |      150.5  |   213.4  |     86.57 |
| SegStitch (Ours) |       60.81 |    67.36 |     89.8  |

## E. SegStitch computational complexity comparisons

The SegStitch method demonstrates significant advantages in computational complexity, primarily due to the substantial reduction in the number of learnable parameters. It is impressive that, despite having only about 6.08 million parameters, SegStitch has outperformed many other algorithms. As shown in Table VI, SegStitch has significantly improved efficiency by reducing the number of learnable parameters and the number of floating-point operations (FLOPs), yet it still maintains a first-class performance compared to other methods. SegStitch achieves a perfect balance between the number of parameters and excellent segmentation results, proving its efficiency in 3D medical image segmentation tasks. Furthermore, with only 67.36 giga floating-point operations, the compactness of SegStitch greatly simplifies the deployment and application of this method in actual clinical settings. The reduction in model size not only enhances the efficiency of medical imaging analysis but also strengthens its effectiveness in practical applications, thereby facilitating seamless integration in clinical practice.

## F. Ablation study of Synapse dataset

We conducted an in-depth ablation study on the core components of SegStitch using the Synapse dataset, aiming to explore the effectiveness of the proposed self-attention mechanism and to assess the specific impact of image input size on the model's segmentation accuracy. In our experimental design, we deliberately removed the self-attention mechanism or ODE block to verify the contribution of these key components to overall performance. The results are summarized in Table III. The findings indicate that both the self-attention mechanism and ODE block are crucial for maintaining the model's high performance, with the absence of either significantly degrading the model's performance. Overall, SegStitch achieves superior segmentation performance through its expanded attention mechanism and self-learning structure. The experiments further emphasize the importance of information exchange between samples, with the absence of self-attention mechanism leading to increased instability in the results. The ODE block, on the other hand, exhibits lower variability in results, which we speculate may be due to the global perception capabilities of the Transformer architecture, significantly enhancing the model's robustness.

## G. Discussion

1) clinical impact: Automated localization and segmentation of the heart and abdominal organs in MRI are crucial for various clinical applications, as they directly impact subsequent quantitative analysis and diagnosis. Although fully automated systems have not yet fully integrated into routine clinical practice, developing robust frameworks can significantly enhance diagnostic workflows. For instance, cardiac MRI segmentation is essential for measuring blood pool and heart volumes, which helps derive functional indices such as ventricular volume and ejection fraction. These quantification metrics are used for patient diagnosis. Therefore, this paper proposes an end-to-end network for automatic segmentation of the heart and abdominal organs. The network effectively delineates cardiac structures and, due to its lightweight design and long-range segmentation capabilities, offers significant advantages in clinical settings (see Figure 7).

2) Performance Analysis of the Self-Attention Mechanism: As shown in Figure 8, we extract the weight values of the feature maps in the self-attention module and visualize them in the form of a heatmap. experiments indicate that the self-attention module focuses more on presenting distant features. Leveraging this advantage, it significantly concentrates on accurately segmenting regions with elongated shapes. In summary, this method effectively learns global semantic

Fig. 8. Visualization of weights for SegStitch self-attention in feature maps. Darker colors indicate higher weight values. This figure demonstrates the regions of focus of the self-attention on the feature maps. From the figure, it can be observed that the self-attention module pays more attention to the elongated parts in the feature maps. This is because the ViT block with spatial positional information tends to learn global features.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00496v1-figure-88a98fc0-bbf8-4778-b393-ff403020ca99.jpg)

## TABLE VII

COMPARISON ANALYSIS OF THE IMPACT OF ODE BLOCK ON LOSS AND DICE SIMILARITY COEFFICIENT IN DIFFERENT ORDINARY DIFFERENTIAL EQUATIONS AND ALTERNATIVE MODULES ON THE ADCD DATASET.THE BEST SCORES ARE INDICATED IN BOLD .

| Basemodel   | Function   | train     |    loss | validation loss   |   mDSC(%) |
|-------------|------------|-----------|---------|-------------------|-----------|
| SegStitch   | sin        | 2 -0.9249 | -0.8895 |                   |     93.32 |
| SegStitch   | ODE        | sin       | -0.9289 | -0.8830           |     93.15 |
| SegStitch   |            | tan       | -0.9287 | -0.8832           |     93.07 |
| SegStitch   | AE         | -         | -0.9273 | -0.8753           |     92.38 |
| SegStitch   | sigmoid    | -         | -0.9275 | -0.8711           |     92.15 |
| SegStitch   | -          | -         | -0.9191 | -0.8741           |     92.21 |

information and local spatial detail features, with particularly outstanding segmentation of elongated structures.

To vividly demonstrate the performance of our SegStitch model, we conducted qualitative comparisons with UNETR++ (as shown in Figure 7). It is evident that on the Synapse dataset, our model's organ segmentation results are generally closer to reality. For specific organs, our model can accurately identify and outline the contours of the kidneys, capturing some detailed contours, which aligns with the conclusions of our experimental results. On the ACDC dataset, our model excels at delineating the contours of the right ventricle. Since the right ventricle typically exhibits a long and narrow shape, this indicates that our model possesses excellent long-range dependency capture and detail processing capabilities.

- 3) Robustness Analysis of ODE block: Developing a method endowed with robust capabilities is of paramount importance, particularly in the dynamic field of medical image analysis. The robustness of a model is indicative of its superior ability to effectively learn from complex and novel datasets. In this section, our objective is to demonstrate the exceptional resilience of SegStitch by conducting a comprehensive evaluation of its segmentation performance on the ACDC dataset.

For the segmentation task on the ACDC dataset, we conducted a comparative analysis of the performance between the model's blocks: one version incorporating the ODE block and another without it. The comparative results are presented in

Table VII. Notably, upon the incorporation of the ODE block, the training and validation loss curves during the training process exhibited a smoother descent. This robust performance effectively addresses the complexities inherent in medical image data, underscoring the efficacy of our approach.

In addition to the aforementioned, we have conducted an investigation into the impact of various modules on the performance of SegStitch. The findings are delineated in Table VII. It should be noted that ' sin 2 ', 'sin', and 'tanh' represent the ablation studies conducted on different ordinary differential equations within the ODE block framework. 'Sigmoid' refers to the scenario where the ODE block is supplanted by a sigmoid activation function, while 'AE' signifies the substitution of the ODE block with a more parameter-rich autoencoder. Our results indicate that the ODE block achieves optimal robustness when employing ' sin 2 ' as the ODE. Compared to the 'sigmoid' activation function, the ODE-based neural network demonstrates superior performance in data fitting. Furthermore, we conducted comparative experiments using 'AE'. During the training phase, the 'AE' experiment's training and validation losses exhibited more pronounced fluctuations, indicating that the model's performance is not solely a result of parameter stacking.

## V. CONCLUSION

In this paper, we introduce SegStitch, an innovative framework for medical image segmentation. Initially, it fuses 3D image features with spatial positional cues, thereby maintaining the integrity of spatial information within the images. Subsequently, by incorporating both coarse and fine-grained attention mechanisms into the Transformer architecture, we have markedly amplified the model's capacity to manage longrange dependencies and discern intricate details. Furthermore, the integration of an ODE block into SegStitch marks a pioneering step, significantly bolstering the model's robustness. To encapsulate, this paper unveils SegStitch, a novel framework for 3D medical image segmentation. Empirical validation demonstrates that our model not only streamlines parameter count and computational load but also surpasses existing benchmarks in segmentation performance on the public datasets Synapse and ACDC. This study heralds a significant leap in medical image analysis, with the potential to profoundly impact future medical image processing endeavors.

## REFERENCES

- [1] Y. Wang, Y. Zhou, W. Shen, S. Park, E. K. Fishman, A. L. Yuille, Abdominal multi-organ segmentation with organ-attention networks and statistical fusion, Medical image analysis 55 (2019) 88-102.
- [2] J. Long, E. Shelhamer, T. Darrell, Fully convolutional networks for semantic segmentation, in: Proceedings of the IEEE conference on computer vision and pattern recognition, 2015, pp. 3431-3440.
- [3] O. Ronneberger, P. Fischer, T. Brox, U-net: Convolutional networks for biomedical image segmentation, in: Medical Image Computing and Computer-Assisted Intervention-MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, Springer, 2015, pp. 234-241.
- [4] Z. Zhou, M. M. Rahman Siddiquee, N. Tajbakhsh, J. Liang, Unet++: A nested u-net architecture for medical image segmentation, in: Deep Learning in Medical Image Analysis and Multimodal Learning for Clinical Decision Support: 4th International Workshop, DLMIA 2018, and 8th International Workshop, ML-CDS 2018, Held in Conjunction with MICCAI 2018, Granada, Spain, September 20, 2018, Proceedings 4, Springer, 2018, pp. 3-11.

- [5] H. Huang, L. Lin, R. Tong, H. Hu, Q. Zhang, Y. Iwamoto, X. Han, Y.-W. Chen, J. Wu, Unet 3+: A full-scale connected unet for medical image segmentation, in: ICASSP 2020-2020 IEEE international conference on acoustics, speech and signal processing (ICASSP), IEEE, 2020, pp. 1055-1059.
- [6] X. Xiao, S. Lian, Z. Luo, S. Li, Weighted res-unet for high-quality retina vessel segmentation, in: 2018 9th international conference on information technology in medicine and education (ITME), IEEE, 2018, pp. 327-331.
- [7] ¨ . O C ¸ icek, A. Abdulkadir, S. S. Lienkamp, T. Brox, O. Ronneberger, 3d ¸ u-net: learning dense volumetric segmentation from sparse annotation, in: Medical Image Computing and Computer-Assisted InterventionMICCAI 2016: 19th International Conference, Athens, Greece, October 17-21, 2016, Proceedings, Part II 19, Springer, 2016, pp. 424-432.
- [8] H. K. Putra, B. Suprihatin, F. Ramadhini, et al., ybrid clahe-gamma adjustment and densely connected u-net for retinal blood vessel segmentation using augmentation data., Engineering Letters 30 (2) (2022).
- [9] C. Yu, J. Wang, C. Gao, G. Yu, C. Shen, N. Sang, Context prior for scene segmentation, in: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, 2020, pp. 12416-12425.
- [10] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al., An image is worth 16x16 words: Transformers for image recognition at scale, arXiv preprint arXiv:2010.11929 (2020).
- [11] Y. Ji, H. Saratchandran, C. Gordon, Z. Zhang, S. Lucey, Sine activated low-rank matrices for parameter efficient learning, arXiv preprint arXiv:2403.19243 (2024).
- [12] B. Wu, Y. Xie, Z. Zhang, M. H. Phan, Q. Chen, L. Chen, Q. Wu, Xlip: Cross-modal attention masked modelling for medical language-image pre-training, arXiv preprint arXiv:2407.19546 (2024).
- [13] Z. Zhang, X. Qi, M. Chen, G. Li, R. Pham, A. Zuhair, E. Berry, Z. Liao, O. Siggs, R. Mclaughlin, et al., Jointvit: Modeling oxygen saturation levels with joint supervision on long-tailed octa, arXiv preprint arXiv:2404.11525 (2024).
- [14] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, I. Polosukhin, Attention is all you need, Advances in neural information processing systems 30 (2017).
- [15] R. Shao, Z. Shi, J. Yi, P.-Y. Chen, C.-J. Hsieh, On the adversarial robustness of vision transformers, arXiv preprint arXiv:2103.15670 (2021).
- [16] A. Radford, K. Narasimhan, T. Salimans, I. Sutskever, et al., Improving language understanding by generative pre-training (2018).
- [17] H. Yan, J. Du, V. Y. Tan, J. Feng, On robustness of neural ordinary differential equations, arXiv preprint arXiv:1910.05513 (2019).
- [18] Z. Yi, nmode: neural memory ordinary differential equation, Artificial Intelligence Review (2023) 1-36.
- [19] O. Bernard, A. Lalande, C. Zotti, F. Cervenansky, X. Yang, P.-A. Heng, I. Cetin, K. Lekadir, O. Camara, M. A. G. Ballester, et al., Deep learning techniques for automatic mri cardiac multi-structures segmentation and diagnosis: is the problem solved?, IEEE transactions on medical imaging 37 (11) (2018) 2514-2525.
- [20] B. Landman, Z. Xu, J. Igelsias, M. Styner, T. Langerak, A. Klein, Miccai multi-atlas labeling beyond the cranial vault-workshop and challenge, in: Proc. MICCAI Multi-Atlas Labeling Beyond Cranial Vault-Workshop Challenge, Vol. 5, 2015, p. 12.
- [21] A. Hatamizadeh, Y. Tang, V. Nath, D. Yang, A. Myronenko, B. Landman, H. R. Roth, D. Xu, Unetr: Transformers for 3d medical image segmentation, in: Proceedings of the IEEE/CVF winter conference on applications of computer vision, 2022, pp. 574-584.
- [22] F. Isensee, P. F. Jaeger, S. A. Kohl, J. Petersen, K. H. Maier-Hein, nnunet: a self-configuring method for deep learning-based biomedical image segmentation, Nature methods 18 (2) (2021) 203-211.
- [23] B. Wu, Y. Xie, Z. Zhang, J. Ge, K. Yaxley, S. Bahadir, Q. Wu, Y. Liu, M.S. To, Bhsd: A 3d multi-class brain hemorrhage segmentation dataset, in: International Workshop on Machine Learning in Medical Imaging, Springer, 2023, pp. 147-156.
- [24] Z. Zhang, B. Zhang, A. Hiwase, C. Barras, F. Chen, B. Wu, A. J. Wells, D. Y. Ellis, B. Reddi, A. W. Burgan, et al., Thin-thick adapter: Segmenting thin scans using thick annotations.
- [25] Z. Zhang, X. Qi, B. Zhang, B. Wu, H. Le, B. Jeong, M.-S. To, R. Hartley, Segreg: Segmenting oars by registering mr images and ct annotations, arXiv preprint arXiv:2311.06956 (2023).
- [26] N. Ibtehaz, M. S. Rahman, Multiresunet: Rethinking the u-net architecture for multimodal biomedical image segmentation, Neural networks 121 (2020) 74-87.
- [27] J. Chen, Y. Lu, Q. Yu, X. Luo, E. Adeli, Y. Wang, L. Lu, A. L. Yuille, Y. Zhou, Transunet: Transformers make strong encoders for medical image segmentation, arXiv preprint arXiv:2102.04306 (2021).
- [28] R. Azad, M. T. Al-Antary, M. Heidari, D. Merhof, Transnorm: Transformer provides a strong spatial normalization mechanism for a deep segmentation model, IEEe Access 10 (2022) 108205-108215.
- [29] G. Sun, Y. Pan, W. Kong, Z. Xu, J. Ma, T. Racharak, L.-M. Nguyen, J. Xin, Da-transunet: integrating spatial and channel dual attention with transformer u-net for medical image segmentation, Frontiers in Bioengineering and Biotechnology 12 (2024) 1398237.
- [30] H. Sun, J. Xu, Y. Duan, Paratranscnn: Parallelized transcnn encoder for medical image segmentation, arXiv preprint arXiv:2401.15307 (2024).
- [31] Y. Chang, H. Menghan, Z. Guangtao, Z. Xiao-Ping, Transclaw u-net: Claw u-net with transformers for medical image segmentation, arXiv preprint arXiv:2107.05188 (2021).
- [32] R. Azad, M. Heidari, M. Shariatnia, E. K. Aghdam, S. Karimijafarbigloo, E. Adeli, D. Merhof, Transdeeplab: Convolution-free transformer-based deeplab v3+ for medical image segmentation, in: International Workshop on PRedictive Intelligence In MEdicine, Springer, 2022, pp. 91-102.
- [33] R. Azad, Y. Jia, E. K. Aghdam, J. Cohen-Adad, D. Merhof, Enhancing medical image segmentation with transception: A multi-scale feature fusion approach, arXiv preprint arXiv:2301.10847 (2023).
- [34] X. Huang, Z. Deng, D. Li, X. Yuan, Y. Fu, Missformer: An effective transformer for 2d medical image segmentation, IEEE Transactions on Medical Imaging 42 (5) (2022) 1484-1494.
- [35] Y. Wu, K. Liao, J. Chen, J. Wang, D. Z. Chen, H. Gao, J. Wu, D-former: A u-shaped dilated transformer for 3d medical image segmentation, Neural Computing and Applications 35 (2) (2023) 1931-1944.
- [36] A. Hatamizadeh, Z. Xu, D. Yang, W. Li, H. Roth, D. Xu, Unetformer: A unified vision transformer model and pre-training framework for 3d medical image segmentation, arXiv preprint arXiv:2204.00631 (2022).
- [37] Y. Jiang, Y. Zhang, X. Lin, J. Dong, T. Cheng, J. Liang, Swinbts: A method for 3d multimodal brain tumor segmentation using swin transformer, Brain sciences 12 (6) (2022) 797.
- [38] L. Xia, H. Zhang, Y. Wu, R. Song, Y. Ma, L. Mou, J. Liu, Y. Xie, M. Ma, Y. Zhao, 3d vessel-like structure segmentation in medical images by an edge-reinforced network, Medical Image Analysis 82 (2022) 102581.
- [39] A. Shaker, M. Maaz, H. Rasheed, S. Khan, M.-H. Yang, F. S. Khan, Unetr++: delving into efficient and accurate 3d medical image segmentation, arXiv preprint arXiv:2212.04497 (2022).
- [40] P. Dong, H. Niu, Z. Yi, X. Xu, nmpls-net: Segmenting pulmonary lobes using nmode, Mathematics 11 (22) (2023) 4675.
- [41] S. Wang, Y. Chen, Z. Yi, nmode-unet: A novel network for semantic segmentation of medical images, Applied Sciences 14 (1) (2024) 411.
- [42] J. Hu, C. Yu, Z. Yi, H. Zhang, Enhancing robustness of medical image segmentation model with neural memory ordinary differential equation., International Journal of Neural Systems (2023) 2350060-2350060.
- [43] H. Pinckaers, G. Litjens, Neural ordinary differential equations for semantic segmentation of individual colon glands, arXiv preprint arXiv:1910.10470 (2019).
- [44] R. Valle, F. Reda, M. Shoeybi, P. Legresley, A. Tao, B. Catanzaro, Neural odes for image segmentation with level sets, arXiv preprint arXiv:1912.11683 (2019).
- [45] R. T. Chen, Y. Rubanova, J. Bettencourt, D. K. Duvenaud, Neural ordinary differential equations, Advances in neural information processing systems 31 (2018).
- [46] X. Liu, T. Xiao, S. Si, Q. Cao, S. Kumar, C.-J. Hsieh, How does noise help robustness? explanation and exploration under the neural sde framework, in: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2020, pp. 282-290.
- [47] Q. Kang, Y. Song, Q. Ding, W. P. Tay, Stable neural ode with lyapunovstable equilibrium points for defending against adversarial attacks, Advances in Neural Information Processing Systems 34 (2021) 1492514937.
- [48] W. Wang, C. Chen, M. Ding, H. Yu, S. Zha, J. Li, Transbts: Multimodal brain tumor segmentation using transformer, in: Medical Image Computing and Computer Assisted Intervention-MICCAI 2021: 24th International Conference, Strasbourg, France, September 27-October 1, 2021, Proceedings, Part I 24, Springer, 2021, pp. 109-119.
- [49] G. Xu, X. Wu, X. Zhang, X. He, Levit-unet: Make faster encoders with transformer for medical image segmentation, arXiv preprint arXiv:2107.08623 (2021).
- [50] H.-Y. Zhou, J. Guo, Y. Zhang, L. Yu, L. Wang, Y. Yu, nnformer: Interleaved transformer for volumetric segmentation, arXiv preprint arXiv:2109.03201 (2021).

- [51] Y. Zhang, H. Liu, Q. Hu, Transfuse: Fusing transformers and cnns for medical image segmentation, in: Medical Image Computing and Computer Assisted Intervention-MICCAI 2021: 24th International Conference, Strasbourg, France, September 27-October 1, 2021, Proceedings, Part I 24, Springer, 2021, pp. 14-24.
- [52] J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, L. Fei-Fei, Imagenet: A large-scale hierarchical image database, in: 2009 IEEE conference on computer vision and pattern recognition, Ieee, 2009, pp. 248-255.
- [53] L. Bottou, Stochastic gradient descent tricks, in: Neural Networks: Tricks of the Trade: Second Edition, Springer, 2012, pp. 421-436.
- [54] O. Ronneberger, P. Fischer, T. Brox, U-net: Convolutional networks for biomedical image segmentation, in: Medical Image Computing and Computer-Assisted Intervention-MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, Springer, 2015, pp. 234-241.
- [55] J. Schlemper, O. Oktay, M. Schaap, M. Heinrich, B. Kainz, B. Glocker, D. Rueckert, Attention gated networks: Learning to leverage salient regions in medical images, Medical image analysis 53 (2019) 197-207.
- [56] J. Chen, Y. Lu, Q. Yu, X. Luo, E. Adeli, Y. Wang, L. Lu, A. L. Yuille, Y. Zhou, Transunet: Transformers make strong encoders for medical image segmentation, arXiv preprint arXiv:2102.04306 (2021).
- [57] H. Cao, Y. Wang, J. Chen, D. Jiang, X. Zhang, Q. Tian, M. Wang, Swin-unet: Unet-like pure transformer for medical image segmentation, in: European conference on computer vision, Springer, 2022, pp. 205218.
- [58] X. Huang, Z. Deng, D. Li, X. Yuan, Missformer: An effective medical image segmentation transformer, arXiv preprint arXiv:2109.07162 (2021).
- [59] F. Milletari, N. Navab, S.-A. Ahmadi, V-net: Fully convolutional neural networks for volumetric medical image segmentation, in: 2016 fourth international conference on 3D vision (3DV), Ieee, 2016, pp. 565-571.
- [60] O. Oktay, J. Schlemper, L. L. Folgoc, M. Lee, M. Heinrich, K. Misawa, K. Mori, S. McDonagh, N. Y. Hammerla, B. Kainz, et al., Attention u-net: Learning where to look for the pancreas, arXiv preprint arXiv:1804.03999 (2018).

## APPENDIX A ABLATION STUDY OF ODE BLOCK

Fig. 9. On the ADCD dataset, we perform channel mean statistics on features processed by the encoder to compare the performance of different ordinary differential equations and alternative modules.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00496v1-figure-6acea9e0-61ab-4df4-99c2-7075a79c2787.jpg)

To further investigate the rationale behind these observations, we conducted zero-mean feature experiments following the ODE block. Using 40 samples, we calculated the mean of the encoded features for each channel by the model. The results suggest that the model's feature distribution becomes more uniform with the ODE in use. Moreover, when using ' sin 2 ' as the ODE, the channel means are more concentrated and align more closely with the normal distribution, enhancing the model's convergence and stability. The experimental outcomes are depicted in Figure 9.

Fig. 10. Conducting an ablation study on the nmODE model using the ACDC dataset, with a comparative analysis of the model's training loss and validation loss throughout the training process.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00496v1-figure-d2993177-3337-41a2-86ec-7d6bfb72db03.jpg)

We performed a comparative analysis on the ACDC dataset between models that include an ODE block and those that do not. The results indicate that, with the introduction of the ODE block, the training and validation loss curves exhibit a smoother descending trend during the training process. The comparative results are shown in Figure 10. This enhanced robustness effectively addresses the inherent complexity of medical image data, further validating the effectiveness of the ODE block in our model.

Fig. 11. Comprehensive Analysis of Loss Performance for the nmODE Module under Various Ordinary Differential Equations and Alternative Modules.

![image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00496v1-figure-a0cecfd6-b65d-4c91-9b5c-ee4eabebd981.jpg)

In addition to the previously discussed aspects, we also investigated the impact of different modules on the performance of SegStitch. The results of this study are detailed in Figure 11. The terms ' sin 2 ', 'sin', and 'tanh' refer to our ablation studies on different ordinary differential equations (ODEs) used in the ODE block. 'Sigmoid' indicates that in this case, the ODE block was replaced with a sigmoid activation function, while 'AE' denotes that the ODE block was replaced with a more parameter-rich autoencoder. Our results show that using ' sin 2 ' as the ODE provided the best robustness for the ODE block module. Compared to the 'sigmoid' activation function, the ordinary differential equation neural network demonstrated superior performance in data fitting. Furthermore, we conducted comparative experiments with 'AE'. During the training phase, the training and validation losses for the 'AE' experiments exhibited more pronounced fluctuations, confirming that the model's performance issues were not merely due to parameter stacking.
Shengbo Tan Zeyu Zhang Ying Cai Daji Ergu Lin Wu Binbin Hu Pengzhang Yu Yang Zhao
SegStitch: Multidimensional Transformer for Robust and Efficient Medical Imaging Segmentation
Medical imaging segmentation plays a significant role in the automatic
recognition and analysis of lesions. State-of-the-art methods, particularly
those utilizing transformers, have been prominently adopted in 3D semantic
segmentation due to their superior performance in scalability and
generalizability. However, plain vision transformers encounter challenges due
to their neglect of local features and their high computational complexity. To
address these challenges, we introduce three key contributions: Firstly, we
proposed SegStitch, an innovative architecture that integrates transformers
with denoising ODE blocks. Instead of taking whole 3D volumes as inputs, we
adapt axial patches and customize patch-wise queries to ensure semantic
consistency. Additionally, we conducted extensive experiments on the BTCV and
ACDC datasets, achieving improvements up to 11.48% and 6.71% respectively in
mDSC, compared to state-of-the-art methods. Lastly, our proposed method
demonstrates outstanding efficiency, reducing the number of parameters by 36.7%
and the number of FLOPS by 10.7% compared to UNETR. This advancement holds
promising potential for adapting our method to real-world clinical practice.
The code will be available at https://github.com/goblin327/SegStitch