Runtime Freezing: Dynamic Class Loss for Multi-Organ 3D Segmentation
James Willoughby, Irina Voiculescu
TL;DR
The paper tackles severe class imbalance in multi-organ 3D CT segmentation by introducing dynamic, per-class loss weighting guided by per-class performance scores $S_i$ and weights $w_i \in [0,1]$. It proposes three update rules—Threshold Class Freezing (TCF), Plateau Class Freezing (PCF), and Class Boost Strategy (CBS)—to adaptively emphasize harder classes while maintaining training stability, and demonstrates these strategies across multiple 3D segmentation architectures. A key methodological contribution is the 3D extension of the Lambda network, Lambda3D, which is compared against 3DUNet and UNETR on public abdominal datasets. Across networks, CBS yields the most consistent improvements for poorly performing classes and enhances overall mean Dice, highlighting the approach's generality and practical impact for imbalanced multi-organ segmentation; Lambda3D also proves to be a viable, efficient 3D alternative with competitive performance. The work suggests that dynamic, class-aware loss adjustment can be a broadly applicable technique to improve per-class segmentation performance in medical imaging, especially when data are unevenly distributed.
Abstract
Segmentation has become a crucial pre-processing step to many refined downstream tasks, and particularly so in the medical domain. Even with recent improvements in segmentation models, many segmentation tasks remain difficult. When multiple organs are segmented simultaneously, difficulties are due not only to the limited availability of labelled data, but also to class imbalance. In this work we propose dynamic class-based loss strategies to mitigate the effects of highly imbalanced training data. We show how our approach improves segmentation performance on a challenging Multi-Class 3D Abdominal Organ dataset.
