[Paper Review] Region Mutual Information Loss for Semantic Segmentation

Shuai Zhao, Yang Wang+2 more|arXiv (Cornell University)|Oct 26, 2019

Robotics and Automated SystemsEngineering83 citations

TL;DR

RMI loss models region-level pixel dependencies to improve semantic segmentation performance, yielding consistent gains on VOC2012 and CamVid without extra inference overhead.

ABSTRACT

Semantic segmentation is a fundamental problem in computer vision. It is considered as a pixel-wise classification problem in practice, and most segmentation models use a pixel-wise loss as their optimization riterion. However, the pixel-wise loss ignores the dependencies between pixels in an image. Several ways to exploit the relationship between pixels have been investigated, \eg, conditional random fields (CRF) and pixel affinity based methods. Nevertheless, these methods usually require additional model branches, large extra memories, or more inference time. In this paper, we develop a region mutual information (RMI) loss to model the dependencies among pixels more simply and efficiently. In contrast to the pixel-wise loss which treats the pixels as independent samples, RMI uses one pixel and its neighbour pixels to represent this pixel. Then for each pixel in an image, we get a multi-dimensional point that encodes the relationship between pixels, and the image is cast into a multi-dimensional distribution of these high-dimensional points. The prediction and ground truth thus can achieve high order consistency through maximizing the mutual information (MI) between their multi-dimensional distributions. Moreover, as the actual value of the MI is hard to calculate, we derive a lower bound of the MI and maximize the lower bound to maximize the real value of the MI. RMI only requires a few extra computational resources in the training stage, and there is no overhead during testing. Experimental results demonstrate that RMI can achieve substantial and consistent improvements in performance on PASCAL VOC 2012 and CamVid datasets. The code is available at https://github.com/ZJULearning/RMI.

Motivation & Objective

• Motivate improving segmentation by incorporating pixel dependencies beyond pixel-wise losses.
• Propose a region-based mutual information loss to enforce high-order consistency between predictions and ground truth.
• Make RMI training-efficient with minimal additional memory and no extra inference cost.
• Enable straightforward integration into existing segmentation frameworks without changing base models.

Proposed method

• Represent each pixel by a region of neighboring pixels (e.g., 3x3) to form a high-dimensional point.
• Cast an image into a distribution of these high-dimensional points for predictions and ground truth.
• Derive a tractable lower bound of mutual information I(Y;P) and maximize this bound during training.
• Approximate the posterior Y|P variance using a second-order independence assumption and a closed-form covariance expression.
• Normalize and stabilize the MI bound computation with a tractable matrix M and Cholesky decomposition.
• Combine RMI with standard cross-entropy in a joint loss with a balancing parameter.

Experimental results

Research questions

• RQ1 Can a region-based mutual information objective improve segmentation accuracy over pixel-wise losses?
• RQ2 How to compute a lower bound on mutual information that is practical for deep learning training?
• RQ3 What are the trade-offs in downsampling and region size for RMI in terms of performance and resource usage?
• RQ4 Does RMI generalize across different segmentation backbones and datasets?

Key findings

• RMI yields substantial and consistent mIoU improvements on VOC2012 val/test sets across DeepLabv3 and DeepLabv3+ baselines.
• RMI outperforms CRF post-processing and affinity-field losses under the reported settings, with no additional inference cost.
• RMI also provides notable gains on the CamVid dataset, indicating broad applicability across datasets.
• A downsampling strategy (with average pooling) and moderate region sizes balance performance and memory usage effectively.
• Ablation shows larger region sizes and smaller downsampling factor generally improve performance but increase compute
• Per-class results show improved segmentation for several categories, reflecting better boundary and detail capture.