Generalized Wasserstein Dice Loss

• The paper introduces GWDL as a semantically informed extension of Dice loss that leverages Wasserstein distance to incorporate inter-class relationships.
• It replaces traditional overlap measures with hierarchy-aware penalties, improving performance in challenging multi-class segmentation tasks such as in neuroimaging.
• Implementation via standard tensor operations ensures computational efficiency while reducing anatomically implausible segmentation errors.

The Generalized Wasserstein Dice Loss (GWDL) is a semantically informed extension of the classic Dice loss, designed to provide robust training supervision for multi-class segmentation tasks, especially in contexts with severe class imbalance and inter-class semantic structure. By incorporating the Wasserstein distance with a user-defined ground metric over the label space, GWDL encodes prior knowledge about class relationships, favoring reasonable misclassifications and mitigating implausible ones. GWDL has become particularly prominent in neuroimaging challenges such as BraTS, where tumor subregions possess a natural hierarchy and standard losses underperform on rare or ambiguous classes (Fidon et al., 2017, Fidon et al., 2021, Fidon et al., 2020).

1. Mathematical Foundations

Let $\mathcal{X}$ denote the spatial domain (e.g., voxels or pixels), and let $\mathcal{L} = \{0, 1, \ldots, L-1\}$ denote the finite set of labels. For each voxel $i \in \mathcal{X}$, the ground truth is a one-hot vector $g^i \in \{0,1\}^{|\mathcal{L}|}$, while the network produces a probability vector $p^i \in \Delta^{|\mathcal{L}|}$, the simplex over $\mathcal{L}$.

A symmetric ground-distance matrix $M \in \mathbb{R}^{L \times L}$, with entries $M_{l,l'} \geq 0$ and $M_{l,l} = 0$, encodes the semantic cost of confounding label $l$ for $\mathcal{L} = \{0, 1, \ldots, L-1\}$0. For instance, $\mathcal{L} = \{0, 1, \ldots, L-1\}$1 can reflect a tree structure among tumor sub-regions, or simple $\mathcal{L} = \{0, 1, \ldots, L-1\}$2-$\mathcal{L} = \{0, 1, \ldots, L-1\}$3 loss for flat class weighting.

The core of GWDL is to replace the Dice overlap by a Wasserstein-aware term. The (1-)Wasserstein distance at voxel $\mathcal{L} = \{0, 1, \ldots, L-1\}$4 between network output $\mathcal{L} = \{0, 1, \ldots, L-1\}$5 and label $\mathcal{L} = \{0, 1, \ldots, L-1\}$6 (assuming $\mathcal{L} = \{0, 1, \ldots, L-1\}$7 is one-hot for class $\mathcal{L} = \{0, 1, \ldots, L-1\}$8) is: $\mathcal{L} = \{0, 1, \ldots, L-1\}$9 The total “error” across all voxels, $i \in \mathcal{X}$0, is the sum of $i \in \mathcal{X}$1 over $i \in \mathcal{X}$2.

The Dice-like overlap substitutes the “true positive” mass by

$i \in \mathcal{X}$3

where $i \in \mathcal{X}$4 indexes the background class, and $i \in \mathcal{X}$5 is the cost of transporting label $i \in \mathcal{X}$6 to background. These are weighted by $i \in \mathcal{X}$7 (usually $i \in \mathcal{X}$8 or $i \in \mathcal{X}$9) and summed: $g^i \in \{0,1\}^{|\mathcal{L}|}$0 The Generalized Wasserstein Dice Score is

$g^i \in \{0,1\}^{|\mathcal{L}|}$1

The training loss is

$g^i \in \{0,1\}^{|\mathcal{L}|}$2

Expansions, differentiation w.r.t. probabilities, and batched pseudocode are provided in implementation guides (Fidon et al., 2017, Fidon et al., 2021).

2. Integration of the Wasserstein Metric with Dice-Type Overlap

Classic Dice loss rewards exact overlap, penalizing all misclassifications equally. GWDL generalizes this by discounting overlap terms proportionally to their semantic discrepancy. Each “true positive” becomes $g^i \in \{0,1\}^{|\mathcal{L}|}$3, where $g^i \in \{0,1\}^{|\mathcal{L}|}$4 is small if the prediction places mass on similar classes, and large if it is distributed to semantically distant ones. In the case of soft labels or fuzzy annotations, the optimal transport formulation can be regularized via Sinkhorn entropy, though for one-hot vs. softmax the dot-product reduction is exact and computationally efficient (Fidon et al., 2021).

This mechanism makes GWDL hierarchy-aware and tolerant of minor, semantically reasonable errors, aligning training gradients with clinically meaningful distinctions.

3. Implementation and Computational Considerations

GWDL is implemented via simple tensor operations without custom solvers for moderate $g^i \in \{0,1\}^{|\mathcal{L}|}$5:

• Compute per-voxel Wasserstein distances by dot products between prediction vectors and relevant rows of $g^i \in \{0,1\}^{|\mathcal{L}|}$6.
• Aggregate numerator and denominator for the generalized Dice ratio.
• Differentiate automatically through all steps: $g^i \in \{0,1\}^{|\mathcal{L}|}$7.
• Practical recommendations include adding a small $g^i \in \{0,1\}^{|\mathcal{L}|}$8 to denominators, warm-starting with mean Dice loss for "stiff" $g^i \in \{0,1\}^{|\mathcal{L}|}$9, and vectorizing batched operations using frameworks like PyTorch or TensorFlow (Fidon et al., 2017).

The dominant computational cost scales as $p^i \in \Delta^{|\mathcal{L}|}$0, where $p^i \in \Delta^{|\mathcal{L}|}$1 is the number of voxels and $p^i \in \Delta^{|\mathcal{L}|}$2 the number of classes—negligible for standard medical segmentation tasks.

4. Hyperparameters and Ground Distance Matrix Design

GWDL introduces no scalar hyperparameters beyond the cost matrix $p^i \in \Delta^{|\mathcal{L}|}$3 and possible weighting vector $p^i \in \Delta^{|\mathcal{L}|}$4. The structure of $p^i \in \Delta^{|\mathcal{L}|}$5 is critical for encoding domain knowledge:

• $p^i \in \Delta^{|\mathcal{L}|}$6 for all $p^i \in \Delta^{|\mathcal{L}|}$7 yields standard multiclass Dice loss.
• Hierarchical (tree-based) $p^i \in \Delta^{|\mathcal{L}|}$8 allows tuning confusion penalties (e.g., adjacent tumor subregions given lower cost than background/tumor confusions).
• Clinically informed matrices have been empirically validated on BraTS datasets, with choices such as

$p^i \in \Delta^{|\mathcal{L}|}$9

Use of class weights $\mathcal{L}$0 as $\mathcal{L}$1 further increases attention to rare or non-background structures (Fidon et al., 2021, Fidon et al., 2017). Pre-training on mean Dice—in effect, a warm-start—mitigates optimization difficulties especially when $\mathcal{L}$2 imposes strong constraints (Fidon et al., 2017).

5. Empirical Results and Robustness Analyses

Quantitative results across BraTS challenges show GWDL offers consistent improvement over mean-class Dice and Dice+cross-entropy loss combinations, particularly for rare structures and worst-case scenarios:

• On BraTS 2021, GWDL + CE achieved higher mean Dice and lower Hausdorff 95th percentile distances than alternatives (e.g., mean Dice + CE) (Fidon et al., 2021).
• GWDL improved the 5th-percentile Dice for enhancing tumor from $\mathcal{L}$3 to $\mathcal{L}$4.
• Qualitative analyses indicate GWDL reduces anatomically implausible errors and yields smoother segmentations.
• No architectural modifications or optimizer pairings are required. The benefit of GWDL is orthogonal to network design (e.g., U-Net vs. U-Net with Transformer bottleneck) (Fidon et al., 2021, Fidon et al., 2020).

A summary of empirical performance is given in the table below:

Dataset/Challenge	Loss Function	Mean Dice (WT/TC/ET)	Hausdorff95 (WT/TC/ET)
BraTS 2021	Dice+CE	92.5/86.5/81.7	-
	GWDL+CE	92.5/86.4/82.6	Lower than baseline
BraTS 2015 (test)	mean Dice	83/70/68	-
	GWDL (tree)	88/73/70	-
	GWDL (tree, pretrain)	89/73/74	-

GWDL also reduces the occurrence of semantically implausible confusions, such as mislabeling enhancing tumor as edema, which is discouraged by appropriate structure in $\mathcal{L}$5 (Fidon et al., 2017).

6. Theoretical Motivation and Limitations

The motivation for GWDL is to align training signals with the semantic structure of the segmentation task. Standard Dice loss treats all misclassifications equivalently, whereas GWDL can weight errors according to hierarchy or anatomical relationships, offering more robust learning under imbalance, boundary ambiguity, and clinical uncertainty (Fidon et al., 2017, Fidon et al., 2021, Fidon et al., 2020).

Limitations include:

• Sensitivity to the design of the cost matrix $\mathcal{L}$6, requiring domain expertise or empirical tuning.
• For large $\mathcal{L}$7, potential computational overhead, although mitigated by sparsity or low-rank approximations.
• Non-convexity inherited from the Dice structure, with associated risks of local minima.

Extension to soft labels or fuzzy annotation regimes is possible by using the full optimal transport LP (with entropic regularization if smoothness is needed), though typical medical applications use one-hot ground truth.

7. Practical Usage and Recommendations

GWDL is modular and integrates into standard CNN segmentation pipelines. Key practical guidelines include:

• Use GWDL as a plug-in replacement for Dice loss in PyTorch/TensorFlow segmentation codebases; per-batch loss computation is vectorized with batched matrix operations.
• Select or design $\mathcal{L}$8 in consultation with clinical domain experts or via cross-validation.
• Warm-start with standard Dice loss before switching to GWDL for best convergence in strongly constrained $\mathcal{L}$9 scenarios.
• Monitor per-class performance and tail metrics (e.g., Hausdorff percentiles) to capture robustness improvements.

Empirical results from BraTS 2020 and 2021 challenges demonstrate GWDL’s advantage as a principled, lightweight means for adding semantic structure-awareness to Dice-style segmentation losses, yielding improved accuracy and reliability on multi-class imbalanced problems encountered in clinical imaging (Fidon et al., 2021, Fidon et al., 2020, Fidon et al., 2017).