STC-ViT: Spatio Temporal Continuous Vision Transformer for Weather Forecasting

Neural ODEs are the continuous time models which learn the evolution of a system over time using Ordinary Differential Equations (ODE) (Chen et al., 2018). The key idea behind Neural ODE is to model the derivative of the hidden state using a neural network. Consider a hidden state $h(t)$ at time $t$ In a traditional neural network, the transformation from one layer to the next could be considered as moving from time $t$ to $t+1$. In Neural ODEs, instead of discrete steps, the change in $h(t)$ over time is defined by an ordinary differential equation parameterized by a neural network:

$$\frac{dh(t)}{dt}=f(h(t),t,\theta)$$

(1)

where $\frac{dh(t)}{dt}$ is the derivative of the hidden state with respect to time, $f$ is a neural network with parameters $\theta$ and $t$ is the time variable, allowing the dynamics of $h(t)$ to change with time.

ResNets vs Neural ODEs To see the connection between ResNets and Neural ODEs, consider a ResNet with layers updating the hidden state as:

$$h_{t+1}=h_{t}+f(h_{t},\theta_{t})$$

(2)

In the limit, as the number of layers goes to infinity and the updates become infinitesimally small, this equation resembles the Euler method for numerical ODE solving, where:

$$h(t+\triangle t)=h(t)+\triangle t\cdot f(h(t),t,\theta)$$

(3)

Reducing $\triangle t$ to an infinitesimally small value transforms the discrete updates into a continuous model described by the ODE given earlier. To compute the output of a Neural ODE, integration is used as a backpropagation technique. This is done using numerical ODE solvers, such as Euler, Runge-Kutta, or more sophisticated adaptive methods, which can efficiently handle the potentially complex dynamics encoded by $f$.

In weather and climate modeling, incorporating physical constraints ensures that the model adheres to the governing physical laws. These constraints can be added in two ways: hard constraints and soft constraints. Hard Constraints: For a given physical constraint $f(x)=0$ where $f(x)$ is the governing physical law, a hard constraint would mean that the machine learning model’s prediction $\hat{y}$ must always satisfy equation 4.

$$f(\hat{y})=0$$

(4)

This constraint can be embedded into model’s architecture as a constrained layer or optimizer.

Soft Constraints: Soft constraints adds a penalty term to the loss function that minimizes the violation of a physical law $f(\hat{y})$ as shown in equation 5

$$min_{\theta}=L(y,\hat{y})+\alpha||f(\hat{y})||^{2}$$

(5)

where $\alpha$ controls the weight of the penalty for violating the physical constraint $f(\hat{y})$ and $||f(\hat{y})||^{2}$ measures the degree of violation (e.g., deviation from mass conservation).

In the STC-ViT architecture, we enhance vision transformer-based weather forecasting by introducing Temporal Continuous Attention (TCA) and Spatial Attention (SA) mechanisms to capture dynamically evolving weather patterns. We build upon the variable tokenization and aggregation scheme proposed by (Nguyen et al., 2023a) followed by Continuous attention mechanism to enhance interpolated feature learning. The detailed architecture of STC-ViT is shown in figure 1.

Temporal Continuous Attention (TCA). We model the temporal dynamics with attention by incorporating derivatives (temporal changes) over time directly into the query, key, and value representations.

We formulate Query (Q) as the "current" time step’s information and is designed to seek relevant patterns or changes reflecting the model’s focus on upcoming changes equation 8.

$$Q_{t}=\frac{dx_{t}}{dt}$$

(8)

Key (K) models the "context" of prior time step states, which provide historical dynamics incorporating past time step derivatives equation 9.

$$K_{t-1}=f(x_{t-1},\frac{dx}{dt}|_{t-1})$$

(9)

Value (V) represents the updated information allowing the model to interpret them in the context of previous and current changes. This modification allows the model to learn the transitional changes from one time step to another which is important for capturing unprecedented changes in weather.

To compute TCA and capture the temporal change for the same variable, we calculate the attention between $Q_{t_{0},i}$ and $K_{t-1,i}$ focusing on the temporal evolution. Assume we have two input samples of size $(N\times W\times H)$ denoted as $X_{t0}$ at time $t_{0}$ and $X_{t-1}$ at time $t_{-1}$ where $N$ is the number of weather variables such as temperature, wind speed, pressure, etc., $H\times W$ refers to the spatial resolution of the variable. The attention mechanism to capture the temporal continuity between time steps $t_{0}$ and $t_{-1}$ is given by:

$$\frac{d}{dt}(Q_{t_{0},i}\cdot K_{t-1,i})=Q_{t_{0},i}\frac{d(K_{t-1,i})}{dt}+K_{t-1,i}\frac{d(Q_{t_{0},i})}{dt}$$

(10)

$$Y_{i}=Q_{t_{0},i}\frac{d(K_{t-1,i})}{dt}+K_{t-1,i}\frac{d(Q_{t_{0},i})}{dt}$$

(11)

$$Att(Q_{t_{0},i},K_{t-1,i})=Softmax(\frac{Y_{i}}{\sqrt{d_{k}}})$$

(12)

$$TCA_{t_{0},t-1}=\sum\limits^{N}_{k=1}Att(Q_{t_{0},i},K_{t-1,i})\cdot V_{t-1,i}$$

(13)

The resulting output represents the attention weighted sum of values for similar variables across input samples at time $t_{0}$ and $t_{1}$. This approach models the dynamics of each variable independently over time, allowing the model to capture temporal dependencies and changes effectively.

Spatial Attention (SA). For each variable i, we calculate the attention between spatial positions $(x,y)$ and $(x^{\prime},y^{\prime})$ given by equation 15:

$$Attention(Q_{x,y,i},K_{x^{\prime},y^{\prime},i})=softmax(\frac{Q_{x,y,i}\cdot K_{x^{\prime},y^{\prime},i}}{\sqrt{d_{k}}})$$

(14)

$$SA_{t_{0},i}=\sum_{x^{\prime},y^{\prime}}Attention(Q_{x,y,i},K_{x^{\prime},y^{\prime},i})\cdot V_{x^{\prime},y^{\prime},i}$$

(15)

where Q, K are matrices formed from all queries, keys, and values, respectively, $d_{k}$is the dimensionality of the keys and queries, and the division by $\sqrt{d_{k}}$ is a scaling factor to deal with gradient vanishing problems during training. Use the spatial attention weights to obtain the spatially-enhanced representation for each variable i by weighting the values:

Concatenation and Fusion. We concatenate the outputs of the TCA and SA along the feature dimension (as shown in equation 16) to provide a comprehensive representation, considering both temporal and spatial dependencies. This dual attention approach allows the Vision Transformer to effectively model complex spatio-temporal dependencies in the weather reanalysis data.

$$STC=concat(TCA_{t_{0},i},SA_{t_{0},i})$$

(16)

This output captures both the temporal evolution of each variable over time and the spatial interactions within each time step, enhancing the model’s ability to understand the continuous, spatiotemporal dynamics of the data.

Neural ODEs as Add operation in Transformers. Neural ODEs model the transformation of data through a continuous dynamical system. Instead of discrete layers, the transformation is modeled by a differential equation. The idea is to treat the depth of the network as a continuous time variable t. As explained in section 2.1, Neural ODEs can be seen as a continuous form of ResNets.

In STC-ViT, instead of adding the output of an attention mechanism discretely as $h^{\prime}=LayerNorm(h+STC(h))$, we model the transformation using a Neural ODE. The attention output is treated as a continuous evolution of the input state $h$ transforming equation 1 as equation 17. Solving this ODE from $t=0$ to $t=1$ gives us the updated state $h(1)$, which incorporates the attention mechanism’s effect continuously over time.

$$\frac{dh(t)}{dt}=STC(h(t),t,\theta)$$

(17)

To compute the output of a Neural ODE, we solve the initial value problem:

$$h(t_{1})=h(t_{0})+\int^{t_{1}}_{t_{0}}f(h(t),t,\theta)dt$$

(18)

This integral is computed using Runge-Kutta Runge (1895) numerical ODE solver. The solver adaptively determines the number of steps required, making the process efficient and accurate. This continuous approach allows for smoother transformations and improve model performance.

We use latitude weighted mean squared error to compute loss for the predicted variables.

$$L_{lat-weight}=\frac{1}{V\times H\times W}\sum_{v=1}^{V}\sum_{i=1}^{H}\sum_{j=1}^{W}(L_{i})(\hat{X}_{t+\Delta t}^{n,i,j}-X_{t+\Delta t}^{n,i,j})$$

(19)

where $L_{i}$ accounts for Latitude weights:

$$L(i)=\frac{\cos(lat(i))}{\frac{1}{H}\sideset{}{{}_{i^{\prime}=1}^{H}}{\sum}\cos(lat(i^{\prime}))}$$

Additionally, we account for three fundamental physical laws i.e. Potential Energy, Kinetic Energy and Thermodynamic Balance in our loss function. Kinetic energy in atmospheric science refers to the energy associated with the motion of air masses. Geopotential height is often used in meteorology to express the potential energy of an air parcel in the Earth’s gravitational field. It represents the energy per unit mass of an air parcel at height z above sea level and thermodynamic balance equation describes the evolution of temperature in an air parcel due to processes like heat addition and pressure changes. These fundamental principles, coupled with wind components and thermodynamic variables like temperature and geopotential height, form the core of atmospheric dynamics used in climate and weather models.

$$L_{kinetic}=\left|\frac{1}{2}(u^{2}_{pred}+v^{2}_{pred})-\frac{1}{2}(u^{2}_{true}+v^{2}_{true})\right|$$

(20)

where u is the eastward wind component (m/s) and v is the northward wind component (m/s).

$$L_{potential}=|g\cdot z_{pred}-g\cdot z_{true}|$$

(21)

where g is the acceleration due to gravity approximately $9.81m/s^{2}$ and z is the height or geopotential height (m).

$$L_{thermo}=\left|\frac{dT}{dt}+u\frac{dT}{dx}+v\frac{dT}{dy}\right|$$

(22)

where $\frac{dT}{dt}$ is the temporal change of temperature (model output), $u\frac{dT}{dx},v\frac{dT}{dy}$ represent the advection of temperature by wind in the x and y directions, respectively.

$$Loss_{physics}=L_{lat-weight}+\alpha\times L_{kinetic}+\beta\times L_{potential}+\gamma\times L_{thermo}$$

(23)

where $\alpha$, $\beta$ and $\gamma$ are the weight factors for physics based loss. The resulting loss aims to ensure that the predicted temperature and wind fields satisfy the physical laws of thermodynamics (like energy conservation) on their own. By setting up the equation to balance energy and temperature gradients, this loss constrains the predictions to maintain internal consistency with physical laws, regardless of the specific values in the ground truth.

We train STC-ViT on ERA5 dataset Hersbach et al. (2018; 2020) provided by the European Center for Medium-Range Weather Forecasting (ECMWF). We compare STC-ViT against several weather forecasting models by training it at two different resolutions of $5.625^{\circ}$ (32 x 64 grid points) provided by WeatherBench (Rasp et al., 2020) and $1.5^{\circ}$ (121 x 240 grid points) provided by WeatherBench2 (Rasp et al., 2024).

We consider weather forecasting as a continuous spatio temporal forecasting problem i.e a tensor of shape $N\times H\times W$ at time $t$ is fed to the pre-processing layer of the model where it passes through pre-derivation to STC-ViT and outputs a tensor of $N^{\prime}\times H^{\prime}\times W^{\prime}$ at future time step $t+\triangle t$. Complete training details of the model are given in the in Appendix A.

We train STC-ViT on hourly data with following set of variables: Land Sea Mask (LSM), Orography, 10-meter U and V wind components (U10 and V10) and 2-meter temperature (T2m) in addition to 6 atmospheric variables: geopotential (Z), temperature (T), U and V wind components, specific humidity (Q) and relative humidity (R) at 7 pressure levels: 50 250 500 600 700 850 925. We use data from 1979-2015 for training, 2016 for validation and 2017-2018 for testing phase. We compare STC-ViT with ClimaX, and ClimODE on ERA5 dataset at $5.625^{\circ}$ resolution provided by WeatherBench (Rasp et al., 2020). To ensure fairness, we retrained ClimaX from scratch without any pre-training.

STC-Vit outperforms ClimaX and ClimODE at all lead times which shows that replacing regular attention with continuous attention in ViT architecture derives improved feature extraction by mapping the changes occurring between successive time steps. Additionally, enforcing physical constraints in the model lead to improved prediction scores. The RMSE and ACC results are shown in Table 1.

To keep training consistent with WeatherBench 2, we utilize the training data from 1979 to 2018, validation data from 2019, and test data from 2020 year. We train STC-ViT on 6 hourly data for following variables: T2m, u10 and v10 wind components and mean sea-level pressure (MSLP) along with five atmospheric variables: geopotential height (Z), temperature (T), U and V wind components, and specific humidity (Q). These atmospheric variables are considered at 13 pressure levels: 50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000 hPa. We also compare our results with both versions of IFS Andersson (2022), IFS-HRES which is state-of-the-art forecasting model at high resolution of $0.1^{\circ}$ and IFS-ENS, ensemble version trained at $0.2^{\circ}$. Additionally, we compare STC-ViT against GraphCast and PanguWeather which are trained at a higher resolution of $0.25^{\circ}$ and finally with Neural GCM trained at $0.7^{\circ}$.

STC-ViT shows competitive performance, outperforming Pangu and GraphCast particularly for geopotential, temperature and wind variables. This is due to the physics loss penalizing the model for not obeying physics and encouraging outputs that balance temperature evolution and respect thermodynamic laws by penalizing deviations from these energy fluxes. We find out that Neural GCM is the best performing model on atmospheric variables which could be due to the physics embedded in the dynamical core. We believe that STC-ViT could also benefit from high resolution training and physics directly incorporated in the model architecture which we will explore in future. The RMSE and ACC results are shown in Figures 2 and 4 respectively.

Vanilla Vision Transformer Vision Transformer has emerged as a powerful architecture which captures long-term dependencies better than any model. For this ablation study, we simply train a basic ViT architecture on ERA5 dataset. Compared with STC-ViT, ViT under performs for prediction accuracy showing the superiority of continuous models in weather forecasting systems.

Vanilla Neural ODE Network. Another ablation study is done using vanilla Neural ODE model. We replace the transformer with Neural ODE architecture as proposed in the original paper Chen et al. (2018). While Neural ODEs are computationally efficient, they only aid in interpolating the temporal irregularities and ignore the spatial continuity. This study proves that STC-ViT learns both spatio-temporal continuous features from the discrete data and is better at representing dynamical forecast systems.

Continuous Attention. To understand the importance of continuous attention, we remove the neural ODE layer and provide the output of attention mechanism to the basic feed forward network. While continuous attention only model does not outperform the STC-ViT, it provides better prediction accuracy than the traditional attention model.

Continuous Attention + Neural ODE layer. We perform another ablation study to compare how well Neural ODE capture continuity in weather information when simply added as a layer in the transformer architecture. Neural ODE Chen et al. (2018) is a continuous depth neural network designed to learn continuous information in the process as well. For this study, we replace the continuous attention with vanilla attention and add a Neural ODE layer in the feed forward block. The feed forward solution is approximated using Runge-Kutta (RK) numerical solver.