ODE-UNet: Continuous-Depth U-Net Architecture

• ODE-UNet is a neural network architecture that integrates neural ODEs to replace traditional U-Net convolutional blocks, offering adaptive receptive fields and parameter efficiency.
• It employs continuous-depth modeling via ODE blocks with adaptive Runge–Kutta integration, allowing the network to adjust its effective depth locally during segmentation tasks.
• Empirical evaluations on colon gland segmentation show ODE-UNet achieves superior Dice and F1 scores with fewer parameters and lower memory usage compared to conventional U-Net models.

ODE-UNet (often denoted "U-Node") is a neural network architecture that replaces the traditional discrete convolutional blocks of a U-Net with continuous-depth blocks parameterized by neural ordinary differential equations (NODEs). This approach enables dynamically adaptive receptive fields and significant parameter savings, as demonstrated in semantic segmentation tasks such as individual colon gland segmentation. ODE-UNet leverages continuous-time modeling, allowing the network to tune its depth locally during inference and training, without incurring the parameter cost of conventional deep architectures.

1. Architectural Formulation

The ODE-UNet framework begins with a standard 4-level U-Net. In classic U-Net, each encoder and decoder stage consists of a stack of two 3×3 convolutional layers, followed by nonlinearities. These stacks are replaced in ODE-UNet by ODE blocks—neural ODE solvers that model the transformation as integration of a parameterized continuous-time dynamics.

Encoder and Decoder Implementation:

• Standard U-Net (encoder at level $\ell$):

$$h_\ell = \text{Conv}_{3 \times 3}(\text{Conv}_{3 \times 3}(h_{\ell-1}))$$

$\tilde{h}_\ell = \text{Downsample}(h_\ell)$

• ODE-UNet (encoder at level $\ell$):

$h_\ell = \text{ODEBlock}(f_\theta, h_{\ell-1})$

$\tilde{h}_\ell = \text{Downsample}(h_\ell)$

• Standard U-Net (decoder at level $\ell$):

$g_\ell = \text{Upsample}(g_{\ell+1})$

$$g_\ell = \text{Conv}_{3 \times 3}(\text{Conv}_{3 \times 3}(\text{concat}(g_\ell, h_\ell)))$$

• ODE-UNet (decoder at level $\ell$):

$g_\ell = \text{Upsample}(g_{\ell+1})$

$$g_\ell = \text{ODEBlock}(f_\theta, \text{concat}(g_\ell, h_\ell))$$

The ODEBlock is defined by integrating a learned function $f_\theta$ from $t=0$ to $t=1$ over the hidden state $h(t)$:

def ODEBlock(f, h0): # Integrate dh/dt = f(h(t), θ) from t=0 to t=1 h1 = ODESolve(f, h0, t0=0, t1=1, method="dopri5", atol=1e-3, rtol=1e-3) return h1

2. Mathematical Foundations

ODE-UNet substitutes the discrete residual update $h_{t+1} = h_t + f(h_t, \theta)$ with a continuous initial value problem:

$$\frac{d h(t)}{dt} = f(h(t), \theta), \quad t \in [0,1], \quad h(0) = h_{\text{previous}}$$

Here, $h(t) \in \mathbb{R}^d$ represents the network activation at "depth" $t$, and $f$ is parameterized as two 3×3 convolutions interleaved with nonlinearities. The same convolutional parameters $\theta$ are reused across the multiple "virtual" layers generated by the ODE solver, effectively producing a continuum of transformations.

3. ODE Solver Configuration and Regularization

The ODEBlock in ODE-UNet employs the explicit 5th-order Runge–Kutta method (Dormand–Prince, "dopri5") with absolute and relative tolerances of $1 \times 10^{-3}$. The solver adaptively chooses the number and size of integration steps, expanding the effective receptive field in challenging regions of the input and contracting it in easier regions. This adaptivity is reflected in the number of function evaluations (NFE) per data patch, which serves as a proxy for the receptive field size.

No additional explicit regularization beyond standard Adam $L_2$ weight decay is introduced; regularization is largely an implicit result of parameter sharing (via repeated application of $f_\theta$) and the error constraint controlled by adaptive integration.

4. Quantitative Behavior and Resource Utilization

A comparative evaluation on the GlaS colon gland segmentation dataset provides the following results:

Method	Params	Obj. Dice (A, B)	F1 (A, B)	Hausdorff* (A, B)
U-Net	30 M	0.868 (0.884, 0.819)	0.841 (0.865, 0.768)	69.6 (55.6, 111)
U-ResNet†	2 M	0.757 (0.789, 0.660)	0.689 (0.743, 0.523)	122 (97.3, 199)
U-Node	2 M	0.881 (0.893, 0.842)	0.861 (0.882, 0.801)	59.5 (48.6, 92.4)

Lower Hausdorff distance is better. †"U-ResNet" is U-Node forced to use exactly one solver step per block.

Key implications:

• ODE-UNet achieves higher Dice (0.881 vs 0.868) and F1 scores than standard U-Net with only $\sim$1/15th as many parameters (2 M vs 30 M).
• Peak GPU memory during training is $\sim$6 GB for ODE-UNet (compared to $\sim$10 GB for U-Net), due to the use of the continuous-adjoint method for backpropagation, which avoids storing all intermediate activations.
• Training time per epoch is approximately 1.5$\times$ slower for ODE-UNet, attributed to repeated convolution evaluations by the adaptive solver.

5. Dynamic Receptive Field and Parameter Efficiency

The ODE solver's dynamic step selection allows receptive fields to grow automatically in spatial regions where the segmentation problem is harder. In these regions, more integration steps (i.e., more repeated applications of $f_\theta$) effectively enlarge the receptive field. Conversely, in simpler regions, fewer steps are taken, yielding computational savings.

Parameter efficiency arises from the reuse of $f_\theta$ across all virtual layers within a block, in contrast to standard U-Net, which allocates new parameters to each convolutional layer to increase receptive field.

6. Training Stability and Hyperparameter Selection

Training of ODE-UNet uses continuous-adjoint sensitivity analysis for gradient computation, allowing activation histories to be reconstructed on demand during backward passes. Stable optimization was observed with an Adam optimizer at a learning rate of $1 \times 10^{-3}$ and standard $L_2$ weight decay to prevent weight drift.

Solver tolerance (atol/rtol) directly controls the tradeoff between segmentation accuracy and computational efficiency. Excessively tight tolerances result in slow training, while loose tolerances degrade segmentation performance. In especially deep architectures or extremely tight tolerance regimes, the adjoint method may accumulate numerical errors, necessitating sensitivity to solver configuration.

7. Practical Implications and Limitations

ODE-UNet demonstrates that dynamically adaptive receptive fields can be leveraged to match or exceed state-of-the-art semantic segmentation with substantially reduced parameter count and memory footprint. This adaptivity is particularly beneficial for segmenting large or morphologically complex structures without the need for bespoke architecture tuning.

However, the increased computational cost due to multiple function evaluations per ODE block leads to slower training and inference. Additionally, solver tolerance emerges as a novel hyperparameter requiring empirical optimization. In deep networks or under stringent tolerances, numerical stability of backpropagation may be compromised, making systematic evaluation of solver settings essential for robust deployment.