PDE-Constrained Optimization for Neural Image Segmentation with Physics Priors

TL;DR

The paper tackles ill-posed microscopy image segmentation by introducing PDE-constrained optimization that injects physics priors into deep learning. It implements a Reaction–Diffusion prior and a Phase-Field interface energy as differentiable residual losses, integrated with a UNet backbone in a two-stage training regime. Evaluated on LIVECell across in-distribution and out-of-distribution morphologies, the approach yields notably improved boundary fidelity and stability in low-data regimes, while maintaining strong segmentation quality in abundant-data settings. This work demonstrates that physics-informed deep learning can act as a principled bridge between variational methods and neural networks, with potential extensions to higher-order PDEs and multi-phase interfaces for more complex biomedical imaging tasks.

Abstract

Segmentation of microscopy images constitutes an ill-posed inverse problem due to measurement noise, weak object boundaries, and limited labeled data. Although deep neural networks provide flexible nonparametric estimators, unconstrained empirical risk minimization often leads to unstable solutions and poor generalization. In this work, image segmentation is formulated as a PDE-constrained optimization problem that integrates physically motivated priors into deep learning models through variational regularization. The proposed framework minimizes a composite objective function consisting of a data fidelity term and penalty terms derived from reaction-diffusion equations and phase-field interface energies, all implemented as differentiable residual losses. Experiments are conducted on the LIVECell dataset, a high-quality, manually annotated collection of phase-contrast microscopy images. Training is performed on two cell types, while evaluation is carried out on a distinct, unseen cell type to assess generalization. A UNet architecture is used as the unconstrained baseline model. Experimental results demonstrate consistent improvements in segmentation accuracy and boundary fidelity compared to unconstrained deep learning baselines. Moreover, the PDE-regularized models exhibit enhanced stability and improved generalization in low-sample regimes, highlighting the advantages of incorporating structured priors. The proposed approach illustrates how PDE-constrained optimization can strengthen data-driven learning frameworks, providing a principled bridge between variational methods, statistical learning, and scientific machine learning.

Figures (7)

• Figure 1: UNet architecture with the PDE-constrained optimization framework. Continuous outputs are evaluated against soft PDE residuals and interface energies alongside data-driven fidelity terms.
• Figure 2: Resilience of PDE-constrained learning in data-sparse environments. The plots illustrate the relative improvement in (a) Boundary F1 Score, (b) Dice Score, and (c) IoU across varying training data fractions. Note the substantial performance leap at the 10% fraction, where physical residuals provide a critical structural blueprint that compensates for the scarcity of supervised signals, particularly for out-of-distribution (OOD) morphologies.
• Figure 3: Influence of PDE constraints at full data capacity (100% fraction). The bar plots quantify the refinement effect of individual and combined PDE priors (Reaction-Diffusion and Phase-Field) across three metrics: (a) Boundary F1 Score, (b) Dice Score, and (c) IoU Score. While the unconstrained baseline performs strongly at this data volume, the physical priors yield consistent improvements, particularly for out-of-distribution (OOD) samples.
• Figure 4: Influence of PDE constraints in the low-sample regime (10% data fraction). The bar plots quantify the refinement effect of individual and combined PDE priors (Reaction-Diffusion and Phase-Field) across three metrics: (a) Boundary F1 Score, (b) Dice Score, and (c) IoU Score. In this data-sparse environment, physical priors provide a critical structural foundation, yielding massive relative gains over the unconstrained baseline, particularly for out-of-distribution (OOD) morphologies.
• Figure 5: Reaction Threshold Sensitivity Analysis (10% Data). Line plots illustrate the sensitivity of segmentation performance to the reaction threshold parameter ($a$) within the reaction-diffusion equation. The model incorporates both Reaction-Diffusion and Phase-Field PDE constraints. Performance is evaluated across (a) Boundary F1 Score, (b) Dice Score, and (c) IoU Score for in-distribution (green) and out-of-distribution (red) morphologies. The vertical dashed line indicates the symmetric threshold ($a = 0.5$), which maintains optimal performance across the primary overlap metrics.