TEMPO-PDE: Trajectory-Aware Post-Training for PDE Foundation Models

Partial differential equation foundation models (PDE-FMs) are neural predictors pretrained across physical systems to learn reusable representations of PDE dynamics. However, accurate early predictions do not ensure that an adapted model remains accurate over the full forecast horizon. Reference trajectories provide ground-truth fields at multiple future times, but the corresponding loss gradients can differ in magnitude or conflict with one another. To address this problem, we propose TEMPO-PDE, a post-training framework with two complementary components. (i) Trajectory-aware supervision scores future predictions using field-value errors, spectral discrepancies, and local spatial differences, and organizes these measures into primary and auxiliary objectives. (ii) Primary-guided gradient composition removes opposing auxiliary components and bounds the remaining auxiliary norm relative to the primary gradient. Both components update the existing predictor without adding an inference module. Within the registered main comparison, TEMPO achieves the lowest observed time-averaged and final-step errors on the five main PDE benchmarks.