🤖 machine learning

(Sparse) Attention to the Details: Preserving Spectral Fidelity in ML-based Weather Forecasting Models

The paper introduces Mosaic, a probabilistic weather forecasting model that utilizes block-sparse attention and learned functional perturbations to preserve spectral fidelity and generate well-calibrated ensembles at 1.5° resolution, achieving state-of-the-art performance and rapid inference speeds while overcoming common spectral degradation issues in machine learning-based predictions.

Original authors: Maksim Zhdanov, Ana Lucic, Max Welling, Jan-Willem van de Meent

Published 2026-04-21

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Maksim Zhdanov, Ana Lucic, Max Welling, Jan-Willem van de Meent

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to predict the weather for the next 10 days. For decades, scientists have used massive supercomputers to solve complex physics equations (like fluid dynamics) to do this. It's accurate, but it's slow and expensive—like trying to paint a masterpiece by hand, one tiny brushstroke at a time.

Recently, AI models have stepped in. They are like lightning-fast artists who can guess the weather in seconds. But there's a catch: they tend to "blur" the picture. They predict the average weather perfectly, but they miss the sharp, chaotic details like the exact eye of a hurricane or the sudden formation of a storm front. They smooth out the rough edges, making the world look too calm.

This paper introduces MOSAIC, a new AI weather model designed to fix this "blurry vision" problem. Here's how it works, explained with everyday analogies.

The Two Problems MOSAIC Solves

The authors identified two main reasons why AI weather models get the details wrong:

The "Average" Trap: Most AI models are trained to predict the average of all possible outcomes. If you ask 100 people to guess the height of a wave, and you average their answers, you get a smooth, safe number. But nature isn't smooth; it's jagged. By aiming for the average, the AI creates a "flat" forecast that misses the sharp peaks and valleys of real weather.
The "Compression" Bottleneck: To save time, many AI models squish high-resolution weather data into a smaller, coarser grid before processing it. Imagine trying to describe a high-definition 4K movie by only looking at a low-resolution thumbnail. You lose the fine details. The model tries to guess the missing details later, but the information is already gone.

The MOSAIC Solution: A New Way to Look at the World

MOSAIC tackles these issues with two clever tricks:

1. The "Group Chat" Attention (Block-Sparse Attention)

Traditional AI models look at every single piece of data (every pixel of the map) and try to figure out how it relates to every other piece. This is like trying to have a conversation with 10 million people at once—it's impossible and slow.

MOSAIC uses a method called Block-Sparse Attention.

The Analogy: Imagine you are organizing a massive global conference. Instead of asking every single attendee to talk to everyone else, you group them into "neighborhoods" (blocks).
How it works: A group of people in Florida (a "block") talks to other groups in the Atlantic and the Arctic. They don't talk to every single person individually; they talk to the groups.
The Benefit: This allows the model to see long-distance connections (like how a storm in the Pacific affects Europe) without getting bogged down in math. Crucially, it does this without squashing the data. It looks at the map at its full, high-definition resolution, preserving every tiny detail of the storm's structure.

2. The "Creative Ensemble" (Probabilistic Forecasting)

Instead of giving you one single "best guess" forecast, MOSAIC generates a whole team of 24 different forecasters (an "ensemble").

The Analogy: Imagine asking 24 different meteorologists to predict the weather. Instead of averaging their answers immediately, you let them each add a tiny bit of their own "creative chaos" based on what they see.
How it works: MOSAIC injects a little bit of learned randomness into its brain. This forces it to generate 24 slightly different, realistic scenarios.
The Benefit: When you look at these 24 scenarios, they all have sharp, realistic details (like a distinct hurricane eye). When you average them after the fact, you get a smooth average, but the individual members retain the "crispness" of real weather. This means the model doesn't just predict the average weather; it predicts the range of possible wild weather.

Why This Matters: The Hurricane Test

The paper tested MOSAIC on Hurricane Ian, a massive Category 4 storm.

The Challenge: Hurricanes are chaotic and happen on a scale that is hard for computers to see clearly.
The Result: Even though MOSAIC was trained on a grid that wasn't the highest possible resolution, it managed to track the hurricane's path and intensity with incredible accuracy. It didn't just say "a storm is coming"; it showed the storm swirling, intensifying, and making landfall just like the real thing.

The Speed and Efficiency

Perhaps the most impressive part is the speed.

Old Way: Traditional supercomputers take hours to run these complex physics simulations.
Other AI: Some AI models are fast but need massive amounts of memory (like a supercomputer in your pocket).
MOSAIC: It can generate a 10-day forecast with 24 different scenarios in under 12 seconds on a single, standard graphics card (like the ones used for gaming). It's like going from a slow, blurry photocopy to a high-definition, instant snapshot.

The Bottom Line

MOSAIC is like upgrading from a fuzzy, averaged-out weather map to a high-definition, 3D simulation that runs in your pocket.

It stops the AI from "smoothing out" the world. Instead, it learns to see the jagged edges, the swirling storms, and the chaotic details that actually save lives during extreme weather events. By looking at the data in groups (blocks) and generating multiple realistic possibilities, it brings the "spectral fidelity" (the sharpness of the signal) back to machine learning weather forecasting.

1. Problem Statement

Machine Learning Weather Prediction (MLWP) models have achieved significant speedups over traditional Numerical Weather Prediction (NWP) systems but suffer from spectral degradation. They systematically suppress power at fine spatial scales, failing to resolve sharp phenomena like frontal zones and tropical cyclones. The authors identify two primary causes for this degradation:

Deterministic Training: Most models are trained to predict the conditional mean (ensemble mean), which is inherently smoother than individual realizations, leading to a loss of high-frequency details.
Compressive Encoding: Many architectures compress high-resolution data into coarser meshes (e.g., via patchification) before processing. This creates an information bottleneck where fine-scale spatial variations become entangled in channel dimensions, causing irreversible information loss.

2. Methodology: MOSAIC

The authors introduce MOSAIC (MOSAIC), a probabilistic weather forecasting model designed to address both sources of spectral degradation through two core innovations:

A. Block-Sparse Attention (BSA)

To eliminate the information bottleneck caused by compression, MOSAIC operates directly on native-resolution grids without compressive encoding. To handle the computational cost of global attention on high-resolution grids (which scales quadratically), the authors propose Block-Sparse Attention:

Mechanism: Instead of selecting key-value pairs for every individual token, tokens are grouped into spatially adjacent blocks. These blocks jointly select which other blocks to attend to.
Three-Branch Structure:
1. Compression: Computes coarse-grained attention between block representations to capture global context.
2. Selection: Dynamically selects the top- $n$ key blocks for each query block and computes fine-grained attention at full resolution within those selected blocks.
3. Local: Computes standard attention within local blocks to capture fine-grained spatial details.
Efficiency: This approach reduces complexity to linear cost ( $O(N)$ ) while maintaining the expressivity of softmax attention. It is hardware-aligned, utilizing contiguous memory access patterns on GPUs.
Grid Representation: The model operates on the HEALPix mesh (Hierarchical Equal Area isoLatitude Pixelation). This ensures spatially close points occupy contiguous memory, facilitating efficient block-based computation and avoiding the pole-convergence issues of latitude-longitude grids.

B. Probabilistic Forecasting via Functional Perturbations

To address the spectral smoothing caused by deterministic training, MOSAIC generates ensemble members using learned functional perturbations:

Mechanism: Instead of adding noise to inputs or outputs, the model injects noise into the SwiGLU (Gated Linear Unit) layers as a bias term inside the gate activation.
Effect: This creates a single weight perturbation that affects the entire forecast in a globally consistent way. The noise is conditioned on the input, allowing the model to learn an adaptive covariance structure where uncertainty varies based on the weather state (e.g., higher uncertainty in storm regions).
Training Objective: The model is trained using the Continuous Ranked Probability Score (CRPS), a proper scoring rule that encourages both accuracy and well-calibrated ensemble spread.

3. Key Contributions

Block-Sparse Attention: A novel attention mechanism that captures long-range dependencies at linear cost, enabling weather models to operate on native-resolution grids without the information loss associated with compressive encoding.
MOSAIC Architecture: A probabilistic U-Net model combining BSA and learned functional perturbations to simultaneously solve architectural and statistical sources of spectral degradation.
Performance Efficiency: The model achieves state-of-the-art results among 1.5° resolution models while being an order of magnitude faster than comparable probabilistic models.

4. Experimental Results

The model was evaluated on two benchmarks: a 1.5° resolution comparison against other 1.5° models, and a cross-resolution comparison against 0.25° state-of-the-art (SOTA) models.

Forecast Skill (RMSE/CRPS):
- At 1.5° resolution, MOSAIC matches or outperforms models trained on 6× finer data (0.25°) on key upper-air variables (e.g., Geopotential at 500 hPa, Wind at 850 hPa).
- It achieves competitive CRPS scores against probabilistic SOTA models like GenCast and FGN.
Spectral Fidelity:
- Spectral analysis shows that individual ensemble members of MOSAIC exhibit near-perfect spectral alignment with ground truth across all resolved frequencies.
- In contrast, deterministic models and ensemble means show significant spectral suppression at fine scales.
Extreme Event Forecasting (Hurricane Ian):
- A case study on Hurricane Ian (Category 4) demonstrated that MOSAIC, despite operating at 1.5° resolution (which cannot resolve the inner core), successfully captured the hurricane's track, recurvature, and landfall.
- The ensemble spread narrowed physically as the lead time decreased, indicating well-calibrated uncertainty quantification.
Efficiency:
- Speed: A 24-member, 10-day forecast is generated in under 12 seconds on a single NVIDIA H100 GPU.
- Memory: Peak GPU memory usage is approximately 3.26 GB, significantly lower than many SOTA models (e.g., GraphCast uses ~36 GB).
- Training: Trained in 2 days on 8 H100 GPUs (16 GPU-days total), compared to weeks of training for other models.

5. Significance

Breaking the Resolution Barrier: MOSAIC demonstrates that high-resolution native processing is feasible without compression, provided efficient attention mechanisms (BSA) are used. This challenges the prevailing trend of compressing data to reduce compute costs.
Probabilistic Efficiency: It proves that probabilistic forecasting does not need to be computationally prohibitive. MOSAIC generates high-quality ensembles at speeds comparable to deterministic models.
Operational Viability: The combination of high spectral fidelity, fast inference, and low memory footprint makes MOSAIC a strong candidate for real-time ensemble nowcasting, rapid data assimilation cycles, and operational decision support for extreme weather events.
Future Scaling: The linear scaling of BSA suggests that MOSAIC can be easily scaled to even finer resolutions (e.g., 0.25°) where its advantages over compression-based models would be even more pronounced.