Mollifier Layers: A New Way to Stabilize Neural Networks Solving Inverse PDEs

Physics-informed machine learning has a load-bearing dependency that rarely gets examined: automatic differentiation. The whole approach rests on being able to take derivatives of a network’s output with respect to its inputs, plug them into a PDE residual, and minimize. Autodiff makes that mechanical, which is why it became the default and why its limitations get treated as implementation details rather than as the constraint they are.

For forward problems with clean data and low-order operators, the default holds up. For inverse problems — recovering spatially or temporally varying parameters from observations — with high-order operators and real measurement noise, it does not.

Mollifier Layers: Enabling Efficient High-Order Derivatives in Inverse PDE Learning (Bhartari, Vinayak, and Shenoy, arXiv:2505.11682, May 2025) proposes replacing the derivative computation entirely, and does it with a tool from the 1930s.

Why recursive autodiff degrades

Autodiff computes an exact derivative of the network as implemented. That exactness is often mistaken for accuracy with respect to the underlying physical field, and the two come apart in three ways.

Error compounds with order. A fourth-order derivative means differentiating the computational graph four times. Each pass builds a larger graph over the previous one, and whatever high-frequency structure the network has learned — including structure it learned from noise — is amplified at every step. A network can fit observations well and still produce a fourth derivative that is numerically meaningless.

Memory grows with the graph. Recursive differentiation constructs progressively larger graphs. For fourth-order operators on non-trivial architectures, memory becomes the binding constraint before accuracy does, which is why practitioners quietly cap the order of problems they attempt.

Noise gets differentiated too. This is the central issue for inverse problems. Differentiation is a high-pass operation: it amplifies exactly the high-frequency content that measurement noise contributes. A network trained to fit noisy observations has encoded some of that noise, and each derivative pass amplifies it further. The result is a residual dominated by noise-derived terms, and parameter estimates that inherit the corruption.

That last point is why the problem is specific to inverse problems. In a forward problem, the network is fit to a PDE and derivatives are taken of a comparatively smooth object. In an inverse problem, the network is fit to data, and the data is noisy.

What a mollifier does

A mollifier is a smooth, compactly supported function that integrates to one — a bump. Convolving a rough function with a mollifier produces a smooth approximation that converges back to the original as the mollifier narrows. Mollifiers were introduced in the 1930s by Kurt Friedrichs and Sergei Sobolev precisely to make analysis possible on functions that are not differentiable.

The property this work exploits is the one that makes mollifiers useful in analysis: differentiation can be moved onto the mollifier. Because the mollifier is defined analytically, its derivatives are known exactly. So instead of differentiating the network and hoping the result means something, you convolve the network output with the analytic derivative of the mollifier.

The derivative is never computed by differentiating the network. It is computed by integrating against a kernel whose derivative you already know in closed form. Differentiation becomes smoothing integration — a low-pass operation rather than a high-pass one.

That inversion is the whole idea, and it explains each of the claimed benefits without requiring anything else to be true:

• High-order derivatives cost a convolution, not a recursive graph expansion, so memory stops scaling with derivative order.
• Noise is attenuated by the smoothing rather than amplified by the differentiation.
• The smoothing width becomes an explicit, tunable regularization parameter instead of an implicit consequence of architecture and training.

The engineering property that matters

Mollifier layers attach at the output layer and require no architectural modification. The authors describe the module as lightweight and architecture-agnostic.

That is a larger deal than it sounds. A great many proposed improvements to physics-informed networks require adopting a particular architecture, loss weighting scheme, or training procedure, which means evaluating them costs a rewrite. A module that sits after an existing network and changes how derivatives are obtained can be tested against a working baseline in an afternoon — and, more importantly, can be removed if it does not help.

It also composes cleanly with the rest of the stack. Nothing about it conflicts with adaptive loss weighting, domain decomposition, or whatever sampling strategy is already in use.

What was tested

The evaluation spans first-, second-, and fourth-order PDEs: Langevin dynamics, heat diffusion, and reaction-diffusion systems. The paper reports improvements in memory efficiency, training time, and accuracy of parameter recovery relative to autodiff-based approaches.

The application that carries the most weight is not synthetic: inferring epigenetic reaction rates from chromatin imaging data. Imaging data is exactly the regime the method is aimed at — noisy, indirect measurements of a field whose governing parameters are the actual object of interest. A method that works on a clean manufactured benchmark and fails on imaging noise would not be interesting; the inclusion of a real inference task is the meaningful part of the evaluation.

Where the judgment calls sit

Anyone adopting this should be clear about what it trades.

The smoothing width is a real hyperparameter with real consequences. Too narrow and noise survives; too wide and genuine features of the field get smoothed away, biasing the recovered parameters. It is more interpretable than most hyperparameters — it has units, and it corresponds to a physical length or time scale — but it does not tune itself, and the appropriate value depends on the noise level and the scale of the structure you are trying to recover.

Bias is being traded for variance, deliberately. Convolution with a finite-width kernel introduces bias. In the noisy-data regime that is a good trade, because the variance of an autodiff-computed fourth derivative on noisy data is enormous. In a clean, high-signal setting it may not be.

Boundaries need care. Compactly supported kernels behave differently near domain edges, where the support extends past available data. This is a well-understood problem in numerical analysis with standard remedies, but it is a place where a naive implementation will produce artifacts precisely where boundary conditions matter most.

The broader point

The interesting thing about this line of work is not the specific module. It is the observation that automatic differentiation was adopted into scientific machine learning as infrastructure — a thing you use without evaluating — when it is in fact a modeling choice with a specific error profile that happens to be badly matched to inverse problems on noisy data.

Reframing derivative computation as smoothing integration, with a tool that predates neural networks by half a century, is a reminder that “how do we take this derivative” is a question worth reopening rather than delegating to the framework.