[Paper Review] Continuous U-Net: Faster, Greater and Noiseless

Chun-Wun Cheng, Christina Runkel+4 more|arXiv (Cornell University)|Feb 1, 2023

Radiomics and Machine Learning in Medical ImagingMedicine3 citations

TL;DR

This paper introduces Continuous U-Net, a novel continuous deep learning architecture for medical image segmentation that models network dynamics via second-order ordinary differential equations (ODEs). By replacing discrete U-Net blocks with continuous dynamic blocks, the method achieves faster convergence, greater robustness to noise, and constant memory cost via adjoint sensitivity methods, outperforming existing U-Net variants across six benchmark datasets.

ABSTRACT

Image segmentation is a fundamental task in image analysis and clinical practice. The current state-of-the-art techniques are based on U-shape type encoder-decoder networks with skip connections, called U-Net. Despite the powerful performance reported by existing U-Net type networks, they suffer from several major limitations. Issues include the hard coding of the receptive field size, compromising the performance and computational cost, as well as the fact that they do not account for inherent noise in the data. They have problems associated with discrete layers, and do not offer any theoretical underpinning. In this work we introduce continuous U-Net, a novel family of networks for image segmentation. Firstly, continuous U-Net is a continuous deep neural network that introduces new dynamic blocks modelled by second order ordinary differential equations. Secondly, we provide theoretical guarantees for our network demonstrating faster convergence, higher robustness and less sensitivity to noise. Thirdly, we derive qualitative measures to tailor-made segmentation tasks. We demonstrate, through extensive numerical and visual results, that our model outperforms existing U-Net blocks for several medical image segmentation benchmarking datasets.

Motivation & Objective

• To address the limitations of discrete U-Net architectures, including fixed receptive fields, high computational cost, and sensitivity to noise.
• To develop a theoretically grounded, continuous alternative to discrete U-Net blocks using second-order ODEs.
• To achieve faster convergence and constant memory cost through adjoint sensitivity methods in a continuous framework.
• To improve robustness against data noise and enhance generalization across diverse medical imaging datasets.
• To establish a foundation for future continuous U-Net variants with advanced components like attention or transformers.

Proposed method

• The network is structured as a continuous U-Net, where encoder and decoder blocks are replaced by dynamic blocks governed by second-order ODEs.
• The dynamic blocks are modeled using second-order neural ODEs, enabling smoother and more flexible feature learning trajectories.
• The adjoint sensitivity method is applied to achieve O(1) memory cost, regardless of model depth, by reversing the ODE integration.
• Theoretical analysis proves faster convergence and reduced sensitivity to noise compared to standard U-Net blocks.
• Different ODE solvers (Euler, AB, ABM, RK4) are evaluated to determine optimal integration methods for segmentation tasks.
• Qualitative measures are derived to tailor the network to specific segmentation tasks, improving performance and stability.

Experimental results

Research questions

• RQ1 Can a continuous U-Net architecture with second-order ODEs outperform standard discrete U-Net blocks in medical image segmentation?
• RQ2 Does modeling U-Net dynamics via continuous ODEs lead to faster convergence and lower memory consumption?
• RQ3 How does the proposed method perform under data noise compared to standard U-Net variants?
• RQ4 Can the continuous framework provide theoretical guarantees on robustness and convergence?
• RQ5 Can continuous U-Net achieve state-of-the-art performance without additional mechanisms like attention or transformers?

Key findings

• Continuous U-Net converges significantly faster than discrete U-Net blocks, requiring fewer iterations to reach optimal performance.
• The model achieves O(1) memory cost via adjoint sensitivity, making it computationally efficient regardless of depth.
• On all six benchmark datasets, Continuous U-Net outperforms standard U-Net blocks and even surpasses models with attention mechanisms.
• Under additive Gaussian noise (σ = 0.2), Continuous U-Net maintains a Dice score within 0.1209 of its clean performance, while the U-Net with inception blocks drops from 0.7434 to 0.1930 on the GlaS dataset.
• The fourth-order Runge-Kutta (RK4) solver outperforms Euler, Adams-Bashforth, and Adams-Bashforth-Moulton methods across all datasets and metrics.
• Continuous U-Net achieves stable and high performance across all datasets and metrics, unlike the FO block which shows inconsistent results, especially on the STARE dataset.