RankSEG-RMA: An Efficient Segmentation Algorithm via Reciprocal Moment Approximation
Zixun Wang, Ben Dai
TL;DR
RankSEG-RMA tackles the inefficiency of previous metric-consistent segmentation approaches by introducing Reciprocal Moment Approximation (RMA), which reduces the complexity of optimizing IoU and Dice from quadratic or near-quadratic to linear in the number of pixels $d$. It also provides a scalable, non-overlapping multiclass extension via an RMA-based pixel-wise scoring mechanism, aligning predictions with Bayes-optimal Dice/IoU while maintaining practical performance. The authors prove a theoretical error bound for the RMA approximation and demonstrate substantial speedups (e.g., $48\times$ on LiTS) with competitive accuracy across diverse vision datasets and models. Overall, RankSEG-RMA delivers a principled, efficient framework for metric-consistent segmentation that scales to large images and supports non-overlapping predictions, offering clear practical advantages for real-world segmentation tasks.
Abstract
Semantic segmentation labels each pixel in an image with its corresponding class, and is typically evaluated using the Intersection over Union (IoU) and Dice metrics to quantify the overlap between predicted and ground-truth segmentation masks. In the literature, most existing methods estimate pixel-wise class probabilities, then apply argmax or thresholding to obtain the final prediction. These methods have been shown to generally lead to inconsistent or suboptimal results, as they do not directly maximize segmentation metrics. To address this issue, a novel consistent segmentation framework, RankSEG, has been proposed, which includes RankDice and RankIoU specifically designed to optimize the Dice and IoU metrics, respectively. Although RankSEG almost guarantees improved performance, it suffers from two major drawbacks. First, it is its computational expense-RankDice has a complexity of O(d log d) with a substantial constant factor (where d represents the number of pixels), while RankIoU exhibits even higher complexity O(d^2), thus limiting its practical application. For instance, in LiTS, prediction with RankSEG takes 16.33 seconds compared to just 0.01 seconds with the argmax rule. Second, RankSEG is only applicable to overlapping segmentation settings, where multiple classes can occupy the same pixel, which contrasts with standard benchmarks that typically assume non-overlapping segmentation. In this paper, we overcome these two drawbacks via a reciprocal moment approximation (RMA) of RankSEG with the following contributions: (i) we improve RankSEG using RMA, namely RankSEG-RMA, reduces the complexity of both algorithms to O(d) while maintaining comparable performance; (ii) inspired by RMA, we develop a pixel-wise score function that allows efficient implementation for non-overlapping segmentation settings.
