清华+中国气象局Nature主刊NowcastNet全解读

导读：

1. 生成模型：简要介绍什么是生成对抗网络，以便部分读者阅读后文。

2. 尺度分离，对症下药：简述基本思想和创新点，为不需要了解细节的领导和导师们准备。

3. Evolution network：NowcastNet的两部分之一，深入技术细节。

4.Generative network:NowcastNet的两部分之二，深入技术细节。

5.如何训练：从损失函数入手，深入NowcastNet的训练过程。

6.结语：保留曲目，一点愚见。

短临降水预报一直是大气科学领域内深度学习方法发挥的首选舞台。继DeepMind DGMR后，清华龙明盛老师团队联合中央气象台，在Nature正刊发表了短临降水预报模型的重磅论文“NowcastNet”。我们细细拆解。

1

生成模型

先打个比方：你手握巨资，但对市场没有信心。此时明智的选择是分散投资，甚至把钱存在银行。

AI显然也是会这一套的：当它被用于短临预报，且对未来的对流尺度过程的预报没有信心，它便倾向于提供平滑甚至保守的降水预报。下图的PredRNN就是这样的“投资者”。

从DGMR开始，学者们提出了一类解决方案：利用“生成模型”预报降水。抛开什么是生成模型，先看一个有趣的例子：

引自https://github.com/dyelax/Adversarial_Video_Generation

上图左的“Ground Truth”是游戏“吃豆人”的画面。上图右为利用神经网络对游戏画面的预测。仔细看，发现有些角色走着走着就逐渐消失，有些角色甚至分裂为两个分道扬镳，这显然不符合游戏的真实场景——但这是不是很像我们前文提到的情况？因为对预测角色下一步动作没有信心，就倾向于保守（让它消失）的，或分散（让他分裂）的预测。

要解决这个问题，我们要做到两点：

1. 约束模型，不要产生看起来很假的图像；
2. 对于多种可能的随机事件，弱水三千只取一瓢。

好了，轮到GAN隆重登场：Generative Adversarial Network，生成对抗网络，即上图中间所用的技术。

引自https://www.researchgate.net/figure/Typical-Generative-Adversarial-Networks-GAN-architecture_fig2_349182009

图中，生成器（Generator）生成图像，而判别器（Discriminator）则既接收生成的图像，又接收真实图像，给出图片来自真实世界还是生成器的判断。在不断轮换迭代训练生成器和判别器的过程中，生成器的目标是学会骗过判别器，而判别器则不断加强识别能力。这是个道高一尺魔高一丈的角逐，也是“对抗”一词的来源。尽管过程是对抗性的，但目的却一致：生成看起来真实的图像，前文第一个问题解决了。天才的设计。

图的最左边还没有提到——我们从一个随机分布（比如正态分布）中抽样出一个随机向量z，作为generator的输入，这样，我们解决了前文的第二个问题——从万千可能性中确定一种，作为随机性的确定性来源。

从另一个角度理解，生成模型的目标不再是一个确定值，而是一种分布——比如猫的图片千奇百怪，但其数值规律符合某种复杂的分布，而生成器的目标就是将一个随机分布映射到这种分布上。

当然，如果你玩过流行的图片生成服务（比如Bing Image Creator），可能会想到一个问题：在生成图片时，我们还要输入提示词（prompt），来描述具体想生成什么样的图片，而上述的GAN结构似乎没有这样的输入。其实这里通常的做法很直接：将提示词作为条件（condition），同样输入模型。本文具体采用的形式为SPADE（Spatially-Adaptive Normalization），将在后文简单介绍。

2

尺度分离，对症下药

有了以上基础，可以仔细看看这篇论文了。首先是保留曲目——犀利地指出“友商”的问题：

1. 如前文所述，预报结果容易趋向模糊，缺乏对强对流的预报能力。
2. DGMR加强了第1点，但作为单纯的生成模型，缺乏物理机制的确定性约束，从而在强度、位置上可能有较大偏差。

有没有可能同时解决以上两个问题？上帝的归上帝,恺撒的归恺撒。用传统的确定性模型结合物理来解决中尺度系统的预报，在此基础上用生成模型以集合预报的方式模拟相对随机的对流尺度信号。在此精神指导下，模型结构确定了。原文的插图要素太多，这里重新制作了一个简化版：

1. 先通过物理机制约束下的Evolution network，从过去的雷达图预报其20km尺度的确定性变化。
2. 利用过去和预报的雷达图，通过GAN的encoder和decoder捕捉1~2km尺度的对流细节，生成最终的预报。其中，20km尺度的预报既通过encoder编码为decoder的输入，又作为“条件”在decoder中约束输出。

现在可以无障碍地本文最重要的公式了：

• 等式左边是条件概率，给定Evolution network Φ和generative network θ，输入过去时次的雷达降水，预报之后的雷达降水。
• 右边则将概率写为集合预报的积分形式：
  • 先通过Evolution network Φ求解中尺度预报；
  • 再给定随机抽样z和generative network θ，在中尺度预报结果Φ(x)约束下求解最终包含对流尺度的预报；
  • 该结果乘以P(z)，为该结果的出现概率，积分后得到多次抽样下的集合预报。

闲言碎语：深度学习论文里，可能为了彰显高贵的数学血统，往往包含一些复杂的公式，但本质往往只是为了给复杂的设计一个抽象的表示，而不是几页纸的推导，唯一的难点可能是小心翼翼不能把林林总总的逗号分号括号写错位置……

接下来就是具体的问题：

1. Evolution network如何融合了物理机制？
2. Generative network的具体结构如何？
3. 模型如何训练？

3

Evolution network

2023年男篮世界杯即将开赛，中国队的成色相比2019年如何？问题可以拆解为两个部分：

• 周琦等仍在队内的球员，四年里实力有何变化？
• 以李凯尔为代表的新球员会带来哪些改变？

我没有错按Ctrl+V，只是想以此类比如何分解雷达降水的变化——同样将其拆解为两个部分：

• 降水系统生、消导致的变化；
• 降水系统平流移动带来的变化。

对流发展，雨势增大，属于前者，NowcastNet用Intensity decoder模拟；对流云从海淀追到朝阳，属于后者，NowcastNet用Motion decoder追踪。

作者以连续方程的形式表示其中的物理机制：

左侧第一项为局地变化，第二项为平流，s在文中被称为“intensity residual”，这是更合适的说法，因为它本质上是考虑了近似轨迹后的强度残差，降水的生消只是s的一部分。

这部分的实际操作也许用作者开源的代码能够解释得更清楚：

# Evolution Network
intensity, motion = self.evo_net(input_frames)
motion_ = motion.reshape(batch, self.pred_length, 2, height, width)
intensity_ = intensity.reshape(batch, self.pred_length, 1, height, width)
series = []
last_frames = all_frames[
    :, (self.configs.input_length - 1):self.configs.input_length, :, :, 0
]

grid = self.grid.repeat(batch, 1, 1, 1)
for i in range(self.pred_length):
  last_frames = warp(
      last_frames,
        motion_[:, i],
        grid.cuda(),
        mode="nearest",
        padding_mode="border"
    )
  last_frames = last_frames + intensity_[:, i]
  series.append(last_frames)
evo_result = torch.cat(series, dim=1)

只看几个关键的步骤：

intensity, motion = self.evo_net(input_frames)

可以看到，将过去的雷达降水输入Evolution network，将预报出前文所述的生消场和移动场。

last_frames = warp(
  last_frames,
  motion_[:, i],
  grid.cuda(),
  mode="nearest",
  padding_mode="border"
)

该部分用移动场处理上一时次雷达降水场last frames。在warp函数中，有了移动场motion，自然可以回溯每个格点位置当前时次的信号是从哪里来的，即其在上一时次的对应位置，并在last_frames中找到对应位置的值更新为该格点当前时刻的值。当然，这里回溯的位置很难正好落到某个格点上，为了防止双线性插值导致结果趋于平滑，这里使用最近邻插值（为了求解梯度，在计算loss时换用双线性插值）。

last_frames = last_frames + intensity_[:, i]

在此基础上叠加生消场intensity，一帧的预报就完成了。

如何描述天气的变化？我们可以始终将目标拆分为一个一个“单位”，比如一个个气团，追踪他们的轨迹和物理性质变化，这种方案称为拉格朗日（Lagrange）方案；或者也可以只描述每个固定位置的变化，称为欧拉（Euler）方案。而这里既以固定格点为预报目标，又对每个格点追溯其单位时间内的轨迹，这是论文中多次提及的“semi-Lagrangian”的由来。

多说一句，6年前就有论文提出了类似的想法实现融合物理约束的深度学习，有兴趣的朋友可以查阅Deep Learning for Physical Processes: Incorporating Prior Scientific Knowledge（https://arxiv.org/abs/1711.07970v2）。

对于Evolution Network的结构，原文Extended Data Fig. 1a已经展示得非常清楚：two-path U-Net的结构，即图像先通过一个encoder，再分解为两个decoder，分别用于预报motion和intensity。这里不再赘述。

4

Generative network

先介绍NowcastNet的GAN结构SPADE（https://arxiv.org/abs/1903.07291）。

引自https://nvlabs.github.io/SPADE/

熟悉计算机视觉的朋友对左图应该不陌生：一般称它为语义分割，用不同色彩代表图片中的不同部分。修改左图的语义，将其作为“提示”，SPADE据此在右图做出对应的改动（更极端地，想象下，随手画几个轮廓，AI就帮你画出了精美的图像：Stable Doodle（https://clipdrop.co/stable-doodle）。

上图是SPADE的核心结构：将原图和随机量输入生成器，以语义分割图作为条件生成新图片（如果你读过NowcastNet，侧着头看上图会觉得有些似曾相识……）。

其中，生成器如何用语义约束输出结果？其基本模块则如下图所示：

归一化是神经网络常见的步骤，如图中的BatchNorm。作者认为，语义图所包含的信息在这类过程中极易被“洗”掉。同时，如果去查阅各类归一化的教程，发现它们都离不开一个统一的步骤：y=ax+b，即对中间层进行放缩和偏置。针对这一点，将语义图剥离开，通过两个卷积网络生成反映语义分布的a和b（注意这里a和b显然是和feature map同样尺寸的张量），通过element-wise的放缩和偏置，实现针对语义信息而与原图无关的特殊归一化处理。这就是语义图在生成器中发挥作用的过程。

理解了SPADE，NowcastNet中生成部分的原理也就很清晰了：20km分辨率的中尺度预报就是“语义”，约束了最后的1~2km分辨率预报结果。NowcastNet中的Nowcast encoder对应SPADE中的image encoder，而Nowcast decoder就是SPADE中的generator。

实际上，SPADE可以被应用于提升图片分辨率，类似本文的情景。唯一的区别是，NowcastNet涉及时序预测，而非单张图片的处理，其encoder的输入（即前文导图中的contextual representation）是从过去观测到未来预报雷达降水图。

5

如何训练

写到这里，其实还有最大的问题没有提及——这么复杂的模型，该怎么训练？从损失函数入手，我们一探究竟。

# Evolution network的损失函数

Evolution network的损失函数包含两个部分：J_accum和J_motion。

直接贴公式：

其中，xt为观测，x't为last_frames根据motion移动后的预报，x''t为在x't基础上叠加了intensity后的预报。分别将二者与观测计算L1距离并优化。此外，这里还仿照DGMR，根据降水强度对损失做了加权处理，具体请参见原文。

因此，NowcastNet并没有针对观测使用数值方法计算出观测的motion场作为motion decoder部分的标签，而是让Evolution network的预报，和只考虑了motion的中间产物均尽可能靠近观测。

此外，对于motion场本身，作者仍做了一些特别处理，在损失函数中增加了正则惩罚：

v1和v2分别是motion的经向和纬向分量，∇v表示对v进行卷积操作，卷积核采用固定的参数，不需要训练：

图中的数字即为卷积核，称为Sobel filter，是一种常用的卷积核，借用Andrew Ng的课件简要介绍下它的作用：

引自Andrew Ng 《Deep Learning》课件

上图左是一个二维数组，左半部分的数值均为10，右半部分的数值均为0，如果转化为一张图片，显然左边是一种颜色，右边是另一种，中间是两种颜色的交界线。如果我们用图中间类似Sobel filter的卷积核对其进行卷积操作，将得到图右的特征图。特征图中为30，即信号较强的地方在正中央，正对应原图中两色交界线的位置。这也是卷积操作被认为能够检测图片边缘的原因。

再回头看NowcastNet中的J_motion：这一步针对motion应用Sobel filter，检测motion的边缘，提取的其实是类似水平风切变的物理量，并作为损失函数，使其尽量小。重要的是，这里的w(x)为降水强度决定的权重，实际上是强迫模型对极端降水预报更小的水平风切变，以使其在预报中具有更大更稳定的时空影响范围。

讲完损失函数，最后是处理误差在多个时间步上累积的问题。之前介绍过Pangu的解决方案是Hierarchical temporal aggregation，例如对于24小时后的天气，直接由跨越24小时端到端的模型来预报。NowcastNet融入了半拉格朗日方案，每个时次的预报需要通过motion追踪上一个时次的预报，显然只能做连续的自回归预报，无法像Pangu一样处理。

作者的方案是切断每个预报时次间的梯度传递，即正向传播时前一时次的预报影响后一时次的预报，但反向传播时后一时次的梯度不影响前一时次。

# Generative network的损失函数

这部分的损失函数容易说明白：

1. 生成器的目标是生成“逼真”的预报，而判别器的目标是正确区分预报和观测。与SPADE类似，对应的过去时次的雷达观测也作为判别器的输入，因此作者将判别器称为“Temporal discriminator”。二者均使用交叉熵（cross-entropy）即可。
2. 对于生成器，为了进一步增强极端降水的预报能力，生成器的损失函数还包含另一个部分。该部分将预报和观测分别做最大池化（max-pooling）的处理，并计算L1损失。注意这里的预报是指集合平均，这样做想必也能抵消一些集合平均导致极端信号被平滑掉的问题。

6

结语

微波炉的发明：1945 年，雷达工程师珀西-斯宾塞在磁控管（一种产生高频无线电波的电子真空管）前停留了一分钟。他突然感到一种奇怪的感觉，发现口袋里的巧克力棒正在融化。斯宾塞很好奇，他把一个生鸡蛋放在磁控管前的锅里。爆炸的鸡蛋溅到了附近的同事身上，证实了微波可以快速烹饪食物。

引自https://invention.si.edu/node/1145/p/431-percy-spencer-microwave-inventor

Swin Transformer和SPADE的作者想必也没有想到，他们的工作助力Pangu和NowcastNet这样的大气科学领域工作在顶级期刊大放异彩。从基于连续方程设计的融入物理的模型结构，到借用适当的生成网络结构，将科学领域的已有知识和任务场景与计算机视觉、自然语言处理等领域的最新进展相联系，可能给了单学科背景的朋友一个挑战，但也是交叉学科和交叉人才，或不同领域展开合作的重要机遇。

我猜想，次季节-季节尺度的预报在类似的思路下有相当的潜力。最后，从模型验证到实际应用，从分析数据、雷达数据到实际观测，其间差距仍不可以道里计，欢迎一线预报员谈一谈体会。

NowcastNet的主要创新在两个方面。第一，将预报分解为物理约束下的确定性中尺度预报和生成模型驱动的对流尺度预报。第二，将确定性预报分解为对流和扩散的预报。这某种程度上是深度学习模型“端到端”策略的妥协，拟合一切的神经网络仍需要合理的设计来减轻模型的负担。尽管本文的设计甚至可以用物理方程来背书，但由于最终的评判仍依赖试验结果的评估而不是过程解释，因此这种设计更像是物理先点亮直觉，直觉驱动设计，颇有些Attention机制的味道。

本文当然还有很多论文细节没有提及，补充材料中作者还做了大量工作（建议阅读），比如针对各项设计的消融实验。可以说，NowcastNet的复杂性和含金量极大，它更像一个“系统”而不是“模型”。清华团队和中国气象局的合作也已有数年。不得不感叹这样的工作非常不容易。

不过在DGMR发表时，大概很多人没想到仅仅2年后就又有一篇同一主题且大量参考了DGMR的论文再次登上Nature正刊。道阻且长，行则将至。读者朋友们，期待你们的Nature。