[Paper Review] A skeletonization algorithm for gradient-based optimization

Martin J. Menten, Johannes C. Paetzold+7 more|arXiv (Cornell University)|Sep 5, 2023

Medical Image Segmentation TechniquesComputer Science3 citations

TL;DR

This paper introduces the first 3D topology-preserving skeletonization algorithm compatible with gradient-based optimization, using differentiable operations like convolutions, matrix operations, and stochastic sampling. It achieves superior topological and geometric accuracy over non-differentiable and neural network-based baselines, enabling seamless integration into deep learning pipelines for medical image analysis tasks such as vessel segmentation and multimodal registration.

ABSTRACT

The skeleton of a digital image is a compact representation of its topology, geometry, and scale. It has utility in many computer vision applications, such as image description, segmentation, and registration. However, skeletonization has only seen limited use in contemporary deep learning solutions. Most existing skeletonization algorithms are not differentiable, making it impossible to integrate them with gradient-based optimization. Compatible algorithms based on morphological operations and neural networks have been proposed, but their results often deviate from the geometry and topology of the true medial axis. This work introduces the first three-dimensional skeletonization algorithm that is both compatible with gradient-based optimization and preserves an object's topology. Our method is exclusively based on matrix additions and multiplications, convolutional operations, basic non-linear functions, and sampling from a uniform probability distribution, allowing it to be easily implemented in any major deep learning library. In benchmarking experiments, we prove the advantages of our skeletonization algorithm compared to non-differentiable, morphological, and neural-network-based baselines. Finally, we demonstrate the utility of our algorithm by integrating it with two medical image processing applications that use gradient-based optimization: deep-learning-based blood vessel segmentation, and multimodal registration of the mandible in computed tomography and magnetic resonance images.

Motivation & Objective

• Address the lack of differentiable skeletonization methods compatible with gradient-based optimization in deep learning.
• Overcome the topological inaccuracies of morphological and neural network-based skeletonization methods, which often introduce breaks in the skeleton.
• Develop a method that preserves the true medial axis topology while being fully compatible with backpropagation and deep learning frameworks.
• Demonstrate utility in real-world medical image processing applications requiring optimization, such as blood vessel segmentation and multimodal image registration.
• Establish a foundation for integrating topological inductive biases into deep learning pipelines through a differentiable, domain-agnostic skeletonization approach.

Proposed method

• The algorithm uses a differentiable formulation based on matrix additions and multiplications, convolutional operations, and basic non-linear activation functions.
• It employs a parallelizable scheme for identifying and removing simple points in a topology-preserving manner, ensuring the medial axis remains geometrically and topologically accurate.
• A stochastic sampling strategy is introduced to discretize non-binary input images, enabling differentiability while maintaining fidelity to the original object structure.
• The method leverages Euler characteristic-based criteria to ensure topological correctness during skeletonization, avoiding spurious breaks common in morphological approaches.
• The entire pipeline is implemented using operations native to deep learning frameworks like PyTorch and TensorFlow, enabling end-to-end backpropagation.
• The algorithm supports both Euler characteristic-based and Boolean-based simplification rules, allowing flexibility in implementation and performance tuning.

Experimental results

Research questions

• RQ1 Can a 3D skeletonization algorithm be designed that is both differentiable and topology-preserving, enabling integration into gradient-based deep learning pipelines?
• RQ2 How does the proposed method compare to non-differentiable, morphological, and neural network-based skeletonization baselines in terms of topological and geometric accuracy?
• RQ3 To what extent can differentiable skeletonization improve performance in deep learning-based medical image segmentation and multimodal registration tasks?
• RQ4 Can the algorithm be implemented using only standard deep learning operations, ensuring broad compatibility with existing frameworks?
• RQ5 Does the use of skeleton-based loss functions in registration lead to improved alignment accuracy compared to standard mask-based losses?

Key findings

• The proposed algorithm achieves superior topological and geometric accuracy compared to non-differentiable, morphological, and neural network-based baselines in benchmarking experiments.
• In multimodal mandible registration, using the skeleton as a loss function improved the Dice Similarity Coefficient (DSC) to 0.37 ± 0.02 (Euler variant) and 0.37 ± 0.01 (Boolean variant), outperforming the baseline without skeletonization (0.38 ± 0.01) and morphological baseline (0.32 ± 0.01).
• The Hausdorff distance (HD) was reduced to 28.0 ± 1.1 mm with the proposed method (Boolean variant), compared to 29.9 ± 0.9 mm in the baseline, indicating improved spatial alignment.
• The average surface distance (ASD) remained comparable across methods (6.5–6.7 mm), confirming consistent surface-level accuracy while improving topological fidelity.
• The algorithm enables seamless integration into deep learning pipelines via autograd frameworks like PyTorch, demonstrating compatibility with gradient-based optimization in real-world applications.
• The method is domain-agnostic and does not require fine-tuning on specific data distributions, unlike learned skeletonization approaches susceptible to domain shift.