Resizer Swin Transformer based Classification using sMRI for Alzheimer’s Disease

1. Introduction

Classification algorithms applied for medical imaging have been a major field of research over the past few decades. The recent success of deep learning techniques has inspired new research and development efforts to improve classification performance and develop novel models for various complex clinical tasks [1,2].

Alzheimer’s disease (AD) is an irreversible chronic neurodegenerative disease that progressively impairs cognitive and behavioral functions. Because of the high definition and noninvasive visualization, structural magnetic resonance imaging (sMRI) is one of the most common imaging techniques for AD identification in both clinical and research [3].

Therefore, deep learning algorithms have been applied increasingly frequently for AD classification using sMRI images since they can replace time and labor-consuming procedures such as feature extraction. In particular, there is no need for feature selection in the deep learning model, but it can also automatically learn sophisticated features by itself [4].

In this study, we propose a deep learning model for AD classification called Resizer Swin Transformer (RST) using sMRI images. The Resizer module creates a learnable image scaling, acquiring characteristics supporting Swin Transformer classification [5]. The cross-window connection realized by the moving window mechanism of the Swin Transformer and the patch merging provides multi-scale learning. Additionally, a convolutional neural network (CNN) model is employed to colorize the sMRI images to enable multi-channel learning [6].

Furthermore, a majority of current AD-related classification methods demand pre-processing procedures like calibration, skull stripping, and alignment to standard templates [7]. Particularly, skull stripping is regarded as a crucial pre-processing step that cannot be skipped since it can lessen the influence of irrelevant data on classification outcomes and simplify the computing process of the model [8]. However, insufficient skull stripping or over-processing caused by the toolkit might result in the loss of edge information, and the data must then be manually checked after processing, which takes time and significantly lowers the already limited data set relevant to AD.

2. Background and related work

2.1. CNN-based classification for AD

The main types of existing AD classification techniques based on the CNN model include the region of interest (ROI), voxel, patch, and attention mechanism. Noticeably, this categorization does not imply that the four approaches mentioned above are entirely distinct.

The ROI-based methods demand the pre-segmentation of brain regions based on prior knowledge, such as brain atlases. For instance, Wang et al. segmented the hippocampus area, one of the most AD-sensitive regions, and employed a dense convolutional neural network (Dense CNN) model to classify NC and AD [9]. On the other hand, one study tried each brain region to train a 3D-CNN ensemble model [10].

The voxel-based technique, which does not require any prior knowledge or laborious preparation, obtains features directly from sMRI images and fully exploits the global characteristics [11]. For instance, Hazarika et al. employed LeNet, AlexNet, VGG, DenseNet, and other models to classify AD while evaluating their effectiveness [12]. In addition, the classification accuracy may be increased by combining the CNN model with transfer learning [13] and data augmentation [14].

Others noticed that only localized brain areas in early AD patients exhibit minor structural abnormalities, which leads to a possibility that features obtained at the voxel or region level cannot contribute to AD identification completely. Patch, the intermediate level between voxel and region, has gained more attention. Flexibility in size and location, patch-based models improve classification accuracy and avoid laborious pre-processing procedures. However, the choice of patches significantly impacts the categorization outcomes. By using anatomical marker detectors, Liu et al. identified the patches discriminative to AD first, then trained a CNN model to learn from those patches [15]. The landmark-based deep multi-instance learning (LDMIL) system was introduced the following year to learn local patch information as well as global information from all patches [16].

Another frequently employed module recently is the attention mechanism, which is quite useful for pinpointing AD-sensitive areas. For example, Zhang et al. added the attention mechanism to the ResNet framework, which effectively enhances the gray matter features information and increases the accuracy of AD diagnosis [17].

2.2. Transformer-based classification for AD

One of the strongest deep learning models available today is Transformer, and a key component is the attention mechanism [18]. Despite the original purpose of the Transformer being for natural language processing (NLP), abundant studies have now demonstrated that Transformer-based models may reach superior performance in computer vision (vision Transformer, ViT). Because of the multi-headed attention mechanism of ViT, Li et al. integrated CNN to capture the relationships between distant brain areas. In this study, ViT received input from the feature maps extracted using the convolutional layer [19]. Another approach proposed by Jang et al. combined the ViT with a CNN structure that has an inductive bias, and the feature maps were generated using 3D ResNet. The 3D information provided by sMRI images can efficiently assess local aberrant characteristics associated with AD and link markers from multiplanar and multilayer slices to gather distant details [20]. Due to the inherent lack of inductive bias in the ViT-related study, a significant quantity of data is needed to train the model. Natural images share similar fundamental properties with brain sMRI images, including texture, edges, shape, etc. Hence, Lyu et al. applied the ImageNet [21] dataset to pre-train the ViT model using joint transfer learning first to address the limited brain imaging data [22].

2.3. Limitations of current research

1. ImageNet is a natural image dataset that each image contains three colors, although the coronal slices taken from 3D sMRI images are only gray-scale. Thereby, using sMRI images directly as the input of the ViT model implies delivering the same images into all three channels, which is a total waste of computational resources. In addition, the gray-scale image deviates significantly from the original RGB color image on each channel.

2. The majority of the current studies demand strict procedures for sMRI image preparation. Skull stripping is one of the crucial components. However, there are still some issues with employing SPM12 for skull stripping, such as partial stripping or loss of edge information due to over-processing. Therefore, the data must be visually checked and manually selected afterward. While eliminating the background (black region) of sMRI images has also become a critical step of pre-processing, the capacity of the deep learning models is hampered by the growing complexity of the pre-processing process.

3. The ViT model can only extract features at the same scale because its computational complexity is proportional to the square of the image size. Besides, it can only extract features on one scale at one time. However, the damaged brain areas of AD patients are typically subtle, making it possible to overlook some regionally specific details.

3. Methods

3.1. Slices options

The center slices of the sMRI images were chosen for this study because it contains most of the brain information. The classification effect of the coronal plane was found to be slightly inferior to that of the sagittal and axial planes in the comparison experiments. Still, the classification performance of the sagittal and axial planes was not significantly different. Therefore, 30 consecutive 2D slices of size 127×181 in the axial plane were employed in this study.

3.2. Model structure

3.2.1. Overlook of Resize Swin Transformer Network

Li et al. applied joint transfer learning approaches to classify AD by first using 3D gray matter volume images as input, extracting local information via ResNet to generate feature maps, and finally adding location encoding to the maps that were input to ViT [19]. Swin Transformer performs multi-scale learning through cross-window connection and patch merging, as well as decreases the computational complexity of ViT from square level to linear level [23]. Therefore, this study proposes the ReSwin Trans-former (RST), a novel network structure that combines CNN and Swin Transformer to realize the AD classification using sMRI images.

The RST model proposed in this study is shown in Figure 1. Axial plane sMRI slices X are the input to the RST. A CNN model is first utilized to map the gray-scale image from one channel to three channels to generate color images. The ReSizer module then scales the input image proportionally and accentuates information essential to classification according to the Swin Transformer during the training procedure. The images are separated into non-overlapping patches using the patch-splitting module. Each patch is regarded as a separate token. A linear embedding layer then maps the token to a size C channel as the Swin Transformer input. Finally, a SoftMax layer summarizes the AD and NC classification predictions from all the slices.

3.2.2. CNN module

sMRI scans only contain gray-scale information, generally known as a single-channel imaging technique. On the other side, various widely used CNN models, such as AlexNet [24], ResNet [25], and EfficientNet [26], are three-channel models. Hence, the single-channel gray-scale images are often repeated three times as input when the joint transfer learning algorithms are involved. Inspired by the colorization of lung MRI images to restore the color images of lungs observed by human eyes proposed in [27] and to address the issue of wasted computational resources, our study converts brain sMRI images to color images using the CNN networks to achieve cross-channel learning from single-channel to three-channels [6]. Figure 2 illustrates the CNN structure.

According to the given brightness Y on the gray-scale image, two chroma channels a and b were generated based on the CIELab color space. Then, the brightness-chroma color space is transformed to RGB color space using the image and OpenCV library. For a given brightness channel $X\in {\mathbb{R}}^{H\times W\times 1}$, we converted it to $Y\in {\mathbb{R}}^{H\times W\times 2}$ via the map $\widehat{Y}=\mathcal{F}\left(X\right)$, where $H$ and $W$ are the dimensions of the image. In addition, for a given $X$, its probability distribution was also obtained $\widehat{Z}=H\left(X\right)$. Where $\widehat{Z}\in {\left[\mathrm{0,1}\right]}^{H_x{W}_{x}Q}$ and $Q$ represents the number of output spaces from the channel a and b. The following equation compares the true value with $\widehat{Z}$ through polynomial cross-entropy loss $L_{cl}(·,·)$:

$$L_{cl}\left(\widehat{Z},Z\right)=-\sum _{h,w}v\left(Z_{h,w}\right)\sum _q{Z}_{h,w,q}\mathit{log}\left({\widehat{Z}}_{h,w,q}\right)$$

In particular, the real color $Y$ is converted into a vector $Z$ via $Z={\mathcal{g}}_{gt}^{-1}\left(Y\right)$, $v\left(·\right)$ is the weight that measures the rarity of the color class and thus rebalances the loss. Finally, the probability distribution $\widehat{Z}$ is mapped to color values by the function $\widehat{Y}=\mathcal{g}\left(\widehat{Z}\right)$.

3.2.3. Resizer Module

In the field of deep learning for image processing, the input image size is typically scaled to 224×224, and both training and inference are carried out at that resolution. Image scaling often uses both bilinear and trilinear interpolation currently. In actuality, this modification does not improve the image of the network, which somewhat even reduces the performance of the model [28]. Therefore, this experiment substitutes the learnable Resizer module for the original linear interpolation after CNN transfers the single-channel to the three-channel color space. Resizer seeks to significantly improve classification performance by learning attributes favorable to Swin categorization by collaborative training with the backbone network, in contrast to other approaches that resize images to enhance human eye perception. The Resizer module is shown in Figure 3.

3.2.4. Swin Transformer

Swin Transformer utilizes within-window calculation of self-attention to increase modeling efficiency. Beginning in the top left corner, the window evenly and non-overlappingly divides the image into sections. Assuming that there are $M$ × $M$ patches in a window, the next module utilizes a different window than the previous layer and moves ($⌊\frac{\mathit{M}}{2}⌋$, $⌊\frac{\mathit{M}}{2}⌋$) patch from the original window when the window-based self-attention module completes its computation. The calculation process of the Swin Transformer block is as follows:

$$\begin{gathered} {\widehat{\mathit{Z}}}^{\mathit{l}}=\mathrm{W}-\mathrm{MSA} \left(\mathit{L}\mathit{N}\left({\mathit{Z}}^{\mathit{l}-1}\right)\right)+\left({\mathit{Z}}^{\mathit{l}-1}\right), \\ {\mathit{Z}}^{\mathit{l}}=\mathrm{MLP} \left(\mathit{L}\mathit{N}\left({\widehat{\mathit{Z}}}^{\mathit{l}}\right)\right)+{\widehat{\mathit{Z}}}^{\mathit{l}}, \\ {\widehat{\mathit{Z}}}^{\mathit{l}+1}=\mathrm{SW}-\mathrm{MSA} \left(\mathit{L}\mathit{N}\left({\mathit{Z}}^{\mathit{l}}\right)\right)+{\mathit{Z}}^{\mathit{l}}, \\ {\mathit{Z}}^{\mathit{l}+1}=MLP\left(\mathit{L}\mathit{N}\left({\widehat{\mathit{Z}}}^{\mathit{l}+1}\right)\right)+{\widehat{\mathit{Z}}}^{\mathit{l}+1} \end{gathered}$$

Where ${\widehat{\mathit{Z}}}^{\mathit{l}}$ and ${\mathit{Z}}^{\mathit{l}}$ are multi-headed self-attentive based on the window and shifted window, respectively, and characterized by the output of the multilayer perceptron (MLP) module. Window movement enables the reciprocal learning of patches between several windows, so achieving the goal of global modeling. Figure 4 depicts the model of the Swin Transformer structure as well as the calculating procedure of the block.

Windows multi-head self-attention (W-MSA) adds relative position bias $\mathit{B}\in {\mathbb{R}}^{{\mathit{M}}^2\times {\mathit{M}}^2}$ for each head in the calculation of multi-headed self-attention, as follows:

$$Attention(\mathit{Q},\mathit{K},\mathit{V})=SoftMax\left(\mathit{Q}\frac{{\mathit{K}}^{\mathit{T}}}{\sqrt{\mathit{d}}}+\mathit{B}\right)\mathit{V}$$

Where, $\mathit{Q}$, $\mathit{K}$, $\mathit{V}\in {\mathbb{R}}^{{\mathit{M}}^2\times \mathit{d}}$ are the query, key, and value matrices, respectively. $\mathit{d}$ is the dimension of the key. The number of patches in the window is ${\mathit{M}}^2$. Given that the relative location is between $[-\mathit{M}+1,\mathit{M}-1]$, the bias matrix $\widehat{\mathit{B}}\in {\mathbb{R}}^{\left(2\mathit{M}-1\right)\times \left(2\mathit{M}-1\right)}$ is set. $\widehat{\mathit{B}}$ from whence the values for $\mathit{B}$ are derived. $\mathit{Q}$, $\mathit{K}$ and $\mathit{V}$ are calculated from ${\mathit{W}}_{\mathit{Q}}$, ${\mathit{W}}_{\mathit{K}}$, and ${\mathit{W}}_{\mathit{v}}$ by applying a linear transformation of ${\mathit{Z}}^{\mathit{l}-1}$.

On the other hand, shifted windows multi-head self-attention (SW-MSA) uses the circular window movement rule to compute the multi-headed self-attentiveness.

4. Evaluation

4.1. Introduction of the datasets

The Australian Imaging, Biomarker & Lifestyle (AIBL, aibl.csiro.au) [29] and the Alzheimer’s Disease Neuroimaging Initiative (ADNI, adni.loni.usc.edu) [30] were used in this study. The National Institutes of Health and the National Institute on Aging provided funds to establish the ADNI, the premier data center for AD research. Both datasets gather sMRI, fMRI (functional MRI), and positron emission computed tomography (PET) from AD and NC participants.

The data used in the study are collected from an MRI scanner built by MRI manufacturer SIEMENS. The slice thickness is 1.2 mm; the field strength is 3.0 Tesla. All the sMRI images were downloaded from ADNI-GO, ADNI1, and AIBL. The standard sMRI image pre-processing was conducted using SPM12 (fil.ion.ucl.ac.uk) on Matlab (R2022a), including format conversion, AC-PC correction, non-parametric non-uniform intensity normalization (N3), and alignment to the MNI standard template. The reconstructed images are 181×217×217, the voxel size is 1×1×1 mm³, and normalized intensity values are in the range of [0,1]. Table 1 displays the demographic information of the datasets.

The ADNI and AIBL data sets are openly accessible, but access to the information still needs official authorization. Additionally, no data may be shared without consent, and only approved researchers may utilize it for study.

4.2. Training Setup

This experiment used the PyTorch deep learning framework to construct the proposed network model. Two NVIDIA 3080 TI GPUs were implemented on a server for training the classification task. The model is first transferred to the sMRI dataset after being pre-trained on the ImageNet-1K dataset. Since the Resizer module may shrink the images to eliminate a tiny amount of incorrect information (dark parts of sMRI images), cropping the input images is unnecessary. With a learning rate of 5×10-5, a patch size of 4×4, a batch size of 16, the number of epochs set to 50, and using cross-entropy loss as a loss function, all networks were trained using the Adam optimizer. The results of this study are based on the test set, with epoch.

4.3. Experimental Results

The classification performance was assessed using the following four common specificity indicators: accuracy (ACC), sensitivity (SEN), specificity (SPE), and precision (PRE). We compared our approach to every type of recent study in this area to demonstrate its superiority, and the findings are displayed in Table 2. Our strategy obviously outperforms earlier research in every index.

4.4. Ablation Experiments

We carried out several ablation experiments to maximize the experimental outcomes.

Experiment 1: We tested various batch sizes and learning rates to determine the settings that would produce the best experimental results. The experimental outcomes are displayed in Figure 5. Most curves are generally stable at epoch=50, yet some still vary remarkably. It reaches the most stable when the batch size is 16, and the learning rate is 5×10-5. Therefore, the batch size is set to 16, and the learning rate is set at 5×10-5, considering parameters like processing speed and classification performance. set at 5×10-5, considering parameters like processing speed and classification performance.

Experiment 2: We used sMRI slices in three different orientations. In Table 3, the experimental findings are displayed. The worst result was using the coronal plane slice without the skull stripping pre-processing procedure, which differed greatly from the findings of the other two orientations. Since most experiments were centered on the axial plane and the differences between the experimental outcomes in the axial and sagittal planes were not particularly apparent, axial plane slices were employed in this experiment.

Experiment 3: Tests were run to show how well the various parts of our proposed RST model worked together. The test results are listed in Table 4. The RST model, as implemented by CNN, performs noticeably better than the other experiments, as shown in the table. In contrast, the step of skull stripping only slightly improves the performance of the RST model and is not proportional to the additional expense it incurs.

Experiment 4: Investigations were carried out on several datasets to confirm the robustness of our proposed model. The outcomes of our trials are shown in Table 5. The table illustrates how our model performs better for various datasets. The analysis may be caused by the unevenness of the datasets and the image discrepancies between the datasets because the findings of AIBL, on the other hand, are worse than those of ADNI.

5. Conclusion and future work

This paper presents a Resizer Swin Transformer architecture that combines cross-channel learning with CNN and extracts features from two-dimensional axial plane slices of brain MRI. The analysis and visualization of the experimental data demonstrated the accuracy with which RST architecture can accomplish the categorization of AD as well as its strong adaptability to various data sets. This experiment does still have some flaws, though. The RST model has comparatively high model parameters, and cross-channel learning in conjunction with CNN enhances the classification performance using the model. Meanwhile, we discovered that the accuracy of the experimental findings is lower than 85% when the model trained on ADNI is directly assessed by AIBL. The model may undergo further development to overcome these problems and provide a lighter and more scalable model.