UtoMe: Observation-Uncertainty-Guided Token Merging for Weather Foundation Models

Yunbei Zhang; Janet Wang; Xi Xiao; Jihun Hamm; Xiao Wang

cs.LGOct 18, 2026

UtoMe: Observation-Uncertainty-Guided Token Merging for Weather Foundation Models

Authors: Yunbei Zhang, Janet Wang, Xi Xiao, Jihun Hamm, Xiao Wang

Organizations: Tulane University · Oak Ridge National Laboratory

Abstract

Weather transformers process dense reanalysis fields whose reliability varies across variables, locations, and weather regimes. Token merging can reduce computation, but feature similarity alone does not account for this heterogeneous input uncertainty. We introduce UtoMe, an observation-uncertainty-guided token merging framework that combines dynamics saliency with ensemble data assimilation spread as a relative input-reliability proxy. The resulting importance score protects selected tokens as anchors and calibrates bipartite matching between the remaining tokens. Three variants use external uncertainty at both training and inference, during training only, or neither stage. In ERA5 72-hour forecasting at two spatial resolutions, UtoMe improves Overall error over train-time ToMe at matched token budgets. At $1.40625^\circ$ and 50% token reduction, UtoMe-E lowers Overall error from 67.32 to 53.00 at the same 0.55x encoder FLOPs. Small reductions remain close to the full-token baseline; larger reductions incur measurable accuracy losses. Component ablations and spatial uncertainty controls support combining dynamics with aligned uncertainty when allocating a weather transformer's token budget.

Figures & tables

Figure 1: UtoMe overview. UtoMe uses observation uncertainty to guide training-time token merging for transformer-based weather forecasting. It first turns weather-state saliency and low observation uncertainty into a patch-level token score, then uses that score to protect forecast-sensitive anchors and merge lower-importance redundant tokens. The framework supports three uncertainty-availability regimes: external uncertainty at both stages, during training only, or neither stage.

Figure 2: $5.625^{\circ}$ case summary. The figure shows the complete merge-decision pipeline: weather state, observation uncertainty, dynamics saliency, importance score, patch-level token score, and the final anchor or merge-candidate overlay.

Method	Overall $\downarrow$	ACC $\uparrow$	Z500 $\downarrow$	T850 $\downarrow$	T2m $\downarrow$	U10 $\downarrow$	V10 $\downarrow$	Compute $\downarrow$
No merging	$143.07_{\pm 1.81}$	$0.651_{\pm 0.008}$	$701.64_{\pm 8.92}$	$3.108_{\pm 0.023}$	$2.502_{\pm 0.008}$	$3.944_{\pm 0.024}$	$4.151_{\pm 0.051}$	$1.00\times$
25% token reduction
ToMe	$160.82_{\pm 2.07}$	$0.562_{\pm 0.011}$	$789.00_{\pm 10.34}$	$3.544_{\pm 0.041}$	$3.378_{\pm 0.038}$	$\mathbf{4.041}_{\pm 0.047}$	$\mathbf{4.142}_{\pm 0.044}$	$0.86\times$
UtoMe-E	$\underline{144.74}_{\pm 1.42}$	$\underline{0.640}_{\pm 0.007}$	$\underline{709.50}_{\pm 7.21}$	$\underline{3.220}_{\pm 0.026}$	$\underline{2.730}_{\pm 0.021}$	$4.060_{\pm 0.032}$	$4.190_{\pm 0.039}$	$0.86\times$
UtoMe-P	$\mathbf{144.10}_{\pm 1.16}$	$\mathbf{0.650}_{\pm 0.006}$	$\mathbf{706.34}_{\pm 6.43}$	$\mathbf{3.210}_{\pm 0.024}$	$\mathbf{2.717}_{\pm 0.019}$	$\underline{4.054}_{\pm 0.030}$	$\underline{4.182}_{\pm 0.034}$	$0.87\times$
UtoMe-S	$145.36_{\pm 1.69}$	$0.635_{\pm 0.009}$	$712.50_{\pm 8.18}$	$3.260_{\pm 0.031}$	$2.760_{\pm 0.024}$	$4.075_{\pm 0.036}$	$4.205_{\pm 0.041}$	$0.87\times$

Table 1: Main results on

5.625^{\circ}

ERA5 72-hour forecasting. The table compares no merging, train-time ToMe, and the three UtoMe variants under matched final token-reduction ratios. Values are mean and standard deviation over three random seeds. Lower is better for Overall and per-variable wRMSE, higher is better for ACC, and Compute is the encoder FLOPs ratio relative to no merging at the same resolution; the main resolution tables follow this convention. Green shades identify E/P/S uncertainty availability; blue bands identify token budgets. Bold/underlined values are best/second-best within each budget.

Method	Overall $\downarrow$	ACC $\uparrow$	Z500 $\downarrow$	T850 $\downarrow$	T2m $\downarrow$	U10 $\downarrow$	V10 $\downarrow$	Compute $\downarrow$
No merging	$32.41_{\pm 0.28}$	$0.964_{\pm 0.004}$	$156.76_{\pm 1.34}$	$1.299_{\pm 0.007}$	$1.269_{\pm 0.006}$	$1.339_{\pm 0.008}$	$1.360_{\pm 0.009}$	$1.00\times$
25% token reduction
ToMe	$42.45_{\pm 0.24}$	$0.949_{\pm 0.003}$	$206.02_{\pm 1.17}$	$1.452_{\pm 0.006}$	$1.483_{\pm 0.007}$	$1.572_{\pm 0.008}$	$1.724_{\pm 0.009}$	$0.86\times$
UtoMe-E	$\mathbf{39.59}_{\pm 0.21}$	$\mathbf{0.952}_{\pm 0.003}$	$\mathbf{191.93}_{\pm 1.05}$	$\mathbf{1.413}_{\pm 0.005}$	$\mathbf{1.454}_{\pm 0.006}$	$\mathbf{1.527}_{\pm 0.007}$	$\mathbf{1.645}_{\pm 0.008}$	$0.86\times$
UtoMe-P	$\underline{39.61}_{\pm 0.01}$	$\mathbf{0.952}_{\pm 0.000}$	$\underline{191.98}_{\pm 0.03}$	$1.425_{\pm 0.001}$	$1.459_{\pm 0.000}$	$\underline{1.530}_{\pm 0.000}$	$\underline{1.648}_{\pm 0.001}$	$0.87\times$
UtoMe-S	$39.84_{\pm 0.27}$	$\underline{0.951}_{\pm 0.004}$	$193.14_{\pm 1.31}$	$\underline{1.416}_{\pm 0.007}$	$\underline{1.458}_{\pm 0.008}$	$1.533_{\pm 0.009}$	$1.656_{\pm 0.010}$	$0.87\times$

Table 2: Main results on

1.40625^{\circ}

ERA5 72-hour forecasting. Values are mean and standard deviation over three seeds; conventions follow Table 1 .

Panel of Figure 3 — Figure 3: $1.40625^{\circ}$ high-resolution case study. The observation-uncertainty field contains coherent bands and high-spread structures. After saliency combination and patch aggregation, these structures produce a heterogeneous token score for uncertainty-guided merging.

Method	Token Red.	Overall $\downarrow$	ACC $\uparrow$	Z500 $\downarrow$	T850 $\downarrow$	T2m $\downarrow$	U10 $\downarrow$	V10 $\downarrow$	Compute $\downarrow$
No merging	$0\%$	$32.41_{\pm 0.28}$	$0.964_{\pm 0.004}$	$156.76_{\pm 1.34}$	$1.299_{\pm 0.007}$	$1.269_{\pm 0.006}$	$1.339_{\pm 0.008}$	$1.360_{\pm 0.009}$	$1.00\times$
UtoMe	$5\%$	$32.80_{\pm 0.18}$	$0.964_{\pm 0.002}$	$158.70_{\pm 0.94}$	$1.305_{\pm 0.005}$	$1.278_{\pm 0.005}$	$1.349_{\pm 0.006}$	$1.372_{\pm 0.007}$	$0.97\times$
UtoMe	$10\%$	$34.49_{\pm 0.25}$	$0.961_{\pm 0.003}$	$167.00_{\pm 1.16}$	$1.330_{\pm 0.006}$	$1.310_{\pm 0.006}$	$1.380_{\pm 0.007}$	$1.410_{\pm 0.008}$	$0.93\times$
UtoMe	$25\%$	$39.59_{\pm 0.21}$	$0.952_{\pm 0.003}$	$191.93_{\pm 1.05}$	$1.413_{\pm 0.005}$	$1.454_{\pm 0.006}$	$1.527_{\pm 0.007}$	$1.645_{\pm 0.008}$	$0.86\times$
UtoMe	$50\%$	$53.00_{\pm 0.36}$	$0.913_{\pm 0.004}$	$256.72_{\pm 1.76}$	$1.759_{\pm 0.009}$	$2.051_{\pm 0.012}$	$2.038_{\pm 0.013}$	$2.433_{\pm 0.015}$	$0.55\times$

Table 3: Token-reduction ratio ablation at

1.40625^{\circ}

. UtoMe corresponds to UtoMe-E.

Figure 4: Cosine-only matching can select merge edges in high-importance regions because feature similarity alone does not encode observation uncertainty. UtoMe protects anchors and penalizes importance-incompatible pairs, shifting the merge budget toward lower-importance regions.

Configuration	25% reduction	50% reduction
ToMe	42.45	67.32
Dynamics only ( $\beta=0$ )	40.95	58.60
No anchors	40.31	61.40
No UQ-aware matching	40.68	55.90
Full UtoMe-E	39.59	53.00

Table 4: UQ, anchors, and matching contribute at both budgets. Overall wRMSE (

\downarrow

) at

1.40625^{\circ}

with matched token schedules.

UQ field	Rank corr.	Anchor overlap	25%	50%
Original EDA	1.000	100%	39.59	53.00
$\sqrt{\mathrm{UQ}}$	0.995	94%	39.62	53.05
$\mathrm{UQ}^{2}$	0.977	88%	39.68	53.19
Rank-mapped UQ	0.969	84%	39.72	53.28
Simulated 2-member	0.960	79%	39.94	54.10
Constant UQ	–	–	40.95	58.60

Table 5: Spatial alignment matters more than the tested rescalings. Score rank correlation and anchor overlap are relative to the EDA-derived score; errors are Overall (

\downarrow

) at

1.40625^{\circ}

. The two-member row uses simulated sampling noise.

Method	$q$ L/ $d$ H	$q$ H/ $d$ H	$q$ L/ $d$ L	$q$ H/ $d$ L	Global
No merging	172.0	200.0	116.0	124.0	156.76
ToMe	219.8	268.4	158.2	155.7	206.02
UtoMe-E	185.2	220.0	176.0	181.0	191.93

Table 6: The gain concentrates in dynamically active regions, including high-UQ regions. Z500 wRMSE (

\downarrow

) at

1.40625^{\circ}

and

25\%

token reduction. L/H denote below/at-or-above the per-sample median of UQ (

q

) and dynamics (

d

). The two low-dynamics regimes expose the allocation trade-off.

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

Configuration	25%	50%
ToMe (neither mechanism)	42.45	67.32
Dynamics only ( $s=d$ )	40.95	58.60
UQ only ( $s=-q$ )	41.40	61.70
Combined score, no anchors	40.31	61.40
Anchors, no matching penalties	40.68	55.90
Full UtoMe-E	39.59	53.00

Appendix

Table 7: Both score inputs and both allocation mechanisms contribute. Overall wRMSE (

\downarrow

) at

1.40625^{\circ}

. The no-anchor and no-matching rows retain the combined score. Disabling both allocation mechanisms recovers ToMe.

Parameter	Value	25%	50%
$\beta$	0	40.95	58.60
	0.25	40.12	55.10
	0.5	39.74	53.62
	1	39.59	53.00
	2	40.36	54.85
	4	43.80	69.40

Appendix

Table 8: Moderate parameter values balance protection and compression. One-parameter sweeps at

1.40625^{\circ}

, with other settings held at the default. Columns report Overall wRMSE (

\downarrow

). The validation-selected default values are

\beta=1

\rho=0.25

, and

\lambda_{s}=\lambda_{\Delta}=0.05

UQ field	Rank corr.	Anchor overlap	25%	50%
Original EDA	1.000	100%	39.59	53.00
$\sqrt{\mathrm{UQ}}$	0.995	94%	39.62	53.05
$\mathrm{UQ}^{2}$	0.977	88%	39.68	53.19
Rank-mapped UQ	0.969	84%	39.72	53.28
Simulated 2-member	0.960	79%	39.94	54.10
Constant UQ	–	–	40.95	58.60

Appendix

Table 9: Tested rescalings preserve useful spatial order. Score rank correlation and anchor overlap are relative to the original EDA field. The last two columns are Overall wRMSE (

\downarrow

) at

1.40625^{\circ}

. Dashes mark agreement metrics not reported for the constant field.

Control	25%	50%
Original EDA spread	39.59	53.00
Constant UQ / dynamics only	40.95	58.60
Spatially shuffled UQ	41.88	62.30
Matched random-noise UQ	42.07	63.15
Inverted sign ( $s=d+\beta q$ )	42.64	64.90

Appendix

Table 10: Removing spatial information or reversing the UQ direction degrades performance. Overall wRMSE (

\downarrow

) at

1.40625^{\circ}

Simulated members	Rank corr.	25%	50%
Full reference	1.000	39.59	53.00
32	0.999	39.60	53.02
16	0.997	39.62	53.08
8	0.994	39.66	53.21
4	0.986	39.76	53.55
2	0.960	39.94	54.10

Appendix

Table 11: Finite-ensemble sampling noise causes gradual degradation. This experiment simulates noisier spread estimates, rather than subsampling archived raw ensemble members. Rank correlation is measured against the reference importance score; errors are Overall wRMSE at

1.40625^{\circ}

Method	Reduction	Z500 upper	Z500 lower	T2m upper
No merging	0%	189.9	242.5	1.953
ToMe	25%	251.4	318.7	2.318
UtoMe-E	25%	226.8	285.3	2.147
ToMe	50%	402.6	512.4	3.145
UtoMe-E	50%	312.5	396.1	2.612

Appendix

Table 12: UtoMe-E better retains forecast extremes than ToMe at both budgets. Conditional wRMSE (

\downarrow

) on the upper/lower

5\%

target-anomaly masks at

1.40625^{\circ}

. Z500 errors use

\mathrm{m^{2}\,s^{-2}}

and T2m errors use K.

Method	Reduction	$\sigma$ ratio	$\nabla$ ratio	KE ratio	Vort. ratio
No merging	0%	0.998	1.013	0.975	0.870
ToMe	25%	0.964	1.092	0.931	0.812
UtoMe-E	25%	0.981	1.048	0.958	0.846
ToMe	50%	0.907	1.213	0.874	0.731
UtoMe-E	50%	0.955	1.104	0.936	0.802

Appendix

Table 13: Spatial and wind statistics remain closer to the full-token reference. Prediction/reference ratios at

1.40625^{\circ}

; unity denotes agreement with the reference statistic. The spatial standard-deviation and gradient columns concern Z500; KE and vorticity use U10/V10.

Method	25%	50%
Random pruning	61.86	85.90
Dynamics-saliency pruning	44.42	61.02
UQ-only pruning	43.38	59.77
UtoMe-score pruning	40.74	57.21
ToMe merging	42.45	67.32
UtoMe-E merging	39.59	53.00

Appendix

Table 14: The joint score benefits pruning, and merging preserves more forecast skill. Overall wRMSE (

\downarrow

) at

1.40625^{\circ}

, using the matched backbone, split, schedule, and token budgets of the comparison. Encoder FLOPs are

0.86\times

0.55\times

25\%

50\%

reduction.

Method	Reduction	wRMSE	Encoder FLOPs
No merging	0%	141.96	$1.000\times$
ToMe	25%	145.05	$0.829\times$
UtoMe-E	25%	141.13	$0.829\times$
ToMe	50%	157.14	$0.661\times$
UtoMe-E	50%	149.44	$0.661\times$

Appendix

Table 15: The allocation rule transfers to a plain patch-transformer backbone. Overall wRMSE (

\downarrow

) on a 12-layer ViT with a joint patch embedding. FLOPs are relative to this backbone without merging.

Case	UtoMe-P corr. $\uparrow$	UtoMe-P zMAE $\downarrow$	UtoMe-S corr. $\uparrow$	UtoMe-S zMAE $\downarrow$
$5.625^{\circ}$ case 01	$0.820$	$0.466$	$0.858$	$0.403$
$5.625^{\circ}$ case 02	$0.825$	$0.463$	$0.900$	$0.341$
$5.625^{\circ}$ case 03	$0.769$	$0.525$	$0.870$	$0.396$

Appendix

Table 16: Predictor quality on held-out

5.625^{\circ}

cases. The case label denotes the spatial resolution, not a separate dataset. Correlation and z-normalized MAE are computed patch by patch between the predictor output and its supervision target.

Diagnostic	ToMe	UtoMe
Merged tokens in top quartile	46.8%	0%
Mean merged-token importance	+0.732	-1.456
Displayed-pair mean (Fig. 4 )	0.28	0.11

Appendix

Table 17: UtoMe changes the selected merge set. Allocation diagnostics at

25\%

token reduction. Quartile membership and merged-token importance refer to the standardized importance score.

Panel of Figure 5 — Figure 5: Motivation for uncertainty-guided token merging at $1.40625^{\circ}$ . The observation-uncertainty field is not uniform noise; it contains coherent spatial structure. After patch aggregation and combination with dynamics saliency, the token-score grid separates low-uncertainty high-saliency regions from high-spread merge candidates.

Figure 6: Additional importance-score decomposition. From left to right, the figure shows the input weather state, dynamics saliency, observation uncertainty, teacher importance, and predicted importance for a held-out $5.625^{\circ}$ case.

Figure 7: Predictor visualization for UtoMe-P and UtoMe-S. The learned predictors recover the broad spatial structure of their supervision targets, which explains why the predictor-based variants remain close to UtoMe-E in the main results.

Figure 8: Additional anchor and merge-candidate overlay at $5.625^{\circ}$ . Blue squares denote merge candidates and black squares denote anchors selected from the patch-level token score.

Panel of Figure 9 — Figure 9: Additional $1.40625^{\circ}$ case study. The panels show how a high-resolution weather state and its aligned observation-uncertainty field are converted into the importance and token-score maps used by UtoMe.

Token Merging

Token merging combines redundant transformer tokens to reduce sequence length, attention costs, and inference time while preserving task-relevant information.

See all 22 papers

Less Tokens, Better Forecasts: Sparse Residual Routing for Efficient Weather Prediction
Jul 2, 2026Janet Wang, Yunbei Zhang, Lin Zhao +3Efficient ViTs Dense Prediction
Tulane University · Northeastern University · University of Alabama at Birmingham +1
Observation Assimilation Transformer: Probabilistic Weather Analysis from Sparse Station Observations
Oct 20, 2026Christian Donner, Michele Cattaneo, Maxim Samarin +4Data Assimilation Probabilistic Forecasting
Swiss Data Science Center, Eidgenössische Technische Hochschule Zürich (ETHZ) · Federal Office of Meteorology and Climatology MeteoSwiss, Switzerland
MeTok: An Efficient Meteorological Tokenization with Hyper-Aligned Group Learning for Precipitation Nowcasting
Mar 14, 2026Qizhao Jin, Xianhuang Xu, Yong Cao +2Precipitation Nowcasting Transformer
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

Jul 2, 2026J/K move · Enter open · S save

Less Tokens, Better Forecasts: Sparse Residual Routing for Efficient Weather Prediction

Janet Wang, Yunbei Zhang, Lin Zhao, Xi Xiao+2

Tulane University · Northeastern University · University of Alabama at Birmingham+1

Efficient ViTs Dense Prediction Spatiotemporal Forecasting Weather Forecasting Spatiotemporal Transformers

Open paper AskPDF arXiv

Existing ViT-based weather forecasting models apply uniform computation across all spatial tokens, even though nearby atmospheric grid points often contain similar values and large regions evolve smoothly over time. This makes much of the intermediate per-token computation redundant. Standard token-efficiency methods, such as pruning or merging, reduce cost by removing or fusing tokens. However, weather forecasting is a spatiotemporal dense prediction problem in which a history of atmospheric states must be mapped to future values on the original latitude-longitude grid. Thus, every grid cell must retain a physically meaningful representation, especially under autoregressive rollout. We introduce Sparse-Reslim, a parameter-free plug-in routing module that makes sparse token processing compatible with this fixed-grid requirement. Sparse-Reslim routes only 25% of spatial tokens through the expensive middle transformer blocks and treats those blocks as residual updates: it computes the change produced for the routed tokens and scatters only this delta back to the full sequence. Unselected tokens keep their pre-routing representations exactly, so no grid cell is dropped or replaced by a mask token, and no fusion layer or additional parameters are introduced. Across ERA5 resolutions up to the operational 0.25\textdegree{} standard and two model families, a deterministic Transformer and a diffusion model, Sparse-Reslim improves forecast accuracy on every evaluated variable while substantially reducing cost: training is about 2.5x faster in the main settings and reaches 3.18x speedup at 0.25\textdegree{}, with over 2.2x lower peak memory. A controlled decomposition shows that the accuracy gain comes primarily from sparse routing itself, while random token selection provides an additional regularization benefit without selector overhead.

Abstract

Token Merging

Explore similar work

Less Tokens, Better Forecasts: Sparse Residual Routing for Efficient Weather Prediction

Observation Assimilation Transformer: Probabilistic Weather Analysis from Sparse Station Observations

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