A Deep Learning Approach Using Masked Image Modeling for Reconstruction of Undersampled K-spaces

This study makes use of the open source fastmri dataset (Zbontar et al 2019). Specifically, this study uses only knee MR images and k-spaces from the validation set of images. The dataset includes Contains fully sampled knee MRIs from 1500 patients obtained from 1.5 and 3.0T scanners. The dataset also contains more than 10,000 DICOM images with coronal and axial proton density with and without fat suppression (Figure 2), T2 weighted with fat suppression and sagittal proton density. The dataset was downloaded as multiple h5 files, meaning a custom dataloader had to be built to pass the data into the model. The model was unable to process h5 datasets or images, meaning the images were saved and passed through as PNG images. These images were not split into seperate categoires during training or validation. Only slight augmentations such as cropping for resizing were applied to the data for the purpose of preserving image details. The data was split 80/20 for train and test.

The model used in this study is a modified version of the SimMIM model. The model features two primary encoders for classification; the vanilla Vision Transformer (ViT), and then Shifted-Window transformer (Swin). The SimMIM utilizes an encoder-decoder architecture, with the encoder being either of the aforementioned models, and the decoder being a sequential model with basic 2D convolutional layers.

The Swin Transformer is a type of vision transformer that classifies based on a shifted window algorithm. The Swin creates a hierarchical representation by starting from small sized patches and gradually merging neighboring patches. As the layers progress (Figure 3), the window for self-attention calculation (red) ”shifts” , resulting in new windows. These new windows cross the boundaries of the old windows, forming attention between them (Liu et al 2021). This method works well for classification, but is not optimized for prediction. The baseline model performed reasonably well for reconstruction (See Figure 11), but changes were made to improve prediction performance.

The ViT transformer works similarly to the Swin transformer. It receives input in the form of a 1D sequence of token embeddings (Dosovitskiy et al 2021). The same methodology of patch embeddings as the Swin transformer are applied (Figure 4). Position embeddings are added to the patch embeddings to retain positional information. The difference between this model and the Swin is that the window for self-attention caluclation does not shift between layers (Liu et al 2021). While this approach suffices for classification, it becomes a drawback for image reconstruction since the details of the image are not preserved during extraction.

For experimentation in this study, the primary changes were made with the hyperparameters in the encoders. To begin, the model first had to be transferred from classification optimized to prediction optimized. The number of classes was reduced to 1, batch size reduced to 1, and the several patch sizes and encoder strides were experimented with. The model found decent success with a patch size of 1 for original embeddings, and 1 for model initialization, a significant decrease from the original values of 4 for each. The encoder stride was also changed, and the model found the most success with an encoder stride of 32 (default) since the model had already been optimized with that value. The number of epochs was held at a steady 100 for all trials.
Before the data was passed into the model, it was resized to fit the model. Originally, the data had the shape of 384 $\times$ 680, but the model only uses square images. To control for this, this study only used center crop to resize the image to 192 $\times$ 192. Augmentations such as random flip or recolor were removed for training, since they are for classification purposes and do not preserve any of the necessary details. For normalization, trials were first run with default ImageNet normalization values (model default), then run with image specific mean and standard deviation values, before being removed completely. Ultimately, we decided to remove normalization from the augmentations since the dataset was relatively small and normalization removes some of the fine details. After the augmentations were made, the images were passed through a masking function. This function was used to simulate under sampling of the k-space, since the removal of input channels would train the model in similar methods.
It is important to note that for this study, the saved images were in the format of ’RGBA’ instead of the typical ’RGB’. For the models, which use ’RGB’ base, many of the functions (especially weights) were optimized and even hard coded for 3 input channels only. This meant we had to create a custom dataloader and modify each function, making sure to scale each augmentation, as well as rewriting the weights and gradients to accommodate the batch size and input channel change. The model was evaluated on the metrics of loss, gradient normalization, and structural similarity. Specifically, this study used L1-loss to determine base similarities between the reconstructed images. Gradient normalization was used to indicate the training performance of the model, making sure that during pretrain the model reached optimal state.

Structural similarity (SSIM) was used as the primary statistic for evaluation since it demonstrates a direct comparison between the image and the reconstructed image. Because it’s primary non machine learning use is to determine drops in image quality (as well as image similarity), it functioned for both image quality and similarity measurement (Blaimer et al 2004). The primary drawback of skimage’s SSIM function is that it requires the images to be passed in as grayscale; however, in this case, our images are already grayscale, meaning minimal detail was lost.

For the Swin transformer, the model saved reconstructed images with structural similarity values over 0.995, and for the ViT transformer, the model saved images with values over 0.60. Originally, the model was also going to evaluated based on low validation loss values ($<$0.001), but by image and SSIM comparison, we concluded that low validation loss was not a proper measurement for image similarity. Based on images similar to the ones shown in Figure 5, it was determined that the patch size of the model and the encoder stride were too large for detail reconstruction (as evidenced by the appearance of solid color patches).

For the above images, the specific SSIM value was 0.9956, showing that even a 99.56$\%$ structural similarity can miss key features. However, the Vision transformer performed significantly worse.

It was to be expected that the Swin would outperform the ViT transformer in reconstruction (Figure 6); however, it was still interesting to see by how much it under performed in comparison to the Swin.

Before the model reached the fine-tuning stage, metrics such as training and gradient norm were collected and graphed to ensure the model never overfit and reached a steady state before being saved for fine tuning.

From the gradient normalization data shown in Figure 7, it was determined the models reached a steady state while completing the training cycle. Investigations into overfitting were not undertaken until the fine-tuning portion of the model. For this, a graph of validation and training loss was constructed during the fine-tuning cycles. These graphs were used to determine if the model was overfitting by examining the trend of the validation loss.

Figure 8 illustrates the the validation and training loss of the Swin encoder-based SimMIM. Since the validation loss continued to change with the training loss, it can be ascertained that the model was not memorizing the data, but rather using the features extracted by the training during each new validation cycle. Figure 9 (below) illustrates the validation and training loss of the ViT encoder-based SimMIM. Unlike the Swin base, the validation accuracy does not significantly shift with the training loss. This shows that since the validation loss is not responsive to the training, it is instead just memorizing the data without improving.

For the SSIM portion of data collection, graphs were made between both validation and loss, as well as between both of the encoders. It is important to note that while SSIM seemed to have an inverse correlation with L1-loss, their relationship is less substantial than expected. While both the Swin and ViT encoder models approached validation loss values of $<$0.01, only the Swin encoder exceeded SSIM values of 99$\%$, while the ViT encoder hovered around 60$\%$.

Figure 11 illustrates the ViT encoder’s maximum SSIM values. While these values hovered consistently around 0.60, Figure 10 illustrates the validation loss values to be $<$0.05 during the same epochs. In comparison, the Swin encoder has similar loss values (though fluctuating more) with significantly higher sturctural similarity. The Swin transformer outperformed the ViT by SSIM values of over 49$\%$. This structural similarity value illustrates a reasonable success in reconstruction, although some of the fine details were lost.

Attempts were made to incorporate other, more reconstruction focused models as encoders for this study. One such model was the SwinRecNet developed in the study by Pan et al 2021 (Figure 12).

The SwinRec functions very similarly to the SwinTransformer, but it optimizes reconstruction instead of classification. It does this by using an edge enhancement reconstruction sub-network for recovering the initially-reconstructed MR image. The problem with using this architecture is that it is also optimized for ’RGB’ images, and was developed for a different model MistNet. This means transferring the model to to be compatible with the SimMIM made it lose certain features, or compatibility as a whole. Ultimately the model was not included in this study due to the loss of performance it found when being made compatible.

This study also examined the effects of augmentation on reconstruction. Data augmentation such as recoloring and interpolation were removed for detail preservation reasons. These augmentations cause significant changes to the images, becoming counterproductive in balance with overfitting reduction. However, slighter augmentations such as normalization, randomflip, and randomcrop were tested to see if the training metrics would improve as a result. Data was gathered for both augmented data models and models where the data faced no augmentation except center crop (for size constraints).

Figure 13 illustrates the relationship between the SSIM of model with augmented data and the SSIM of a model without. Interestingly, the model’s SSIM appears to increase much quicker with augmentation, but improves less over time than the model with non-augmented data. Since augmentation is used typically to reduce overfitting chances and keep the model learning, this is an interesting result that the augmented data model appears to have a constant SSIM value for epochs 10-50. This finds that for reconstruction, not augmenting the data appears to have a greater effect on training the model. This is inversely true in comparison to many classification networks, although it makes sense because augmentation should not be applied for many basic reconstruction networks due to image alterations.

Because of the fine details required between k-space and image conversion, the reconstructed image shown in Figure 5 (right) will likely produce an inaccurate MR image. Additionally, this study only utilized a relatively small dataset of  7500 images, in comparison to other studies that use datasets with millions of images. Of course, this would probably help the overall reconstruction, since the model would have more images to train on, but it is still important to note that this study was conducted on a small sample set. Lastly, the masking in this study was only simulated through a masking function. In a real world masked k-space, the masking might function differently, meaning the model would need to be tested on actual masked k-spaces to confirm the results for reconstruction.