Multi-Modal Foundation · AMLab-UvA
Mosaic
A 214M-parameter transformer for fast, probabilistic global weather forecasting with preserved spectral fidelity.
What Mosaic is
Mosaic is a probabilistic weather forecasting model developed by researchers at AMLab-UvA. It addresses spectral degradation common in machine learning weather prediction models, specifically the spectral damping that arises from deterministic training against ensemble means and the high-frequency aliasing caused by compressive coarse-grid bottlenecks. Mosaic bypasses compressive latent encoding by operating directly on native-resolution representations via mesh-aligned block-sparse attention, alongside learned functional perturbations that maintain realistic spectral variability across ensemble members.
The model operates on a 1.5° equiangular grid across 82 atmospheric channels. Despite being trained at 1.5° resolution, Mosaic matches or exceeds the performance of models trained on 0.25° data on key variables, while generating ensemble members whose kinetic-energy spectra align with ERA5 reanalysis across resolved frequencies. Highly optimised for inference, it can generate a 24-member, 10-day probabilistic forecast in under 12 seconds using a single H100 GPU.
How it works
Mosaic is structured as a U-Net of transformer blocks mapped onto a HEALPix mesh. The mesh organisation places spatially neighbouring points into contiguous memory, enabling coalesced GPU reads and block operations. Its core mechanism is a tri-branch block-sparse attention kernel combining three components via learned gates: coarse block-to-block compression to capture broad synoptic dynamics, top-k selection for fine-scale key blocks, and full local attention within blocks. The model incorporates Grouped-Query Attention with a ratio of 4, two-dimensional Rotary Position Embeddings (RoPE) over longitude and latitude, and additive noise injection in the SwiGLU gates to produce stochastic ensemble members.
The model takes 82 atmospheric variables (4 surface and 78 pressure-level variables across 13 pressure levels) and 3 conditioning static channels alongside sinusoidal day and year encodings. Two model variants exist: an 'era5' variant trained purely on ERA5 reanalysis taking two 24-hour historical states to predict at a 24-hour step, and an 'hres' variant pretrained on ERA5 and fine-tuned on ECMWF HRES analysis taking four 6-hour historical states to predict at a 6-hour step.
Key facts
- Model architecture
- Transformer U-Net with block-sparse attention on a HEALPix mesh
- Parameter count
- 214M
- Spatial resolution
- 1.5° (240 longitude × 121 latitude)
- Atmospheric channels
- 82 predicted channels, 3 static channels
- Input context length
- 2 states / 48 h (era5 variant); 4 states / 24 h (hres variant)
- Forecast step cadence
- 24 h (era5 variant); 6 h (hres variant)
- Training data
- ERA5 reanalysis; ECMWF HRES analysis
What you can use it for
- Medium-range global weather forecasting
- Roll out 10-day forecasts across 82 surface and upper-air atmospheric variables from standard 1.5° reanalysis or operational initial conditions.
- Probabilistic ensemble generation
- Produce multi-member ensemble forecasts using learned SwiGLU noise perturbations to quantify forecast uncertainty and compute spread-skill metrics.
- Extreme weather tracking
- Track synoptic-scale events such as tropical cyclones over multi-day horizons, resolving ensemble track spreads as lead time progresses.
- Atmospheric kinetic-energy spectral analysis
- Examine spectral dynamics without artificial smoothing, ensuring simulated fields match physical reanalysis spectra across all resolved spatial scales.
Run it here
We list this model but cannot run it yet: no pipeline_tag — needs a human.
Everything on this page still applies. Running it needs a loader for this repository, not a bigger machine.
Strengths and limits
Good at
- Retains realistic atmospheric kinetic-energy spectra in individual ensemble members, eliminating spectral damping and aliasing without blurring fine-scale features.
- Generates a 24-member, 10-day forecast in under 12 seconds on a single H100 GPU.
- Block-sparse attention implementation scales linearly with sequence length and runs up to 61.8 times faster than dense attention.
- Matches or outperforms 0.25° baseline models on key variables despite operating at a coarser 1.5° resolution.
Watch out for
- Operates at 1.5° resolution (~166 km), which cannot resolve small-scale mesoscale phenomena such as convective thunderstorms or the inner cores of tropical cyclones.
- Requires high-end CUDA hardware with compatible Triton and flash-attn installations, needing approximately 9 GB of GPU RAM for a single-member 10-day rollout.
- Input datasets must provide exact 82-variable configurations and 3 static fields formatted to a 1.5° equiangular grid with poles.
Inputs and outputs
You provide
Sequence or text
A protein or nucleotide sequence, or free text — whichever this model was trained on.
You get back
Mean-pooled representation, one vector per input.
Licence
Check the licence before commercial use · apache-2.0
The apache-2.0 licence carries terms worth reading before you depend on it.
Questions
Can I use Mosaic commercially?
Check the licence before commercial use. The apache-2.0 licence carries terms worth reading before you depend on it.
What is the difference between the era5 and hres variants?
The era5 variant is trained solely on ERA5 reanalysis, runs on a 24-hour step, and requires 48 hours of past context. The hres variant is pretrained on ERA5 and fine-tuned on HRES analysis, running on a 6-hour step with 24 hours of input context.
Which atmospheric variables does Mosaic predict?
Mosaic outputs 82 channels: 4 surface variables (2m temperature, 10m u/v wind components, mean sea level pressure) and 6 upper-air variables (geopotential, specific humidity, temperature, u/v wind components, and vertical velocity) evaluated across 13 pressure levels ranging from 1000 hPa to 50 hPa.
How are probabilistic ensembles generated?
Ensemble members are produced by injecting additive noise into the SwiGLU activation gates of the transformer, which acts as a learned functional perturbation preserving spectral characteristics.
What input format does Mosaic require?
It accepts 1.5° equiangular data formatted as zarr stores matching WeatherBench2 conventions, with dimensions of 240 longitude by 121 latitude and containing all 10 base atmospheric variables.
How to cite it
Maksim Zhdanov, Ana Lucic, Max Welling, Jan-Willem van de Meent. (Sparse) Attention to the Details: Preserving Spectral Fidelity in ML-based Weather Forecasting Models. Proceedings of the 43rd International Conference on Machine Learning (ICML), 2026. https://arxiv.org/abs/2604.16429
Written from Mosaic's own documentation (model card). Statements citing a figure that does not appear there are removed before publication. Licence and running details come from our registry.