BeamVQ: Beam Search with Vector Quantization to Mitigate Data Scarcity in Physical Spatiotemporal Forecasting

TL;DR

BeamVQ tackles data scarcity in physical spatiotemporal forecasting by introducing a probabilistic framework that combines a deterministic base predictor with a Top-K VQ-VAE to generate diverse futures. It uses beam search over the continuous state space to explore multiple trajectory variants and employs a domain-specific metric to guide selection, enabling a self-ensemble that augments training data. Across meteorological, fluid dynamics, and PDE-based benchmarks, BeamVQ achieves substantial mean squared error reductions (up to 39%) and improves extreme-event detection and physical plausibility. The approach enhances long-horizon forecasting and provides a robust, uncertainty-aware toolkit for data-limited physical systems, with broad applicability to climate, ocean, and engineering simulations.

Abstract

In practice, physical spatiotemporal forecasting can suffer from data scarcity, because collecting large-scale data is non-trivial, especially for extreme events. Hence, we propose \method{}, a novel probabilistic framework to realize iterative self-training with new self-ensemble strategies, achieving better physical consistency and generalization on extreme events. Following any base forecasting model, we can encode its deterministic outputs into a latent space and retrieve multiple codebook entries to generate probabilistic outputs. Then BeamVQ extends the beam search from discrete spaces to the continuous state spaces in this field. We can further employ domain-specific metrics (e.g., Critical Success Index for extreme events) to filter out the top-k candidates and develop the new self-ensemble strategy by combining the high-quality candidates. The self-ensemble can not only improve the inference quality and robustness but also iteratively augment the training datasets during continuous self-training. Consequently, BeamVQ realizes the exploration of rare but critical phenomena beyond the original dataset. Comprehensive experiments on different benchmarks and backbones show that BeamVQ consistently reduces forecasting MSE (up to 39%), enhancing extreme events detection and proving its effectiveness in handling data scarcity.

Figures (6)

• Figure 1: The visualization of extreme marine heatwave events shows that BeamVQ enhances Backbone models and improves their ability to capture extreme events. Detailed experimental results are provided in the experiments section.
• Figure 2: Architecture Overview of BeamVQ. (a) Stage $1$: Base Model Training: A deterministic predictor (FNO/ViT/ConvLSTM) learns single-step mappings $\mathbf{X}_t \xrightarrow{f_{\Theta_f}} \hat{\mathbf{Y}}_{t+1}$ via MSE minimization. (b) Stage $2$: Top-K VQ-VAE: Latent code $\mathbf{z}$ from encoder $e_{\Phi_h}$ is quantized to $K$ nearest codebook vectors $\{\mathbf{q}^{(k)}\}$, decoded to diverse predictions $\{\tilde{\mathbf{Y}}_{t+1}^{(k)}\}$. (c) Joint Optimization: The optimal reconstruction $\tilde{\mathbf{Y}}_{t+1}^*$ (selected by metric $M$) guides base model refinement, while top-$K'$ ensemble $\bar{\mathbf{Y}}_{t+1}$ enables self-training.
• Figure 3: The prediction results of marine extreme heatwave events include: A visual comparison (from left to right): ground truth labels, SimVP+BeamVQ prediction results on day 10, and SimVP prediction results on day 10. The cumulative changes of RMSE over prediction time. The cumulative changes of CSI over prediction time.
• Figure 4: 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.
• Figure 5: The t-SNE visualization in (a), (b), and (c) shows the Ground-truth, ConvLSTM and ConvLSTM+BeamVQ predictions, respectively. (d) shows the analysis of the Codebank parameters.

Theorems & Definitions (1)

• Theorem 1