# Optimal Transport Metric Learning for Feature Alignment in Partially Supervised Segmentation

arXiv: https://arxiv.org/abs/2609.19176 · HTML: https://mlsanity.com/paper/2609.19176/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 16 of its 28 references that are in the map.
- First posted on arXiv 15 Sep 2026.

## Builds on (most cited first, 16 of 16)

- [FaceNet: A Unified Embedding for Face Recognition and Clustering](https://mlsanity.com/paper/1503.03832/) (2015)
- [Matching Networks for One Shot Learning](https://mlsanity.com/paper/1606.04080/) (2016)
- [SuperGlue: Learning Feature Matching with Graph Neural Networks](https://mlsanity.com/paper/1911.11763/) (2019)
- [The Medical Segmentation Decathlon](https://mlsanity.com/paper/2106.05735/) (2021)
- [The Liver Tumor Segmentation Benchmark (LiTS)](https://mlsanity.com/paper/1901.04056/) (2019)
- [DoDNet: Learning to segment multi-organ and tumors from multiple partially labeled datasets](https://mlsanity.com/paper/2011.10217/) (2020)
- [A Siamese Long Short-Term Memory Architecture for Human Re-Identification](https://mlsanity.com/paper/1607.08381/) (2016)
- [Multi-organ Segmentation via Co-training Weight-averaged Models from Few-organ Datasets](https://mlsanity.com/paper/2008.07149/) (2020)
- [Wasserstein regularization for sparse multi-task regression](https://mlsanity.com/paper/1805.07833/) (2018)
- [COSST: Multi-organ Segmentation with Partially Labeled Datasets Using Comprehensive Supervisions and Self-training](https://mlsanity.com/paper/2304.14030/) (2023)
- [Topology-Aware Focal Loss for 3D Image Segmentation](https://mlsanity.com/paper/2304.12223/) (2023)
- [Universal Segmentation of 33 Anatomies](https://mlsanity.com/paper/2203.02098/) (2022)
- [Cross-hospital Sepsis Early Detection via Semi-supervised Optimal Transport with Self-paced Ensemble](https://mlsanity.com/paper/2106.10352/) (2021)
- [Omni-Seg: A Single Dynamic Network for Multi-label Renal Pathology Image Segmentation using Partially Labeled Data](https://mlsanity.com/paper/2112.12665/) (2021)
- [OTCXR: Rethinking Self-supervised Alignment using Optimal Transport for Chest X-ray Analysis](https://mlsanity.com/paper/2404.11868/) (2024)
- [Labeled-to-Unlabeled Distribution Alignment for Partially-Supervised Multi-Organ Medical Image Segmentation](https://mlsanity.com/paper/2409.03228/) (2024)

## Abstract (arXiv)

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.