[Paper Review] Neural Ordinary Differential Equations for Semantic Segmentation of Individual Colon Glands

Hans Pinckaers, Geert Litjens|arXiv (Cornell University)|Oct 23, 2019

Radiomics and Machine Learning in Medical Imaging20 citations

TL;DR

This paper proposes U-Node, a neural ordinary differential equation (NODE)-enhanced U-Net for semantic segmentation of colon glands in histopathology images. By replacing standard residual blocks with continuous-depth NODEs, the method achieves improved segmentation accuracy with 93% fewer parameters (2M vs. 30M) and reduced memory usage, while enabling adaptive receptive fields through variable function evaluations during inference and training on the GlaS dataset.

ABSTRACT

Automated medical image segmentation plays a key role in quantitative research and diagnostics. Convolutional neural networks based on the U-Net architecture are the state-of-the-art. A key disadvantage is the hard-coding of the receptive field size, which requires architecture optimization for each segmentation task. Furthermore, increasing the receptive field results in an increasing number of weights. Recently, Neural Ordinary Differential Equations (NODE) have been proposed, a new type of continuous depth deep neural network. This framework allows for a dynamic receptive field at a fixed memory cost and a smaller amount of parameters. We show on a colon gland segmentation dataset (GlaS) that these NODEs can be used within the U-Net framework to improve segmentation results while reducing memory load and parameter counts.

Motivation & Objective

• To address the limitations of fixed receptive fields in standard U-Net architectures, which require manual architecture tuning and increase parameter counts with depth.
• To explore whether neural ordinary differential equations (NODEs) can enable memory-efficient, adaptive receptive fields in medical image segmentation.
• To evaluate the performance of NODE-based networks on the GlaS colon gland segmentation benchmark, comparing against standard U-Net and residual baselines.
• To demonstrate that NODEs can reduce model size and memory usage while improving segmentation accuracy through dynamic, context-aware feature learning.

Proposed method

• Integrates neural ODEs into the U-Net encoder and decoder paths, replacing standard residual blocks with continuous-depth blocks defined by an ODE system.
• Uses the adjoint method for efficient backpropagation through the ODE solver, avoiding storage of intermediate activations and reducing memory consumption.
• Employs the fifth-order Dormand-Prince (dopri5) ODE solver with a tolerance of 1e-3 to balance accuracy and computational cost.
• Applies variable number of function evaluations (NFE) per image and layer, enabling adaptive receptive fields that adjust based on input complexity.
• Trains the model using Adam optimizer with data augmentation (translation, flipping, rotation, elastic distortion, color jitter) and cross-entropy loss.
• Uses test-time augmentation with flipping and averaging to improve robustness and generalization.

Experimental results

Research questions

• RQ1 Can neural ODEs improve semantic segmentation performance in medical imaging while reducing model size and memory usage?
• RQ2 Does the adaptive receptive field of NODEs lead to better generalization and segmentation accuracy on complex histopathology images?
• RQ3 How does the performance of a NODE-enhanced U-Net compare to standard U-Net and residual U-Net variants in colon gland segmentation?
• RQ4 Can the continuous-depth architecture of NODEs maintain high accuracy with significantly fewer parameters than standard U-Net architectures?

Key findings

• The U-Node model achieved an object Dice score of 0.881 on the full GlaS test set, outperforming the standard U-Net (0.868) and U-ResNet (0.757).
• U-Node achieved a Hausdorff distance of 59.5, significantly lower than U-Net (69.6) and U-ResNet (122), indicating better shape similarity.
• Despite using only 2 million parameters (93% fewer than U-Net’s 30 million), U-Node matched or exceeded the performance of the original U-Net on all metrics.
• The model demonstrated improved segmentation quality with less noise and better boundary delineation, especially on larger glands, as shown in qualitative results.
• The adaptive receptive field, enabled by variable function evaluations, allowed the network to dynamically adjust context usage, particularly benefiting from higher NFE on difficult patches.
• Training was slower than standard U-Net due to repeated ODE solver calls, but the method enables trade-offs between speed and accuracy at inference via NFE control.