PaperlayerView exact revision on arXiv

Source-linked AI summary

Derivative-Informed Training of Neural Operators On-the-Fly via Sketched Tangent Consistency

Xinhan Yang, Lu Lu, Shancong Mou

arXiv:2609.32797v1cs.LG

TL;DR

Neural-operator Jacobians matter in downstream optimization and control, but derivative-informed methods rely on costly offline tangent labels. The paper introduces sketched tangent consistency with operator-aware loss conditioning, and reports solution and Jacobian accuracy comparable to offline DIFNO across five PDE benchmarks while avoiding offline derivative-data generation and storage.

  • Problem

    Neural-operator Jacobians matter for downstream gradients, while the paper asks whether derivative-informed training can avoid offline tangent solves, stored labels, and architecture changes.

    Source: §1 · pp. 1–2

  • Method

    sTCL samples input perturbation directions to enforce forward-sensitivity consistency on the fly, with loss conditioning chosen for the tangent operator and known differentiable PDE residual.

    Source: §3.1 · p. 4, §3.3 · p. 6, §3.4 · p. 6

  • Results

    Across five PDE benchmarks, sTCL achieves solution and Jacobian accuracy comparable to DIFNO while eliminating offline tangent-data generation and derivative-label storage.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    For PDEs with known differentiable residuals, sketching and PDE-aware loss design provide an effective on-the-fly alternative to offline derivative supervision.

    Source: §1 · pp. 2–3

    sTCL requires a known PDE with a differentiable residual and accessible tangent actions, excluding purely black-box and unknown-physics settings.

    Source: §3.1 · p. 4, §5 · p. 11

Abstract

from arXiv · show

Derivative-informed training improves neural operators by directly supervising their input-output sensitivities, which is crucial when neural operators are used as differentiable surrogates for inverse problems, PDE-constrained optimization, design, and control. However, existing methods rely on offline-generated derivative labels, making data generation slow, storage-intensive, and difficult to adapt across datasets, resolutions, or perturbation bases. We propose sketched tangent consistency loss (sTCL), an on-the-fly derivative-informed training objective that, for PDEs with a known and differentiable residual, enforces sensitivity consistency directly from the governing equation without offline tangent labels or neural-operator architecture changes. sTCL uses randomly sketched input perturbations to provide a lightweight derivative-level physics constraint during training. However, raw forward-sensitivity residual penalties can fail for stiff, ill-conditioned, indefinite, or coupled saddle-point tangent operators. To address this, we introduce lightweight operator-aware loss-conditioning mechanisms selected by a simple tangent-operator decision rule. Across Helmholtz, nonlinear diffusion-reaction, Burgers, Allen-Cahn, and Navier-Stokes, with the same neural-operator backbone for all methods, the PDE-specific sTCL losses achieve solution and Jacobian accuracy comparable to offline derivative-informed training (DIFNO) while eliminating the offline derivative-data generation and storage pipeline. These results show that on-the-fly derivative-informed training need not merely amortize offline tangent-solve cost into training; with appropriate sketching and loss design, sTCL provides an effective drop-in path to derivative-informed neural operators. Code is available at https://github.com/yang9579/Derivative-informed-traning-on-the-fly.

1 Introduction

The introduction motivates derivative supervision for neural operators used as differentiable surrogates, then asks whether it can be made on-the-fly without offline tangent solves or architecture changes. It presents sTCL and operator-aware conditioning, reporting accuracy comparable to offline DIFNO across five PDE benchmarks.

1 Introduction

  • Accurate neural-operator Jacobians matter because downstream objective gradients depend on them, so solution accuracy alone may not yield reliable optimization or control directions.Source: §1 · p. 1
  • Offline derivative-informed methods require tangent solves and stored derivative tensors, creating a costly pipeline that must be regenerated when data, resolution, or perturbation bases change.For moderate nonlinear-PDE datasets, the offline stage can take hours to days and tens of GB per training set.Source: §1 · p. 2
  • The paper asks whether known PDE equations can provide on-the-fly derivative supervision without offline tangent solves, stored labels, architecture changes, or full-Jacobian materialization.Source: §1 · p. 2
  • sTCL sketches random input perturbations, computes surrogate Jacobian–vector products, and enforces a PDE-derived forward-sensitivity consistency loss during training.The derivative penalty uses a few Jacobian–vector products and tangent-operator applications, with per-step cost of the same order as data fitting.Source: §1 · p. 2
  • Raw forward-sensitivity residual penalties can give poor optimization signals when the tangent operator is stiff, ill-conditioned, or indefinite, motivating operator-aware loss conditioning.The paper uses DST-based inverse-Laplacian loss for symmetric positive-definite elliptic operators and short-iteration MINRES target matching for symmetric indefinite operators.Source: §1 · p. 2
  • Across Helmholtz, nonlinear diffusion–reaction, Burgers, Allen–Cahn, and steady Navier–Stokes, sTCL matches DIFNO’s solution and Jacobian accuracy while avoiding offline tangent-data generation and storage.The comparison uses the same FNO backbone, data, and training protocol; training time is comparable even before counting DIFNO’s offline label-generation stage.Source: §1 · pp. 2–3

2 Related Work

Related work connects sTCL to neural operators, derivative-informed training, physics-informed residual losses, and randomized sketching. The paper combines derivative supervision with governing-equation residuals and argues that tangent-operator conditioning determines whether the resulting signal is useful.

Neural operators and physics-informed operator learning

  • Neural operators learn PDE solution maps, while physics-informed operator methods enforce the PDE residual on predicted solution values rather than on solution-operator derivatives.Source: §2 · p. 3

Sobolev and derivative-informed training

  • Sobolev-style neural-operator methods such as DINO, DIFNO [9], and sensitivity-constrained FNO [10] supervise Jacobian–vector products using offline-generated tangent labels.sTCL retains derivative supervision but enforces the forward sensitivity equation during training instead of generating and storing labels beforehand.Source: §2 · p. 3

Physics-informed residual losses and conditioning

  • Raw physics-informed residual losses can be difficult to optimize, and this conditioning problem also arises for forward-sensitivity residuals governed by the tangent operator.Source: §2 · p. 3

Randomized sketching and Jacobian-vector products

  • sTCL uses randomized directional sketching and forward-mode Jacobian–vector products instead of materializing the full surrogate Jacobian.Its physics-informed interpretation samples input-space perturbation directions and enforces the forward sensitivity equation in expectation.Source: §2 · pp. 3–4

Position of this work

  • The paper combines governing-equation enforcement with derivative supervision, and identifies conditioning—not sketching alone—as central to whether the on-the-fly derivative signal is useful.Source: §2 · p. 4

3 Methodology

The methodology uses the forward sensitivity equation as a derivative-label-free consistency condition, estimates its residual through random directional sketches, and conditions the resulting loss according to tangent-operator structure.

3 Methodology

  • sTCL uses the same forward sensitivity equation as offline derivative-informed methods, but enforces it as a stochastic training constraint rather than solving it to produce labels beforehand.Source: §3 · p. 4

3.1 Assumptions

  • sTCL assumes a known governing PDE with a differentiable residual and accessible state and input tangent actions; black-box simulators and unknown-physics systems are outside its scope.Source: §3.1 · p. 4

3.2 Neural operators and forward sensitivity equation

  • For a perturbation direction v, the forward sensitivity w = DS(f)[v] satisfies the linearized PDE residual equation, with A and B denoting the residual’s state and input derivatives.The surrogate Jacobian is the target of derivative-informed training because downstream gradients depend on it.Source: §3.2 · pp. 4–5

3.3 From offline Sobolev supervision to on-the-fly sketched tangent consistency

  • Instead of precomputing and storing tangent labels, sTCL samples input directions and enforces the forward sensitivity equation in expectation, avoiding full-Jacobian materialization.Each sketch requires a Jacobian–vector product and tangent-operator applications, keeping the derivative penalty at the same computational order as data fitting up to the sketch factor.Source: §3.3 · pp. 5–6

3.4 Conditioning and operator-aware loss design

  • The raw residual loss squares the tangent operator’s singular-value spectrum, which can impair optimization; operator-aware conditioning uses cheap preconditioning for suitable elliptic operators and short MINRES target matching for indefinite or near-singular cases.The MINRES target uses a small fixed iteration count as a conditioning aid, not as an accurate tangent solve.Source: §3.4 · pp. 6–7

4 Experiments

Experiments compare sTCL with offline DI and pure FNO across five PDE benchmarks, using shared FNO backbones and matched training protocols. Results cover function-value and Jacobian accuracy, training-time costs, and a Helmholtz MINRES ablation.

4.1 On-the-fly sTCL versus offline DI on function values

  • Below 1.5× in 14/15 cells and below 1.2× in 11/15, the sTCL/offline-DI function-value relative-error ratio remained comparable.sTCL reduced pure-FNO mean function-value error in all 15 cells; training costs were of the same order, while offline-DI timings excluded tangent-data generation.Source: §4.1 · p. 8, Table 2 · p. 9, Table 3 · p. 9

4.2 On-the-fly sTCL versus offline DI on Jacobians

  • 73.53 ± 0.60% versus 178.08±15.45% on Burgers N=512 (sTCL versus offline DI); overall, Jacobian accuracy was comparable, with PDE-dependent differences.sTCL did not uniformly beat offline DI, but recovered much of its improvement over pure FNO without an offline tangent dataset.Source: §4.2 · pp. 9–10, Table 4 · p. 10

4.3 Preconditioning ablation

  • 3 MINRES iterations reduced Helmholtz N=512 function-value error from 3.50% to 2.52% and Jacobian error from 23.03% to 15.52%.Beyond m=5, both errors changed only slightly relative to seed variability while training time increased; m=5 was a practical cost–accuracy balance. The main comparisons matched loss conditioning to tangent-operator class.Source: Table 5 · p. 10, §4.3 · p. 10

5 Conclusion

Across five PDE benchmarks, sTCL achieves solution and Jacobian accuracy comparable to offline derivative-informed training while eliminating the offline derivative-data pipeline. The conclusion identifies sketching and PDE-aware loss design as essential and notes scope limitations.

5 Conclusion

  • Across five PDE benchmarks, sTCL achieves solution and Jacobian accuracy comparable to offline derivative-informed training under the same backbone, without uniformly outperforming it.The benchmarks span parabolic/advective, phase-field parabolic, SPD elliptic, symmetric indefinite, and coupled incompressible saddle-point regimes.Source: §5 · p. 11
  • sTCL sketches perturbation directions and applies PDE-specific forward-sensitivity consistency losses, avoiding precomputed tangent labels and their regeneration cost.The study identifies sketching and PDE-aware loss design as the two ingredients needed across tangent-operator classes.Source: §5 · p. 11
  • sTCL requires a known PDE with a differentiable residual and access to its linearization or tangent actions; black-box and unknown-physics settings are outside scope.In steady Navier–Stokes, evaluating additional momentum, incompressibility, and boundary tangent terms raises per-update cost relative to PDEs with fewer residual components.Source: §5 · p. 11

A.1 Experiment protocol and benchmark details

The experiments use a controlled protocol with shared model and training settings, independent data splits, and held-out Jacobian evaluation. Benchmark-specific training budgets and probe-direction banks vary by PDE.

A.1 Experiment protocol and benchmark details

  • FNO runs use four Fourier layers, width 32, GELU, Adam at learning rate 10^-3, batch size 32, and sketch count q=4, with fixed seeds and gradient clipping.Training lasts 2000 epochs for Helmholtz, nonlinear diffusion, and Burgers; 1500 for Allen–Cahn; and 5000 for steady Navier–Stokes.Source: §A.1 · p. 13
  • Each benchmark has independent training, validation, and test splits, with Nval=128 and Ntest=128; Jacobian evaluation uses independently generated probe directions.The direction-bank sizes are 8 for Helmholtz, 10 for nonlinear diffusion, 200 for Burgers, 200 Matérn directions for Allen–Cahn, and 16 solenoidal directions for steady Navier–Stokes.Source: §A.1 · p. 13
  • Jacobian errors are evaluated post hoc on held-out tests using the same direction bank within each PDE, and the test JVP banks are used only for evaluation.Each benchmark uses its best-epoch checkpoint for both function-value and Jacobian results; the reference JVP pipeline is benchmark-specific.Source: §A.1 · p. 13

A.1.1 Nonlinear diffusion–reaction.

For nonlinear diffusion–reaction, the tangent operator is symmetric positive definite, but stiffness makes the raw residual loss difficult to condition. The reported DST-preconditioned loss improves both function and Jacobian errors over pure FNO at N=1024.

A.1.1 Nonlinear diffusion–reaction

  • The benchmark uses a 65 × 65 uniform grid, with reference states computed on the 63 × 63 interior grid using centered finite differences and Newton iteration.The source is the sum of four Gaussian bumps, and the parameter field uses the 800 largest-variance mode pairs.Source: §A.1.1 · p. 13
  • The state tangent operator is symmetric positive definite, and the conditioning challenge is stiffness rather than indefiniteness.Its eigenvalues grow like those of the discrete Laplacian, and the raw residual loss squares the elliptic spectrum, over-weighting high-frequency components while slowing low-frequency sensitivity learning.Source: §A.1.1 · p. 14
  • For a unit-normalized probe drawn from the same 800-mode cosine KLE as the inputs, the method forms a discrete tangent residual at the surrogate state.The surrogate state is detached when used as a coefficient in the residual.Source: §A.1.1 · p. 14
  • At N=1024, the DST-preconditioned loss lowers mean per-sample function error from 2.98% for pure FNO to 1.65%, and Jacobian error from 10.88% to 6.07%.The inverse Dirichlet Laplacian is applied exactly in the DST-I basis, changing residual geometry without solving the tangent equation.Source: §A.1.1 · p. 14

A.1.2 Helmholtz.

For Helmholtz, the symmetric tangent operator is generally indefinite, so the method matches surrogate JVPs to approximate tangent solutions rather than relying on raw residual minimization. A five-iteration preconditioned MINRES target is reported as a lower-cost choice within seed variability of longer runs.

A.1.2 Helmholtz

  • The Helmholtz benchmark uses a 65 × 65 Dirichlet grid and filters out the most severe near-resonant response outliers from the sampled data.The filter does not make the tangent operator positive definite.Source: §A.1.2 · p. 14
  • For eigenvalues close to zero, residual-loss curvature scales with the squared eigenvalue, so raw residual minimization can miss a physically important sensitivity mode despite a small PDE residual.A positive residual preconditioner cannot remove the squared near-zero factor without acting like an inverse on the resonant subspace.Source: §A.1.2 · p. 14
  • Because Helmholtz’s symmetric tangent operator is indefinite, sTCL uses MINRES target matching rather than a raw residual penalty near resonant eigenvalues.A shifted-Laplacian preconditioner is applied exactly by DST-I; the approximate tangent solution supplies the target for the surrogate JVP.Source: §A.1.2 · p. 14
  • Five preconditioned MINRES iterations are within seed variability of 10 and 25 iterations on both metrics, at substantially lower wall-clock cost.The approximate targets are computed for the current mini-batch and discarded, so tangent labels are not stored.Source: §A.1.2 · p. 14

A.1.3 Burgers.

For Burgers, sTCL evaluates tangent consistency along Gaussian time-grid probes using an RHS-normalized residual. The loss keeps the causal space–time and advective terms together, without a spatial preconditioner.

A.1.3 Burgers.

  • The loss disables spatial preconditioning because the tangent operator is a space–time initial-value operator with time-derivative and transport terms.Its strong-form residual keeps the causal parabolic and advective terms together.Source: §A.1.3 · p. 15
  • For each Gaussian time-grid probe, sTCL computes the surrogate Jacobian–vector product and evaluates its tangent residual relative to the corresponding right-hand-side energy.The residual is evaluated on interior spatial nodes and noninitial frames.Source: §A.1.3 · p. 15
  • At N=512, sTCL and offline DI have similar function errors, while their Jacobian errors are 73.53% and 178.08%, respectively.The sTCL Jacobian error is lower in this comparison.Source: §A.1.3 · p. 15

A.1.4 Allen–Cahn.

The Allen–Cahn and steady Navier–Stokes sections pair benchmark-specific tangent-loss designs with empirical checks, followed by the per-cell hyperparameter settings used for reported comparisons.

A.1.4 Allen–Cahn.

  • Allen–Cahn uses a raw space–time L2 tangent residual, retaining the temporal, diffusive, and state-dependent reaction terms together without additional preconditioning.The tangent equation is a forward parabolic initial-value problem.Source: §A.1.4 · p. 16
  • At N=512, Allen–Cahn sTCL lowers mean field error from 0.94% to 0.59% and Jacobian error from 15.64% to 3.82%, close to offline-DI's 3.64%.The neural operator's hard initial-condition ansatz gives exact state and sensitivity values at t=0.Source: §A.1.4 · p. 16

A.1.5 Steady Navier–Stokes.

  • The Navier–Stokes momentum loss projects out the Fourier gradient component and weights by an inverse Oseen-symbol magnitude; separate continuity and boundary terms remain essential.The periodic projector is only approximate on the bounded Dirichlet domain.Source: §A.1.5 · pp. 16–17
  • At N=512, sTCL reduces Navier–Stokes velocity error from 0.619% to 0.400% and velocity-JVP error from 1.084% to 0.570%, without tangent labels.At N=2048, the corresponding reported errors are 0.198% and 0.282%.Source: §A.1.5 · p. 17

A.2 Per-cell hyperparameter settings

  • All derivative-informed cells use EMA-normalized weighting, while sTCL uses an operator-specific loss form for each PDE benchmark.The settings do not include a broader joint search over learning rate and λ.Source: §A.2 · p. 17

A.3 Sketch count, resolution, and gradient variance

The sketch count trades computation against gradient-estimator variance, but increasing it is not needed to preserve accuracy across the tested resolutions. A single direction improves field and JVP errors over solution-only training, while resolution-related error growth persists.

Accuracy versus resolution

  • q = 1 reduces field error by 50.7%, 49.3%, and 36.9%, and JVP error by 44.8%, 43.5%, and 32.1% versus q = 0 on the 322, 642, and 1282 grids, respectively.Among q ≥1, neither metric improves monotonically with q; the lowest mean errors occur at q = 4 on 322 and q = 16 on the larger grids.Source: §A.3 · p. 18, Table 7 · p. 18
  • Both field and JVP errors grow moderately with resolution for every q ≥1, and increasing q from 1 to 32 does not remove that growth.The sweeps do not isolate resolution from factors including model capacity, training-set size, and training budget.Source: §A.3 · p. 18

Runtime

  • On the 642 grid, training wall time grows approximately linearly with q in the controlled N=128 sweep.Table 8 reports single-run profiling values, not across-seed averages.Source: §A.3 · p. 18, Table 8 · p. 18

Gradient variance

  • From q = 1 to q = 32, mean V weighted falls by 94.3%, 92.2%, and 93.1% on the 322, 642, and 1282 grids, but remains 0.7–14.5% of V total across tested settings.At q = 4, V weighted is 3.3–4.6% of V total; because the expected objective is q-independent, greater q need not improve final accuracy.Source: §A.3 · p. 19

Sketch count and resolution

  • A fixed small q remained effective across a 16-fold increase in grid dimension; q = 4 is a conservative empirical operating point, not a dimension-independent guarantee.Inputs and probes share an 800-mode cosine KLE, so the tested input manifold has dimension at most 800 even as grid dimension rises from 1024 to 16,384.Source: §A.3 · pp. 19–20

A.4 Long-run convergence check

Long-run validation trajectories stabilize within the tested budgets without changing the qualitative ordering of pure FNO, offline DI, and sTCL. The discussion recommends selecting the sensitivity loss by tangent-operator structure and weighing offline cost against the on-the-fly approach.

A.4 Long-run convergence check

  • Validation curves for pure FNO, offline DI, and sTCL stabilize within the long-run budgets, with small late fluctuations and no qualitative change in method ordering.Figure 1 shows trajectories for Burgers, Helmholtz, nonlinear diffusion, Allen–Cahn, and steady Navier–Stokes.Source: §A.4 · p. 20, Figure 1 · p. 21

B Discussion and practical recommendations

  • The recommended loss depends on tangent-operator structure: space–time residuals for the tested parabolic problems, DST inverse-Laplacian for SPD elliptic problems, MINRES-m for symmetric indefinite problems, and Leray–Oseen-preconditioned residuals for steady NS.For small problems with cheap offline tangent generation, offline DI and sketched sensitivity are competitive; sketched sensitivity is attractive when offline cost dominates.Source: B Discussion and practical… · p. 20
  • sTCL’s derivative-label-free advantage is greatest when new data would require rerunning offline tangents, disk storage is constrained, sketch directions may change, or forward solves are expensive.Source: B Discussion and practical… · p. 20
  • Five MINRES iterations remain adequate in the reported Helmholtz setting: m = 5 is within seed variance of m = 25 at substantially lower per-step cost.Increasing the inner count from m = 1 to m = 3 gives the dominant improvement; the paper argues that on-the-fly solves need only support useful descent directions at the surrogate’s current accuracy.Source: B Discussion and practical… · p. 20
Loading 2609.32797v1…