PaperlayerView exact revision on arXiv

Source-linked AI summary

Physics-Informed Foresight Pruning for Sparse PINN Solvers of Nonlinear PDEs

Ahmad Ishaque Karimi, Uvini Balasuriya Mudiyanselage, Kookjin Lee

arXiv:2608.25564v1cs.LGcs.AI

TL;DR

PINN pruning must preserve both predicted solutions and the derivatives entering PDE residuals, but standard output-side spectral pruning may miss residual-sensitive parameters. The paper introduces PI-SAP for sparse PirateNet solvers and compares it with NTK-SAP across four PDE families. PI-SAP improves Gray–Scott residual fidelity across tested sparsities and is competitive at aggressive sparsity, while diagnostics show that solution accuracy, residual fidelity, and kernel conditioning remain distinct objectives.

  • Problem

    PINNs use substantial capacity for coupled solution and differential-residual objectives, leaving unclear how much capacity is necessary and what pruning should preserve.

  • Method

    The paper applies initialization-time NTK-SAP and residual-sensitivity-based PI-SAP to sparse residual-adaptive PirateNet PDE solvers.

  • Results

    PI-SAP improves Gray–Scott residual fidelity across the complete sparsity sweep and is competitive at high sparsity, while no criterion is uniformly optimal across equations or sparsity levels.

  • Takeaways & Limitations

    Pruning should balance solution-side and residual-side training dynamics and evaluate solution error, physics residual, and task-aware diagnostics separately.

  • Takeaways & Limitations

    The small-batch PINN-NTK diagnostics approximate full training-set NTKs, and their coefficients were not extensively tuned.

Abstract

from arXiv · show

Physics-informed neural networks (PINNs) often rely on over-parameterized models to optimize coupled solution and differential-residual objectives, leaving unclear how much capacity is necessary and what pruning should preserve. We study foresight pruning at initialization for sparse PirateNet PDE solvers. Standard neural tangent kernel spectrum-aware pruning (NTK-SAP) aims to preserve output-side training dynamics but may overlook parameters whose main influence arises through derivatives in the governing equations. We introduce physics-informed spectrum-aware pruning (PI-SAP), which assigns saliency using sensitivity of the PDE residual. Experiments on the Gray-Scott equations, complex Ginzburg-Landau equation, Burgers' equation, and linear convection equation show that PI-SAP more consistently preserves Gray-Scott residual fidelity and is competitive under aggressive sparsity. However, no criterion is uniformly optimal across equations or sparsity levels. Small-batch PINN-NTK diagnostics further show that residual fidelity, solution accuracy, and kernel conditioning are distinct objectives, motivating pruning methods that explicitly balance solution-side and residual-side training dynamics during optimization.

I. INTRODUCTION

The paper asks whether initialization-time pruning can reduce PINN capacity while preserving both solution and PDE-residual training behavior. It introduces residual-aware saliency and finds a qualified, equation- and sparsity-dependent advantage rather than a universal winner.

  • PINNs must represent solution fields and their spatial and temporal derivatives, making pruning difficult for nonlinear PDEs with coupled fields or sharp gradients.These demands contribute to conditioning problems, gradient imbalance, and increased computational cost.
  • NTK-SAP preserves an output-side training-dynamics proxy, but parameters with limited output influence can strongly affect derivatives and PDE residuals.PI-SAP therefore scores weights through sensitivity of the residual network.
  • The evaluation spans Gray–Scott, complex Ginzburg–Landau, Burgers’, and convection equations using residual-adaptive PirateNets.Gray–Scott provides the most complete benchmark, including derivative-sensitive interfaces.
  • PI-SAP improves Gray–Scott residual fidelity throughout the sparsity sweep and is particularly competitive at high sparsity, while NTK-SAP can be stronger at intermediate sparsity.The criteria protect complementary aspects of training: output-side dynamics and residual-sensitive physics.
  • Solution accuracy, residual fidelity, and kernel conditioning are distinct objectives that can evolve at different rates during PINN training.The paper consequently motivates pruning criteria that balance solution-side and residual-side dynamics.

II. TECHNICAL BACKGROUND

PINNs train differentiable neural solution maps by combining initial or boundary constraints with governing-equation residuals. The study uses residual-adaptive PirateNets and globally masks selected kernel parameters while keeping gates, biases, and output coefficients dense.

  • A PINN differentiates predicted fields with respect to input coordinates, substitutes the derivatives into the governing PDE, and minimizes residual and initial or boundary errors.This distinguishes its training signal from ordinary supervised learning.
  • The neural solution qθ maps space–time coordinates to fields, while the pointwise residual is rθ(z) = N[qθ](z).For coupled systems, qθ and rθ contain one component for each field.
  • PirateNets begin as shallow mappings and progressively introduce nonlinear depth during training to avoid unstable initialization of deep PDE residual networks.The architecture serves as the backbone for each PINN solver.
  • The two-dimensional implementation uses coordinate encodings, auxiliary feature streams, and three residual-adaptive blocks of width 256 with Swish activations.The representation is projected to predicted PDE fields.
  • Global pruning covers eligible kernel parameters, while biases, adaptive gates αℓ, and physics-informed output coefficients remain dense.The eligible set includes the Fourier-feature, auxiliary-stream, and residual-block kernels, including both factors from random weight factorization.
  • After thresholding, the binary mask is applied to parameter values and gradients after every Adam update, preventing pruned parameters from regrowing.NTK-SAP and PI-SAP share this architecture, eligible set, and enforcement procedure.

C. NTK View

The NTK view links parameter-to-output Jacobians and kernel eigenspectra to gradient-descent dynamics. For PINNs, separate solution and residual kernel blocks mean output-side spectral preservation may not preserve residual-side behavior.

  • The empirical NTK is K(X, X) = Jθ(X)Jθ(X)⊤, where Jθ(X) contains derivatives of stacked outputs with respect to trainable parameters.The kernel is formed from the parameter-to-output Jacobian.
  • NTK eigenvalues control linearized squared-loss training rates: small eigenvalues create slow modes, while wide spectral spread produces uneven learning.Conditioning here refers to empirical-kernel eigenspectral conditioning, including κ(K) = λmax(K)/λmin(K).
  • NTK-SAP prunes connections with little influence on an NTK-spectrum proxy to preserve important aspects of dense-model optimization behavior.The approach is supported by the connection between spectral properties and convergence dynamics.
  • PINN solution and residual terms can converge at different rates because training is governed by different NTK blocks.This motivates pruning criteria that account for the differential operator rather than matching only output-side dynamics.

D. NTK-SAP Baseline

NTK-SAP selects a sparse mask before training by ranking eligible connections with a tractable output-dynamics proxy. The mask is globally thresholded and enforced throughout optimization.

  • NTK-SAP computes the sparse mask at initialization and enforces it during training by masking both gradients and parameters.This prevents pruned connections from regrowing through optimizer momentum.
  • The method scores mask variables using sensitivity of a perturbed output difference over a pruning input batch.The perturbation is applied to eligible kernel parameters.
  • Only ordinary kernel weights are eligible for pruning; biases, PirateNet α gates, and other initialization coefficients retain mask value one.The same scope is used for the global threshold.
  • For target sparsity s, weights at or below the global 100s-th percentile of eligible scores are zeroed, while higher-scoring weights are retained.The rule removes weights judged least influential to the output-dynamics proxy.

III. PROPOSED METHOD

PI-SAP ranks parameters by their influence on PDE residuals, while diagnostic extensions evaluate conditioning and separately balance solution-side and residual-side kernel behavior. Cross-kernel terms remain future work.

  • PI-SAP: PI-SAP uses the same global ranking and target sparsity as NTK-SAP, but protects connections that most affect governing-PDE satisfaction.For reaction–diffusion systems, the intended saliency targets include weights influencing derivative-sensitive diffusion and reaction terms.
  • Diagnostic Extensions: Conditioning-aware NTK-SAP favors masks with larger average spectral scales while penalizing condition-number and log-spectrum spread terms.Its saliency is computed from |θ_j ∂ϕ/∂θ_j|.
  • Diagnostic Extensions: PINN-block-aware SAP constructs separate solution and residual Jacobians and their corresponding empirical kernel blocks on a shared small batch.For two-field systems, the Jacobians stack field-wise solution and residual component derivatives.
  • Diagnostic Extensions: The block-aware objective combines solution- and residual-kernel criteria and penalizes mismatch between their mean log-eigenvalue scales.Cross terms are left for future work.

IV. EXPERIMENTAL SETUP

The experiments evaluate pruned PirateNet solvers across four PDE settings, with the Gray–Scott study using a fixed, detailed training configuration and multiple accuracy, residual, and runtime metrics.

  • Gray–Scott Equations: The Gray–Scott study uses a scaled system on Ω = [−1, 1]^2 with t ∈ [0, 2], 101 snapshots, and a 200×200 reference grid.It uses specified diffusion, feed, kill, PirateNet, optimization, weighting, and causal-time-marching settings.
  • Cross-PDE Evaluation: Complex Ginzburg–Landau runs compare dense, NTK-SAP, and PI-SAP models across five pruning levels.Burgers’ and convection experiments use prior cross-PDE studies averaged over seeds 0–4 with 50k Adam steps.
  • Metrics: Evaluation reports relative L2 error for each field, mean solution error for two-field systems, mean PDE residual for Gray–Scott and Ginzburg–Landau, and runtime.Runtime is wall-clock training time for the masked implementation.
  • Metrics: Because pruning masks dense arrays instead of invoking sparse kernels, proportional wall-clock speedups are not expected.This implementation choice limits runtime interpretation.

A. Gray–Scott Equations: Main Case Study

In the Gray–Scott stripe regime, PI-SAP more consistently preserves physics residuals and high-frequency structure, while solution-error leadership varies with pruning level and both methods degrade sharply at high sparsity.

  • Main Results: PI-SAP has lower mean PDE residual than NTK-SAP at every pruning level, while NTK-SAP is slightly better for mean solution error at 50% pruning.PI-SAP has lower mean solution error at 10%, 30%, 70%, and 90% pruning.
  • Solution Error: 34.1%: PI-SAP reduces mean error relative to NTK-SAP at 10% pruning.This is the reported relative reduction for the Gray–Scott equations.
  • Sparsity and Runtime: Both methods show a sharp solution-error increase between 50% and 70% pruning, placing the practical accuracy threshold for this configuration in that interval.Runtime remains 16.4–18.0 hours without a systematic wall-clock advantage from masking.
  • High-Frequency Structure: PI-SAP roughly halves terminal high-frequency error relative to NTK-SAP for both species at 70% pruning.The Fourier analysis targets spatial modes with radial frequency at least 0.35 of the maximum discrete frequency.

C. NTK-Block Diagnostics for the Gray–Scott Equations

The 70% Gray–Scott diagnostics reveal a trade-off: different pruning criteria favor solution accuracy, residual loss, or conditioning, and results vary across PDE comparisons.

  • Diagnostic scope: The diagnostic runs use 70% pruning and a shorter 600k-step budget, so they should not replace the completed sparsity sweep.The table is explicitly labeled as a diagnostic extension rather than the main benchmark.
  • Gray–Scott diagnostics: Original NTK-SAP has the best solution accuracy among pruned models, while PINN-block-aware SAP has the lowest residual losses but worst solution error.These outcomes show that residual-side optimization does not necessarily preserve the full solution field.
  • Gray–Scott diagnostics: Conditioning-aware NTK-SAP does not improve over original NTK-SAP in this setting.The conditioning-aware variant is therefore informative diagnostically but not superior in this experiment.
  • Ginzburg–Landau comparison: Both pruning methods improve over dense at 10% in Ginzburg–Landau, NTK-SAP is better at 30% and 50%, and PI-SAP is better at 70% and 90%.PI-SAP improves mean error by about 11.5% relative to NTK-SAP at 70% pruning and about 7.4% at 90%.
  • Interpretation: The cross-equation comparison rejects a universal winner: pruning criteria protect complementary output-side and residual-sensitive training dynamics.The balance between these objectives depends on the PDE and sparsity level.

E. Burgers’ Equation and Convection Equation

Burgers’ and convection-equation results show that pruning can outperform dense models, but the preferred saliency rule changes with width, transport strength, and sparsity.

  • Convection equation: For the convection equation, the preferred method varies with both β and network width.The table is organized across transport parameter β and network width rather than a single fixed configuration.
  • Experimental setup: Burgers’ and convection studies average results over seeds 0–4 for 50k Adam steps and report only the best pruned result in each table.The compact reporting differs from the full sparsity sweeps used elsewhere.
  • Burgers’ equation: Best pruned Burgers’ models outperform dense baselines at all three widths, with the strongest effect at width 512.The results support pruning as a possible form of structural regularization.
  • Burgers’ equation: PI-SAP is best for Burgers’ equation at widths 128 and 256, whereas NTK-SAP is best at width 512.Thus, the preferred saliency criterion changes with network width.
  • High-β stress test: At β = 15, NTK-SAP at 70% sparsity reduces mean relative L2 error from 5.3298×10^-2 to 2.5010×10^-2.At β = 20, NTK-SAP at 50% sparsity reduces the error from 1.3302×10^-1 to 7.9886×10^-2.
  • Discussion: Across four PDE families, the useful criterion depends on which part of physics-constrained training dynamics is under pressure.Solution error, physics residual, and task-aware spectral or morphological diagnostics should be evaluated together.

VII. LIMITATIONS AND REPRODUCIBILITY

The study reports reproducibility and scope limitations alongside evidence that pruning behavior varies across PDEs and sparsity regimes. Its diagnostics are preliminary, and its runtime measurements do not establish hardware-level sparse acceleration.

  • Runtime interpretation: Dense JAX masks make reported runtimes measure training stability and overhead rather than hardware-level sparse speedup.Gray–Scott and complex Ginzburg–Landau runtimes remain nearly constant across sparsity levels.
  • Statistical scope: Gray–Scott and complex Ginzburg–Landau results are not based on large multi-seed sweeps, limiting statistical assessment of initialization effects.Burgers’ and convection studies average over five seeds, while the other benchmarks are primarily completed runs.
  • Future reproducibility: The high-sparsity Gray–Scott and complex Ginzburg–Landau settings should be repeated across seeds near the accuracy-degradation threshold.These differences may be sensitive to initialization effects.
  • PINN-block diagnostics: Small-batch Ksol and Kres are practical approximations, and untuned diagnostic coefficients make the experiments an initial objective-design stress test.Constructing full NTK blocks would be computationally prohibitive.
  • Scope of conclusions: The conclusion is PDE- and sparsity-dependent: PI-SAP improves Gray–Scott residuals across tested sparsities, while NTK-SAP can remain competitive or superior elsewhere.The paper therefore frames its results as a pruning principle rather than a universally optimal criterion.
Loading 2608.25564v1…