Thank you for your questions. There clearly seems to be some confusion and we clarify it below. We will also add these clarifications to the final revised manuscript.

1. We completely disagree here because our choice to coarsegrain the fluxes is based on robust scientific principles and utmost care was ensured while creating the model training data. We arrived at the decision only after consulting multiple domain experts on gravity waves. We should clarify that the momentum fluxes used to the train the models are conservatively coarsegrained (and not simply interpolated) using Python's xESMF library, so they preserve all the information from the 25 km fine-grid by averaging along the longest-resolved wavelengths. Even if we were using a 1 km climate model output, the correct way to define the fluxes would be to coarse grain them onto a coarser grid appropriately chosen according to the longest resolved wavelength, otherwise, we would have to deal with ringing effects associated with wave phases. Simply put, even as the final grid seems coarse, it preserves all the knowledge from the high-resolution dataset and contains all the relevant information for climate models. All this would lead to high errors in ML-based predictions and would be rendered fruitless for climate science applications. Since climate models are typically coarse, the coarse resolution selected here (100-300 km) is exactly the optimal fit. Too fine of a resolution would mean that we are predicting wave phases - which would be a poorly-defined problem. In this 'coarsegrained' form, the fluxes could be used optimally by climate models as the traditional parameterizations too use a similar approach to compute momentum fluxes.

2. Yes, efforts are underway to test the online performance of this scheme. We understand the importance of online testing and have clearly acknowledged in the conclusion section that we are working on it. Since climate models are complex fortran codes, coupling the torch model to a climate model is a challenging technical problem - especially if we want the coupled ML model to provide the optimal speed (which we do). We had to overcome some key logistics problem to build a team that can effectively help us with this and we are now making quick progress on this. Moreover, most climate model parameterizations are tested offline first. So, this does not reduce the significance of this work at all as rigorous offline testing (the three proposed tests) is key to ensuring stable online performance later. Also, we appreciate that this may not have been clear from the text but we are not “forecasting climate” here, rather representing missing physics in climate models at any given instance by learning from high-resolution reanalysis data. Otherwise, this line of thought would seem to suggest that most downstream tasks proposed in foundation models like ClimaX, Aurora, or Prithvi have no value because they are offline; this is clearly not the case. While our downstream tasks are not online yet, its rigorous offline testing has inherent value to ensure optimal online performance and the analysis presented here comprises more than half of the entire problem.

We thank you for your thoughtful comments. Most points raised are not actually weaknesses and we explain why:

1. Choice of Foundation Model: Aurora’s stable code was publicly released on August 22 2024. By then we had already completed our analysis and started writing the manuscript, so it made little sense to go back and redo the analysis just because there is a new foundation model in town. Also, it seems like Aurora pretty much focuses on weather forecasting and air quality forecasting as the only downstream task, so science-wise there's not much difference between Aurora and Prithvi. AtmoRep does not discuss any downstream tasks whatsoever in their non-peer reviewed paper. If the authors themselves have not discussed many downstream tasks in their paper, it makes us a bit skeptical using the model for our analysis as well. ClimaX is too coarse resolution and trained on potentially-biased CMIP6 model output to serve as a meaningful/good choice for climate science applications. Thus, we find Prithvi WxC to be more balanced in terms of both pre-training and fine-tuning applications. Despite not using the models, we reiterate that intercomparing the foundations models is simply not the focus of this study. That said, the goal is also not to suggest that Prithvi is the optimal foundation model for the task of climate model parameterization development. The goal is to open the field of climate science/modeling open to the broader AI community by devising novel use cases in the domain and providing a robust roadmap or proof-of-concept. In principle, any well-designed FM could be used for this ML parameterization development task. We just demonstrate it using Prithvi since it was developed by a mix of both ML and climate-science domain experts and has a well-defined encoder-decoder architecture.
2. Features and data selection: we should clarify that using different input variables for the two models does not create any discrepancies whatsoever because (as also mentioned in the manuscript) the potential temperature (𝛉) used as input for one model is simply a product of temperature (T) and pressure (p): 𝛉 = T*(p)$^{constant}$. Thus, both models contain the same amount of physical information, just in different dimensions, and these models are robust enough to not be affected by these differences. We like your point though, and we would be happy to add the loss curves for model training using different feature combinations in the revised manuscript.
3. Choice of baseline: we used an Attention UNet as a baseline because it is a state-of-the-art method for dense prediction. Even previous papers like ClimaX compared their forecasting performance to a UNet and ResNet (though not time series). While we are not comparing our study to theirs, we have certainly drawn inspiration from their approach. Also, since Attention UNet has been accepted well by the ML community, we are confident it can serve as a very effective baseline for our instantaneous mapping task. The fine-tuned model is definitely using 2.3B parameters, but to note, we have frozen the weights of the model so these parameters are not changing, only trainable parameters are the convolution around the Prithvi WxC. Sure, we can definitely go ahead and train a transformer for it, but we also want to acknowledge and build on the accepted published baselines for this task (the Gupta et al. (2024), ICML paper, https://arxiv.org/pdf/2406.14775)
4. Sure, we would be happy to share it in the revised version.
5. Log-scaling: Please note the log scaling of the y-axis. The dotted lines are indeed the 2.5th and 97.5th percentile (see linear plot here: https://tinyurl.com/iclr25)
6. Hellinger distance: unfortunately, we could not find a super-relevant study which could act as a reference here to interpret the exact values of the hellinger distance. However, our decision to use 0.05 is explained and supported by monte carlo Gaussian sampling conducted over thousands of Gaussian samples, where we found that a Hellinger of 0.05 for standard Gaussians is explained by a roughly equal, i.e.,~60% perturbation in the mean or a ~60% perturbation of the std. dev. This also makes interpreting results easier in terms of increased spread or increased shift in the mean, i.e. if the shapes are similar, the Hellinger distance above 0.05 can be interpreted as being large due to changes in standard deviation or due to changes in standard deviation. Would be happy to elaborate further in the revised manuscript supported with appropriate figures. Since the Hellinger distance ranges from 0 (identical) to 1 (disjoint), a distance of 0.05 definitely ranges on the lower side indicating the distributions are indeed strongly similar with minimal divergence. When compared to other measures (e.g., KL divergence, Wasserstein distance, etc.), 0.05 often aligns with very small discrepancies. In high-precision domains, a Hellinger distance of 0.05 might be considered excellent.

Thankyou so much for your thoughtful comments. We clarify why the concerns do not really classify as weaknesses.

Weaknesses:

1. Innovativeness: The use of AI to empower climate modeling (not weather forecasting) has been severely limited. Our work opens avenues for the broader ML community to develop AI models to advance climate modeling and climate science (which has a different set of challenges than weather forecasting). We also introduce the tail-Hellinger distance, a novel metric focusing on the accuracy of the predicted tails—crucial for capturing rare but impactful events in any domain where ML is applied to learn distributions. Our work bridges ML advancements and high-impact applications in climate science (not weather), promoting interdisciplinary research and deployment of FMs in other scientific fields, and provides a pathway for ML experts to foray into climate science. It is worthwhile mentioning ClimSim - which won the best paper at NeurIPS 2023. ClimSim did not introduce any new model architecture, but provided a dataset/pathway to allow ML experts to work with ideas in climate modeling. We have appropriately submitted our work to the “Applications to physical sciences” area of the conference.

2. Experimental Setup: We disagree here. The two methods have been used in different contexts. Eqn disc has been used for momentum closures in the ocean (where analytic forms are not known) and the prob. models focus on combining low-fidelity and multi-fidelity datasets for precipitation. Both setups also require different datasets than ours. Moreover, in our case, the analytic forms are already known but not resolved. We searched the literature and found only one study which has looked at (resolved) gravity wave fluxes (Gupta et al. (2024); https://arxiv.org/pdf/2406.14775). Subsequently, here we used their Attention UNet baseline. Using different inputs for the two models does not create any discrepancies because (as also mentioned in the manuscript) the potential temperature (𝛉) is simply a product of temperature (T) and pressure (p): 𝛉 = T*(p)$^{constant}$. So, both models have the same physical info, in different dimensions, and these models are robust enough to be unaffected by these differences.

3. Experiments: Thanks for the comment but this is definitely not the case (and does not qualify as a weakness). MERRA-2 has (a) a factor two coarser grid (0.5 deg x 0.625 deg as opposed to ERA5’s 0.25 deg) and (b) has finite volume numerics. As a result, MERRA2 does not provide a competent GW field (Li et al. 2023, https://doi.org/10.1002/qj.4605). Also, including (weak) GW fluxes from MERRA2 won't be informative because the small value prediction affects the UNet baseline and the finetuning model alike. Otherwise, the UNet baseline, which is pretraining agnostic, would have shown better skill in predicting small values. If anything, we see the opposite, i.e., pretraining on one dataset and finetuning on the other leads to improvements wrt baseline.

Questions:

1. A similar strategy can be applied to develop other params e.g. clouds. We do not suggest using ERA5 to develop these since it under-resolves convection. One can use a mix of high-resolution climate model output and satellite irradiances and provide task-specific input data like specific humidity, saturation pressure, precipitation fluxes, latent heat fluxes etc. from other high-res datasets to develop these ML params. The baseline models would also vary. For GWs, we only found one existing baseline — the Attention UNet (Gupta et al.) — to predict the small-scale fluxes, and so we used that to inform our finetuning. Similarly, in the future, SOTA benchmarks for cloud params & pecip. (and other processes) could be used to compare the performance of finetuning models. Alternatively, in case of a lack of a baseline, encoder-decoder pairs from other FMs can promote effective intercomparison of param. architecture.

2. Data-scarcity. Past ML param. studies: Wang et al. (2022)(https://doi.org/10.1029/2022MS002984) use ~4.5 mil samples as training+validation set. Espinosa et al. (2022)(https://doi.org/10.1029/2022GL098174) use ~11 mil samples. Similarly, Zanna and Bolton use 10 yrs of high-res training data. Temporal coverage: we (purposely) used only 4 yrs of data for training. This could lead to: (1) underrepresentation of tropical convection and hence gravity waves: processes like the El-Niño Southern Oscillation have a typical period of 2-7 yrs but are not fully represented. Similarly, the quasi biennial oscillations (QBO) in the tropical stratosphere have a period of 28 months. Since we sample our data over 2010, 2012, 2014, 2015, the training data does not cover both phases of the QBO.

3. Of 48 months, 47 months training + 1 month (May 2015) validation. We also trained the baseline on 3 yrs and tested it on 1 yr and found similar results & losses. Training data was randomly shuffled.

Thank you for your detailed responses. Your clarification regarding the "weaknesses," especially the explanation of the "experiments," has addressed my concerns about the completeness of the work. However, I remain uncertain about the paper’s contribution to interdisciplinary research, particularly with regard to its innovative aspects. While this is certainly a strong application of machine learning to climate modeling and forecasting, from a research paper’s perspective, it is essential to emphasize the novelty or irreplaceability of the method.

The reason ClimSim received recognition at NeurIPS 2023 is due to the importance of quality datasets in advancing machine learning, much like the role ImageNet has played in computer vision. While excellent algorithmic applications can have a similar impact, this paper does not sufficiently highlight the unique contributions it makes to the field, community, or society at large. Although I understand the authors' frustrations, I believe that, as an interdisciplinary piece, the paper should more clearly address the varying expectations across fields.

In particular, for a submission to a leading machine learning conference, the paper should be structured around the interests of that community. For instance, it would be beneficial to clearly explain the reasons behind the "data scarcity" issue (thank you for addressing this in your response), and to discuss the dataset split—whether using only 2% of the data for testing may affect generalizability.

Based on these points, I will not change my score, unfortunately. I would recommend considering submission to the "Applied Data Science Track" at a top-tier machine learning conference or to a high-impact journal in the climate science field.

We thank the reviewer for putting in valuable time to review our manuscript. It seems like they are mixing two very different things here but that is ok -- we try to answer them below:

• The goal of this study is to demonstrate ONE WAY in which foundation models (FM) can be relevant to the domain of AI for climate science modeling (not weather forecasting). ClimaX and Aurora (and even Prithvi) otherwise only address downstream tasks relevant over weekly to monthly timescales (which is weather, not climate!). The study DOES NOT CLAIM that Prithvi is the only FM fit to do so. That is simply not the point of this study. And we think the reviewer also understands this which is why they suggest “possible extensions”.

• Climax has many loose ends: ClimaX – an ICML workshop paper – was trained on a very coarse dataset of 5 degrees. Even for rollout forecasting the authors used convolution models instead of time series models (UNet and ResNet). For some surprising reasons, ClimaX's authors did not compare their results with (or acknowledge) Fourcastnet, which was the SOTA at that time. However, in this paper we are not creating a comparison between models anyway, rather creating a case for downstream application for climate model (and not weather model) improvement. Given that this is the first study (to our knowledge) to apply FMs to address a physical problem which affects climate (and not weather) models over yearly and multi-year timescales, we strongly disagree that this paper does not have any ML novelty. Prithvi is higher resolution compared to ClimaX and is trained on MERRA2 dataset. We are not presenting Prithvi WxC here, and are merely using it, so we can’t compare its performance for all the tasks. So, we think this concern raised by the reviewer might better belongs to their paper, not ours.

• Aurora code was unavailable before August: Aurora code was released on August 22. We already finished our analysis and started preparing our paper by then. So, it does not make much sense to then include a half-cooked analysis in our paper just for the sake of it when it is not needed. Our conclusion stands on its own. This is also consistent with ICLR policy regarding use of fresh work for comparison. But we appreciate that the reviewer might not be aware of this. Aurora's stable code was released on August 22. So its not practical for us to run a completely new model within 2 weeks of submission due date and write down the results.

WP #3

• This is not correct - it is NOT A "super low resolution" task as you think. The manuscript clearly mentions that the fluxes are “conservatively coarsegrained” and not merely interpolated. This means that the interpolated data contains all the information of the 25 km grid, then projected onto a 300 km grid. Conservatively coarsegraining fluxes is the scientifically accurate way to represent them as the fluxes are only defined in terms of wave averages. The physically accurate way to design a flux prediction experiment is to coarsegrain them to a 100-300 km grid (as that is the scale of the longest gravity waves). This consideration is based on past atmospheric wave dynamics literature (Polichtchouk et al. (2022) for instance; https://doi.org/10.1002/qj.4202). Now, since you mention ClimaX, we should mention that their downstream tasks are actually quite low-resolution because they simply coarsen the input-output pairs to a coarser grid by discarding the finer-scale details. Thus, the task presented here presents a robust test to assess “small-scale” predictions. We appreciate that this subtle detail might not be apparent to reviewer who may be approaching this problem from a pure data-driven perspective.

WP #4

• This comment is a bit ambiguous too, and we strongly disagree with the reviewer. We have used the most recent SOTA baseline to compare our results for gravity waves downstream task. We searched extensively to find the most appropriate baselines for the gravity wave flux task but found only one study (Gupta et al. (2024); ICML 2024- https://arxiv.org/pdf/2406.14775). That study establishes Attention UNet - a SOTA model for dense prediction - as a baseline. We were also inspired by the fact that if ClimaX can use a basic UNet to compare their forecasting results without any temporal component whatsoever, we can definitely use an improved version of their choice to serve as a baseline here.

• Lastly, we feel sorry for the reviewer to be triggered by our reference to Prithvi as SOTA model, even though there is not reason to be. We decided to used Prithvi for our application because it was developed by a team of both ML and climate/weather domain experts. Not sure why the reviewer sounds offensive.

We thank the reviewer for their detailed comments, and for acknowledging that our work has the potential to be useful to both the ML & the climate science community. However, we do disagree on the degree to which this study could be useful to the broader ML community. Below we explain why.

1. Solving climate science use cases can benefit other ML problems too: An increasing number of AI weather forecasting models and foundation models (FMs) have been developed over the past 3-4 years. The release of FourCastNet [1] inspired both the weather forecasting and machine learning community alike to develop more stable weather forecasting models. Likewise, the release of ClimaX [2] proof-of-concept inspired the development of more versatile FMs like Aurora [3], AtmoRep [4], and Prithvi [5], all of which involved heavy collaboration between the weather forecasting and the AI community. Yet, the application of these FMs to climate prediction tasks has been severely limited. Either we wait for these models to be stable over multi-year timescales or we find novel ways to bypass these limitations and use AI to improve climate models (not weather models) now. Literally none of the downstream applications for FMs focus on reducing climate uncertainty – because it typically involves analysis and rollout on longer timescales over which these AI models are not stable, and significant rapid advancements will be needed before they could do so. Our analysis establishes that this should not discourage the use of these models for climate prediction and climate modeling. This is best accomplished by introducing new climate science use cases more accessible to ML experts and by developing new probabilistic metrics potentially useful in broader ML. Therefore, the focus of this study is not just to highlight the prowess of FMs for downstream applications (which has been amply done before), but also pioneer a new research direction to use AI to advance climate modeling and climate prediction, at the same time aligning with societal needs. Presenting this at ICLR and sharing this with the ML community is the best way to inspire more research in this direction.

2. Submitted to “Application to physical sciences” area: we totally appreciate that the idea of using a finetuned model to perform weather-related downstream tasks has been put forth by past studies. However, this is the first study to leverage weather-focused FMs to create a viable downstream task for climate model development and the state of the art Attention Unet baseline (state-of-the-art for gravity wave analysis as least) has been used to create a new but similar finetuning model architecture. By doing so, we confident we are redefining the limits of what problems (domain and timescale) weather FMs can be used to address, potentially motivating more climate-focused applications in the future, also creating room for development of better model architectures. In the same spirit, we have submitted our paper to the “application to physical sciences” sub-area of the conference.

3. New probabilistic metrics: The tail Hellinger metric introduced in the study can be used more broadly across different problems, especially in studies which focus on simulating distributions and their tails, like in extreme event analysis and modeling of intermittent systems.

4. Relevant audience: submitting our paper to ICLR therefore ensures that we get the optimal audience whose focus is both on identifying novel avenues to apply ML to advance long-term climate analysis through development of new models and use new metrics to quantify performance of machine learning models to predict intermittent nonlinear dynamics.

References:

[1] Pathak, Jaideep, et al. "Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators." arXiv preprint arXiv:2202.11214 (2022).

[2] 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.

[3] Bodnar, Cristian, et al. "Aurora: A foundation model of the atmosphere." arXiv preprint arXiv:2405.13063 (2024).

[4] Lessig, Christian, et al. "AtmoRep: A stochastic model of atmosphere dynamics using large scale representation learning." arXiv preprint arXiv:2308.13280 (2023).

[5] Schmude, Johannes, et al. "Prithvi WxC: Foundation Model for Weather and Climate." arXiv preprint arXiv:2409.13598 (2024).

Here we provide detailed answers to your questions 2 and 3.

2. Yes, this is an important issue and is reminiscent of the issues common to AI weather forecasting models like FourCastNet, GraphCast, etc., where the models tend to learn the large-scale features better but struggle to accurately represent the small-scale fronts, filaments, atmospheric rivers, etc. Moreover, the high correlation coefficient in the instantaneous fluxes in Figure 7 show that most of these small values might be appearing outside of the selected hotspots. We propose a couple of solutions to tackle this issue:

• A more physics-informed loss function, for e.g, regularizing the predicted distribution towards the true distribution

• A latitude-weighted loss function – since the correlation coefficient is weaker in the tropics, a latitude weighted loss function can improve/reduce the errors in the tropics

• A different scaling for fluxes - strong tendency for small values to be treated as noise. While the cuberoot scaling of fluxes works very well in yielding a plausible climatology, it may make it difficult to learn the small-scale values as the cuberoot tends to push small values away from zero and large-values towards one.

3. This is a good question, thanks! The traditional Hellinger distance ranges from 0 to 1, 0 indicating that the two distributions are identical for the whole sample space and 1 indicating that the two distributions are completely disjoint. The tail-Hellinger distance envisioned in this study, however, has a different range. We provided some context how to interpret values for the tail-Hellinger metric in the manuscript but would be happy to elaborate more in the revised version.

• Essentially, tail-Hellinger zooms in on the tails assuming that the bulk of the distributions are similar (if not identical). With ½ as a fixed constant, a positive increment to tail-Hellinger come from the second term which is simply the integral of the tail for the predicted distribution. If the bulk is identical for both $p$ and $q$, then the second term (integral of $p$) should be equal to 0.5 as well. If subsequently the tails are disjoint, then the tail-Hellinger will be equal to 1. In this case, the interpretation will be that the bulk of the distributions are identical, yet the tails are totally different/disjoint.

• Any negative contribution, that makes the tail-Hellinger negative comes from the cross-term between $p$ (predicted) and $q$ (truth). If the tails of the two distributions are disjoint then the contribution from the cross term will be zero and the tail-Hellinger will be nudged toward more positive values by the second term (scaled integral of $p$). On the other hand, a negative value would imply a fatter tail of the predicted distribution than the true distribution because $\sqrt{p(x)}\sqrt{q(x)} \geq \sqrt{p(x)}\sqrt{p(x)}$ over the tails and thus.

• So, if the Hellinger distance is (say) -0.25, this means the model is underestimating the tails, i.e. $\sqrt{q(x)} > \sqrt{p(x)}$. If the bulks are identical, this would mean that the range of $q$’s tails is narrower than the range of $p$’s tails but since $q (x)$ > $p(x)$ $\forall x$, the cross-integral $\int \sqrt{q(x)}\sqrt{p(x)} dx$ > 0.5 + $\int \sqrt{p(x)}\sqrt{p(x)} dx$, leading to negative values. The more the prediction underestimated the tails, the more negative is the tail hellinger distance. The fatter the tail of $q$ relative to $p$, the more negative the tail-Hellinger. However, the tail could be fatter on either one side or both, and the tail-Hellinger would not be informative in suggesting which (limitation).

• It is worth mentioning that, negative values of tail-Hellinger are more likely for small values of $\epsilon$ as that increases the probability of the bulk of the distributions to not be identical.

• Ofcourse, this is the net sum of the tail product over both tails of the distributions - sort of like a two-sided student t-test. A more nuanced metric could be defined to focus on just the left or the right tail. Another possibility is to not assume similar bulks and assume systematic biases in the mean and account for such biases when computing the distributions distance. Since, it was not super relevant for this study, we did not discuss these ideas in the Appendix.