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

Xinhan Yang; Lu Lu; Shancong Mou

cs.LGSep 26, 2026

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

Authors: Xinhan Yang, Lu Lu, Shancong Mou

Organizations: Department of Industrial and Systems Engineering, University of Minnesota Minneapolis, MN 55455 · Department of Statistics and Data Science, Yale University New Haven, CT 06511, USA

Abstract

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.

Figures & tables

Benchmark	Equation / map	Tangent class	Tangent feature	sTCL form
Nonlinear diffusion	$-\nabla\!\cdot(e^{{\color[rgb]{0,0,1}a}}\nabla{\color[rgb]{1,0,0}u})+{\color[rgb]{1,0,0}u}^{3}=f$	SPD elliptic	Stiff elliptic spectrum	DST inverse-Laplacian residual
Helmholtz	$(-\Delta-\kappa^{2}e^{2{\color[rgb]{0,0,1}a}}){\color[rgb]{1,0,0}u}=f$	Indefinite elliptic	Near-resonant modes	MINRES target matching
Burgers	${\color[rgb]{1,0,0}u}_{t}+{\color[rgb]{1,0,0}u}\,{\color[rgb]{1,0,0}u}_{x}=\nu{\color[rgb]{1,0,0}u}_{xx}+{\color[rgb]{0,0,1}f(t)}$	Parabolic, non-self-adjoint	Advective dynamics	RHS-normalized space–time residual
Allen–Cahn	$\begin{gathered}{\color[rgb]{1,0,0}u}_{t}=\varepsilon^{2}{\color[rgb]{1,0,0}u}_{xx}+{\color[rgb]{1,0,0}u}-{\color[rgb]{1,0,0}u}^{3},\\ {\color[rgb]{1,0,0}u}(x,0)={\color[rgb]{0,0,1}a(x)}\end{gathered}$	Parabolic, phase-field	Interface-sensitive coarsening	Raw space–time $L^{2}$ residual
Steady NS	$\begin{gathered}-\mu\Delta{\color[rgb]{1,0,0}y}+({\color[rgb]{1,0,0}y}\!\cdot\!\nabla){\color[rgb]{1,0,0}y}+\nabla{\color[rgb]{1,0,0}p}={\color[rgb]{0,0,1}u},\\ \nabla\!\cdot\!{\color[rgb]{1,0,0}y}=0\end{gathered}$	Non-self-adjoint saddle point	Pressure coupling and advection	Leray–Oseen-preconditioned residual

Table 1: Benchmarks, tangent-operator classes, and sTCL forms used in the main experiments. The table combines the loss-design rule with the experimental suite; boundary conditions, data distributions, solvers, tangent equations, and implementation details are in section A.1 . In the equation/map column, blue marks the operator input and red marks the response.

PDE	$N$	Pure FNO	sTCL	Offline DI
Helmholtz	128	$\phantom{0}7.95\pm 0.54$	$\phantom{0}7.48\pm 0.54$	$\phantom{0}7.13\pm 0.90$
	512	$\phantom{0}3.92\pm 0.46$	$\phantom{0}2.25\pm 0.31$	$\phantom{0}3.08\pm 1.37$
	1024	$\phantom{0}2.09\pm 0.25$	$\phantom{0}1.51\pm 0.17$	$\phantom{0}1.83\pm 0.75$
Nonlin. diffusion	128	$\phantom{0}9.71\pm 0.15$	$\phantom{0}4.77\pm 0.17$	$\phantom{0}3.43\pm 0.16$
	512	$\phantom{0}4.49\pm 0.16$	$\phantom{0}2.26\pm 0.07$	$\phantom{0}1.52\pm 0.03$
	1024	$\phantom{0}2.98\pm 0.05$	$\phantom{0}1.65\pm 0.03$	$\phantom{0}1.07\pm 0.02$

Table 2: Function-value test relative error (%), reported as mean

\pm

standard deviation over three seeds. All cells use Adam with lr

=10^{-3}

and batch size

B=32

; derivative-informed cells use sketch count

q=4

. Helmholtz, nonlinear diffusion, and Burgers are trained for

2000

epochs; Allen–Cahn is trained for

1500

epochs; and steady NS is trained for

5000

epochs. All derivative-informed cells use the same EMA-normalized weighting. PDE-level validation grids select

\lambda=1

for Burgers and

\lambda=2

for Allen–Cahn; Helmholtz, nonlinear diffusion, and steady NS use the per-cell validation-selected values reported in table 6 . The steady NS rows report velocity error. Within each PDE, sTCL and offline DI use matched probe-direction distributions and the same per-update direction count, so the comparison does not confound the supervision mode with different probe constructions.

PDE	Pure FNO	sTCL (on the fly)	Offline DI (training only)
Helmholtz	$192s$	$2412s$ (MINRES-5)	$2036s$
Nonlin. diffusion	$192s$	$1976s$ (DST-precond.)	$1796s$
Burgers	$408s$	$2144s$ (normalized residual)	$2208s$
Allen–Cahn	$291s$	$2295s$ (raw $L^{2}$ )	$1911s$
Steady NS	$582s$	$9990s$ (Leray–Oseen)	$4810s$

Table 3: Representative full-budget per-training-job wall-clock at

N{=}512

on a single RTX PRO 6000 Blackwell GPU; times are not averaged over training seeds. Helmholtz, nonlinear diffusion, and Burgers use

2000

epochs; Allen–Cahn uses

1500

epochs; and steady NS uses

5000

epochs. “offline-DI training” excludes offline CPU tangent-data generation.

PDE	$N$	Pure FNO	sTCL	Offline DI
Helmholtz	128	$51.64\pm 1.48$	$47.78\pm 2.41$	$45.99\pm 4.64$
	512	$26.01\pm 3.04$	$13.83\pm 2.64$	$19.52\pm 9.03$
	1024	$12.17\pm 0.92$	$\phantom{0}9.17\pm 1.51$	$10.67\pm 3.75$
Nonlin. diffusion	128	$29.74\pm 0.55$	$16.06\pm 0.62$	$11.65\pm 0.50$
	512	$15.87\pm 0.30$	$\phantom{0}8.38\pm 0.15$	$\phantom{0}5.36\pm 0.05$
	1024	$10.88\pm 0.17$	$\phantom{0}6.07\pm 0.05$	$\phantom{0}3.80\pm 0.12$

Table 4: Jacobian relative error (%), reported as mean

\pm

standard deviation over three seeds across all

15

benchmark cells. The training budgets match table 2 : Helmholtz, nonlinear diffusion, and Burgers are trained for

2000

epochs, Allen–Cahn for

1500

epochs, and steady NS for

5000

epochs. The reported metric is the per-(sample, direction) average

\bar{E}

$m$	Function rel. err.	Jacobian rel. err.	Training time (s)
$\phantom{0}1$	$\phantom{0}3.50\pm 0.45\%$	$23.03\pm 2.52\%$	$2056$
$\phantom{0}2$	$\phantom{0}2.80\pm 0.20\%$	$18.32\pm 1.48\%$	$2132$
$\phantom{0}3$	$\phantom{0}2.52\pm 0.22\%$	$15.52\pm 1.56\%$	$2228$
$\phantom{0}5$	$\phantom{0}2.47\pm 0.28\%$	$14.71\pm 2.09\%$	$2412$
$10$	$\phantom{0}2.46\pm 0.21\%$	$14.61\pm 1.65\%$	$2884$
$25$	$\phantom{0}2.41\pm 0.15\%$	$14.51\pm 1.22\%$	$4216$

Table 5: Effect of the MINRES inner-iteration count

m

on Helmholtz at

N{=}512

with a

2000

-epoch training budget and fixed

\lambda=0.1

; the main

N{=}512

comparison in tables 2 and 4 uses the validation-selected

\lambda=0.5

. Errors are reported as mean

\pm

standard deviation over three seeds.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

PDE	$\lambda$	Offline-DI bank $r$	sTCL form	Notes
Helmholtz	per- $N$ (see below)	$289$	MINRES- $5$ target-matching	validation-selected 5-pt $\lambda$ grid
$N{=}128$ : $\lambda{=}0.1$ for both methods; $N{=}512$ : $\lambda{=}0.5$ for both; $N{=}1024$ : $\lambda{=}0.9$ (offline DI), $\lambda{=}0.7$ (sTCL).
Nonlin. diffusion	$1$	$289$	DST inverse-Laplacian residual	per-cell selected from 6-pt grid
Burgers	$1$	$200$	RHS-normalized space–time residual	PDE-level selected from 5-pt grid
Allen–Cahn	$2$	$200$	raw space–time $L^{2}$ residual	PDE-level selected from 5-pt grid
Steady NS	per- $N$ (see below)	$200$	Leray–Oseen-preconditioned residual	validation-selected 4-pt $\lambda$ grid

Appendix

Table 6: Per-cell hyperparameter settings for table 2 and table 4 . All cells use Adam at

\eta=10^{-3}

, batch size

32

, and three seeds

\{0,1,2\}

. All derivative-informed cells use the common EMA-normalized weighting. Both derivative-informed methods use

q=4

directions per training example at each update, drawn from matched PDE-specific distributions; offline DI draws direction–label pairs from a stored bank of PDE-dependent size

r

, whereas sTCL uses unlabeled directions.

$q$	$32^{2}$		$64^{2}$		$128^{2}$
	Field	JVP	Field	JVP	Field	JVP
$0$	$8.94\pm 0.26$	$28.53\pm 0.53$	$9.19\pm 0.32$	$29.33\pm 0.68$	$9.33\pm 0.28$	$29.84\pm 0.62$
$1$	$4.41\pm 0.23$	$15.74\pm 0.82$	$4.66\pm 0.05$	$16.56\pm 0.45$	$5.88\pm 0.44$	$20.25\pm 1.06$
$2$	$4.33\pm 0.21$	$15.46\pm 0.86$	$4.65\pm 0.29$	$16.70\pm 0.75$	$5.66\pm 0.13$	$19.70\pm 0.26$
$4$	$4.18\pm 0.34$	$15.02\pm 1.05$	$4.65\pm 0.11$	$16.51\pm 0.39$	$5.88\pm 0.09$	$19.91\pm 0.33$
$8$	$4.27\pm 0.54$	$15.56\pm 1.74$	$4.57\pm 0.29$	$16.26\pm 0.85$	$5.72\pm 0.15$	$19.65\pm 0.51$

Appendix

Table 7: Nonlinear-diffusion accuracy versus sketch count and spatial resolution. Each entry is the mean

\pm

sample standard deviation of the relative error (%) over three training seeds

\{0,1,2\}

. For each grid, the paired columns report field and JVP errors.

Benchmark	$q=0$	$q=1$	$q=2$	$q=4$	$q=8$	$q=16$	$q=32$
Nonlin. diffusion	$1.95$	$3.69$	$5.53$	$8.46$	$15.05$	$29.23$	$54.61$

Appendix

Table 8: Single-run nonlinear-diffusion training time (minutes) in the controlled

N{=}128

sketch-count sweep on the

64^{2}

grid. These profiling values are not across-seed averages.

Grid	Dim.	$q$	Raw sTCL	Weighted sTCL	Total
$32^{2}$	$1024$	$0$	$0$	$0$	$(8.710\mathbin{\pm}1.520)\!\times\!10^{-7}$
$32^{2}$	$1024$	$1$	$(1.043\mathbin{\pm}0.297)\!\times\!10^{-2}$	$(9.832\mathbin{\pm}5.880)\!\times\!10^{-7}$	$(6.765\mathbin{\pm}3.170)\!\times\!10^{-6}$
$32^{2}$	$1024$	$2$	$(4.509\mathbin{\pm}0.311)\!\times\!10^{-3}$	$(4.327\mathbin{\pm}0.970)\!\times\!10^{-7}$	$(5.936\mathbin{\pm}1.560)\!\times\!10^{-6}$
$32^{2}$	$1024$	$4$	$(2.578\mathbin{\pm}0.429)\!\times\!10^{-3}$	$(2.627\mathbin{\pm}0.830)\!\times\!10^{-7}$	$(5.712\mathbin{\pm}2.700)\!\times\!10^{-6}$
$32^{2}$	$1024$	$8$	$(1.272\mathbin{\pm}0.125)\!\times\!10^{-3}$	$(1.202\mathbin{\pm}0.279)\!\times\!10^{-7}$	$(4.726\mathbin{\pm}2.400)\!\times\!10^{-6}$
$32^{2}$	$1024$	$16$	$(7.590\mathbin{\pm}1.300)\!\times\!10^{-4}$	$(7.163\mathbin{\pm}2.410)\!\times\!10^{-8}$	$(4.294\mathbin{\pm}1.960)\!\times\!10^{-6}$

Appendix

Table 9: Gradient trace variance

\overline{V}_{a}

for nonlinear diffusion, averaged over epochs

26

–

2000

and reported as mean

\pm

sample standard deviation over three training seeds. The Raw sTCL, Weighted sTCL, and Total columns report

\overline{V}_{\mathrm{raw}}

\overline{V}_{\mathrm{weighted}}

, and

\overline{V}_{\mathrm{total}}

, respectively.

Panel of Figure 1 — Figure 1: Single-seed training and validation trajectories for representative benchmark settings: $4000$ epochs for Burgers, Helmholtz, nonlinear diffusion, and Allen–Cahn, and $5000$ epochs for steady NS. Panels (a)–(c) and (e) show the MSE, and panel (d) shows the relative error. Dashed curves denote training and solid curves denote validation for pure FNO, sTCL, and offline DI; these curves are convergence diagnostics rather than across-seed uncertainty summaries. The dotted vertical line in each panel marks epoch $500$ .

Neural Operators

Neural operators learn mappings between functions, enabling fast surrogates for PDEs and other continuous physical systems.

See all 313 papers

Sensitivity-Constrained Neural Operators for Data-Efficient Forward and Inverse Modeling of Partial Differential Equation Systems
Aug 30, 2026Abdolmehdi Behroozi, Chaopeng Shen, Daniel Kifer +1PDE Surrogate Modeling Jacobian Regularization
Department of Civil and Environmental Engineering, Penn State University, University Park, PA 16802, USA · School of Electrical Engineering and Computer Science, Penn State University, University Park, PA 16802, USA
Preconditioned Physics-Informed Neural Operator Training
Sep 28, 2026Shizheng Wen, Siddhartha Mishra, Marius ZeinhoferPhysics-Informed Neural Operators PDE Operator Learning
ETH Zürich
Fractional is Better: Learnable Derivative Orders in Neural Operator Learning
Oct 14, 2026Fares B. Mehouachi, Saif Eddin JabariPDE Operator Learning Neural Operators
1New York University Abu Dhabi, Abu Dhabi, UAE · 2New York University, New York, NY, USA

Aug 30, 2026J/K move · Enter open · S save

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

Abdolmehdi Behroozi, Chaopeng Shen, Daniel Kifer, Kathryn Lawson

Department of Civil and Environmental Engineering, Penn State University, University Park, PA 16802, USA · School of Electrical Engineering and Computer Science, Penn State University, University Park, PA 16802, USA

PDE Surrogate Modeling Jacobian Regularization PDE Inverse Problems Neural Operators

Open paper AskPDF arXiv

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.

Abstract

Neural Operators

Explore similar work

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

Preconditioned Physics-Informed Neural Operator Training

Fractional is Better: Learnable Derivative Orders in Neural Operator Learning