Volumetric Axial Disentanglement Enabling Advancing in Medical Image Segmentation

1 Introduction

The precise delineation of clinical diagnoses critically hinges on the analysis of volumetric medical imaging data, such as MRI and CT scans. These volumetric data encapsulate the exact spatial configurations of tissues and organs within patients' bodies and may capture dynamic changes over time, including blood flow or organ movement. The inherent volumetric nature of these data enables an in-depth analysis of pathological states and physiological processes. Therefore, accurate segmentation of complex tissue structures is paramount for extracting precise biomorphic information, effectively marking disease, and facilitating subsequent analyses and diagnoses. However, 2D segmentation methods fail to address the hierarchical relationships inherent in these data since they lack the capacity to adequately represent the spatial continuity of volumetric data to incur segmentations with poor 3D cohesion and precision.

∗Corresponding authors.

The symmetrical design of 3D conv extends the convolution kernel across all axes to extract features from the entire volumetric space by treating the time axis (z-axis) information the same as cross-section (xy-plane) information, and hence to optimize the use of continuous volume information. This approach overlooks the distinct physical and scanning characteristics of different axes, especially for the differences between the z-axis and the xy-plane in terms of physical properties and information density. As illustrated in Fig. 1, the visualization of volumetric data highlights the differing information densities across axes, the xy-plane accurately reflects tissue density distribution and shows precise locations of various tissue types, and the z-axis represents the morphological changes of the scanned sequence over time. When z-axis information is rich, the model should focus more on its detailed z-axis features; whilst z-axis information is sparse, the model should prioritize the detailed features of the xy-plane. However, the current 3D conv approaches fail to consider these differences, often leading to inadequate differentiation and ineffective use of the unique characteristics of each axis.

Due to the differences in resolution and information density between the z-axis and xy-plane in volumetric data, the usage of 3D conv may introduce spatial biases and temporal errors by uniformly extracting features across all dimensions. This can lead to the misapplication of xy-plane features to the temporal z-axis, thereby limiting segmentation accuracy. To address this issue, we propose the volumetric axial disentanglement and further design the Post-Axial Refiner (PaR) module, which disentangles specific axial features and applies specialized attention approaches to post-decoding 3D feature maps (coarse masks). PaR module refines temporal and spatial features for accurate mapping of temporal variations and structural differences to enhance segmentation accuracy. This plug-and-play module can be seamlessly integrated into any 3D segmentation network architecture.

The proposed method is validated using publicly available volumetric segmentation datasets (FLARE2021, OIMHS, and SegTHOR) and compared with 11 previous state-of-theart models to verify its effectiveness. The results demonstrate effective improvements in segmentation performance across all models with the integration of PaR, including enhanced overall accuracy and boundary delineation. The results underscore the enhancement of volumetric axial disentanglement for precise medical volumetric segmentation.

Our contributions are summarized as follows:

1. Substantiating the operation of volumetric axial disentanglement as a means to tackle the challenges arising from variations in physical properties and information density across distinct axial directions in 3D medical imaging.
2. Introducing a simple, plug-and-play PaR module that performs volumetric axial disentanglement and attentionbased axial feature extraction to enhance specific axial feature representation in 3Dmedical imaging.
3. By validating on multiple publicly available datasets FLARE2021, OIMHS, and SegTHOR with 11 of the mainstream volumetric segmentation models, the proposed PaR module consistently improves segmentation performance.

2 RelatedWork

2.1 Volumetric Segmentation

Volumetric segmentation methods have significantly advanced due to developments in deep learning and computational resources. The rise and continuous iterations of Convolutional Neural Networks (CNN) and Vision Transformers (ViT) [Dosovitskiy et al., 2020] have propelled these advancements [Liao et al., 2020; Huang et al., 2023; Huang et al., 2024; Huang et al., ]. The 3D U-Net [C¸ ic¸ek et al., 2016], a seminal model in this field, employs 3D conv to capture spatial relationships within volumetric data. Enhancements to the 3D U-Net include attention mechanisms that prioritize relevant regions of the input data [Jiang et al., 2022] to enhance feature extraction capabilities and reduce noise interference. Transformer [Vaswani et al., 2017] models adapted for volumetric segmentation leverage selfattention for long-range dependencies, as seen in the UNETR architecture [Hatamizadeh et al., 2021a; Shaker et al., 2024]. Hybrid models that integrate CNNs and transformers balance spatial feature extraction with global context capture, resulting in further performance gains [Chen et al., 2021; Wang et al., 2021; Xie et al., 2021; Hatamizadeh et al., 2021b]. Notably, the nnU-Net [Isensee et al., 2021], a selfadapting framework, has demonstrated superior performance in a wide range of medical imaging tasks.

Beyond network architecture, numerous other research efforts have emerged for volumetric segmentation. Postprocessing approaches, including morphological operations and conditional random fields (CRFs), are employed to refine segmentation outcomes [Chen et al., 2022]. Semi-supervised and unsupervised segmentation methods enhance robustness by leveraging large volumes of unlabeled data [Bortsova et al., 2019; Huang et al., 2022]. Additionally, Generative Adversarial Networks (GANs) improve the realism and accuracy of segmentation outputs [Kwon et al., 2019; Goodfellow et al., 2020]. Federated learning enables model training across institutions while preserving privacy [Kaissis et al., 2021; Miao et al., 2023]. Overall, continuous innovations in neural network architectures and learning paradigms promise more accurate and reliable medical image analysis.

2.2 SegmentationMask Refinemement

Segmentation mask refinement methods aim to improve the accuracy of masks generated by pre-existing segmentation models. While numerous refinement methods exist for 2D image segmentation, few of them address the refinement of volumetric medical image segmentation masks.

In 2D segmentation mask refinement, PointRend [Kirillov et al., 2020] employs an MLP to correct low-confidence pixel labels within Mask R-CNN [He et al., 2017]. SegFix [Yuan et al., 2020] improves edge predictions by learning a mapping function between edge and inner pixels to replace inaccurate edge predictions with corresponding inner pixel predictions. RefineMask [Zhang et al., 2021] enhances Mask R-CNN by incorporating a semantic head for additional guidance. MaskTransfiner [Ke et al., 2022] uses an FCN to detect incoherent regions and refine their labels with a Transformer.

As for the refinement of the volumetric image segmentation mask, [Duan et al., 2019] designs a refinement step that imposes shape prior knowledge to augment segmentation quality and mitigate image artifacts. [Xue et al., 2023] introduces an approach employing a graph convolutional network (GCN) to refine inaccurate segmentation masks. [Shu et al., 2024] proposes a method utilizing a self-refine module to improve the labels produced by an auxiliary network through progressively accurate predictions from the primary network.

3 Methodology

Firstly, we provide a concise overview of the Post-Axial Refiner (PaR) module, which consists of two main pipelines: Permutative Dimensional Intensification (PDI) and Sequential Axial Intensification (SAI). We then explore the structure and functionality of the PDI pipeline and analyze its role in capturing and learning information of specific axial features. Finally, we present the SAI pipeline, which is dedicated to axial disentanglement on a specific axis for extracting and enhancing distinctive spatio-temporal characteristics.

3.1 Overview of the Post-Axial Refiner

The detailed structure of PaR is illustrated in Fig. 2. The baseline network takes the original volume image as input, and the 3D feature map generated by its decoder is defined as the coarse mask yo ∈ RC×W×H×D , where C represents the number of segmentation categories, and W, H, and D denote the width, height, and depth of the volumetric medical images, respectively. Given the inherently richer spatiotemporal information and texture details of volumetric medical images as opposed to 2D plane images, PaR utilizes volumetric axial disentanglement to extract the key features across the transverse, sagittal, and coronal planes of volumetric data to enhance the model's feature representation capability. To effectively capture the distinct feature variations in each axial plane and extract important feature information, PaR employs two pipelines: PDI and SAI.

The PDI pipeline is designed to capture and learn critical feature information within the coarse masks by applying a dimensional permutation operation and multi-head selfattention [Vaswani et al., 2017] along the specific axis. Initially, various spatial distribution tensors are generated from the coarse masks yo by a permutation operation to enrich the semantic content. Subsequently, multi-head self-attention is employed to learn a unified feature structure across these spatial distributions to ultimately capture the most salient specific axial feature information, represented as yrP ∈ RC×W×H×D .

For the enhancement of specific axial feature representation in volumetric data through volumetric axial disentanglement, we introduce the SAI pipeline. It initiates with 3D Conv, Group Normalization (GN), ReLU, and residual connection, preliminarily extracting features from the coarse masks. These features are then grouped along the specific axis and analyzed both channel-wise and spatial-wise to learn the key characteristic of specific axial features, resulting in yrS ∈ RC×W×H×D .

The outputs of both PDI yrP and SAI yrS are combined through addition, subsequently constrained within the (0,1) range via a sigmoid function, and then element-wise multiplied by the original coarse mask yo . This product is further enhanced through residual connections with the mathematical formulation of this progress as follows:

where yr ∈ RC×W×H×D is the refined segmentation mask.

3.2 Permutative Dimensional Intensification

Permutative Dimensional Intensification (PDI) improves volumetric image segmentation accuracy and robustness by capturing and learning specific axial feature information across different spatial distributions within the coarse masks. Here, we set the W-axis as the specific disentanglement axis.

For the coarse mask yo , the PDI pipeline generates spatial distributions through a dimensional permutation operation. This operation rearranges the dimensions of yo to produce multiple feature maps {yi}NS i=1 under different spatial configurations, such as y1 ∈ RH×D×C×W , y2 ∈ RC×D×H×W , and y3 ∈ RC×H×D×W . The W-axis is permuted to the last dimension to enhance the focus on W-axial information in subsequent attention. These permuted feature maps are then merged into 3D tensors {yα i }NS i=1 to make them suitable as the input for the multi-head self-attention, which is applied to the extraction of time-axial and spatial features. The formula for multi-head self-attention can be expressed as:

where {yγ i }NS i=1 are the attention outputs.

Subsequently, attention outputs {yγ i }NS i=1 are reshaped back into 5D formats {yδ i }NS i=1 through the reshape operation to be aligned with the dimensions of the coarse masks yo by permutation, denoted as enhanced segmentation maps {yr i }NS i=1 . The final step in PDI involves a summative integration of these enhanced segmentation maps to produce the W-axial refined segmentation mask yrP , thereby strengthening the coarse mask feature representation through capturing and learning specific spatio-temporal information.

3.3 Sequential Axial Intensification

The inherent brevity of sequence lengths in volumetric medical images presents a fundamental drawback that is often mitigated through interpolation in the data preprocessing stage, leading to variations in time-axial information. Given that most segmentation networks typically identify segmentation targets based on the slice plane, they struggle to understand the rich spatio-temporal information across the sequence. To address this issue, we introduce the Sequential Axial Intensification (SAI) pipeline, which applies axial disentanglement and two attention groups along the specific axis (W-axis here). This pipeline enhances W-axial spatio-temporal representation and learns unified features across different axial planes of the coarse masks to refine the coarse masks.

For the coarse mask yo , the SAI pipeline begins with an initial feature extraction layer comprising 3D Conv, GN, ReLU, and residual connection. This layer performs the preliminary feature extraction to produce the feature map yθ ∈ RC×W×H×D . To capture the unique patterns and characteristics of the W-axial feature by subsequent attention in yθ , the feature map is divided along the W-axis into subsets {yλ i ∈ RC× WG ×H×D }Gi=1 . These subsets are further divided into two groups: {yλ,c i ∈ RC× W 2G ×H×D }Gi=1 and {yλ,s i ∈ RC× W 2G ×H×D }Gi=1 .

The channel-wise groups. {yλ,c i }Gi=1 is processed through the adaptive max pooling fω and the adaptive mean pooling fη to extract W-axial channel information, which is then combined via addition and passed through an adaptive weighting function Fc(x) = Wcx + bc to generate {yfc i ∈ RC×1×1×1 }Gi=1 :

where Wc ∈ R1×1×1×1 and bc ∈ RC×1×1×1 are trainable parameters. Note that the initial weight of Wc is set to 1, and the initial weight of bc is set to 0.

Then {yfc i }Gi=1 are constrained between (0,1) using a sigmoid function and combined with {yλ,c i }Gi=1 through element-wise multiplication to extract the key channel features within {yλ,c i }Gi=1 .

The spatial-wise groups. {yλ,s i }Gi=1 undergo a 3D Conv layer to capture spatial features within the local neighborhood of the feature map. This operation can reflect the significance of spatial positions across {yλ,s i }Gi=1 . The captured features are then processed by an adaptive weighting function Fs(x) = Wsx + bs to generate {yfs i ∈ R1× W 2G ×H×D }Gi=1 : {yfs i = Fs(Conv3D(yλ,s i )) =Ws(Conv3D(yλ,s i ))+bs}Gi=1 , where Ws ∈ R1×1×1×1 and bs ∈ R1× W 2G ×H×D are trainable parameters, set to 1 and 0, respectively.

Similarly, {yfs i }Gi=1 is processed by sigmoid activation to ensure values between (0,1) and then are combined with {yλ,s i }Gi=1 by element-wise multiplication to extract the key spatial features within {yλ,s i }Gi=1 .

Finally, the outputs of channel-wise groups and spatialwise groups are concatenated and reshaped back to the dimensions of the coarse mask yo , culminate in yµ S ∈ RC×W×H×D , which is redistributed spatially to further enhance specific axial spatio-temporal feature representation.

Spatial redistribution. To further extract and enhance spatial feature representation from the disentangled segmentation map of the W-axis, we employ a spatial redistribution approach to capture common features in different spatial states.

The map yµ S is firstly reshaped into yµ1 S ∈ RC×G× WG ×H×D , followed by a dimensional exchange between G and WG dimensions to generate yµ2 S ∈ RC× WG ×G×H×D . This transformed tensor is then reshaped back to yµ3 S ∈ RC×W×H×D and added to yµ S to generate yrS ∈ RC×W×H×D . Through the supplement of spatial information and learning unified features across different axes, spatial redistribution augments the model's capacity to comprehend axial spatial dynamics, aiding in the capture of complex spatio-temporal relationships.

4 Experiments

4.1 Datasets and Implementation Details

The proposed PaR module is jointly trained end-to-end with the baseline model, using the 3D feature maps output by the baseline decoder as input to the module. To substantiate the efficacy of our method, we conducted experiments across three publicly available datasets: FLARE2021 [Ma et al., 2022], OIMHS [Ye et al., 2023], and SegTHOR [Lambert et al., 2020]. For all datasets, we employ an 8:1:1 random split for the training, validation, and testing sets. To ensure fairness, we employ identical data preprocessing protocols and hyperparameter configurations. Across all training sessions, we utilize LDiceCE as the loss function and the AdamW optimizer with a learning rate of 0.0001, over 80,000 iterations, and a batch size of 2. Validation employs a sliding window approach with a 0.5 overlap. Further details on the datasets and experimental setups are provided in the Appendix.

4.2 Benchmarking Results

We evaluate the computational cost and performance implications of integrating the PaR into 11 baselines on FLARE2021, OIMHS, and SegTHOR datasets, as demonstrated in quantitative results Table 1 and qualitative results Fig. 3.

FLARE2021. In our evaluation of the FLARE2021 dataset, all models exhibit performance improvements across all metrics, including mIoU, Dice, and HD95. The enhancements in mIoU and Dice range from 0.39% to 3.33% and 0.2% to 2.26%, respectively, underscoring PaR's efficacy in refining global segmentation accuracy through multi-head attention on different spatial configurations along the specific axis. Furthermore, the HD95 metric, which evaluates boundary segmentation accuracy, indicated notable improvements, exemplified by the reduction in HD95 for the 3D U-Net model from 16.31 to 2.49. This demonstrates PaR's ability to correct significant segmentation errors by learning spatial dependencies along the axial dimension.

OIMHS. The benchmark results on theOIMHS dataset further demonstrate robust performance enhancements across all baseline models when integrated with PaR. The mIoU and Dice improve range from 0.81% to 4.96% and 0.45% to 3.51%, respectively. Notably, the UNETR model shows a reduction in HD95 from 29.15 to 7.24, illustrating PaR's effectiveness in correcting boundary delineations through centralized axial disentanglement of volumetric information. Even state-of-the-art models like 3D UX-NET experienced improvements, with mIoU and Dice increasing by 1.06% and 0.65%, respectively. By disentangling the specific axis, PaR can extract and enhance spatio-temporal feature representation to surpass the limitations of the isotropic nature of 3D conv and improve performance.

SegTHOR. The effectiveness of PaR in disentangling timeaxial features and enhancing spatio-temporal representation is similarly demonstrated in the SegTHOR dataset, with a stable increase in mIoU and Dice scores by 0.32% to 3.99% and 0.15% to 2.61% across all benchmark models. The boundary correction effects of PaR are also apparent, with reductions in HD95 of the RAUNet model from 14.58 to 2.99. These results consistently demonstrate PaR's ability to improve segmentation performance by volumetric axial disentanglement and focused attention along the specific axis.

Experimental results across three datasets confirm that the integration of PaR can consistently enhance performance across CNN-based segmentation models while introducing only a marginal computational cost. This improvement transcends model architectures and dataset characteristics, indicating that volumetric axial disentanglement along a specific axis could effectively boost overall segmentation accuracy and boundary precision. Even in state-of-the-art models, PaR provides measurable gains, validating the robustness and applicability in various 3D medical imaging contexts.

4.3 Ablation Studies

Key components and fusion strategy. To validate the efficacy of our proposed method across CNN-based and Transformer-based architectures, we conducted ablation experiments using 3D U-Net and UNETR as baselines on the FLARE2021 and OIMHS datasets, as shown in Tables 2 and

3. The integration of either the PDI or SAI pipeline individually improved segmentation performance, with PDI capturing and learning specific axial critical spatio-temporal information and SAI applying specific axial feature disentanglement to enhance feature representation. Fusion strategies of PDI and SAI pipelines are also evaluated, including addition, element-wise multiplication, and cat. The addition strategy demonstrates the most obvious improvements across all metrics, underscoring the complementary nature of the PDI and SAI pipelines. Overall, the integrated PaR module consistently outperformed baseline models in segmentation accuracy and boundary precision, validating the effectiveness of specific axial feature disentanglement.

Axial information density reduction. To better understand and validate the efficacy of PaR in W-axial feature disentanglement, we conducted ablation experiments on the FLARE2021 dataset under different frame sampling scenarios, simulating various physical spatial distances in 3D medical imaging, with results shown in Table 4. The experiment evaluates the model's performance improvements with PaR integration by deleting every second and third frame to assess the impact of varying time-axial information density. The findings indicate that as time-axial information is reduced, the segmentation performance gains from PaR become more pronounced. This demonstrates that baseline models tend to allocate equal attention to all axes, whereas PaR effectively disentangles information along the W-axis, enhancing the baseline model's ability to extract meaningful spatio-temporal features in scenarios with limited axial data. These results underscore the importance of axial disentanglement in volumetric data when time-axial information density is insufficient.

Disentanglement along different axes. We conduct an ablation study to assess the impact of integrating the PaR module for disentangling along different axes on segmentation performance, validated on the FLARE2021 and OIMHS datasets, with results presented in Table 5 and Table 6. The results demonstrate that disentangling along the W, H, or D axes consistently improves the baseline model's segmentation performance, supporting the hypothesis that specific axial disentanglement mitigates the inherent anisotropic limitations of 3D convs in volumetric data.

Heads impact on axial expression. We conducted ablation experiments on the FLARE2021 dataset to evaluate the impact of varying the number of multi-head self-attention heads in the PDI pipeline, as shown in Table 7, to confirm optimal hyperparameter selection. The results suggest that appropriate heads should be selected based on axial information density, aiming to optimize feature mapping across multiple subspaces, thereby reducing computational overhead and enhancing feature expressiveness. Experiments with 4, 8, 16, and 24 heads reveal that 8 heads provide the best performance among all configurations.

Kernel size effects on feature extraction. An ablation study was conducted on the kernel sizes of SAI's preliminary feature extraction layer on FLARE2021 to determine the optimal hyperparameter settings, with results presented in Table 8. Larger kernel sizes cover a broader local area, capturing more contextual information for edge region recognition. However, excessively large kernels might lead to oversmoothing, losing essential details and texture information. Experiments with kernel sizes of 3 and 7 show close Dice scores, yet the HD95 metric of kernel size 3 is better than that of kernel size 7, indicating superior boundary optimization.

5 Conclusion and FutureWork

In this paper, we propose a volumetric axial disentanglement approach to address the pronounced disparities in spatial information across different axial dimensions inherent in volumetric medical imaging. By designing the Post-Axial Refiner (PaR) module, we have enabled a disentangled representation of spatial states and time-axial features, which effectively enhances both the global segmentation performance and the precision of boundaries across several public datasets, including FLARE2021, OIMHS, and SegTHOR. The experimental results underscore the critical importance of axial feature disentanglement in the enhancement of volumetric medical image segmentation, and it only introduces marginal computational overhead. Moreover, the modular design permits adaptation to other volumetric data via axis selection.

The significance of this work lies in the proposal of a concise yet effective construct to address the unique challenge of axial information variability in volumetric medical imaginga challenge that CNN-based models fail to address with their uniform treatment of all dimensions. Looking ahead, our research will focus on refining the proposed PaR module to enhance its efficiency and lightweight design. This continued research is expected to lead to further enhancements in the accuracy of volumetric segmentation, pushing the performance frontier of 3D conv from a fundamental perspective, and contributing to more precise diagnostic capabilities and treatment planning.

Contribution Statement

Xingru Huang, Jian Huang, and Yihao Guo contributed equally to this work.