加利福尼亚洛杉矶大学和卡内基梅隆大学的研究人员推出Stormer:一种可伸缩的Transformer神经网络,用于熟练可靠的中程天气预报

「加利福尼亞洛杉磯大學和卡內基梅隆大學研究人員研發出Stormer:具彈性的Transformer神經網絡,用於準確可靠的中程天氣預報」

当今科学和社会面临的主要问题之一是天气预报。准确的天气预报在帮助人们规划和应对自然灾害和极端天气事件以及帮助研究人员更好地理解环境在日益担忧气候变化的背景下起着至关重要的作用。数值天气预报(NWP)模型一直是大气科学家工作的支柱。这些模型使用解释热力学和流体流动的微分方程组,并可以在时间上进行综合,以产生未来的预测。尽管NWP模型广泛使用,但它们有一些缺点,比如参数化错误,不能准确描述重要的小尺度物理现象,如辐射和云物理。

由于整合一个大型微分方程组的困难,数值方法的计算成本也很高,特别是在精确的空间和时间分辨率下进行建模。此外,由于模型依赖于气候科学家的知识来改进方程、参数化和算法,NWP预报的准确性在增加数据后并没有提高。越来越多的人对基于数据驱动的深度学习天气预报方法感兴趣,以克服NWP模型的问题。通过使用历史数据,如ERA5再分析数据集,可以训练深度神经网络来预测未来的天气条件。这是该技术的主要前提。与传统的NWP模型不同,后者需要数小时来进行预测,而深度神经网络一旦训练完成,可以在几秒钟内做出预测。

在这个领域的早期研究工作中,使用传统的ResNet和UNet等视觉架构进行天气预报,因为气象数据和自然图片具有相似的空间结构。然而,它们的表现不及数值模型。然而,由于改进的模型设计、训练方法和增加的数据和计算能力,最近取得了显著的进展。首个超越运营IFS模型的模型是Pangu-Weather,这是一个基于3D地球特定Transformer模型,采用0.25∘数据(721×1440格点)进行训练。紧接着,Keisler的图形神经网络设计通过GraphCast扩展到0.25∘数据,表现出比Pangu-Weather更好的性能。

尽管预测的准确性出色,但当前的方法有时会使用复杂、高度定制的神经网络拓扑结构,几乎没有进行消融实验,这使得很难确定导致它们有效的确切因素。例如,我们不知道在GraphCast中多网格消息传递对效率的贡献有多大,以及3D地球特定Transformer相比普通Transformer的优势。在这个领域的进一步发展需要更好地理解当前的方法,并最好进行简化。一个统一的框架也将更容易创建超越天气预报的气候和天气基础模型。该研究表明,在与合适的训练方法相结合时,简单的设计可以胜过尖端技术。

来自加州大学洛杉矶分校、卡内基梅隆大学、阿贡国家实验室和宾夕法尼亚州立大学的研究人员提出了Stormer,这是一个简单的Transformer模型,在传统的Transformer骨干上几乎不需要修改即可实现天气预报中的最新性能。研究小组从传统的视觉Transformer(ViT)架构开始,进行了深入的消融研究,确定了影响模型性能的三个关键因素:模型由三个组成部分构成:(1)一个天气特定的嵌入层,用于对大气变量之间的相互作用进行建模,将输入数据转换为一系列令牌;(2)一个具有随机动力学预测目标的模型,训练模型在随机时间间隔内预测天气动力学;以及(3)一种压力加权损失,通过在损失函数中对不同压力级别的变量进行加权来近似每个压力级别的密度。他们提出的随机动力学预测目标,通过使用模型训练的不同间隔的组合,使单个模型能够在推理过程中为给定的提前时间生成多个预测。

例如,通过分布6小时的预测12次或12小时的预测6次,可以得到一个3天的预报。这些预测的结合显著提高了性能,特别是对于较长的提前时间。研究小组使用WeatherBench 2进行了Scalable transformers for weather forecasting(Stormer)的验证,这是一种常用的基于数据驱动的天气预报基准。测试结果显示,Stormer在经过7天后超越了最先进的预测系统,在1-7天的时间内实现了重要大气变量的有竞争力的预测准确性。值得注意的是,Stormer在几乎5倍较低分辨率的数据和数量级更少的GPU小时训练下,超越了性能基准。最后,他们的扩展研究表明,通过证明随着模型容量和数据规模的增加,Stormer的性能不断提高,还有额外的改进空间。