Towards mechanistic understanding in a data-driven weather model: internal activations reveal interpretable physical features

TL;DR

This paper uses interpretability tools to analyze internal features of a data-driven weather model, revealing physically meaningful representations of atmospheric phenomena and demonstrating how interventions can modify predictions.

Contribution

It introduces a method to interpret internal features of a data-driven weather model, linking neuron activations to physical weather phenomena and enabling targeted interventions.

Findings

Identified interpretable features corresponding to weather phenomena like cyclones and atmospheric rivers.

Demonstrated interventions on internal features can produce physically consistent changes in model outputs.

Provided a framework for understanding and trusting data-driven weather models.

Abstract

Large data-driven physics models like DeepMind's weather model GraphCast have empirically succeeded in parameterizing time operators for complex dynamical systems with an accuracy reaching or in some cases exceeding that of traditional physics-based solvers. Unfortunately, how these data-driven models perform computations is largely unknown and whether their internal representations are interpretable or physically consistent is an open question. Here, we adapt tools from interpretability research in Large Language Models to analyze intermediate computational layers in GraphCast, leveraging sparse autoencoders to discover interpretable features in the neuron space of the model. We uncover distinct features on a wide range of length and time scales that correspond to tropical cyclones, atmospheric rivers, diurnal and seasonal behavior, large-scale precipitation patterns, specific geographical coding, and sea-ice extent, among others. We further demonstrate how the precise abstraction of these features can be probed via interventions on the prediction steps of the model. As a case study, we sparsely modify a feature corresponding to tropical cyclones in GraphCast and observe interpretable and physically consistent modifications to evolving hurricanes. Such methods offer a window into the black-box behavior of data-driven physics models and are a step towards realizing their potential as trustworthy predictors and scientifically valuable tools for discovery.