Finetuning Weather Foundation Models to Develop Climate Model Parameterizations

Authors: Aman Gupta, Sujit Roy, Johannes Schmude, Vishal Gaur, Wei Ji Leong, Manil Maskey, Rahul Ramachandran, Aditi Sheshadri
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: Atmospheric Dynamics, Parameterizations, Climate Modelling, Foundation Model, ERA5, Finetuning, Machine Learning

Abstract:
Climate prediction models parameterize a range of atmospheric-oceanic processes like clouds, turbulence, and gravity waves. These physical parameterizations are a leading source of uncertainty and strongly influence future projections of global temperature rise. We present a fresh approach to developing parameterizations for coarse-climate models by leveraging pre-trained AI foundation models (FMs) for weather and climate. A pre-trained encoder and decoder from a 2.3 billion parameter FM (NASA and IBM's Prithvi WxC) --- which contains a latent probabilistic representation of atmospheric evolution --- is fine-tuned to create a data-driven predictor of atmospheric gravity waves (GWs). Current climate models are not fine enough to resolve GWs. We create an ML-based parameterization that learns GW fluxes from high-resolution ``GW resolving" climate models to represent them in "GW missing" coarse-climate models. The fluxes predicted by our fine-tuned model are comprehensively evaluated using a set of three tests. Comparison with a baseline (Attention U-Net) reveals the superior predictive performance of the fine-tuned model throughout the atmosphere. The model outperforms the baseline even in regions excluded from the FM pre-training. This is quantified using the Hellinger distance which is 0.11 for the baseline and 0.06, i.e., roughly half, for the fine-tuned model. FMs are largely unexplored in climate science. Our findings emphasize their versatility and reusability to accomplish a range of weather- and climate-related downstream applications, especially in a low-data regime. These FMs can be further leveraged to create new parameterizations for other earth-system processes.

Decision: reject
Reviewer scores: 3, 3, 5, 3

Reviews and discussion:
<thread>
creation date:
1730373497283

thread_type:
official_review

<thread_object>
creator: reviewer_dybv

<content>
summary:
A recent foundation model for weather and climate (Prithvi WxC) is fine tuned to predict gravity waves. Those predictions are supposed to parameterize subscale processes of coarse climate models and thus improve climate projections. The fine tuned model outperforms a single-task U-Net and generalizes well in regions that are outside Prithvi WxC's training range.

strengths:
_Originality:_ The proposed manuscript touches on relevant topics of the society and applies state-of-the-art foundation models. It is great to see that the fine tuned foundation model generates accurate predictions even outside the regime of the foundation model's training data.

_Clarity:_ The manuscript is mostly clear and well organized. Supplementary material and code for replicating the results are not provided, though.

weaknesses:
_Quality_
The quality of the manuscript lacks in several aspects. Mostly, the comparison of the foundation model vs the baseline seems unfair.
1. Unclear for what reason the authors choose Prithvi WxC as foundation model and not ClimaX or AtmoRep. Comparing these three foundation models can be considered more fair than comparing the fine tuned FM against conventional one-task approaches, like Attention U-Net etc.
2. In the same vein, given that the baseline and fine tuning models are trained on different sets of input variables (lines 203 through 210), it is unclear how to differentiate between the model quality and the data selection.
3. The baseline is a convolutional architecture, whereas the fine tuned model is a transformer, which questions the role of the pretraining vs. the model architecture. It would be great to see how a transformer would compare as baseline model. Similarly, the number of parameters of the U-Net (35M) is not comparable to that of the fine tuned model (2.3B).
4. A plot showing the convergence of the baseline vs the fine tuning model would be informative to better capture their behavior (line 271).
5. In Figure 5, the dotted lines are described to indicate the 2.5th and 97.5th percentiles. The data distributions, however, hardly confirm this. There appears to be substantially more than 2.5 percent of the data left and right of the blue dashed vertical lines.
6. How is the argument in line 325 substantiated, that a Hellinger distance of 0.05 or less is considered pretty good? Is there some literature or data that suggests this decision?

_Significance_
It is hard to assess whether the proposed method is relevant for climate forecasting, mostly as the parametrizations are not tested in numerical climate models.
1. Output fluxes are on fairly coarse resolution and I'm concerned that the coars gravity wave precitions are of limited value for a numerical climate model?
2. The model is introduced to provide parametrizations for climate models, however the study does not test those parametrizations in climate models. As detailed in the limitations section, the actual verification of the parametrization is a key contribution that I consider substantial for assessing the relevance of the proposed method.


_Minor comments_
- lines 37-38: Add details about why SOTA future climate projections are highly uncertain
- line 289: Typo in "evolutionlution"

questions:
1. What spatial and temporal resolution is required to resolve gravity waves? (see abstract and line 60) Please add details to the manuscript.
2. In the caption of Figure 2, what does it mean that encoder and decoder blocks are frozen and used for fine tuning? This reads conflicting, since frozen weights cannot be fine tuned. Also, this figure does not seem to convey much information for the model setup at hand. I have difficulties extracting details about data or architecture. EDIT: In Section 2.4 this is outlined clearly; I suggest to remove Figure 2.
3. What do the models learn effectively? It seems like the models are trained to approximate Equations (1)-(3) and I do not understand why this is done with deep learning models instead of using these equations directly.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730619869051

thread_type:
official_review

<thread_object>
creator: reviewer_xxrg

<content>
summary:
This paper presents a data-driven parameterization scheme for gravity waves, aiming to achieve gravity wave parameterization in coarse climate models by fine-tuning weather foundation models. The work utilizes the newly proposed Prithvi WxC as a pre-trained model and leverages higher-resolution ERA5 data to achieve this goal. Their main contribution is both reducing the costs of training data-driven parameterization models from scratch and improving their generalization capabilities through fine-tuning foundational models.

strengths:
1. Significance: This work introduces a fine-tuning algorithm for weather foundation models into parameterization, achieving a more lightweight and better generalization performance data-driven parameterization scheme.    
2. The manuscript is well-written and comprehensible, adhering to ICLR formatting guidelines, with no discernible errors.    
3. The performance of the proposed algorithm is demonstrated through three different evaluation methods.

weaknesses:
1.**Innovativeness:** This work showcases commendable application performance on specific tasks; however, its contribution to the broader machine learning community may be viewed as somewhat limited. Additionally, the fine-tuning scheme is discussed in detail in Section 3.2 of the cited paper that presents the pre-trained foundational model used here, where it is described as a typical downstream task. Given the context of a high-level conference, the innovativeness of this work might be considered subtle.  

2. **Experimental Setup:** It seems that neither the methods relying on equations nor the mixed probability methods mentioned in the 'Related Work' section are included in the baselines; rather, the comparison is made with Attention U-Net. Furthermore, the input variables for Attention U-Net are one-quarter fewer than those utilized in this work, resulting in a reduction of input information. Could you clarify why it is not possible to fully input the data of [488, 64, 128]? 

3. **Experiments:** In the daily average GW momentum flux experiments (Figures 5 and 10), notable differences are observed in the distribution of small-scale features between the deep learning model and ERA5. It would be beneficial to include the MERRA-2 distribution to clarify that the data-driven parameterization scheme successfully learns GW information from the ERA5 dataset, rather than simply inheriting features from a pre-trained model.

questions:
1. In the "Contributions" section, the authors suggest that this method could potentially be extended to cloud parameterization and precipitation forecasting (line 105). However, these topics are not explored within this work. Given their distinct nature compared to gravity wave parameterization, and considering that the ERA5 dataset may lack reliable data for these variables, could the authors provide additional experiments or insights to substantiate this claim?

2. In lines 201 and 202, the authors state, "This corresponds to roughly 35k training samples, which pretty much classifies as 'data-scarce'." This definition of sample size seems to differ from what is typically observed in other areas of the machine learning community. Could the authors provide context on the typical dataset sizes used for similar climate modeling tasks, and explain why 35k samples is considered data-scarce in this domain compared to other machine learning applications?

3. Could the authors clarify how the fine-tuning and test datasets are partitioned for the experiments? Specifically, please provide the exact split ratios for training, validation, and testing, and indicate whether the data was split randomly or using another method, such as temporal splits for time series data. This would enhance readers' understanding of the experimental setup.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730650575163

thread_type:
official_review

<thread_object>
creator: reviewer_ripd

<content>
summary:
This paper introduces the fine-tuning of weather/climate foundation models to the task of atmospheric gravity waves.

strengths:
- The presentation of the results is interesting and thought-provoking. Very cool to see the difference in the modeling between baseline and fine-tuned model.

weaknesses:
- The novelty / content of the paper is somewhat limited. For example in ClimaX --  an ICML paper -- which introduced a full end-to-end pretraining and finetuning pipeline over multiple datasets and multiple weather and climate finetuning tasks. A possible extension for this paper would be to use different foundation model backbones - any of the large weather models are interesting here.
- The Prithvi WxC model is very low resolution compared to other models, shows hardly no ablations, and for none of the tasks they consider, they beat any SOTA baselines. The current SOTA in weather and climate modeling has moved way beyond the 0.625x0.5 resolution. This of course is not the fault of the authors. Yet for example, the Aurora foundation model (left out in the paper) is trained on many atmospheric datasets on much finer resolution. It would have been nice to put a comparison between these two models.
- The downstream task is super low resolution - which a priori is ok, but not rtoo impressive.
- At least one more baseline would be needed to really gauge the results.


I would advise the authors to consider user other foundation models too. And please stop refering to the Prithvi WxC model as SOTA model.

questions:
-
</content>
</thread_object>

</thread>
<thread>
creation date:
1730676803982

thread_type:
official_review

<thread_object>
creator: reviewer_5f5c

<content>
summary:
This paper explores the use of AI foundation models (FMs) to improve climate model parameterizations, specifically focusing on gravity wave (GW) effects. The authors leverage a pre-trained weather foundation model (Prithvi WxC) by fine-tuning its encoder and decoder components to predict gravity wave momentum fluxes that are typically unresolved in coarse-resolution climate models.


The work demonstrates how a foundation model pre-trained on MERRA-2 reanalysis data can be fine-tuned using ERA5 data to create parameterizations that capture gravity wave physics. The authors develop a model that predicts GW momentum fluxes given background atmospheric conditions, comparing their fine-tuned approach against a baseline Attention U-Net model trained from scratch. They evaluate the models using three tests: predicting global flux distributions, region-specific flux spectra across different atmospheric heights, and temporal evolution of fluxes at known gravity wave hotspots.


The results show that the fine-tuned foundation model approach outperforms the baseline, particularly in predicting stratospheric gravity wave behavior - even across pressure levels where the original foundation model was not pre-trained. The authors employ the tail-Hellinger distance to specifically evaluate how well the models capture extreme events in the flux distributions.


The paper positions this work as a proof-of-concept for using foundation models to develop improved climate model parameterizations more broadly, suggesting that similar approaches could be applied to other unresolved processes like clouds and precipitation. The authors argue that this approach offers advantages in terms of training efficiency, generalization capability, and physical consistency, while acknowledging current limitations and areas for future work.

strengths:
The paper puts forth a fine-tuning framework and a rigorous evaluation methodology in applying foundation models to climate science parameterizations. The work's primary strengths lie in its comprehensive experimental design and clear presentation of results. The authors provide a thorough evaluation framework through three increasingly stringent tests - from global distributions to temporal evolution at specific hotspots - which could serve as a valuable template for assessing other climate model parameterizations.


From a technical perspective, the paper employs the tail-Hellinger distance metric for evaluating extreme events in flux distributions, showing careful consideration of evaluation methodology appropriate for climate science applications. The empirical results demonstrate that their fine-tuned approach outperforms a strong baseline (Attention U-Net), particularly in an interesting case where the model generalizes well to stratospheric gravity wave behavior even in regions where the original foundation model wasn't pre-trained.


The paper is very clear and accessible through a well-structured presentation. The authors have done a good job at bridging machine learning and climate science concepts, making the work comprehensible to both communities. The figures are particularly well-designed, with Figure 1 effectively illustrating the gravity wave prediction task and Figures 5-8 systematically presenting the evaluation results across different atmospheric conditions.


While the work doesn't advance machine learning methodology significantly, it provides a well-executed case study demonstrating how foundation models can be leveraged for climate model parameterizations. The successful demonstration with limited fine-tuning data (four years of ERA5) suggests a practical pathway for developing similar parameterizations, though these contributions are more relevant to climate science than machine learning. The clear presentation and thorough validation make the work's climate science implications accessible to a broad audience, even if the core innovations lie primarily in the application domain rather than methodological advancement.


To summarize, I acknowledge the paper's strong execution and clarity while being more explicit about its primary contributions being to climate science rather than machine learning methodology.

weaknesses:
**Technical Innovation:**
The paper's primary limitation lies in its modest contribution to machine learning methodology, which is a crucial consideration for ICLR. While the work presents a compelling application of foundation models to climate science, it essentially applies established fine-tuning techniques to a new domain without introducing significant methodological innovations. The approach largely follows the pre-training/fine-tuning paradigm that has been well-documented in recent literature, including applications in weather and climate modeling [1,2].

**Connection to prior works:**
Similar applications of foundation models in Earth system science have already been demonstrated by works such as ClimaX [1] and Aurora [2], which showed that weather-trained models can be effectively fine-tuned for various downstream tasks, including data-scarce scenarios. The current paper's findings, while valuable for climate science, align with these expected outcomes and don't present surprising methodological insights for the machine learning community.

**Limited contributions to ML:**
From a technical perspective, the paper primarily describes a straightforward application of fine-tuning the Prithvi WxC model to predict gravity wave momentum fluxes. I do acknowledge the engineering effort required to carry out this study, but from a machine learning standpoint I see limited contributions. While the authors have conducted thorough experiments, these contributions are more relevant to climate science evaluation than advancing machine learning methodology. The neural network architecture modifications described in Section 2.4 are relatively standard adaptations rather than novel technical contributions.

**Venue fit:**
The paper's strengths - particularly its thorough evaluation of gravity wave predictions and implications for improving climate model parameterizations - would be better suited for climate science venues where domain experts could properly evaluate the scientific implications. Journals such as Journal of Advances in Modeling Earth Systems (JAMES) or Geophysical Research Letters (GRL) would provide a more appropriate audience and review process for assessing the work's primary contributions to climate modeling.

**Recommendation for improvement:**
While the paper effectively demonstrates the potential of machine learning in climate science, its core innovations lie in the application domain rather than in advancing machine learning methodology. The work would benefit from either substantial enhancement of its machine learning contributions for ICLR or redirection to a venue better aligned with its primary contributions to climate science.

**Impact considerations:**
This assessment isn't meant to diminish the paper's value but rather to highlight that its strengths may be better appreciated and more impactful in a different academic venue. The thorough experimental validation and careful consideration of climate science implications would likely generate more meaningful discussion and follow-up work in the climate modeling community.

[1] Nguyen, T., Brandstetter, J., Kapoor, A., Gupta, J. K., & Grover, A. (2023). ClimaX: A foundation model for weather and climate. arXiv preprint arXiv:2301.10343.

[2] Bodnar, C., Bruinsma, W. P., Lucic, A., Stanley, M., Brandstetter, J., Garvan, P., ... & Perdikaris, P. (2024). Aurora: A foundation model of the atmosphere. arXiv preprint arXiv:2405.13063.

questions:
1. Could the authors elaborate on why they chose to freeze both the encoder and decoder during fine-tuning? Would allowing some layers to be trainable, particularly in the decoder, potentially improve performance? At 2.3 billion parameters, the model size is not too prohibitive for full-parameter fine-tuning.

2. The daily flux predictions show difficulty with small flux values (Figure 5b). Could the authors discuss potential approaches to address this limitation?

3. The tail-Hellinger distance is an interesting metric - could the authors provide more intuition about how to interpret different values, particularly negative ones?

4. Have the authors considered evaluating the model's performance during extreme weather events or seasonal transitions where gravity wave behavior might be particularly challenging to predict?

5. Could the authors provide more details about the data preparation, particularly how the coarse-graining from 25km to 280km resolution was implemented?

6. What was the rationale behind choosing 4 convolutional blocks before and after the frozen encoder-decoder? How sensitive is the model performance to this architectural choice?

7. The authors mention plans to couple their scheme to a coarse-climate model. Could they elaborate on the technical challenges they anticipate in this integration? As climate models require long-time integration of the underlying PDEs do you foresee any stability issues when an ML parametrization is used?

8. The authors chose Prithvi WxC as their foundation model, but recent work has shown Aurora achieving state-of-the-art performance in weather prediction, especially at high resolutions (see [2] above). Could the authors discuss why they selected Prithvi over Aurora, and whether they expect their findings would generalize or potentially improve when using Aurora as the base model? This discussion would be particularly relevant given Aurora's demonstrated superior accuracy in weather forecasting and its potential advantages for learning atmospheric dynamics. It would also help readers understand whether the choice of Prithvi was primarily due to practical considerations (like availability or computational constraints) or if there were specific architectural features that made it more suitable for this particular application.

9. No code was provided as part of the submission. This would help to further assess the technical correctness and reproducibility of the results. Are the authors planning to open source their framework?
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes using foundation models to better the climate models predictions and specifically focusing on modeling the gravity wave effects via fine tuning. Almost all the authors agree that the technical contribution from ML perspective is minimal. I agree with the assessment of the reviewers that ICLR is perhaps not the best venue for this work.