Masked LoGoNet: Fast and Accurate 3D Image Analysis for Medical Domain

Figure 1(b) illustrates an overview of our feature extractor (ULKANet), which is a U-shaped model and has an encoder and a decoder. The encoder consists of a sequence of blocks. Each block consists of a repeating sequence of three components: a patch embedding component, a chain of transformer-like modules that employ LKA ($L_{i}$ modules for $i^{th}$ block of the encoder), and a layer normalization component. For conciseness, Figure 1(b) only shows the top-level blocks, while a detailed illustration of the model architecture and inner components is provided in the appendix section 7.

The Patch Embedding component plays a crucial role in the processing of input data within the encoder block, transforming the input into a tensor that is subsequently passed to the next component in the sequence. Throughout the current encoder block, the dimension of the embedding vectors remains constant, denoted as $dim$.

The mathematical representation of the projection operation is defined as follows:

To enable our model to efficiently extract complex feature dependencies that are often present in medical images, we opt for using transformer modules. However, instead of using the regular transformers with self-attention that is slow and needs more memory (Singh et al., 2020), we use LKA (Guo et al., 2022) in the attention layer. This type of attention mechanism decomposes large convolution kernels into spatial dependencies and channel convolutions. It enables our model to go deeper and remain memory efficient. The attention module is implemented as follows:

where $X$ is the input tensor and ChConv3D is a depth-wise convolution operating on a single channel. DiConv3D is a dilated depth-wise convolution to broaden the receptive field and to enable the extraction of long-range dependencies. The point-wise convolution $\textit{Conv3D}_{1\times 1}$ is applied to aggregate the information across the channels. The final activations are obtained as follows:

where $\odot$ is the element-wise product. The remaining components of the transformer block follow the conventional structure of typical transformers.

The decoder in our model aims to restore the spatial resolution of the input using a sequence of blocks (green circles in Figure 1(b)). Each decoder block consists of a chain of three convolution modules followed by an upsampling operation. The convolution modules are responsible for volumetric convolution operation. They consist of a Conv3D layer and a batch normalization layer, followed by a LeakyReLU activation function. The upsampling operation scales the resolution by a factor of two. As we stated earlier, a second larger illustration of our architecture that shows the inner modules can be found in the appendix section 7.

For each individual block in the encoder, the decoder has one corresponding block. There is also an additional decoder block in the bottleneck layer, as shown in Figure 1(b). The input to each decoder block is supplied by the block in the previous layer and also the corresponding encoder block through a skip connection. In order to enhance the reconstruction of input, we use the skip connections to facilitate the transfer of high-level features (Yosinski et al., 2014) to the layers that are responsible for the reconstruction task.

One of the difficulties in medical image segmentation is the presence of organs that have complex shapes. For instance, the human gallbladder and spleen have an elongated structure. Hence, to achieve satisfactory performance in the segmentation task, the model should be able to detect and extract relevant features in multiple regions of the input images, heavily relying on global features. On the other hand, this organ has irregular corners. This characteristic requires the model to be able to detect local features in multiple regions of the input. While increasing the model capacity by adding more layers, and also composing larger training sets, will potentially enable the model to automatically learn these regularities, this will likely increase costs during both the deployment and development stages.

To reduce the burden of automatic feature mining and, consequently, to reduce the costs, we propose to impose an inductive bias (Mitchell, 1997) on the feature extraction process. We propose to have two feature extractors in parallel, one focusing on the global scale and another one focusing on the local scale–as shown in Figure 1(a). The global module is able to extract long-range dependencies due to access to the original data cube. On the other hand, the local module focuses on short-range dependencies. This is accomplished by partitioning the input cube into smaller ones, allowing for a more focused analysis and resulting in finer-grained features.

To implement our idea, we use one instantiation of ULKANet in the global module, and a sequence of $N$ instantiations of ULKANet in the local module. In the analysis section, we show that while using only one ULKANet can reduce the model size and speed up inference, it will also significantly deteriorate prediction accuracy. Additionally, we show that alternative strategies, used in comparable models, are either slower or achieve lower prediction accuracy. To prepare the data for the local module, the input 3D image is split into $N$ smaller cubes of size $B\times B\times B$. Given an image of size $S\times H\times W$, the value of $B$ is obtained by $B=\sqrt[3]{\frac{S\times H\times W}{N}}$.

To reconstruct the input data cube, the outputs of the local module are concatenated, as shown in Figure 1(a). In order to aggregate the outputs of the global and local modules, we use an element-wise summation operator. The resulting tensor is expected to represent both global and local range dependencies.

Before fine-tuning our model on labeled data, we utilize a multi-task pre-training technique to put the model weights in a favorable state. This self-supervised approach allows the model to learn general information from 3D medical images, without the necessity of ground-truth labels.

Pre-training of our model is done in three stages. First, we methodically mask certain regions of the input images. In this stage, the goal is to capture long-range and short-range feature dependencies. Second, we generate pseudo-labels for the masked images. The model later learns to generalize to unseen cases by predicting the pseudo-labels of the masked data. Finally, the masked images, along with their pseudo-labels, are used to pre-train the model. Below we explain each step.

Table 1 compares our model to the baseline methods in terms of inference time (FLOPs) and the number of trainable parameters in the BTCV dataset. We see that our model has the lowest inference time after SegResNetVAE. Tables 2 and 3 compare the accuracy of our method to the baselines. We observe that the performance of SegResNetVAE is significantly lower than that of ours. Taking into account both the inference speed and the prediction accuracy, our model seamlessly ranks first among all the models. Appendix section 8 reports more experiments about the training time, test time, and memory consumption.

Table 1 shows that our model is considered an average-sized network. One noteworthy observation is that in some cases, e.g., nnUNet or DiNTS Instance, even though the number of trainable parameters is on a par or smaller than ours, their inference speed is substantially slower. Tables 2 and 3 show that our model exhibits superior performance compared to the baseline methods. Specifically, when evaluating our proposed model without pre-training, it outperforms the baselines across $13$ out of $19$ tasks. Furthermore, incorporating our pre-training strategy into LoGoNet enhances its performance even further, surpassing the baselines in $18$ out of $19$ tasks. These findings underscore the effectiveness and versatility of our approach in tackling a diverse range of tasks with notable efficacy.

In the BTCV dataset, LoGoNet outperforms the top three baseline models on average by $2.7\%$, $3.0\%$, and $3.2\%$, respectively. Regarding the inference time, our model outperforms the top three models by $17.6\%$, $14.8\%$, and $118.2\%$, respectively.

In this section, we demonstrate the properties of our model from multiple aspects. Specifically, we report a qualitative comparison between our model and the best baseline model, evaluate our strategy for extracting local and global features, evaluate our pre-training approach, show the impact of model size on performance, analyze the hyper-parameter sensitivity of our model, and finally, report an ablation study on the steps in our pre-training method. The experiments in this section are carried out in the BTCV dataset unless stated otherwise.

We begin by qualitatively inspecting our model. Figure 3 compares the output of LoGoNet to the best performing baseline model in BTCV dataset, i.e., DiNTS Search (more qualitative comparisons can be found in appendix section 11). We see that our model particularly excels in segmenting organ boundaries. This can be attributed to our effective strategy for extracting local-range dependencies, which plays a crucial role in extracting details from input data. Our model’s adeptness in capturing long-range dependencies allows it to grasp contextual information that extends over significant distances within the data. Simultaneously, its proficiency in handling short-range dependencies ensures precision in capturing localized patterns.

To further quantitatively support our strategy for extracting local and global features in parallel, in the next experiment, we report the performance of our model compared to the regular method for extracting features from medical images, which is relying on a single feature extractor. This translates into comparing LoGoNet to our feature extractor ULKANet. Table 4 reports the results. We observe that our strategy enables our model to outperform the alternative method.

In the next experiment, we report the efficacy of our pre-training method. To carry out this experiment, we use the algorithm proposed in Section 3.3 to initialize the weights of our model, and then, we follow the regular fine-tuning steps. In Tables 2 and 3 (the last rows), we report the results of this model for both datasets, indicated by postfix PRE. We see that the improvements achieved by pre-training are consistent across both datasets.

In the next experiment, we compare the effectiveness of our self-supervised pre-training approach to the alternative methods. In particular we compare to SimMIM (Xie et al., 2022), Rubuk’s Cube (Tao et al., 2020), and SimCLR (Chen et al., 2020) strategies. Table 5 reports the result. The numbers are obtained by initializing LoGoNet. Notably, our proposed model exhibits superior performance in three out of four experiments, showcasing its effectiveness in a diverse set of tasks. The complete results are available in the appendix 11. The comparison in Table 5 highlights the competitive edge of our model. This indicates the robustness and efficacy of our multi-task self-supervised learning methodology in capturing meaningful representations across various domains.

An inherent advantage of our pre-training approach lies in its versatility, as it is designed to be compatible with both CNN and ViT-based models. This flexibility broadens the applicability of our approach, allowing it to seamlessly integrate with different architectural paradigms commonly used in computer vision tasks.

To understand the impact of model size on the prediction accuracy, we report the performance of our default model compared to a larger variant. Our larger variant uses four encoder blocks with 3, 3, 24, and 3 transformer modules, respectively. The dimensions of the embedding vectors in this model are 96, 192, 384, and 768, respectively. Table 6 reports the results. Upon increasing the dimensions of our model, we observed an improvement in results, though it fell short of our initial expectations. We attribute this to the limited number of labeled data available. However, upon integrating our pre-training methodology into our standard and larger variants of LoGoNet, we noted a significant enhancement in performance, particularly noticeable in the larger LoGoNet.

In Section 3.3, we claimed that having multiple clusterers serves as a multi-task training approach. In order to demonstrate the benefit of having multiple clusterers, and also show the sensitivity of our model to the number of these learners in our algorithm, we report the results of our model with varying numbers of clusterers in Table 7. We see that as the number of clusterers increases, the performance improves. The results support our hypothesis regarding the ability of our model to extract broader knowledge from the unlabeled data in the presence of multi-tasking.

Finally, we report an ablation study on the effectiveness of our masking approach during the pre-training stage. In Section 3.3, we argued that by distorting input images, the model must learn the properties of neighboring pixels in order to predict the correct labels. We then argued that this exploration task enables the model to faster learn the domain and to generalize better. The results reported in Table 8 supports our claim. We see that by incorporating the masking step, the performance noticeably improves signifying a better generalizablity of our method.

In summary, we demonstrated the efficacy of our model in two datasets across 19 segmentation tasks. We also compared our method to eight recent baseline models, including those that use Visual Transformers. Our results testify to the effectiveness of our novel feature extraction techniques. Our analysis shows that our pre-training method is successfully able to exploit unlabeled data to improve parameter initialization. We also showed that our method significantly speeds up inference time compared to the best-performing models.

The computer vision domain is a rapidly evolving research field. It seems unrealistic to expect long-term plans. However, with the existing challenges in the medical domain, this community will invest more in developing methods for mitigating the lack of large labeled sets. Therefore, in the next step, we plan to explore Domain Adaptation, which is one of the well-known methods for addressing this challenge.

We used the scikit-learn implementation⁸⁸8Available at: https://scikit-learn.org/stable/ of the Mini Batch KMeans algorithm as the clusterers in our pre-training pipeline. During the training phase of the k-means models, we adopted a transformation process that converted the input image from a $Channel\times X\times Y\times Z$ format to a vector representation of dimensions $Z\times T$, where $T$ is equivalent to $Channel\times X\times Y$. This transformation enabled us to generate a label for each cluster per image slice, resulting in a sequence of labels for a sequence of images. Subsequently, the model underwent $350$ iterations of training, with each iteration utilizing a randomly selected $10\%$ subset of the unlabeled data. To introduce diversity and enhance robustness, we employed a stochastic approach in determining the value of $K$, randomly sampling from a range spanning $80$ to $500$.

The information pertaining to pre-training is encapsulated in Table 12. To train the pre-trained model, we leveraged the $AdamW$ optimizer and fine-tuned the process by configuring specific parameters. In particular, we assigned values of $0.1$ and $0.7$ to $\phi_{1}$ and $\phi_{2}$ respectively. Additionally, the sequence of distorted images, denoted as $M$, was set to $5$.

Table 11 presents the outcome of selecting hyperparameter values, with results obtained from the BTCV dataset using the ULKANet model. This tabulated information sheds light on the meticulous decision-making process involved in determining specific values for key hyperparameters, providing valuable insights into our experimental configuration.

Our observations reveal that augmenting both the values of $M$ (length of sequenced mask images) and $\phi_{1}$ (rate of sampled images) results in an increased rate of masked images. However, this heightened rate poses challenges for our model, making it more intricate to exploit dependencies between successive slices for effectively capturing information related to missing voxels. This delicate interplay between hyperparameters emphasizes the necessity of finding an optimal balance to enhance model performance, as an excessive increase in masked images may impede the model’s ability to leverage contextual dependencies within the data.

Furthermore, we introduced randomness in the selection of patch sizes, choosing from the set $(1,2,4,8,16,32,96)$.

Our innovative pre-training approach involves the incorporation of a header designed to adapt the model output to align with the requirements of our pseudo-labeling. Figure 2 shows the structure of our proposed pre-training. The structure of this header can be found in algorithm 4. This header modification serves as a crucial step in optimizing the model’s output for seamless integration with our pseudo-labeling methodology during the training process.

Table 14 provides a comprehensive overview of the specifics pertaining to our training or fine-tuning procedures. This tabulated information encapsulates crucial details, offering insights into the intricacies of our training regimen. By examining the contents of this table, readers can gain a nuanced understanding of the parameters, configurations, and methodologies employed during the training or fine-tuning phase of our experiments.

Our experimental setup involved the utilization of 8 nodes, each equipped with dual-GPU chipsets. Each GPU within these chipsets boasted a substantial 35GB of memory, ensuring ample resources for running our experiments efficiently. To adhere to established standards and foster equitable comparisons, we employed a comprehensive array of augmentation techniques to augment data variability. It’s noteworthy that these augmentations were uniformly applied to all models, encompassing both our proposed models and the baseline models. This meticulous approach ensures a fair and unbiased comparative analysis.

For the implementation of our models and the baseline models, we leveraged the $MONAI$ framework, which provided a robust and versatile foundation for our experimentation. This framework facilitated the seamless integration of our methodologies, streamlining the implementation process and contributing to the reliability of our experimental results. ⁹⁹9Available at: https://monai.io/ In the course of each iteration, we strategically implemented a randomized cropping strategy, extracting two images for each case during the training phase. This deliberate approach was employed with the intent of diversifying the training dataset for each input case within every epoch, thereby enhancing the overall richness of the training process. The randomness introduced by the cropping strategy contributed to a more robust and varied training experience.

Furthermore, to fine-tune and optimize our training procedure, we incorporated the advanced $AdamW$ optimizer. This optimizer played a pivotal role in adjusting the model’s weights during training, ensuring an efficient convergence towards optimal performance. The synergistic combination of the cropping strategy and the utilization of the $AdamW$ optimizer exemplifies our commitment to refining the training dynamics for improved model efficacy.

To refine our model using labeled data, we employed the DiceCELoss as the objective function during the fine-tuning or training process. The DiceCELoss function serves as a crucial metric, enabling us to strike a balance between the Dice coefficient and Cross-Entropy, optimizing the model’s performance on the labeled dataset. This choice of objective function reflects our commitment to a meticulous fine-tuning or training process, ensuring that the model is adept at capturing the nuanced patterns present in the labeled data. The DiceCELoss is articulated by the following formulation:

where

and

Thus, our fine-tuning and training loss term is the weighted summation between the regular dice loss term and the cross entropy term. $p_{i}$ represents the predicted probability for the $i$-th class. $t_{i}$ represents the ground truth label for the $i$-th class. $N$ represents the number of classes. $\epsilon$ is a small constant (e.g., 1e-5) added to the denominator to avoid division by zero.

Our experimentation has revealed that assigning equal weights to both CELoss and DiceLoss yields more favorable outcomes, surpassing the performance achieved with other weight ratios. The results of various weight configurations for losses are presented in Table 13. By according equal significance to both Cross-Entropy Loss (CELoss) and Dice Loss, we strike a balance that enhances the model’s ability to effectively capture diverse patterns in the data. The results highlight the critical role of assigning equal weights in attaining optimal outcomes and underscore the robustness of this approach within the framework of our experimental setup. This observation emphasizes the significance of maintaining a balanced weighting scheme, showcasing its effectiveness in ensuring stability and reliability across various experimental conditions. The consistent performance across diverse scenarios further validates the efficacy of the equitable weighting strategy adopted in our study.

As unlabeled data, we use a set of 4500 distinct data points re-published as a meta-dataset by Tang et al. (2022). Each of these data points represents a volumetric 3D scan specifically focusing on key anatomical areas such as the chest, abdomen, head, and neck. This deliberate selection of regions ensures a thorough examination of the anatomical aspects relevant to our research. This meta-dataset consists of the following dataset:

• •

  HNSCC ¹⁰¹⁰10Available at: https://wiki.cancerimagingarchive.net/display/Public/HNSCC: imaging, radiation therapy, and clinical data of head and neck squamous cell carcinoma (HNSCC) patients at MD Anderson Cancer Center, the dataset comprises 954 samples.

• •

  LUNA16 ¹¹¹¹11Available at: https://luna16.grand-challenge.org/Data/: is a publicly available data set and challenge designed to advance the development of computer-aided detection (CAD) algorithms for the accurate identification of pulmonary nodules in low-dose computed tomography (CT) scans, contributing to the improvement of lung cancer detection, and the data set consists of 842 samples.

• •

  CT Colonography¹²¹²12Available at: https://www.cancerimagingarchive.net/collections/: part of the National CT Colonography Trial, this dataset comprises 1532 cases of CT colonography imaging accompanied by polyp descriptions and their respective locations within colon segments. It serves as a valuable resource for validating the use of CT colonography in detecting colorectal neoplasia.

• •

  CT Images in COVID-19 ¹³¹³13TCIA COVID-19 Datasets: this collection of medical imaging data includes CT scans and associated clinical information collected during the COVID-19 pandemic. The dataset includes 722 samples.

• •

  LIDC-IDRI ¹⁴¹⁴14Available at: wiki.cancerimagingarchive.net/display/Public/LIDC-IDRI: is a web-accessible collection of thoracic computed tomography (CT) scans with marked-up annotations of lung nodules, designed for the development and evaluation of computer-assisted diagnostic methods for lung cancer detection, the dataset comprises 450 samples.

The provided dataset¹⁵¹⁵15Available at: https://www.synapse.org/#!Synapse:syn3193805/wiki/217789 consists of 40 abdominal CT scans that were collected under Institutional Review Board (IRB) supervision. These scans are sourced from a combination of an ongoing colorectal cancer chemotherapy trial and a retrospective ventral hernia study. The CT scans were obtained during the portal venous contrast phase, capturing the anatomical details of the abdominal region. The dataset showcases variability in terms of volume sizes, with the scans ranging from $512\times 512\times 85$ to $512\times 512\times 198$ voxels. Additionally, the field of view (FOV) varies, spanning approximately $280\times 280\times 280,\text{mm}^{3}$ to $500\times 500\times 650,\text{mm}^{3}$. In-plane resolution fluctuates between $0.54\times 0.54,\text{mm}^{2}$ and $0.98\times 0.98,\text{mm}^{2}$, while slice thickness varies from $2.5,\text{mm}$ to $5.0,\text{mm}$. Standard registration data for the dataset has been generated using NiftyReg.

For the preparation of the data set, we sliced at precise intervals of $1.5$ mm in both the x and y directions and $2$ mm for the isotropic resolution in the z-direction.

The dataset includes manual segmentations of thirteen abdominal organs. The segmentation was performed on a volumetric basis, providing accurate organ boundaries within the CT scans. The list of segmented organs includes the spleen, right kidney, left kidney, gallbladder, esophagus, liver, stomach, aorta, inferior vena cava, portal vein and splenic vein, pancreas, right adrenal gland, and left adrenal gland. It’s important to note that some patients within the dataset may not have a right kidney or gallbladder, and as a result, these organs might not be labeled in those specific cases.

The dataset serves as a valuable resource for advancing automated segmentation approaches in the field of medical imaging. Researchers can utilize this dataset to develop and refine algorithms that accurately identify and segment abdominal organs in CT scans.

This dataset¹⁶¹⁶16Available at: http://medicaldecathlon.com/ contains ten different tasks. Within this data set, a wide variety of organs and structures are known to encompass vital anatomical regions of significance, such as the heart, liver, prostate, and pancreas.

To use the dataset, we divided it into exact 1.0 mm intervals along the x, y, and z directions, maintaining isotropic resolution for all tasks.

We evaluated our models using six tasks from the MSD dataset. Table 15 displays the case numbers for the various tasks. The initial task focuses on ”Colon Cancer Primaries” and poses a challenge due to its varied and irregular appearance, referred to as ”heterogeneous” (heterogeneous indicating dissimilar components or elements, causing irregular or variegated appearances; for instance, a dermoid cyst exhibits heterogeneous attenuation on CT scans).

The subsequent task revolves around ”Spleen,” involving a challenge related to the significant variation in foreground size. The following task targets ”Hepatic vessels and tumor,” presenting a challenge in dealing with small tubular structures adjacent to the heterogeneous tumor. Moving on, we have the ”Pancreas Tumor” task, which aims to segment the liver and tumor, posing an unbalanced labeling challenge with large (background), medium (pancreas), and small (tumor) structures.

The ”Lung Tumors” task comes next, presenting the challenge of segmenting a small target (cancer) within a large image. Lastly, we have the ”Cardiac” task, which focuses on segmenting the Left Atrium and faces the challenge of a small training dataset with significant variability.