VMNN: An Accurate and Practical Neural Network for Ocean Vertical Mixing
摘要
Accurate and efficient ocean vertical mixing simulation is critical for Earth System modeling. Traditional numerical methods often trade computational efficiency for accuracy, while existing AI-driven approaches, though efficient, overlook two key aspects in system design. Specifically, spatial context from adjacent columns is omitted in dataset construction, and existing models fail to adequately capture local spatial correlations and vertical dependencies. The absence of reasonable physical modeling hinders the accuracy of these systems and limits their practical application. To address these issues, this study first constructs a high-quality dataset that incorporates spatial information. Building upon this, we propose the Vertical Mixing Neural Network (VMNN), which employs convolutional layers to extract and compress local spatial information, followed by a combination of Transformer and Bi-LSTM layers to capture vertical sequential dependencies. We then demonstrate the superiority of VMNN in both experimental scenarios and practical applications. In experiments, evaluation results on the dataset demonstrate that our model outperforms existing methods in accuracy, further validated by ablation studies that confirm the rationale behind the model architecture. In practical applications, to our knowledge, VMNN is the first deep learning model integrated into Earth System Model (ESM) that achieves comparable numerical accuracy and physical consistency with traditional methods, while enabling at least a 2 × speedup through GPU acceleration. These findings demonstrate the feasibility of AI-driven approaches for Earth System modeling and their potential to improve computational efficiency in climate simulations.