ClimaX: A Foundation Model for Weather and Climate

Aditya Grover, UCLA
LANS Seminar Graphic

Most state-of-the-art approaches for weather and climate modeling are based on physics-informed numerical models of the atmosphere. These approaches aim to model the non-linear dynamics and complex interactions between multiple variables, which are challenging to approximate. Additionally, many such numerical models are computationally intensive, especially when modeling the atmospheric phenomenon at a fine-grained spatial and temporal resolution. Recent data-driven approaches based on machine learning instead aim to directly solve a downstream forecasting or projection task by learning a data-driven functional mapping using deep neural networks. However, these networks are trained using curated and homogeneous climate datasets for specific spatiotemporal tasks, and thus lack the generality of numerical models. In this talk, I will present ClimaX, a flexible and generalizable deep learning model for weather and climate science that can be trained using heterogeneous datasets spanning different variables, spatio-temporal coverage, and physical groundings. ClimaX extends the Transformer architecture with novel encoding and aggregation blocks that allow effective use of available compute while maintaining general utility. The pre-trained ClimaX can then be fine-tuned to address a breadth of climate and weather tasks, including those that involve atmospheric variables and spatio-temporal scales unseen during pretraining. Compared to existing data-driven baselines, we show that this generality in ClimaX results in superior performance on benchmarks for weather forecasting and climate projections, even when pretrained at lower resolutions and compute budgets. Towards the end of the talk, I will present ClimateLearn, our open-sourced library to standardize machine learning for climate science.

Speaker Bio: Aditya Grover is an assistant professor of computer science at UCLA. His goal is to develop efficient machine learning approaches that can interact and reason with limited supervision with a focus on deep generative models and their intersection with sequential decision making and causal inference. He is also an affiliate faculty at the UCLA Institute of the Environment and Sustainability, where he grounds his research in real-world applications in climate science. Aditya's 40+ research works have been published at top venues including Nature, deployed in production at major technology companies, and covered in popular press venues. His research has been recognized with two best paper awards, four research fellowships, four faculty awards, the ACM SIGKDD doctoral dissertation award, and the AI Researcher of the Year Award by Samsung. Aditya received his postdoctoral training at UC Berkeley, PhD from Stanford, and bachelors from IIT Delhi, all in computer science.