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

NowcastDiT: Diffusion Transformers are Effective Precipitation Nowcasters
Sep 29, 2026Haoran Xu, Xingzhuo Guo, Yuchen Zhang +3Precipitation Forecasting Diffusion Transformer
School of Software, BNRist, Tsinghua University, China
PA-Net: Precipitation-Adaptive Mixture-of-Experts for Long-Tail Rainfall Nowcasting
Mar 14, 2026Xinyu Xiao, Sen Lei, Eryun Liu +5Precipitation Nowcasting Mixture-of-Experts Models
Fudan University · the School of Information Science and Technology, Southwest Jiaotong University · Zhejiang University +2
Probabilistic Precipitation Nowcasting with Rectified Flow Transformers
May 29, 2026Johannes Schusterbauer, Jannik Wiese, Nick Stracke +2Precipitation Forecasting Rectified Flow

Sep 29, 2026J/K move · Enter open · S save

NowcastDiT: Diffusion Transformers are Effective Precipitation Nowcasters

Haoran Xu, Xingzhuo Guo, Yuchen Zhang, Jincheng Zhong+2

School of Software, BNRist, Tsinghua University, China

Precipitation Forecasting Diffusion Transformer Probabilistic Forecasting Spatiotemporal Forecasting Precipitation Nowcasting

Open paper AskPDF arXiv

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.

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