[Paper Review] Tree Energy Loss: Towards Sparsely Annotated Semantic Segmentation

Zhiyuan Liang, Tiancai Wang+3 more|arXiv (Cornell University)|Mar 21, 2022

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes Tree Energy Loss (TEL), a novel loss function for sparsely annotated semantic segmentation that uses minimum spanning trees (MSTs) to model low-level (color) and high-level (feature) affinities, enabling dynamic, coarse-to-fine generation of soft pseudo labels for unlabeled pixels. TEL achieves state-of-the-art performance across point-, scribble-, and block-wise annotation settings without multi-stage training, extra data, or post-processing, outperforming prior methods by up to 7.3% mIoU gain.

ABSTRACT

Sparsely annotated semantic segmentation (SASS) aims to train a segmentation network with coarse-grained (i.e., point-, scribble-, and block-wise) supervisions, where only a small proportion of pixels are labeled in each image. In this paper, we propose a novel tree energy loss for SASS by providing semantic guidance for unlabeled pixels. The tree energy loss represents images as minimum spanning trees to model both low-level and high-level pair-wise affinities. By sequentially applying these affinities to the network prediction, soft pseudo labels for unlabeled pixels are generated in a coarse-to-fine manner, achieving dynamic online self-training. The tree energy loss is effective and easy to be incorporated into existing frameworks by combining it with a traditional segmentation loss. Compared with previous SASS methods, our method requires no multistage training strategies, alternating optimization procedures, additional supervised data, or time-consuming post-processing while outperforming them in all SASS settings. Code is available at https://github.com/megvii-research/TreeEnergyLoss.

Motivation & Objective

• Address the challenge of training accurate semantic segmentation models with minimal annotations, such as points, scribbles, or blocks.
• Overcome limitations of existing SASS methods, including reliance on auxiliary tasks, time-consuming proposals, or inconsistent regularization and consistency learning.
• Provide a unified, end-to-end framework that enables dynamic online self-training using unlabeled pixels without alternating optimization or external supervision.
• Introduce a block-wise annotation setting to better evaluate and grade performance across different annotation sparsity levels.

Proposed method

• Construct two minimum spanning trees (MSTs) from image color (RGB) and deep CNN features to model low-level and high-level pixel affinities, respectively.
• Use the MST structure to preserve essential pixel relationships by iteratively removing high-dissimilarity edges, ensuring structural consistency in affinity matrices.
• Generate affinity matrices by accumulating edge weights along MST paths, capturing both local and global contextual relationships.
• Apply a cascaded filtering mechanism: first refine predictions using low-level affinity, then further refine with high-level affinity to produce increasingly accurate soft pseudo labels.
• Combine the tree energy loss with a standard cross-entropy segmentation loss in a single-stage training setup, enabling online self-training.
• Integrate TEL into existing segmentation networks with minimal modification, making it plug-and-play and computationally efficient.

Experimental results

Research questions

• RQ1 Can minimum spanning trees effectively model both low-level and high-level affinities to guide pseudo-label generation in sparse semantic segmentation?
• RQ2 Does a cascaded, coarse-to-fine refinement of pseudo labels using MST-based affinities improve segmentation performance compared to simultaneous or sequential alternatives?
• RQ3 Can TEL achieve state-of-the-art performance without requiring multi-stage training, additional supervised data, or complex post-processing?
• RQ4 How does the proposed block-wise annotation setting compare to point- and scribble-wise settings in evaluating SASS methods?

Key findings

• TEL achieves state-of-the-art performance across all SASS settings—point, scribble, and block annotation—without requiring multi-stage training or external data.
• On the Pascal VOC 2012 dataset, TEL improves mIoU by 7.3% over the baseline during training, demonstrating strong knowledge distillation from unlabeled data.
• The method achieves 1.7% higher mIoU than Non-Local-based high-level affinity, showing superior feature-level affinity modeling.
• The cascading aggregation strategy (LH-C) outperforms parallel and other cascading variants, confirming the benefit of refining predictions in a coarse-to-fine manner.
• The hyperparameter sensitivity analysis shows that λ = 0.4 and σ = 0.02 yield optimal performance, with minimal sensitivity to σ values.
• Pseudo label precision exceeds network prediction precision early in training, indicating effective and reliable self-supervision.