[Paper Review] Globally Optimal Segmentation of Mutually Interacting Surfaces using Deep Learning

Hui Xie, Zhe Pan+6 more|arXiv (Cornell University)|Jul 2, 2020

Anomaly Detection Techniques and Applications24 references4 citations

TL;DR

This paper proposes a deep learning framework that unifies graph-based optimal surface segmentation with end-to-end differentiable optimization using the primal-dual Interior Point Method (IPM). By parameterizing surface cost functions and learning them via a U-Net backbone, the method enables globally optimal segmentation of multiple interacting surfaces in OCT and IVUS images, achieving state-of-the-art accuracy with explicit modeling of surface interactions and efficient backpropagation through the IPM layer.

ABSTRACT

Segmentation of multiple surfaces in medical images is a challenging problem, further complicated by the frequent presence of weak boundary and mutual influence between adjacent objects. The traditional graph-based optimal surface segmentation method has proven its effectiveness with its ability of capturing various surface priors in a uniform graph model. However, its efficacy heavily relies on handcrafted features that are used to define the surface cost for the "goodness" of a surface. Recently, deep learning (DL) is emerging as powerful tools for medical image segmentation thanks to its superior feature learning capability. Unfortunately, due to the scarcity of training data in medical imaging, it is nontrivial for DL networks to implicitly learn the global structure of the target surfaces, including surface interactions. In this work, we propose to parameterize the surface cost functions in the graph model and leverage DL to learn those parameters. The multiple optimal surfaces are then simultaneously detected by minimizing the total surface cost while explicitly enforcing the mutual surface interaction constraints. The optimization problem is solved by the primal-dual Internal Point Method, which can be implemented by a layer of neural networks, enabling efficient end-to-end training of the whole network. Experiments on Spectral Domain Optical Coherence Tomography (SD-OCT) retinal layer segmentation and Intravascular Ultrasound (IVUS) vessel wall segmentation demonstrated very promising results. All source code is public to facilitate further research at this direction.

Motivation & Objective

• To address the challenge of segmenting multiple interacting surfaces in medical images where weak boundaries and mutual dependencies reduce accuracy.
• To overcome the limitations of traditional graph-based methods that rely on handcrafted features for surface cost modeling.
• To integrate deep learning’s feature representation capability with the global optimality of convex optimization for multi-surface segmentation.
• To enable end-to-end training by making the constrained optimization process differentiable via a neural network layer implementation of the IPM.
• To demonstrate the method’s effectiveness on real-world medical imaging tasks—retinal layer and vessel wall segmentation—using SD-OCT and IVUS data.

Proposed method

• Parameterize the surface cost function in a graph-based segmentation model using a deep neural network (U-Net) to learn image-dependent cost parameters.
• Formulate the multi-surface segmentation as a constrained convex optimization problem that enforces surface interaction constraints and global optimality.
• Implement the primal-dual Interior Point Method (IPM) as a differentiable layer to solve the optimization problem, enabling backpropagation through the optimization process.
• Derive closed-form gradients for the IPM layer with respect to network weights, input features, and constraints using KKT conditions and matrix calculus.
• Use PyTorch to implement the IPM layer with GPU-accelerated batch processing and pseudo-inverse handling for singular matrices.
• Integrate the IPM layer into a deep learning pipeline, enabling joint training of the U-Net and optimization components with full backpropagation.

Experimental results

Research questions

• RQ1 Can deep learning be effectively combined with convex optimization to achieve globally optimal segmentation of multiple interacting surfaces in medical images?
• RQ2 How can the surface cost function in a graph-based segmentation model be parameterized and learned end-to-end using deep neural networks?
• RQ3 To what extent can differentiable optimization layers like IPM improve the robustness and accuracy of multi-surface segmentation in the presence of weak boundaries?
• RQ4 Can the proposed method generalize across different medical imaging modalities such as SD-OCT and IVUS?
• RQ5 How does explicit modeling of surface interactions through constraints improve segmentation performance compared to independent surface detection?

Key findings

• The proposed method achieved state-of-the-art performance on retinal layer segmentation in SD-OCT images, significantly outperforming baseline U-Net and FCN-based models.
• The method demonstrated robustness in segmenting vessel walls in IVUS images, even in regions with poor image contrast and weak boundaries.
• The integration of the IPM layer enabled end-to-end training with stable convergence, requiring at most 7 IPM iterations per forward pass.
• Quantitative evaluation showed improved Dice scores and reduced boundary errors compared to non-optimized and non-interaction-aware baselines.
• The ablation study confirmed that modeling surface interactions significantly enhanced segmentation accuracy, especially in ambiguous or low-contrast regions.
• The open-source release of all code and data at [26] facilitates reproducibility and further research in differentiable optimization for medical image segmentation.