Verbesserte Interpretierbarkeit von KI-Wettermodellen durch intermediate decoding

Trotz der schnellen Fortschritte und der starken Performance von KI-Wettermodellen wissen wir noch wenig darüber, wie sie ihre Vorhersagen intern aufbauen, weshalb sie oft als ”Black Boxes“ gelten. Das liegt daran, dass sie in einem hochdimensionalen latenten Zwischen-Raum arbeiten, der sich nur schwer direkt interpretieren lässt. Während es viele Studien zur Interpretierbarkeit von anderen KI-Modellen gibt, vor allem für Large Language Models und Vision Models, ist vergleichsweise wenig zur Interpretierbarkeit von KI-Wettermodellen bekannt.

Für Sprachmodelle ist ein gängiger Ansatz zur Modell-Interpretation die "logit lens". Dabei werden Zwischenzustände des latenten Raums direkt in den Decoder des Modells eingespeist, der die sogenannte Logit-Verteilung erzeugt, aus der der wahrscheinlichste nächste Token, also ein Wort oder ein Wortteil, abgelesen werden kann. Beobachtet man, wie sich dieser vorhergesagte Token verändert, wenn vor der Dekodierung weitere Prozessorblöcke einbezogen werden, bekommt man einen Einblick, wie das Modell seine Vorhersage Schritt für Schritt aufbaut. Anders gesagt: Man sieht, was das Modell in verschiedenen Tiefen im Netzwerk als nächstes Wort ”erwartet“.

Um zu verstehen, wie KI-Wettermodelle Vorhersagen erstellen, haben wir die logit lens-Methode für das GraphCast KI-Wettermodell von Google angepasst und leiten Zwischenzustände des Modells direkt an den Decoder weiter. Der Decoder ist allerdings darauf trainiert, Vorhersagen erst am Ende der Modellpipeline zu erzeugen, und die Zwischenzustände verändern ihre Form mit zunehmender Tiefe des Modells. Um dies zu lösen, trainieren wir Translatoren, einfache lineare Transformationen, die jeden Zwischenzustand in die vom Decoder erwartete Form überführen. Eine ähnliche Erweiterung der logit lens-Methode unter Einsatz von Translatoren wurde auch bei Sprachmodellen angewendet, wo ebenfalls beobachtet wurde, dass sich die latente Basis mit zunehmender Modelltiefe verändert.

Die Abbildung zeigt die Ergebnisse des intermediate decodings unter Einbeziehen unterschiedlich vieler Prozessorblöcke in einer 6-stündigen GraphCast-Vorhersage. Anstatt alle 16 GraphCast-Prozessorblöcke zu durchlaufen, können auch Zwischenzustände des Modells direkt dekodiert und dargestellt werden: Gezeigt sind schrittweise die Entstehung der  Vorhersage des 500 hPa-Geopotentialfeld in der Mitte der Atmosphäre für den 1. April 2020 um 18:00 UTC, ohne Prozessorblöcke oben (nur Encoding und Decoding), in der Mitte mit der Hälfte der Prozessorblöcke und schließlich für die vollständige Vorhersage.

Der Startzustand der Vorhersage wurde subtrahiert, sodass nur die vorhergesagte 6-Stunden-Änderung sichtbar ist. Rote Bereiche zeigen 6-Stunden-Zunahme des Geopotentials in m² s^-2 (und damit des Drucks), blaue Abnahmen.  Bemerkenswerterweise erzeugt das Modell selbst bei der Umgehung sämtlicher Prozessorblöcke und der Nutzung ausschließlich von Encoder und Decoder bereits ein physikalisch plausibles Muster, einschließlich einer sinnvollen Veränderung an einer klar ausgeprägten Rossby-Wellen-Struktur auf der Südhalbkugel. Mit zunehmender Anzahl an Prozessorböcken werden die Vorhersagen schrittweise präziser, und die Vorhersage-Qualität wird besser, was zeigt, dass jeder Block inkrementell zur Konstruktion der endgültigen Vorhersage beiträgt. Mit zunehmender Anzahl an Prozessorblöcken werden die Muster sukzessive detaillierter und die Vorhersage wird genauer.