Differentiable clDice Loss
- Differentiable clDice Loss is a topological loss function that directly optimizes centerline overlap, ensuring connectivity preservation in segmentation tasks.
- It employs soft skeletonization via differentiable operations, allowing standard backpropagation in deep networks for tubular structure segmentation.
- Empirical studies show that integrating clDice loss improves topology fidelity and structural accuracy in medical imaging, neural mapping, and road network extraction.
The differentiable clDice loss is a topological loss function for segmentation of network-like tubular structures, designed to directly optimize for centerline (skeleton) overlap between predicted and ground-truth segmentations within deep learning frameworks. Unlike traditional volumetric losses such as Dice, which prioritize foreground voxel accuracy and can disproportionately reward thick structures while neglecting topology, the clDice loss ensures that the resulting predictions more faithfully preserve the connectivity and branching crucial in domains such as vascular imaging, neural mapping, and road network extraction. Differentiable clDice loss, or "soft-clDice," is constructed to be fully differentiable via soft approximations of skeletonization, enabling optimization by standard neural network training algorithms. The loss and its architectural extensions have demonstrated state-of-the-art topological fidelity and reduced errors in structural segmentation tasks.
1. Mathematical Formulation of clDice Loss
Let denote the ground-truth binary segmentation mask, and the predicted soft mask produced by a neural network. Skeletons of the ground truth and prediction, denoted and , are obtained via an operator that is either morphological (e.g., thinning) or a differentiable approximation.
Topology precision and sensitivity are defined as
The clDice score is then the harmonic mean:
The corresponding loss used in training is
In the "soft-clDice" setting (Shit et al., 2020), skeletons are computed through differentiable soft morphological operations such as min-pooling (erosion), max-pooling (dilation), and their compositions to enable backpropagation.
2. Topological and Mathematical Guarantees
The theoretical basis of clDice ensures topology preservation up to homotopy equivalence, provided the conditions and are satisfied for prediction 0 and label 1. This topological guarantee extends to both 2D and 3D binary segmentations (Shit et al., 2020).
By focusing the metric on skeleton overlap, clDice equalizes influence across tubular branches of varying thickness, overcoming Dice's bias toward volumetrically larger (thicker) regions. The condition 2 is satisfied if and only if every predicted centerline voxel lies within the reference mask, and every reference centerline voxel is covered by the prediction, thus ensuring homotopy equivalence.
3. Differentiable Skeletonization: Algorithms and Implementation
Differentiable clDice's core lies in "soft skeletonization," which renders the discrete, non-differentiable skeleton operation compatible with gradient-based optimization. This is achieved by repeated application of differentiable soft-erosion (min-pooling) and soft-dilation (max-pooling) followed by elementwise ReLU to accumulate medial axis voxels:
7 Here, 3 is set to at least the maximum expected vessel radius. This soft skeletonization (and its derivatives) is fully compatible with frameworks such as PyTorch and TensorFlow.
More recently, accuracy improvements for 3D data involve learning the skeletonization via dedicated 3D U-Net architectures, where a secondary U-Net replaces soft morphology and outputs a soft skeleton 4 from the predicted mask and underlying image (Rougé et al., 2023). This alternative addresses the suboptimal thickness and disconnectivity observed in the classic soft-skeleton approach for 3D images.
4. Integration into Segmentation Networks
The cascaded multitask U-Net workflow leverages differentiable clDice within a two-stage process (Rougé et al., 2023):
- A primary 3D U-Net predicts a segmentation 5 from the imaging input.
- A secondary U-Net, trained to emulate skeletonization, refines or produces a soft skeleton 6 from 7 (optionally concatenated with the image patch).
The composite training loss combines segmentation Dice, skeleton Dice, and clDice loss: 8 Empirically, 9 yielded optimal results. The skeleton network is pretrained independently and then frozen during joint fine-tuning, ensuring that all gradients from clDice propagate to the segmentation network only. This framework efficiently maintains topological integrity while improving skeleton localization accuracy.
5. Empirical Performance, Efficiency, and Applications
Extensive experiments demonstrate that adding differentiable clDice loss consistently improves metrics reflecting structural correctness. In 3D time-of-flight MRA volume datasets, for example (Rougé et al., 2023):
| Method | DSC | clDice | Euler error | 0 error | 1 error |
|---|---|---|---|---|---|
| Baseline U-Net (Dice only) | 0.76±0.02 | 0.85±0.02 | 37.8±20.3 | 9.5±6.5 | 33.3±24.4 |
| U-Net + soft-clDice | — | — | 33.7±19.3 | 8.2±6.8 | 30.0±22.6 |
| Cascaded U-Net (learned skeleton) | — | — | 23.5±14.8 | 7.0±6.7 | 25.5±16.5 |
Inference, implemented via sliding-window and Gaussian weighting, uses standard deep learning techniques. Soft-clDice increases computational cost marginally, and the learned-skeleton approach remains competitive in runtime (2128–170 s/epoch depending on variant and hardware).
Domains of application include vessel segmentation and centerline extraction in medical imaging, neuron tracing, and road segmentation, where maintaining correct branch and loop structures is critical (Rougé et al., 2023, Shit et al., 2020, Shi et al., 2024).
6. Limitations and Generalizations
The soft-skeleton approach is limited by thick, noisy skeletons and decreased reliability in capturing small vessels or complex 3D topologies. The learned-skeleton U-Net significantly mitigates these issues for 3D images but introduces dependence on accurate supervised skeletons for training.
Recent generalizations, such as centerline-boundary Dice (cbDice) and cl-X-Dice, introduce additional geometric and scale-aware terms (e.g., boundary distances, vessel radii, inverse-radii) to modulate sensitivity to translations and vessel diameters (Shi et al., 2024). Variants both preserve differentiability and allow tuning of the loss to prioritize topology, geometry, or thickness according to application constraints.
7. Practical Recommendations
- Pretrain segmentation and skeleton networks independently for stability.
- Use 3 for loss balancing in multitask settings.
- Freeze the skeleton extraction network during joint fine-tuning.
- Choose skeletonization iterations 4 according to the maximal vessel radius present.
- Monitor not only Dice coefficients but also Euler characteristic, 5 and 6 errors to avoid overfitting to volumetric score at the expense of topology.
- The differentiable clDice loss integrates natively with most DL frameworks via custom or built-in min/max pooling operations and supports efficient end-to-end training across 2D and 3D segmentation tasks.
The differentiable clDice loss and its architectural/algorithmic extensions provide a robust framework for enforcing and optimizing topological fidelity in the segmentation of thin, connected structures, achieving marked improvements in both classical and learned approaches across a wide range of datasets (Rougé et al., 2023, Shi et al., 2024, Shit et al., 2020).