Deep Auxiliary Learning via Bilevel Feature Optimization for Label-Efficient Medical Image Segmentation

Fuping Wu; Xi Jia; Le Zhang; Jiahua Li; Rima Mustafa; Thomas E. Nichols; Bartlomiej W. Papiez

cs.CVOct 18, 2026

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 ·

Abstract

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(\cdot;\theta)$ built from bidirectional communication modules (BCMs), followed by a main-task decoder $h_{m}$ for segmentation and an auxiliary-task decoder $h_{a}$ . 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 $\oplus$ (step 2, Eq. ( 2a )–( 2b )). Conv Block: convolution, batch normalization, and ReLU; $C^{l}$ : number of channels at layer $l$ .

BCM	BLO	Auxiliary task: Boundary detection				Auxiliary task: Distance-map prediction
		Dice (%) $\uparrow$		ASSD (pixels) $\downarrow$		Dice (%) $\uparrow$		ASSD (pixels) $\downarrow$
		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.

Single auxiliary task
Auxiliary task(s)	REFUGE (3 seeds)				ACDC (5 folds)
	Dice (%) $\uparrow$		ASSD (pixels) $\downarrow$		Dice (%) $\uparrow$			ASSD (mm) $\downarrow$
	OC	OD	OC	OD	RV	Myo	LV	RV	Myo	LV
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 (%) $\uparrow$		ASSD (HD95) $\downarrow$		Dice (%) $\uparrow$		ASSD (HD95) $\downarrow$
	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).

Panel of Figure 7 — 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 (%) $\uparrow$			ASSD (HD95) $\downarrow$			Dice (%) $\uparrow$			ASSD (HD95) $\downarrow$
	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 (%) $\uparrow$		ASSD (HD95) $\downarrow$		Dice (%) $\uparrow$		ASSD (HD95) $\downarrow$
	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 (%) $\uparrow$		ASSD (pixels) $\downarrow$
Method	(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.

(a) SSL without auxiliary task (results cited from [ 52 ] )
Method	Dice (%) $\uparrow$
Method	RV	Myo	LV	Average
Independent	67.92	74.97	83.22	75.37
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.

Using 4 labeled scans for training
Method	Auxiliary task: boundary detection						Auxiliary task: distance-map prediction
	Dice (%) $\uparrow$			ASSD (mm) $\downarrow$			Dice (%) $\uparrow$			ASSD (mm) $\downarrow$
	RV	Myo	LV	RV	Myo	LV	RV	Myo	LV	RV	Myo	LV
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\times$
Training memory	6.8 GB	9.9 GB	$1.46\times$
Time / iteration	0.68 s	3.81 s	$5.6\times$
Iterations to best	5040	2880	$0.57\times$
Wall-clock to best	0.95 h	3.05 h	$3.2\times$
Inference latency	$\approx$ 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.

(a) CCM , U-Net-L (1024 ch., $\approx$ 31.05 M params)
Methods	Dice (%) $\uparrow$				ASSD (mm) $\downarrow$
Methods	RV	Myo	LV	Avg	RV	Myo	LV	Avg
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.

(a) CCM , U-Net-S (128 ch., $\approx$ 0.49 M params)
Methods	Dice (%) $\uparrow$		ASSD (pixels) $\downarrow$
Methods	OC	OD	OC	OD
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

Medical image segmentation delineates organs, lesions, and other structures in scans to support measurement, diagnosis, treatment planning, and research.

See all 788 papers

SSS: Semi-Supervised SAM-2 with Efficient Prompting for Medical Imaging Segmentation
Jun 10, 2025Hongjie Zhu, Xiwei Liu, Rundong Xue +5Promptable Image Segmentation Semi-Supervised Medical Image Segmentation
SWUN · MBZUAI · XJTU +2
SemiSAM-O1: How far can we push the boundary of annotation-efficient medical image segmentation?
Apr 27, 2026Yichi Zhang, Le Xue, Bichun Xu +6Image Segmentation Medical Imaging Foundation Models
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
RADA: Region-Aware Dual-encoder Auxiliary learning for Barely-supervised Medical Image Segmentation
Apr 13, 2026Shuang Zeng, Boxu Xie, Lei Zhu +6Medical Image Segmentation 3D Medical Image Segmentation

Jun 10, 2025J/K move · Enter open · S save

SSS: Semi-Supervised SAM-2 with Efficient Prompting for Medical Imaging Segmentation

Hongjie Zhu, Xiwei Liu, Rundong Xue, Zeyu Zhang+4

SWUN · MBZUAI · XJTU+2

Promptable Image Segmentation Semi-Supervised Medical Image Segmentation Medical Image Segmentation Semi-Supervised Learning Segment Anything Models

Open paper AskPDF arXiv

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.

Abstract

Medical Image Segmentation

Explore similar work

SSS: Semi-Supervised SAM-2 with Efficient Prompting for Medical Imaging Segmentation

SemiSAM-O1: How far can we push the boundary of annotation-efficient medical image segmentation?

RADA: Region-Aware Dual-encoder Auxiliary learning for Barely-supervised Medical Image Segmentation