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

TL;DR

MeTok introduces a distribution-centric meteorological tokenization and a novel transformer architecture that better captures synergistic weather interactions, significantly improving precipitation nowcasting accuracy, especially for extreme events.

Contribution

The paper proposes MeTok and HyAGTransformer, novel methods that enhance weather pattern modeling by focusing on distribution rather than position, leading to more robust precipitation predictions.

Findings

Improves IoU by at least 8.2% for extreme precipitation prediction

Demonstrates scalability and stability with more data and parameters

Outperforms traditional methods in 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.