Xia et al. apply a multi-head self-attention (MHSA) transformer machine-learning model for aerosol chemistry for the first time in the WRF-Chem CTM. The results are promising and fascinating. However, many issues with language and the presentation of the results need refinement and tempering. My comments are below:

Abstract:

I would not use the descriptor "Remarkably" if these are similar speedups we are seeing in the literature for CPU and GPU speedups of ML chemical solvers. 

"While global long-term simulations have not yet been implemented, AIMACI’s robust generalization capability, coupled with our easily plug-and play solution, paves the way for its coupling into global climate models for further testing in near future. This advancement promises to enhance the precision and efficiency of atmospheric aerosol simulations in climate modeling."

-- I would temper this assessment. The results do indeed seem promising but you have not shown stable, global, year-long simulations using this fully ML-learned replacement and past studies have shown that this is not guaranteed at all even if shorter-term simulations work.

Intro:

How do aerosol schemes compare with gas-phase chemical mechanisms in terms of computational cost? Aren't many schemes dealing with heterogeneous chemistry separated from the gas-phase mechanism and incur lower computational overhead? 

Perhaps slightly more information/discussion is needed to contextualize the cost of the aerosol scheme compared to other components of chemistry/climate models. The background on the previous AI models is great, however. 

L54: "methodologies for describe the evolution of PSD", should be "describing"

"Unlike photochemistry which only involves chemical reactions between species, the full aerosol chemistry and interactions encompasses numerous other intricate processes such as nucleation, coagulation, thermodynamics" --> Aren't aerosol schemes also coupled to the chemical mechanism as well? Or are you discussing only microphysics/thermodynamic calculation of aerosol schemes? It is a little vague so far

"These advancements hold great promise for the future of climate modeling, enabling fast, accurate, and stable simulations of aerosol chemistry and interactions, thereby reducing uncertainties stemming from simplified representations of these processes." -- I don't think you can claim that this is stable because the time scales are not long enough for climate-relevant (or even seasonal) time scales. They are stable for 1 month, which is good to see but no guarantee that they are stable beyond this (as you have not shown it)

Methods:

"Therefore, here we innovate by pioneering a novel Artificial Intelligence Model for Aerosol Chemistry and Interactions (AIMACI), leveraging the MHSA algorithm. The MHSA algorithm, serving as the foundational architecture of state-of-theart transformer models, has been successfully" -- Please tone down this language, it reads as very pretentious. Multi-head attention transformers have existed in the CS/NLP literature for a while. You can claim this is new for aerosol chemistry. 

There does not seem to be much information on how you trained the MHSA transformer either. Did you fine-tune on a pre-trained model? Did you build the MHSA from scratch or use an out-of-the-box GPT model? Or was this completely borrowed and reworked from the other Xia et al paper? Either way, more information is needed on the model development information, and how the training was done (what resources/wall time it took). E.g., what does the tokenization of the WRF-Chem data look like? How were these patches determined/designed, etc? 

I don't find Figure 1 very informative or helpful. Doesn't situate where the ML model is within WRF-Chem, does not give a sense of the dimensionality either. Table 1 is more useful than this, so please rethink this use of space

"After training, the AIMACI scheme was flexibly coupled into WRF-Chem, utilizing TorchScript and Libtorch tools officially provided by PyTorch. This coupling approach encapsulates the AIMACI scheme within a static library, minimizing alterations to the original codebase while offering a lightweight, adaptable, and easily plug-and-play solution" -- How does this approach compare to the other common approach of using C Foreign Function Interface (CFFI) to create C-style bindings for Python scripts? The latter also seems flexible and lightweight without altering the code base. So if there is an inherent advantage of the TorchScript approach that should be stated. 

"the initial 16 days from 2019-03-02 00:00 UTC, were designated as the training set, the penultimate day served as the validation set, and the final day constituted the test set." --Does the training or the results change if these days are random? This seems like a very short training time window to draw generalizations for aerosol chemistry across an entire domain. Is this long enough to cycle through all the species residence times/lifetimes? 

Results:

"The results are promising with an average R² of 0.99 for all 37 evaluated species." -- Could you have an SI figure plotting the concentration distribution of all species? I wonder if sectional aerosols are easier or harder to emulate if their distributions are normal, flat, or uniform to each other. 

"3.2 Offline Single-step Simulations with the AIMACI Scheme " -- Should this be Online?

Figure 3: Can you provide relative error metrics as well? Seems like AIMACI over/underpredicts SO4 depending on the bin compared to the numerical model. Is there a reason that drives this difference? 

"A notable aspect of the AIMACI scheme is its grid-based training and prediction methodology, which contrasts with existing AI large models such as Pangu (Bi et al., 2023) and Fengwu (Chen et al., 2023) that operate on entire fields" -- This was not stated in the methods, it would help to give more context to the model. Do you not tokenize the entire field?

"In Figure 3, the AIMACI scheme has successfully captured and reproduced the intricate spatial patterns of sulfate column concentrations across different particle sizes with R2 values all exceeding 0.88, even after a prolonged 10-day simulation. This achievement underscores the AIMACI scheme's exceptional stability and accuracy. " --But the hotspots are over/under-predicted and it seems like RMSE increases after 5 days, and oscillates within each day. Perhaps this error would grow outside of the training time horizon. 

"suggesting that these discrepancies do not lead to runaway error growth. This sustained performance further substantiates the AIMACI scheme's reliability, positioning it as a robust tool for extended atmospheric and climate simulations. " -- The language throughout is too strong, championing this model as revolutionary yet it's tested on short time windows. It very well may be, but the results presented in this paper do not warrant this kind of treatment. We do not necessarily expect runaway error growth in 10-day simulations and do not see that is the case in other studies as well. It is over longer-term time scales (e.g., 1 month +) where long-term stability is a concern, whether due to chemical lifetimes cycling over, seasonal weather patterns changing, etc. 

Figure 4. Results here look promising but please provide both absolute error and relative error plots that accompany the figure. Looking at raw concs is not helpful beyond a simple eye test

Figure 5. Seems like the AIMACI starts to drift over time (though very small) wonder what happens if simulation time extends longer and longer

Figure 7/8: Should also have absolute and relative error plots. Need to see where/why AIMACI is over/under-predicting. It is remarkable that the model has learned the fields well enough to solve for different months. But again: 1) a one month simulation is not long enough to claim that this is long-term stable, especially when we see in the time series there starts to be a drift over time, and 2) there are obvious mismatches of hotspots in the map

"Our training dataset, which comprises only partial data from March, may not adequately encompass the unique meteorological conditions associated with extreme events such as typhoons." --Points like this need to be made rather than saying you have 'pioneering' results. Are the errors in the time series caused by meteorology or by chemical conditions not encountered before? Is there a causal way to determine what is happening here rather than speculating? 

If you could simulate AIMACI for 3 months starting in Spring that would help determine if this model is truly stable. But honestly, a year minimum is what is necessary in the field of atmospheric chemistry and climate, but I do realize that that incurs large computational costs and may be outside of the scope here. But if that is the case, then the discussion of stability needs to be tempered. 

"Given that the WRF-Chem, written in Fortran, is not conducive to GPU acceleration, we conducted offline tests of the AIMACI scheme's computational speed on a GPU and compared it with the numerical scheme on a CPU, where the AIMACI scheme was coupled into the WRF-Chem. " --Some may argue that this is not an apples-to-apples comparison but I am ok with it. These speedups seem similar to Kelp et al (2020) and Liu et al (2021). 

This work introduces the Artificial Intelligence Model for Aerosol Chemistry and Interactions (AIMACI), trained on 16 days of WRF-Chem simulation data. When integrated with the online WRF-Chem, AIMACI demonstrates high consistency and accuracy in modeling inorganic aerosols, maintaining stability over one-month scale online simulations. Furthermore, AIMACI exhibits significant computational speedup compared to conventional numerical schemes, highlighting its potential to overcome the computational challenges of traditional methods. The comments listed below are minor clarifications. Once these points are addressed satisfactorily, I believe the paper will be suitable for publication in ACP.

1.    Line 23：“8 aerosol species including water content in aerosols”

But the main text mainly discusses four species (i.e., SO42-, NO3-, aerosol water, and number concentration of aerosol).

2.    Line 28-29：“paves the way for its coupling into global climate models for further testing in near future. This advancement promises to enhance the precision and efficiency of atmospheric aerosol simulations in climate modeling”

Currently, the AIMAC demonstrates stability only on a one-month scale. Additionally, the training data is derived from a climate model, meaning that AIMAC only reproduces model simulations and does not enhance the precision of the climate model.

3.    Line 38-39 “aerosol chemistry and interactions”

Please specify which physical and chemical processes are included.

4.    Line 93-99 “The results demonstrate that, … thereby reducing uncertainties stemming from simplified representations of these processes”

Results are not recommended to be included in the introduction section.

5. Line 116-119: It is recommended to introduce the chemical schemes for sulfate and nitrate. Additionally, how are the emissions of primary aerosols and precursor gases configured in the model?

6. Line 150-151: “Input Embedding Layer: This initial layer receives meteorological variables and chemical species as input features”

Is the training conducted separately for each grid point? How does the model account for interactions between different grid points?

7. Line 166-167: “The simulation result was segmented as follows: the initial 16 days from 2019-03-02 00:00 UTC, were designated as the training set”

Why was such a short training period chosen? Is this the minimum period necessary for training, or does the training cost increase significantly with a longer period?

8. Line 172-173: “Each training sample included 65 input features (4 meteorological variables, 5 gas species, and 14 aerosol species with 4 size bins)”

Precursor gases and atmospheric oxidants (e.g., the hydroxyl radical, ozone) are crucial for the secondary aerosol chemical process. Why are these not included as input features? Additionally, since the training data is hourly, why is the solar zenith angle not considered? Furthermore, the latitude and longitude are also absent.

9. Line 175-176: “the concentrations of other inorganic mass (OIN), mineral dust, black carbon (BC), organic carbon (OC), calcium (Ca2+), and carbonate (CO32-) are not altered”

But coagulation between aerosols can potentially lead to changes in the particle size distribution and also affect the aerosol number concentrations.

10. Line 177: “they play a significant role in affecting the acidity or alkalinity of the atmospheric environment”

Black carbon and organic carbon do not affect acidity. Why are they included as inputs?

11. Line 183: “minimizing alterations to the original codebase while offering a lightweight”

It is not clearly specified where the AIMACI scheme is integrated within WRF-Chem. Additionally, it is unclear which specific model processes the AIMACI scheme replaces. What modifications are needed for the interaction and feedback between the AIMACI scheme and subsequent model processes within WRF-Chem?

12. Line 186: “Furthermore, we conducted three sets of additional experiments”

It is recommended to distinguish between offline and online AIMACI simulations. Additionally, since there are several numerical experiments, it is suggested to list these experiments clearly in a table.

13. Line 232-237: “Atmospheric aerosols significantly impact the climate system … climate models for precise climate simulations”

It does not add much value. It is recommended to provide a more detailed discussion on the results, such as comparing the spatial distribution correlation between the AIMACI and numerical schemes.

14. Line 244-252: “Sulfate, derived primarily from … role in atmospheric processes” 

This is not a description of the results but rather a lot of background information. It is recommended to not include it in this section.

15. Line 266: “3.2 Offline Single-step Simulations with the AIMACI Scheme”

Online AIMACI simulation?

16. Line 286: “Figure 3” & Line 307 “Figure 4” 

It is recommended to supplement the spatial distribution of relative errors.

17. Line 294-295 “A notable aspect of the AIMACI scheme is its grid-based training and prediction methodology, which contrasts with existing AI large models such as Pangu (Bi et al., 2023) and Fengwu (Chen et al., 2023) that operate on entire fields.”

What is the grid-based methodology? How does it differ from large models?

18. Line 305-306 “This sustained performance further substantiates the AIMACI scheme's reliability, positioning it as a robust tool for extended atmospheric and climate simulations.”

Please be careful in your discussions. The results in Figure 3 are based on only 10 days of testing. They do not prove reliability for climate simulations.

19. Line 321 Figure 5: It is recommended to add hourly ticks in Figure 5. Additionally, why is there no apparent diurnal cycle pattern?

20. Line 328 “Despite these pronounced fluctuations”

It seems that fluctuations increase after the fifth day (e.g., sulfate). What could be the possible reasons for this?

21. Line 372 “Figure 7”

Is it an online AIMACI simulation? Please clarify.

22. Line 381-383 “The AIMACI scheme, despite being trained on data from only 16 days in March, demonstrates a remarkable ability to reproduce these distribution characteristics across different environmental conditions”

Why can AIMACI scheme trained from March generalize to other seasons? Does your training data sufficiently cover these scenarios, or does your model have some extrapolation capability?

23. Line 403-404 “This phenomenon is closely related to the interactions between aerosols and other processes within the physics-AI hybrid model, which differ significantly from offline simulation scenarios”

I would like the authors to provide a more detailed explanation, such as comparing the differences between offline and online AIMACI schemes.