Deep Auxiliary Learning via Bilevel Feature Optimization for Label-Efficient Medical Image Segmentation
Authors: Fuping Wu, Xi Jia, Le Zhang, Jiahua Li, Rima Mustafa, Thomas E. Nichols, Bartlomiej W. Papiez
Organizations: Big Data Institute, Nuffield Department of Population Health, University of Oxford · Imperial College London · Nuffield Department of Population Health, University of Oxford · University of Birmingham · Big Data Institute, Nuffield Department of Population Health, and Centre for Integrative Neuroimaging, FMRIB, Nuffield Department of Clinical Neurosciences, University of Oxford
Auxiliary learning (AL), which can improve the performance of a primary task by leveraging related auxiliary tasks, is commonly regarded as a special case of multi-task learning. Existing AL approaches typically require manually annotated auxiliary labels, increasing annotation costs, particularly in medical imaging where expert labeling is expensive. In this work, we investigate whether segmentation performance can be improved using auxiliary tasks whose labels are automatically derived from existing segmentation annotations, requiring no additional manual labeling. To address this, we propose a novel auxiliary learning framework, formulated as a bilevel optimization problem. Our method enables effective knowledge transfer between tasks via a feature communication module, enhancing the primary segmentation task. We evaluate the method on three medical image segmentation benchmarks under limited-label settings. The results demonstrate consistent improvements over state-of-the-art methods, achieving superior label efficiency. The code is publicly available at https://github.com/FupingWu90/ALBO.
Figures & tables
Fig. 1 : Visualization of feature maps from the encoder of Independent, Joint, and ALBO for a random REFUGE sample.
Fig. 2 : Overview of ALBO (auxiliary learning via bilevel optimization). (a) Network architecture: a U-Net-shaped shared encoder g(⋅;θ) built from bidirectional communication modules (BCMs), followed by a main-task decoder hm for segmentation and an auxiliary-task decoder ha . The auxiliary label (e.g., a boundary or distance map) is selected from an auxiliary task pool and is derived automatically from the ground-truth segmentation. All parameters are trained with the bilevel optimization (BLO) of Sec. III-B . (b) Conventional convolutional module (CCM), i.e., two stacked Conv Blocks, used in place of the BCM in the ablation study. (c) BCM: the main (blue, top) and auxiliary (green, bottom) branches first learn task-specific features with two Conv Blocks each (step 1, Eq. ( 1a )–( 1b )), and then exchange features through cross-branch Conv Blocks fused by ⊕ (step 2, Eq. ( 2a )–( 2b )). Conv Block: convolution, batch normalization, and ReLU; Cl : number of channels at layer l .
BCM
BLO
Auxiliary task: Boundary detection
Auxiliary task: Distance-map prediction
Dice (%) ↑
ASSD (pixels) ↓
Dice (%) ↑
ASSD (pixels) ↓
OC
OD
OC
OD
OC
OD
OC
OD
✗
✗ ( Independent )
74.77
91.18
21.94
14.51
74.77
91.18
21.94
14.51
✗
✗ ( Joint )
74.65
90.93
22.08
13.79
74.46
90.55
20.72
13.84
✓
✗ ( Independent )
81.54
93.17
10.28
7.74
81.54
93.17
10.28
7.74
✓
✗ ( Joint )
82.91
93.09
9.62
7.86
81.41
93.00
10.90
7.87
TABLE I : Ablation study on the REFUGE 2018 dataset using U-Net as the backbone with 128 (U-Net-S) bottom channels. In the ‘BCM’ column, ✗ means using the CCM module instead. In ‘BLO’ column, ✗ ( Method ) means using Method to train the model instead of BLO. OC: optic cup; OD: optic disc. Best and second-best values in each column are highlighted in red and blue , respectively.
Fig. 3 : The four auxiliary-task labels used in our study, shown for one representative sample per dataset (REFUGE, ACDC, Hippocampus). From left to right: input image, ground-truth segmentation, and the boundary, signed-distance, saliency, and Canny-edge auxiliary labels. The boundary and distance labels are derived from the segmentation mask, the saliency label is its Gaussian-smoothed foreground, and the Canny label is derived from the image intensity.
Auxiliary task(s)
REFUGE (3 seeds)
ACDC (5 folds)
Dice (%) ↑
ASSD (pixels) ↓
Dice (%) ↑
ASSD (mm) ↓
OC
OD
OC
OD
RV
Myo
LV
RV
Myo
LV
Single auxiliary task
Boundary
73.58
82.42
18.76
28.32
70.74
80.21
87.48
4.29
1.49
2.19
Distance
65.36
78.28
24.31
34.62
68.46
78.68
86.23
5.03
1.79
2.55
Saliency
67.71
81.73
18.71
27.03
65.96
76.87
85.32
5.53
1.83
2.96
TABLE II : Effect of auxiliary-task selection and combination, with 10 labeled images. All variants use the identical ALBO bilevel optimizer and backbone; only the auxiliary set changes. Best and second-best per column are in red and blue .
Method
400 labeled images
10 labeled images
Dice (%) ↑
ASSD (HD95) ↓
Dice (%) ↑
ASSD (HD95) ↓
OC
OD
OC
OD
OC
OD
OC
OD
Independent+BCM
81.54 ∗∗
93.17 ∗∗
10.28 ∗∗ (23.24 ∗∗ )
7.74 ∗∗ (18.82 ∗∗ )
67.54 ∗∗
80.91 ∗∗
22.83 ∗∗ (56.33 ∗∗ )
30.50 ∗∗ ( 89.08 )
Joint+BCM
82.91
93.09 ∗∗
9.62 (21.59 ∗∗ )
7.86 ∗∗ (18.18 ∗∗ )
65.01 ∗∗
77.33 ∗∗
23.93 ∗∗ (56.42 ∗∗ )
36.07 ∗∗ (92.71 ∗∗ )
GCS+BCM [ 20 ]
83.39 ∗∗
93.78 ∗∗
9.30 ∗ ( 20.51 )
6.98 ∗∗ ( 16.07 ∗∗ )
64.75 ∗∗
77.30 ∗∗
26.66 ∗∗ (60.92 ∗∗ )
36.74 ∗∗ (97.79 ∗∗ )
AuxiLearn+BCM [ 22 ]
82.75
93.32 ∗∗
9.93 (22.72 ∗∗ )
7.53 ∗∗ (18.64 ∗∗ )
63.53 ∗∗
78.11 ∗∗
27.49 ∗∗ (64.67 ∗∗ )
36.50 ∗∗ (91.98 ∗∗ )
TABLE III : Performance comparison on the REFUGE dataset with 400 and 10 labeled images for training. The number of parameters (# Params) for each model is around 0.5M. Best and second-best values in each column are shown in red and blue , respectively. Each ASSD (pixels) value is followed by the 95% Hausdorff Distance (HD95, pixels) in parentheses. ∗ and ∗∗ indicate a statistically significant difference from ALBO with p<0.05 and p<0.01 , respectively (paired Wilcoxon signed-rank test on the per-image scores averaged over the three seeds).
Fig. 4 : Performance change of all compared methods on the REFUGE and ACDC with varying numbers of training samples.
Fig. 5 : Visualization of segmentation predictions from compared methods trained with different numbers of labeled samples. Here, each sample for visualization is the median case of the Independent method (trained with full data) in terms of Dice.
Method
160 labeled scans
10 labeled scans
Dice (%) ↑
ASSD (HD95) ↓
Dice (%) ↑
ASSD (HD95) ↓
RV
Myo
LV
RV
Myo
LV
RV
Myo
LV
RV
Myo
LV
Independent
89.16 ∗∗
89.47 ∗∗
94.01 ∗
1.03 ∗ ( 5.32 )
0.44 ∗ (2.47)
0.63 ∗ (3.46)
67.92 ∗∗
74.97 ∗∗
83.22 ∗∗
5.75 ∗∗ ( 22.87 ∗∗ )
2.21 ∗∗ ( 10.58 ∗∗ )
3.01 ∗∗ (12.95 ∗∗ )
Joint
88.93 ∗∗
89.56 ∗∗
94.19
1.11 ∗∗ (6.01 ∗ )
0.45 ∗∗ (2.42)
0.63 (3.44)
66.22 ∗∗
75.67 ∗∗
81.98 ∗∗
6.25 ∗∗ (25.79 ∗∗ )
2.55 ∗∗ (13.71 ∗∗ )
3.42 ∗∗ (14.69 ∗∗ )
GCS [ 20 ]
89.20 ∗∗
89.49 ∗∗
94.14 ∗∗
1.04 ∗∗ (5.51)
0.46 ∗∗ (2.43 ∗ )
0.61 ∗∗ (3.39)
66.70 ∗∗
75.14 ∗∗
82.08 ∗∗
6.35 ∗∗ (27.03 ∗∗ )
2.71 ∗∗ (13.14 ∗∗ )
3.60 ∗∗ (15.52 ∗∗ )
AuxiLearn [ 22 ]
89.09 ∗∗
89.47 ∗∗
94.15
1.03 (5.46)
0.45 ∗ (2.37)
0.56 ( 2.97 )
65.53 ∗∗
74.00 ∗∗
80.78 ∗∗
6.82 ∗∗ (27.46 ∗∗ )
2.55 ∗∗ (12.71 ∗∗ )
3.38 ∗∗ (13.72 ∗∗ )
TABLE IV : Performance comparison on the ACDC dataset using 160 and 10 labeled scans for training. Best and second-best values in each column are shown in red and blue , respectively. Each ASSD (mm) value is followed by the 95% Hausdorff Distance (HD95, mm) in parentheses. ∗ and ∗∗ indicate a statistically significant difference from ALBO with p<0.05 and p<0.01 , respectively (paired Wilcoxon signed-rank test over all samples).
Method
208 labeled scans
20 labeled scans
Dice (%) ↑
ASSD (HD95) ↓
Dice (%) ↑
ASSD (HD95) ↓
HH
HB
HH
HB
HH
HB
HH
HB
Independent
88.56 ∗∗
86.87
0.39 ∗∗ (1.75 ∗∗ )
0.36 ∗∗ (1.54)
85.02 ∗∗
82.81 ∗∗
0.66 ∗∗ (2.79 ∗∗ )
0.52 ∗∗ (2.20 ∗∗ )
Joint
88.89
86.95
0.35 (1.50)
0.36 (1.71)
85.63
83.31 ∗∗
0.51 ( 2.03 )
0.48 ∗∗ (1.94 ∗∗ )
GCS [ 20 ]
88.74
86.99
0.36 ∗∗ (1.65 ∗∗ )
0.35 (1.53)
85.76
83.74 ∗∗
0.53 ( 2.07 )
0.47 ∗ (1.94)
AuxiLearn [ 22 ]
89.02
86.99
0.36 (1.58)
0.36 (1.60)
84.93 ∗∗
83.63 ∗∗
0.72 ∗∗ (3.94 ∗∗ )
0.48 ∗∗ (2.05 ∗∗ )
TABLE V : Performance comparison on Hippocampus with 208 and 20 labeled scans for training. Best and second-best values in each column are shown in red and blue , respectively. Each ASSD (mm) value is followed by the 95% Hausdorff Distance (HD95, mm) in parentheses. ∗ and ∗∗ indicate a statistically significant difference from ALBO with p<0.05 and p<0.01 , respectively (paired Wilcoxon signed-rank test over all samples).
Method
#Lab
#Unl
#Par
Time
Dice (%) ↑
ASSD (pixels) ↓
(img)
(img)
(M)
(ms)
OC
OD
OC
OD
1 labeled image
Labeled-only (few-shot)
BRBS [ 49 ]
1
0
3.93
7
34.61 ∗∗
50.31 ∗∗
68.01 ∗∗
79.98 ∗∗
DSPNet [ 50 ]
1
0
59.28
77
21.74 ∗∗
28.95 ∗∗
/
/
UPRENet [ 51 ]
1
0
51.60
781
45.47 ∗∗
69.29 ∗∗
/
31.74 ∗∗
TABLE VI : Comparison with recent few-shot and semi-supervised segmentation methods on REFUGE with 1, 4 and 10 labeled images. #Lab/#Unl: number of labeled/unlabeled training images; #Par: parameters (M); Time: per-image inference latency (ms) at the corresponding number of labeled images. ALBO and Independent use no unlabeled data. ∗ and ∗∗ indicate a statistically significant difference from ALBO at p<0.05 and p<0.01 (paired Wilcoxon signed-rank test on per-image scores over three seeds). Best and second-best per column within each block are in red and blue . “ / ” denotes no meaningful ASSD is achieved overall; these failures are reflected by the low Dice.
Method
Dice (%) ↑
RV
Myo
LV
Average
Independent
67.92
74.97
83.22
75.37
(a) SSL without auxiliary task (results cited from [ 52 ] )
PL
67.03
77.87
84.88
76.59
AdvSemiSeg
73.11
80.18
86.47
79.92
CPS
69.31
76.37
84.08
76.59
TABLE VII : Performance comparison for SSL methods on the ACDC dataset using 5-fold cross-validation with 10 labeled and 150 unlabeled scans for training in each fold. Best and second-best values in each column are shown in red and blue , respectively.
Method
Auxiliary task: boundary detection
Auxiliary task: distance-map prediction
Dice (%) ↑
ASSD (mm) ↓
Dice (%) ↑
ASSD (mm) ↓
RV
Myo
LV
RV
Myo
LV
RV
Myo
LV
RV
Myo
LV
Using 4 labeled scans for training
Independent
75.77
76.94
85.12
4.26
2.10
2.63
75.77
76.94
85.12
4.26
2.10
2.63
Joint
73.98
77.25
83.99
5.08
2.03
2.47
74.88
76.87
85.76
3.98
1.79
2.16
CAGrad [ 14 ]
49.09
58.81
71.36
9.13
3.71
5.15
42.11
53.04
67.10
11.76
5.44
6.44
TABLE VIII : Performance comparison on the ACDC dataset using DINOv3 as the backbone with 4 and 10 labeled training scans. Best and second-best values in each column are shown in red and blue, respectively.
Fig. 6 : ALBO with a convolutional (U-Net, blue) versus a Mamba (VSS, orange) backbone, using boundary detection as the auxiliary task. (a,b) REFUGE with 10 and 400 labeled images; (c,d) ACDC with 10 and 160 labeled scans. “ ∗ ” and “ ∗∗ ” denote p<0.05 and p<0.01 (paired Wilcoxon signed-rank test, Mamba vs. CNN); “n.s.” = not significant.
Metric
Independent
ALBO
Ratio
Parameters
0.504 M
0.504 M
1.0×
Training memory
6.8 GB
9.9 GB
1.46×
Time / iteration
0.68 s
3.81 s
5.6×
Iterations to best
5040
2880
0.57×
Wall-clock to best
0.95 h
3.05 h
3.2×
Inference latency
≈ 9–14 ms (same net)
—
TABLE IX : Computational and resource overhead of ALBO versus the single-level baseline ( Independent ) using the identical U-Net, on REFUGE ( N=10 ).
Fig. 7 : Training loss convergence on REFUGE ( N=10 ). ALBO’s main-task (segmentation) loss decreases faster and to a lower value than the single-level baseline ( Independent ).
Fig. 8 : Representative failure cases of ALBO under the 10-labeled setting on REFUGE (optic cup), ACDC (abnormal right ventricle), and Hippocampus (posterior). Each row shows the input image, ground truth, and ALBO prediction. The smallest, anatomically atypical, or lowest-contrast structure is frequently omitted and merged into the dominant neighboring structure.
Methods
Dice (%) ↑
ASSD (mm) ↓
RV
Myo
LV
Avg
RV
Myo
LV
Avg
(a) CCM , U-Net-L (1024 ch., ≈ 31.05 M params)
Independent
89.16
89.47
94.01
90.88
1.03
0.44
0.63
0.71
Joint
88.93
89.56
94.19
90.89
1.11
0.45
0.63
0.73
GCS [ 20 ]
89.20
89.49
94.14
90.94
1.04
0.46
0.61
0.71
AuxiLearn [ 22 ]
89.09
89.47
94.15
90.91
1.03
0.45
0.56
0.68
TABLE X : Performance of SOTA AL and MTL methods on the ACDC dataset using U-Net as the backbone and boundary detection as the auxiliary task. We employed five-fold cross-validation, and used 160 labeled scans for training at each fold. U-Net uses 128 (or 1024) bottom channels with either the CCM or BCM module. Models tend to underfit with smaller networks, whereas ALBO is more robust than the compared methods when using the proposed BCM module. Best and second-best values in each column are shown in red and blue , respectively.
Methods
Dice (%) ↑
ASSD (pixels) ↓
OC
OD
OC
OD
(a) CCM , U-Net-S (128 ch., ≈ 0.49 M params)
Independent
74.77
91.18
21.94
14.51
Joint
74.65
90.93
22.08
13.79
GCS [ 20 ]
76.97
92.16
18.58
11.37
AuxiLearn [ 22 ]
75.65
88.94
18.72
17.87
TABLE XI : Performance of compared methods on the REFUGE 2018 dataset using U-Net as the backbone with 128 (or 1024) bottom channels and CCM (or BCM) modules. We conducted all experiments three times with different random seeds, and used 400 labeled images for training at each run. We adopted boundary detection as the auxiliary task. The results show that models tend to overfit with larger networks, while our model performed more robustly than others. Additionally, other models showed better performance when using the proposed BCM module. #Params denotes the parameter count of the model. Best and second-best values in each column are shown in red and blue , respectively.
Medical image segmentation delineates organs, lesions, and other structures in scans to support measurement, diagnosis, treatment planning, and research.
In the era of information explosion, efficiently leveraging large-scale unlabeled data while minimizing the reliance on high-quality pixel-level annotations remains a critical challenge in the field of medical imaging. Semi-supervised learning (SSL) enhances the utilization of unlabeled data by facilitating knowledge transfer, significantly improving the performance of fully supervised models and emerging as a highly promising research direction in medical image analysis. Inspired by the ability of Vision Foundation Models (e.g., SAM-2) to provide rich prior knowledge, we propose SSS (Semi-Supervised SAM-2), a novel approach that leverages SAM-2's robust feature extraction capabilities to uncover latent knowledge in unlabeled medical images, thus effectively enhancing feature support for fully supervised medical image segmentation. Specifically, building upon the single-stream "weak-to-strong" consistency regularization framework, this paper introduces a Discriminative Feature Enhancement (DFE) mechanism to further explore the feature discrepancies introduced by various data augmentation strategies across multiple views. By leveraging feature similarity and dissimilarity across multi-scale augmentation techniques, the method reconstructs and models the features, thereby effectively optimizing the salient regions. Furthermore, a prompt generator is developed that integrates Physical Constraints with a Sliding Window (PCSW) mechanism to generate input prompts for unlabeled data, fulfilling SAM-2's requirement for additional prompts. Extensive experiments demonstrate the superiority of the proposed method for semi-supervised medical image segmentation on two multi-label datasets, i.e., ACDC and BHSD. Notably, SSS achieves an average Dice score of 53.15 on BHSD, surpassing the previous state-of-the-art method by +3.65 Dice. Code will be available at https://github.com/AIGeeksGroup/SSS.
Semi-supervised learning (SSL) has become a promising solution to alleviate the annotation burden of deep learning-based medical image segmentation models. While recent advances in foundation model-driven SSL have pushed the boundary to extremely limited annotation scenarios, they fail to maintain robust competitive performance in complex imaging modalities. In this paper, we propose SemiSAM-O1, an annotation-efficient framework using only one annotated template image for segmentation. SemiSAM-O1 extends the specialist-generalist collaborative learning framework to the extreme one-label setting by fully exploiting the foundation model's feature representation capability beyond its prompting interface. SemiSAM-O1 operates in two stages. In the first stage, the foundation model's encoder extracts dense features from all volumes, and class prototypes derived from the single annotated template are propagated to the unlabeled pool via feature similarity to produce coarse initial pseudo-labels. In the second stage, an iterative training-and-refinement loop progressively improves both the segmentation model and the pseudo-labels over multiple rounds, where each round trains the model from scratch on current pseudo-labels and generates updated predictions with voxel-wise uncertainty estimates. An uncertainty-guided refinement step further leverages the foundation model's global feature space to correct high-uncertainty regions by aggregating labels from their most similar confident neighbors, establishing a virtuous cycle of mutual improvement. Extensive experiments on a wide range of segmentation tasks across different modalities and anatomical targets demonstrate that SemiSAM-O1 significantly narrows the performance gap between one-label semi-supervised learning and full supervision, while significantly reducing the computational overhead of online foundation model inference.
Yichi Zhang, Le Xue, Bichun Xu +6
Artificial Intelligence Innovation and Incubation Institute, Fudan University, Shanghai, China. · Shanghai Academy of Artificial Intelligence for Science, Shanghai, China. · Department of Radiotherapy, Tongji Hospital, School of Medicine, Tongji University, Shanghai, China. +2
Deep learning has greatly advanced medical image segmentation, but its success relies heavily on fully supervised learning, which requires dense annotations that are costly and time-consuming for 3D volumetric scans. Barely-supervised learning reduces annotation burden by using only a few labeled slices per volume. Existing methods typically propagate sparse annotations to unlabeled slices through geometric continuity to generate pseudo-labels, but this strategy lacks semantic understanding, often resulting in low-quality pseudo-labels. Furthermore, medical image segmentation is inherently a pixel-level visual understanding task, where accuracy fundamentally depends on the quality of local, fine-grained visual features. Inspired by this, we propose RADA, a novel Region-Aware Dual-encoder Auxiliary learning pipeline which introduces a dual-encoder framework pre-trained on Alpha-CLIP to extract fine-grained, region-specific visual features from the original images and limited annotations. The framework combines image-level fine-grained visual features with text-level semantic guidance, providing region-aware semantic supervision that bridges image-level semantics and pixel-level segmentation. Integrated into a triple-view training framework, RADA achieves SOTA performance under extremely sparse annotation settings on LA2018, KiTS19 and LiTS, demonstrating robust generalization across diverse datasets.