Mollifier Layers: Enabling Efficient High-Order Derivatives in Inverse PDE Learning

TL;DR

The paper tackles the memory and numerical instability challenges of computing high-order derivatives in inverse PDE learning with physics-informed neural networks. It introduces Mollifier Layers, an output-layer module that replaces recursive autodiff with convolutions against analytic mollifier kernels, enabling stable, memory-efficient computation of derivatives via $\hat{u}(n) = \hat{g} * \eta(n)$ and $\hat{u}_j(n) = \hat{g} * \eta_j(n)$. Across 1D Langevin, 2D heat, and 2D reaction–diffusion problems, plus a biophysical application inferring spatially varying epigenetic rates from STORM images, mollified variants consistently improve parameter recovery, reduce training time, and lower memory usage by 6–10×. The approach is architecture-agnostic, scalable to forward solvers and neural ODEs, and opens avenues for broader PhiML applications where stable high-order gradients are essential.

Abstract

Parameter estimation in inverse problems involving partial differential equations (PDEs) underpins modeling across scientific disciplines, especially when parameters vary in space or time. Physics-informed Machine Learning (PhiML) integrates PDE constraints into deep learning, but prevailing approaches depend on recursive automatic differentiation (autodiff), which produces inaccurate high-order derivatives, inflates memory usage, and underperforms in noisy settings. We propose Mollifier Layers, a lightweight, architecture-agnostic module that replaces autodiff with convolutional operations using analytically defined mollifiers. This reframing of derivative computation as smoothing integration enables efficient, noise-robust estimation of high-order derivatives directly from network outputs. Mollifier Layers attach at the output layer and require no architectural modifications. We compare them with three distinct architectures and benchmark performance across first-, second-, and fourth-order PDEs -- including Langevin dynamics, heat diffusion, and reaction-diffusion systems -- observing significant improvements in memory efficiency, training time and accuracy for parameter recovery across tasks. To demonstrate practical relevance, we apply Mollifier Layers to infer spatially varying epigenetic reaction rates from super-resolution chromatin imaging data -- a real-world inverse problem with biomedical significance. Our results establish Mollifier Layers as an efficient and scalable tool for physics-constrained learning.

Figures (26)

• Figure 1: Limitations of autodiff and overview of PhiML+Mollifier architecture. (a) Training time comparison for PINNs with and without PDE residual loss. (b) PINN-predicted vs. actual Laplacian in a reaction-diffusion system. (c) PINN-predicted vs. actual forcing term in a Langevin system. (d) PhiML+Mollifier Layer architecture replacing autodiff with mollifier-based convolution.
• Figure 2: Parameter inference for the Langevin equation using PINN and Mollified PINN. a) Ground truth and predicted $u(t)$. b–c) $MSE_u$ Data and $MSE_f$ PDE residual learning curves. d) Ground truth and inferred forcing term $\lambda(t)$ under Gaussian noise. e–f) Actual vs. predicted noise trends in $\lambda(t)$. g) Inferred time-varying $\lambda(t)$ without noise. h) Inferred time-varying $\lambda(t)$ with added noise.
• Figure 3: Parameter estimation for the Heat and Reaction–Diffusion equations using Mollified PINN. a) Ground truth and predictions for constant and spatially varying thermal diffusivity $\lambda(x, y)$ and Laplacians for the 2D heat equation. b) Predicted vs. actual chromatin order parameter $\phi_d$, recovered Laplacian, and training curves for the reaction–diffusion system.
• Figure 4: Mollified PINN capture spatial noise trends and extract reaction rate statistics in the DNA Reaction–Diffusion system. a) Spatial noise trends in inferred reaction rates from synthetic data. b,c) Predicted mean and variance of reaction rates from super-resolution images of human nuclei.
• Figure 5: Constant forcing term inference in Langevin equation using PINN and Mollified PINN