BeamVQ: Aligning Space-Time Forecasting Model via Self-training on Physics-aware Metrics

TL;DR

BeamVQ addresses the mismatch between data-driven space-time forecasts and physical laws by introducing a self-training loop guided by physics-aware metrics. It converts any encoder-decoder forecast into a probabilistic model using a discrete code bank and employs beam search to generate multiple candidate futures, which are filtered by physics scores and used to augment training data. Across five benchmarks and ten backbones, BeamVQ yields substantial gains in both statistical accuracy (lower MSE) and physical alignment (divergence, TKE, energy spectrum), with notable improvements in long-term forecasting. The method is flexible, scalable, and ready to integrate with additional physical constraints and larger datasets to further enhance real-world predictive reliability.

Abstract

Data-driven deep learning has emerged as the new paradigm to model complex physical space-time systems. These data-driven methods learn patterns by optimizing statistical metrics and tend to overlook the adherence to physical laws, unlike traditional model-driven numerical methods. Thus, they often generate predictions that are not physically realistic. On the other hand, by sampling a large amount of high quality predictions from a data-driven model, some predictions will be more physically plausible than the others and closer to what will happen in the future. Based on this observation, we propose \emph{Beam search by Vector Quantization} (BeamVQ) to enhance the physical alignment of data-driven space-time forecasting models. The key of BeamVQ is to train model on self-generated samples filtered with physics-aware metrics. To be flexibly support different backbone architectures, BeamVQ leverages a code bank to transform any encoder-decoder model to the continuous state space into discrete codes. Afterwards, it iteratively employs beam search to sample high-quality sequences, retains those with the highest physics-aware scores, and trains model on the new dataset. Comprehensive experiments show that BeamVQ not only gave an average statistical skill score boost for more than 32% for ten backbones on five datasets, but also significantly enhances physics-aware metrics.

Figures (8)

• Figure 1: Overview of BeamVQ: Input data $\mathcal{X}$ is fed to the encoder $E_{\phi}$ to produce the latent state $\mathcal{Z}_t = E_{\phi}(\mathcal{X})$. BeamVQ inserts code bank quantization layers between the encoder $E_{\phi}$ and decoder $D_{\varphi}$. $K$ candidate latent states are sampled via top-$K$ beam search: $s_{all} = \{s_1, s_2, \ldots, s_{n}\}$. Each state $s_i$ is decoded to the output sequence: $\mathcal{Y}_i = D_{\varphi}(s_{i}), \forall i \in [N]$. These candidate outputs are filtered based on the physics-aware score $F(\mathcal{Y}_i)$, in which the high-score samples are added back to the training dataset. BeamVQ iteratively generates new samples with the updated model weights to shift the data distribution to better match the physical law.
• Figure 2: I. The top row shows the actual and predicted distributions of height and speed with a white box highlighting significant discrepancies. Dotted patterns reflect the direction and magnitude of velocity vectors. II. The second row shows the average energy spectrum, depicting energy distribution across wavenumbers (k) and highlighting fluctuations in the dynamic system over different scales. Although CNO and FNO achieve similar MSE, the energy spectrum of CNO is closer to the ground-truth and is thus more physically realistic.
• Figure 3: The closed-loop of BeamVQ: (1) Beam search. it generates different output sequences for each input, by sampling different continuous states with the probabilistic model. (2) Dataset Update: it filters out the output sequences with high physics-aware Scores to expand the training dataset. (3) Self-training: it trains on the shifted data distribution for better physical consistency.
• Figure 4: Visualisation of the embedding space.
• Figure 5: The BeamVQ plugin improves physical consistency and prediction accuracy.(a) shows a visual comparison of the actual target, predicted results, and errors at different time steps. (b) displays the changes in SSIM, RMSE, and relative L2 error over time steps. (c) compares the turbulent TKE. (d) presents the energy spectrum at different wavenumbers.