Paper deep dive
Optimal Transport Metric Learning for Feature Alignment in Partially Supervised Segmentation
Dakini Mallam Garba, Salim Abdou Daoura
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 96%
Last extracted: 9/18/2026, 3:43:28 AM
Summary
The paper proposes a two-stage learning framework for multi-organ segmentation using partially annotated datasets to address domain shifts. Stage 1 learns from available annotations using standard losses. Stage 2 introduces learnable organ prototypes and a Sinkhorn-triplet loss based on Optimal Transport to enforce feature consistency across datasets, aligning latent embeddings of the same organ while separating different organs. The method achieves performance comparable to state-of-the-art methods on the BTCV dataset with improved computational efficiency and robustness to domain shift.
Entities (15)
Relation Signals (12)
Dakini Mallam Garba β authored β Optimal Transport Metric Learning
confidence 99% Β· Dakini Mallam Garba and Salim Abdou Daoura
Salim Abdou Daoura β authored β Optimal Transport Metric Learning
confidence 99% Β· Dakini Mallam Garba and Salim Abdou Daoura
Optimal Transport Metric Learning β uses β Sinkhorn-triplet loss
confidence 97% Β· we introduce learnable organ prototypes and a Sinkhorn-triplet loss to enforce organ-wise feature consistency
Optimal Transport Metric Learning β implementedwith β nnUNet
confidence 96% Β· All experiments were conducted with the nnUNet framework
Optimal Transport Metric Learning β evaluatedon β MSD
confidence 95% Β· We evaluated our approach using four widely adopted organ-labeled datasets: ... MSD
Optimal Transport Metric Learning β evaluatedon β Lits19
confidence 95% Β· We evaluated our approach using four widely adopted organ-labeled datasets: ... Lits19
Optimal Transport Metric Learning β evaluatedon β KiTS19
confidence 95% Β· We evaluated our approach using four widely adopted organ-labeled datasets: Kits19
Optimal Transport Metric Learning β evaluatedon β BTCV dataset
confidence 95% Β· Our approach achieves performance comparable to state-of-the-art methods on the BTCV dataset
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Multi-organ segmentation is often challenged by partially annotated datasets and domain shifts across different imaging sources. To address these limitations, we propose a two-stage learning framework that efficiently leverages partial supervision. In the first stage, the model learns from available annotations to produce accurate segmentations of annotated organs, establishing robust feature representations. In the second stage, we introduce learnable organ prototypes and a Sinkhorn-triplet loss to enforce organ-wise feature consistency across datasets. This encourages latent embeddings of the same organ to remain close, while increasing separation between different organs, even when annotations are missing. Our approach achieves performance comparable to state-of-the-art methods on the BTCV dataset, while remaining computationally efficient. By explicitly aligning feature distributions rather than relying solely on pseudo-labels, the framework effectively mitigates domain shift, making it particularly suitable for medical image segmentation tasks with limited annotation resources.
Tags
Links
- Source: https://arxiv.org/abs/2609.19176v1
- Canonical: https://arxiv.org/abs/2609.19176v1
Trouble viewing inline? Open PDF directly β
Full Text
37,306 characters extracted from source content.
Expand or collapse full text
Optimal Transport Metric Learning for Feature Alignment in Partially Supervised Segmentation Dakini Mallam Garba [0009β0007β5198β1506] and Salim Abdou Daoura [0009β0005β7128β552ν] Abstract Multi-organ segmentation is often challenged by partially annotated datasets and domain shifts across different imaging sources. To address these limitations, we propose a two-stage learning framework that efficiently leverages partial supervision. In the first stage, the model learns from available annotations to produce accurate segmen- tations of annotated organs, establishing robust feature representations. In the second stage, we introduce learnable organ prototypes and a Sinkhorn-triplet loss to enforce organ-wise feature consistency across datasets. This encourages latent embeddings of the same organ to remain close, while increasing separation between different organs, even when annotations are missing. Our approach achieves performance comparable to state-of-the-art methods on the BTCV dataset, while remaining computationally effi- cient. By explicitly aligning feature distributions rather than relying solely on pseudo- labels, the framework effectively mitigates domain shift, making it particularly suitable for medical image segmentation tasks with limited annotation resources. 1 Introduction Automatic organ segmentation plays a central role in various applications, such as diag- nostic imaging, radiotherapy planning, and pathology analysis. Deep Learning methods have emerged as the dominant paradigm for this task, but they traditionally require a large number of fully annotated scans. Since manual annotation is inherently time- consuming, often only partially annotated datasets are available. As a result, many existing datasets provide only partial annotations focused on a single target organ. Training separate models for each dataset is computationally ineffi- Dakini Mallam Garba ML Collective, e-mail: dakinimallam@gmail.com Salim Abdou Daoura ML Collective, e-mail: salim.adaoura@gmail.com 1 arXiv:2609.19176v1 [eess.IV] 15 Sep 2026 2Dakini Mallam Garba and Salim Abdou Daoura Fig. 1 Our training setting: From partially annotated datasets, we aim to learn a single unified model that will produce predictions for all organs. Pancreas Liver Kidneys Net TrainingInference cient, while approaches that share an encoder but use task-specific decoders fail to fully exploit inter-organ relationships. To address this limitation, we propose a two-stage training framework that learns from partially annotated datasets to build a single unified model, as illustrated in Figure 1. In the first stage, learning is performed using only the available ground truth of the partial datasets. We employ adapted segmentation losses that encourage agreement between ground truth and prediction. This stage establishes reliable feature representations. In the second stage, we employ an optimal transport-based triplet loss that encourages organ-wise feature consistency across datasets. This regularization helps the network learn coherent representations for both annotated and unannotated domains, mitigating domain shift. For instance, when the liver is annotated in one dataset but unannotated in another, our model regularizes their feature distributions to remain close in the latent space, ensuring semantic consistency. Simultaneously, it enforces separation from the distributions of other organs to improve inter-class discriminability. We compare features predicted for unlabeled organs against a set of learned prototype features derived from the ground truth annotations. The proposed Sinkhorn-triplet loss replaces pointwise distances with optimal transportβbased divergences, aligning entire feature distributions rather than individual samples. Our contributions are summarized as follows: β’ We propose a novel application of the Sinkhorn-triplet loss to partially supervised medical image segmentation. By replacing pointwise distances in the triplet formu- lation with Sinkhorn divergences between feature distributions, our method enforces distribution-level alignment between predicted and ground truth feature spaces. β’ We design a training strategy where pseudo-labels are used solely to determine anchor class membership, rather than as direct supervision signals. This prevents the direct propagation of pseudo-label errors and promotes robust feature alignment through optimal transportβbased regularization. β’ Our approach enhances feature consistency and cross-dataset generalization, en- abling accurate segmentation of unannotated organs by aligning features across domains at the distributional level. Optimal Transport Metric Learning in Partially Supervised Segmentation3 2 Related Work 2.1 Partially Supervised Segmentation Early efforts to train unified models from partially annotated datasets often focused on modifying the loss function. [1] introduced a masked Dice loss that leverages only available ground truth labels during backpropagation. [2] further refined this idea by incorporating a class-asymmetric objective function (CAL) that penalizes incorrect predictions under the assumption of mutual exclusivity. While these approaches effec- tively ignore missing annotations, they do not exploit the pseudo-labels generated by the model itself, thus overlooking potentially valuable semantic information that could enhance segmentation performance. An alternative research direction focused on architectural design. DoDNet [3] and OmniSeg [4] implemented a dynamic segmentation head, where specific convolutional kernels are assigned to each task for the segmentation of a specific organ and tumors. The kernels are adaptively generated by a task controller, conditioned on both the input image and the assigned task. This mechanism enables the network to adapt its decoding path to the segmentation of specific organs or lesions, offering greater flexibility and computational efficiency compared to training multiple independent models or using multi-head decoders. However, these methods still rely solely on ground truth super- vision and require multiple forward passes at inference time to segment all organs, limiting their scalability. More recent approaches make use of the pseudo-labels generated by the model via a two-stage training strategy. In COSST [5], the network first learns from avail- able annotation, then generates pseudo-labels for missing organs. The quality of these pseudo-labels is assessed by outlier detection in the latent space. Finally, these reliable pseudo-labels are iteratively used for finetuning. PSSNET [6] also adopts a two-stage learning process where in the second stage, the authors retrain the segmentation model using an adversarial network as a means to reduce the discrepancy between annotated and unannotated domains. Although these strategies filter pseudo-labels, they still risk propagating noisy supervision when the filtering mechanism fails, potentially desta- bilizing training and degrading generalization. More recently, [29] showed promising results by introducing a labeled-to-unlabeled distribution alignment framework. 2.2 Deep Metric Learning Deep Metric Learning has been used in various applications such as Person RE- Identification [21], few shot learning [23] and face recognition [22]. Its goal is to train an embedding model to measure the similarity between data points. Images are mapped into feature vectors through a deep neural network encoder. A loss function is tailored to minimize the distance between similar vectors, and maximize the distance 4Dakini Mallam Garba and Salim Abdou Daoura between vectors of different classes. Some standard losses used are the contrastive loss [24] or the triplet loss [22]. The triplet loss operates on triplets of samples : an anchor, a positive point (from the same class) and a negative point (from a different class). The distance of the positive pair is pulled to be smaller than the negative pair. It creates an embedding space where features of the same class are close while the inter-class dis- tance is high. In our framework, we adapt this principle to medical image segmentation by introducing a modified triplet objective that enforces organ-wise feature separability across datasets. 2.3 Optimal Transport Optimal transport (OT) determines the cost of transforming one distribution into an- other. It captures geometric correspondences between distributions, making it particu- larly effective for comparing structured feature spaces. OT has been used in a variety of computer vision tasks: domain adaptation [11], multitask learning [12], feature matching [13] or generative model [8]. Recent advances have also leveraged OT for self-supervised representation alignment [7] and topology-aware segmentation losses based on OT distances between persistent diagrams [14]. In this work, we exploit an OT-based distance to compare feature distributions across organs. Unlike pointwise metrics, OT accounts for the underlying geometry of the embedding space, providing a more robust measure of feature separability and alignment across partially annotated datasets. 3 Method In this section, we formalize our two-stage training framework. Figure 2 provides an overview of our proposed method. 3.1 Problem Formulation We consider ν partially labeled datasets ν· νν =ν· 1 ,..., ν· ν s.t. ν· ν =ν₯, ν¦, βν β [1,..., ν]. For each dataset ν· ν , only one organ νΆ ν is annotated in the ground truth mask ν¦. Our goal is to learn a single unified model F ν : [1, ν, ν»,ν] β¦β [ν, ν, ν»,ν] from the collection of partially labeled data ν· νν . Optimal Transport Metric Learning in Partially Supervised Segmentation5 Net L seg Ground Truth (a) Stage 1 Net L SinkTriplet Block L seg (b) Stage 2 (c) L SinkTriplet Block- class c MLP projector [1,Z,H,W][1,Z,H,W] [N,Z,H,W] [B,1,Z,H,W] [B,1,Z,H,W] [B,N,Z,H,W] [B,D,Zβ,Hβ,Wβ] Hadamard product [B,D,Zβ,Hβ,Wβ] [B,N,Z,H,W] downsampling [B,1,Zβ,Hβ,Wβ] class filtering Prototype Classifier Positive prototypes Negative Prototypes ΞΌ a ν + c ν - c L SinkTriplet Sampling Pseudo Labels from stage 1 Intra-class Clustering Protoype classifer L proto Fig. 2 Overview of the proposed two-stage framework. (a) In Stage 1, the model is trained using only the available ground truth annotations available in partially labeled datasets. At the end of stage 1, pseudo-labels are generated and filtered. (b) In Stage 2, training is regularized with the Sinkhorn-triplet loss that enforces inter-class feature alignement. A prototype classifier learns each class prototypes (c) The Sinkhorn-triplet loss is computed. Anchor distribution is sampled based on the pseudo-labels from stage 1 while the positive and negative distributions are derived from the prototype classifier. 3.2 Stage 1: Supervised Learning with Ground Truth The first stage of our framework consists in learning exclusively from the available annotations. The objective is to maximize the agreement between the predictions of the labeled organs and the corresponding ground truth. Given a prediction Μν¦ = νΉ ν (ν₯), the segmentation objective is defined as: L ν νν =L ν·ννν ( Μν¦ ν‘ , ν¦)+L νΆνΈ ( Μν¦ ν‘ , ν¦)(1) Here, Μν¦ ν‘ denotes the predicted probabilities restricted to the labeled organ.L ν·ννν and L νΆνΈ correspond to the standard Dice loss and cross-entropy loss, respectively. After stage 1, pseudo-labels are generated for the training set, and a confidence threshold ν is applied, based on the softmax predictions, to retain only the most reliable predictions. 3.3 Stage 2: Feature consistency with the Sinkhorn-triplet loss Approaches that directly leverage pseudo-labels are often prone to noise propagation, as unreliable predictions may be treated as ground truth and destabilize training. Instead, we use pseudo-labels solely to infer class membership in unlabeled regions, without using them as explicit supervision targets. This allows us to define positive and negative pairs while relying only on prototypes derived from ground-truth annotations to guide semantic alignment across datasets. 6Dakini Mallam Garba and Salim Abdou Daoura 3.3.1 Multimodal prototype learning In the second stage of training, we define a set of learnable prototypes that serve as compact and multimodal representations of each organ class. Each class ν is associated with νΎ prototypes Pν = ν ν,1 ,..., ν ν,νΎ , initialized randomly and optimized jointly with the network parameters. During training, feature embeddings νΉ are extracted from a mid-level layer of the network and projected into a lower-dimensional space. For each class ν, we compute local centroidsC ν from the features corresponding to that class us- ing an online clustering step. The prototypes are then updated to minimize their distance to these centroids, ensuring that they capture the underlying modes of the class distribu- tion. To prevent excessive fragmentation of the feature space, we introduce a variance regularization term that enforces intra-class compactness, encouraging prototypes of the same class to remain close to their mean representation. The prototype learning loss is formulated as : νΏ νννν‘ν = 1 νΆ νΆ βοΈ ν=1 β₯ P ν βC ν β₯ 2 + ν ννννννν‘ νΆ νΆ βοΈ ν=1 ννν(P ν )(2) 3.3.2 The triplet loss The classical triplet loss operates on triplets of samples (ν₯ ν ,ν₯ + ,ν₯ β ) corresponding respectively to an anchor, a positive example from the same class, and a negative example from a different class. It is defined for a triplet as: L ν‘νννννν‘ = max ν(ν₯ ν ,ν₯ + )β ν(ν₯ ν ,ν₯ β )+ ν, 0 ,(3) where ν is typically the Euclidean distance, and ν is a fixed margin encouraging inter-class separation. At each iteration, the stage-1 pseudo-label mask Μ ν ννν‘νβ selects feature vectors as- signed to class ν. Each selected vector ν₯ ν serves as an anchor, while the positive and negative examples are drawn from the learned prototype sets: β’ ν₯ + is the closest prototype inP ν , and β’ ν₯ β is the closest prototype from another class ν β² β ν. This ensures that the alignment process is driven by ground-truth derived class proto- types while pseudo-labels are only used to guide anchor selection in unlabeled regions. 3.3.3 The Sinkhorn-triplet loss Following [20], we generalize the triplet formulation to operate on distributions rather than individual embeddings, replacing the pointwise distance by the Sinkhorn diver- gence [9]. We define each anchor, positive, and negative example as a discrete probability measure supported in the feature spaceR ν· . The anchor distribution ν ν is constructed Optimal Transport Metric Learning in Partially Supervised Segmentation7 from the pseudo-labeled region predicted as class ν. Instead of using all voxels, we randomly sample νΎ ν feature vectors ν ν ν from this region and assign them uniform weights: ν ν = 1 νΎ ν νΎ ν βοΈ ν=1 νΏ ν ν ν , ν ν ν β ν | Μ ν ννν‘νβ ( ν) = ν.(4) This sampling reduces the impact of noisy or redundant pseudo-labeled points. The positive distribution ν + ν corresponds to the ground-truth prototype distribution of the same class : ν + ν = 1 νΎ ν νΎ ν βοΈ ν=1 νΏ ν ν,ν , ν ν,ν β P ν ,(5) and the negative distribution ν β is defined as the union of prototype distributions from all other classes: ν β ν = 1 Γ ν β² β ν νΎ ν β² βοΈ ν β² β ν νΎ ν β² βοΈ ν=1 νΏ ν ν β² ,ν , ν ν β² ,ν β P ν β² .(6) The Sinkhorn divergence between two empirical distributions ν and ν is defined as: ν ν,ν (ν,ν) = νν ν,ν (ν,ν)β 1 2 νν ν,ν (ν, ν)β 1 2 νν ν,ν (ν,ν),(7) where νν ν,ν (ν,ν) is the optimal transport cost regularized by entropy: νν ν,ν (ν,ν) =min νβΞ (ν,ν) β¨νΆ ν , νβ©+ ν KL(ν|| νβ ν),(8) with cost νΆ ν (ν₯, ν¦) = 1 ν β₯ν₯β ν¦β₯ ν and regularization parameter ν . The resulting Sinkhorn-triplet loss is then expressed as: L ννννννννννν‘ = max ν ν,ν (ν ν ,ν + ν )β ν ν,ν (ν ν ,ν β ν )+ ν, 0 ,(9) This formulation enforces that the predicted feature distribution of a pseudo-labeled region aligns with its corresponding class prototype distribution while remaining distant from other classes in the Wasserstein space. The overall objective function used in the second training stage combines segmen- tation and alignment objectives: L =L ν νν + ν ν‘νννννν‘ βL ννννννννννν‘ + ν νννν‘ν βL νννν‘ν (10) where ν ν‘νννννν‘ and ν νννν‘ν are weighting coefficients. 8Dakini Mallam Garba and Salim Abdou Daoura 4 Experiments 4.1 Datasets We evaluated our approach using four widely adopted organ-labeled datasets: Kits19 [16] containing kidney and kidney tumor annotations, Lits19 [19] focusing on liver segmentation MSD [15] Pancreas and Spleen datasets. Tumor masks were merged with their corresponding organ masks. In addition to the held-out test sets of each dataset, we employed the BTCV dataset [17] as an external benchmark to assess cross-dataset generalizability. Table 4 in the appendix summarizes the different dataset splits. 4.2 Experiments Setup 4.2.1 Implementation details All experiments were conducted with the nnUNet framework [18], chosen for its strong and robust performance in biomedical image segmentation. For preprocessing, all CT scans are resampled to 2Γ0.80Γ0.80 m, clipped to [-1024,1024] Hounsfield Units (HU), and rescaled to [0,1]. The patch size used by nnUNet is 96 x 160 x 160. For stage 2, we used the Geomloss library [9] to compute the Sinkhorn divergence. Our feature projector is a small MLP composed of two Conv3d layers with kernel size of 1, separated by a batch normalization and ReLU activation layers. A complete list of the hyperparameters used is available in the appendix. 4.2.2 Evaluation metrics We report three widely used metrics: the Dice Similarity Coefficient (DSC), 95th Haus- dorff distance (HD95) and the Average Surface Distance (ASD). DSC measures spatial overlap between the prediction mask and the ground truth. HD95 is the 95th percentile of the maximum distances between the boundary points in the prediction and the ground truth. ASD measures the average distance between the predicted segmentation and the ground truth along the surface of the segmented object. 4.3 Comparison with the State-of-the-art We evaluated our method against seven state-of-the-art (SOTA) partial-label segmen- tation methods, with all baseline results reported from [5]. For a fair comparison, our method uses the same architecture and training strategy. SOTA methods include those using individual networks trained on each partial dataset (denoted Multi-Net), those Optimal Transport Metric Learning in Partially Supervised Segmentation9 Table 1 Performance comparison (DSC, %, higher is better; HD95, m, lower is better; ASD, m, lower is better) of partially labeled segmentation methods. P: held-out testing sets of Kits19, MSD and Lits19. F: BTCV. Our method achieves the best DSC on the partial datasets while staying comparable to SOTA methods on BTCV dataset. MethodsSpleen (P)Pancreas (P)Liver (P) DSCHD95ASDDSCHD95ASDDSCHD95ASD Multi-Nets91.5926.153.7277.216.881.4094.817.541.28 TAL93.722.970.8276.366.371.4394.866.101.31 ME94.092.380.7775.996.010.7794.707.451.90 PLT93.4134.853.9777.696.821.3295.246.131.10 Co-train93.0942.065.2477.506.221.2195.245.100.96 DoDNet93.7515.442.2379.095.841.1793.6511.262.03 COSST94.401.450.5477.985.821.2995.225.781.18 Ours95.5215.032.7283.839.441.8396.0012.302.16 MethodsSpleen (F)Pancreas (F)Liver (F) DSCHD95ASDDSCHD95ASDDSCHD95ASD Multi-Nets87.7229.366.2173.1122.064.1494.807.191.18 TAL90.0212.152.2873.867.771.7695.998.011.05 ME90.5814.252.9174.8721.773.6595.9610.151.60 PLT88.0055.878.1775.3811.242.1395.823.650.62 Co-train88.5352.438.8275.9310.811.9095.825.180.77 DoDNet89.9422.014.5076.517.341.6194.4520.443.58 COSST92.344.131.2577.175.901.3696.283.860.71 Ours90.0320.774.7381.277.731.2694.3215.822.44 using only ground-truth supervision (TAL [26] and ME [10]), those using pseudo-label supervision (COSST [5], PLT [1], and Co-training [28]), and the conditioned network DoDNet [3]. We exclude kidney results due to a preprocessing discrepancy: the baseline methods treat left and right kidneys as two separate labels, whereas our method treats them as a single class. A summary of the comparison is presented in Table 1. On the held-out test sets corresponding to the partial datasets, our proposed method demonstrates superior segmentation accuracy in terms of the Dice Score across all evaluated organs. Notably, an improvement of over 4% in DSC is observed for the pan- creas compared to all competing methods, highlighting our modelβs effectiveness on challenging structures. Furthermore, when considering boundary localization metrics, our approach achieves a lower HD95 for the spleen compared to methods relying on multiple networks (Multi-Nets, PLT, Co-train). This suggests that our unified model not only performs well but also offers greater computational efficiency compared to these approaches. However, we observe relatively high HD95 values across several organs, particularly for the liver, which is the largest organ evaluated. This suggests a poten- tial limitation where the model prioritizes volumetric overlap leading to high DSC, especially for large organs, over precise boundary delineation. This behavior is likely 10Dakini Mallam Garba and Salim Abdou Daoura Table 2 Ablation study on our method components on the BTCV dataset. The addition of the second training stage improves the model segmentation capibility, with all losses. The Sinkhorn-triplet loss gets the highest improvements overall. MethodsSpleenPancreasLiverKidney DSC HD95 ASD DSC HD95 ASD DSC HD95 ASD DSC HD95 ASD Baseline89.61 21.59 4.33 80.87 9.201.44 94.06 17.70 2.81 81.55 27.32 5.13 w/ L ν‘νννννν‘ 90.03 28.17 5.38 81.03 8.541.35 94.30 16.95 2.66 82.67 26.45 4.92 w/ L ννννβννν 89.93 20.76 5.06 81.05 8.561.39 94.08 15.65 2.63 83.82 24.26 4.75 w/ L ν νννννννννν‘ 90.03 20.77 4.73 81.27 7.731.26 94.32 15.82 2.44 83.50 26.67 4.83 exacerbated by our proposed regularization loss, which primarily focuses on ensuring feature distribution coherence rather than explicitly penalizing fine-grained contour in- accuracies. On the BTCV dataset, our method achieves performance comparable to state-of-the- art techniques across most organs and metrics, while demonstrating notably superior results for pancreas segmentation. More importantly, our model exhibits improved robustness to domain shift compared to baseline methods. For example, the Multi-Nets approach experiences a performance degradation of 4.1% in DSC for the pancreas when moving from the partial test sets to BTCV (77.21% to 73.11%). In contrast, our method shows a significantly smaller drop of only 2.56% for the same organ under the same domain shift. This indicates that the feature consistency enforced by our regularization mechanism enhances the modelβs ability to generalize to unseen data distributions. 4.4 Ablation Studies We conducted an ablation study to measure the improvements induced by our different method components. 4.4.1 Quantitative results Table 2 summarizes the quantitative results on the BTCV dataset. Baseline is the model at the end of stage 1, the results of stage 2 are compared with a Sinkhorn divergence loss, the standard triplet loss and the proposed Sinkhorn-triplet loss. Both the Triplet loss and the Sinkhorn loss, when added individually, demonstrate improvements in average DSC compared to the baseline. Specifically,L ννννβννν yields noticeable gains in Kidney DSC ie. 81.55% vs 83.82%, along with improvement in HD95 for Spleen and Liver. L ννννννν‘ also shows modest DSC improvements across most organs but less consistently improves boundary metrics. L ννννννννννν‘ achieves the highest DSC scores across most organs, surpassing both the baseline and the models trained with individual losses. Notably, Liver DSC increases Optimal Transport Metric Learning in Partially Supervised Segmentation11 Background Pancreas Kidney Liver Spleen Background Pancreas Kidney Liver Spleen 0 0.520.55 0.240.4 0.52 0 0.510.6 0.8 0.550.51 00.8 0.66 0.24 0.6 0.800.21 0.40.8 0.66 0.210 Background Pancreas Kidney Liver Spleen Background Pancreas Kidney Liver Spleen 0 0.57 0.59 0.25 0.4 0.57 0 0.59 0.78 0.92 0.590.59 00.830.71 0.25 0.780.8300.13 0.4 0.92 0.710.130 0 0.2 0.4 0.6 0.8 1 Fig. 3 Heatmaps of inter-class distances after stage 1 (left) and after stage 2 (right) on the BTCV dataset. Overall, stage 2 achieves higher inter-class distances, yielding better separability. from 94.06% (Baseline) to 94.32%, and Kidney DSC improves from 81.55% to 83.50%. This strongly suggests that simultaneously optimizing for feature separability and intra- class compactness leads to the most robust volumetric segmentation performance when generalizing to unseen data. 4.4.2 Separability in the embedding space To quantify the quality of the feature space, we measured the Calinski-Harabasz (CH) Index, which evaluates both intra-class compactness and inter-class separability. The results for Stage 1 and Stage 2 are presented in Table 3. We also computed the dis- tances between class centroids in the embedding space. As shown in Figure 3, the results highlight consistently higher distances between the majority of classes in the second stage. Notably, the centroid distance between the Spleen and Pancreas, two anatomically adjacent organs, improved from 0.80 to 0.92. This suggests that the pro- posed Sinkhorn-Triplet loss effectively enhances inter-class separability, yielding more accurate segmentation results. Table 3 Comparison of Calinski-Harabasz index score in embedding space between Stage 1 and Stage 2 (Higher is better). Stage 2 demonstrates better intra-class compactness and inter-class separability. ModelCalinski-Harabasz Index Baseline743.44 w/ L ν νννννννννν‘ 958.20 12Dakini Mallam Garba and Salim Abdou Daoura (a)(b)(c)(d) Fig. 4 Visualization of segmentation results after Stage 1 and Stage 2 of our framework. (a) Original scan, (b) Ground truth annotation, (c) Stage 1 result, (d) Stage 2 result. Our framework enables more accurate voxel predictions and corrects errors from Stage 1. 0.1 0.5 11.2 1.5 86.9 87 87.1 νννννν ν·ννν ννννν 10 β2 10 β1 10 0 10 1 86.9 87 87.1 ν (blur) ν·ννν ννννν Fig. 5 Ablation study showing the effect of margin (left) and Sinkhorn blur parameter ν (right) on mean Dice score. Our method is overall robust to these hyperparamters, with slightly better result for a margin of 0.5 and small blur ν values. 4.4.3 Qualitative Results Figure 4 provides a qualitative comparison between the segmentation outputs of the baseline model and our proposed solution. The visualizations highlight more accurate segmentations achieved by our method after Stage 2, with particularly evident improve- ments for structures like the kidneys. We attribute this enhanced segmentation quality to the better-defined class clusters in the learned feature space. Optimal Transport Metric Learning in Partially Supervised Segmentation13 4.4.4 Sensitivity to hyperparameters We investigated the impact of the margin hyperparameter on the overall segmentation performance. The margin controls the desired separation between positive and negative distributions in the learned embedding space. We trained our full model with varying margin values, keeping all other hyperparameters constant, and evaluated the mean Dice score on the BTCV validation set. As shown in Figure 5 (left), the modelβs performance exhibits relative robustness to the choice of margin, particularly within the range of 0.1 to 1.0. The optimal mean Dice score of 87.15% was achieved with a margin of 0.5. Margins greater than or equal to 1.2 led to a more noticeable drop in performance, suggesting that enforcing an overly large separation might hinder the optimization process. We also made several experiments with varying blur value ν which controls the entropy regularization in the Sinkhorn divergence calculation. Results shown in Figure 5 (right) indicate that the highest performance is achieved with the smallest tested blur valueν=0.03. Performance consistently decreases slightly as the blur increases, suggesting that a more precise OT approximation is marginally beneficial. However, the overall performance variation across the tested range is minimal (less than 0.25%), demonstrating that the method is relatively robust to this hyperparameter. 5 Limitations We acknowledge some limitations that open avenues for future research. The proposed Sinkhorn-triplet loss is designed to minimize the divergence between feature distribu- tions. While effective for semantic consistency, this global alignment objective does not explicitly penalize local boundary errors. Future work could investigate integrating boundary-aware constraints, such as active contour losses, directly into the optimal transport formulation. Another Limitation is that our two-stage approach relies on the initial supervised model to generate pseudo-labels that define the anchor distributions for Stage 2. Exploring frameworks that can dynamically refine pseudo-labels without a fixed multi-stage pipeline is a promising direction to mitigate this dependency. 6 Conclusion We introduced a novel framework for multi-organ segmentation under partial supervi- sion, designed to effectively handle datasets with incomplete annotations. Our method reduces the representation gap between labeled and unlabeled domains through feature- level alignment. We believe that this approach provides a promising direction towards improving domain generalization and robustness in medical image segmentation tasks. Competing Interests The authors have no conflicts of interest to declare that are relevant to the content of this chapter. 14Dakini Mallam Garba and Salim Abdou Daoura Appendix Datasets splits used for experiments For training and validation, we used the four public partial datasets Kits19 [16], Lits19 [19] and Spleen and Pancreas from MSD [15]. The BTCV dataset is only used for test, to assess our method generalization capabilities. The different splits used are showed on Table 4. Table 4 Dataset splits used in our experiments. DatasetTraining Validation Test KiTS198020105 LiTS19521360 Pancreas (MSD)11228140 Spleen (MSD)16421 BTCVβ30 Experiments setting experiments were conducted using a NVIDIA A100 40Go. Table 5 and 6 gather the hyperparameters used for stage 1 and stage 2 respectively. The model with the best dice performance during stage 1 was used to initialize stage 2. Table 5 Stage 1 hyperparameters. HyperparameterValue OptimizerSGD Learning rate (νν )1e-3 Number of epochs1000 Momentum0.99 Confidence threshold (ν) 0.9 Table 6 Stage 2 hyperparameters. HyperparameterValue Learning rate (νν )5e-5 Learning rate (MLP)1e-4 Learning rate (prototype parameters) 1e-4 Number of epochs100 Momentum0.99 Number of prototypes per class (νΎ )5 Triplet margin (ν)0.5 Prototype loss weight (ν proto )1 Compactness weight (ν compact )0.1 Triplet loss weight (ν triplet )1 Sinkhorn divergence blur (ν )0.03 Sinkhorn scaling factor0.9 Optimal Transport Metric Learning in Partially Supervised Segmentation15 References 1. Gonz Μ alez, G., Washko, G.R., San Jos Μ e Est Μ epar, R.: Multi-structure segmentation from partially labeled datasets: application to body composition measurements on CT scans. In: International Workshop on Reconstruction and Analysis of Moving Body Organs, p. 215β224. Springer (2018) 2. Schutera, M., Rettenberger, L., Pylatiuk, C., Reischl, M.: Methods for the frugal labeler: multi- class semantic segmentation on heterogeneous labels. Plos One, 17(2), e0263656 (2022) 3. Zhang, J., Xie, Y., Xia, Y., Shen, C.: Dodnet: learning to segment multi-organ and tumors from multiple partially labeled datasets. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, p. 1195β1204 (2021) 4. Deng, R., Liu, Q., Cui, C., Asad, Z., Yang, H., Huo, Y.: Omni-Seg: a single dynamic network for multi-label renal pathology image segmentation using partially labeled data. arXiv:2112.12665 (2021) 5. Liu, H., Xu, Z., Gao, R., Li, H., Wang, J., Chabin, G., Oguz, I., Grbic, S.: Cosst: multi-organ segmentation with partially labeled datasets using comprehensive supervisions and self-training. IEEE Transactions on Medical Imaging, 43(5), 1995β2009. IEEE (2024) 6. Liu, Q., Zeng, H., Sun, Z., Li, X., Zhao, G., Liang, Y.: Many birds, one stone: medical image segmentation with multiple partially labeled datasets. Pattern Recognition, 155, 110636. Elsevier (2024) 7. Gorade, V., Sing, A., Mishra, D.: OTCXR: rethinking self-supervised alignment using optimal transport for chest X-ray analysis. In: 2025 IEEE/CVF Winter Conference on Applications of Computer Vision (WACV), p. 7143β7152. IEEE (2025) 8. Arjovsky, M., Chintala, S., Bottou, L.: Wasserstein generative adversarial networks. In: Interna- tional Conference on Machine Learning, p. 214β223. PMLR (2017) 9. Feydy, J., S Μ ejourn Μ e, T., Vialard, F.-X., Amari, S., Trouv Μ e, A., Peyr Μ e, G.: Interpolating between optimal transport and MMD using Sinkhorn divergences. In: The 22nd International Conference on Artificial Intelligence and Statistics, p. 2681β2690. PMLR (2019) 10. Shi, G., Xiao, L., Chen, Y., Zhou, S.K.: Marginal loss and exclusion loss for partially supervised multi-organ segmentation. Medical Image Analysis, 70, 101979. Elsevier (2021) 11. Ding, R., Zhou, Y., Xu, J., Xie, Y., Liang, Q., Ren, H., Wang, Y., Chen, Y., Wang, L., Huang, M.: Cross-hospital sepsis early detection via semi-supervised optimal transport with self-paced ensemble. IEEE Journal of Biomedical and Health Informatics, 27(6), 3049β3060. IEEE (2023) 12. Janati, H., Cuturi, M., Gramfort, A.: Wasserstein regularization for sparse multi-task regression. In: The 22nd International Conference on Artificial Intelligence and Statistics, p. 1407β1416. PMLR (2019) 13. Sarlin, P.-E., DeTone, D., Malisiewicz, T., Rabinovich, A.: SuperGlue: learning feature matching with graph neural networks. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, p. 4938β4947 (2020) 14. Demir, A., Massaad, E., Kiziltan, B.: Topology-aware focal loss for 3D image segmentation. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, p. 580β589 (2023) 15. Antonelli, M., Reinke, A., Bakas, S., Farahani, K., Kopp-Schneider, A., Landman, B.A., Litjens, G., Menze, B., Ronneberger, O., Summers, R.M., et al.: The medical segmentation decathlon. Nature Communications, 13(1), 4128 (2022) 16. Heller, N., Isensee, F., Maier-Hein, K.H., Hou, X., Xie, C., Li, F., Nan, Y., Mu, G., Lin, Z., Han, M., et al.: The state of the art in kidney and kidney tumor segmentation: results of the KiTS19 challenge. Medical Image Analysis, 67, 101821. Elsevier (2021) 17. Landman, B., Xu, Z., Igelsias, J., Styner, M., Langerak, T., Klein, A.: MICCAI multi-atlas labeling beyond the cranial vaultβworkshop and challenge. In: Proc. MICCAI Multi-Atlas Labeling Beyond Cranial VaultβWorkshop Challenge, vol. 5, p. 12 (2015) 18. Isensee, F., Jaeger, P.F., Kohl, S.A.A., Petersen, J., Maier-Hein, K.H.: nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation. Nature Methods, 18(2), 203β211 (2021) 16Dakini Mallam Garba and Salim Abdou Daoura 19. Bilic, P., Christ, P., Li, H.B., Vorontsov, E., Ben-Cohen, A., Kaissis, G., Szeskin, A., Jacobs, C., Mamani, G.E., Chartrand, G., et al.: The liver tumor segmentation benchmark (LiTS). Medical Image Analysis, 84, 102680. Elsevier (2023) 20. Dou, J.X., Luo, L., Yang, R.M.: An optimal transport approach to deep metric learning (student abstract). In: Proceedings of the AAAI Conference on Artificial Intelligence, 36(11), 12935β 12936 (2022) 21. Varior, R.R., Shuai, B., Lu, J., Xu, D., Wang, G.: A siamese long short-term memory architecture for human re-identification. In: European Conference on Computer Vision, p. 135β153. Springer (2016) 22. Schroff, F., Kalenichenko, D., Philbin, J.: FaceNet: a unified embedding for face recognition and clustering. In: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, p. 815β823 (2015) 23. Vinyals, O., Blundell, C., Lillicrap, T., Wierstra, D., et al.: Matching networks for one-shot learning. Advances in Neural Information Processing Systems, 29 (2016) 24. Hadsell, R., Chopra, S., LeCun, Y.: Dimensionality reduction by learning an invariant mapping. In: 2006 IEEE Conference on Computer Vision and Pattern Recognition (CVPRβ06), vol. 2, p. 1735β1742. IEEE (2006) 25. Kulis, B., Jordan, M.I.: Revisiting k-means: new algorithms via Bayesian nonparametrics. arXiv:1111.0352 (2011) 26. Fang, X., Yan, P.: Multi-organ segmentation over partially labeled datasets with multi-scale feature abstraction. IEEE Transactions on Medical Imaging, 39(11), 3619β3629. IEEE (2020) 27. Liu, P., Deng, Y., Wang, C., Hui, Y., Li, Q., Li, J., Luo, S., Sun, M., Quan, Q., Yang, S., et al.: Universal segmentation of 33 anatomies. arXiv:2203.02098 (2022) 28. Huang, R., Zheng, Y., Hu, Z., Zhang, S., Li, H.: Multi-organ segmentation via co-training weight- averaged models from few-organ datasets. In: International Conference on Medical Image Com- puting and Computer-Assisted Intervention, p. 146β155. Springer (2020) 29. Jiang, X., Zhang, D., Li, X., Liu, K., Cheng, K.-T., Yang, X.: Labeled-to-unlabeled distribu- tion alignment for partially supervised multi-organ medical image segmentation. Medical Image Analysis, 99, 103333. Elsevier (2025)