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

Weather foundation models are trained on dense reanalysis fields, but these inputs are not noiseless measurements of the atmosphere. Observation and representativeness errors vary across variables, locations, and weather regimes, so a transformer can spend substantial compute fitting token-level fluctuations that are less reliable than the surrounding signal. We present UtoMe, an observation-uncertainty-guided token merging framework for training-efficient weather transformers. Motivated by a three-way variance view that separates observation uncertainty from aleatoric and epistemic uncertainty, UtoMe builds a patch-level importance score from dynamics saliency and ensemble-derived spread, protects forecast-sensitive low-uncertainty tokens as anchors, and merges low-importance feature-similar tokens with an uncertainty-aware matching cost. Three variants cover different deployment settings: external uncertainty at both training and inference, external uncertainty only during training, and no external uncertainty. On ERA5 forecasting at two spatial resolutions, UtoMe preserves forecast skill under moderate token reduction and degrades much more gracefully than train-time ToMe under aggressive merging. The method reduces FLOPs by up to $45\%$, showing that observation uncertainty can be used as a direct compute-allocation signal for weather foundation models.