Gradient-Informed Temporal Sampling Improves Rollout Accuracy in PDE Surrogate Training

Wenshuo Wang; Fan Zhang

cs.LGMar 18, 2026

Gradient-Informed Temporal Sampling Improves Rollout Accuracy in PDE Surrogate Training

Authors: Wenshuo Wang, Fan Zhang

Organizations: School of Future Technology, South China University of Technology, China · State Key Laboratory of Ocean Sensing & Ocean College, Zhejiang University, China · Kavli Institute for Astrophysics and Space Research, Massachusetts Institute of Technology, USA

Abstract

Researchers train neural simulators on uniformly sampled numerical simulation data. But under the same budget, does systematically sampled data provide the most effective information? A fundamental yet unformalized problem is how to sample training data for neural simulators so as to maximize rollout accuracy. Existing data sampling methods either tend to collapse into locally high-information-density regions, or preserve diversity but remain insufficiently model-specific, often leading to performance that is no better than uniform sampling. To address this, we propose a data sampling method tailored to neural simulators, Gradient-Informed Temporal Sampling (GITS). GITS jointly optimizes pilot-model local gradients and set-level temporal coverage, thereby effectively balancing model specificity and dynamical information. Compared with multiple sampling baselines, the data selected by GITS achieves lower rollout error across multiple PDE systems, model backbones and sample ratios. Furthermore, ablation studies demonstrate the necessity and complementarity of the two optimization objectives in GITS. In addition, we analyze the successful sampling patterns of GITS as well as the typical PDE systems and model backbones on which GITS fails.

Figures & tables

Dataset	Backbone	uniform	loss-only	coverage-only	GradMatch	GLISTER	PRISM	GITS
diff-react	ConvLSTM	0.480	0.504	0.240	0.472	0.549	0.435	0.191
	FNO	0.592	0.887	0.418	0.885	0.898	0.742	0.361
	Transformer	0.353	0.652	0.264	0.646	0.889	0.653	0.310
	U-Net	0.246	0.135	0.080	0.101	0.399	0.053	0.071
diff-sorp	ConvLSTM	0.229	0.712	0.381	0.410	0.386	0.369	0.026
	FNO	0.303	0.543	0.213	0.188	0.890	0.144	0.139

Table 1: Test rollout nRMSE at sampling ratio

0.05

. The best (lowest) value in each row is bolded.

Dataset	Backbone	loss-only	grad-only	loss-div	GITS
diff-react	ConvLSTM	0.504	0.891	0.474	0.191
	FNO	0.966	0.900	0.899	0.361
	Transformer	0.652	0.943	0.642	0.310
	U-Net	0.135	0.135	0.389	0.071
diff-sorp	ConvLSTM	2.017	2.017	1.101	0.026
	FNO	0.543	0.592	0.253	0.139

Table 2: Ablation at sampling ratio

0.05

Case	$K$	LO $\cap$ GO	LD $\cap$ GITS	Entropy (LO/GITS)	Coverage (LO/GITS)
Success: rdb / FNO /0.10	10	9/10	10/10	0.141 / 0.278	0.20 / 0.30
Failure A: diff-sorp / ConvLSTM /0.10	10	10/10	10/10	0.000 / 0.141	0.10 / 0.20
Failure A: diff-sorp / ConvLSTM /0.20	20	18/20	19/20	0.463 / 0.542	0.20 / 0.30
Failure B: diff-react / Transformer /0.05	5	0/5	2/5	0.000 / 0.311	0.20 / 0.40

Table 3: Geometry in the success and failure regimes of GITS. LO

\cap

GO denotes the overlap size between loss-only and grad-only , and LD

\cap

GITS denotes the overlap size between loss-div and GITS . Higher entropy and temporal coverage indicate more dispersed temporal selections.

Case	Spearman(grad, utility)	Spearman(loss, utility)
Success: rdb / FNO /0.10	0.774	0.501
Failure A: diff-sorp / ConvLSTM /0.10	-0.641	-0.654
Failure A: diff-sorp / ConvLSTM /0.20	-0.369	-0.539
Failure B: diff-react / Transformer /0.05	0.876	0.936

Table 4: Score–utility alignment in representative regimes. Entries report three-seed mean Spearman correlations between pilot scores and empirical single-start utility from the local probe updates.

Appendix figures & tables19 assets

Supplementary material from the paper’s appendix.

Appendix

Component	Symbol	Role	Final setting	Candidate values / rule	Scope / notes
Pilot start pool	$\mathcal{P}$	start indices used to train the pilot	full candidate set	not swept	fixed for all datasets/backbones/ratios
Pilot epochs	$E_{p}$	pilot training budget	$5$	$\{1,2,5,10,20\}$	swept globally on the development suite
Short rollout horizon	$H$	local rollout length for pilot scoring	$10$	$\{1,5,10,20,50,100\}$	swept globally on the development suite
Global coverage weight	$\lambda_{\mathrm{cov}}$	strength of global temporal coverage	$1.0$	$\{0.25,0.5,1.0,2.0\}$	swept jointly with $c_{\mathrm{win}}$
Sliding-window coverage weight	$c_{\mathrm{win}}$	strength of local sliding-window coverage	$0.5$	$\{0,0.25,0.5,1.0\}$	swept jointly with $\lambda_{\mathrm{cov}}$
Global kernel scale	$\tau$	decay scale in $F_{\mathrm{cov}}$	$\lfloor T_{c}/K\rfloor$	derived from budget-implied target spacing	not swept independently

Appendix

Table 5: Final GITS hyperparameters, rollout protocol, candidate ranges, and selection rules. Parameters marked as “derived” are not swept independently in order to limit tuning freedom.

Development configuration	$\mathbf{H=1}$	$\mathbf{H=5}$	$\mathbf{H=10}$	$\mathbf{H=20}$	$\mathbf{H=50}$	$\mathbf{H=100}$
diff-react / FNO /ratio $=0.10$	0.521	0.447	0.393	0.418	0.461	0.502
diff-sorp / Transformer /ratio $=0.10$	0.218	0.179	0.154	0.166	0.183	0.207
rdb / FNO /ratio $=0.10$	0.148	0.099	0.074	0.087	0.116	0.143
rdb / U-Net /ratio $=0.10$	0.243	0.197	0.158	0.172	0.196	0.224
Average development nRMSE	0.283	0.231	0.195	0.211	0.239	0.269
Mean sampling time (s)	0.31	0.82	1.39	2.54	5.84	11.02

Appendix

Table 6: Sensitivity of GITS to the short rollout horizon

H

on the fixed development suite. Lower is better. The last row reports mean subset-selection time per development configuration for the full GITS pipeline under each

H

Development configuration	$\mathbf{E_{p}=1}$	$\mathbf{E_{p}=2}$	$\mathbf{E_{p}=5}$	$\mathbf{E_{p}=10}$	$\mathbf{E_{p}=20}$
diff-react / FNO /ratio $=0.10$	0.512	0.441	0.393	0.401	0.412
diff-sorp / Transformer /ratio $=0.10$	0.219	0.181	0.154	0.158	0.163
rdb / FNO /ratio $=0.10$	0.131	0.096	0.074	0.079	0.083
rdb / U-Net /ratio $=0.10$	0.219	0.186	0.158	0.165	0.171
Average development nRMSE	0.270	0.226	0.195	0.201	0.207
Mean sampling time (s)	0.47	0.73	1.39	2.43	4.16

Appendix

Table 7: Sensitivity of GITS to the number of pilot training epochs

E_{p}

on the fixed development suite. Lower is better. The last row reports mean subset-selection time per development configuration for the full GITS pipeline under each pilot budget.

$\boldsymbol{\lambda_{\mathrm{cov}}}$	$\boldsymbol{c_{\mathrm{win}}}$	Average development nRMSE	Selection note
0.25	0.00	0.314	–
0.25	0.25	0.287	–
0.25	0.50	0.261	–
0.25	1.00	0.249	–
0.50	0.00	0.268	–
0.50	0.25	0.234	–

Appendix

Table 8: Sensitivity of GITS to the two coverage weights

(\lambda_{\mathrm{cov}},c_{\mathrm{win}})

on the fixed development suite. Lower is better. Because changing these weights has negligible effect on sampling cost, we report only average development nRMSE.

Backbone	Architectural setting
U-Net	Compact two-level encoder–decoder with double-convolution blocks, max-pooling downsampling, transposed-convolution upsampling, skip connections, and base width 32.
FNO	Width 64 and four spectral blocks, with 16 retained Fourier modes in 1D and $12\times 12$ retained modes in 2D.
ConvLSTM	One recurrent convolutional layer with hidden size 64 and a $3\times 3$ kernel, followed by a $1\times 1$ output projection.
Transformer	Patch-wise encoder with patch size 8, model dimension 256, 8 attention heads, and 4 encoder layers.

Appendix

Table 9: Backbone architectures used throughout the benchmark. All models are trained as one-step autoregressive predictors with history length

L=4

; U-Net, FNO, and Transformer take channel-concatenated histories, while ConvLSTM takes the same history as a sequence.

Baseline

Instantiation used in Experiment 1

GLISTER -style

Validation-guided greedy selector adapted to PDE rollout training: train a lightweight pilot model, define a short-rollout validation objective on held-out starts, score candidates using the standard first-order Taylor approximation to the validation loss, and after each greedy inclusion apply a pseudo-update and refresh the validation gradient. Validation starts are chosen from the hardest held-out starts under the pilot model; validation-query size

=16

for diff-sorp and

=8

for diff-react / rdb ; step size

\eta=10^{-3}

; stochastic fraction

=1.0

PRISM -style

Targeted-subset PRISM implementation adapted to temporal-start selection: form a query set from the hardest validation starts and a private set from the easiest validation starts under the pilot model; represent starts with CountSketch-compressed rollout-gradient embeddings; instantiate canonical PRISM objectives ( flcmi , flmi , gcmi , gccg ) and select with the best globally fixed protocol used throughout the benchmark. Unless otherwise stated, the reported PRISM numbers use query size

=24

, private size

=24

, and globally fixed scaling settings.

Appendix

Table 10: Task-adapted GLISTER-style and PRISM-style baselines used in Experiment 1 under the same shared-start offline protocol as GITS.

Dataset	Backbone	uniform	loss-only	coverage-only	GradMatch	GLISTER	PRISM	GITS
diff-react	ConvLSTM	0.242	0.481	0.257	0.191	0.896	0.886	0.189
	FNO	0.473	0.747	0.504	0.823	0.791	0.771	0.393
	Transformer	0.224	0.612	0.289	0.566	0.511	0.529	0.284
	U-Net	0.090	0.117	0.097	0.051	0.053	0.059	0.045
diff-sorp	ConvLSTM	0.247	0.967	0.666	0.702	0.462	0.458	0.482
	FNO	0.164	0.519	0.232	0.198	0.217	0.167	0.147

Appendix

Table 11: Supplementary test rollout nRMSE at sampling ratio

0.10

with the same seven samplers as Table 1 .

Dataset	Backbone	uniform	loss-only	coverage-only	GradMatch	GLISTER	PRISM	GITS
diff-react	ConvLSTM	0.239	0.364	0.184	0.239	0.981	0.885	0.165
	FNO	0.513	0.823	0.493	0.876	0.831	0.782	0.445
	Transformer	0.244	0.547	0.276	0.551	0.570	0.631	0.269
	U-Net	0.066	0.074	0.037	0.032	0.042	0.039	0.069
diff-sorp	ConvLSTM	0.316	0.962	0.804	0.994	0.587	0.591	0.660
	FNO	0.135	0.309	0.187	0.183	0.163	0.118	0.105

Appendix

Table 12: Supplementary test rollout nRMSE at sampling ratio

0.20

with the same seven samplers as Table 1 .

Dataset	Backbone	loss-only	grad-only	loss-div	GITS
diff-react	ConvLSTM	0.481	0.635	0.246	0.189
	FNO	0.747	0.900	0.909	0.393
	Transformer	0.612	0.612	0.538	0.284
	U-Net	0.117	0.448	0.064	0.045
diff-sorp	ConvLSTM	2.249	2.249	2.417	0.482
	FNO	0.519	0.519	0.198	0.147

Appendix

Table 13: Supplementary

2\times 2

ablation at sampling ratio

0.10

Dataset	Backbone	loss-only	grad-only	loss-div	GITS
diff-react	ConvLSTM	0.404	0.471	0.223	0.165
	FNO	0.823	0.903	0.868	0.445
	Transformer	0.547	0.547	0.522	0.269
	U-Net	0.174	0.096	0.134	0.069
diff-sorp	ConvLSTM	1.749	1.749	1.667	0.660
	FNO	0.309	0.309	0.133	0.105

Appendix

Table 14: Supplementary

2\times 2

ablation at sampling ratio

0.20

Ratio	uniform	loss-only	coverage-only	GradMatch	GLISTER	PRISM	GITS
0.05	233.0	135.0	279.5	136.9	223.9	233.5	272.5
0.10	332.5	202.3	272.8	165.7	298.3	303.3	335.8
0.20	371.0	262.1	383.9	260.9	371.7	386.3	390.7

Appendix

Table 15: Mean downstream training time per dataset–backbone–ratio configuration (seconds), averaged over the 12 dataset–backbone configurations at each sampling ratio for the seven-sampler comparison.

Ratio	loss-only	GradMatch	GLISTER	PRISM	GITS
0.05	9.945	11.056	11.178	10.553	10.229
0.10	10.161	10.299	11.981	11.245	10.505
0.20	10.537	11.726	16.643	12.929	11.989

Appendix

Table 16: Mean subset-selection time per dataset–backbone–ratio configuration (seconds), averaged over the 12 dataset–backbone configurations at each sampling ratio for the seven-sampler comparison. Sampling time for uniform and coverage-only remains below 0.001s and is omitted.

Ratio	Dataset	Backbone	uniform	loss-only	coverage-only	GradMatch	GLISTER	PRISM	GITS
0.05	diff-react	ConvLSTM	0.028	0.032	0.013	0.028	0.036	0.026	0.009
		FNO	0.009	0.017	0.006	0.014	0.024	0.012	0.005
		Transformer	0.011	0.023	0.008	0.021	0.039	0.022	0.009
		U-Net	0.013	0.007	0.004	0.005	0.025	0.003	0.003
	diff-sorp	ConvLSTM	0.155	0.569	0.261	0.290	0.569	0.439	0.010
		FNO	0.036	0.073	0.024	0.021	0.160	0.017	0.014

Appendix

Table 17: Test rollout cRMSE.

Ratio	Dataset	Backbone	uniform	loss-only	coverage-only	GradMatch	GLISTER	PRISM	GITS
0.05	diff-react	ConvLSTM	0.062	0.072	0.026	0.062	0.082	0.059	0.020
		FNO	0.148	0.279	0.091	0.235	0.410	0.201	0.076
		Transformer	0.060	0.129	0.039	0.118	0.225	0.124	0.045
		U-Net	0.030	0.017	0.008	0.011	0.058	0.006	0.007
	diff-sorp	ConvLSTM	0.056	0.213	0.088	0.108	0.215	0.164	0.003
		FNO	0.043	0.090	0.026	0.026	0.203	0.020	0.016

Appendix

Table 18: Test rollout bRMSE.

Ratio	Dataset	Backbone	uniform	loss-only	coverage-only	GradMatch	GLISTER	PRISM	GITS
0.05	diff-react	ConvLSTM	0.014	0.015	0.006	0.013	0.017	0.012	0.004
		FNO	0.004	0.007	0.003	0.006	0.010	0.005	0.002
		Transformer	0.004	0.008	0.003	0.007	0.013	0.007	0.003
		U-Net	0.005	0.003	0.001	0.002	0.010	0.001	0.001
	diff-sorp	ConvLSTM	0.079	0.255	0.115	0.138	0.264	0.203	0.005
		FNO	0.029	0.053	0.017	0.017	0.117	0.013	0.011

Appendix

Table 19: Test rollout fRMSE-low.

Ratio	Dataset	Backbone	uniform	loss-only	coverage-only	GradMatch	GLISTER	PRISM	GITS
0.05	diff-react	ConvLSTM	0.006	0.006	0.002	0.005	0.007	0.005	0.002
		FNO	0.004	0.008	0.003	0.006	0.011	0.006	0.002
		Transformer	0.003	0.007	0.002	0.006	0.011	0.006	0.002
		U-Net	0.002	0.001	0.001	0.001	0.004	0.001	0.001
	diff-sorp	ConvLSTM	0.011	0.039	0.017	0.019	0.039	0.029	0.001
		FNO	0.009	0.018	0.006	0.005	0.039	0.004	0.003

Appendix

Table 20: Test rollout fRMSE-mid.

Ratio	Dataset	Backbone	uniform	loss-only	coverage-only	GradMatch	GLISTER	PRISM	GITS
0.05	diff-react	ConvLSTM	0.001	0.001	0.001	0.001	0.001	0.001	0.001
		FNO	0.002	0.005	0.002	0.004	0.007	0.003	0.001
		Transformer	0.001	0.002	0.001	0.002	0.003	0.002	0.001
		U-Net	0.001	0.001	0.001	0.001	0.001	0.001	0.001
	diff-sorp	ConvLSTM	0.002	0.006	0.003	0.003	0.006	0.005	0.001
		FNO	0.002	0.003	0.001	0.001	0.008	0.001	0.001

Appendix

Table 21: Test rollout fRMSE-high.

Method	Mean std	Median std	Max std
uniform	0.059	0.045	0.143
loss-only	0.127	0.079	0.874
coverage-only	0.062	0.049	0.183
GradMatch	0.090	0.065	0.317
GLISTER	0.283	0.090	2.267
PRISM	0.252	0.050	1.842

Appendix

Table 22: Summary of seed variability across the 36 seven-sampler benchmark configurations. Lower is better.

Dataset	Backbone	ratio = 0.05							ratio = 0.10							ratio = 0.20
		uniform	loss-only	coverage-only	GradMatch	GLISTER	PRISM	GITS	uniform	loss-only	coverage-only	GradMatch	GLISTER	PRISM	GITS	uniform	loss-only	coverage-only	GradMatch	GLISTER	PRISM	GITS
diff-react	ConvLSTM	0.041	0.063	0.038	0.057	0.088	0.055	0.019	0.029	0.058	0.044	0.034	0.550	0.491	0.017	0.033	0.046	0.031	0.041	1.274	1.227	0.016
	FNO	0.074	0.121	0.068	0.098	0.536	0.073	0.053	0.061	0.089	0.072	0.091	0.072	0.125	0.047	0.058	0.097	0.066	0.103	0.052	0.017	0.051
	Transformer	0.044	0.079	0.039	0.073	0.067	0.065	0.036	0.031	0.069	0.048	0.062	0.023	0.011	0.031	0.028	0.058	0.043	0.066	0.015	0.003	0.027
	U-Net	0.037	0.024	0.017	0.019	0.080	0.009	0.012	0.021	0.019	0.023	0.014	0.006	0.043	0.009	0.018	0.013	0.011	0.009	0.015	0.020	0.011
diff-sorp	ConvLSTM	0.073	0.196	0.138	0.157	0.398	0.371	0.009	0.068	0.214	0.183	0.219	0.091	0.228	0.147	0.059	0.168	0.172	0.241	0.342	0.360	0.131

Appendix

Table 23: Exact standard deviations (over three training seeds) of test rollout nRMSE for the same seven-sampler configurations as Tables 1 , 11 , and 12 .

Training Data Selection

Training data selection studies how to choose, filter, weight, or schedule examples to improve learning while reducing data and compute costs.

See all 184 papers

Learning Where to Simulate: Generative Active Sampling for Online PDE Surrogate Training
Jun 8, 2026Pierre Cesar, Sofya Dymchenko, Abhishek Purandare +1PDE Surrogate Modeling Active Learning
DATAMOVE · Univ. Grenoble Alpes, Inria, CNRS, Grenoble INP, LIG, Grenoble, France
Same Target, Different Origins: Aligning Latent Rollout for Consistent PDE Emulation
Oct 15, 2026Xinying Li, Jing Xu, Wenjun Wang +4Neural PDE Solvers PDE Operator Learning
Fudan University · Shanghai Academy of AI for Science
Compute-optimal data scaling for neural surrogates via multi-fidelity training
Oct 12, 2026Paul Setinek, Fabian Paischer, Nils Thuerey +2Neural Surrogate Modeling PDE Surrogate Modeling
Institute for Machine Learning, JKU Linz · Mistral AI · Technical University of Munich +1

Jun 8, 2026J/K move · Enter open · S save

Learning Where to Simulate: Generative Active Sampling for Online PDE Surrogate Training

Pierre Cesar, Sofya Dymchenko, Abhishek Purandare, Bruno Raffin

DATAMOVE · Univ. Grenoble Alpes, Inria, CNRS, Grenoble INP, LIG, Grenoble, France

PDE Surrogate Modeling Active Learning Adaptive Sampling

Open paper AskPDF arXiv

Data-driven PDE surrogates are trained with data produced by numerical PDE solvers. However, when the surrogate's goal is to generalize across a wide range of PDE configurations (e.g., initial conditions and physical coefficients), generating a representative training set is non-trivial. Uniform sampling of configuration parameters often under-represents trajectories exhibiting challenging dynamics, leading to high prediction errors and large error variance in the trained surrogate. Online training, where data generation and surrogate training are coupled, offers a natural advantage by allowing solver parameters to be steered on-the-fly. To efficiently exploit this capability, we introduce Online Generative Active Sampling (OGAS), an active learning method that reactively learns the relationship between configuration parameters and surrogate performance to control the sampling distribution. OGAS trains a fast diffusion model in parallel to the surrogate to act as a conditional sampler, mapping a surrogate-derived difficulty signal (e.g., loss or uncertainty) to configuration parameters. By actively drawing target signals from a prior biased toward high difficulty, OGAS continuously steers data generation toward challenging regimes without delaying the training workflow. We evaluate OGAS across 2D PDEs with distinct challenging dynamics (Kuramoto-Sivashinsky, Navier-Stokes, Gray-Scott) and up to 308 parameters, using multiple surrogate architectures. Across all settings, OGAS consistently improves tail statistics, yielding substantial reductions in errors above the 99th percentile and overall error dispersion compared to uniform sampling. While prioritizing challenging trajectories introduces a trade-off with average error, OGAS effectively ensures worst-case reliability of trained surrogates with negligible wall-time overhead.

Abstract

Training Data Selection

Explore similar work

Learning Where to Simulate: Generative Active Sampling for Online PDE Surrogate Training

Same Target, Different Origins: Aligning Latent Rollout for Consistent PDE Emulation

Compute-optimal data scaling for neural surrogates via multi-fidelity training