PaperlayerView exact revision on arXiv

Source-linked AI summary

Sensitivity-Constrained Neural Operators for Data-Efficient Forward and Inverse Modeling of Partial Differential Equation Systems

Abdolmehdi Behroozi, Chaopeng Shen, Daniel Kifer, Kathryn Lawson

arXiv:2608.29888v1cs.LG

TL;DR

High-dimensional PDE inputs can leave state-only neural operators with inaccurate input–output sensitivities, limiting reliability in forward, inverse, and repeated inference. The paper introduces SC-NOs, which add sampled solver-derived Jacobian supervision to existing operators. Across controlled benchmarks and a tsunami proof of concept, sensitivity supervision improves forward and inverse behavior and the observed accuracy–cost tradeoff, while remaining an empirical result within the tested settings.

  • Problem

    State-only neural-operator training constrains solution values but can learn inaccurate sensitivities for high-dimensional inputs, which matters for inverse and repeated inference.

  • Method

    SC-NOs augment existing neural operators with sampled solver-derived Jacobian supervision, resampling sensitivity entries across minibatches instead of imposing the full Jacobian each update.

  • Results

    Across controlled advection–diffusion and RANS benchmarks, sensitivity supervision improved forward prediction and produced larger gains in gradient-based inverse reconstruction, while improving the observed accuracy–cost tradeoff.

  • Takeaways & Limitations

    Sampled Jacobian supervision is presented as a practical strategy for neural PDE surrogates used where forward accuracy, inverse stability, robustness, and computational cost matter together.

  • Takeaways & Limitations

    The evidence is empirical over the tested PDE classes, resolutions, and training budgets, and the tsunami workflow is not a complete operational warning system.

Abstract

from arXiv · show

Neural operators provide fast surrogates for partial differential equation (PDE) solvers, but their reliability can degrade for high-dimensional spatial inputs and inverse or repeated inference. State-only training constrains solution values but not the learned input--output response. We study sensitivity-constrained neural operators (SC-NOs), which augment standard training with sampled solver-derived Jacobian supervision. Selected sensitivities from differentiable solvers or discrete adjoints are matched during training, allowing response information to be amortized across minibatches without imposing the full Jacobian at every update. We evaluate SC-NO on advection--diffusion and RANS--Spalart--Allmaras benchmarks, input-dimensionality scaling tests, long-horizon autoregressive rollout, and a shallow-water Tohoku tsunami source-inversion case. Sensitivity supervision improves forward prediction and yields larger gains in gradient-based inverse reconstruction of distributed fields. Scaling experiments show an improved accuracy--cost tradeoff for high-dimensional gridded inputs, while ablations indicate that state values and Jacobian information provide complementary supervision. In the tsunami case, SC-FNO reconstructs gridded seafloor deformation from sparse early gauge observations and forecasts subsequent wave propagation in a near-real-time proof-of-concept workflow. These results support sampled sensitivity supervision as a practical way to improve neural PDE surrogates when forward accuracy, inverse stability, robustness, and computational cost must be considered together.

1 Introduction

The paper addresses unreliable neural-operator responses for high-dimensional PDE inputs by augmenting existing architectures with sampled solver-derived Jacobian supervision. SC-NOs match both solution fields and selected input–output sensitivities while limiting per-update cost.

  • Motivation: High-dimensional gridded inputs make it difficult for state-only neural operators to learn how distributed parameters affect PDE solutions.This can require large datasets and degrade performance when inputs change substantially.
  • Motivation: Incorrect sensitivities can undermine robustness, out-of-distribution behavior, autoregressive rollout, and gradient-based inverse reconstruction despite acceptable forward predictions.The issue is inaccurate local derivatives with respect to the input.
  • Method: SC-NOs augment existing neural-operator families with selected solver-derived Jacobian entries from differentiable solvers or discrete adjoints.The method samples and resamples Jacobian subsets instead of imposing the full Jacobian at every optimization step.
  • Scope: The study evaluates forward prediction, inverse reconstruction, input-dimensionality scaling, out-of-distribution behavior, long-horizon rollout, and tsunami source inversion across existing operator families.The contribution is a training strategy rather than a new neural-operator architecture.
  • Method: The training objective combines state-prediction and Jacobian-matching losses, with learnable weights balancing the two supervision sources.The model Jacobian is obtained by automatic differentiation and matched against solver-derived sensitivities.
  • Method: Sampled Jacobian supervision reduces each update from O(NuNp) full-Jacobian entries to O(mNp) entries when m ≪ Nu.The paper treats these asymptotic quantities as cost scaling, while empirical wall-clock costs include sensitivity preparation.

3 Experimental Design

The experiments compare state-only and sensitivity-constrained versions of established neural operators across controlled PDE benchmarks, scaling and rollout tasks, and a Tohoku tsunami application. Costs include data generation, Jacobian preparation, and model training.

  • Benchmark PDEs: Three PDE classes are evaluated: advection–diffusion, RANS–Spalart–Allmaras, and shallow-water Tohoku tsunami dynamics.The first two provide controlled tests; the tsunami case serves as the larger application benchmark.
  • Benchmark PDEs: The advection–diffusion benchmark uses spatially varying velocity and initial concentration fields as high-dimensional inputs for concentration-trajectory prediction.Sensitivity supervision targets the distributed inputs used in the operator task.
  • Benchmark PDEs: The RANS–Spalart–Allmaras benchmark varies initial vorticity and distributed forcing, including a multi-step trajectory task and input-dimensionality scaling experiments.Intrinsic input resolution is varied while the training grid remains fixed to isolate dimensionality effects.
  • Benchmark PDEs: The tsunami benchmark reconstructs earthquake-induced seafloor deformation from sparse early gauge observations and forecasts subsequent water-stage evolution.The source-inference setting uses observations from the first 30 minutes and a differentiable neural-operator surrogate.
  • Model comparisons: FNO, WNO, and DeepONet are each compared with a paired sensitivity-constrained variant trained using the same state objective plus Jacobian matching.The comparisons test whether benefits persist across operator parameterizations rather than introducing a new architecture.
  • Evaluation tasks: Evaluation covers forward prediction, inverse reconstruction, out-of-distribution inputs, and long-horizon autoregressive rollout.The rollout task recursively feeds predictions back into a one-step transition model.
  • Computational cost: Total reported cost includes reference-data generation, solver-derived Jacobian preparation, and neural-operator training; state-only baselines set CJ = 0.Sensitivity overhead is included in wall-clock accounting rather than treated as free information.

4 Controlled PDE Benchmarks

Controlled benchmarks show that sensitivity-constrained neural operators improve forward prediction, inverse reconstruction, rollout stability, and scaling efficiency relative to state-only baselines. Ablations indicate that Jacobian supervision complements state-value supervision, while partial sampling provides a practical accuracy–cost compromise.

  • Forward prediction: Sensitivity-constrained models achieve lower forward relative L2 error than state-only counterparts, especially when training data are limited.The performance gap narrows as additional solution samples are added.
  • Forward prediction: The forward benefit of sensitivity supervision is consistent across FNO, WNO, and DeepONet, with FNO showing the largest improvement.This indicates the gain is not restricted to one neural-operator architecture.
  • Inverse reconstruction: Sensitivity-constrained models produce lower inverse reconstruction errors than state-only baselines for initial concentration and vorticity fields across sample sizes and computational costs.The separation is larger than in forward prediction because inverse reconstruction depends directly on surrogate gradients.
  • Inverse reconstruction: Inverse improvements are consistent across FNO, WNO, and DeepONet, while FNO achieves the strongest inverse performance after sensitivity supervision.The method improves optimization gradients but does not make the inverse problems well posed.
  • Long-horizon rollout: Sensitivity supervision slows cumulative error growth during autoregressive rollout beyond the training horizon.For FNO, final rollout error falls from 0.276 to 0.086, an approximately 3.2× reduction; similar reductions occur for WNO and DeepONet.
  • Input-dimensionality scaling: As input-field resolution increases, SC-FNO requires lower computational cost than FNO to reach the tested relative L2 target of 0.075.The authors characterize this as a finite-range empirical scaling observation rather than a universal complexity law.
  • Sensitivity ablations: State-value and Jacobian supervision are complementary: denser state data improve models, but Jacobian-supervised models remain more accurate across tested state resolutions.Jacobian supervision constrains responses to distributed-input perturbations rather than only solution values.
  • Sensitivity ablations: Partial Jacobian sampling substantially reduces error before improvements begin to saturate at denser coverage, so full Jacobian enforcement is unnecessary for benefit.The ablation varies sampled final-time Jacobian locations from 0 × 0 through 1 × 1, intermediate subsets, and the full 64 × 64 Jacobian.

6 Near-Real-Time Inverse–Forward Inference for Tsunami Forecasting

The Tohoku tsunami case tests sensitivity-constrained neural operators for sparse-observation source reconstruction and subsequent wave forecasting. SC-FNO improves forward prediction, inverse recovery, forecast reconstruction, and robustness to tested observation noise, with a near-real-time inference cost.

  • Case setup: The benchmark uses earthquake-induced seafloor deformation inputs and shallow-water reference simulations over a 1000 km × 1000 km domain.The reference solver uses an unstructured triangular mesh with 308,001 cells and 154,598 nodes.
  • Forward evaluation: Sensitivity supervision lowers tsunami forward error at comparable computational cost for both in-distribution and tested OOD events.The tested OOD setting concerns an epicenter shift, not unrestricted extrapolation to arbitrary rupture geometries.
  • Forward evaluation: SC-FNO remains closer to reference wave fields and gauge signals, while baseline FNO shows dissipative drift and errors in later-arrival timing and amplitude.The comparison spans reported time snapshots and a marked gauge location.
  • Source inversion: Using sparse gauge histories from the first 30 min, SC-FNO reconstructs the deformation’s main location, polarity, and large-scale geometry more accurately than FNO.The inverse framework uses pretrained neural operators as differentiable surrogates for gradient-based source inference.
  • Forecast reconstruction and robustness: SC-FNO produces more accurate post-observation gauge histories and maintains more informative source reconstructions than FNO across tested low-to-moderate noise levels.The full three-stage inversion and forward reconstruction require approximately 4.5 minutes on a single NVIDIA A100 GPU, excluding simulation generation, adjoint computation, and training.
  • Robustness limits: The noise experiment demonstrates robustness to tested synthetic observation perturbations but is not a full uncertainty-quantification study.An operational setting would additionally require uncertainty in gauge availability, rupture geometry, bathymetry, boundary conditions, and nearshore physics.

7 Discussion

Sensitivity supervision complements state supervision by improving forward prediction, inverse reconstruction, and observed accuracy–cost tradeoffs for high-dimensional PDE surrogates. Its practical value is strongest for repeated or inverse use, while the evidence remains empirical and subject to solver, cost, and deployment constraints.

  • What Sensitivity Supervision Changes: Sensitivity supervision complements state supervision: forward gains narrow with more state data, while Jacobian information constrains responses to distributed-input perturbations.The ablation separates solution-value supervision from sensitivity supervision rather than attributing gains to additional state values alone.
  • What Sensitivity Supervision Changes: Sampled Jacobian supervision improves the observed accuracy–cost tradeoff for the tested high-dimensional gridded-input PDE tasks, without establishing a universal scaling law.The paper characterizes the result as empirical and computational rather than a formal resolution of the curse of dimensionality.
  • Limitations and Practical Requirements: Applying SC-NO to new PDE systems requires solver-derived sensitivities, while sampling density and net benefit depend on solver, adjoint cost, output dimension, and downstream task.Sensitivity generation, storage or access, and model-side differentiation add non-negligible computational cost.
  • Scope of Sensitivity-Constrained Operator Learning: SC-NO targets operator-level sensitivities for distributed spatial fields by sampling solver-derived Jacobian entries and resampling them across minibatches.The method augments solution-field matching rather than introducing a new neural-operator architecture.
  • Why Sensitivity Supervision Matters More for Inversion: Across controlled advection–diffusion and RANS benchmarks, sensitivity supervision improves forward prediction and produces larger gains in gradient-based inverse reconstruction.Inverse reconstruction errors show greater separation between sensitivity-constrained and state-only models than corresponding forward errors.
  • Applications and Practical Boundary: The tsunami benchmark is a near-real-time inverse–forward proof of concept that reconstructs gridded seafloor deformation from sparse early gauges and forecasts later waves.Operational warning deployment would require broader uncertainty treatment, sensor robustness, and validation across additional events.

D. PDE2: Turbulent Navier–Stokes equations with Spalart–Allmaras closure

The PDE2 setup combines neural-operator learning with adjoint-based sensitivities for turbulent Navier–Stokes equations using Spalart–Allmaras closure. The workflow also includes an Okada deformation model, fine-to-coarse grid mapping, and three-stage source inversion.

  • D. PDE2: Turbulent Navier–Stokes equations with Spalart–Allmaras closure: Neural-operator learning is used as the core PDE2 modeling setup.
  • D. PDE2: Turbulent Navier–Stokes equations with Spalart–Allmaras closure: Adjoint-based sensitivities provide the sensitivity component of the PDE2 workflow.
  • D. PDE2: Turbulent Navier–Stokes equations with Spalart–Allmaras closure: The workflow includes an Okada deformation model and fine-to-coarse grid mapping.
  • D. PDE2: Turbulent Navier–Stokes equations with Spalart–Allmaras closure: Source reconstruction is organized as a three-stage source inversion procedure.

H. Computational cost analysis

Neural operators learn PDE solution maps on function spaces, but standard state-only training does not constrain parameter sensitivities. SC-NOs add sampled Jacobian supervision to improve generalization, stability, and inversion under stated assumptions.

  • Neural operators approximate mappings from input functions to solution functions and can transfer solutions across resolutions.
  • Standard solution-based losses do not explicitly account for how solutions change with respect to parameters, which matters for generalization and inverse problems.
  • Sensitivity training is formulated as L(θ) = Lu(θ) + λLs(θ), with λ controlling sensitivity enforcement.
  • Under uniform-convergence assumptions, reducing sensitivity error constrains the learned Lipschitz constant and tightens generalization bounds on unseen parameters.
  • For inversion, minimizing both solution and sensitivity errors yields convergence to an O(εu + εs) neighborhood of the true optimum under local strong convexity.

B Data Generation

Training data pair spatially varying inputs with high-fidelity PDE solutions and Jacobians. Final-time Jacobians capture accumulated spatial-parametric effects while reducing sensitivity-supervision cost.

  • Data Generation: Each sample contains an input function, its corresponding PDE solution, and solver-computed sensitivities for sensitivity-aware training.
  • Data Generation: The input combines initial solution time steps with a spatially varying parameter field p(x).
  • Data Generation: The neural operator receives tu initial time steps and predicts the additional tG = nt − tu steps through the final simulation time.
  • Data Generation: For spatial parameters, the Jacobian is matrix-valued and is evaluated at the final time step t = T.
  • Data Generation: Each Jacobian entry quantifies the final solution response to a parameter change, including local and nonlocal PDE-induced dependencies.
  • Data Generation: Final-time sensitivity supervision reduces computational cost while preserving critical spatial-parametric relationships for optimization.
  • Data Generation: Random spatial parameters are generated as bounded, spatially coherent Gaussian random fields using a tanh-based transformation.

C PDE1: Advection-Diffusion Equation with Spatially Distributed Velocity Field

The experiments cover advection–diffusion, turbulent flow, and tsunami source reconstruction. They use differentiable or adjoint sensitivity computation to assess forward accuracy, inversion, robustness, and physical reconstruction constraints.

  • C PDE1: Advection-Diffusion Equation with Spatially Distributed Velocity Field: The advection–diffusion benchmark models concentration transport with spatially varying velocity, constant diffusion D = 0.005, and periodic boundaries.
  • C PDE1: Advection-Diffusion Equation with Spatially Distributed Velocity Field: A differentiable solver uses finite differences on a 150 × 150 grid, downsampled to 50 × 50 divisions, with central-difference diffusion and upwind advection.
  • C PDE1: Advection-Diffusion Equation with Spatially Distributed Velocity Field: The SC-NO maps initial solution steps and velocity fields to subsequent concentration states while supporting velocity sensitivities.
  • C PDE1: Advection-Diffusion Equation with Spatially Distributed Velocity Field: The turbulent-flow case uses a vorticity–stream function formulation with Spalart–Allmaras turbulence modeling and Reynolds numbers from 500–1000.
  • C PDE1: Advection-Diffusion Equation with Spatially Distributed Velocity Field: Continuous adjoints compute final-vorticity sensitivities with respect to initial vorticity and forcing, with gradient cost independent of input-parameter count.
  • C PDE1: Advection-Diffusion Equation with Spatially Distributed Velocity Field: Physically informed bathtub upscaling is needed near shorelines because interpolation fails to enforce depth non-negativity and shoreline emergence constraints.
  • C PDE1: Advection-Diffusion Equation with Spatially Distributed Velocity Field: In the tsunami inversion, early gauge observations guide reconstruction of earthquake-induced seafloor deformation for near-real-time forecasting.
  • C PDE1: Advection-Diffusion Equation with Spatially Distributed Velocity Field: SC-FNO refinement converges toward the true Okada deformation, recovering rupture location and polarity with minimal spatial artifacts, whereas baseline FNO refinement fails to become meaningful.

F Performance Metrics

The paper evaluates continuous-field errors with relative L2 error and MAE, and uses R2 plus wet–dry classification metrics for inverse and tsunami-impact tasks.

  • F Performance Metrics: Relative L2 error measures normalized discrepancy between predicted and reference solution fields over evaluated spatial-temporal degrees of freedom.
  • F Performance Metrics: MAE reports average absolute error by normalizing absolute discrepancies by each sample’s evaluated spatial-temporal degrees of freedom.
  • F Performance Metrics: R2 is reported for reconstructed input fields, predicted solution trajectories, or tsunami source reconstructions depending on the experiment.
  • F Performance Metrics: Wet–dry classification labels cells dry below h < 0.01 m and wet at h ≥ 0.01 m for tsunami-impact evaluation.
  • F Performance Metrics: Inundation accuracy uses TP, TN, FP, and FN counts as Accuracy = (TP + TN) / (TP + TN + FP + FN).

G Hyperparameters and Settings

The section specifies shared training settings and points to model-specific hyperparameter tables for FNO, WNO, and DeepONet.

  • All models use Adam with batch size 16, initial learning rate 0.001, and a scheduler decaying by γ = 0.95 every 100 epochs.
  • Training spans an nx × ny spatial grid, with nt = tu + tG combining input and predicted time steps.
  • FNO, WNO, and DeepONet hyperparameters are provided in separate tables covering all PDEs.
  • A settings table organizes branch and trunk network layers, training epochs, grid dimensions, and temporal-step parameters.

H Computational Cost Analysis

The computational-cost analysis compares data-generation and training expenses across PDE1–PDE3, including preprocessing and Jacobian computation where applicable.

  • Cost reporting covers average wall-clock time per epoch, total training time over 500 epochs, and preprocessing overhead.
  • All experiments ran on a single NVIDIA A100 GPU with 32 GB memory, and every reported experiment fit within that capacity.
  • Separate tables provide total computational-cost breakdowns for PDE1, PDE2, and PDE3 over 500 epochs.

CRediT authorship contribution statement

The CRediT statement assigns conceptual, methodological, supervisory, investigative, software, writing, and project-support contributions across the listed authors.

  • Abdolmehdi Behroozi contributed to conceptualization, methodology, software, analysis, investigation, data curation, visualization, and the original draft.
  • Chaopeng Shen contributed conceptualization, methodology, supervision, project administration, funding acquisition, and manuscript revision.
  • Daniel Kifer contributed methodology, supervision, and manuscript revision.
  • Kathryn Lawson contributed methodology, investigation, and manuscript revision.
Loading 2608.29888v1…