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

Introduction

Deep learning has achieved remarkable success in medical image segmentation. However, its performance is highly dependent on large-scale annotated datasets, the creation of which is time-consuming and requires domain expertise [1, 2]. This limitation has prompted research into label-efficient learning strategies.

Auxiliary learning (AL), a specialized form of multi-task learning (MTL), has shown strong potential to improve the performance of a primary task under limited supervision by exploiting related auxiliary tasks, and therefore is well-suited to weakly supervised settings such as semi-supervised learning [3]. Unlike MTL, which treats all tasks equally, AL prioritizes a designated main task through auxiliary supervision.

Despite this encouraging progress, especially on natural image benchmarks [4], two key challenges remain unresolved for segmentation. The first is the design of label-efficient auxiliary tasks. Existing methods typically rely on either (i) additional annotated auxiliary tasks such as depth estimation [5], or (ii) self-supervised objectives that benefit from large corpora of unlabeled images [6, 7]. Although effective, these strategies often require additional annotation, auxiliary modalities, or extensive unlabeled datasets, which may be difficult to obtain in medical imaging because of annotation cost and privacy constraints [8]. In this work, we address a fundamental question: Through AL, can models improve the label efficiency of the target (main) task using auxiliary tasks, whose labels are derived solely from the existing data? And if so, how? Previous studies have employed handcrafted or algorithmically generated labels, such as object boundaries, to regularize segmentation outputs [9, 10, 11]. However, these approaches are evaluated only when abundant labeled images are available, and their robustness under limited-label settings is not discussed. Notably, our setting assumes only the existing annotations of the main task; the auxiliary labels are generated from these annotations, and the training set is not expanded with additional unlabeled images. This is the key difference from semi-supervised learning (SSL), which improves label efficiency by exploiting a pool of unlabeled data; the two strategies are therefore complementary, as discussed in § 5-F.3.

The second challenge is to optimize the AL models such that the auxiliary task consistently benefits, rather than interferes with, the primary segmentation task. Simultaneously learning multiple tasks can result in negative transfer, where the performance of the main task deteriorates [12]. Existing approaches primarily attempt to mitigate this either by dynamically adjusting task importance during training [13], or modifying gradients to reduce conflict between tasks [14, 15]. However, the underlying causes of negative transfer remain underexplored, particularly in medical imaging settings where tasks are often highly imbalanced in difficulty and representation [16]. As illustrated in Figure 1, without an appropriate optimization technique, the features learned for boundary detection by methods such as Joint and Independent can be inconsistent with this auxiliary task.

To address these challenges, we propose a new auxiliary learning approach for medical image segmentation in which

This work was supported by the funding from Novo Nordisk.

Fuping Wu (Email: [email protected]) was with Big Data Institute, Nuffield Department of Population Health, University of Oxford. He is now with Imperial College London.

Xi Jia (Email: [email protected]), Jiahua Li and Rima Mustafa are with the Nuffield Department of Population Health, University of Oxford. Le Zhang is with the University of Birmingham.

Thomas E. Nichols is with Big Data Institute, Nuffield Department of Population Health, and Centre for Integrative Neuroimaging, FMRIB, Nuffield Department of Clinical Neurosciences, University of Oxford.

Bartłomiej W. Papież is with Big Data Institute, Nuffield Department of Population Health, University of Oxford. auxiliary labels are automatically derived from the main-task data, thereby eliminating the need for additional manual labeling or dataset expansion. The intuition is that models can benefit from learning auxiliary tasks that capture complementary structural cues (e.g., boundaries or foreground masks) to improve feature representations for the main segmentation task. Building on this idea, we propose ALBO: auxiliary learning via bilevel optimization. Our main contributions are summarized as follows:

• We propose ALBO, an auxiliary learning framework for medical image segmentation that exploits auxiliary tasks whose labels are automatically derived from existing segmentation annotations. A bidirectional communication module enables feature exchange between the primary and auxiliary branches.

• We formulate AL as a bilevel optimization problem and employ a finite-difference approximation to make the optimization tractable during training while enhancing representation learning.

• The proposed method is evaluated on two 2D and one 3D benchmark datasets and three types of backbones (convolutional U-Net, ViT-based DINOv3, and Mamba architecture), demonstrating superior label efficiency.

Related Works

Auxiliary learning (AL)

AL enhances the performance of a main task by leveraging knowledge from related auxiliary tasks [17]. Prior AL research has primarily focused on two fronts: auxiliary task selection and optimization strategies. Several works have investigated how to identify effective auxiliary tasks by modeling task relationships, in which a common strategy is to group tasks based on their positive transfer potential [18], [19]. Another line of research focuses on improving the optimization of AL models. In GCS [20], gradient cosine similarity is employed to assess task-relatedness and to determine whether an auxiliary loss is beneficial for the main task. In [21], gradient similarity is used to guide the adaptive selection of auxiliary tasks in reinforcement learning. In [16], task weights are automatically adjusted by minimizing validation error. In AuxiLearn [22], nonlinear combinations of losses are learned dynamically to improve optimization. In AuxiNash [4], the AL problem is formulated as a generalized bargaining game to resolve task interactions.

Unlike prior work, we propose a framework applicable to multiple auxiliary tasks, rather than methods tailored to a specific task. For example, Fan et al. [23] design a task-specific scheme using auxiliary image registration for one-shot segmentation, which is less flexible in terms of auxiliary task selection. Sobotka et al. [24] require paired multimodal data as auxiliary labels, which falls outside our annotation-free setting. Huang et al. [25] do not study auxiliary learning; instead, they introduce an auxiliary model for expert correction.

Registration-based methods have also been explored for label-efficient segmentation, including atlas-based unsupervised segmentation via contrastive registration that warps a single reference mask (CLMorph [26]) and few-shot segmentation via self-supervised reconstruction (EFS-MedSeg [27]); unlike these, ALBO neither relies on registration nor on unlabeled data, deriving its auxiliary labels for free from the existing main-task labels.

The optimization strategies above steer the auxiliary influence only through the loss or gradient of a shared encoder, by gating the auxiliary gradient with its similarity to the main gradient [20], [21], reweighting tasks by validation error [16], learning a nonlinear loss combination [22], or bargaining over task weights [4]. They therefore neither model an explicit bidirectional exchange of features between the main and auxiliary branches, nor structurally guarantee that the main task is prioritized during optimization. ALBO addresses both limitations: the bidirectional communication module lets the main and auxiliary branches exchange features in both directions (Sec. III-A), and the bilevel formulation casts the main task as the upper level so the auxiliary task is optimized only insofar as it serves the main one (Sec. III-B), using annotation-free auxiliary labels rather than task-specific extra data [23], [24].

Multi-task learning (MTL)

MTL, unlike AL, aims to improve the performance of all tasks simultaneously [28], [29], rather than prioritizing the main task. Based on architectural design strategies, previous methods can be broadly categorized into soft parameter-sharing [30, 31, 32] and hard parameter-sharing [33, 34]. Soft parameter-sharing methods design separate encoders for each task while enforcing constraints on their associated parameters. For example, MTAN [35] employs an attention mechanism among task-specific networks to facilitate the sharing of a general feature pool. Hard parameter-sharing methods, in contrast, primarily focus on finding a local update vector to reduce task conflict. Liu et al. introduce an approach to compute an update vector with equal projections on each task gradient [36]. PCGrad [37] projects each task’s gradient onto the normal plane of other task gradients, combining them to form the final update vector. GradDrop [38] randomly drops task gradients based on the extent of their conflicts. CAGrad [14] minimizes the average loss function and is theoretically proven to converge to the optimal point. Additionally, mixture of experts (MoE) architectures has also been widely adopted in MTL with an example of modeling the global task relationships by adding a generic convolution path to the original MoE structure [39].

Beyond architectural design, a large body of MTL work targets the optimization of shared parameters, which can be grouped into loss-balancing schemes that reweight per-task losses, such as random loss weighting (RLW) [40] and AMTL [41], and gradient-manipulation schemes that resolve conflicting task gradients, such as PCGrad [37], GradDrop [38], CAGrad [14], and Nash-MTL [42]. However, these methods are designed for the symmetric MTL objective of improving all tasks jointly; they treat every task equally and do not privilege a designated main task, so a freely derived auxiliary signal can dilute rather than strengthen the primary segmentation. In contrast, ALBO adopts an asymmetric bilevel formulation that places the main segmentation task at the upper level and updates the auxiliary branch only to benefit it, which is better matched to the label-efficient auxiliary-learning setting studied here. We empirically compare against representative methods from both groups in § 5-E.

Methods

Methods

Because the auxiliary labels here are derived for free from the main-task data, our aim is not to solve the two tasks symmetrically but to make the auxiliary task serve the main one; we therefore take AL as a bilevel problem in which the main task occupies the upper level. Given an image $x$, its corresponding main task label $y_m$, and an auxiliary label $y_a$ automatically generated from $(x,y_m)$, AL aims to learn a prediction model $f$ for the main task with the help of $y_a$. Throughout the paper, we refer to $y_a$ as the auxiliary label and to the task of predicting it as the auxiliary task. Our model is composed of three functional modules: a shared encoder $g(\cdot;\theta)$, a primary (main-task) head $h_m(\cdot;\psi_m)$, and an auxiliary head $h_a(\cdot;\psi_a)$, so that the main- and auxiliary-task predictors are $f=h_m\circ g$ and $f_a=h_a\circ g$, respectively. Because the encoder is built from the bidirectional communication module (BCM) in Sec. III-A, it carries two interacting branches, and its parameters split into a main branch $\theta_m$ and an auxiliary branch $\theta_a$, i.e., $\theta=(\theta_m,\theta_a)$; the two heads carry parameters $\psi_m$ and $\psi_a$. Grouped by the role they play in optimization, the main (upper-level) variables are $\{\theta_m,\theta_a,\psi_m\}$ and the auxiliary (lower-level) variables are $\{\theta_a,\psi_a\}$; the shared branch $\theta_a$ is the variable through which the two levels interact.

As shown in Figure 2(a), in this work, our main task is to segment anatomical structures. We employ the U-Net [43] as the encoder. To better exploit auxiliary supervision, we propose a novel auxiliary learning framework via bilevel optimization, ALBO, which aligns with our motivation from both network architecture and optimization perspectives, as detailed in the following sections.

Bidirectional Communication Module

Generally, suppose that at the $l$-th layer of the encoder, we extract the main and auxiliary features from image $x$, denoted by $z_m^l\in\mathbb{R}^{\frac{C^l}{2}\times H\times W}$ and $z_a^l\in\mathbb{R}^{\frac{C^l}{2}\times H\times W}$, respectively. $C^l$, $H$, and $W$ denote the number of channel, height, and width of the feature map, respectively.

To enable better auxiliary feature exploitation by the main task, we design the BCM to transfer knowledge from the auxiliary branch to the main one and vice versa, thereby exploiting the synergistic potential of reciprocally improving both tasks through information passing. As shown in Figure 2(a), we replace the conventional convolutional module (CCM) (Figure 2(b)) with the BCM (Figure 2(c)) as the core module of the encoder. Specifically, given $z_m^l$ and $z_a^l$, we generate $z_m^{l+1}$ and $z_a^{l+1}$ at the $l$-th layer with the following two steps.

step 1 (Task-specific feature learning):

step 2 (Feature communication):

where ConvB $\coloneqq$ ReLU $\circ$ BN $\circ$ Conv denotes convolutional block consisting of a convolutional layer (Conv), a batch normalization (BN), and a ReLU activation.

For simplicity, we denote the process from $z_m^l$ to $z_m^{l+1}$ as $g_m^{l+1}$, which is the composition of operations in Eqs. (1a), (1b), and (2a) with learnable parameters $\theta_m^{l+1}$, i.e., $z_m^{l+1}=g_m^{l+1}(z_m^l|\hat{\theta}_m^{l+1},\theta_m^{l+1})$. Similarly, we have $z_a^{l+1}=g_a^{l+1}(z_a^l|\hat{\theta}_a^{l+1},\theta_a^{l+1})$ for the auxiliary branch with learnable parameters $\theta_a^{l+1}$.

By setting $z_m^0=z_a^0=x$, and using BCM in the encoder, the predictions of the main and auxiliary tasks become

where $1\leq l\leq L$ is the encoding layer index.

Formulation for AL using BCM

When using BCM to build the encoder, the loss function for the main task is formulated as $\mathcal{L}_m(\theta_m,\theta_a,\psi_m)=\mathbb{E}_x[\ell_m(\hat{y}_m,y_m)]$, and the loss function for the auxiliary task as $\mathcal{L}_a(\theta_m,\theta_a,\psi_a)=\mathbb{E}_x[\ell_a(\hat{y}_a,y_a)]$, where $\ell_{m/a}$ measures the dissimilarity between the prediction and ground truth, $\theta_m=(\theta_m^0,\ldots,\theta_m^L)$ and $\theta_a=(\theta_a^0,\ldots,\theta_a^L)$. In this work, we adopt the cross-entropy loss [44] for $\ell_{m/a}$.

For AL, we consider the main task as our target, meaning we can combine the auxiliary and main branches as a whole for main task learning when minimizing $\mathcal{L}_m$, rather than treating the features from the auxiliary task as given conditions. Hence, we formulate the AL problem as

To solve the above bilevel optimization problem, we employ an approximation scheme by alternately updating the parameters. Specifically, we have

where $\alpha$ and $\beta$ are learning rates.

Optimization intuition. Rather than minimizing $\mathcal{L}_m$ at the current auxiliary branch, the main update (7) evaluates it at the look-ahead branch $\theta_a'=\theta_a-\alpha\nabla_{\theta_a}\mathcal{L}_a$, i.e., one gradient step ahead on the auxiliary task. Differentiating through this step introduces a second-order coupling term, $-\alpha\nabla_{\theta_m,\theta_a}^2\mathcal{L}_a\cdot\nabla_{\theta_a'}\mathcal{L}_m$. This term pairs the sensitivity of the main loss to the shared auxiliary parameters, $\nabla_{\theta_a'}\mathcal{L}_m$, with the curvature of the auxiliary objective. It thus measures, to first order, how the one-step auxiliary update perturbs the shared representation. This perturbation, in turn, alters the main objective. Intuitively, it reinforces auxiliary updates whose direction agrees with the main-task descent and discounts those that would increase $\mathcal{L}_m$, providing a built-in safeguard against negative transfer. This is the mechanism by which ALBO utilizes the auxiliary task to benefit, rather than compete with, the main task.

The outer update Eq. (5b) requires the second-order gradient $\nabla_{\theta_m/\theta'_a}\mathcal{L}_m(\theta_m,\theta'_a,\psi_m)$, which is memory- and time-consuming. To avoid explicitly forming the Hessian, we adopt a finite-difference approximation [45] that needs only two additional forward–backward passes, giving the final update gradient

where $\theta_a^\pm=\theta_a\pm\epsilon\cdot\nabla_{\theta'_a}\mathcal{L}_m(\theta_m,\theta'_a,\psi_m)$ and $\epsilon$ is a small positive scalar. The complete step-by-step derivation of Eq. (9) is given in Appendix § 1.

In terms of complexity, this finite-difference scheme computes the required Hessian–vector product using only first-order gradients (one look-ahead step and two extra forward–backward passes per outer iteration), so each iteration retains the $O(P)$ time and memory of single-level training in the number of shared parameters $P$, avoiding both the $O(P^2)$ explicit Hessian and the heavier second-order automatic differentiation; the bilevel update thus adds only a constant-factor overhead, confined to training.

Experiments

Datasets and pre-processing

We evaluated the proposed method with three segmentation tasks: (1) Optic cup/disc (OC/OD) segmentation using the REFUGE 2018 dataset [46]. REFUGE consists of 400 color fundus photography images for training, 400 for validation, and 400 for testing. All images were cropped to a size of $512 \times512$. For REFUGE, each model was trained three times with different random initializations to ensure reliability. (2) Cardiac segmentation using the ACDC dataset [47], which contains 200 short-axis cardiac MR-cine images with the labels of the right ventricle (RV), myocardium (Myo), and left ventricle (LV). All slices were cropped to a size of $128 \times128$. For ACDC, each model was evaluated with five-fold cross-validation. For each fold, we used 40 cases for validation, and the remaining 160 cases for the training data selection. (3) The Hippocampus dataset from Medical Segmentation Decathlon Challenge [48] consists of 260 3D scans with annotations of the hippocampus head (HH) and body (HB) available. All volumes were resized to $48 \times64 \times48$ with zero-padding. For the Hippocampus, we also employed five-fold cross-validation. For each fold, we used 52 scans for validation, and the remaining 208 scans for the training set selection. For evaluation, the average Dice and the average symmetric surface distance (ASSD) were reported.

Implementation details

We implemented nine representative methods for comparison: (1) Independent: this is the conventional independent learning for the main segmentation tasks without involving any auxiliary tasks; (2) Four AL methods, including Joint, GCS [20], AuxiLearn [22], and AuxiNash [4]. (3) Four MTL methods, including CAGrad [14], Nash [42], RLW [40], and AMTL [41]. Concretely, Independent trains only the main task, minimizing $\mathcal{L}_m(\theta_m,\theta_a,\psi_m)$ over $\{\theta_m,\theta_a,\psi_m\}$ with no auxiliary loss, whereas Joint minimizes the equal-weighted sum $\mathcal{L}_m+\mathcal{L}_a$ over all parameters $\{\theta_m,\theta_a,\psi_m,\psi_a\}$ by conventional single-level gradient descent, i.e., without the bilevel look-ahead of ALBO. For comparison, all models were evaluated using the same U-Net shaped network (as shown in Figure 2(a)), and run on a single 40GB NVIDIA A100 GPU. We employed 3D models for the Hippocampus and 2D models for the others. For simplicity, we denote the U-Net with 128 and 1024 bottom channels as U-Net-S and U-Net-L, respectively. For all datasets, we trained U-Net-S using the SGD optimizer with the initial learning rate of $\alpha=\beta=0.1$ and decay of 0.9 every 200 iterations. All models were trained for a maximum of 15,000 iterations with a batch size of 20.

Ablation study

We conducted an ablation study to evaluate the effectiveness of the two key components of ALBO, i.e., the BCM module and the bilevel optimization (BLO) algorithm, using the whole REFUGE dataset. Since the effectiveness of different automatically generated auxiliary tasks is unclear, we evaluated two widely used tasks, i.e., boundary detection and distance-map prediction [3], for comparison.

As presented in Table I, the proposed method achieved the best results across most metrics when both BCM and BLO are enabled, and remained robust across different auxiliary tasks. Removing BLO resulted in a noticeable performance degradation for both auxiliary tasks. For example, with distance-map prediction as the auxiliary task, the OC Dice decreased from 83.99% to 81.41%, while with boundary detection the OD Dice decreased from 94.32% to 93.09%. Without BCM, the parameters $\theta_m$ and $\theta_a$ in Eq. (4) merged into a shared encoder, i.e., $\theta$, and the performance of the penultimate row degraded significantly (paired t-test: $p < 0.01$). In this configuration the bilevel update acts on a fully shared encoder without a dedicated auxiliary branch, so the boundary auxiliary task is imposed directly on the shared main-task features. This can over-constrain the higher-contrast optic disc representation, reducing its Dice, while the smaller optic cup still benefits from the auxiliary signal. These results suggest that BLO is most effective when combined with BCM, since branch separation and bidirectional communication help route the auxiliary signal without causing negative transfer.

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
$\times$	$\times$ (Independent)	74.77	91.18	21.94	14.51	74.77	91.18	21.94	14.51
$\times$	$\times$ (Joint)	74.65	90.93	22.08	13.79	74.46	90.55	20.72	13.84
$\checkmark$	$\times$ (Independent)	81.54	93.17	10.28	7.74	81.54	93.17	10.28	7.74
$\checkmark$	$\times$ (Joint)	82.91	93.09	9.62	7.86	81.41	93.00	10.90	7.87
$\times$	$\checkmark$	77.73	88.97	14.71	17.58	78.29	91.44	14.06	11.90
ALBO (Ours)		82.30	94.32	9.36	6.38	83.99	93.52	8.90	7.34

Comparing the first four rows of Table I further shows that BCM alone improves the performance of Independent and Joint on REFUGE. However, it does not mean that BCM alone can boost the performance in all cases. For example, experiments show that on the ACDC dataset, all nine compared methods obtained inferior results when using the BCM module instead of the CCM module. Remarkably, the utility of BCM can be exploited robustly by BLO. More detailed discussion about the effect of BCM and CCM can be found in Appendix II.

Effect of auxiliary-task selection and combination

We study how the choice and number of auxiliary tasks affect segmentation performance by evaluating four auxiliary tasks, namely boundary, distance, saliency, and Canny edge (Figure 3), both individually and in combination.

Table II reports results with 10 labeled samples from REFUGE and ACDC. First, the boundary task is the strongest and most robust single auxiliary task: it attained the best score on all six ACDC metrics and the best OC Dice on REFUGE, and no other single task was competitive across both datasets. The remaining single tasks were weaker and less consistent. The Canny edge was competitive on ACDC but markedly worse on REFUGE, while the distance and saliency tasks were the weakest overall. Second, combining auxiliary tasks was not systematically beneficial. On REFUGE, using all four tasks gave the best OD Dice and the lowest ASSD, marginally ahead of the boundary task; on ACDC, however, every combination was inferior to the single boundary task, and the four-task set fell well below it.

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
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
Canny	66.24	78.57	20.50	33.72	69.66	79.98	87.22	4.60	1.40	2.27
Two auxiliary tasks
Boundary+Canny	72.18	80.89	18.74	29.45	70.36	78.99	87.14	5.13	1.54	2.20
Boundary+Distance	66.82	75.97	24.67	38.75	68.08	76.34	84.43	5.27	1.88	2.68
Three auxiliary tasks
Boundary+Canny+Distance	68.42	79.60	18.52	31.29	68.10	79.47	86.85	4.41	1.56	2.18
Boundary+Distance+Saliency	68.10	81.34	19.11	27.03	66.53	75.57	83.88	4.98	1.97	3.08
Four auxiliary tasks
Boundary+Distance+Saliency+Canny	72.81	83.49	16.02	23.67	66.85	76.17	84.59	4.76	2.00	2.75

These results indicate that the bilevel feature communication mechanism extracts most of the informative structure from a single, well-matched auxiliary task. Adding more tasks provides, at best, a marginal and dataset-dependent gain, and may harm accuracy. Thus, increasing the number of auxiliary tasks is not intrinsically beneficial. This justifies our choice of a single, mask-derived boundary auxiliary task: it is the most robust option across both modalities, and avoids the task-balancing burden and negative-transfer risk of multi-auxiliary training.

Label efficiency study

We evaluated the label efficiency of ALBO by training the model with different numbers of samples and using boundary detection as the auxiliary task.

Comparison on REFUGE.

For OC and OD segmentation, we randomly selected 10, 20, 40, 80, and 160 retinal images for training. Since all compared methods achieved significantly better performance ($p < 0.01$) with the BCM than with the CCM when training with the whole 400 labeled images, we used BCM-embedded U-Net-S for all methods to provide a fairer comparison of label efficiency.

Table III presents the average performance on OC and OD of all methods, trained with either 400 or 10 labeled images. With sufficient training data, the performance of all methods was close. When the labeled data were reduced to 10 samples, ALBO demonstrated the best robustness, whereas all other AL and MTL methods performed worse than Independent. We additionally report the 95% Hausdorff Distance (HD95) and, using a paired Wilcoxon signed-rank test, mark in Table III the metrics that differ significantly from ALBO. With 400 labeled images, although the Dice scores of all methods were similar, ALBO attained the best HD95 on both OC and OD, indicating more accurate boundary delineation. With only 10 labeled images, ALBO yielded the best Dice and ASSD on both structures, and its Dice and ASSD improvements over the competing methods were statistically significant ($p < 0.01$).

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).

Method	400 labeled images				10 labeled images
	Dice (%)↑ OC	Dice (%)↑ OD	ASSD (HD95) ↓ OC	ASSD (HD95) ↓ OD	Dice (%)↑ OC	Dice (%)↑ OD	ASSD (HD95) ↓ OC	ASSD (HD95) ↓ 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**)
AuxiNash +BCM [4]	81.37**	93.07**	10.36** (23.38**)	7.81** (19.75**)	59.74**	73.52**	38.84** (87.78**)	44.29** (111.63**)
RLW+BCM [40]	81.05**	93.42**	10.31** (22.61**)	7.44** (18.77**)	67.48**	73.86**	20.17** (47.79)	43.89** (101.17**)
Nash+BCM [42]	80.97**	93.11**	10.76** (23.93**)	7.84** (19.40**)	63.64**	74.24**	24.93** (63.93**)	42.37** (110.77**)
AMTL+BCM [41]	81.16**	93.66**	10.49** (23.68**)	7.07** (17.74**)	63.61**	77.94**	24.13** (59.05**)	35.77** (101.58**)
CAGrad+BCM [14]	82.55	93.62**	9.60 (21.15**)	7.15** (16.98**)	67.20**	79.77**	22.23** (54.56**)	31.25** (92.15**)
ALBO (Ours)	82.30	94.32	9.36 (19.81)	6.38 (15.44)	73.58	82.42	18.76 (48.76)	28.32 (87.66)

Figure 4 (a) illustrates the trend of Dice scores averaged on OC and OD obtained by all methods as the number of labeled training samples decreased. Overall, ALBO is more robust to changes in training data size. When the number of labeled training samples is $N \le80$, ALBO significantly outperforms the competing methods in average Dice ($p < 0.05$). A similar trend can be observed for ASSD. The top two rows of Figure 5 illustrate predictions from each method on a representative test sample, trained with 400 and 10 retinal images, respectively. These results show that ALBO achieves the best performance when the labeled training data are limited.

Comparison on ACDC.

For cardiac segmentation, besides using the full 160 scans in each fold, we also randomly selected 10, 20, 40, and 80 scans for training. In contrast to REFUGE, the models tended to underfit on ACDC, and CCM provided more favorable results than BCM. Accordingly, we used U-Net-L for all compared methods to ensure a fair comparison.

Table IV presents the average performance on RV, Myo, and LV of all methods, trained with either 160 or 10 labeled scans. ALBO consistently obtained the best results across all metrics, particularly when the labeled training data was limited to 10. Table IV also reports the HD95 and the significance of each method relative to ALBO (using a paired Wilcoxon signed-rank test). ALBO obtained the best Dice and ASSD on every structure and the best or second-best HD95; with only 10 labeled scans it significantly outperformed all competing methods on all three metrics for every structure ($p < 0.01$). Similarly, Figure 4 (b) illustrates the trend of Dice of all methods averaged on all structures with varying numbers of labeled training data. As the number of labeled scans decreases, the performance gap between ALBO and the competing methods widens, demonstrating the superior label efficiency of ALBO. Specifically, ALBO significantly outperforms others in average Dice across all training sample sizes ($p < 0.05$). A similar trend is observed for ASSD. The middle two rows of Figure 5 illustrate representative predictions from each method, trained with 160 and 10 scans, respectively. ALBO produces more accurate predictions for the sample, especially when labeled data were limited.

Comparison on Hippocampus.

For hippocampus segmentation, we used the full set of 208 scans and randomly selected 20 scans for training, and employed U-Net-L for all compared methods. Table V presents the average results on the HH and HB. Most methods outperformed the Independent when only 20 labeled scans were available, and ALBO remained robust in this low-label setting. Table V additionally reports the HD95 and the per-metric significance relative to ALBO (using a paired Wilcoxon signed-rank test). Because this benchmark was close to saturation, all methods performed comparably and many pairwise differences were not statistically significant; ALBO nevertheless remained among the most competitive methods, achieving the best HD95 and the tied-best ASSD on the HH with 208 scans and the best HB Dice with 20 scans, while a few of the strongest auxiliary-learning methods such as AuxiNash and Nash were only marginally better on individual metrics. The bottom two rows of Figure 5 illustrate representative predictions from each method, trained with 208 and 20 scans, respectively. ALBO produces anatomically consistent segmentations.

Discussion

Can ALBO learn task-consistent feature representations?

In Figure 1, we visualize encoder feature maps from the main and auxiliary branches for a representative retinal image using Independent, Joint and ALBO. Compared with Independent and Joint, ALBO produces feature representations that are more aligned with the corresponding tasks: the main-branch features are more focused on the segmentation target, while the auxiliary-branch features more clearly capture boundary-related structure. This suggests that ALBO encourages more task-consistent feature learning through bidirectional communication and bilevel optimization.

How does ALBO compare with few-shot methods?

ALBO is an auxiliary-learning method and does not require unlabeled data, which distinguishes it from semi-supervised learning (SSL). To position ALBO in the low-data setting, we compared it against three recent few-shot medical segmentation methods, including BRBS [49], DSPNet [50] and UPRENet [51], on REFUGE with 1, 4 and 10 labeled images. Since these methods can also exploit unlabeled image, each was run in two modes: a labeled-only (few-shot) mode, and a semi-supervised mode that additionally used the remaining 390–399 unlabeled training images. Table VI reports accuracy together with model size and inference time.

Method	#Lab (img)	#Unl (img)	#Par (M)	Time (ms)	Dice (%)↑ OC	Dice (%)↑ OD	ASSD (pixels)↓ OC	ASSD (pixels)↓ 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**
Independent	1	0	0.50	14	39.81**	62.44**	/	59.95**
ALBO (Ours)	1	0	0.50	14	60.42	66.90	24.44	54.28
Semi-supervised (uses unlabeled data)
BRBS [49]	1	399	3.93	7	72.76**	77.25**	18.75**	44.66**
DSPNet [50]	1	399	59.28	77	30.19**	63.48**	/	30.38**
UPRENet [51]	1	399	51.60	781	22.30**	65.45*	/	47.91**
4 labeled images
Labeled-only (few-shot)
BRBS [49]	4	0	3.93	7	41.86**	52.80**	68.85**	85.75**
DSPNet [50]	4	0	59.28	370	66.02**	76.57*	20.88**	23.41**
UPRENet [51]	4	0	51.60	3132	66.61**	85.22**	23.04**	20.41**
Independent	4	0	0.50	14	56.29**	74.68**	21.82**	42.39**
ALBO (Ours)	4	0	0.50	14	68.96	77.99	18.24	37.90
Semi-supervised (uses unlabeled data)
BRBS [49]	4	396	3.93	7	71.90**	74.49**	20.09**	49.28**
DSPNet [50]	4	396	59.28	370	54.84**	72.57**	26.06**	26.43**
UPRENet [51]	4	396	51.60	3132	40.69**	67.70**	40.62**	48.57**
10 labeled images
Labeled-only (few-shot)
BRBS [49]	10	0	3.93	7	57.63**	68.68**	36.02**	64.80**
DSPNet [50]	10	0	59.28	632	69.00**	84.75**	19.96**	16.05**
UPRENet [51]	10	0	51.60	7850	68.89**	85.64**	24.78**	23.48**
Independent	10	0	0.50	14	67.54**	80.91**	22.83**	30.50
ALBO (Ours)	10	0	0.50	14	73.58	82.42	18.76	28.32
Semi-supervised (uses unlabeled data)
BRBS [49]	10	390	3.93	7	66.84**	75.23**	27.61**	41.58**
DSPNet [50]	10	390	59.28	632	67.78**	81.10**	24.64**	25.82
UPRENet [51]	10	390	51.60	7850	57.22**	78.39**	35.20**	37.49**

The three baselines cover the two dominant paradigms of few-shot medical segmentation. BRBS is a registration-based method following the registration-to-segmentation paradigm: it trains a deformable registration network on unlabeled images and generates a large pseudo-labeled training set by warping the few labeled atlases onto them, after which a compact segmentation network is trained on the synthetic set. It uses two small from-scratch U-Nets (3.93M parameters) and runs only the segmentation network at test time, so its inference is fast ($\sim7$ ms/image) and independent of the label budget. DSPNet and UPRENet are prototypical episodic methods with a COCO-pretrained ResNet-101 encoder, which explains their much larger parameter counts of 52–59M. They embed the labeled support images, extract class prototypes, and segment each query by nearest-prototype matching; DSPNet refines pooling-based prototypes into detail-preserving cluster prototypes, whereas UPRENet additionally filters global and local prototypes by an uncertainty estimation and fuses two augmented support branches through an evidential module.

Crucially, their $N$-shot setting requires extracting prototypes from every labeled support image and averaging one prediction per support image at inference, so their latency grows roughly linearly with the number of labeled images, namely from 77 ms and 0.78 s at one label to 0.63 s and 7.8 s at ten labels for DSPNet and UPRENet, respectively. ALBO, BRBS and Independent instead run a single fixed network with constant latency ($\sim14$ ms for ALBO), giving ALBO up to a near three-order-of-magnitude speed advantage at larger label budgets.

These architectural differences help explain the observed performance. BRBS is designed to exploit unlabeled images of magnitude longer inference time. By comparison, ALBO attained the highest OC Dice among all labeled-only methods across all label budgets. With 10 labeled images it also surpassed every few-shot and SSL variant on OC segmentation, including those that consume hundreds of unlabeled images, while using only 0.50M parameters and $\sim14$ ms per image inference.

We emphasize that ALBO addresses a different problem from SSL: ALBO improves label efficiency through auxiliary labels derived from the existing annotations rather than through unlabeled data, making the two approaches complementary rather than competing. As shown in the following subsection, ALBO can be combined with SSL methods to achieve further improvements.

Can ALBO be used for semi-supervised learning (SSL)?

For SSL, the performance of models trained solely on labeled data often serves as a crucial foundation. This is particularly evident in methods that utilize pseudo-labeling, where the predictions on unlabeled data significantly influence model updates. Given that ALBO demonstrates high label efficiency, it is expected to have the potential for SSL applications.

Generally, the objective of SSL can be formulated as $\mathcal{L}_{sup}(\mathcal{D}_l) + \lambda\mathcal{R}(\mathcal{D}_{ul})$, where $\mathcal{D}_l$ and $\mathcal{D}_{ul}$ denote the labeled and unlabeled data sets, $\mathcal{L}_{sup}$ represents the supervised loss, $\mathcal{R}$ denotes a regularization term on unlabeled data, and $\lambda$ is a balancing parameter. Similar to Eq. (4), ALBO can be extended to SSL as follows,

Therefore, SSL methods that are capable of computing $\mathcal{R}_{m/a}$ can be integrated with ALBO. As a proof of concept, we use the standard Pseudo-label (PL) method [53] to validate this integration with ALBO on the ACDC dataset, denoted as PL + ALBO. We follow the evaluation protocol established in MERU [52], and compare our approach with six SSL methods: PL [53], AdvSemiSeg [54], CPS [55], CrossTeaching [56], URPC [57], and MERU.

Table VII-(a) summarizes the performance of these SSL methods without the incorporation of any auxiliary tasks. Among them, only AdvSemiSeg, CPS, and MERU consistently outperformed the baseline Independent across all three anatomical structures, whereas PL degrades performance on RV. Table VII-(b) presents the results where boundary detection is used as the auxiliary task. In this setting, PL + ALBO improved performance, while the PL + Joint model was consistently worse than the original PL model in terms of Dice across all structures. Table VII-(c) shows the results when the distance-map prediction is used as the auxiliary task. Both PL + Joint and PL + ALBO performed better than when boundary detection was used. Notably, PL + ALBO achieved the highest average Dice score of 81.69% among all evaluated methods. These results suggest that ALBO can complement SSL methods effectively. Future work may extend this analysis to additional SSL frameworks to further evaluate its generality and effectiveness.

Method	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
CrossTeaching	72.04	72.91	80.86	75.27
URPC	60.29	66.30	74.67	67.09
MERU	74.90	81.11	88.35	81.45
(b) PL using boundary detection as the auxiliary task
PL + Joint	63.77	77.04	84.56	75.12
PL + ALBO (Ours)	74.76	80.02	85.91	80.23
(c) PL using distance-map prediction as the auxiliary task
PL + Joint	69.23	78.69	86.02	77.98
PL + ALBO (Ours)	76.85	80.71	87.50	81.69

Relationship between ALBO and SSL. Both ALBO and SSL target the limited-annotation regime, but they rely on different resources. ALBO assumes only the existing annotations of the main task and generates its auxiliary labels from them without expanding the training set, whereas SSL improves label efficiency by additionally exploiting a pool of unlabeled images. This leads to three practical observations. (i) Complementarity. Because the two strategies draw on different sources of information, they are complementary rather than competing. As shown in Table VII, combining ALBO with PL improved the average Dice of PL from 76.59% to 80.23% (boundary detection) and 81.69% (distance-map prediction), whereas PL + Joint yielded only 75.12% and 77.98%, respectively. (ii) Exploiting unlabeled data. ALBO itself does not require unlabeled data, but it can exploit them when they are available: as formulated in Eq. (7), any SSL regularizer that can be computed on unlabeled images can be added to the main and auxiliary levels, while the bilevel structure that prioritizes the main task remains unchanged. (iii) Practical scenarios. ALBO is preferable when a sizable and representative pool of unlabeled images is unavailable or cannot be used, for example, in small single-center cohorts or under data-sharing and privacy restrictions that limit the collection of additional images. Using only the 10 labeled ACDC scans, ALBO reached an average Dice of 79.48% (Table IV), which is higher than four of the six SSL methods in Table VII that additionally used 150 unlabeled scans. When abundant unlabeled data from the same distribution are available, SSL is the natural choice, and ALBO can then serve as a stronger supervised component within it, as in PL + ALBO.

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

160 labeled scans

10 labeled scans

Dice (%) ↑ RV

Dice (%) ↑ Myo

Dice (%) ↑ LV

ASSD (HD95) ↓ RV

ASSD (HD95) ↓ Myo

ASSD (HD95) ↓ LV

Dice (%) ↑ RV

Dice (%) ↑ Myo

Dice (%) ↑ LV

ASSD (HD95) ↓ RV

ASSD (HD95) ↓ Myo

ASSD (HD95) ↓ LV

Independent

89.16**

89.47**

94.01*

1.03* (5.32)

0.44* (2.47)

0.62* (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** (13.71**)

3.38** (13.72**)

AuxiNash [4]

89.06**

89.85

94.31

1.14* (6.21*)

0.42 (2.29)

0.54 (3.07)

66.86**

75.62**

82.37**

6.96** (29.23**)

2.41** (11.95**)

3.24** (14.24**)

RLW [40]

89.18**

89.43**

94.22

1.04* (5.65)

0.45* (2.56)

0.59 (3.51)

67.33**

74.58**

82.38**

4.63** (27.37**)

2.55** (12.16**)

3.17** (13.55**)

Nash [42]

89.21

89.82

94.32

1.01 (5.51)

0.42 (2.25)

0.58 (3.39)

68.40**

76.17**

82.35**

6.19** (26.81**)

2.68** (12.96**)

3.45** (14.54**)

AMTL [41]

89.02**

89.85

94.24

1.15* (6.47**)

0.42 (2.23)

0.58 (3.14*)

66.88**

75.54**

81.53**

6.76** (27.97**)

2.52** (12.74**)

3.23** (13.58**)

CAGrad [14]

88.93**

89.70

94.18

1.17* (6.63**)

0.42* (2.25)

0.54 (3.05)

66.88**

75.29**

83.79**

6.52** (28.26**)

2.76** (13.76**)

2.78* (12.86*)

ALBO (Ours)

89.91

89.87

94.62

0.95 (5.35)

0.40 (2.23)

0.47 (2.69)

70.74

80.21

87.48

4.29 (16.03)

1.49 (6.69)

2.19 (9.94)

Can ALBO be applied with pre-trained visual foundation models (FMs)?

The preceding experiments evaluated the effectiveness of ALBO when using a convolutional U-Net trained from scratch. Given the remarkable transferability of FMs, we next investigated whether auxiliary tasks can further enhance segmentation performance when ALBO is built on pre-trained models. To this end, we use the recently released DINOv3 (ViT-B/16) model [58] as the backbone for evaluation.

As DINOv3 is a pre-trained ViT-based model, directly reformulating its transformer blocks into a BCM-style architecture may lead to catastrophic forgetting of the learned knowledge. Instead, given $z_m^l$ and $z_a^l$, we generate $z_m^{l+1}$ and $z_a^{l+1}$ at the $l$-th ViT block ($l = 1, \ldots, 12$) through a dual-path formulation:

where $\operatorname{LoRABlock}^l$ denotes the $l$-th transformer block wrapped with LoRA [59] for efficient fine-tuning, and $\tau_{m/a}$ are learnable scalars controlling feature communication via additive fusion. Two BCM blocks are appended after the modified DINOv3 backbone to produce the final predictions.

We compared ALBO against Independent, Joint, and CAGrad, using both boundary detection and distance-map prediction as auxiliary tasks on the ACDC dataset with 4 and 10 labeled scans, respectively. Following the standard LoRA fine-tuning protocol, all methods were optimized using the AdamW optimizer with an initial learning rate of 0.001 and evaluated under five-fold cross-validation. The results are summarized in Table VIII.

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
ALBO (Ours)	76.08	79.11	87.14	3.82	1.50	2.28	76.50	78.03	86.79	3.65	1.80	2.83
Using 10 labeled scans for training
Independent	83.00	82.53	89.70	1.85	1.25	1.56	83.00	82.53	89.70	1.85	1.25	1.56
Joint	81.87	82.31	88.96	2.20	1.24	1.75	83.92	81.94	89.45	1.86	1.20	1.51
CAGrad [14]	66.69	70.94	79.80	5.15	2.55	3.59	63.74	70.09	81.24	4.85	2.29	2.94
ALBO (Ours)	83.25	83.08	90.28	1.75	1.23	1.83	83.59	82.13	89.42	1.74	1.41	1.95

With 10 labeled scans, ALBO achieved performance comparable to Independent, particularly when boundary detection was used as the auxiliary task. Under the more challenging 4-label setting, ALBO outperformed both Independent and Joint, especially under the boundary detection setting. In contrast, CAGrad exhibited a substantial performance degradation compared with Independent.

Can ALBO be extended to the Mamba architecture?

The preceding experiments evaluated ALBO on convolutional (U-Net) and transformer (DINOv3) backbones. We further examined whether ALBO extends to the Mamba [60] (state-space) architecture. To this end, we replaced the convolutional blocks in BCM with Vision-Mamba (VSS) blocks [61, 62].

We compared the Mamba- and convolution-based ALBO using boundary detection as the auxiliary task, under both low and high label settings: REFUGE with 10 and 400 labeled images, and ACDC with 10 and 160 labeled scans. Figure 6 reports the per-structure Dice with paired Wilcoxon significance tests. On ACDC the Mamba backbone achieved performance comparable to the convolutional one. With 10 labeled scans, the two models were statistically indistinguishable on all structures (mean foreground Dice: 79.6% vs. 79.5%), and with 160 labeled scans the difference remained below 1%. On REFUGE the Mamba backbone performed worse under the

10-label setting, but the gap shrank as the number of labels increased, from 6.5% mean foreground Dice at 10 labels to 1.7% at 400 labels.

These results indicate that ALBO generalizes naturally to Mamba-based architectures. The larger performance gap observed on REFUGE under extremely limited supervision is likely attributable to the backbone rather than the proposed ALBO, since state-space vision models are generally more data-hungry than convolutional networks and are not consistently superior across vision tasks. The key finding is that ALBO remains effective across convolutional, transformer, and state-space architectures.

Computational and resource analysis

We analyze the cost of the bilevel training scheme. Table IX reports the training and inference cost of ALBO and of a single-level baseline (Independent) on REFUGE with $N = 10$ labeled images, measured on one NVIDIA RTX 6000 (24 GB) with batch size 20.

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)		—
Inference memory/params	identical		—

The additional computational cost arises from the bilevel optimization, which introduces one virtual inner update and a finite-difference approximation of the Hessian–vector product [45], requiring approximately two additional forward–backward passes per iteration. Consequently, the per-iteration training time increases by approximately 5.6×, and the peak training memory increases by 46%. This overhead is incurred only during training. At inference, ALBO performs a single forward pass through the same segmentation network as the baseline, resulting in identical model size, memory footprint, and real-time inference latency ($< 15$ ms per image). Furthermore, ALBO converges in fewer iterations (2880 versus 5040) while achieving a lower training loss (Figure 7), reducing the wall-clock time to the best model from the expected 5.6× to 3.2×. Overall, training completes in approximately three hours on a single GPU with a peak memory footprint below 10 GB.

Failure cases and limitations

Figure 8 shows representative failure cases of ALBO under the 10-labeled setting on all three datasets. Most errors occur on the smallest, lowest-contrast, or anatomically atypical target, and ALBO tends to miss the target structure and assign its pixels to the neighboring structure (cup→disc, posterior→anterior, RV→myocardium/background). On REFUGE, the OC is the dominant failure: its mean Dice (73.58%) is well below the OD (82.42%), and 10.4% of test images have an OC Dice below 0.5, including complete misses when the cup–disc boundary exhibits weak contrast. On ACDC, the RV is the most difficult structure and the most severe errors occur in the abnormal-RV pathology subgroup, where ten labeled scans are insufficient to capture the large anatomical variability. In addition, because the cardiac volumes are processed slice-by-slice using a 2D model, the lack of through-plane context further causes inter-slice inconsistencies on basal and apical slices. On the Hippocampus dataset, failures are concentrated in the posterior subdivision on thin head and tail slices, where the boundary between the anterior and posterior regions is ambiguous. In the most difficult cases, the posterior Dice decreases to 0.09.

Despite its promising performance, ALBO has several limitations.

• (i) Training cost. The bilevel update requires a virtual inner step and a finite-difference Hessian–vector approximation, raising the per-iteration cost to $\approx5.6\times$ and the peak training memory by 46% relative to single-level training (Table IX). This cost is confined to the one-time, offline training stage, and inference is a single forward pass with identical model size and latency; nevertheless, it makes training more resource-intensive than conventional baselines.

• (ii) Limited gains with abundant labels. The advantage of ALBO is most pronounced in the low-label regime; as the number of labeled samples grows, the gap to competing methods narrows. On near-saturated benchmarks (e.g., Hippocampus, and the $N=160/400$ settings of ACDC/REFUGE) ALBO remains competitive but is no longer uniformly best, and several differences do not reach statistical significance (Table III–Table V).

• (iii) Performance on rare structures. Like other methods trained with very limited supervision, ALBO cannot span the full anatomical variability from ten labeled scans; it therefore has a ceiling on rare, atypical, or very small, low-contrast targets, such as the optic cup at $N=10$.

• (iv) Cross-domain generalization. Our validation is conducted within each dataset; robustness under distribution shift across sites, scanners, or modalities is not evaluated here. As ALBO is a label-efficiency method that is orthogonal to domain-generalization and domain-adaptation techniques, integrating it with such techniques and assessing cross-domain robustness remain open.

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**)
AuxiNash [4]	89.10	87.40**	0.34 (1.49**)	0.33** (1.44)	86.24*	84.03	0.54* (2.22)	0.45 (1.78)
RLW [40]	88.70*	87.09	0.37** (1.73**)	0.34 (1.54)	86.35	83.58**	0.48 (2.10)	0.47** (1.95**)
Nash [42]	88.96	86.75	0.33 (1.39)	0.35 (1.49*)	86.40	84.05**	0.51 (2.27)	0.43 (1.71)
AMTL [41]	88.52**	86.92	0.37** (1.59*)	0.35 (1.57)	86.04	83.74**	0.55 (2.49)	0.46* (1.87*)
CAGrad [14]	88.81	86.54**	0.35* (1.64**)	0.37** (1.73)	85.91	83.54**	0.55 (2.64)	0.47** (2.01**)
ALBO (Ours)	89.01	87.08	0.33 (1.38)	0.33 (1.54)	86.25	84.51	0.52 (2.39)	0.44 (1.76)

• (v) Combination with other label-efficient methods. ALBO improves label efficiency solely through auxiliary labels automatically derived from the labeled data, which is complementary to both few-shot and semi-supervised learning. We show that ALBO is competitive with such methods (§ 5-F.2) and can be paired with a semi-supervised objective (§ 5-F.3); however, how best to unify auxiliary labels, few-shot priors, and unlabeled data within a single framework remains largely unexplored.

Future work will focus on improving robustness to small and anatomically atypical structures, extending ALBO to stronger 3D architectures and complementary auxiliary tasks, evaluating cross-domain generalization, and integrating ALBO more closely with few-shot and semi-supervised learning frameworks.

Conclusion

In this work, we investigated AL for medical image segmentation, focusing on auxiliary tasks whose labels can be automatically derived from existing segmentation annotations without additional manual labeling. We proposed ALBO, a bilevel auxiliary learning framework that combines bidirectional feature communication with bilevel optimization to improve the primary segmentation task. Extensive experiments on three medical imaging benchmarks demonstrated that ALBO consistently improves label efficiency, particularly in low-label settings, while generalizing across convolutional, transformer, and state-space backbones. These results suggest that automatically generated auxiliary tasks, when coupled with appropriate architectural and optimization strategies, provide an effective and practical approach to label-efficient medical image segmentation.

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
AuxiNash [4]	89.06	89.85	94.31	91.07	1.14	0.42	0.54	0.70
RLW [40]	89.18	89.43	94.22	90.94	1.04	0.45	0.59	0.70
Nash [42]	89.21	89.82	94.32	91.12	1.01	0.42	0.58	0.67
AMTL [41]	89.02	89.85	94.24	91.04	1.15	0.42	0.58	0.72
CAGrad [14]	88.93	89.70	94.18	90.94	1.17	0.42	0.54	0.71
(b) BCM, U-Net-L (1024 ch., ≈32.04 M params)
Independent	89.49	89.68	94.44	91.20	0.99	0.43	0.58	0.67
Joint	89.62	89.94	94.52	91.36	1.04	0.42	0.53	0.67
GCS [20]	89.87	89.89	94.70	91.49	0.96	0.42	0.48	0.62
AuxiLearn [22]	89.53	89.82	94.43	91.26	1.01	0.42	0.52	0.65
AuxiNash [4]	89.80	90.09	94.64	91.51	1.11	0.41	0.52	0.68
RLW [40]	89.56	89.83	94.35	91.25	0.97	0.44	0.60	0.67
Nash [42]	89.62	89.98	94.48	91.36	1.01	0.42	0.54	0.66
AMTL [41]	89.41	89.88	94.50	91.26	1.05	0.43	0.59	0.69
CAGrad [14]	89.70	90.18	94.87	91.58	0.99	0.38	0.44	0.61
ALBO (Ours)	89.91	89.87	94.62	91.47	0.95	0.40	0.47	0.60
(c) CCM, U-Net-S (128 ch., ≈0.49 M params)
Independent	87.15	87.31	93.08	89.18	1.30	0.56	0.72	0.86
Joint	87.36	87.34	93.28	89.33	1.32	0.56	0.73	0.87
GCS [20]	86.75	87.09	92.99	88.94	1.46	0.59	0.72	0.93
AuxiLearn [22]	86.96	86.98	92.82	88.92	1.38	0.60	0.77	0.91
AuxiNash [4]	87.26	87.84	93.55	89.55	1.35	0.57	0.71	0.88
RLW [40]	86.95	87.09	93.12	89.05	1.47	0.57	0.69	0.91
Nash [42]	86.82	87.67	93.12	89.20	1.46	0.55	0.70	0.90
AMTL [41]	85.76	87.29	92.58	88.54	1.80	0.60	0.96	1.12
CAGrad [14]	86.75	87.26	93.23	89.08	1.47	0.58	0.79	0.95
(d) BCM, U-Net-S (128 ch., ≈0.50 M params)
Independent	85.52	84.58	91.25	87.12	1.61	0.72	0.87	1.06
Joint	85.77	85.47	91.51	86.77	1.70	0.72	0.94	1.22
GCS [20]	84.02	85.49	91.82	87.11	1.99	0.65	0.89	1.18
AuxiLearn [22]	86.10	85.52	92.22	87.94	1.44	0.70	0.90	1.01
AuxiNash [4]	86.61	86.83	93.04	88.83	1.58	0.60	0.75	0.98
RLW [40]	84.93	84.83	90.65	86.80	1.82	0.73	0.99	1.18
Nash [42]	85.42	86.24	92.07	87.91	1.76	0.70	0.90	1.12
AMTL [41]	86.08	86.02	91.81	87.97	1.57	0.67	0.86	1.03
CAGrad [14]	85.42	85.82	91.29	87.51	1.65	0.69	0.97	1.10
ALBO (Ours)	88.26	88.49	93.91	90.22	1.18	0.49	0.56	0.74

Appendix I: Derivation of the finite-difference approximation

We provide the full derivation of the look-ahead gradient summarized in Eq. (9). Recall the inner (auxiliary) update $\theta'_a = \theta_a - \alpha\nabla_{\theta_a}\mathcal{L}_a(\theta_m,\theta_a,\psi_a)$. Applying the chain rule through $\theta'_a$ and approximating the resulting Hessian–vector products by central finite differences yields, for the main branch,

Methods	Dice (%) $\uparrow$		ASSD (pixels) $\downarrow$
	OC	OD	OC	OD
(a) CCM, U-Net-S (128 ch., $\approx$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
AuxiNash [4]	72.58	86.40	21.18	22.93
RLW [40]	78.02	92.79	18.24	10.64
Nash [42]	76.10	91.42	19.05	12.25
AMTL [41]	73.53	90.37	22.47	15.04
CAGrad [14]	78.07	92.60	17.99	10.56
(b) BCM, U-Net-S (128 ch., $\approx$0.50 M params)
Independent	81.54	93.17	10.28	7.74
Joint	82.91	93.09	9.62	7.86
GCS [20]	83.39	93.78	9.30	6.98
AuxiLearn [22]	82.75	93.32	9.93	7.53
AuxiNash [4]	81.37	93.07	10.36	7.81
RLW [40]	81.05	93.42	10.31	7.44
Nash [42]	80.97	93.11	10.76	7.84
AMTL [41]	81.16	93.66	10.49	7.07
CAGrad [14]	82.55	93.62	9.60	7.15
ALBO (Ours)	82.30	94.32	9.36	6.38
(c) CCM, U-Net-L (1024 ch., $\approx$31.05 M params)
Independent	71.98	84.87	25.30	25.14
Joint	69.91	82.68	30.38	28.95
GCS [20]	71.89	86.47	23.98	22.71
AuxiLearn [22]	49.43	60.13	60.43	70.48
AuxiNash [4]	44.60	60.52	70.18	68.21
RLW [40]	72.65	82.35	21.75	28.08
Nash [42]	74.27	84.69	19.05	24.61
AMTL [41]	71.92	83.97	28.92	25.47
CAGrad [14]	71.05	83.86	24.13	25.66
(d) BCM, U-Net-L (1024 ch., $\approx$32.04 M params)
Independent	77.92	90.82	13.89	11.41
Joint	71.43	89.73	20.32	12.66
GCS [20]	75.78	90.63	15.33	12.76
AuxiLearn [22]	44.59	65.12	64.02	57.98
AuxiNash [4]	44.97	58.49	66.91	74.29
RLW [40]	74.43	88.74	15.82	14.21
Nash [42]	75.05	88.50	16.67	15.10
AMTL [41]	68.30	85.95	24.66	19.49
CAGrad [14]	72.43	87.53	21.43	16.85
ALBO (Ours)	80.47	93.60	11.45	7.53

and, for the auxiliary branch,

where $\theta_a^\pm=\theta_a\pm\epsilon\cdot\nabla_{\theta'_a}\mathcal{L}_m(\theta_m,\theta'_a,\psi_m)$ and $\epsilon$ is a small positive value. Writing the two cases compactly with the shorthand $\theta_{m/a}$ recovers Eq. (9) in the main text.

Effect of the BCM and CCM Modules across Network Scales

To support the discussion of the bidirectional communication module (Sec. III-A) and the bilevel optimization (Sec. III-B), this appendix reports the full comparison between the BCM and CCM modules under two network scales, U-Net-S (128 bottom channels) and U-Net-L (1024 bottom channels), on the ACDC and REFUGE datasets. For ACDC, we used five-fold cross-validation with 160 labeled scans per fold. For REFUGE, each model was trained three times with different random initializations using 400 labeled images. Tables X–XI summarize the results.

On ACDC, the models are prone to underfitting (Table X). With U-Net-L, CCM and BCM give comparable results for all methods. With the smaller U-Net-S, CCM is preferable to BCM for the nine compared methods; nevertheless, ALBO with BCM still attains the best average Dice (90.22%).

On REFUGE, embedding the BCM consistently helps at the small scale: with U-Net-S, all nine compared methods improve significantly over CCM (paired $t$-test, $p\ll0.01$; tbl:XI, panels (a)–(b)), opposite to the trend on ACDC. When the network is enlarged to U-Net-L, all nine methods degrade markedly due to overfitting, under both CCM and BCM, whereas ALBO remains the most robust (panels (c)–(d)).

These observations indicate that the benefit of the BCM is not automatic: it depends on the dataset and the network scale, and the bilevel optimization is what enables the BCM to be exploited robustly across settings.

References