Derivative Computation in PINNs: Automatic Differentiation, Finite Differences and Beyond

arXiv:2608.11020 · cs.LG, cs.NA, math.NA, physics.comp-ph · Submitted 2026-08-11 · Read on arXiv

Maciej J. Mikulski, Tadeusz Uhl

AGH University of Krakow

cs.LG, cs.NA, math.NA, physics.comp-ph

Submitted: 2026-08-11

Updated: 2026-08-12

Comments: 22 pages, 5 figures

Code: https://github.com/google/jax

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 75/100

The gist: The paper systematically investigates finite-difference (FD) derivative computation in Physics-Informed Neural Networks (PINNs) as an alternative to automatic differentiation (AD).

Terminology

Summary

The paper systematically investigates finite-difference (FD) derivative computation in Physics-Informed Neural Networks (PINNs) as an alternative to automatic differentiation (AD). On three benchmark PDEs, the authors show that with a properly calibrated step size, FD matches AD in accuracy on every problem while running faster across the full tested batch-size range and using substantially less GPU memory. They also propose a stochastic variant (sFD) that outperforms AD on a stationary problem. The paper further demonstrates that for neural architectures with inter-sample dependencies (e.g., BatchNorm, self-attention), the standard PyTorch autograd idiom is silently incorrect; the correct per-sample alternative is computationally infeasible at PINN-relevant batch sizes, while FD provides a forward-only approximation that is empirically an order of magnitude closer to the true per-sample derivative.

The paper notes that PINNs embed governing equations directly into the loss function of a neural network, unifying data and physical priors within a single optimization framework. Among practical challenges facing PINNs, the computation of derivatives stands out as particularly critical: the PDE residual loss requires evaluating partial derivatives of the network output with respect to its inputs, often up to second or higher order.

The authors state that modern deep learning frameworks like PyTorch and JAX provide AD capabilities, but "these mechanisms were designed and optimized primarily for first-order derivatives of the scalar loss with respect to neural networks weights, as used in the context of backpropagation for supervised training. They do not natively support efficient per-sample gradient computation required for PDE residuals at collocation points. AD-based gradient computation in PINNs leads to very quick graph build-up, especially for second order derivatives," which becomes prohibitive for large batch sizes or higher order differential equations.

The paper's contributions are:

  1. Demonstrating that the perturbation parameter ε is crucial for numerical accuracy, proposing a principled empirical method for selecting optimal ε based on the interplay between truncation and floating-point errors, and introducing a scheme that periodically recalibrates ε during training.

  2. Providing theoretical analysis of how optimal ε scales with derivative order and proposing three practical strategies for mixed-order PDEs: geometric-mean step size (FD), order-specific step sizes (eFD), and stochastic finite differences (sFD).

  3. Noting that a common way to compute derivatives for neural networks containing inter-sample dependencies leads to silently incorrect results, with theoretical and experimental evidence.

The paper focuses on 3-point central difference schemes. For the first derivative:

Dε(1) f (x) = [f (x + ε) − f (x − ε)] / (2ε)

For the second derivative:

Dε(2) f (x) = [f (x + ε) − 2f (x) + f (x − ε)] / ε2

In practice, this amounts to evaluating the neural network on an enlarged batch of perturbed points, i.e., the concatenation of [X − ε, X, X + ε]. This contrasts with Lim et al. [2022], where a grid is fixed and ε is determined by grid spacing.

The paper analyzes the trade-off between truncation error (from Taylor series truncation) and roundoff error (from floating-point arithmetic limitations).

For the first derivative:

  • Truncation error: E(1) trunc(ε) = (ε2/6)f‴(x) + O(ε3)

  • Roundoff error: E(1) round(ε) = κ1 · εmf(x)/ε, where κ1 is an unknown constant

  • Optimal step size: ε(1) opt = C1 · εm(1/3)

For the second derivative:

  • Truncation error: E(2) trunc = (ε2/12)f⁗(x) + O(ε4)

  • Roundoff error: E(2) round = κ2 · εmf(x)/ε2

  • Optimal step size: ε(2) opt = C2 · εm(1/4)

The paper provides a table showing theoretical scaling across floating-point formats:

  • FP64 (εm = 2.2×10−16): εm(1/3) = 6.1×10−6, εm(1/4) = 1.2×10−4

  • FP32 (εm = 1.2×10−7): εm(1/3) = 4.9×10−3, εm(1/4) = 1.9×10−2

  • FP16 (εm = 9.8×10−4): εm(1/3) = 9.9×10−2, εm(1/4) = 1.8×10−1

  • BF16 (εm = 7.8×10−3): εm(1/3) = 2.0×10−1, εm(1/4) = 3.0×10−1

The paper proposes a procedure that evaluates n logarithmically spaced candidate step sizes and selects the one that minimizes the RMSE between the FD approximation and the AD-computed reference. For FP32, the heuristic range is ε min = 1e-6, ε max = 1e-1 with 50 candidates. Periodic recalibration of ε is proposed at fixed intervals during training.

Three FD-based strategies are proposed:

  1. FD: Uses a single compromise step size computed as the geometric mean of order-specific optima: εGM = sqrt(ε(1) opt · ε(2) opt)

  2. eFD: Uses ε(1) opt for first derivatives and ε(2) opt for second derivatives, achieving optimal accuracy at the cost of 40% higher batch size

  3. sFD: For each training step and collocation point, samples ε log-uniformly from the range of optimal step sizes, maintaining a single 3-point stencil while introducing stochasticity

The paper identifies that architectures with Batch Normalization or Attention layers introduce dependencies between samples within a mini-batch. The common trick of setting grad outputs=ones like(outputs) in PyTorch is silently wrong. The only correct way is to compute per-sample gradients with O(B) separate autograd calls, which is prohibitively expensive in time and memory at the batch sizes needed for PINN training. FD reduces the discrepancy against the true per-sample reference by roughly an order of magnitude for BatchNorm and by a factor of three for self-attention.

All experiments use a fully connected MLP with 4 hidden layers and tanh activation. Hidden width is 128 for Burgers 1D and Poisson 2D, and 256 for Heat2D-CG. Training uses Adam optimizer with learning rate 10−3 and cosine annealing. The loss is L = LPDE + λIC·LIC + λBC·LBC with λIC = λBC = 10 for Burgers and Heat2D-CG, and λBC = 1000 for Poisson. Collocation points are resampled uniformly at each epoch with batch size 8192. FD step size ε is recalibrated every 4000 epochs. Each configuration is trained for 20k epochs with 25 independent random seeds. All experiments run on an NVIDIA GeForce RTX 4070 (12 GB) using PyTorch with FP32 precision.

The paper validates the theoretical error analysis by empirically measuring FD error as a function of ε on a fixed, randomly initialized MLP (4 layers, 64 units, tanh). Figure 2 shows the characteristic V-shaped curve of FD error as a function of step size. Figure 3 compares error landscapes across FP64, FP32, and BF16, with experimental minima agreeing with theoretical predictions, implying constant factor C1 ≈ 2–3.

Poisson 2D on a domain with holes (stationary, second-order only, so eFD coincides with FD):

  • sFD: SMAPE 13.10 ± 0.64%, L2 Relative 1.11 ± 0.39%, Max Error 0.026 ± 0.010

  • AD: SMAPE 13.67 ± 0.55%, L2 Relative 1.41 ± 0.31%, Max Error 0.032 ± 0.008

  • FD: SMAPE 13.66 ± 0.63%, L2 Relative 1.40 ± 0.39%, Max Error 0.033 ± 0.014

sFD achieves 21% lower L2 error than AD and 20% lower max error. AD and FD are statistically indistinguishable.

1D viscous Burgers equation (mixed first and second order derivatives):

  • AD: SMAPE 6.81 ± 0.04%, L2 Relative 1.36 ± 0.11%, Max Error 0.088 ± 0.012

  • eFD: SMAPE 6.81 ± 0.04%, L2 Relative 1.38 ± 0.17%, Max Error 0.091 ± 0.019

  • FD: SMAPE 6.82 ± 0.04%, L2 Relative 1.38 ± 0.14%, Max Error 0.093 ± 0.019

  • sFD: SMAPE 7.77 ± 0.34%, L2 Relative 4.23 ± 1.21%, Max Error 0.233 ± 0.120

AD, FD, and eFD are statistically indistinguishable. sFD performs poorly, with roughly 3× the L2 error of deterministic methods.

Heat2D-CG (3D problem with irregular geometry and Robin boundary conditions):

  • AD: SMAPE 14.79 ± 0.78%, L2 Relative 2.56 ± 0.32%, Max Error 0.571 ± 0.025

  • eFD: SMAPE 14.82 ± 0.66%, L2 Relative 2.60 ± 0.31%, Max Error 0.577 ± 0.024

  • FD: SMAPE 14.97 ± 0.79%, L2 Relative 2.66 ± 0.38%, Max Error 0.564 ± 0.018

  • sFD: SMAPE 19.88 ± 4.03%, L2 Relative 4.65 ± 2.24%, Max Error 0.581 ± 0.089

AD and eFD are tied as best methods. The AD baseline is competitive with the original PINNacle vanilla PINN (3.64%) and close to the best method (LAAF at 2.39%).

FD is faster than AD across the entire tested range. For the plain MLP, FD beats the memory-efficient batched AD idiom by roughly 2–3×. For architectures with inter-sample dependencies, the gap widens to 6–14×, but more importantly the batched idiom is silently incorrect there, so the only valid AD baseline is per-sample AD — against which FD is two to three orders of magnitude faster (360× for plain MLP, 1600× for MLP with BN, 1400× for MLP with attention at B = 27). Memory advantage is 2× for plain MLP and one-to-three orders of magnitude for inter-sample architectures, capping AD's feasible batch size at values too small for converged PINN training.

All FD methods are robust to recalibration frequency. L2 error is essentially flat from interval 1 to 4000 epochs. Only at interval 10000 (calibrated twice during entire training) is there substantially worse L2 error for sFD, and slightly worse for FD and eFD.

The paper's most practical finding is that on MLP-based PINNs with properly calibrated step size, AD, FD, and eFD produce statistically indistinguishable solutions on every problem tested. In the regime in which PINNs actually operate, the approximation error of optimally chosen ε is below the seed-to-seed noise of the training itself.

From a computational standpoint, AD of a scalar output with respect to its inputs builds a computational graph that must be retained through the backward pass, and for second-order derivatives this graph grows quadratically with the number of nonlinearity layers and linearly with batch size. Finite differences require only forward passes and have essentially constant memory per point.

sFD is interpreted as a form of noise-induced regularization. The empirical benefit is most visible on the stationary Poisson problem, where a common failure mode is solution collapse. The paper explains: "A residual evaluated at x no longer depends only on the infinitesimal behavior of uθ at x, but on values in a finite neighborhood of radius ε. Points near the boundary therefore become sensitive to narrow transition layers even when the collocation point itself does not lie inside the layer. In effect, FD increases the visible width of such pathological boundary layers."

sFD strengthens this effect by sampling ε across a range of scales. However, on Burgers and Heat2D-CG, where the solution contains physically meaningful sharp gradients, stochastic variation of ε can inject excessive variance or blur derivative information.

  • Standard MLP PINNs: Use FD (geometric-mean ε) as the default. It matches AD accuracy on every problem tested, is approximately 2× faster, and has essentially constant memory as derivative order grows. Recalibrate ε a few times during training.

  • BatchNorm models: Avoid the grad outputs=ones idiom — it is silently wrong; use FD. Alternatively implement a non-standard form of BatchNorm that uses running statistics also in training without introducing inter-sample dependencies.

The paper concludes that finite difference is a viable alternative to automatic differentiation. The choice of optimal step size is crucial for accuracy and can be done easily with the proposed method. As the characteristics of the neural network model change, it may be worth recalibrating the step size a few times during training. The common grad outputs=ones idiom is silently wrong in neural architectures that use inter-sample dependencies. The stochastic FD method introduces regularization that can help models suffering from solution collapse, but the exact mechanism requires further investigation.

For architectures with inter-sample dependencies, the paper reports median max-absolute discrepancy of Laplacian estimates across 5 random initializations:

  • Plain MLP: ∥batched AD − per-sample AD∥∞ = 0, ∥FD − per-sample AD∥∞ = 8.2×10−4

  • MLP with BN: ∥batched AD − per-sample AD∥∞ = 8.5, ∥FD − per-sample AD∥∞ = 0.79

  • MLP with attention: ∥batched AD − per-sample AD∥∞ = 7.4, ∥FD − per-sample AD∥∞ = 2.1

On the plain MLP all three methods agree at the level of FD truncation error. On MLP with BN and MLP with attention, the grad outputs=ones idiom is wrong by 7-8 absolute units, while FD's discrepancy is roughly an order of magnitude smaller for BN and more than a factor of three smaller for attention.

Improvements for AI systems

Based on this paper, here are specific improvements for AI systems:

  • Improvement: Implement finite-difference derivative computation with periodically recalibrated optimal step sizes (ε) instead of automatic differentiation in PINN training loops.

  • Capability: The improved system can train PINNs 2–3× faster with 2× lower GPU memory usage while achieving statistically identical accuracy to AD-based training on standard MLP architectures.

  • Improvement: Build a calibration routine that evaluates 50 logarithmically spaced candidate step sizes (range 1e-6 to 1e-1 for FP32) against an AD reference at initialization and periodically (every 4000 epochs) during training.

  • Capability: The system automatically maintains optimal numerical accuracy without manual tuning, robust to changes in network weights and loss landscape during optimization.

  • Improvement: For PDEs with mixed first and second derivatives, use separate optimal ε values for each derivative order (ε(1) opt = C1·ε m(1/3), ε(2) opt = C2·ε m(1/4)) rather than a single compromise value.

  • Capability: The system achieves AD-equivalent accuracy on complex PDEs (e.g., Burgers, Heat2D-CG) with only 40% larger batch size, while remaining faster and more memory-efficient than AD.

  • Improvement: Implement a training mode that samples ε log-uniformly from the optimal range at each step for each collocation point, providing noise-induced regularization.

  • Capability: The system can overcome solution collapse in stationary PDE problems (e.g., Poisson with holes), achieving 21% lower L2 error than AD while maintaining computational efficiency.

  • Improvement: Add a diagnostic check that compares batched AD (grad outputs=ones idiom) against FD or per-sample AD for architectures with BatchNorm or attention layers; automatically switch to FD when discrepancies exceed a threshold.

  • Capability: The system avoids silently incorrect derivative computations that can produce errors of magnitude 7–8 in Laplacian estimates, instead achieving 10× better accuracy with FD.

  • Improvement: Scale optimal ε according to floating-point format (FP64: 6e-6 for 1st derivative; FP32: 5e-3; FP16: 1e-1) using the theoretical scaling laws ε opt ∝ ε m(1/(order+1)).

  • Capability: The system maintains accuracy across mixed-precision training regimes, enabling efficient low-precision PINN training without sacrificing derivative accuracy.

  • Improvement: Replace graph-based backward passes with forward-only evaluations on perturbed batches [X−ε, X, X+ε], avoiding quadratic graph growth for second derivatives.

  • Capability: The system can train PINNs with batch sizes 1–3 orders of magnitude larger than AD-based approaches for architectures with inter-sample dependencies, enabling converged training where AD becomes infeasible.

  • Improvement: Implement adaptive recalibration that monitors loss landscape changes and recalibrates ε only when needed (robust from interval 1 to 4000 epochs, degrading only at 10000+).

  • Capability: The system minimizes computational overhead from unnecessary recalibration while maintaining accuracy across long training runs.

  • Improvement: Use FD for training but periodically validate with AD (or per-sample AD for dependent architectures) to confirm equivalence and detect any drift.

  • Capability: The system provides confidence guarantees that FD-based training produces solutions statistically indistinguishable from AD, with quantitative error bounds.

Sources

Related papers