MeTok: An Efficient Meteorological Tokenization with Hyper-Aligned Group Learning for Precipitation Nowcasting

Abstract

Recently, Transformer-based architectures have advanced meteorological prediction. However, this position-centric tokenizer conflicts with the core principle of meteorological systems, where the weather phenomena undoubtedly involve synergistic interactions among multiple elements while positional information constitutes merely a component of the boundary conditions. This paper focuses primarily on the task of precipitation nowcasting and develops an efficient distribution-centric Meteorological Tokenization (MeTok) scheme, which spatially sequences to group similar meteorological features. Based on the rearrangement, realigned group learning enhances robustness across precipitation patterns, especially extreme ones. Specifically, we introduce the Hyper-Aligned Grouping Transformer (HyAGTransformer) with two key improvements: 1) The Grouping Attention (GA) mechanism uses MeTok to enable self-aligned learning of features from different precipitation patterns; 2) The Neighborhood Feed-Forward Network (N-FFN) integrates adjacent group features, aggregating contextual information to boost patch embedding discriminability. Experiments on the ERA5 dataset for 6-hour forecasts show our method improves the IoU metric by at least 8.2% in extreme precipitation prediction compared to other methods. Additionally, it gains performance with more training data and increased parameters, demonstrating scalability, stability, and superiority over traditional methods.

Figures (7)

• Figure 1: The weather phenomena essentially involve nonlinear interactions among multiple factors across multi-scale spatiotemporal domains. The MeTok scheme introduces the historical distribution of each embedding to gather information with similar patterns.
• Figure 2: Patch embeddings are rearranged according to the historical precipitations.
• Figure 3: (a) The encoder-translator-decoder framework. The encoder models the spatial features with the rearranged patch features. The translator learns spatiotemporal dynamics. The decoder integrates the spatiotemporal information and local features. (b) The meteorological embeddings are rearranged based on the historical precipitations in the MeTok scheme. (c) Within the HyAGTransformer, the GA mechanism captures dependencies between embeddings and grouped features.
• Figure 4: The ablation study of the proposed framework. We conducted a comprehensive analysis of the proposed framework, examining it from various facets.
• Figure 5: The visualization of prediction results on the ERA5 dataset.