PINN-Cast: Exploring the Role of Continuous-Depth NODE in Transformers and Physics Informed Loss as Soft Physical Constraints in Short-term Weather Forecasting

TL;DR

PINN-Cast introduces a continuous-depth transformer with Neural ODEs and physics-informed loss for improved short-term weather forecasting, combining data-driven efficiency with physical consistency.

Contribution

It presents a novel continuous-depth transformer architecture with an auxiliary derivative attention branch and a physics-informed training objective for weather prediction.

Findings

Outperforms standard discrete transformers in short-term weather forecasting.

Demonstrates the effectiveness of Neural ODE integration in transformer encoders.

Shows that physics-informed loss improves physical consistency of forecasts.

Abstract

Operational weather prediction has long relied on physics-based numerical weather prediction (NWP), whose accuracy comes at the cost of substantial compute and complex simulation workflows. Recent transformer-based forecasters offer efficient data-driven alternatives, however transformers are physics-agnostic models. Additionally, standard transformer encoders evolve representations through discrete layer updates that may be less suited to modeling smooth latent dynamics. In this work, we propose a continuous-depth transformer encoder for weather forecasting that integrates Neural Ordinary Differential Equation (Neural ODE) dynamics within each encoder block. Specifically, we replace discrete residual updates with ODE-based updates solved using adaptive numerical integration. We also introduce a two-branch attention module that combines conventional patch-wise self-attention with an auxiliary branch that applies a derivative operator to attention logits, providing an additional change-sensitive interaction signal. To further align forecasts with governing principles, we propose a customized physics-informed training objective that enforces physical consistency as a soft constraint. We evaluate the proposed method against a standard discrete transformer baseline and an existing continuous-time Neural ODE forecasting variant, demonstrating the importance of PINN-Cast in short term weather forecasting.