Test-Time Scaling for PDE Surrogates via Solution-Space Refinement Flows

Neural PDE surrogates have emerged as powerful tools for accelerating scientific and engineering simulations, but are typically deployed as one-step predictors and remain largely decoupled from the test-time scaling paradigm that has driven recent advances in generative models and large language models. Inspired by flow-based learning and on-policy training, we introduce Flow Corrector (FOCO), a computationally-efficient general framework that improves frozen PDE surrogates through iterative test-time refinement. FOCO learns a shared corrector from self-generated on-policy states in solution space and stochastic detached rollouts, avoiding backpropagation through long correction trajectories while substantially reducing training computation and activation-memory requirements. Across global weather super-resolution and next-state prediction for 2D and 3D spatiotemporal PDE systems, FOCO consistently improves frozen surrogates and outperforms matched baselines, reducing normalized MSE by \(78.6%\) on 3D magnetohydrodynamic turbulence and by \(84.8%\) over a 20-step rollout on the Wave Layer benchmark. The learned corrector further transfers zero-shot across neural-operator architectures without retraining.