1 Introduction↩︎

Medical image segmentation is a fundamental task in medical image analysis and plays an important role in clinical applications such as disease diagnosis, treatment planning, preoperative assessment, and therapy monitoring[1], [2]. In 3D multi-organ segmentation, models are required to delineate multiple anatomical structures from volumetric CT or MRI scans, where accurate voxel-wise predictions are essential for reliable clinical decision-making. Despite the remarkable progress of deep learning-based segmentation methods, their success typically relies on dense voxel-level annotations, which are expensive and time-consuming to acquire and require substantial expertise from trained clinicians[3]–[5].

To mitigate the limitations caused by insufficient annotations, semi-supervised medical image segmentation has received increasing attention in recent years. It aims to reduce annotation cost by learning from a small set of labeled volumes together with abundant unlabeled data[6]–[9]. Representative methods include consistency regularization[10]–[12], pseudo-labeling[13], and co-training[14], [15], which encourage models to learn from unlabeled images by enforcing stable or mutually agreed predictions. As illustrated in Fig. 1 (a), a typical co-training framework feeds the same unlabeled image into two student networks and encourages their predicted masks to be consistent with each other. In this way, unlabeled data provides additional supervision through cross-model agreement, making the model less dependent on dense voxel-level annotations. Despite promising results, most of them primarily improve prediction-level robustness by enforcing consistency in the output space, while leaving the feature space insufficiently organized with respect to the complex anatomical variations within each semantic class.

This limitation is particularly pronounced for small organs and ambiguous boundary regions. Anatomical structures from the same class often exhibit strong intra-class variability across patients, slice locations, scanning conditions, and surrounding anatomical contexts. Their appearances can differ substantially in shape, scale, intensity, local texture, and boundary clarity[16], [17]. Moreover, small or elongated structures usually occupy only a limited number of voxels and are easily confused with adjacent tissues[18]–[20]. As a result, models may suffer from inter-class confusion, unstable boundary predictions, and insufficient recognition of small structures. As shown in Fig. 1 (e), organs in Synapse[21] exhibit clear disparities in both anatomical size and segmentation difficulty. Several small or high-variation organs, such as the esophagus, pancreas, and adrenal glands, are associated with lower segmentation ranks, indicating that these structures are particularly vulnerable to representation collapse and boundary errors. Notably, the stomach also shows a discrepancy between its relatively large anatomical size and its segmentation performance, suggesting that a large organ volume does not necessarily alleviate the difficulty caused by complex shape variations and ambiguous boundaries[18], [20]. Therefore, multi-organ segmentation requires careful consideration of both organ size and anatomical complexity. This observation motivates us to move beyond prediction-level consistency and investigate feature-space modeling for anatomically variable organs.

To better organize the feature space, proxy- and prototype-based representation learning provides a natural solution by encouraging intra-class compactness and inter-class separability. However, conventional single-prototype designs, as shown in Fig. 1 (b), represent each class with one deterministic center, which may over-compress diverse anatomical patterns into an averaged representation[22]. This is insufficient for organs with heterogeneous shapes, textures, and boundary appearances. Recent distributional proxy methods further model each class as a learnable distribution to capture class-level uncertainty, as illustrated in Fig. 1 (c)[23]. Nevertheless, such global mean-variance representations mainly describe the overall class distribution and still lack explicit mechanisms to encode fine-grained intra-class variations, such as organ-specific sub-modes, ambiguous boundary patterns, and local contextual differences. Since recent proxy-based segmentation studies have shown the importance of intra-class variation modeling for resolving ambiguous regions[24], we argue that an effective medical segmentation proxy should jointly preserve distributional class semantics and expose multiple variation-aware directions within each class.

To this end, we propose Variation-Conditioned Distributional Proxy Learning (VCDP), a training-only regularization module for semi-supervised 3D medical image segmentation. As illustrated in Fig. 1 (d), VCDP models every anatomical class as a distributional proxy conditioned by multiple intra-class variation prototypes. The distributional proxy, parameterized by a learnable Gaussian distribution, captures shared class-level semantics, while class-specific learnable variation prototypes model fine-grained intra-class patterns in the embedding space. By integrating these two complementary representations into a unified compatibility score, VCDP provides proxy-guided soft assignment for voxel embeddings and encourages them to match both the global identity of an organ and its local anatomical variations. Meanwhile, labeled voxels provide reliable semantic calibration for class representatives, preventing proxy drift during training. In this way, VCDP enables unlabeled data to participate in feature-space organization without requiring manual annotations, while improving the representation of diverse anatomical appearances and ambiguous boundary regions. As a plug-and-play auxiliary module, VCDP can be attached to decoder features during training and removed entirely during inference, without modifying the base semi-supervised framework or introducing additional computational cost.

Our contributions are summarized as follows:

• We propose a plug-and-play variation-conditioned distributional proxy learning module for semi-supervised 3D medical image segmentation, which represents each anatomical class with a stochastic Gaussian proxy and multiple learnable variation prototypes to explicitly model class-level semantics and fine-grained anatomical variability in the feature space.

• We formulate a unified variation-conditioned class compatibility score that fuses distributional proxy similarity with soft variation aggregation, producing class-wise soft assignments that guide all voxel embeddings toward both global anatomical semantics and fine-grained intra-class patterns.

• We design a training-only proxy-guided regularization strategy that uses variation-conditioned soft assignments to organize all voxel embeddings, improving representation learning under limited annotations without adding inference cost.

2 Related Work↩︎

2.1 Semi-supervised Medical Image Segmentation↩︎

Semi-supervised medical image segmentation reduces the reliance on dense annotations by exploiting unlabeled volumes together with limited labeled data. Existing methods mainly include consistency regularization, teacher-student learning, pseudo-labeling, uncertainty estimation, and co-training [16], [25]–[30]. These methods improve unlabeled learning by enforcing stable predictions, generating pseudo labels, filtering uncertain supervision, or performing mutual supervision between networks. Despite their effectiveness, most of them focus on prediction-level constraints, leaving explicit feature-space representation learning relatively underexplored.

2.2 Prototype and Proxy-based Representation Learning↩︎

Prototype and proxy-based representation learning introduce class-level representatives to structure the embedding space and have been widely studied in metric learning, classification, and dense prediction [22], [31]. In semantic segmentation, prototypes regularize pixel- or voxel-level embeddings by encouraging intra-class compactness and inter-class separability [32], [33]. Recent methods further enhance representation capacity with multiple prototypes, prototype refinement, or class-wise memory mechanisms [34], demonstrating the value of feature-space modeling for dense prediction.

2.3 Distributional & Variation-aware Representation Learning↩︎

Distributional modeling captures uncertainty beyond deterministic feature points and has been explored for dense prediction with diverse appearances or ambiguous annotations. In medical image analysis, probabilistic representations have been used for uncertainty estimation, ambiguous segmentation, and prototype-based learning [20], [24]. Variation-aware modeling further captures fine-grained intra-class differences, such as local appearance changes, boundary ambiguity, and structural sub-modes. These representations are particularly useful for 3D medical segmentation, where voxel appearances vary across patients, spatial locations, and imaging conditions.

3 Methodology↩︎

3.1 Overview↩︎

Figure 2 illustrates the overall framework of the proposed VCDP for semi-supervised 3D medical image segmentation. Unlike methods that mainly rely on prediction-level consistency or pseudo-label supervision, VCDP formulates SSMIS as a variation-conditioned anatomical representation learning problem, where each class is represented by a Gaussian proxy for shared semantics and multiple variation prototypes for intra-class structural sub-modes.

Let \(\mathcal{X}=\{\mathcal{X}_l,\mathcal{X}_u\}\) denote a dataset of 3D medical images, where \(\mathcal{X}_l=\{(x_i,y_i)\}_{i=1}^{N_l}\) is the labeled set and \(\mathcal{X}_u=\{x_j\}_{j=N_l+1}^{N_l+N_u}\) is the unlabeled set, with \(N_l \ll N_u\). Here, \(x_i \in \mathbb{R}^{D \times H \times W}\) and \(y_i \in \{0,1,\dots,C-1\}^{D \times H \times W}\) denote a 3D image and its voxel-wise label, respectively. Given an input image \(x\), the segmentation network \(f_{\theta}\) produces the prediction \(p\) and an intermediate decoder feature map \(F \in \mathbb{R}^{B \times C_f \times D \times H \times W}\). A lightweight projection head \(\phi(\cdot)\) maps \(F\) into normalized voxel embeddings \(Z=\{z_i\}_{i=1}^{N}\), where \(z_i \in \mathbb{R}^{d}\). Specifically, VCDP represents each class with a Gaussian proxy and multiple variation prototypes, as detailed in Section 3.2, and uses their fused compatibility scores to regularize dense voxel embeddings from both labeled and unlabeled data, as described in Section 3.3. It is jointly optimized with the original segmentation objective during training, while all auxiliary components are removed during inference with no additional cost.

3.2 Distributional and Variation-aware Class Modeling↩︎

Medical images often exhibit substantial anatomical ambiguity and inter-subject variability, making it insufficient to represent each anatomical class with a single deterministic prototype. Voxels belonging to the same class may present diverse appearances due to differences in organ shape, boundary uncertainty, lesion extent, and imaging contrast. Therefore, VCDP models each class with a variation-conditioned representation, where a Gaussian proxy captures shared class-level semantics under uncertainty and multiple variation prototypes describe fine-grained intra-class structural sub-modes.

For each anatomical class \(c\), VCDP maintains a variation-conditioned class representation: \[\mathcal{P}_c = \{(\mu_c,\sigma_c), \mathbf{v}_{c,1}, \dots, \mathbf{v}_{c,K}\},\] where \(K\) = 5 and \(\mu_c\) denotes the learnable Gaussian proxy mean, \(\sigma_c\) denotes a positive dispersion scale for stochastic proxy sampling, and \(\{\mathbf{v}_{c,k}\}_{k=1}^{K}\) denotes learnable variation prototypes for intra-class structural sub-modes.

Given the voxel embeddings \(z_i\), VCDP first constructs a distributional proxy for each anatomical class. Specifically, the Gaussian proxy of class \(c\) is parameterized as \[u_c \sim \mathcal{N}(\mu_c, \mathrm{diag}(\sigma_c^2)),\]

where \(\mu_c\) represents the learnable semantic center of the class, and \(\sigma_c\) denotes a positive dispersion scale used to generate stochastic proxy representatives. In our implementation, \(\sigma_c\) is kept fixed during training and is only used to control the sampling spread of the Gaussian proxy.

To estimate the distributional compatibility between a voxel embedding \(z_i\) and class \(c\), we draw \(S\) stochastic class representatives from the Gaussian proxy: \[u_{c,s} = \mathrm{Normalize}(\mu_c + \sigma_c \odot \epsilon_s), \quad \epsilon_s \sim \mathcal{N}(0,I), \quad s=1,\dots,S.\] The distributional similarity is then computed by averaging the cosine similarities between \(z_i\) and all sampled representatives: \[s_{\mathrm{dist}}(z_i,c) = \frac{1}{S} \sum_{s=1}^{S} \cos(z_i,u_{c,s}).\] Compared with a deterministic class prototype, the stochastic proxy provides a distributional description of class semantics and allows the model to account for uncertainty in voxel-level representations. This is particularly beneficial for semi-supervised 3D medical image segmentation, where unlabeled data may contain ambiguous boundaries and diverse anatomical variations.

In addition to the Gaussian proxy, each class is equipped with \(K\) learnable variation prototypes to capture fine-grained intra-class structural patterns, such as local appearance changes, shape variations, and boundary ambiguity. For a voxel embedding \(z_i\), VCDP computes its variation-aware similarity to class \(c\) by softly aggregating the similarities over all class-specific variation prototypes: \[s_{\mathrm{var}}(z_i,c) = \frac{1}{\tau} \log \sum_{k=1}^{K} \exp\left( \tau \cos(z_i,\mathbf{v}_{c,k}) \right),\]

where \(\tau\) controls the sharpness of the soft aggregation. This log-sum-exp formulation serves as a differentiable approximation to hard prototype selection, allowing each voxel embedding to adaptively match the most relevant intra-class variation patterns while still preserving gradient contributions from other possible structural modes.

After modeling distributional class semantics and intra-class variations, VCDP fuses them into a unified class compatibility score. For each voxel embedding \(z_i\) and class \(c\), the fused score is defined as \[g(z_i,c) = s_{\mathrm{dist}}(z_i,c) + \lambda_{\mathrm{var}} s_{\mathrm{var}}(z_i,c),\] where \(\lambda_{\mathrm{var}}\) balances the contribution of the variation-aware term. This score jointly measures the compatibility between a voxel embedding and both the global semantic distribution and fine-grained structural variations of an anatomical class.

4pt

3.3 Soft Assignment-guided Embedding Regularization↩︎

Given the projected embedding map \(Z\), we flatten its spatial dimensions into a set of voxel embeddings \(\{z_i\}_{i=1}^{N}\), where \(N=BD'H'W'\). VCDP regularizes these embeddings from both labeled and unlabeled volumes through class-wise soft assignments derived from the detached fused compatibility score.

To make the optimization stable, we use two complementary training paths: a dense embedding regularization path and a labeled proxy calibration path.

In the dense embedding regularization path, each voxel embedding is compared with every semantic class using the fused compatibility score. However, to prevent noisy unlabeled embeddings from directly shifting the class-level Gaussian proxy means, we stop the gradient of \(\mu_c\) when computing this score: \[\bar{g}(z_i,c) = g\left( z_i,c; \mathrm{sg}(\mu_c), \sigma_c, \{\mathbf{v}_{c,k}\}_{k=1}^{K} \right).\]

where \(\mathrm{sg}(\cdot)\) denotes the stop-gradient operation. Here, \(\bar{g}(z_i,c)\) is the stop-gradient version of the fused compatibility score used only in the dense regularization path. The distributional term \(s_{\mathrm{dist}}\) still measures the compatibility between \(z_i\) and the stochastic Gaussian representatives of class \(c\), but its Gaussian parameters \((\mu_c,\sigma_c)\) are not updated by this path. In contrast, the variation prototypes \(\{\mathbf{v}_{c,k}\}_{k=1}^{K}\) and the projected voxel embeddings remain trainable. Therefore, this path mainly shapes the embedding space and refines the intra-class variation prototypes, while preserving stable class-level Gaussian semantics.

Based on \(\bar{g}(z_i,c)\), we compute a class-wise soft assignment over all semantic classes: \[q_i(c) = \frac{ \exp(\bar{g}(z_i,c)) }{ \sum_{c'=0}^{C-1} \exp(\bar{g}(z_i,c')) }.\] This assignment reflects the relative compatibility between voxel embedding \(z_i\) and each variation-conditioned class representation. Unlike confidence-based pseudo-labeling strategies that only supervise selected voxels with hard labels, this formulation provides dense feature-space regularization over all projected voxel embeddings in the current mini-batch.

The first regularization term encourages each voxel embedding to align with compatible class representations according to its soft assignment: \[\mathcal{L}_{align} = \frac{1}{N} \sum_{i=1}^{N} \sum_{c=0}^{C-1} q_i(c) \left[ 1-\bar{g}(z_i,c) \right].\] This term promotes distribution-aware feature alignment by pulling each embedding toward class representations with high compatibility scores. To further enhance class discriminability, we introduce a distribution discrimination term: \[m_c = \frac{1}{N} \sum_{i=1}^{N} \left(2q_i(c)-1\right)\bar{g}(z_i,c),\] \[\mathcal{L}_{dis} = \frac{1}{sec:C} \sum_{c=0}^{C-1} \exp(-m_c).\] This term encourages embeddings strongly assigned to class \(c\) to have high compatibility with that class, while suppressing excessive responses from weakly assigned embeddings. The dense embedding regularization objective is then defined as \[\mathcal{L}_{reg} = \mathcal{L}_{align} + \mathcal{L}_{dis}.\]

In parallel, the labeled proxy calibration path provides a lightweight stabilization signal for the Gaussian proxy means by aligning them with reliable labeled voxel features.

For class \(c\), let \[\Omega_c=\{i\mid y_i=c\}\] denotes the set of labeled voxel indices belonging to class \(c\) in the current mini-batch. We compute the labeled class feature anchor as \[a_c = \mathrm{Normalize} \left( \frac{1}{|\Omega_c|} \sum_{i\in\Omega_c} z_i \right).\] The proxy calibration loss is formulated as \[\mathcal{L}_{cal} = \frac{1}{|\mathcal{sec:C}_L|} \sum_{c\in\mathcal{sec:C}_L} \left[ 1- \cos \left( \mathrm{Normalize}(\mu_c), \mathrm{sg}(a_c) \right) \right],\] where \(\mathcal{sec:C}_L\) denotes the set of classes present in the labeled mini-batch. The labeled anchor \(a_c\) is detached and serves as a reliable feature anchor, so this calibration path stabilizes the Gaussian class centers without directly imposing additional gradients on the backbone embeddings.

The final training objective is \[\mathcal{L} = \mathcal{L}_{base} + \lambda_{reg}\mathcal{L}_{reg} + \lambda_{cal}\mathcal{L}_{cal},\] where \(\mathcal{L}_{base}\) denotes the original objective of the underlying semi-supervised segmentation framework, and \(\lambda_{reg}\) and \(\lambda_{cal}\) balance the two auxiliary objectives. During inference, the projection head, Gaussian proxies, and variation prototypes are removed, and the final prediction is produced only by the original segmentation network.

4 Experiment↩︎

4.1 Experimental Setup↩︎

Dataset. We evaluate VCDP on two multi-organ CT segmentation benchmarks, Synapse [21] and AMOS [37]. Synapse contains 30 abdominal CT scans with 13 foreground organs, including spleen (Sp), right/left kidneys (RK/LK), gallbladder (Ga), esophagus (Es), liver (Li), stomach (St), aorta (Ao), inferior vena cava (IVC), portal and splenic veins (PSV), pancreas (Pa), and right/left adrenal glands (RAG/LAG). Following previous semi-supervised settings, we split Synapse into 20 training, 4 validation, and 6 testing scans, and use 20% of the training scans as labeled data. AMOS contains 360 CT scans with 15 foreground organs. In addition to most Synapse organs, AMOS includes duodenum (Du), bladder (Bl), and prostate/uterus (P/U), while excluding portal and splenic veins. We use 216 scans for training, 24 for validation, and 120 for testing, and adopt a more challenging 5% labeled setting with the remaining training scans used as unlabeled data.

Implementation Details. We implement VCDP as a training-time plug-in module and integrate it into CPS[27], MagicNet[28], DHC[30], GenSSL[29], SS-Net[25], Adsh[35], and DCMamba[36]. For fair comparison, we keep the original prediction branches and training objectives of all baselines unchanged, and jointly optimize VCDP as an auxiliary feature-space regularizer during training. All experiments are implemented in PyTorch and conducted on an NVIDIA A40 GPU with a batch size of 4. The weight decay for the proposed proxy/prototype module is set to \(1\times10^{-4}\). At inference time, the projection head, Gaussian proxies, and variation prototypes are removed, thus introducing no additional computational cost. We report DSC, NSD, and HD95, where higher DSC/NSD and lower HD95 indicate better performance.

3.0pt

2.4pt

4.2 Comparisons with State-of-the-art Methods↩︎

As shown in Table 1, VCDP brings broad improvements across different semi-supervised segmentation frameworks on Synapse under the 20% labeled setting. When integrated into CPS, VCDP improves Dice to 67.51% \((\uparrow 1.25\%)\), increases NSD from 75.48% to 78.06% \((\uparrow 2.58\%)\), and reduces HD95 from 13.41 mm to 9.35 mm \((\downarrow 4.06\text{ mm})\). Similar gains are observed for MagicNet, where Dice improves to 68.84% \((\uparrow 2.98\%)\). More substantial gains are observed for baselines with lower initial performance, with Dice improvements of 3.61% for DHC, 6.46% for GenSSL, 2.46% for SS-Net, 6.96% for Adsh, and 9.68% for DCMamba. In particular, DCMamba+VCDP achieves a large Dice gain \((\uparrow 9.68\%)\), while GenSSL+VCDP reduces HD95 to 5.86 \((\downarrow 1.83\text{ mm})\). These results show that VCDP is not tied to a specific semi-supervised strategy, but provides complementary feature-space supervision for diverse prediction-space frameworks.

On AMOS under the more challenging 5% labeled setting, VCDP also improves most evaluated baselines. For MagicNet, VCDP improves Dice to 63.65% \((\uparrow 0.47\%)\) and reduces HD95 from 11.08 mm to 10.14 mm \((\downarrow 0.94\text{ mm})\), achieving the best overall performance among the evaluated methods. The gains are more promising for GenSSL and SS-Net: GenSSL+VCDP improves Dice/NSD by 14.29%/20.41% and reduces HD95 by 26.29 mm, while SS-Net+VCDP improves Dice/NSD by 20.39%/22.39%. These large improvements indicate that VCDP provides strong representation-level guidance when the baseline feature space is insufficiently organized. For DHC and DCMamba, Dice improvements are less uniform on AMOS, but VCDP still reduces HD95, suggesting potential benefits in boundary robustness.

Tables 2 and 3 further report the per-class Dice scores on Synapse and AMOS, respectively. The gains are mainly concentrated on challenging organs with small sizes, ambiguous boundaries, and large intra-class variations. On Synapse, VCDP notably improves difficult structures, including the gallbladder, esophagus, stomach, and pancreas. On AMOS, VCDP brings substantial gains to under-performing baselines such as GenSSL and SS-Net. Overall, these results indicate that VCDP improves organ-level robustness, with particularly strong benefits for anatomically difficult categories requiring fine-grained variation modeling.

4.3 Qualitative Analysis↩︎

Fig. 3 compares the segmentation results of the ground truth, MagicNet, and MagicNet+VCDP on the pancreas region. MagicNet produces fragmented and inaccurate predictions around the pancreas, reflecting its difficulty in handling small and ambiguous anatomical structures. In contrast, MagicNet+VCDP better preserves the pancreas region and generates a more complete mask closer to the ground truth. These results show that VCDP improves the segmentation of challenging organs through variation-conditioned feature-space modeling.

4.4 Ablation Study↩︎

To validate the effectiveness of Gaussian and variation modeling in VCDP, we conduct ablation studies on the Synapse dataset using MagicNet as the baseline framework, as shown in Table 4. The baseline MagicNet achieves 65.86% mean Dice and 78.60% NSD. With only the Gaussian proxy, the performance increases to 68.06% Dice (\(\uparrow\) 2.20%) and 79.25% NSD (\(\uparrow\) 0.65%), showing the benefit of distributional class modeling. With only variation prototypes, the model also improves the baseline Dice to 66.96% (\(\uparrow\) 1.10%), indicating that intra-class variation modeling provides useful feature-space guidance. Combining both components achieves the best results, reaching 68.84% Dice (\(\uparrow\) 2.98%) and 79.65% NSD (\(\uparrow\) 1.05%). These results verify that Gaussian proxies and variation prototypes are complementary for modeling shared class semantics and fine-grained anatomical variations.

5 Conclusion↩︎

In this paper, we proposed VCDP, a plug-and-play training-only regularization module for semi-supervised 3D medical image segmentation. By representing each anatomical class with a Gaussian proxy and multiple variation prototypes, VCDP explicitly organizes voxel embeddings to capture both shared class semantics and fine-grained intra-class variations. Experiments on Synapse and AMOS demonstrate broad improvements across multiple semi-supervised frameworks, especially for anatomically challenging organs. Since all auxiliary components are removed during inference, VCDP introduces no additional inference cost.

References↩︎

---

1. Corresponding author.↩︎