Table of Contents
Fetching ...

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.

RankSEG-RMA: An Efficient Segmentation Algorithm via Reciprocal Moment Approximation

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 . 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., 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.
Paper Structure (28 sections, 5 theorems, 23 equations, 3 figures, 8 tables, 4 algorithms)

This paper contains 28 sections, 5 theorems, 23 equations, 3 figures, 8 tables, 4 algorithms.

Key Result

Theorem 2.1

Assume that $Y_i \perp Y_j | \mathbf{X}$. A segmentation rule $\bm{\delta}^*$ is a global maximizer of $\mathbb{E}(\text{Dice}^\text{I}(\bm{\delta}))$ if and only if $\delta^*_j = \mathbbm{1}(p_j \geq p_{j_{\tau^*}})$, where $j_{\tau}$ is the index with the $\tau$-th largest probability. The optimal where $\mathcal{J}_{\tau}=\{j: \sum_{j^{\prime}=1}^d \mathbbm{1}(p_{j^{\prime}} \ge p_{j_{\tau}})\}

Figures (3)

  • Figure 1: Comparison of Argmax-prob and RankSEG-RMA. (a) Argmax-prob: each pixel is predicted to the class with the highest probability. (b) RankSEG-RMA: segmentation masks $\widehat{\mathcal{I}}_c^+$ for each class are obtained independently; non-overlapping parts $\widehat{\mathcal{I}}_c$ are accepted, while overlaps $\widehat{\mathcal{I}}^{\text{overlap}}$ are resolved by applying argmax over RMA-scores.
  • Figure 2: Image-wise performance and an example of a worst-case segmentation on KiTS. The left plot presents the IoU for each image, with indices sorted in ascending order according to IoU under Argmax-prob. The right plot displays the segmentation results for a slice of a worst-case image.
  • Figure 3: Qualitative visualizations on Pascal VOC. From left to right: input image, ground truth, prediction by Argmax-prob, and prediction by RankSEG-RMA. The key differences are highlighted in red boxes.

Theorems & Definitions (8)

  • Theorem 2.1: The Bayes rule for Dice$^\text{I}$-segmentation dai2023rankseg
  • Example
  • Theorem 3.1: Reciprocal moment approximation to RankSEG
  • Definition A.1: Calibration
  • Theorem B.1: The Bayes rule for IoU$^\text{I}$-segmentation dai2023rankseg
  • Lemma C.1: chao1972negative
  • Lemma C.2: wooff1985bounds
  • proof