MindSpore AI科学计算系列 | 初探海洋大模型AI-GOMS,打开大模型在海洋方面的大门
作者：昇思MindSpore点击蓝字
MindSpore
关注我们
作者：于璠 来源：知乎
背景
天气和海洋有着难以分割的关系，海洋环流与大气之间的热量交换对天气的变化有着巨大的影响。目前在中期天气预报领域已经有大量的AI大模型展现出了预报高精度天气的能力，其中代表性的工作有GraphCast、Pangu Weather。在短临预报方面，DeepMind最先开始尝试使用生成式模型进行降雨预报，以及清华大学龙明盛团队提出引入物理机制的NowcastNet。我们认为，海洋会是AI在气象领域发力的下一个战场。近期，清华大学黄小猛团队与华为先进计算与存储实验室合作提出AI-GOMS海洋基础大模型，打开了大模型在海洋方面的大门。
海洋预报与大气要素的预报有相似点也有不同点。在AI的做法上，海洋领域的天气预报在形式上类似于中期预报模型，也依赖于再分析资料。大气领域常用的再分析资料有ERA5，这也是Pangu等天气大模型实验的数据，而海洋领域同样有HYCOM等海洋模式。因此，在形式上我们完全可以迁移部分的中期预报模型在海洋的预报上进行尝试。
而不同的是，海洋与大气的物理机制不同，倘若我们希望在模型中融入物理机制，需对模型进行相应的调整。此外，海洋除了海面温度、海浪高度等物理层面的预报，也会有如含氮量、叶绿素含量等生化相关的要素预报，这同样需要在模型下游采取不同的策略来得到相应的预报。目前海洋方面的AI探索还处于初期阶段，这为大模型提供了巨大的机会。本文将介绍最新的AI-GOMS海洋大模型。
1、AI-GOMS
AI-GOMS是一个基于傅里叶算子的自编码器基础模型，它的Backbone以初始场的海洋模式基础要素以及边界条件为输入，对5类基础海洋要素进行预报，包括sea temperature(T), sea salinity(S), zonal velocity(U), meridional velocity(V)和sea surface height(SSH)。AI-GOMS以HYCOM全球再分析数据为输入，进行未来30天的天级全球预报。
AI-GOMS的模型架构如图1所示。由于海洋模式是格点状的输入，作者在这里采取类似ViT的操作，把2维的格点数据划分成non-overlapping patch来转变成1维序列token，以便于后期进行的编码。在图1 a子图中，我们可以看到模型对输入做了随机掩码(Random Mask)的操作，这是为了让模型能够学习到数据内在的特征，缓解过拟合来得到更好的长期预测效果。经过随机掩码之后的token，将经过Fourier Module来把输入映射到频域空间。Fourier Module的设计参考了AFNO的做法，其优点是可以让输入的序列长度是可变的，这使得前面的随机掩码操作变得可行。后面的Decoder Module会在频域空间对Fourier Module输出的特征进行适当的筛选，选取有利于预报的要素特征，最后再经由Projection得到基础的海洋要素预报。
图1. AI-GOMS模型架构
基于Backbone的基础要素预报，文章还提出了3个下游的轻量级预报任务(Downstream Module)，包括生化要素预报、下采样、海浪预报。Downstream Module的设计如图1 b子图所示，Backbone的部分是a子图的输入直到Decoder Module的部分，这部分的输出做为Feature Tensor。
在Backbone的模型中，Feature Tensor是经由Projection输出基础要素的值。而在Downstream Module中，Feature Tensor将和下游任务的额外输入(Downstream Input)进行结合，再通过一个轻量级的微调网络从而得到下游的输出。对于不同的下游任务而言，Downstream Input和相应的轻量级微调网络是不同的。在下采样中，Downstream Input是1/4度的Kuroshio的5个要素变量，经由Residual Convolution blocks和upscaling ConvTranspose2d layer输出1/12度的高分辨率Kuroshio格点数据。生化要素预报和海浪预报的微调网络均是1维的AFNO算子和Projection layer。
海浪预报所需的Downstream Input是从ERA5数据中提取的初边始条件，而生化要素预报采用的是NASA提供的卫星同化数据，最终预报8种生化要素(total chlorophyll aconcentration, chlorophyte concentration, diatom concentration, coccolithophores concentration,cyanobacteria concentration, iron concentration, nitrate concentration and mixed layer depth)。三种下游任务的预报效果展示如图2所示。
图2. AI-GOMS下游任务预报效果（a图是Kuroshio地区的降尺度效果，图中以海面流速为例；b图是海浪高度的预报；c图是叶绿素浓度的预报）
2、总结
目前海洋预报领域的大模型应用正处在起步阶段，ViT的改进模型以及生成式的模型都可以与这个领域的问题进行结合来做一些工作，AI-GOMS则做了一个初步的尝试。在未来，海洋物理机制和深度学习的结合可以成为一个有意义的研究课题。此外，由于海洋上的观测数据获取成本较高，海洋的观测比较依赖卫星数据。随着我国风云卫星的发射，结合卫星、海洋预报模式及大模型或将是以后的一个重要的方向，并在实际应用中带来巨大的经济效益。
参考文献
Remi Lam, Alvaro Sanchez-Gonzalez, et al. GraphCast: Learning skillful medium-range global weather forecasting. https://arxiv.org/abs/2212.12794
Bi, K., Xie, L., Zhang, H. et al. Accurate medium-range global weather forecasting with 3D neural networks. Nature 619, 533–538 (2023). https://doi.org/10.1038/s41586-023-06185-3
Ravuri, S., Lenc, K., Willson, M. et al. Skilful precipitation nowcasting using deep generative models of radar. Nature 597, 672–677 (2021). https://doi.org/10.1038/s41586-021-03854-z
Zhang, Y., Long, M., Chen, K. et al. Skilful nowcasting of extreme precipitation with NowcastNet. Nature 619, 526–532 (2023). https://doi.org/10.1038/s41586-023-06184-4
Xiong, W., Xiang, Y., Huang, X. et al. AI-GOMS: Large AI-Driven Global Ocean Modeling System. https://arxiv.org/abs/2308.0315
往期回顾
MindSpore AI科学计算系列 | “没有最快，只有更快”，快速傅里叶变换详解
MindSpore AI科学计算系列 | 有效改善预报模糊问题，生成式模型在短临降水预报中大放异彩
MindSpore AI科学计算系列 | 基于深度学习模型来替代传统DFT模型以及DeephE3模型的分析综述
MindSpore AI科学计算系列 | 以MindSpore Elec为例的智能电磁计算若干进展综述
End
MindSpore官方资料
官方QQ群 : 871543426
官网：https://www.mindspore.cn/
Gitee : https://gitee.com/mindspore/mindspore
GitHub : https://github.com/mindspore-ai/mindspore
论坛：https://www.hiascend.com/forum/forum-0106101385921175002-1.html
MindSpore套件模型库：https://github.com/mindspore-lab
Openl启智社区：https://openi.org.cn
页:
[1]