背景
本文主要分为下面三个大章节:
Diffusion经典工作介绍
GraphCast模型
GenCast模型
1、Diffusion经典工作介绍
Diffusion model也就是所谓的扩散模型,是当下非常火热的生成式模型。本节的内容并不是跟踪diffusion系列最新的进展,而是对diffusion最经典的工作进行梳理。Diffusion模型需要的数学理论涉及统计、随机微分方差等,哪怕是数学专业的研究生,如果不是相关领域的可能也不会了解这方面的知识,所以还是有一定的门槛。下面的理论基础会涉及到部分数学推导和公式,难免会有些枯燥,但尽管如此本文也已经省略了很多细节的推导。如果对理论背景不感兴趣的同学可以跳过理论基础这一小节。
1.1 理论基础
这类模型核心在“扩散”两个字。diffusion模型整体上分为两个过程,正向扩散过程和逆向扩散过程。以图像为例,我们给清晰的图像不断地加入噪声,经过有限步长之后,把这张清晰的图像变成一个噪声的分布(一般是高斯分布),这个过程就是正向扩散的过程。那么可以猜到,逆向扩散过程则是以采样一个随机噪声为起点,不断地“去噪”,经过有限步之后生成一张清晰的图像。虽然当前diffusion有大量的变种,但基本都不会脱离这两个过程。初看这个思路的时候会产生一个问题,这个逆向过程凭什么可以做到这一点呢?为了解释这个问题,我们需要对diffusion背后的数学原理有一些了解。因此,下面用形式化的数学语言来表示diffusion的建模思路。
1.2 DDPM
在上一小节,我们概括性地介绍了diffusion模型的数学建模背景。这个工作[1]在2015年就已经发表,但这篇文章当时是在MNIST、cifar-10这一类的数据集上进行实验,并没有引出很大的轰动。直到2020年DDPM(Denoising Diffusion Probabilistic Models)生成了较高质量的图像,才让更多人把关注点放在了diffusion这个模型上面(从侧面反映了有的时候严格的数学推导其实并没有很多的人愿意去关注,也并不代表一定能产生好的效果)。但一个坚实的理论基础是技术发展的根本,了解背后的原理能让我们在繁杂的前沿工作中看到方向。事实上,DDPM在理论方面并没有更多的推进,而是在schedule的选取、网络架构及参数化上做出了改进。下面我们看看DDPM在此基础上到底做了什么。
1.3 Score-based生成模型
斯坦福大学的宋飏教授在21年发表在ICLR上的文章[3]从理论上统一了他前期的工作Score matching with Langevin dynamics (SMLD)[4]与上一节提到的DDPM。文章[3]以Stochastic Differential Equations(SDE)为工具,从更高的视角给出了SMLD与DDPM的理论框架,也为后续diffusion系列的建模奠定了一个坚实的理论基础。这篇工作也是当年ICLR的oral,是我认为难得的好paper之一。文章[3]所使用的SDE需要的数学基础比起前两节有过之而无不及,严谨的描述会使本文过于冗长,一些细节我也并没有独立去验证。为增加本文的可读性,我尽量用一些大白话去描述这篇文章的主要思想,对理论有追求的同学强烈建议去阅读原文[3]。值得一提的是,这篇文章和一些别的AI paper不同,并没有在模型架构的设计上有创新。而把更多的篇幅用来阐述理论基础,并在算法上提出了PC sampler,提升生成图片的质量。下面我们开始聊聊这篇文章的核心思路。
对于sampling的过程,除了上述对逆向SDE进行sampling之外,作者还提出了一个probability flow ODE的概念。对于初始SDE(1.17),我们也同样可以写出对应的逆向形式,可以参考文献[3]的Eq17. 对应的递归公式与上面Algorithm 2/3相似,在predictor中
的递归公式中,网络部分前面增加1/2做为系数。
2、GraphCast
GraphCast[6]是DeepMind团队在2023年登上Science正刊的中期预报大模型。按个人的观点,2022年可以算做AI气象元年,如ForeCastNet、Pangu(是盘古气象大模型,而非LLM)等都是在2022年挂出来的,并引起了极大的反响。虽然用神经网络做天气预报的场景从很早之前就有,譬如2015年施行健团队提出的ConvLSTM,但Pangu让气象领域的人真正意识到大模型有希望在实际预报业务中使用。
ForeCastNet、Pangu等模型主要的框架都是transformer架构,而GraphCast的主体架构是更加传统的encoder-processor-decoder架构,预报未来10天的天气,时间分辨率为6小时。值得一提的是模型的建模思路是用Graph的方式对地球建模(但实际的网络中并没有采用GNN的模块),具体如下图所示。
3.2 模型结构
关于GenCast的模型结构我在这里放一下原文推理过程的模型结构图,这是一个典型的diffusion去噪的过程图。GenCast的思路是利用过去2个时间步的输入来预测下一个时刻的预报结果,然后再把下一个时刻的预报结果做为输入,滚动预报,直至预报未来15天的天气。单个时刻的预报前向20次,也就是下图中N=20,而时间分辨率是6小时,故T=15*4=60。
3.3 实验结果
论文在不同的维度对比了GenCast和基线(传统的集合预报ENS)的效果。在温压湿风等44个常规要素的预报上,对比了RMSE的指标。GenCast相对于ENS在未来15天的预报中,RMSE整体优于ENS。其中预测的前2\~3个时效,RMSE指标优于ENS 10%-20%左右,在高空的要素上劣于ENS。而后几个时效的预测中RMSE指标略优于ENS0\~5%左右。
在集合预报(GenCast的集合预报主要通过不同的随机初始噪声为起点,得到不同的Sampling的方式进行集合预报)的可靠性上,实验主要对比2米温度、10米风速u、850hPa温度、500hPa地势高度、700hPa,通过Rank histograms进行对比。GenCast相对于ENS在未来15天中不同预报时刻的曲线都相对比较平坦。
在极端情况的对比中,实验主要对比极端高温、低温、强风速、海平面低气压,通过Brier skill score进行对比。在以上四种极端天气要素中,GenCast在未来15天的预报结果,Brier评分都优于ENS。
此外,GenCast还预测了一个下游任务,也就是台风路径预测。论文以2019年登陆日本的Hagibis台风做为case进行预测,预测路径如下图所示。可以看出GenCast预测的台风路径随着预报时间的临近越来越“收敛”于EAR5的结果。
参考文献
[1] Deep unsupervised learning using nonequilibrium thermodynamics, Sohl-Dickstein, Jascha and Weiss, Eric A. and
Maheswaranathan, Niru and Ganguli, Surya. JMLR, 2015.
[2] Denoising Diffusion Probabilistic Models, Jonathan Ho, Ajay Jain and Pieter Abbeel. NIPS, 2020.
[3] Score-Based Generative Modeling through Stochastic Differential Equations, Yang Song, Jascha Sohl-Dickstein and
Diederik P. Kingma) et al. https://arxiv.org/abs/2011.13456
[4] Generative modeling by estimating gradients of the data distribution, Yang Song and Stefano Ermon. In Advances in
Neural Information Processing Systems, pp. 11895–11907, 2019.
[5] Elucidating the Design Space of Diffusion-Based Generative Models, Tero Karras, Miika Aittala, Timo Aila and Samuli
Laine. NIPS, 2022.
[6] Learning skillful medium-range global weather forecasting. Remi Lam et al. Science, 2023. DOI:10.1126/
science.adi2336