Computing high-order mixed derivatives in physics-informed neural networks using multi-index Bell polynomials

Technical Summary: Computing High-Order Mixed Derivatives in PINNs Using Multi-Index Bell Polynomials

Problem Statement
Physics-informed neural networks (PINNs) solve partial differential equations (PDEs) by minimizing residuals that depend on network derivatives with respect to inputs. While first and second derivatives are standard, high-order PDEs—such as the Kawahara, Zakharov–Kuznetsov (ZK), and Lax equations—require mixed derivatives of order five, seven, or higher. Standard implementations rely on repeated automatic differentiation (AD), which generates rapidly growing computational graphs and expressions. This approach often leads to memory exhaustion and inefficiency when computing high-order mixed partials and their gradients with respect to network parameters. Existing alternatives, such as Taylor-mode AD or jet-based methods, often compute directional derivatives or require multiple passes to recover a full set of mixed partials, and they may still suffer from nested graph complexity.

Methodology
The paper proposes a deterministic method to compute high-order mixed derivatives and their parameter gradients in a single forward and backward pass. The approach is built on three core components:

1. Forward Recursion via Faà di Bruno and Bell Polynomials:
   The method organizes the multivariate Faà di Bruno formula over a prescribed "downward-closed" set of multi-indices $A$. For a fully connected network with affine maps and smooth elementwise activations, the $K$-th order derivatives are computed using partial Bell polynomials.

   • The forward pass propagates derivatives $T^{(l, \alpha)}$ through the network layers.
   • Crucially, the Bell polynomial convolutions (which depend only on the multi-index set and the order) are tabulated once and reused, avoiding the need to re-evaluate partition sums during the forward pass.
   • The method supports a "residual seeding" strategy where the multi-index set $A$ is automatically restricted to only those indices required by the specific differential operator, omitting unnecessary derivatives.

2. Explicit Backpropagation Recursion:
   Instead of differentiating through a forward-mode pass or relying on nested AD graphs, the paper derives an explicit backward sweep.

   • The adjoint variables (gradients of the loss with respect to intermediate derivatives) are propagated using a recurrence relation derived from the derivative of the Bell polynomials.
   • This recursion computes the gradient of any loss formed from the derivatives (including nonlinear products and coupled fields) in a single backward pass.
   • The method handles products of up to four derivative factors, covering advective terms, body forces, and higher-order nonlinearities.

3. Implementation and Verification:
   The algorithm is implemented in Fortran 2008 (DNNF90) using contiguous allocatable arrays to minimize memory overhead. Verification is performed via:

   • Comparison with an independent multivariate Taylor-jet implementation (agreement to $< 4 \times 10^{-19}$).
   • Symbolic checks of test problems.
   • Finite difference checks for input derivatives and loss gradients up to order seven.

Key Contributions

• Efficiency in High-Order Derivatives: The method avoids the memory failures and exponential growth of computational graphs associated with nested AD. On a single CPU core, it evaluates 330 mixed derivatives (order 7, 4 inputs) and the corresponding loss gradient without memory exhaustion, whereas several standard AD implementations fail at order 5 or 6.
• Downward-Closed Optimization: By restricting calculations to a downward-closed set of indices derived from the specific PDE residual, the method reduces the computational load significantly (e.g., reducing the ZK test from 330 to 89 indices) without altering the solution trajectory.
• Closed-Form Gradient Computation: The explicit backpropagation recursion provides a closed-form reverse sweep for parameter gradients, distinct from differentiation through a forward-mode pass.
• Broad Applicability: The implementation supports orders up to 15 and handles coupled multi-field systems and nonlinear products of derivative factors.

Results
The paper validates the method on several benchmarks:

• Performance: In timing tests on a 4–8–8–1 network, the proposed method computes gradients for order-7 mixed derivatives in ~9 ms. In contrast, nested AD implementations (PyTorch, JAX) exhaust memory beyond order 4 or 5, and CoDiPack takes ~38 seconds for order 7. The method's time complexity grows by a factor of ~1.3 per order, compared to 3.1–8.5 for nested methods.
• Accuracy:
  • Dispersive Equations: The method solves third-, fifth-, and seventh-order dispersive equations (KdV, Kawahara, ZK) with relative solution errors ranging from $10^{-4}$ to $10^{-2}$. The seventh-order ZK test in 3+1 dimensions achieves a relative error of $6 \times 10^{-4}$.
  • Coupled Systems: It successfully solves the Kovasznay flow (incompressible Navier-Stokes) and a five-field electrohydrodynamic system.
  • Lax Equation: A seventh-order Lax equation with quartic nonlinear terms is solved with a relative $L^2$ error of $8.8 \times 10^{-4}$.
• Optimization Behavior: The paper notes that while residuals can be minimized effectively, solution accuracy depends heavily on the activation function and optimization strategy. For high-order derivatives, bounded activation functions (sine, Bessel functions $J_0, J_1$) outperform tanh in terms of loss and derivative accuracy, particularly when the highest-order term is essential.
• Failure Modes: The study highlights that minimizing residuals alone does not guarantee solution accuracy. In the Kovasznay flow without boundary data and in Kalman filter updates with forgetting factors, the residual decreased by orders of magnitude while the solution diverged from the exact field.

Significance and Limitations
The paper claims significance in providing a specialized, memory-efficient, and exact (up to roundoff) method for high-order mixed derivatives in PINNs, specifically addressing the bottleneck of repeated automatic differentiation. It demonstrates that restricting the derivative set to the necessary downward-closed indices yields substantial speedups without loss of accuracy.

The author modestly notes the following limitations:

• The method is specialized to networks composed of affine maps and smooth elementwise activations; it does not replace general AD for arbitrary computational graphs.
• It requires derivatives of the activation function up to order $K+1$.
• The computational cost for a dense set of derivatives still grows combinatorially as $\binom{D_0+K}{K}$.
• The reported timing comparisons are implementation-dependent and based on single-core execution.
• The numerical examples are controlled tests with exact or manufactured solutions, not new physical simulations.

The work concludes that while the method effectively solves the derivative computation bottleneck, careful attention to activation functions and optimization strategies is required to ensure that residual minimization translates to accurate physical solutions.