autoregressive-neural-pde-solver

name	autoregressive-neural-pde-solver
description	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.
category	physics
version	1.0.0
author	Synthetic Sciences
license	MIT
tags	["Neural Operator","Autoregressive","PDE","Rollout","FNO","Loss Function","Normalization"]
dependencies	["torch>=2.1.0","numpy>=1.24.0"]

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 in range(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 in range(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, , , ))