Differentiable clDice: Topology-Aware Segmentation Loss
- Differentiable clDice is a loss function that preserves topology by comparing soft skeletons of predicted and ground truth masks using differentiable morphological operations.
- It integrates with deep learning models through iterative soft-skeletonization and efficient gradient computation, ensuring connectivity in segmentation tasks.
- Empirical evaluations show improved connectivity metrics in tubular structure segmentation, while extensions like cbDice address geometric sensitivity limitations.
Differentiable clDice, also known as soft-clDice, is a topology-preserving loss function specifically designed for training neural networks to segment tubular or network-like structures, such as blood vessels, neurons, or roads, in 2D and 3D images. Unlike conventional overlap-based losses, differentiable clDice penalizes topological errors by comparing skeletonized versions of the predicted and ground truth masks, ensuring preservation of connectivity and homotopy equivalence. The formulation employs differentiable morphological operators that allow integration into end-to-end deep learning pipelines, providing efficient gradient computation for training (Shit et al., 2020).
1. Mathematical Formulation and Skeletonization
The central construct in differentiable clDice is the comparison of soft skeletons extracted from the prediction and the ground truth %%%%1%%%%, with denoting the spatial domain (image grid or volume). The soft-clDice metric is based on two terms: topology precision () and topology sensitivity (),
- ,
- ,
where denotes the differentiable skeletonization operator. The differentiable clDice loss is then given by
$\mathrm{soft\mathchar`-clDice}(P, G) = \frac{2 \cdot \mathrm{Tprec} \cdot \mathrm{Tsens}}{\mathrm{Tprec} + \mathrm{Tsens}}$
The soft skeletonization operator is implemented using iterative min- and max-pooling operations simulating morphological erosion and dilation, followed by open operations and ReLU thresholding. This operation preserves differentiability at every stage, making the entire loss function suitable for gradient-based optimization (Shit et al., 2020).
2. Differentiability and Backpropagation
Differentiable clDice is constructed using sums, products, and rational functions of the network's output, and all skeletonization steps are composed of neighborhood-wise pooling and point-wise functions. The gradients of the loss with respect to network predictions can be analytically written using the chain rule. For loss $\mathcal{L}_c = 1 - \mathrm{soft\mathchar`-clDice}$:
- For topology precision and sensitivity ,
Downstream differentiation with respect to or the output of skeletonization is recursively obtained by chaining these partials with those for the skeletonization and linear mask terms, all of which are differentiable (Shit et al., 2020). This allows implementation in modern deep learning frameworks via automatic differentiation.
3. Theoretical Guarantees and Complexity
In the limiting case where predictions and ground truth are binary (), it is shown that
where denotes the morphological skeleton. Under mild topological conditions (foreground and background both being homotopy-equivalent to 1D graphs), this criterion guarantees homotopy equivalence between and , i.e., their underlying graphs are topologically identical up to continuous deformation (Shit et al., 2020).
The operational complexity of soft-skeletonization is , where is the maximum number of skeletonization iterations (typically matching the maximum tubular radius in pixels or voxels), and is the number of voxels. This is computationally efficient: typically –$25$, yielding overhead relative to a standard Dice loss. By contrast, persistent homology–based losses are (Shit et al., 2020).
4. Implementation Details and Training Integration
The soft-clDice loss function is realized as follows:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 |
def soft_skeleton(I, K): S = 0 cur = I for j in range(K+1): opened = soft_dilate(soft_erode(cur)) delta = relu(cur - opened) S = S + (1 - S) * delta cur = soft_erode(cur) return S def soft_cldice(P, G, K, eps): S_P = soft_skeleton(P, K) S_G = soft_skeleton(G, K) A = sum(S_P * G) B = sum(S_P) C = sum(S_G * P) D = sum(S_G) a = (A+eps)/(B+eps) b = (C+eps)/(D+eps) return 2*a*b/(a+b+1e-12) def loss_fn(P, G, alpha): dsc = 2*sum(P*G)/(sum(P)+sum(G)+eps) cl = soft_cldice(P,G,K=20,eps=1e-6) return (1-alpha)*(1-dsc) + alpha*(1-cl) |
All operations are vectorized, leveraging GPU acceleration, and require hyperparameter tuning:
- (skeleton iterations): maximum tube radius in pixels
- (blend between Dice and clDice):
- (numerical stability): typically to
Combining soft-clDice with volumetric losses (e.g., Dice) stabilizes training. Pure soft-clDice loss () can result in slower convergence, especially under high prediction noise (Shit et al., 2020).
5. Comparative Performance and Limitations
Empirical evaluations on five public benchmarks covering 2D and 3D vascular, road, and neuron segmentation demonstrate that soft-clDice:
- Improves connectivity-related metrics, e.g., StreetMoverDistance (15–20%), Opt-Junction F1 (5–10%)
- Reduces Betti numbers errors (, drop by 20–50%)
- Increases topology-aware clDice by 2–6 points absolute and volumetric Dice by 1–2 points
- Recovers small network branches and avoids introducing spurious bridges
The approach notably produces segmentations with more accurate network topology and fewer isolated or missing branches (Shit et al., 2020). However, clDice, by design, is insensitive to boundary shifts within the vessel radius and treats all skeleton points equally, which introduces sensitivity to skeleton parallel translations and may confer a bias toward larger-diameter structures. This leads to limitations in precise geometric quantification, such as diameter imbalance or insensitivity to subtle boundary errors—motivating recent extensions that introduce boundary and diameter weighting (Shi et al., 2024).
6. Extensions and Relation to Other Metrics
The centerline-boundary Dice (cbDice) and the more general cl-X-Dice family were developed to address the geometric and diameter sensitivity shortcomings of clDice. These introduce distance or inverse-radius weighting to the skeleton and mask overlaps, thereby extending clDice to be translation- and diameter-aware (Shi et al., 2024). Such modifications retain the core topological guarantees of clDice while adding sensitivity to geometric detail and uniform treatment over varying vessel radii.
A comparison of related loss functions appears below:
| Loss Function | Topology Preservation | Boundary/Geometric Sensitivity | Diameter Balance |
|---|---|---|---|
| Dice | No | Yes | No (area-biased) |
| clDice (soft-clDice) | Yes | No (insensitive within radius) | No (skeleton-biased) |
| cbDice/cl-X-Dice | Yes | Yes | Yes |
Soft-clDice thus occupies a central position within topology-aware segmentation losses, balancing computational efficiency with provable topological correctness, and forms the basis for more finely resolved boundary- and structure-aware losses (Shit et al., 2020, Shi et al., 2024).