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

Titel: (Sparse) Attention to the Details: Behoud van Spectrale Fideliteit in ML-gebaseerde Weermodellen

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

---

1. Het Probleem

Machine Learning Weather Prediction (MLWP) modellen hebben de afgelopen jaren enorme stappen gemaakt in snelheid en efficiëntie ten opzichte van traditionele Numerieke Weermodellen (NWP). Ze kunnen echter twee fundamentele beperkingen niet overwinnen die leiden tot spectrale degradatie: het systematisch onderdrukken van energie op fijne ruimtelijke schalen (kleine golflengten). Dit resulteert in voorspellingen die te "glad" zijn en geen realistische structuren (zoals frontale zones of tropische cyclonen) kunnen weergeven.

De twee hoofdoorzaken van dit probleem zijn:

1. Deterministisch trainen tegen ensemble-middelen: De meeste modellen worden getraind om de conditionele gemiddelde van toekomstige toestanden te voorspellen. Het gemiddelde van vele mogelijke uitkomsten is inherent gladder dan een individuele realisatie, waardoor fijne details verloren gaan.
2. Compressieve encoding (Informatie-flesnek): Veel modellen comprimeren hoge-resolutie weerdata naar een grover rooster (mesh) voordat ze verwerkingen uitvoeren. Deze compressie creëert een informatie-flesnek waarbij onomkeerbare informatieverlies optreedt, vooral op fijne schalen.

2. Methodologie: MOSAIC

De auteurs introduceren MOSAIC (MOSAIC), een probabilistisch weermodel dat specifiek is ontworpen om beide bronnen van spectrale degradatie aan te pakken.

A. Architectuur: Block-Sparse Attention op Native Resolutie

In plaats van data te comprimeren, werkt MOSAIC direct op het native-resolutie rooster (1.5° in dit geval) zonder compressie. Om de kwadratische complexiteit van standaard attention-mechanismen te omzeilen bij zulke grote datasets, gebruiken ze een nieuwe variant genaamd Block-Sparse Attention (BSA):

• HEALPix Mesh: De data wordt gemapt naar een HEALPix-rooster (een hiërarchische, gelijke-oppervlakte tesselatie van de sfeer). Dit zorgt ervoor dat ruimtelijk naburige punten ook contigu in het geheugen liggen, wat efficiënte GPU-toegang mogelijk maakt.
• Block-structuur: In plaats van dat elke query-token onafhankelijk kiest waar hij naar kijkt, worden tokens gegroepeerd in blokken. Deze blokken selecteren gezamenlijk welke andere blokken (keys/values) ze moeten "attenderen".
• Drie takken: BSA behoudt de structuur van Native Sparse Attention (NSA) met drie takken:
  1. Compressie: Berekent attention op een grof niveau tussen blokken om globale context te vangen.
  2. Selectie: Selecteert de top-$n$ blokken met de hoogste scores en berekent daar fijne attention op (voor lange-afstandsafhankelijkheden).
  3. Lokaal: Berekent standaard attention binnen lokale blokken voor fijne details.
• Efficiëntie: Deze aanpak heeft een lineaire complexiteit ($O(N)$) en maakt het mogelijk om lange-afstandsafhankelijkheden te vangen zonder de geheugenkosten van volledige attention.

B. Probabilistische Voorspelling: Learned Functional Perturbations

Om het deterministische probleem aan te pakken, genereert MOSAIC ensembleleden via geleerde functionele perturbaties:

• In plaats van alleen een gemiddelde te voorspellen, injecteert het model ruis in de bias van de SwiGLU-lagen (de feed-forward netwerken).
• Deze ruis wordt gestuurd door een klein aantal latente variabelen, waardoor de onzekerheid gestructureerd en input-afhankelijk is.
• Dit resulteert in een ensemble van 24 leden die elk een realistische, scherpere realisatie van het weer vertegenwoordigen, met spectrale eigenschappen die dicht bij de waarheid liggen.

C. Trainingsstrategie

• Loss Functie: Het model wordt getraind met de CRPS (Continuous Ranked Probability Score), een probabilistische loss die zowel nauwkeurigheid als goed gekalibreerde onzekerheid belooft.
• Data: Getraind op ERA5 data (1.5° resolutie), maar presteert op een niveau dat vergelijkbaar is met modellen die op 6x fijnere data (0.25°) zijn getraind.

3. Belangrijkste Bijdragen

1. Block-Sparse Attention: Een hardware-geoptimaliseerd mechanisme dat lange-afstandsafhankelijkheden vangt tegen lineaire kosten, waardoor weermodellen direct op native-resolutie roosters kunnen werken zonder compressie.
2. MOSAIC Model: Een probabilistisch model dat zowel de architecturale (compressie) als statistische (deterministisch trainen) oorzaken van spectrale degradatie adresseert.
3. Spectrale Fideliteit: Het bewijs dat individuele ensembleleden van MOSAIC een "bijna perfecte" spectrale uitlijning hebben met de grondwaarheid over alle opgeloste frequenties, iets wat deterministische modellen en ensemble-middelen niet halen.

4. Resultaten

De resultaten tonen aan dat MOSAIC concurrerend is met de state-of-the-art (SOTA) modellen, ondanks dat het op een grover rooster (1.5°) werkt dan veel concurrenten (0.25°).

• Voorspellende Vaardigheid (Skill):
  • Op de 1.5° benchmark presteert MOSAIC gelijk aan of beter dan modellen zoals NeuralGCM en ArchesWeatherGen op sleutelvariabelen (zoals geopotentiële hoogte op 500 hPa en windsnelheid).
  • Op de 0.25° benchmark (waar het model getest wordt tegen 0.25° SOTA-modellen) behaalt MOSAIC vergelijkbare RMSE en CRPS scores als GenCast, GraphCast en Aurora, ondanks dat het op 1.5° data is getraind.
• Spectrale Analyse:
  • Figuur 2 toont dat deterministische modellen en ensemble-middelen de spectrale macht op fijne schalen sterk onderdrukken (ratio < 0.5).
  • Individuele leden van MOSAIC behouden echter de spectrale macht (ratio ~1.0), wat betekent dat ze de fijne structuren van het weer correct weergeven.
• Efficiëntie:
  • Een 24-leden ensemble van 10 dagen wordt gegenereerd in minder dan 12 seconden op één NVIDIA H100 GPU.
  • Dit is ongeveer 20x sneller dan ArchesWeatherGen en meer dan 150x sneller dan NeuralGCM.
• Extreme Gebeurtenissen (Hurricane Ian):
  • Een case study toont aan dat MOSAIC de track van Hurricane Ian correct voorspelt, inclusief de kromming en landingsplaats, ondanks de lagere resolutie. Het ensemble toont een fysiek redelijke onzekerheidskwalificatie (de spreiding neemt af naarmate de voorspelling dichter bij de gebeurtenis komt).

5. Betekenis en Toekomst

Dit paper is significant omdat het aantoont dat compressie niet nodig is voor hoogwaardige ML-weermodellen. Door slimme attention-mechanismen (Block-Sparse) en probabilistische technieken te combineren, kan men modellen bouwen die:

1. Sneller zijn dan huidige SOTA-modellen.
2. Realistischer zijn in het weergeven van fijne schalen en extreme gebeurtenissen.
3. Betrouwbaarder zijn in hun onzekerheidsindicaties.

De auteurs suggereren dat dit pad leidt naar toekomstige modellen op nog fijnere resoluties (bijv. 0.25°) die binnen real-time limieten kunnen opereren, wat cruciaal is voor snelle ensemble-data-assimilatie en operationele besluitvorming bij extreme weersomstandigheden.