Source-linked AI summary
Multi-Task Attention-Based Semi-Supervised Learning for Medical Image Segmentation
Shuai Chen, Gerda Bortsova, Antonio Garcia-Uceda Juarez, Gijs van Tulder, Marleen de Bruijne
TL;DR
Medical image segmentation often has few labeled examples, motivating a method that uses unlabeled images without relying solely on noisy pseudo-labels. MASSL couples supervised segmentation with attention-guided, class-specific reconstruction and outperformed several baselines across brain-tumor and white-matter-hyperintensity tasks. Alternating training was generally more effective than joint training, while performance remained below state of the art on the reported test sets.
Problem
Medical image segmentation has expensive, tedious manual annotation, while reconstruction is difficult to combine with U-Net skip-connections and pseudo-labels can be lower quality than ground truth.
Method
MASSL combines supervised segmentation and unsupervised reconstruction through attention that separately reconstructs foreground and background regions from soft segmentation predictions.
Results
MASSL outperformed same-label fully supervised CNNs, pre-training methods, and the no-attention MSSL approach across brain-tumor and white-matter-hyperintensity applications.
Takeaways & Limitations
The attention mechanism improves the link between segmentation and reconstruction, while alternating training generally performs better and more stably than joint training.
Takeaways & Limitations
On BraTS18 and WMH17 testing sets, MASSL performance was lower than reported state-of-the-art results.
Abstract
from arXiv · showhide
We propose a novel semi-supervised image segmentation method that simultaneously optimizes a supervised segmentation and an unsupervised reconstruction objectives. The reconstruction objective uses an attention mechanism that separates the reconstruction of image areas corresponding to different classes. The proposed approach was evaluated on two applications: brain tumor and white matter hyperintensities segmentation. Our method, trained on unlabeled and a small number of labeled images, outperformed supervised CNNs trained with the same number of images and CNNs pre-trained on unlabeled data. In ablation experiments, we observed that the proposed attention mechanism substantially improves segmentation performance. We explore two multi-task training strategies: joint training and alternating training. Alternating training requires fewer hyperparameters and achieves a better, more stable performance than joint training. Finally, we analyze the features learned by different methods and find that the attention mechanism helps to learn more discriminative features in the deeper layers of encoders.
1 Introduction
MASSL addresses limited medical segmentation annotations by combining supervised segmentation with an unlabeled reconstruction task. Its attention mechanism creates class-specific reconstruction targets to learn more discriminative features, outperforming several comparison methods across two segmentation problems.
- Manual annotations for medical image segmentation are expensive and often available for only a small subset of training data.
- Combining reconstruction with U-Net-like segmentation is difficult because skip-connections can bypass the dimensionality reduction needed by an autoencoder.
- MASSL combines an autoencoder with a U-Net-like network and reconstructs attention-generated segmentation labels rather than the original input.
- The attention mechanism separates reconstruction targets by class, guiding the auxiliary task toward more discriminative latent representations than traditional reconstruction.
- MASSL was evaluated on brain tumors and white matter hyperintensities, outperforming fully supervised, pre-training, and no-attention baselines.
2 Methods
The method shares an encoder between segmentation and reconstruction networks, using attention to reconstruct foreground and background separately. Training can be joint or alternating, with each strategy updating the shared multi-task model differently.
- Architecture and loss functions: MASSL uses segmentation and reconstruction networks that share an encoder, with an attention mechanism connecting the two tasks.
- Architecture and loss functions: The segmentation network uses U-Net-like skip-connections and Dice loss on labeled images, while the reconstruction network uses a no-skip decoder and MSE on labeled and unlabeled images.
- Architecture and loss functions: Unlike MSSL, MASSL reconstructs foreground and background separately using soft segmentation predictions from the segmentation network.
- Architecture and loss functions: The attention-based reconstruction objective blocks gradient propagation through segmentation predictions and weights foreground and background terms by mask size.
- Training strategy: Joint training combines segmentation and reconstruction losses in each minibatch using equal labeled and unlabeled samples and a weighting hyperparameter γ.
- Training strategy: Alternating training samples labeled and unlabeled batches separately and updates the segmentation and reconstruction paths according to the corresponding loss.
3 Experiments
Experiments use public BraTS18 and WMH17 MRI datasets with separate training, validation, and testing splits. The network is a five-level 3D U-Net-inspired architecture, and encoder feature discrimination is assessed by voxel-wise linear regression.
- Datasets: Experiments use public BraTS18 and WMH17 challenge data for brain-tumor and white-matter-hyperintensity segmentation.
- Datasets: BraTS18 contains 220 scans split into 120 training, 50 validation, and 50 testing scans, with binary whole-tumor segmentation from FLAIR images.
- Datasets: WMH17 contains 60 scans from three sites, split into 30 training, 10 validation, and 20 testing scans with site-balanced sets.
- Network and training: The network uses five resolution levels, 3D convolutions, instance normalization, LeakyReLU activations, and skip-connections only in the segmentation decoder.
- Feature analysis: Voxel-wise linear regression evaluates how well encoder features discriminate foreground from background at each resolution level.
4 Results
MASSL with alternating training achieved the strongest segmentation performance, while joint training was less stable. Feature analysis further indicated that attention produced more discriminative deeper encoder features.
- Segmentation performance: MASSL(alter) achieved the best segmentation performance across the evaluated methods.The comparison covered supervised CNN, pre-training approaches, and multi-task methods with and without attention.
- Training strategies: Joint training performed slightly worse than alternating training and varied substantially across labeled/unlabeled data splits.This variation reflected instability in joint training and difficulty tuning γ.
- Evaluation: The BraTS18 results report Dice similarity coefficients averaged over all cross-validations, with significance tests comparing methods within each column.The final column uses all labeled images as unlabeled data except for CNN.
- Feature analysis: MASSL learned more discriminative features in deeper encoder layers than CNN and MSSL, as indicated by higher R2 scores.The feature analysis supports the proposed role of attention in improving deeper encoder representations while retaining the reconstruction objective.
5 Discussion and Conclusion
MASSL improves semi-supervised segmentation across brain-tumor and WMH applications, with attention linking reconstruction to segmentation and alternating training generally outperforming joint training. The discussion also identifies dataset sparsity, simplified comparisons, limited feature analysis, and restricted MRI data as important boundaries.
- Results: MASSL outperforms supervised CNN, pretraining methods, and MSSL without attention across the evaluated segmentation applications.In fully supervised training, improvement over baseline CNN was statistically significant only for BRATS, partly because sparse WMH foreground makes attention maps less effective.
- Mechanism: The attention mechanism mainly drives improvement by introducing segmentation information into reconstruction and linking the two tasks more effectively.The mechanism is described as integrable into CNN architectures and extendable to multi-class segmentation.
- Training strategy: Alternating training generally performs better than joint optimization and avoids fine-tuning γ, but still requires suitable initial learning rates and can be unstable.Experiments found greater stability when reconstruction used a smaller initial learning rate than segmentation.
- Limitations: The comparisons use a regular autoencoder instead of the variational autoencoder used by Sedai et al., limiting direct comparability between pretraining strategies.The authors suggest MASSL could also work with a VAE, potentially fusing the tasks more effectively.
- Limitations: Feature analysis uses a simple linear regression model, so it measures linear rather than nonlinear discriminative power.The authors identify more complex nonlinear analysis as a direction for future work.
- Limitations: Using one MRI sequence and a subset of scans leaves BraTS18 and WMH17 performance below reported state-of-the-art results.The cited best testing-set Dice performances are 0.8839 for BraTS18 whole tumor and 0.80 for WMH17.