Label tree semantic losses for rich multi-class medical image segmentation

Technical Summary: Label Tree Semantic Losses for Rich Multi-Class Medical Image Segmentation

Problem Statement
Rich and accurate medical image segmentation is critical for pre-operative planning, intra-operative navigation, and post-operative assessment. However, standard learning methods for segmentation tasks typically penalize all classification errors equally, failing to exploit the semantic relationships within the label space. This limitation becomes increasingly problematic as the cardinality of labels grows to include subtly different classes (e.g., specific brain sub-regions or tissue types). Furthermore, in surgical imaging and large-scale annotation efforts, data is often sparsely annotated with "positive-only" labels, where unlabelled regions cannot be assumed to be background. Existing approaches often lack mechanisms to leverage label hierarchies or handle the ambiguity of unlabelled pixels in a semantically meaningful way.

Methodology
The authors propose a framework that incorporates a hierarchical organization of labels (a label tree) into the training process via two complementary semantic loss functions. These losses are designed to work with both dense supervision (e.g., whole-brain parcellation) and sparse, background-free annotations (e.g., neurosurgical hyperspectral imaging).

1. Wasserstein Distance-Based Segmentation Loss ($\mathcal{L}_{wass+seg}$):
   This loss utilizes the Wasserstein distance (optimal transport) to measure the cost of transforming predicted probability distributions into ground truth distributions. The "ground distance" matrix ($M$) is derived from the label tree, where the distance between two leaf nodes is defined by the path length between them in the hierarchy. This formulation penalizes errors based on semantic severity; misclassifying a class as a semantically distant sibling incurs a higher penalty than misclassifying it as a close relative. The final loss is a compound of the Wasserstein distance and a generic segmentation loss (e.g., Dice or Cross-Entropy).

2. Tree-Weighted Semantic Cross-Entropy Loss ($\mathcal{L}_{twce+seg}$):
   This approach extends standard Cross-Entropy (CE) by aggregating probabilities across all nodes in the label tree, not just the leaf nodes. Using an adjacency matrix to represent parent-child relationships, the model computes aggregated probabilities for internal nodes. The loss is then calculated as a weighted sum of CE across all nodes, where weights are assigned based on the edge structure. This allows the model to receive supervision signals at multiple levels of abstraction (coarse and fine).

3. Integration with Sparse Positive-Only Learning:
   The proposed losses are integrated into a framework for training with sparse, positive-only annotations (based on Wang et al., 2026). In this setting, the model is trained only on annotated positive classes. During inference, an Out-of-Distribution (OOD) detection mechanism flags pixels with low maximum predicted probability as background. The authors extend this by aggregating probabilities at a chosen hierarchy level (e.g., top-level nodes) to determine the OOD threshold, allowing for more robust background detection without degrading performance on positive classes.

Key Contributions

• Novel Loss Functions: Introduction of two tree-based semantic losses (Wasserstein-based and tree-weighted CE) that explicitly encode hierarchical label relationships during optimization.
• Versatile Framework: Demonstration of how these losses apply to both densely supervised whole-brain parcellation (WBP) and sparsely supervised neurosurgical hyperspectral imaging (HSI), including integration with OOD detection for background identification.
• Performance Improvement: Empirical evidence showing that incorporating expert-defined or anatomically derived label hierarchies leads to more semantically meaningful supervision and improved segmentation performance compared to standard baselines.
• Robustness on Small Structures: Specific improvements observed on small anatomical structures with limited annotations, where standard baselines often struggle.

Experimental Results
The authors evaluated their approach on two distinct tasks:

1. Whole Brain Parcellation (WBP): Using 3D T1-weighted MRI data from three datasets (Mindboggle42, AOMIC, IXI) with full supervision.
   • Findings: The Wasserstein-based compound loss ($\mathcal{L}_{wass+seg}$) consistently outperformed the standard nnU-Net baseline (Dice + CE) and the Generalised Wasserstein Dice loss. The improvements were most significant for metrics related to small structures (Dicesmall, NSDsmall) and surface overlap (NSD). Statistical analysis confirmed significant gains in NSD and small-structure metrics across nearly all train-test settings.
2. Neurosurgical Hyperspectral Imaging (HSI): Using sparse, positive-only annotations from intra-operative surgical videos.
   • Findings: The proposed framework improved performance on top-level classes (e.g., distinguishing normal vs. abnormal tissue) compared to baselines trained only on leaf nodes. The hierarchy-weighted top-level supervision ($\mathcal{L}_{twce+seg}$ with specific edge weights) showed strong support in the sparse setting. Confusion matrix analysis revealed that the proposed methods reduced semantically implausible errors (e.g., confusing normal and abnormal tissues) compared to the baseline.

Significance and Claims
The paper claims that leveraging the structure of the label space as a source of information is crucial for rich segmentation tasks, particularly when dealing with subtle class variations and sparse data. By moving beyond uniform error penalization, the proposed methods enable models to make "better mistakes" (i.e., errors that are semantically closer to the truth). The authors emphasize that their approach improves upon standard baselines without requiring architectural modifications to the segmentation networks (e.g., nnU-Net or U-Net), making it a practical extension for existing pipelines. The work bridges a gap in surgical imaging, where such hierarchical semantic supervision has previously been unexplored, and provides a scalable solution for training on large, sparsely annotated datasets.