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

AI summary

Overview

Research area: Machine learning for weather forecasting — specifically efficient transformer architectures (vision-transformer-style weather foundation models) and adaptive token reduction (token merging/pruning).

Technical level: Intermediate. The paper assumes familiarity with transformer tokenization, attention cost scaling, and the basics of data assimilation, but the core idea (spend compute on unreliable, dynamically active regions) is intuitive.

Scope: The paper proposes and evaluates UtoMe, a token-merging policy for ERA5 72-hour global forecasting that combines atmospheric dynamics saliency with ensemble data assimilation spread, across two spatial resolutions, three uncertainty-availability variants, and matched token budgets.

What This Paper Is About

Weather transformers process thousands of dense spatial tokens per forecast, which makes training and inference expensive, and token merging is a standard way to cut that cost by combining similar tokens. The problem is that similarity between hidden features can mean either "this region is redundant" or "this region is poorly observed," and a similarity-only rule cannot tell the two apart. UtoMe's goal is to allocate a limited token budget using a weather-specific signal — how dynamically active a patch is, minus how uncertain the input observation of that patch is — so that forecast-relevant structure survives compression.

Key Contributions

1. A formulation of observation-uncertainty-guided token allocation for weather transformers, using ECMWF Ensemble of Data Assimilation (EDA) spread as a relative proxy for input reliability rather than an absolute error estimate, motivated by a three-way variance decomposition that separates aleatoric, epistemic, and observation-uncertainty terms.
2. A concrete allocation mechanism — a patch-level importance score s_i = d_i − β q_i combining dynamics saliency (gradients plus anomaly amplitude over five forecast variables) with normalized spread, plus anchor protection of the top ρ = 0.25 fraction of tokens and importance-calibrated bipartite matching with penalties λ_s = λ_Δ = 0.05 — with the pieces isolated through controlled ablations.
3. Three variants covering different uncertainty-availability regimes: UtoMe-E (external EDA spread at training and inference), UtoMe-P (a lightweight predictor distilled from the EDA-derived score, used at inference without uncertainty fields), and UtoMe-S (learned from a local 3×3 spatial-inconsistency proxy, no external uncertainty at all).
4. An empirical evaluation on ERA5 72-hour forecasting at 5.625° and 1.40625° including spatial uncertainty controls, difficult-regime diagnostics, token-reduction sweeps, pruning comparisons, and a second transformer backbone.

Main Findings

• Joint scoring beats similarity-only merging at matched budgets. At 1.40625° and 50% token reduction, UtoMe-E lowers Overall error from ToMe's 67.32 to 53.00 at the same 0.55× encoder FLOPs; at 25% reduction the corresponding numbers are 42.45 (ToMe) versus 39.59 (UtoMe-E), both at 0.86×. The paper states UtoMe-E lowers ToMe's Overall error by 2.86 and 14.32 at the two budgets.
• Low-resolution results follow the same pattern. At 5.625°, 25% reduction: no merging 143.07, ToMe 160.82, UtoMe-P 144.10, UtoMe-E 144.74, UtoMe-S 145.36. At 50%: ToMe 220.05 versus UtoMe-E 151.95, UtoMe-P 152.42, UtoMe-S 167.63.
• Gains are not parity with the full-token model. At 1.40625° and 25% reduction, Z500 wRMSE rises from 156.76 to 191.93 (22.4%), and the paper explicitly notes the gains are not uniform across targets — ToMe has lower U10/V10 errors at low-resolution 25% reduction and lower V10 error than E/P at high-resolution 50% reduction.
• Compression has a measurable accuracy cost. At 1.40625°: 5% reduction raises Overall from 32.41 to 32.80 (1.2%) for a 3% FLOPs cut; 10% raises error by 6.4% for 7% FLOPs; 25% gives a 14% FLOPs saving for a 22.2% error increase; 50% saves 45% of encoder FLOPs but increases Overall error by 63.5%.
• Every component contributes. At 1.40625°, removing UQ increases Overall wRMSE by 1.36/5.60 at 25%/50% reduction, removing anchors costs 0.72/8.40, and removing the matching penalties costs 1.09/2.90 (full UtoMe-E: 39.59/53.00; ToMe: 42.45/67.32).
• Spatial alignment matters more than magnitude rescaling. Square-root, squared, and rank-mapped versions of the UQ field change wRMSE by at most 0.13/0.28, while spatially shuffling UQ degrades wRMSE to 41.88/62.30 and drops anchor overlap to 26%.
• External spread is optional at deployment. On three held-out 5.625° cases, UtoMe-P's predicted scores correlate 0.769–0.825 with the EDA-derived teacher and UtoMe-S reaches 0.858–0.900 against its own proxy target — score agreement, not calibrated uncertainty.
• Gains concentrate in dynamically active regions. At 1.40625° and 25% reduction, UtoMe-E cuts Z500 wRMSE in the high-UQ/high-dynamics regime from ToMe's 268.4 to 220.0 (18.0%), but errors rise in the two low-dynamics regimes (e.g., 158.2 → 176.0 in low-UQ/low-dynamics). At 50% reduction it lowers upper-tail Z500, lower-tail Z500, and upper-tail T2m errors by 22.4%, 22.7%, and 16.9% relative to ToMe.
• The policy demonstrably changes which tokens are merged. At 25% reduction, merge-edge Jaccard overlap with cosine-only matching is 0.139, and merged-token membership in the top importance quartile falls from 46.8% to 0%.
• The approach transfers. On a plain 12-layer ViT, UtoMe-E improves over ToMe by 2.7%/4.9% at matched FLOPs, and it also beats score-based pruning at both budgets.

Methodology in Plain English

The researchers start from a standard ViT-style weather transformer that chops the globe into patches, turns each patch into a token, and runs every token through every layer. Token merging shrinks the sequence by fusing similar tokens, which saves compute.

Their key argument is that in weather, "similar features" can be misleading, because reanalysis inputs are estimates built from uneven observations. They write down a variance decomposition showing that forecast variance has three parts: irreducible noise, model uncertainty, and uncertainty about the true input state — the third being the one standard treatments ignore.

To act on that, they build a per-patch score with two ingredients. The first is dynamics saliency, computed only from input-time fields: horizontal and vertical grid differences plus deviation from the spatial mean, averaged over the five forecast variables. The second is observation uncertainty, taken from EDA ensemble spread, normalized per variable, averaged over each patch, and standardized per sample. The spread is subtracted from saliency with weight β = 1, so sharp, confidently observed patches get the strongest protection while high-spread patches become better merge candidates unless their dynamics are strong enough to compensate.

The merge step keeps the standard ToMe bipartite matching but changes two things: the top 25% of tokens by score become permanent anchors excluded from both merge pools, and the matching score for the remaining tokens is cosine similarity minus a penalty on the pair's summed importance and a penalty on their importance difference. The highest-scoring edges merge until the layer's budget is met, with features averaged and sizes tracked as in ToMe.

For cases where uncertainty fields are unavailable, UtoMe-P trains a small predictor on early hidden tokens against the EDA-derived score with a pairwise-ranking loss (weight 0.2), and UtoMe-S substitutes a hand-built proxy — each pixel's deviation from its 3×3 average — and trains the same predictor on that target. Experiments use two ERA5 resolutions (512 and 2048 initial tokens), 41 input channels from the ClimaX variable set, 72-hour lead times, five target variables, three random seeds, and up to 64 AMD MI250X GPUs.

Why This Matters

Impact on research. Efficient transformer work in vision and language treats every token as equally trustworthy. This paper argues that in geoscience the data source itself carries spatially structured reliability information that should shape the compute budget — a coupling between uncertainty quantification and architecture efficiency that has not been explored much. It also provides a careful, honest accounting: improvements over ToMe coexist with real losses versus full-token models, and the authors flag that Overall mixes physical units and is numerically dominated by Z500.

Real-world applications (as implications of cheaper, budget-aware forecasting):

• Operational medium-range forecasting, where encoder cost limits ensemble size or update frequency.
• Renewable energy planning and grid operations, which depend on 10 m wind and 2 m temperature forecasts of the kind used here as targets.
• Agricultural and food-security monitoring, which relies on repeated global temperature and circulation forecasts.
• Disaster preparedness and aviation, where mid-tropospheric circulation (Z500) and lower-tropospheric temperature (T850) drive the guidance products downstream users consume.

Industry relevance. Any organization training or serving large Earth-system foundation models faces the same cost wall; a drop-in merge policy that reduces encoder FLOPs to 0.55× at 50% token reduction, or 0.86× at 25%, is attractive for both training budgets and inference latency. The three variants map cleanly onto real deployment constraints: some pipelines have ensemble spread available, some have it only in the training archive, and some have neither.

Future Directions

• Recovering the low-dynamics losses. UtoMe-E improves high-UQ/high-dynamics and low-UQ/high-dynamics regions but degrades the two low-dynamics regimes; a policy that protects quiet regions from over-compression is an open problem the authors flag.
• Better uncertainty signals. EDA spread is a relative proxy that may miss shared model biases and fine-scale extremes, and the two-member simulation in Table 5 already degrades results (39.94/54.10 versus 39.59/53.00). Calibrated or learned token-level uncertainty estimators are untested.
• Beyond 72-hour direct forecasts. The paper covers direct 72-hour ERA5 forecasts on two transformer backbones, without establishing transfer to other architectures or to long autoregressive rollouts.
• End-to-end cost accounting. Encoder FLOPs are the efficiency metric here; wall-clock training time, memory, and full-pipeline latency at higher resolutions remain unmeasured, as does the cost of the predictor probe in the P/S variants.

Target Audience

Researchers and engineers working on efficient transformer architectures who want a domain-grounded alternative to similarity-only token merging; weather and climate ML practitioners building or serving foundation models on ERA5-style reanalysis; and uncertainty-quantification researchers interested in turning a diagnostic quantity into an architectural control signal. Readers should be comfortable with attention mechanics and the vocabulary of data assimilation, but the paper's framing — protect what is important and what is uncertain — is accessible to graduate students entering either field.

Authors’ 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.