Hierarchical Cautious Optimization for Semi-Supervised Medical Image Segmentation with Limited Labeled Data

ICLR 2026 Conference Submission13920 Authors

Keywords: Semi Supervised Learning, MRI, Medical imaging, Label-efficient learning, Optimization algorithms

TL;DR: We propose Hierarchical Cautious Optimization (HCO), which builds momentum from labeled gradients and only incorporates aligned unlabeled ones, boosting semi-supervised segmentation across datasets with minimal cost.

Abstract: Semi-supervised learning (SSL) effectively addresses limited labeled data challenges in volumetric medical image segmentation by leveraging both ground-truth labels and pseudo-labels from unlabeled data. However, conventional optimizers treat gradients from labeled and unlabeled sources equally, often leading to either over-trust or over-rejection of pseudo-label signals. We introduce Hierarchical Cautious Optimization (HCO), which establishes a trust hierarchy between gradient sources. HCO computes momentum estimates using only gradients from labeled data and incorporates unlabeled gradients only when they align with this trusted direction. Our approach integrates into existing momentum-based optimizers with minimal implementation effort and computational cost. Evaluations across three datasets demonstrate consistent performance improvements, particularly on a challenging fetal MRI dataset where Dice scores for fetal lungs and liver increased from 0.68 to 0.84 and 0.71 to 0.82, respectively. The consistent gains across optimizers and datasets, combined with minimal implementation overhead, position HCO as a practical enhancement for existing SSL medical segmentation pipelines.

Primary Area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Submission Number: 13920