Source-linked AI summary
Few-shot Medical Image Segmentation using a Global Correlation Network with Discriminative Embedding
Liyan Sun, Chenxin Li, Xinghao Ding, Yue Huang, Guisheng Wang, Yizhou Yu
TL;DR
Limited medical annotations make unseen-organ segmentation difficult for conventional supervised models. The paper proposes episodically trained GCN-DE, combining efficient support-query global correlation with discriminative embedding, and reports state-of-the-art accuracy across CT and MRI experiments.
Problem
Conventional segmentation models require extensive annotations and can overfit or generalize poorly when unseen organs have only a few reference labels.
Method
GCN-DE uses episodic support-query training, efficient long- and short-range global correlation, and discriminative embedding for same-class attraction and different-class separation.
Results
Experiments on CT and MRI report state-of-the-art few-shot medical image segmentation accuracy, with GCN-DE achieving the best accuracy for three of four organs and the highest mean DC score.
Takeaways & Limitations
Global correlation, discriminative regularization, and slice selection each improve segmentation performance in the reported experiments.
Abstract
from arXiv · showhide
Despite deep convolutional neural networks achieved impressive progress in medical image computing and analysis, its paradigm of supervised learning demands a large number of annotations for training to avoid overfitting and achieving promising results. In clinical practices, massive semantic annotations are difficult to acquire in some conditions where specialized biomedical expert knowledge is required, and it is also a common condition where only few annotated classes are available. In this work, we proposed a novel method for few-shot medical image segmentation, which enables a segmentation model to fast generalize to an unseen class with few training images. We construct our few-shot image segmentor using a deep convolutional network trained episodically. Motivated by the spatial consistency and regularity in medical images, we developed an efficient global correlation module to capture the correlation between a support and query image and incorporate it into the deep network called global correlation network. Moreover, we enhance discriminability of deep embedding to encourage clustering of the feature domains of the same class while keep the feature domains of different organs far apart. Ablation Study proved the effectiveness of the proposed global correlation module and discriminative embedding loss. Extensive experiments on anatomical abdomen images on both CT and MRI modalities are performed to demonstrate the state-of-the-art performance of our proposed model.
I. INTRODUCTION
Few-shot medical image segmentation addresses limited annotations by transferring knowledge from annotated support images to unseen organ classes. GCN-DE uses efficient global correlation and discriminative embedding to exploit spatial regularity and separate organ representations.
- Motivation: Full supervision requires many labeled image pairs, which are difficult to obtain for specialized medical segmentation tasks.Precise organ and lesion segmentation supports diagnosis, treatment, and prognosis.
- Few-shot setting: Few-shot segmentation episodically learns feature correlations between query images and annotated support images to segment unseen classes.The model transfers knowledge from seen classes using one or a few labeled support images.
- Medical-image regularity: Medical foregrounds exhibit consistent intensity, morphology, and structure, but query-support spatial nonalignment makes correlation modeling challenging.The supplied passage specifically contrasts medical-image regularity with larger natural-image intra-class variation.
- Motivation: Prototype approaches discard contextual or shape information, while sSE attention requires spatial overlap between support and query foregrounds.Global correlation is proposed to connect spatially distant objects with the same semantic identity.
- Proposed direction: GCN-DE computes long- and short-range foreground correlations and pulls same-class features together while pushing different-class features apart.Experiments report improved segmentation from the nonlocal module, discriminative regularization, and slice selection, with state-of-the-art accuracy.
II. RELATED WORKS
Prior work established convolutional medical segmentation and few-shot transfer, while nonlocal architectures modeled long-range dependencies. The paper positions its contribution against limitations in unseen-organ generalization and prior correlation strategies.
- Medical image segmentation: Convolutional medical segmentation methods have been applied to brain parcellation, tumor, multi-organ, and liver segmentation.These applications span several tissue, organ, and lesion segmentation tasks.
- Positioning: The proposed architecture combines a support/query network with sSE attention, global correlation, and discriminative embedding.The support image and one-class annotation form a two-channel support input.
- Medical image segmentation: When an unseen organ must be segmented from few reference images, conventional segmentation models can overfit and generalize poorly.This motivates methods that address scarce annotations for target classes.
- Few-shot learning: Few-shot learning methods use episodic or metric-learning strategies to transfer knowledge from annotated examples to new tasks.The related work includes memory augmentation, model-agnostic meta-learning, and Siamese-network approaches.
- Global correlation: Nonlocal and criss-cross architectures capture long-range dependencies, but prior nonlocal models overlook dependencies between objects of the same class.Related medical applications include efficient feature extraction and channel-information exchange.
III. METHODS
The task separates annotated training classes from unseen testing classes and uses episodic support-query training for few-shot generalization. Volumetric CT and MRI data are handled through 2D slice matching.
- Problem definition: A grayscale medical image and multiclass annotation are drawn from a dataset, with disjoint training and testing class sets.The model receives one or a few labeled reference images for an unseen testing class.
- Training scheme: Episodic training samples densely annotated support-query pairs from training classes to supervise transfer to unseen classes.A different class can be sampled in the next episode, and multiclass labels are converted into binary labels.
- Problem definition: Each image presents only a subset of the available training classes, so the number of labels in an image does not exceed the training-class count.The excerpt identifies background as class 0 in the training-class set.
B. Global Correlation Network
The global correlation network combines support-guided attention with efficient short- and long-range correlation between support and query features. It enhances query foregrounds while reducing the quadratic cost of full spatial correlation.
- Network architecture: The support branch concatenates a support image with its foreground annotation, while support features guide query processing through sSE attention and global correlation.The two branches have similar structure but different parameters, and discriminative embedding uses query and support features.
- Global correlation: Global correlation operates on concatenated query-support features to weight representations according to relations among spatial positions.The correlation function compares positions in the concatenated feature and weights the representation mapped by g.
- Efficiency motivation: Full nonlocal correlation has quadratic complexity in image height and width, limiting its use on large feature maps.The paper also notes that small deep semantic maps distort spatial information and limit nonlocal relations.
- Efficient module: The efficient module approximates global correlation by processing long-range and short-range dependencies separately, then merging their outputs with a query-feature shortcut.Long-range processing samples fixed-step pixels, whereas short-range processing divides features into local subregions.
- Implementation: The correlation module projects features into lower-dimensional embeddings, restores channel dimensionality, and applies a residual connection.The supplied implementation passages describe the θ, ϕ, g, ω, unvec, and residual operations.
- Visualization: Foreground organs in query features are enhanced after global correlation for left kidney, right kidney, and spleen examples.The module is applied at shallow scales where localization features are preserved.
C. Discriminative Embedding
The discriminative embedding regularizes support and query features so same-organ representations cluster while different-organ representations remain separated, helping suppress false positives during unseen-organ segmentation.
- C. Discriminative Embedding: Medical-organ contexts form regular latent-space clusters because foreground appearance and surrounding-organ intensity and position vary consistently across scans.This regularity motivates discriminative feature organization for few-shot segmentation.
- C. Discriminative Embedding: The learned embedding is intended to improve segmentation of unseen organ classes and suppress false positives in other major-organ regions.These false positives can diminish the utility of few-shot segmentation in medical imaging practice.
- C. Discriminative Embedding: The embedding regularizer draws features from the same foreground class closer and keeps features from different foreground classes farther apart.It uses intra-class and inter-class L2 distances with a non-negative thresholding operation.
- C. Discriminative Embedding: The regularization is imposed on support and query features at the network backend to preserve discriminability for pixel-wise class prediction.Applying it to both feature sets links the support-query representation used during few-shot prediction.
D. Training Scheme
Training uses episodic support-query pairs from the training dataset, with segmentation losses combined with the discriminative embedding loss to update the model.
- D. Training Scheme: Each training episode samples a class and two annotated images, which form a support-query pair for few-shot learning.The support and query images contain binary annotations for the sampled class and background.
- D. Training Scheme: The support and query images contain the same annotated classes, enforcing structural similarity between them.The number of annotated classes is kept equal across the pair.
- D. Training Scheme: The model is trained for N episodes with T iterations per episode, updating predictions from annotated support and query images by gradient descent.The query prediction at each location is used within the training procedure.
- D. Training Scheme: The combined loss uses Dice and binary cross-entropy terms together with the discriminative embedding loss.These components are combined into the overall training objective.
IV. IMPLEMENTATIONS
The GCN-DE model is trained episodically with different epoch counts for MRI and CT datasets using SGD.
- IV. IMPLEMENTATIONS: Training takes 25 epochs on MRI and 40 epochs on CT, with 25 episodes in each epoch.Each episode contains an annotated support-query image pair with equal numbers of classes.
- IV. IMPLEMENTATIONS: Stochastic gradient descent is used with learning rate 1e-2, implemented in PyTorch on one Titan XP with 12G memory.
A. Datasets
Evaluation uses abdomen MRI and CT datasets with five-fold cross-validation, while the ablation study measures the effects of discriminative embedding and global correlation using DC scores.
- A. Datasets: The MRI evaluation uses CHAOS, containing 20 annotated 3D T2-SPIR scans covering liver, left kidney, right kidney, and spleen.Five-fold cross-validation uses 16 scans for training and 4 for testing in each fold.
- A. Datasets: The CT evaluation uses 30 annotated scans from the MICCAI 2015 Multi-Atlas Abdomen Labeling challenge, with five-fold cross-validation.Each fold uses 24 training scans and 6 testing scans.
- A. Datasets: The ablation study removes discriminative embedding to form GCN and then replaces global correlation with spatial SE to form the baseline without the global-correlation module.DC scores are reported to assess these components.
- A. Datasets: CT testing designates one of six testing scans as support and averages results from the remaining five query scans.
B. Evaluation Metrics
The study evaluates segmentation with Dice Coefficient scores and uses 2D slice matching for volumetric CT and MR images. An ablation compares the effects of discriminative embedding and global correlation.
- Dice Coefficient scores evaluate segmentation accuracy, with higher values approaching 1 indicating more accurate predictions.The metric compares the segmentation prediction with the ground truth.
- CT and MR volumes are evaluated through 2D slice matching for volumetric image segmentation.
- Removing discriminative embedding produces GCN, while replacing global correlation with spatial SE produces the baseline model.
- The ablation reports improved Dice Coefficient scores from adding both discriminative embedding and the global correlation module.
- GCN-DE reduces false positives over GCN for CT spleen segmentation, while GCN improves true-positive predictions over the baseline.
D. Comparison with State-of-the-art Methods
GCN-DE is compared with established few-shot segmentation methods on four abdominal organs across CT and MRI. It achieves the strongest overall quantitative performance, with qualitative advantages in false-positive reduction and organ shape fidelity.
- The comparison evaluates GCN-DE against OSLSM, co-FCN, PANet, SG-ONE, and the medical-image method SE-FSS.The models use authors’ publicized codes with adjusted parameters and full training.
- GCN-DE achieves the best prediction accuracy for spleen, right kidney, and left kidney on both CT and MRI.
- GCN-DE achieves second-best liver accuracy on CT and the highest mean Dice Coefficient across the four organs.The mean score outperforms the second-best method by a large margin.
- Qualitatively, GCN-DE produces fewer false-positive predictions than other compared models on CT and MRI.
- PANet and SG-ONE fail to reproduce precise target-organ shapes, whereas SE-FSS and co-FCN produce better results than those methods.
- The modified episodic-training baseline outperforms SE-FSS, attributed to a better support-query image pairing approach.