Building an AI Foundation Model for Weather and Climate

A recent workshop brought together experts from NASA, Oak Ridge National Laboratory (ORNL), IBM Research, NVIDIA, and several universities to develop a plan to create an artificial intelligence foundation model (AI FM) for weather and climate.

The saying, "If you want to go fast, go alone; if you want to go far, go together," perfectly encapsulates the essence of open science and the necessity of collaboration in addressing complex issues. In this spirit of open science, a workshop was convened at NASA’s Marshall Space Flight Center in Huntsville, Alabama, on September 20 and 21, 2023.

The workshop brought together members of NASA's Interagency Implementation and Advanced Concepts Team (IMPACT), representatives from NASA Headquarters, scientists and engineers specializing in artificial intelligence (AI) model development, high-performance computing experts, machine learning and data engineers, and weather and climate subject matter experts. Participants were from NASA, Oak Ridge National Laboratory (ORNL), IBM Research, NVIDIA, and several universities, all of whom were already engaged in various aspects of AI or large scale computing and weather/climate science.

The primary objective of the workshop was to align this diverse array of expertise in collectively creating an AI foundation model (AI FM) for weather and climate in the next six to eight months (the initial concept is outlined in the paper linked at the bottom of this announcement). This weather FM will contain parameters such as wind speed and direction, air temperature, specific humidity, cloud mass variables, and longwave and shortwave radiation variables. To be valuable to the broader science community, an AI FM for weather and climate should enable many different types of downstream applications in addition to forecasting. The workshop served as a kick-off meeting to establish a tentative project plan and provided the opportunity for the participants to volunteer for tasks and assume responsibilities.

Workshop Work

The first morning of the workshop was dedicated to providing a common foundational grounding for all participants, who brought varying expertise and perspectives. The overview sessions covered the concept of AI foundation models and their potential value in scientific endeavors, the available infrastructure resources for the team, and an initial list of science use cases to validate the AI FM. The sessions explored potential AI model architectures and the process for designing, building, testing, and deploying these models.

The afternoon featured three breakout sessions for in-depth exploration of specific areas. The first breakout concentrated on refining the model architecture based on requirements derived from the science use cases. The second breakout focused on consolidating existing and new science use cases to evaluate and demonstrate the value of the pretrained AI foundation model. Participants ranked these use cases according to their impact and the feasibility of creating labeled datasets within a six-month timeframe. Participants in the third breakout session identified and evaluated potential computational infrastructure to support this effort, determined the necessary software stacks, and selected the stacks suitable for pretraining, fine-tuning, and inference.

The following morning, each breakout team reported their discussions and presented a tentative plan. This plenary session allowed all participants to collectively combine and refine each breakout group's approach into a final plan. Ultimately, participants reached a consensus on the overall approach, a tentative timeline, and their respective roles and responsibilities.

Workshop Results and Next Steps

The objective of this work is to pretrain an autoregressive AI foundation model capable of handling multi-resolution and multi-temporal data. The team will also explore different architectures to achieve these goals. This model will be tested against seven selected science use cases, carefully chosen from a list of 40. These use cases cover a broad range of downstream applications ranging from regression and forecasting to classification and superresolution. The model, the source code for pretraining and fine tuning, the use case datasets, and documentation will be made openly available.

This workshop serves as a model example of setting the participation ground rules based on open science principles at the outset. The approach facilitated active participation by all attendees, enabling them to share new ideas and expertise during discussions. Furthermore, the workshop underscored the importance of active partnerships in addressing complex challenges such as developing AI models for scientific purposes, where expertise and resources may not be concentrated within a single organization.

Learn More

Mukkavilli, S.K., et al. (2023). AI Foundation Models for Weather and Climate: Applications, Design, and Implementation. Cornell University arXiv. doi:10.48550/arXiv.2309.10808

Workshop Organizers
  • Rahul Ramachandran, NASA IMPACT
  • Tsengdar Lee, NASA Headquarters
  • Valentine Anantharaj, ORNL
  • Raghu Ganti, IBM Research
  • Kommy Weldemariam, IBM Research
  • Elizabeth Fancher Barrios, NASA IMPACT (logistics and planning) 
Workshop Participants

Model Development 

  • Aristeidis Tsaris, ORNL
  • Sujit Roy, The University of Alabama in Huntsville (UAH)/NASA IMPACT
  • Karthik Mukkavilli, IBM Research
  • Daniel Salles Civitarese, IBM Research
  • David Hall, NVIDIA 
  • Johannes Schmude, IBM Research

Resource Management and Support

  • Bill Thigpen, NASA's Ames Research Center, NASA Advanced Supercomputing Facility (NAS)
  • Mike Little, NASA's Goddard Space Flight Center, NASA Center for Climate Simulation (NCCS)

Science Use Cases/Fine Tuning

  • Simon Pfreundschuh, Colorado State University
  • U.S. Nair, UAH
  • Andrew Molthan, NASA Marshall
  • Alex Melancon, UAH 
  • Chris Phillips, UAH/NASA IMPACT
  • Rajat Shinde, UAH/NASA IMPACT
  • Arlindo Da Silva, NASA Goddard, Global Modeling and Assimilation Office (GMAO) 
  • Emily Berndt, NASA Marshall/Short-term Prediction Research and Transition (SPoRT) center
  • Mark Carroll, NASA Goddard
  • Kwo-sen Kuo, NASA Goddard
  • Hendrik Hamann, IBM Research
  • Steven Wysmuller, IBM Research

Model Interpretability

  • Haonan Chen, Colorado State University 

Programmatic Alignment

  • Tsengdar Lee, NASA Headquarters
  • Manil Maskey, NASA Headquarters

Students and Postdocs

  • Ankur Kumar, UAH
  • Aman Gupta, Postdoctoral Scholar, Stanford University 
  • Takuya Kurihana, University of Chicago/IBM Intern 

ML Engineering/Workflow Support

  • Muthukumaran Ramasubramanian, UAH/NASA IMPACT
  • Iksha Gurung, UAH/NASA IMPACT 
  • Vishal Gaur, UAH/NASA IMPACT
  • Tyler Skluzacek, ORNL