AI Research Portal

This paper introduces CNM-UNet, a novel Continuous Neural Memory ODE UNet for medical image segmentation. It addresses the trade-offs between model compactness, computational accuracy, and efficiency in existing ODE-based segmentation methods by replacing the hierarchical decoder layers of a vanilla UNet with a single Continuous Neural Memory ODEs Block (CNM-Block). This design reduces computation costs and improves training efficiency by leveraging continuous temporal feature extraction. To enhance cross-domain generalization, it incorporates a Dual SElf-updated (DUSE) strategy based on test-time adaptation principles. Experimental results demonstrate CNM-UNet's advantages in computational capacity, convergence speed, and cross-domain adaptability.

The paper addresses the following limitations in existing ODE-based medical image segmentation methods:

Trade-offs in discretization: Current methods often rely on discretization, which introduces local truncation errors and can lead to substantial global deviations and computational instability.

Proposed CNM-UNet, which replaces all hierarchical decoder layers in a vanilla UNet with a single Continuous Neural Memory ODEs Block (CNM-Block) for continuous temporal feature extraction.
Developed a Dual SElf-updated (DUSE) strategy based on test-time adaptation principles to enhance cross-domain generalization.
Demonstrated comprehensive advantages of CNM-UNet in computational capacity, convergence speed, and cross-domain adaptability.

The CNM-UNet architecture consists of three main components: an encoder, a feature fusion module, and a decoder (CNM-Block).

Encoder: This module shares the same architecture as the vanilla UNet, comprising double convolution blocks with Batch Normalization (BN) and ReLU activation, followed by max-pooling layers.
Feature Fusion Module: This module implements a multi-scale feature fusion strategy. Outputs from encoder stages are weighted and fused, then upsampled and transformed via 3x3 convolutions. A weighted summation aggregates these features before being passed to the decoder.
Decoder (CNM-Block): This block replaces the multiple discrete ODE decoders found in methods like DiODE-UNet. It employs a continuous ODE solver to model data distribution characteristics. The core of the CNM-Block is based on Neural Memory ODEs (nmODEs), where the global learnable parameter 'y' is used as the actual network input and acts as a global memory parameter. This design aims to achieve rapid convergence and high segmentation accuracy by directly coupling loss function dynamics with 'y' optimization. The system's dynamics are governed by the equation f = -y + sin²(y + C(y)), where C(·) is a convolution operation. The continuous ODE solver approximates the evolution using the Dormand-Prince-Shampine 5th-order Runge-Kutta (dopri5) method. Backpropagation is handled via an augmented adjoint method.

The DUSE strategy is introduced to improve generalization. It involves:

An adaptive parameter update strategy for the CNM-Block based on test data distribution characteristics. This includes dynamically adjusting the parameter 'y' using a gradient statistics-based adaptation mechanism and generating a low-frequency prompt 'P₁' for adapted image generation.
A prompt-learning-driven Batch Normalization (BN) statistic calibration method. This enhances domain alignment by developing a dynamically weighted BN loss function and applying a differentiated weight allocation strategy to BN layers.

Architecture:
- Encoder: Standard UNet encoder.
- Feature Fusion: Multi-scale feature fusion with weighted fusion, upsampling, and 3x3 convolutions.
- Decoder: Single Continuous Neural Memory ODEs Block (CNM-Block) leveraging nmODEs.
Key Equations:
- ODE dynamics: f = -y + sin²(y + C(y))
- Continuous evolution approximation: y(tn) = y(to) + ∫[to,tn] f(y(t), t)dt
- Absolute-distance-based loss: Ln = Σ(1/2 * (μf - μf)² + 1/2 * (σf - σf)²)
- Dynamically weighted BN loss: Lp = Σ (μ - μ + [ – σ)
- Global statistics update: Mf₁ = (1 - m) Mf + mp, Vf₁ = (1 - m)Vf₂ + mofi
- Adapted image generation: î = F⁻¹ ([OnePad(P₁) FA(xi), FP (xi)])
Loss/Objective: The primary objective is segmentation. The method uses an absolute-distance-based loss (Ln) to adjust convolutional layer weights governing 'y' and a dynamically weighted BN loss (Lp) for adaptive BN statistic calibration.
Optimization:
- Optimizer: AdamW optimizer.
- Scheduler: CosineAnnealingLR scheduler.
- Training epochs: 300.
- Batch size: 8.
- Continuous ODE solver timestep (Δt): 0.05.
- Momentum coefficient (m): 0.2.
- Hbase: 0.5.
- Scaling factor (k) for CNM-Block BN: 500 for OD/OC dataset, 1000 for Polyp dataset.
- Task-specific initial learning rates: 0.005 (OD/OC), 0.01 (Polyp).
Augmentation: Horizontal flipping, vertical flipping, and random rotation.
Data Handling: Images are normalized and resized to specific resolutions for each dataset (ISIC2018: 256x256, OD/OC: 512x512, Polyp: 352x352).
Implementation Details: Experiments conducted on RTX 4090 GPU using PyTorch. Results reported as mean and standard deviation over five runs. DUSE evaluation uses single-source training with multi-target testing. Prompt size α = 0.01.

Datasets: ISIC2018, OD/OC Dataset, Polyp Dataset (combining CVC-ClinicDB, ETIS-LaribPolypDB, Kvasir-Seg, BKAI-IGH-NeoPolyp).
Metrics: Dice Similarity Coefficient (DSC), Mean Intersection over Union (mIoU).
Baselines: Vanilla UNet, Continuous UNet, DiODE-UNet, ResUNet, Att-UNet, UNet (LMD), nmS4M-UNet, nmSSM-UNet, MALUNet, EGE-UNet, Ulite. For TTA comparison: Source Only, AdaBN, Tent, COTTA, EATA, SAR, VPTTA.
Numeric Results:
- CNM-UNet achieved the best DSC on Polyp-A (0.879±0.001), Polyp-C (0.814±0.006), and Polyp-D (0.861±0.004) datasets.
- It ranked second on ISIC2018 (0.885±0.001) behind Continuous UNet.
- CNM-UNet reduced parameters by 3.95M and GFLOPs by 19.20-24.54 compared to vanilla UNet on ISIC2018 and Polyp datasets.
- Compared to Continuous UNet, CNM-UNet showed a 3.94M parameter reduction and a 197.18-371.62 GFLOP decrease.
- DUSE strategy improved DSC over "Source Only" baseline by 0.374 on OD/OC dataset (set A) and 0.214 (set D), and by 0.064 (set A) and 0.084 (set B) on the Polyp dataset.
- DUSE demonstrated robustness to image corruptions (Brightness, Contrast, Defocus Blur, Gaussian Noise, Shot Noise, Saturate) on ISIC18, outperforming "Source Only" baseline.
Runtime/Training Cost: CNM-UNet has significantly lower GFLOPs and training time compared to Continuous UNet. Table 1 shows CNM-UNet training time as 1.80 hours, compared to 14.90 hours for Continuous UNet.

No explicit paper-reported limitations were found.

Method evidence: "CNM-UNet leverages ODEs' dynamic properties to establish continuous temporal feature extraction." (Abstract)
Result evidence: "CNM-UNet achieves the best results on the Polyp-A, Polyp-C, and Polyp-D datasets and ranks second on the ISIC2018 dataset." (Page 6)
Limitation evidence: No explicit paper-reported limitations were found.

CNM-UNet: Continuous Ordinary Differential Equations for Medical Image Segmentation

Summary

Problem

Key Contributions

Method

Technical Details

Results

Weaknesses / Limitations

Evidence

Additional Figures