Class-consistent Contrastive Learning Driven Cross-dimensional Transformer for 3D Medical Image Classification

1 Introduction

Transformers [Vaswani et al., 2017] have achieved performance breakthroughs in capturing long-range token dependencies and global feature extraction. However, it meets three challenges (Figure.1(a)) when applied to 3D medical image data: 1) The flattening operation of Transformer destroys the 3D spatial structure of data, making it difficult to capture spatial structure correlations and brings irrelevant information and ambiguity to feature representation learning for 3D medical image data. 2) Huge computational complexity and memory consumption due to the computational complexity and memory consumption of Transformer is quadratic to the spatial size. 3) Challenge in learning the feature representation with limited positive and negative samples, particularly for discriminative representation learning, due to transformerbased models are only effective when large training datasets are available. Although many Transformer methods [Fu et al., 2024; Guo et al., 2021; Chen et al., 2021b; Zhu et al., 2022; Zhu et al., 2023a] have achieved encouraging performance in 3D data analysis, the aforementioned challenges remain unresolved and limit the effectiveness of Transformer on 3D medical image analysis.

Existing 3D Transformers are also unable to overcome the aforementioned challenges. For example, the 3D Transformer used in [Zhou et al., 2023] employs Local Volume-based Multi-head Self-attention (LV-MSA) and Global Volume-based Multi-head Self-attention (GV-MSA) to construct feature pyramids for learning representations on 3D volumes. Although LV-MSA reduces computational resources, it can only extract local information. Furthermore, GV-MSA faces the challenge of high computational complexity while learning global feature representations. The Multiplane and Multi-slice Transformer [Jang and Hwang, 2022] extracts 3D feature representations by constructing attention relationships among multi-plane (axial, coronal, and sagittal) and multi-slice images, but it is unable to avoid the disruption of 3D spatial structure. Additionally, the Multi-plane and Multi-slice Transformer, built upon the Transformer architecture, also encounters the challenge of high computational complexity for 3D input. Despite the various strategies employed by existing methods to integrate Transformers with 3D data, these 3D Transformers still fail to overcome the challenges inherent in 3D medical data. Hence, there is an urgent need for a computationally efficient, 3D structurally aware Transformer in the field of 3D medical image analysis.

We propose a novel Cross-dimensional Transformer (CdTransformer) that consists of 1) a Cross-dimensional Attention Module (CAM) and 2) an Inter-dimensional Feedforward Network (IdFN) to exploit long-range token dependencies among 3D spatial pixels and extract the global discriminative representations to overcome above challenges and further improve the ability of Transformer on 3D medical image analysis (Figure.1(a)). Specifically, CAM introduces a novel cross-dimensional attention that promotes information sharing and fusion between different dimensions for building connections between 3D spatial pixels, addressing the challenge that Transformer lacking reasonable spatial structure correlation establishes mechanisms. More significantly, CAM converts the quadratic memory consumption of Transformer to linear memory consumption for addressing high complexity and memory consumption problems. IdFN adopts a novel Inter-dimensional attention to activate and fuse the effective features from both spatial and depth views, overcoming the challenge of Feed-forward Network ignoring the depth dimension information.

Apart from architectural novelties, a novel contrastive learning, named Class-consistent Contrastive Learning (CcCL) (Figure.1(b)), is proposed to overcome the limitation that Transformer-based methods are ineffective when meeting limited training datasets [Chen et al., 2020]. Specifically, CcCL innovatively exploits the inter-class feature and intraclass feature from the slice-distorted samples by taking full advantage of Transformer's strengths in long-range dependencies learning. By maximizing the consistent of positive slice-distorted pairs while minimizing the inconsistency of negative slice-distorted pairs, CcCL enables Transformer to learn class-aware discriminative features under limited training datasets. We conduct comprehensive experiments and demonstrate the significance of CcCL in boosting discriminative representation learning with limited positive and negative samples.

Our proposed CcCL and CdTransformer were validated on six 3D medical image classification tasks including five 3D MedMNIST datasets [Yang et al., 2023] and one well-known LIDC-IDRI dataset [Kuan et al., 2017]. Extensive experimental results demonstrate that CdTransformer outperforms state-of-the-art CNNs and Transformers on 3D medical image classification.

Our main contributions include:

• A cross-dimensional Transformer with linear computational complexity and memory consumption has been established for 3D medical image analysis, which addresses three limitations of Transformer in 3D medical image analysis by effectively capturing spatial structure correlations while mitigating the challenges of high computational complexity and memory consumption.
• Our class-consistent contrastive learning innovatively takes full advantage of the inter-class feature and intraclass feature from the slice-distorted samples to boost the effectiveness of the Transformer in learning discriminative representation.
• Our cross-dimensional attention module innovatively establishes the 3D spatial structure correlation in 3D medical image data, where Transformer cannot.
• Our inter-dimensional feed-forward network enables advanced aggregating spatial and depth contexts from both spatial and depth views, which addresses the limitation that the feed-forward Network of Transformer ignores the depth dimension information.
• Experimental results on six 3D medical image datasets demonstrate that our cross-dimensional Transformer serves as a generalized module, exhibiting remarkable capabilities in 3D medical data analysis.

2 RelatedWork

Many authors try to use Transformer [Dosovitskiy et al., 2020] for 3Dmedical image analysis [Zhao et al., 2023; Zhu et al., 2023b; Qin et al., 2022]. For example, Xie et al. [Xie et al., 2021] proposed a hybrid model of CNN Transformer, namely CoTr, for 3D medical image segmentation. Inside the model, the deformable Transformer (DeTrans) that employs the deformable selfattention mechanism is introduced to reduce the computational and spatial complexities of modelling the longrange dependency on multi-scale and high-resolution feature maps. Wang et al. [Wang et al., 2022] proposed an fNIRS classification network based on Transformer, named fNIRS-T, for functional near-infrared spectroscopy classification. To explore the spatial-level and channel-level representation of fNIRS signals, two Transformers, fNIRS Spatial-level Transformer (fNIRSST) and fNIRS Channel-level Transformer (fNIRS-CT), are employed inside model. fNIRS-ST can extract local brain area features and fNIRS-CT for the hemodynamic response of a single channel. Hatamizadeh et al. [Hatamizadeh et al., 2022] proposed a novel architecture, dubbed as UNEt TRansformers (UNETR), that utilizes a transformer as the encoder to learn sequence representations of the input volume and effectively capture the global multi-scale information.

3 Method

3.1 Cross-dimensional Transformer

Cross-dimensional Transformer (CdTransformer) are advanced in 3D medical data analysis through two key innovations: 1) Cross-dimensional Attention Module (CAM); 2) Inter-dimensional Feed-forward Network (IdFN). The key design elements of CAM are cross-dimensional attention and depth attention. The cross-dimensional attention promotes information exchange and fusion between each dimension for building the spatial structure relationship. Depth attention is attached behind CAM which exploits the correlations along depth dimensions. Compared with existing Transformers, CAM has the advantage of capturing spatial structure correlation and converting the quadratic complexity into linear complexity and significantly reduces the computational memory consumption. IdFN is designed for better aggregating and transforming spatial and depth contexts through a novel architecture. IdFN learns local image features from spatially neighboring pixels and exploits the spatial and depthsensitive features via the attention mechanism, which overcomes the challenge that the feed-forward network of Transformer ignores depth contexts. The two modules complement each other and investigate the long-range dependencies among both spatial and depth dimensions with linear memory consumption.

1) Cross-dimensional Attention Module (CAM)

CAM (Figure.2(b)) can investigate the long-range dependencies among pixel's spatial structure and model the global spatial context with linear memory and computation consumption. Specifically, CAM utilizes three independent crossdimensional attention to learn global spatial contexts by modeling dimensional correlation rather than the spatial dimension, which enables CAM to build connections among dimensions and absorb the complementary information from spatial structure.

Formally, given a 3D input feature map X ∈ RC×D×H×W , where C is the number of channels, D, H, and W represent the spatial dimension. The CdTransformer layer first utilizes three convolution layers to project X into three sequences [Q;K; V ] ∈ Rc×D×H×W , where c is the hidden dimension of the input sequences, where Q is the input Query sequence, K and V are the input Key,Value sequence. Afterward, different from the conventional self-attention that computes spatial attention maps (its memory complexity brought by the key-query dot product interaction is quadratic with the spatial resolution of 3D input data), cross-dimensional attention performs attention calculation on three dimension {D,H,W} separately by three independent self-attention modules.

To ensures that the contextualized global relationships between pixels are exploited, the pixel-wise aggregation of cross-dimensional context is fused via an efficient crossdimensional fusion layer. Formally,

Figure 3: The conception of class-consistent contrastive learning, which pushes inter-class features apart and pushes intra-class features close.

where QD ∈ RD×(CHW) , KD ∈ RD×(CHW) , VD ∈ RD×(CHW) . d is a learnable scaling parameter to control the magnitude of the dot product. ˆ Yx is the reshape of Yx. And the size of attention map in YD, YH, YW is RD×D , RH×H , and RW×W , respectively.

To overcome the challenge that Transformer only builds the dependencies along spatial dimensions and ignores the relation between depth dimensions, a depth attention module is attached behind the attention module. Formally, three 3D convolution layers with 1 × 1 × 1 kernel size first project YS = [Y ′D ∥Y ′H ∥Y ′W ] into three sequences [QC;KC; VC] ∈ RDHW×C . Afterward, the depth attention is computed

Thus, the final output of CAM is calculated as YC.

2) Inter-dimensional Feed-forward Network (IdFN)

IdFN (Figure. 2(c)) is designed for boosting spatial and depth contexts aggregating and transforming. Different from the regular feed-forward network (FFN) which consists of two convolution layers, IdFN consists of three parallel streams, one of which is used for learning local image features from spatially neighboring pixels by two stacked convolution layers. The other two streams exploit the spatial and depthsensitive features via attention mechanism. Formally, given an input feature map X ∈ RC×D×H×W , IdFN is formulated as:

where Ws is spatial-wise attention weight, which is generated by a convolutional operation ws ∈ RC×1×1×1 followed by a sigmoid function.

Wd is depth-wise attention weight, which is generated by a global average pooling operation and two convolutional operations. The global average pooling first performed on the input X to generate the depth-wise statistics z ∈ RC . Afterward, a simple gating mechanism with a sigmoid activation is performed on the depth-wise statistics to compute depth-wise attention weights, which is achieved via two convolutional operations and can be formulated as:

where w1 and w2 denote two convolutional operations.

3.2 Class-consistent Contrastive Learning

CcCL (Figure. 3) innovatively exploits the inter-class feature and intra-class feature from the slice-distorted samples by taking full advantage of the Transformer's strengths in long-range dependencies, which boosts Transformer in discriminative feature representation learning and addresses the challenge of Transformer being ineffective in limited training datasets. Specifically, given one sample xi, the corresponding positive sample x′i can be obtained by randomly selecting one dimension and randomly changing the order of slices as shown in Figure.1(b). As the sample and generated sample {xi, x′i } belong to one category, the category-aware features should be consistent. CcCL enables pushing inter-category feature apart and pulling intra-category feature close. With the guidance of CcCL, Transformer can learn the categoryaware discriminate features. The formulation of CcCL can be:

where zx = Proj(Enc(x)), the × symbol denotes the inner product, τ is a scalar temperature parameter, I is the index of an arbitrary sample, A represents sample from different categories. Remarkable, our CcCL enables producing unlimited positive samples, which addresses the challenge of sample choosing in contrastive learning and enhances Transformer in feature representation learning by creatively capturing category-aware correlated features.

4 Experiment

4.1 Datasets and Implementation Details

1) 3D MedMNIST Dataset: Five 3D standardized medical datasets, including AdrenalMNIST3D, NoduleMNIST3D, VesselMNIST3D, SynapseMNIST3D, FractureMNIST3D, with diverse classification tasks (binary classes, multi-classes, and multi-label) of 3D MedMNIST dataset [Yang et al., 2023] are used in this paper. These 3D medical data as the benchmark in medical image classification tasks. AdrenalMNIST3D consists of 1584 left and right adrenal glands CT images that are used to distinguish normal adrenal glands from adrenal masses. NoduleMNIST3D consists of 1849 lung nodule chest CT images used to classify two types of malignancy levels. VesselMNIST3D consists of 1909 brain vessels collected by reconstructing MRI images, which are used for binary classification of the aneurysm and healthy vessel segments. The SynapseMNIST3D is a new 3D volume dataset to classify whether a synapse is excitatory or inhibitory.

2) LIDC-IDRI Dataset: The LIDC-IDRI dataset [Kuan et al., 2017] contains 1018 computed tomography (CT) scans from 1010 patients altogether collated from seven academic centers across the United States (US). The slice thicknesses of the CT scans range from 0.45 to 5.0 mm. Each CT scan was annotated by four experienced thoracic radiologists. In this study, we follow previous works [Al-Shabi et al., 2019a; Al-Shabi et al., 2022], the annotated nodules of size smaller than 3mm, slice spacing inconsistent, or missing slices are removed. Finally, there are a total of 837 nodules left, of which 442 nodules are benign and 395 nodules are malignant.

1. Implementation Details: The 3D medical data classifi-

cation network consists of 4 stages. CdTransformer is used in stage-1 to 4, the number of CdTransformer blocks is [2, 3, 3, 4], the attention heads number in CAM are [2, 4, 8, 16]. The channel expansion factor is 4. Between each stage, a 3D convolution layer with 4×4×4 kernel size and step size 2 is used for downsampling. Behind the CdTransformer blocks, a classification head, and a projection head are attached.The network is pre-trained under the supervision of LCcCL. For the LIDC-IDRI dataset, all the lung nodules were cropped around the centers of the lung nodules with size 32 × 32 × 32 pixels and normalized by the z-score standardization method (mean value is -400, std value is 750). The randomly adding Gaussian noise, horizontal flip, vertical flip, and z-axis flip are utilized for data augmentation. The framework is implemented on Pytorch and four A4000 GPUs with 16 GB memory, Adam optimizer with a minibatch size of 32 was applied for optimization. The learning rate and weight decay were set to 1e-4 and 0.01, respectively.

4.2 Comparison with State-of-the-art Methods

The quantitative results in Table 1 demonstrate that CdTransformer achieved the best quantitative results on 3D MedM- NIST dataset. Specifically, compared with state-of-the-art methods, our CdTransformer achieved the best AUC, and ACC on four datasets, which demonstrates that CdTransformer has an advantage in capturing the category-aware features through establising long-range pixel dependencies among three dimensions. Meanwhile, as shown in Table 1, with the assistance of CcCL, our CdTransformer further improved the performance, which also proves the superiority of our CcCL in exploiting discriminative features.

The results on LIDC-IDRI database are shown in Table 2. As it can be seen from Table 2, our CdTransformer achieves the highest Accuracy, Precision, and F1-score compared with state-of-the-art methods, which represents our model has distinct advantages over compared methods in 3D nodule representation learning. Meanwhile, those results confirmed that 1) absorbing the global spatial and depth contextual information could further improve the performance; 2) building the dependencies correlations along both spatial and depth dimensions could assist the model to capture global and significant contextual information; 3) CcCL has advantages in improving the performance of CdTransformer.

4.3 Comparison with State-of-the-art Transformer

We also compared CdTransformer with with state-of-theart Transformers, inclduing Transformer [Vaswani et al., 2017], MobileViT [Mehta and Rastegari, 2021], Swin Transformer [Liu et al., 2021], CrossViT [Chen et al., 2021a], SSAN [Guo et al., 2021], MMTransformer [Jang and Hwang, 2022], and 3D Transformer [Zhou et al., 2023] in 3D medical image classification from the aspects of effectiveness, memory consumption and computational complexity.

Experimental results demonstrate CdTransformer outperforms state-of-art Transformers. From Table. 3, we can notice that the Transformer (ViT) has a poor performance. The major reason is that the vanilla self-attention module cannot effectively and reasonably model long-range dependencies among both spatial and depth dimensions of 3D medical image data. An interesting phenomenon is that SSAN obtained the second worst results. Although SSAN could investigate the relationship between spatial and depth correlations, it lacks the ability to exploit the correlations between dimensions. MobileViT has less computational complexity and achieved the best performance. These results demonstrate the necessity of reducing memory consumption and computational complexity of Transformer in 3D medical image classification and point to three aspects that Transformer needs to face: 1) The less memory consumption and computational complexity are advanced in improving the performance of the Transformer. 2) Effective and reasonable modeling of longrange dependencies among both spatial and depth dimensions are significant to 3Dmedical image data analysis. 3) It makes perfect sense to boost the effectiveness of the Transformer in learning long-range dependencies and discriminative representation.

CdTransformer significantly reduces computational complexity and memory consumption. As show in Table. 4, given an input data X ∈ RC×D×H×W , the per-layer complexity of the vanilla self-attention module is 2(DHW)2C +

4DHWC 2 and the attention metric size is DHW ×DHW. The per-layer complexity of three dimensions inside the CdTransformer is 6HWCD2 , 6DWCH 2 , 6DHCW 2 and the attention metric size is D×D, H ×H, and W ×W, respectively. The complexity of depth dimension is 6DHWC 2 , the attention metric size is C × C. These results demonstrate that CdTransformer reduces computational complexity and memory consumption. What's more, while both 3D Transformer [Zhou et al., 2023] and MMTransformer [Jang and Hwang, 2022] are tailored for 3D medical image datasets, they adopt different approaches for 3D feature extraction. In the case of 3D Transformer, it employs two types of selfattention, leading to heightened computational complexity and increased memory consumption. On the other hand, MMTransformer focuses on extracting 3D structure information by combining insights from three distinct views of the 3D data. However, this fusion process may compromise the inherent 3D structure, and similar to 3D Transformer, the selfattention mechanism in MMTransformer comes with comparable computational demands and memory

4.4 Effect Analysis using CdTransformer

The effectiveness of each part of the CdTransformer is also demonstrated. We analyze the influence of each part on the classification results by ablation studies: 1) 3D ResNet (baseline): The classification is achieved by directly using 3D ResNet without attention mechanism as the baseline. 2) Baseline + Transformer: The Transformer (ViT) is used in each stage of 3D ResNet for capturing long-range pixel dependencies and global feature extracting. 3) Baseline + CAM + FFN: Different from 2), here the vanilla self-attention module inside Transformer is replaced by a dimension-mutual attention module. Meanwhile, to directly prove that the dimension-mutual attention module has the advancement in investigating the long-range dependencies among pixels and modeling the global spatial context, we use the normal feedforward network (FFN) for feature transformation. 4) Baseline + CAM + IdFN: Based on setting 3), we replaced the FFN module with our proposed Inter-dimensional feedforward network (IdFN). By comparing it with setting 3), we can evaluate the effectiveness of IdFN. 5) CdTransformer: Our proposed Cross-dimensional Transformer network with CcCL.

Table 5 lists the classification performance of the methods described above. From the Table 5, we can notice that the innovations of CAM and IdFN bring significant enhancements, and the baseline has a poor performance. The major reason is that the typical convolution block construed by stacked convolution layers cannot model long-range dependencies among both spatial and depth dimensions of nodule. The spatial attention module assists the baseline model to obtain a 4.55%, 4.82%, 8.18% and 4.39% improvement on AUC, ACC, Precision, and F1-Score, respectively. When utilizing CAM, the AUC, ACC, Precision, Sensitivuty, and F1Score achieve 4.69%, 4.82%, 6.28%, 2.56% and 4.61%, respectively. Those improvements proved the effectiveness of the CAM in modeling long-range dependencies among pixels and capturing global and significant contextual information. What's more, when using CAM and IdFN simultaneously, the AUC, ACC, Precision, and F1-Score are increased by 5.39%, 6.02%, 15.32%, and 5.14%, which reveals the advantage of further aggregating and transforming spatial and depth contexts. The above extensive experiments with promising results reveal the power of the CAM and IdFN and its significance in improving the performance of the Transformer.

4.5 Significance of CcCL

The significance of CcCL is proved in Table 1, Table 2 and Table 5. The CcCL assists CdTransformer to obtain an improvement on AUC, ACC, Precision, and F1-Score on the six 3D medical datasets. The CcCL assists CdTransformer in obtaining a 4.55%, 4.82%, 8.18%, and 4.39% improvement on AUC, ACC, Precision, and F1-Score on LIDC-IDRI dataset. What's more, in the five 3D MedMNIST, CcCL also improved the performance of model. All of experimental results demonstrate that CcCL improves the performance of Transformer.

5 Conclusion

Our CdTransformer addressed the limitation of Transformer in 3D medical image classification through a novel crossdimensional attention module and an inter-dimensional feedforward network. Cross-dimensional attention module and inter-dimensional feed-forward network modules promote information exchange and fusion between dimensions for modeling the long-range dependencies among three dimensions with linear memory consumption and computational complexity. Additionally, our class-consistent contrastive learning boosts Transformer in learning feature representation. Extensive experimental results on six 3D medical image data demonstrate that CdTransformer outperforms the state-ofthe-art CNNs and Transformers in 3D medical image classification. And our class-consistent contrastive learning can significantly improve the performance of Transformer.