REHRSeg: Unleashing the Power of Self-Supervised Super-Resolution for Resource-Efficient 3D MRI Segmentation

MRI segmentation aims to identify and delineate ROI in acquired MR images, playing a crucial role in computer-aided diagnosis and disease progression monitoring [13]. Recent advancements in deep learning technologies have led to significant breakthroughs in automatic MRI segmentation. Among the most widely adopted architectures are U-shaped CNN-based networks [14, 15, 16, 17], with nnUNet [18] emerging as a leading model. nnUNet automatically configures various datasets and focuses on dataset preprocessing and inference strategies, currently dominating the field of medical image segmentation. Furthermore, vision transformer (ViT) techniques have been applied to MRI segmentation, leveraging the transformer’s ability to capture global and long-range dependencies [19, 20, 21, 22, 23].

Despite the high performance of these methods, they often require substantial computational resources for whole-volume segmentation of HR MRI images. Consequently, several approaches have aimed to develop efficient 3D MRI segmentation models with reduced resource demands by designing lightweight architectures. For example, ADHDC-Net [24] uses decoupled convolution, dilated convolution, and attention-based refinement to minimize the number of parameters and operations in brain tumor segmentation. Similarly, UNETR++ [21] employs efficient paired attention to reduce the complexity of self-attention computations and constrains the number of parameters by sharing weights between queries and keys. MISSU [25] incorporates a self-distillation mechanism to refine local 3D features with multi-scale fusion blocks, which are removed during inference to ease computational burdens. However, these lightweight models often struggle to capture rich features, limiting segmentation performance. Our approach addresses this limitation by utilizing the priors embedded in super-resolution models to enhance segmentation performance without increasing computational costs during inference.

Recently, several studies [26, 6, 27, 8, 9] have focused on achieving high-resolution (HR) segmentation from low-resolution (LR) images, primarily designed to efficiently capture global information and further reduce the computational cost, especially during inference. For example, AR-Seg [8] reduces computational costs for video semantic segmentation by lowering the resolution for non-keyframes, while employing a cross-resolution fusion module to prevent accuracy degradation. GDN [7] replaces conventional down-sampling with a learnable procedure to reduce image resolution for real-time segmentation. UHRSNet [28] fuses local features extracted from image patches with global features from downsampled images.

Super-resolution techniques have also been employed in this context to improve segmentation performance by capturing fine-grained details and providing richer semantic information. Lei et al. [29] propose a framework for remote sensing images that simultaneously performs quadruple super-resolution and segmentation. PFSeg [26] integrates super-resolution as an auxiliary task for patch-free 3D medical image segmentation, and incorporates a fusion module for joint optimization. ISDNet [27] introduces super-resolution tasks as guides in the coarse-grained branches of the network, and fuses this information into fine-grained branches. DS2F [9] redesigns the segmentation and super-resolution framework by proposing a shared feature extraction module and a proxy loss for ROI recognition.

While these methods achieve HR segmentation from LR images, they still rely on resource-intensive HR annotations during training. In contrast, our proposed REHRSeg framework only requires LR annotations for training, significantly reducing the annotation burden. Additionally, REHRSeg does not depend on HR images for super-resolution assistance, making it more practical for clinical applications where 2D scanning protocols are commonly employed.

Due to limitations in medical resources and scanning time, the acquisition of MR images employs a 2D scanning protocol in many clinical settings, resulting in low inter-plane resolution with large slice thickness [30, 31]. Therefore, super-resolution technology has been employed to enhance the resolution of such images, particularly in the inter-plane direction. Deep learning methods, especially CNN-based architectures, have become the dominant approach for MRI super-resolution. Conventional deep learning methods typically rely on fully supervised training [32, 33, 34], requiring paired HR and simulated LR images. However, collecting HR images in real-world clinical scenarios is significantly more challenging and costly than acquiring LR images, making it difficult to obtain sufficient paired training samples.

To address this, self-supervised training methods have been increasingly applied to super-resolution tasks. These methods achieve super-resolution of MR images using only LR images for training. For example, Xuan et al. [35] propose a VAE-based generative model trained on LR images, and synthesizes HR images from interpolated latent space to train the super-resolution model. SMORE [36] uses HR and LR data inherent in the high-resolution plane for training, restoring image quality by improving resolution and reducing aliasing. Wang et al. [37] extend SMORE by introducing video frame interpolation as a preliminary task, and design an architecture based on implicit neural representation with a three-stage training protocol to refine the results. TSCTNet [38] incorporates a cycle-consistency constraint to enable self-supervised learning for reducing the slice gap for MR images.

In this work, we explore the capacity of self-supervised learning from annotated LR images in the field of MRI super-resolution to enhance and provide supervision for segmentation tasks, a direction not yet explored in prior super-resolution studies.

The overall framework of REHRSeg is designed to produce high-resolution (HR) segmentation $\mathbf{\hat{Y}}_{HR}$, with an inter-plane resolution that is $r$ times higher than the original low-resolution (LR) image $\mathbf{I}_{LR}$ and its annotation $\mathbf{Y}_{LR}$. As illustrated in Fig. 2, we begin by applying self-supervised super-resolution (self-SR) on the annotated MR images, generating coarse HR annotations $\mathbf{\tilde{{Y}}}_{HR}$ alongside the HR MR images $\mathbf{\tilde{{I}}}_{HR}$. Next, the self-SR model and the pseudo HR data are leveraged to enhance the segmentation model, enabling simultaneous generation of both refined LR segmentation $\mathbf{\hat{Y}}_{LR}$ and HR segmentation $\mathbf{\hat{Y}}_{HR}$. This is achieved through three key components: 1) utilizing self-SR as a pseudo-data generator for segmentation, 2) incorporating an uncertainty-aware super-resolution (UASR) head in self-SR for additional boundary guidance, and 3) applying structural knowledge distillation from the self-SR model.

We introduce a self-supervised super-resolution (self-SR) method as a preliminary task in REHRSeg. This approach is built upon the principles of the most recent self-supervised MRI super-resolution method [37], with modifications to support the super-resolution of annotated MR images. Given a LR image $\mathbf{I}_{LR}$ and its corresponding annotation $\mathbf{Y}_{LR}$, both with a spatial resolution of $a\times a\times c$, we assume that the frequency- and phase-encoding directions have identical spatial resolution. The super-resolution process is applied with a scaling factor of $r$, generating HR images with a spatial resolution of $a\times a\times(c/r)$ in two main steps: First, both the image and annotation are interpolated to achieve isotropic voxel spacing, so that the spatial consistency can be ensured between training and inference stages. Our experiments show that the utilization of SMORE [36] as an alternative to traditional interpolation techniques, such as B-spline interpolation for images and nearest-neighbor interpolation for annotations, can further improve the final results.

In the second step, we generate the LR-HR pairs required for self-SR training by simulating slice separation along the $x$ axes. Specifically, the interpolated images are convoluted with a 1D Gaussian filter $h(x;r)$ as the slice profile with FWHM equal to $r$, followed by downsampling by a factor of $r$ along with the annotations. Unlike SMORE, we avoid introducing aliasing artifacts, as our self-SR model directly operates on $\mathbf{I}_{LR}$ instead of the interpolated one to reduce computational cost. Additionally, the slice gap is not considered in our self-SR method, as many segmentation datasets lack this information. Notably, several SR methods such as SMORE[36], DeepResolve[32], and SAINR[34] also neglect the slice gap without experiencing significant performance degradation. In this scenario, the task can be considered as slice interpolation and deblurring for 3D data.

To accelerate training and ensure faster convergence, we leverage a pretrained model from the domain of video frame interpolation. This strategy has been previously shown to be effective in self-supervised super-resolution [37]. Here we use the FLAVR [39] model pretrained on Vimeo-90K [40] to initialize the backbone of our self-SR model. By applying the trained model to $\mathbf{I}_{LR}$ and $\mathbf{Y}_{LR}$, we obtain the estimated pseudo HR image $\mathbf{\tilde{I}}_{HR}$ and its annotation $\mathbf{\tilde{Y}}_{HR}$. The pseudo data and the self-SR model are then used to enhance the segmentation model in the subsequent stage.

We propose to utilize the self-SR model as a data generator to create pseudo-data for the MRI segmentation task, which is implemented in two distinct ways: data augmentation for LR segmentation and pseudo supervision for HR segmentation. Given the HR image $\mathbf{\tilde{I}}_{HR}$ and its annotation $\mathbf{\tilde{Y}}_{HR}$ produced by self-SR, we can simulate the LR data $\mathbf{\tilde{I}}_{LR}$ and $\mathbf{\tilde{Y}}_{LR}$ with a large slice thickness that matches the spatial resolution of the segmentation dataset, by blurring and downsampling the synthesized data as follows:

where $*_{z}$ denotes 1D convolution along the $z$ axes, and $\downarrow^{r}_{z}$ represents the downsampling operator along the $z$ axes by a factor of $r$. This process creates $r\times$ the amount of LR data compared to the original dataset, as the improved resolution in the synthesized data allows for additional data generation. This is accomplished by applying the downsampling operator while adjusting the corresponding starting index of the target image. Note that this can also be considered a novel form of data augmentation for LR segmentation, and the prior knowledge embedded within the self-SR process can be leveraged to further improve the segmentation performance with this additional data.

Furthermore, by generating the pseudo $\mathbf{\tilde{I}}_{LR}$-$\mathbf{\tilde{Y}}_{HR}$ pair via self-SR, the segmentation model is enabled to perform HR segmentation in parallel with LR segmentation task. This is achieved by using an additional HR segmentation head consisting of a upsampling layer to the segmentation backbone, just before the original LR segmentation head. During training, both tasks share the same input $\mathbf{\tilde{I}}_{LR}$, allowing for multi-task learning of LR and HR segmentation using pseudo-supervision from self-SR.

We further utilize the self-SR model to provide uncertainty guidance for the segmentation task, with the goal of enhancing its performance with uncertainty-aware segmentation. It is widely recognized that the tissues in medical images often overlap and blurriness occur on their boundaries, leading to high uncertainty in the delineation of these regions [41]. This issues not only lead to unreliable segmentation region, but also cause poorly reconstructed regions for the super-resolution tasks. Therefore, we propose to estimate uncertainty in the self-SR model by designing an uncertainty-aware super-resolution (UASR) head, which identifies the regions with high reconstruction error and provides guidance for the segmentation task.

The first step of uncertainty estimation with UASR involves the generation of intermediate results. We extract the 3D feature map $\mathcal{F}^{sr}_{o}\in\mathbb{R}^{C\times D\times H\times W}$ from the backbone of the self-SR model, where $D$, $H$, $W$ represent the slice depth, height, and width of the input LR images, respectively. The features along the depth dimension are concatenated into a 2D feature map, which is then processed to generate intermediate features $\mathcal{F}_{m}$ via Leaky ReLU activation and convolution. These features are split along the channel dimension and merged back into the depth dimension, yielding $N$ intermediate images $\{P_{I}^{j}\}_{j=1}^{N}$ and annotations $\{P_{Y}^{j}\}_{j=1}^{N}$. The whole procedure can be expressed as:

where $\{W_{Y}^{j},b_{Y}^{j}\}_{i=1}^{N}$ and $\{W_{Y}^{j},b_{Y}^{j}\}_{i=1}^{N}$ are the 3D convolutional filters applied to the intermediate features. The operators $\text{merge}(;a;b)$ and $\text{split}(;a;b)$ denote the merging of features along dimension $a$ into dimension $b$, and the splitting of features along dimension $a$ into dimension $b$, respectively. Next, we obtain multiple attention maps by applying another set of convolutions, followed by a channel-wise softmax function:

Then, the attention maps are used to perform channel-wise selection for the intermediate images and annotations, and the final results are obtained via:

where $\otimes$ denotes element-wise multiplication. Similarly, the uncertainty map is also generated from the intermediate attention maps via:

where $[\cdot]$ denotes concatenation along the channel dimension, and $\sigma$ is the sigmoid activation function. The modified loss function for the super-resolution task is used to ensure that the uncertainty map accurately reflects regions with high reconstruction error:

where $L_{pix}^{sr}$ is the original pixel-level loss function for the super-resolution task, which is defined as $L_{1}$ loss for the image. This modification ensures that the uncertainty map is automatically learned to represent the reconstruction error. Note that we do not use the uncertainty map to regularize the super-resolution of annotations, which uses cross entropy and dice loss, to stabilize the learning process.

Since regions with high reconstruction error often correspond to boundaries of anatomical structures, we incorporate the uncertainty map to guide the segmentation task by introducing a segmentation loss regularization term:

where the uncertainty map extracted from UASR is first normalized to have its intensity range of $[0,1]$ before employed as the weight map for the pixel-level segmentation loss.

We also propose to leverage knowledge distillation from the self-SR model to further enhance MRI segmentation. Given the disparity between image reconstruction and segmentation tasks, traditional distillation methods that constrain features with pixel-wise alignment only provide limited benefits. Instead, the structural correlation between regions shares similar patterns between the feature map produced by different tasks, which have also been confirmed in the previous work [42]. Therefore, the segmentation task has potential to benefit from the self-SR task by distilling the correlation patterns from it. Here the feature maps $\mathcal{F}^{sr}$ and $\mathcal{F}^{seg}$ are extracted from the trained self-SR model and current segmentation model. The shape of $\mathcal{F}^{sr}$ is aligned with $\mathcal{F}^{seg}$ with bilinear interpolation, resulting in feature maps of ${C^{\prime}\times D^{\prime}\times H^{\prime}\times W^{\prime}}$. We model the spatial correlation between regions using a fully-connected affinity graph, where the nodes represent different spatial locations, and the edges denote their similarity. In this graph, we aggregate the $\beta$ voxels in a local patch to represent the feature of each node, setting the granularity of the graph to $\beta$. Consequently, the affinity graph comprises $(D^{\prime}\times H^{\prime}\times W^{\prime})/\beta$ fully connected nodes. We compute the similarity $a_{ij}$ between the $i$-th and $j$-th nodes by:

where average pooling is applied to aggregate features before calculating similarity. This process is performed for both the self-SR and segmentation models, with $a_{ij}^{sr}$ and $a_{ij}^{seg}$ denoting the similarities from the self-SR and segmentation models, respectively. The loss for correlation distillation is defined as the squared difference between these similarity measures:

Additionally, we align spatial features to highlight critical regions and improve segmentation performance. This is achieved through a cosine distance loss between the features from the self-SR model and the adapted segmentation features, as described in the previous work [43]:

where $h^{(v)}(\cdot|\Theta^{(v)})$ is the student adaptor for feature vectors with parameters $\Theta^{(v)}$.

The overall loss for training our segmentation model is composed of an uncertainty-aware segmentation loss, a pseudo HR segmentation loss, and a knowledge distillation loss:

where $\lambda$ is used to balance the losses for segmentation and distillation. $\mathcal{L}_{u}^{seg}$ is implemented with cross entropy loss weighted with uncertainty map, and $\mathcal{L}_{HR}^{seg}$ is implemented with cross entropy loss with dice loss.

Our experiments are conducted on a public dataset with HR images and their annotations, as well as an in-house dataset with only LR labeled images. Due to the difference in data availability, we calculate quantitative metrics on the public dataset for HR results, while providing only qualitative HR results for the in-house dataset. Both datasets undergo 5-fold cross-validation to ensure a robust and unbiased evaluation.

We utilize a downsampling factor (DSF) of 4 for both datasets, and initialize our self-SR backbone with the FLAVR model pretrained on $4\times$ video frame interpolation. To stabilize the training of the UASR head, we only include the uncertainty-guided loss during the final 20,000 iterations, following an initial fine-tuning phase of 130,000 iterations with the batch size of 64. For structural distillation, we aggregate only in-plane features with a granularity $\beta$ set to $1\times 2\times 2$. The spatial distillation process employs a convolution layer following the previous work [50], since experiments show that more complex distillers, such as those in [43], do not improve segmentation performance in this context. The segmentation network is initialized using the default setting of nnUNet-3D [18], and its data augmentation strategy is applied throughout our experiments. To allow the segmenter to produce HR results, we interpolate the features before the last layer and use two convolution with a ReLU activation in between, which is light-weighted to produce the final results. All the experiments are performed on an NVIDIA RTX A6000 GPU.

In this section, we evaluate whether the self-SR model can provide effective guidance for the segmentation model, by inspecting whether its performance is improved for LR segmentation. We also conduct experiments for HR segmentation on the Meningioma-SEG-CLAS dataset, as 2D-based segmentation methods can be directly applied to cross-sectional slices extracted from HR volumes. It is important to note that these 2D-based methods use HR images as inputs, whereas our method operates on LR images in the inference stage. Quantitative results presented in Table 1 indicate that 2D-based methods generally underperform 3D-based methods for MRI segmentation, particularly in terms of the HD95 metric. Notably, nnUNet remains a strong baseline and outperforms several transformer-based methods, especially for the in-house dataset where the complex data distribution poses challenges for these methods to achieve satisfactory results. This can be attributed to the limited amount of training data, which prevents transformer-based methods from fully leveraging their advantages.

Despite these challenges, REHRSeg significantly improves performance over baseline ($p<$ 0.05, HD95) and outperforms all the alternatives for both LR and HR segmentation. Visualization results in Fig. 5 also demonstrate that REHRSeg excels in both tumor types, and notably enhances the nnUNet baseline in boundary recognition. Moreover, REHRSeg shows strong performance in difficult cases where other methods fail to accurately capture the tumor shapes, as illustrated in the fourth row of Fig. 5.

Although REHRSeg is trained on LR images and their corresponding annotations, it allows the acquisition of coarse HR results via the proposed self-SR model. Additionally, the segmentation model is enhanced using these self-SR results to produce HR segmentation outputs. The evaluation of these results are further illustrated in the following sections.