导读:
1. Pangu做了什么:通俗地概述Pangu-Weather,标准是新闻专业的也可以看懂,老少咸宜。
2. Pangu是如何做的:Pangu的技术路线,只需要一点点深度学习知识,标准是看完产生“我上我也行”的错觉。
3. 从Peer Review再看Pangu:总结40页的Nature同行评审结果,重点都浓缩在这里。
4. 浅谈“气象大模型”:发表一些关于AI在大气科学领域应用前景的愚见。
1
Pangu做了什么
大气变化可以用方程组来描述。通过观测得到大气当前的状态,将其输入方程得到下一时刻的状态,从而实现天气预报。
因此,天气预报应该包括至少3个要素:
整套流程也叫“数值天气预报”(NWP)。而Pangu所做的则是抛开第2、3步括号内的传统方法,用纯数据驱动的AI模型来完成任务。其效果超越了可能是地球上最好的NWP,且速度提升10000倍。当然,第1步的苦差事仍需要NWP来完成。
2
Pangu是如何做到的
让我们用术语再把故事讲一遍:Pangu在逐小时的ERA5再分析数据集上训练。ERA5提供的初始大气状态可以描述为一个数组,而下一时刻的状态也是一个同样结构的数组,于是上文的第2、3步在深度学习的视角中就变成了一个image-to-image的监督学习任务。
Pangu的结构如下图所示:
可以认为,Pangu完成的更多是工程领域的工作,其关键点可以概括为以下几个方面:
Transformer是当前深度学习领域的”版本答案”,最初被应用于自然语言处理(NLP)领域,其核心是Self-Attention机制。所谓Self-Attention,即对于一个模型输入,计算其各部分两两间的相关性,以理解其结构。以下图为例:
左图句子为“The animal didn't cross the street because it was too tired.”,"it" 指代“animal",因此与”animal"相关性最强,而右图为“The animal didn't cross the street because it was too wide.”,“it"指代”street",因此与“street”相关性最强。
Transformer可以被应用到图像处理领域,然而图像并非是天然的序列形式,最直接的应对方法如ViT (Vision Transformer),将图像切成一个个方块(patch),每个patch可类比于句子中的每个词(实际上应该叫token),再排列为一个序列,应用到Transformer上:
其中,Position Embedding用来区分patch在图像中的位置。
然而,假设图像的长、宽扩大2倍,patch的数量将扩大4倍,再两两计算各个patch间的attention,计算量将扩大16倍。为了解决这个问题,Swin Transformer应运而生:
其核心思想是,将patches分割为不同的窗口,如图左下角,每个最小的方格是一个patch,每个红框窗口包含4x4个patch,分别计算每个窗口内不同patch间的attention,如此,只要保持窗口大小不变,计算量将不再随图像尺寸指数级增长。
当然,这里引入了另一个问题:无法学习窗口间的相互关系。Swin Transformer采用Shifted Windowing和Patch Merging来解决这个问题。前者将在后文简单说明,后者将另文阐述。有了这两项技术,某种程度上,Swin可以理解为披着Transformer外衣的卷积神经网络(CNN)。
可能是受CNN的平移不变性启发(一只猫无论出现在图像的左上角还是右下角,都具有类似的特征),Swin Transformer中的positional bias(即表示patch位置的参数,为何用bias一词,可以在Swin的原论文[1]中找到答案)表示patch在其所属窗口中的相对位置,而忽略了其在整张图中的绝对位置。这对于大气模拟而言并不合适,因为不同纬度的格点间距不同,相同的大气变量分布出现在大气层的不同位置,其物理意义也必然有所区别。因此,Pangu将positional bias设计为由两个patch在大气层的绝对位置决定的参数,以进行区分(实际上只考虑了纬度和高度的不同)。
图像通常由长、宽、通道(如RGB)组成3维数组,而气象场为4维(经度、纬度、高度、变量),因此,在切分patch时,除了长、宽外,Pangu还在高度(等压面)上进行切分,从而可以计算不同等压面间的attention,这也是所谓3D的由来。
此外,Swin Transformer为了增强窗口间的互动,会进行shifted windowing操作,简单的示意图如下:
可以直观看出,通过将窗口的位置统一平移,原先不在一个窗口中的相邻patch将被组合到一个新的窗口中。为了在工程上保持每个窗口大小不变,Swin通过平移,将所有patch仍“凑”为4x4的窗口:
上图已经非常清晰地描述了平移的过程,一个细节是,这样的平移会将原来南辕北辙的patch组合到一起,而他们互相间通常没有很强的关系,因此采用掩码的方式避免它们的互动。与此不同,地球在纬向(东西方向)是个环,即359.75°与0°是相邻的。Pangu据此特性对shifted windowing作了相应修改。
以上就是作者说明的所有“Earth-specific”之处。基本没有吸纳任何物理性质,更像是“Sphere-specific”,本质是“让我们把不确定的东西做成可学习的参数”。
Pangu的另一项特性被作者称为“Hierarchical temporal aggregation”,其实质类似贪心算法:作者训练了4个模型,分别预报1h、3h、6h、24h后的天气。例如预报28h后的天气,将先利用24h模型预报得到24h后的预报场,再以此为初始场,相继利用3h、1h的模型得到28h的天气。显然,这么做是因为误差将在短时模型中快速累积放大。其思想甚至有些类似ResNet。
3
从Peer Review再看Pangu
Nature的Peer Review和作者的回复已在网上公布[2]。可以说初审意见就相当正面,认可了Pangu在再分析资料上超越数值模式的里程碑意义,也没有要求伤筋动骨的改动,而是希望更准确地表述Pangu的局限性,或避免过度表述没有说服力的优势。其内容基本可以归结为以下几点:
由评审意见可以引申出基于AI的天气模型的不足和未来方向:
4
浅谈“气象大模型”
Pangu论文ArXiv版本的摘要里有一句话:
“
Pangu-Weather ends the debate on whether AI-based methods can surpass conventional NWP methods.
”
这句话并不是无的放矢。纯数据驱动的模型是否能在预报性能上超越NWP,一直是业界关心的问题。数年前比较有代表性的可能包括Peter D. Dueben[4]和Jonathan Weyn[5]的系列工作。当时的模型无论是分辨率还是预报技巧都远远无法和数值模式相比,甚至不确定未来深度学习是否可以达到最先进的NWP的水平。
风云突变,由于算力的暴涨和新模型结构的成熟,IFS已经在最近一年被声称击败了很多次(在一些晚于Pangu的文章中,Pangu自然也无法幸免地成为背景板)。
挂一漏万地列举下这些大模型:盘古(华为),FourCastNet(Nvidia),GraphCast(Google),风乌(上海人工智能实验室),伏羲(复旦大学)……
变化的另一个可能原因是计算机领域的人才大批量进军所谓AI4Science,而他们所掌握的大规模深度学习的工程技术带来的质变往往被大气科学领域的人忽视,Pangu的卓越工作就是最好的证明。
当然,目前离声称NWP已经被AI取代还为时过早(实际上我认为NWP仍然是未来),其理由前文已经涉及。这里可以再从Pangu的两大亮点——效率和准确性出发,浅尝辄止地探讨几个问题。
首先是效率。一方面是预报效率。尽管大部分深度学习天气模型的初衷可能并不是解决NWP运行缓慢的问题,但可能不得不承认这些模型最大也最明确的价值来自于指数级的效率提升。但这种提升是否能够跨越其它的风险,使AI模型成为天气预报的主力军,需要时间的检验。
另一方面是训练效率。Pangu用192块V100训练了16天,且最后结果仍未收敛。大模型的训练代价昂贵,显然会带来两大问题:
至于准确性,以Pangu为代表的模型证明了深度学习模型可以在特定条件下超越NWP,其里程碑价值毋庸置疑,但里程碑的后面究竟还有没有路尚未可知。它有可能是:
也可能是:
我们的目标是发明更快的交通工具,而不是仅仅证明骡子也可以快过马。前者是未来,是可持续的新范式,后者则可能是昙花一现的故事。
尽管大气科学领域从可解释性角度反对AI的声浪似乎逐渐消弭,但输入、输出固定,且不包含任何物理约束的AI模型尚难以用于敏感性试验、因果分析等科学研究。因此区别两条路的核心,可能是如何让AI像NWP一样与大气科学相辅相成,以当下的风潮为契机,找到真正由数据驱动的学科可持续发展道路。在此如履薄冰地列举一些可能的方向,欢迎批评指正:
总而言之,纯数据驱动方法的“降维打击”并没有终结一切——“后GPT”时代,一众NLP研究陷入窘境,而与此相反,有扎实理论基础的大气科学领域,AI的应用研究可能仍刚刚开始。
引用链接:
[1] https://arxiv.org/pdf/2103.14030.pdf
[2] https://static-content.springer.com/esm/art%3A10.1038%2Fs41586-023-06185-3/MediaObjects/41586_2023_6185_MOESM1_ESM.pdf
[3] https://www.ecmwf.int/en/about/media-centre/science-blog/2023/rise-machine-learning-weather-forecasting
[4] https://www.ecmwf.int/en/about/who-we-are/staff-profiles/peter-dueben
[5] https://atmos.washington.edu/~jweyn/
[6] https://www.nature.com/articles/s41597-023-02059-5
[7] https://arxiv.org/abs/2301.10022
[8] https://agupubs.onlinelibrary.wiley.com/doi/full/10.1029/2022MS003034
关于“梯度大气”
AI技术在大气科学领域的应用如火如荼,相关的优秀公众号已不在少数,但似乎还缺乏一些主观解读。不主观就没有错误,没错误就无法进步。“梯度大气”立志更慢的更新频率(但稳定),更多的“原创错误”(但有见解),以定期的长文为形式抛砖引玉,向同行学习。
欢迎关注,期待批评!