Performance evaluation and optimization of a deep learning parameterization method trained from submesoscale-permitting ocean simulations

In the ocean, submesoscale physical phenomena O(100m) to O(1km) have been reported to play a key role in ocean oxygen ventilation, nutrient supply to the surface ocean, and carbon export, as well as the transfer of energy to larger scales [1]. However, due to limitations in computational resources, current ocean general circulation models are frequently run at resolutions on the order of O(10km) to O(100km) and cannot directly resolve submesoscale turbulence (i.e., subgrid-scale phenomena). Therefore, parameterization schemes are required to simulate these subgrid-scale phenomena.

Recent advances in machine learning have triggered the active exploration of data-driven approaches to parameterization for subgrid-scale phenomena that utilize data from observations and simulations. In previous studies [2, 3], the neural network is trained directly using the same variables as the neural network's output, such as viscosity and diffusivity coefficients. However, this approach does not guarantee that the inferred model parameters accurately represent the state of subgrid-scale phenomena they aim to reproduce. We propose a novel parameterization method for estimating diffusivity and viscosity parameters to parameterize subgrid-scale phenomena and have implemented this method in MITgcm, an ocean simulator [4, 5]. This method trains a neural network using the state variables (i.e., velocity fields, potential temperature, and salinity) derived from the simulation results at a resolution that can directly resolve subgrid-scale phenomena. Therefore, unlike previous studies, the diffusivity and viscosity parameters inferred by the trained network can reproduce the global state of subgrid-scale phenomena.

The ocean simulator MITgcm is implemented in Fortran, which does not have a built-in package to compute gradients within the neural network, in contrast to deep learning libraries (e.g., PyTorch) like Python. In our previous work [4, 5], we used a quasi-newton optimization method, which does not require computation of these gradients. However, the optimization performance of this method was limited. In this study, we use adjoint code within MITgcm to compute gradients for optimizing neural networks and examine the effect of different optimizers on training performance.

 

Acknowledgments
This work used computational resources of supercomputer Fugaku provided by the RIKEN Center for Computational Science through the HPCI System Research Project (Project ID: hp240394).

References
[1] M. Lévy et. al (2024), The impact of fine-scale currents on biogeochemical cycles in a changing ocean, Annual Review of Marine Science, 16(1), 191–215.
[2] Y. Han et. al (2020), A moist physics parameterization based on deep learning, Journal of Advances in Modeling Earth Systems, 12(9), e2020MS002076.
[3] Y. Zhu et. al (2022), Physics-informed deep-learning parameterization of ocean vertical mixing improves climate simulations, National Science Review, 9(8), nwac044.
[4] R. Irie et. al (2024), Parameterizing ocean vertical mixing using deep learning trained from high-resolution simulations, EGU General Assembly 2024, EGU24-2297.
[5] R. Irie et. al (2024), Optimizing a deep-learning model for parameterizing submesoscale phenomena in an ocean simulator, Workshop on Scientific Machine Learning and Its Industrial Applications.