Training patterns for autoregressive neural PDE solvers (FNO, DeepONet, CNO). Covers rollout training, noise injection for stability, multi-component loss functions (H1, frequency-sensitive, boundary-aware), per-channel normalization for coupled multi-variable systems, and the PDEBench nRMSE metric. Use when training any neural operator that predicts time-dependent PDE solutions.
Install with Codex or Claude Copy this prompt, paste it into Codex, Claude, or another assistant, and let it review the skill page and install it for you.
A direct command skips the review prompt. Inspect the source before running it.
Training patterns for autoregressive neural PDE solvers (FNO, DeepONet, CNO). Covers rollout training, noise injection for stability, multi-component loss functions (H1, frequency-sensitive, boundary-aware), per-channel normalization for coupled multi-variable systems, and the PDEBench nRMSE metric. Use when training any neural operator that predicts time-dependent PDE solutions.
Autoregressive Neural PDE Solvers — Training Patterns
When to Use
Training any neural operator (FNO, DeepONet, CNO, etc.) to predict time-dependent PDE solutions
The model predicts one (or few) timesteps ahead, then feeds predictions back as input
Multi-variable PDE systems (compressible NS, MHD, multi-species reaction, etc.)
Problems where single-step accuracy is insufficient — rollout stability matters
Core Principle: Train How You Infer
The single biggest mistake in neural PDE solving is training on single-step predictions but inferring with multi-step rollout. The model never learns to handle its own errors.
# WRONG: Teacher forcing (single-step)for t inrange(T):
pred = model(ground_truth[:, :, t]) # Always sees perfect input
loss += mse(pred, ground_truth[:, :, t+1])
# At inference, errors compound exponentially because model never saw noisy inputs# RIGHT: Autoregressive rollout training
inp = initial_condition # shape: [B, X, init_steps, C]for t inrange(init_steps, T_total):
pred = model(inp) # Sees its OWN previous predictions
loss += mse(pred, ground_truth[:, :, t:t+1, :])
inp = torch.cat([inp[:, :, 1:, :], pred], dim=-2) # Shift window, append prediction# Model learns to be robust to its own prediction errors
Data Layout Convention
For spectral operators (FNO), use spatial-first layout:
# Standard: [N_samples, X_spatial, T_timesteps, C_channels]# NOT: [N_samples, T_timesteps, X_spatial, C_channels]# When loading from [N, T, X] or [N, T, X, C]:
data = np.transpose(raw[:, ::stride_t, ::stride_x], (0, 2, 1)) # 3D
data = np.transpose(raw[:, ::stride_t, ::stride_x, :], (0, , , ))
2
1
3
# 4D
Noise Injection for Rollout Stability
During autoregressive training, add small noise to predictions before feeding back. This prevents the model from relying on artificially clean inputs.
Partition Fourier spectrum into bands with increasing weights — upweights high-frequency content.
deffrequency_loss(pred, target, weights=[1.0, 2.0, 4.0]):
"""Higher weight on high-frequency errors."""
pred_ft = torch.fft.rfft(pred, dim=1)
tgt_ft = torch.fft.rfft(target, dim=1)
n = pred_ft.shape[1]
bin_size = n // len(weights)
loss = 0.0for i, w inenumerate(weights):
s, e = i * bin_size, (i+1) * bin_size if i < len(weights)-1else n
loss += w * torch.abs(pred_ft[:, s:e] - tgt_ft[:, s:e]).mean()
return loss
# Usage: loss = mse_loss + 0.1 * frequency_loss(pred, target)
Boundary-Aware Loss (for non-periodic BCs)
Upweight loss near domain boundaries where spectral methods struggle.
Per-Channel Normalization for Multi-Variable Systems
When PDE variables have different scales (e.g., density ~O(6), velocity ~O(0.5), pressure ~O(59)), the loss is dominated by the largest-scale variable.
# Compute from TRAINING set only
ch_mean = train_data.mean(dim=(0, 1, 2)) # [C]
ch_std = train_data.std(dim=(0, 1, 2)) + 1e-8# [C]# Normalize all data
data_normalized = (data - ch_mean) / ch_std
# Denormalize for evaluation (handles CPU/GPU mismatch)defdenormalize(x):
return x * ch_std.to(x.device) + ch_mean.to(x.device)
# IMPORTANT: Compute nRMSE in ORIGINAL scale, not normalized
preds_orig = denormalize(preds)
targets_orig = denormalize(targets)
nRMSE Metric (Standard PDE Benchmark Definition)
WARNING: Multiple valid nRMSE formulas exist and can differ by up to 5-10%. Always compute BOTH and verify you beat baselines under both.
Train/Val/Test Split (Critical for Methods Papers)
N_TRAIN, N_VAL, N_TEST = 8000, 1000, 1000
train_data = data[:N_TRAIN]
val_data = data[N_TRAIN:N_TRAIN+N_VAL] # Checkpoint selection
test_data = data[N_TRAIN+N_VAL:] # Final metric (evaluated ONCE)# During training: evaluate on val_loader for model selection# After training: evaluate on test_loader ONCE for final reported number
