Weather model speeds up forecasting by choosing tokens wisely
UtoMe: Observation-Uncertainty-Guided Token Merging for Weather Foundation Models
Machine LearningArtificial IntelligenceComputer Vision and Pattern Recognition
Summary
Weather forecasts use big computer models that look at lots of weather data, but not all data points are equally reliable. The paper introduces UtoMe, a method that helps the computer model pick important pieces of data (tokens) based on how certain the observations are and their weather dynamics. This smarter selection reduces the computing work while keeping forecast accuracy good, especially for big token reductions. The authors tested this on a common weather dataset and found their method improved forecast errors compared to earlier token merging approaches.
What this means in practice
- •For weather forecasting teams: Run weather transformer models faster by merging tokens with uncertainty guidance to save computing resources while maintaining forecast quality.
- •For climate data engineers: Optimize processing of large-scale weather data by selectively focusing on more reliable and dynamically important tokens to balance speed and accuracy.
Authors
Yunbei Zhang, Janet Wang, Xi Xiao, Jihun Hamm, Xiao Wang
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.