Parameter conditioned interpretable U-Net surrogate model for data-driven predictions of convection-diffusion-reaction processes

Technical Summary: Parameter Conditioned Interpretable U-Net Surrogate Model for Convection-Diffusion-Reaction Processes

Problem Statement
The paper addresses the computational bottleneck associated with repeatedly solving nonlinear, instationary convection-diffusion-reaction partial differential equations (PDEs) for parameter identification and sensitivity analysis. Motivated by the macroscopic modeling of cancer cell phenotypic plasticity (specifically the ability of cells to change phenotype in response to environmental stress), the authors focus on a two-dimensional PDE model. While high-fidelity numerical solvers (implemented in C++) provide accurate results, their repeated execution for tasks such as inverse problems or large-scale parameter sweeps is computationally prohibitive. The challenge lies in developing a data-driven surrogate that can efficiently approximate the solution map $(U_0, c) \mapsto U_T$ across a family of PDE instances defined by varying coefficient functions (diffusion, advection, and reaction), while maintaining accuracy and enabling gradient-based inference.

Methodology
The authors propose a combined numerical and data-driven workflow consisting of three main components:

1. High-Fidelity Numerical Solver:

   • A C++ solver was developed using second-order finite-difference methods for spatial discretization and a step-size controlled third-order Runge-Kutta-Fehlberg (RKF) integrator for temporal discretization.
   • The solver handles Neumann boundary conditions using second-order forward/backward stencils.
   • A mesh convergence study confirmed second-order spatial convergence. The solver generates a dataset of 10,000 training pairs and 2,500 test pairs on a $256 \times 256$ grid (downsampled to $64 \times 64$ for training).
   • Initial conditions are randomized sums of "hill-like" functions, and model parameters (coefficients for diffusion matrix $D$, velocity $v$, and reaction term $f$) are sampled from specific distributions to create diverse PDE regimes.

2. Surrogate Model Architecture (Conditioned U-Net):

   • The core of the approach is a U-Net architecture adapted for parameter conditioning.
   • Coordinate Encoding: To address the translation-invariance bias of standard convolutions, the input is augmented with normalized spatial coordinate channels (CoordConv), allowing the network to learn position-dependent features essential for boundary conditions.
   • Parameter Conditioning: The physical parameter vector $c$ (containing 4 randomized coefficients) is embedded via a Multi-Layer Perceptron (MLP). This embedding modulates intermediate feature maps using Feature-wise Linear Modulation (FiLM). FiLM applies affine transformations ($\gamma$ and $\beta$) to feature channels, enabling a single network to adapt to different PDE regimes (diffusion strength, advection speed, reaction rates) without retraining.
   • Residual Blocks: The base building blocks are residual blocks incorporating FiLM, Group Normalization, LeakyReLU activation, and dropout. This structure facilitates multiscale representation learning and training stability.
   • Attention Mechanism: A bottleneck layer includes attention mechanisms to capture long-range dependencies.

3. Training Strategy:

   • The model is trained to minimize the Huber loss, which balances the smoothness of $\ell_2$ loss for small residuals with the robustness of $\ell_1$ loss for outliers.
   • Optimization is performed using the AdamW optimizer with a Cosine Annealing learning rate schedule and gradient clipping.

Key Results

• Accuracy and Convergence: The numerical solver demonstrated second-order convergence. The trained U-Net surrogate achieved low prediction errors on held-out test data, with a final validation loss of $1.75 \times 10^{-4}$. Visual comparisons show the surrogate accurately reproduces spatial localization and transport-induced deformation patterns.
• Efficiency: The surrogate inference time is approximately 11ms per sample on an NVIDIA Tesla V100 GPU, independent of batch size. This represents a significant speedup compared to the C++ solver, which requires seconds to minutes depending on resolution.
• Generalization Analysis: A factorial test dataset (separating initial conditions from parameter vectors) was used to analyze generalization.
  • The error variability was found to be significantly higher across different conditioning vectors (parameter regimes) than across different initial conditions.
  • This indicates that the difficulty of approximation is governed primarily by the induced PDE regime (e.g., specific combinations of diffusion and advection) rather than the initial density distribution.
  • Correlation analysis revealed only weak-to-moderate monotone relationships between individual coefficients and the error, suggesting that approximation difficulty arises from complex interactions between transport, diffusion, and reaction effects rather than a single parameter.

Significance and Claims
The paper claims that the proposed parameter-conditioned U-Net with FiLM and coordinate encoding offers a favorable trade-off between expressivity, computational efficiency, and implementation simplicity for fixed-grid, fixed-horizon PDE surrogate modeling.

• One-Network-for-Many-Regimes: Unlike training separate models for each parameter setting, this approach learns a unified solution operator family capable of generalizing across a prescribed range of coefficient functions.
• Enabling Inverse Problems: Because the surrogate is differentiable with respect to both inputs ($U_0$) and parameters ($c$), it facilitates gradient-based parameter inference and sensitivity analysis. This allows for rapid calibration and uncertainty quantification workflows that would be impractical with repeated numerical solves.
• Modest Scope: The authors acknowledge that the current model is trained for a fixed time horizon and a low-dimensional coefficient vector. They note that while the approach is robust for the studied convection-diffusion-reaction dynamics, more complex phenotypic plasticity mechanisms (requiring PIDE models with integral terms) or spatially varying parameter fields would require further architectural generalization. The study does not claim to solve the biological mechanisms of cancer plasticity itself but provides a robust computational tool to model and analyze such dynamics.