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

Qizhao Jin; Xianhuang Xu; Yong Cao; Shiming Xiang; Xinyu Xiao

cs.AIMar 14, 2026

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

Authors: Qizhao Jin, Xianhuang Xu, Yong Cao, Shiming Xiang, Xinyu Xiao

Organizations: School of Artificial Intelligence, University of Chinese Academy of Sciences · China Meteorological Administration, National Meteorological Centre · State Key Laboratory of Multimodal Artificial Intelligence Systems, Institute of Automation

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 & tables

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.

Panel of Figure 2 — 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.

Model	Params(M)	MACs(G)	TS $\uparrow$	IoU $\uparrow$
ConvLSTM	14.79	463.74	20.68	37.83
ConvGRU	17.19	604.8	20.41	38.03
TrajGRU	12.72	720.64	19.50	36.09
PredRNN	25.84	1103.64	17.95	36.58
PFST	40.65	472.97	18.95	38.05
SCCN	48.1	2990.48	19.77	38.01

Table 1: Overall performances of our framework on ERA5 dataset. The results are the mean values covering all categories, where the values in bold are the best.

Model	Params(M)	MACs(G)	TS $\uparrow$	IoU $\uparrow$
ConvLSTM	14.03	309.16	14.86	28.81
ConvGRU	17.19	403.2	14.34	26.99
TrajGRU	12.72	480.43	14.83	28.69
PredRNN	25.84	735.76	14.52	28.64
PFST	40.65	315.31	11.68	26.32
SCCN	48.1	1993.71	14.82	29.64

Table 2: Overall performances of our framework on the WeatherBench dataset. The results are the mean values covering all categories, where the values in bold are the best.

	ERA5	WB	WB 2
Spatial resolution	$0.25^{\circ}$	$1.40625^{\circ}$	$1.5^{\circ}$
Temporal resolution	1 hr	1 hr	6hr
Vairables	12	12	79
Spatial range	$0^{\circ}$ $\sim$ $55^{\circ}$ N	global	global
Spatial range	$140^{\circ}$ E $\sim$ $70^{\circ}$ W	global	global
Training data	2016 $\sim$ 2018	2014 $\sim$ 2016	1989 $\sim$ 2019

Table 3: The experiment setting of ERA5, WeatherBench, and WeatherBench 2.

Method	Prediction time (h)
Method	0 $\sim$ 6	6 $\sim$ 12	12 $\sim$ 18	18 $\sim$ 24	24 $\sim$ 30	30 $\sim$ 36
IFS HRES	1.39	1.28	1.27	1.45	1.42	1.57
IFS ENS mean	1.29	1.15	1.11	1.30	1.23	1.39
Ours	0.84	1.03	1.22	1.35	1.49	1.58

Table 4: Performance of the proposed method on the WeatherBench 2 dataset.

Training data	# Var.	# Param.	RMSE
2009 $\sim$ 2019	52	4.82M	1.91
1999 $\sim$ 2019	52	8.15M	1.48
1989 $\sim$ 2019	52	8.15M	1.34
1989 $\sim$ 2019	79	8.15M	1.25

Table 5: Study on the scalability of the proposed method on the WeatherBench 2 data set.

Model	RL	LR	MR	HR	RS
PFST	88.13	50.12	22.88	17.27	12.04
SCCN	87.13	51.42	20.16	17.5	13.85
SimVP	87.28	50.97	25.89	21.02	14.04
SimVP+gSTA	87.23	51.71	25.13	19.36	13.88
TAU	88.41	52.30	26.33	19.31	12.58
Ours	88.52	54.25	30.04	25.01	22.25

Table 6: Results on the ERA5 dataset of each precipitation nowcasting IoU for multiple levels. We report IoU (%) for each level.

Method	RL	LR	MR	HR	RS
Dice Loss	87.08	50.31	25.89	20.20	14.79
Seesaw Loss	88.28	54.18	26.59	18.82	17.20
Recall Loss	88.12	53.93	29.22	20.92	19.82
Ours(Dice Loss)	88.52	54.25	30.04	25.01	22.25
Ours(Recall Loss)	88.39	56.89	31.30	24.72	20.73

Table 7: The comparison of IoU metric of different long tail learning methods on the ERA5 dataset.

MeTok	GA	N-FFN	IoU $\uparrow$	TS $\uparrow$
			40.93	21.04
✓			41.19	21.22
✓	✓		42.74	21.51
✓		✓	43.93	21.96
	✓	✓	40.79	20.98
✓	✓	✓	44.41	22.33

Table 8: Ablation studies on the effectiveness of each ingredient of HyAGTransformer on the ERA5 dataset.

Figure 4: The ablation study of the proposed framework. We conducted a comprehensive analysis of the proposed framework, examining it from various facets.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Intensity	RL	LR	MR	HR	RS
Precipitation(mm)	[0, 0.1)	[0.1, 4.0)	[4.0, 13.0)	[13.0, 25.0)	[25.0, $\infty$ )
ERA5	80.76%	18.22%	0.96%	0.06%	0.006%
WeatherBench	83.59%	15.94%	0.44%	0.02%	0.003%

Appendix

Table 9: The distribution of precipitation intensity on the ERA5 dataset and WeatherBench dataset. Precipitation intensity is categorized into five ascending levels: Rainless (RL), Light Rain (LR), Moderate Rain (MR), Heavy Rain (HR), and Rainstorm (RS).

Metric	Method	Prediction time (hours)
Metric	Method	0 $\sim$ 1	1 $\sim$ 2	2 $\sim$ 3	3 $\sim$ 4	4 $\sim$ 5	5 $\sim$ 6
IoU $\uparrow$	Pers.	31.10	28.64	26.98	25.78	24.89	24.22
	Clim.	13.36	13.36	13.36	13.36	13.36	13.36
	W-Clim.	16.68	16.68	16.68	16.68	16.68	16.68
	ConvLSTM	40.52	40.10	38.84	37.39	35.85	34.29
	ConvGRU	40.01	39.71	38.83	37.75	36.56	35.31

Appendix

Table 10: Overall performances of the proposed method on the ERA5 dataset along the temporal axis.

Metric	Method	Prediction time (hours)
Metric	Method	0 $\sim$ 1	1 $\sim$ 2	2 $\sim$ 3	3 $\sim$ 4	4 $\sim$ 5	5 $\sim$ 6
IoU $\uparrow$	Pers.	28.50	27.94	26.20	24.95	24.03	23.33
	Clim.	16.10	16.10	16.10	16.10	16.10	16.10
	W-Clim.	17.04	17.04	17.04	17.04	17.04	17.04
	ConvLSTM	29.56	29.39	29.05	28.69	28.25	27.93
	ConvGRU	27.64	27.55	27.23	26.86	26.53	26.14

Appendix

Table 11: Overall performances of the proposed method on the WeatherBench dataset along the temporal axis.

Method	Prediction time (h)
Method	0 $\sim$ 6	6 $\sim$ 12	12 $\sim$ 18	18 $\sim$ 24	24 $\sim$ 30	30 $\sim$ 36
GraphCast	0.57	0.48	0.6	1.08	1.06	1.23
Ours	0.87	1.08	1.32	1.47	1.59	1.68

Appendix

Table 12: Comparison result of the proposed method with GraphCast on the WeatherBench 2 dataset.

\downarrow

is better.

Method	Prediction time (h)
Method	0 $\sim$ 6 $\downarrow$	6 $\sim$ 12 $\downarrow$	12 $\sim$ 18 $\downarrow$	18 $\sim$ 24 $\downarrow$	24 $\sim$ 30 $\downarrow$	30 $\sim$ 36 $\downarrow$
Baseline	0.95	1.19	1.45	1.71	1.93	2.11
Ours	0.79	0.91	1.05	1.18	1.33	1.49

Appendix

Table 13: Comparison result of the proposed method with the baseline for 10-meter wind speed prediction on the WB2 dataset.

Figure 5: The visualization of prediction results on the ERA5 dataset.

Method	Prediction time (hours)
Method	0 $\sim$ 1	1 $\sim$ 2	2 $\sim$ 3	3 $\sim$ 4	4 $\sim$ 5	5 $\sim$ 6
MeTok	52.15	48.39	44.83	42.05	39.49	37.16
Learnable	40.54	38.30	36.93	34.65	33.31	32.95

Appendix

Table 14: “MeTok" is our method, “Learnable" indicates the learning-based grouping method. The values in bold are better.

Figure 6: The ablation study of the proposed framework. We comprehensively analyzed the proposed framework, examining it from various facets.

Figure 7: The predicted performance across different geographical regions.

Precipitation Nowcasting

Precipitation nowcasting predicts rainfall over minutes to hours from radar, satellite, station, and numerical weather data.

See all 40 papers

Sep 29, 2026

NowcastDiT: Diffusion Transformers are Effective Precipitation Nowcasters

Precipitation nowcasting demands accurate short-term forecasts under strong spatiotemporal variability. Diffusion models are well suited to modeling complex precipitation distributions, yet existing approaches often introduce increasingly specialized designs, leaving the capability of a standard diffusion architecture underexplored. We show that a standard Diffusion Transformer already provides a simple and scalable foundation for precipitation nowcasting, with domain-specific requirements accommodated naturally within its design space. Based on this principle, we develop NowcastDiT and instantiate this flexibility through two complementary adaptations: a dynamics-aware noise prior for temporally coherent forecasts, and end-to-end reinforcement learning with timestep-aware rewards for meteorological skill. Experiments on SEVIR and MRMS benchmarks show that NowcastDiT achieves state-of-the-art performance in both perceptual quality and meteorological skill. These results suggest that standard DiT can serve as an effective foundation for precipitation nowcasting.

Haoran Xu, Xingzhuo Guo, Yuchen Zhang +3

School of Software, BNRist, Tsinghua University, China

Precipitation Forecasting Diffusion Transformer Probabilistic Forecasting Spatiotemporal Forecasting Precipitation Nowcasting

Ask PDF

Mar 14, 2026

PA-Net: Precipitation-Adaptive Mixture-of-Experts for Long-Tail Rainfall Nowcasting

Precipitation nowcasting is vital for flood warning, agricultural management, and emergency response, yet two bottlenecks persist: the prohibitive cost of modeling million-scale spatiotemporal tokens from multi-variate atmospheric fields, and the extreme long-tailed rainfall distribution where heavy-to-torrential events -- those of greatest societal impact -- constitute fewer than 0.1% of all samples. We propose the Precipitation-Adaptive Network (PA-Net), a Transformer framework whose computational budget is explicitly governed by rainfall intensity. Its core component, Precipitation-Adaptive MoE (PA-MoE), dynamically scales the number of activated experts per token according to local precipitation magnitude, channeling richer representational capacity toward the rare yet critical heavy-rainfall tail. A Dual-Axis Compressed Latent Attention mechanism factorizes spatiotemporal attention with convolutional reduction to manage massive context lengths, while an intensity-aware training protocol progressively amplifies learning signals from extreme-rainfall samples. Experiment on ERA5 demonstrate consistent improvements over state-of-the-art baselines, with particularly significant gains in heavy-rain and rainstorm regimes.

Xinyu Xiao, Sen Lei, Eryun Liu +5

Fudan University · the School of Information Science and Technology, Southwest Jiaotong University · Zhejiang University +2

Precipitation Nowcasting Mixture-of-Experts Models Adaptive Model Routing Long-Tail Learning Attention Mechanisms

Ask PDF

May 29, 2026

Probabilistic Precipitation Nowcasting with Rectified Flow Transformers

Accurate weather forecasts are essential across various domains and are safety-critical in extreme weather conditions. Compared to simulation-based forecasting, data-driven approaches show greater efficiency, enabling short-term, high-resolution nowcasting. In particular, diffusion models proved effective in weather nowcasting due to their strong probabilistic foundation. However, existing methods rely on deterministic compression to reduce the complexity of high-dimensional weather data, limiting their ability to capture uncertainty in the decoding process. In this work, we introduce

\textbf{FREUD}

, a

\textbf{Fr}

ame-wise

\textbf{E}

ncoder and

\textbf{U}

nited

\textbf{D}

ecoder model based on rectified flow transformers for efficient compression of spatio-temporal weather data. Frame-wise encoding enables continuous forecast updates, while the unified video decoder ensures temporal consistency. Our uncertainty-preserving first stage allows us to capture aleatoric uncertainty via ensembling, which is particularly beneficial for extreme weather events with high decoding variability. We achieve state-of-the-art performance in precipitation nowcasting with a compact latent-space rectified flow transformer on the SEVIR benchmark and show further performance gains by model and test-time scaling. Code available here: https://github.com/CompVis/weather-rf

Johannes Schusterbauer, Jannik Wiese, Nick Stracke +2

Precipitation Forecasting Rectified Flow Probabilistic Forecasting Precipitation Nowcasting

Ask PDF

Abstract

Precipitation Nowcasting

Explore similar work

NowcastDiT: Diffusion Transformers are Effective Precipitation Nowcasters

PA-Net: Precipitation-Adaptive Mixture-of-Experts for Long-Tail Rainfall Nowcasting

Probabilistic Precipitation Nowcasting with Rectified Flow Transformers