基于双水平路由注意力机制的时序Transformer临近预报研究

A Study of Temporal Transformer Nowcast Based on Bilevel Routing Attention Mechanism

摘要: 0~2小时的临近预报是预报强对流天气的重要手段，准确的临近预报是气象防灾的重要屏障。临近预报是气象预报的一个重要业务。本文提出了一个基于时序Transformer的全局特征提取的Encoder-Forecasting临近预报模型结构。该模型的Transformer利用双级路由注意力机制定位特征中的几个最相关的键值对来提高计算效率。模型同时采用MaxPool和AvgPool提取目标对象的局部特征，并将局部特征和全局特征进行有效融合，进而充分提取目标对象的特征用于临近预报。提出的模型与几个模型在一个公共的临近预报数据集上进行实验验证。实验结果充分证明了我们提出的模型的有效性和准确性。

Abstract: The nowcast of 0~2 hour is an important toolth for forecasting severe convective weather. Accurate nowcast is an important barrier to prevent from weather disaster. Therefore, nowcast is a very important field in the meteorological society. This paper proposes a nowcast model which based on Encoder-Forecasting architecture and temporal Transformer. The transformer utilizes bilevel routing attention to improve the computational efficiency. At the same time, the proposed model also use the operations of Maxpool and AvgPool to extract the local features. The global features and local features are effectively fused to represent the objects’ features. The comprehensive experiments based on the proposed model and four state-of-art models are conducted on a public datasets. The experimental results effectively show the correctness and effectiveness of the proposed model.

1. 引言

强对流天气由于破坏性强，给人们的生命财产安全造成比较大的损失。同时由于它形成的突发性强，持续发生时间比较短，并且影响范围相对比较小等特点，因此对强对流天气的预报变得比较困难[1]。强对流天气是中小尺度天气系统的一部分。准确的预报强对流天气的发生又是气象防灾减灾及气象预警工作的重要步骤。气象雷达主要用于大气探测，它是探测中小尺度天气的重要工具，雷达回波能够比较准确的反映实际天气状况。因此准确的外推雷达回波是预报强对流天气的重要手段。0~2小时的雷达回波外推得到气象雷达探测目标未来2小时的气象信息，从而能够比较准确的预测未来天气的情况。

2. 相关工作

2.1. 传统外推方法

传统的临近预报技术主要包括雷达回波外推[2]和数值预报[3]。然而在0~2小时的临近预报应用中，数值预报由于模拟复杂的天气系统需要较多的计算资源及太多的假设初始条件而得到的结果次于雷达回波外推得到的结果[4]。因此对临近预报的研究主要集中于雷达回波外推。传统的雷达回波外推技术主要采用交叉相关法[5]和光流法[6]。但由于雷达回波运动的轨迹受多种气象要素的约束变化无常，很难用数学公式进行描述表示。传统的光流法和交叉相关法主要利用相邻的几张雷达回波进行回波轨迹的预测，由于使用的回波数量少，很难实现对回波运动轨迹规律的描述，同时随着外推时间的变长，对雷达回波的生消演变结果的准确度急剧降低，因此外推效果需要进行极大的改进以满足人们对气象预报准确性的要求。

2.2. 基于深度学习的临近预报

近年来，随着图形处理器(Graphic Process Unit, GPU)计算力的快速增长及各种大数据的出现，深度学习[7]的兴起及其在许多领域得到了广泛的应用并取得了重要的成功。Shi等人基于传统的长短期记忆网络(LSTM)并结合雷达回波的空间特征属性提出了卷积长短期记忆网络(Convolutional LSTM, ConvLSTM) [8]用于雷达回波外推并取得了比传统光流法更好的外推效果。针对卷积操作的局部不变性特点和雷达回波运动是局部变化的两种状况，Shi等人又改进提出了TrajGRU [9]用于主动学习回波序列中的局部变化结构用于临近预报。Wang等人提出的PredRNN (predictive recurrent neural network) [10]模型使用一个统一的存储池学习预测目标的空间外形和时序变化特征。Trebing等人[11]基于UNet框架使用深度可分离卷积减少模型的参数，同时在编码器中加入CBAM (Convolutional Block Attention Module)注意力机制，SmartAt-UNet在简化网络模型参数、提高模型推理的情况下仍然得到了较好的预测效果。Wang等人在论文[12]中通过时空记忆流的锯齿状循环转换机制和一对解耦记忆单元的新卷积循环单元改进了PredRNN。DeepMind团队[13]基于生成对抗网络开发了一套降水概率预报模型。

以上这些预报模型都是基于卷积操作进行临近预报，然而受限于卷积操作固有的平均属性，这些模型在预测高强度回波值时的准确度就比较低。同时卷积操作的效果受其卷积核的限制，它只能提取目标对象的局部特征。然而随着Transformer在自然语言[14]、图像处理[15]等多种应用中取得了令人瞩目的成功，与传统的卷积操作相对比，Transformer在提取目标对象的全局及长时序特征方面表现更加出色。基于Transformer的深度学习模型也逐渐应用于临近预报中。论文[16]提出了基于窗口的多头自注意力机制的Rainformer用于提升高强度降水预报的准确性。Li等人[17]提出了一种采用LPT-QPN模型用于临近降水预报。LPT-QPN模型采用新颖的Transformer模块来学习提取降水的长期演变规律特征，它同时采用多头平方注意力(multihead squared attention, MHSA)构建降水的非线性关系，使用多头平方注意力同时能够减少LPT-QPN模型的计算复杂度。实验验证表明LPT-QPN模型能够提高对某些特定降水时间的预测精度。为了减少临近预报外推中的累计误差，一个时序空间并行的Transformer [18]被提出用于临近预报并取得了很好的效果。

虽然深度学习在临近预报应用中取得了重要的成功，但其在临近预报中的研究也处于发展阶段。目前研究主要是利用大量的回波数据让深度学习模型进行自动学习其运行规律，同时深度学习模型也在快速的发展中，本文根据目前发展的趋势及深度学习最新技术，提出并改进了一种基于Transformer的深度学习模型并将其用于临近预报研究，努力尝试提高模型在预报高强度回波的预报精度。

3. 研究方法

Figure 1. Network architecture of the proposed nowcast model

图1. 提出的临近预报模型网络结构图

本文中，我们充分考虑雷达回波序列的时序特征和空间特征及Vision Transformer的全局特征提取属性，采用Encoder-Forecaster架构用于临近预报。整个模型架构如图1所示。Encoder和Forecaster都是由四个模块组成，Encoder中的每个模块最开始执行PatchMerging操作将目标对象的尺寸降低1/2，Forecaster中的四个模块开始则执行PatchExpanding操作将目标对象的尺寸放大2倍，这样最终恢复目标到原始的尺寸大小。其中Encoder和Forecaster模型的组成构件都是一样的，但它们处理回波数据的方式不同。Encoder中处理输入数据包含了回波序列的时序性和空间特性，因此我们将其称为时序空间Encoder，提取回波空间时序特征的模块称为空间时序特征提取模块(Spatialtemporal Feature Extract Module, STFEM)。Forecaster处理回波的空间特性，我们将其称其为空间Forecaster。Forecaster的4个模块中都分别包含Transformer组件和卷积组件用于提取全局特征与局部特征并将其进行融合，我们称这4个模块为全局与局部特征融合模块(Global and Local Feature Fusion Module)。我们称提出的模型为RainbiFormer。我们在随后的两个小节部分对时序空间Encoder和空间Forecaster进行详细的介绍。

3.1. 时序空间Encoder

在提出模型的时序空间Encoder部分，我们引入双级路由注意力机制的Transformer模块对输入的每个回波进行全局特征提取，这样不但学习到回波的空间特征、同时也能够提取到整个回波序列的时序特征。双级路由注意力的结构如图2中的Global Module所示。双级路由注意力主要是为了降低Transformer中multihead self-attention (MHSA)计算的复杂度[19]。双级路由注意力首先将输入的2D特征$X\in R^{H\times W\times C}$ 划分为$S\times S$ 的非重叠区域，这些区域包含$\frac{HW}{S^2}$ 个特征向量。通过这种划分方式得到的query，key，value向量为$Q，K，V\in R^{S^2\times \frac{HW}{S^2}\times C}$ 。这些向量执行的线性投影操作如公式1所示。

$Q=X^r{W}^q,K=X^k{W}^k,V=X^v{W}^v$ (1)

其中$W^q,W^k$ 和$W^v\in R^{C\times C}$ 分别为query，key，value的投影权值矩阵。然后通过有向图在区域之间进行交互。其中区域级的queries和keys为$Q^r,K^r\in R^{S^2\times C}$ ，区域之间的仿射图通过$Q^r$ 和$K^r$ 的转置矩阵乘得到邻接矩阵$A^r\in R^{S^2\times S^2}$ ，操作定义如公式2。

$A^r=Q^r{\left(Q^r\right)}^T$ (2)

然后仅保留每个区域的仿射图中前k个连接，从而使用行相关操作得到路由索引矩阵$I^r\in N^{S^2\times k}$ 。根据现代GPU合并操作原理，将key和value向量首先聚合起来，聚合的key和value向量表示为公式3所示。

$K^g=gather\left(K,I^r\right), V^g=gather\left(V,I^r\right)$ (3)

聚合的key和value向量在聚集的key和value值对上使用Attention操作，同时使用了一个局部场景增强项。整个双级路由注意力操作如公式4所定义

$O=Attention\left(Q,K^g,V^g\right)+LCE$ (4)

局部特征提取单元则是将整个雷达回波序列作为一个时空数据整体进行处理和学习，它的结构如图2所示。局部特征提取单元首先得到特征图的最大值和平均值，然后将其对应乘积到特征图，得到转换后的两个特征图通过卷积核为3的卷积操作进行降维融合操作。基于全局特征提取的特征与基于局部特征提取单元的特征再通过特征融合单位进行信息的融合处理，这样Encoder模块能够高效的提取雷达回波序列的时空特征，学习得到雷达回波的运动轨迹。

Figure 2. Spatial timing Encoder module structure framework diagram

图2. 空间时序Encoder模块结构框架图

在本文中采用Batch为12，输入6个回波，预测未来12个回波的方式。也就是根据当前的半小时回波情况预测未来1小时的回波运行状态。时序空间Encoder中的四个模块中的全局特征提取模块都是采用独立的处理方式对每个回波进行单独提取。Local Module采用Max Pooling与Avg Pooling两种操作方式分别提取回波的局部特征，然后再将得到特征图进行concat操作。而Local Model对回波进行整体操作，不再使用时序方式进行。Global Module对6个时序回波进行特征提取后再在channel上进行concat操作(Batch channel Height, width)的方式，这样与Local Module得到的结果形式一样，再进行融合操作。

3.2. 空间Forecaster

空间Forecaster同样是由基于全局特征提取的空间Transformer组件和局部特征提取组件所构成，他们的输入首先执行PatchExpanding操作放大特征图的尺寸。其输入数据格式为B × C × H × W，表示为Batch、channel、height和width。Forecaster中模块的结构如图3所示。Forecaster中模块使用的组件与Encoder模块中的组件是一样的。但在Forecaster模块中只考虑回波序列的空间特征，因此他们的处理是将输入的回波数据作为一个整体进行提取操作，但不考虑输入数据的时序性。Forecaster输出的特征图的通道数为96，然后再加入一个卷积核为1的卷积操作将其维度降为所需要预测的回波数目12。

Figure 3. Structural framework diagram of the space Forecaster module

图3. 空间Forecaster模块结构框架图

4. 实验结果与分析

为了验证提出算法模型的有效性，我们选择 Storm Event Imagery (SEVIR) [20]数据集中的NEXRAD天气雷达数据，该数据的空间分辨率是1公里，时间分辨率是5分钟，图像大小是384*384。同时该数据集包含了4个小时为一个时间段的20,393个事件。在实际验证过程中根据GPU和我们将这个数据集按照论文[21]中的划分方式将其划分为训练集、验证集和测试集三部分。模型输入6帧，预测未来的12帧，即根据过去的半小时回波情况，外推未来的1小时回波信息。在4小时的数据集中，采用时间跨度为12，训练集、验证集和测试集的个数分别为35,718，9060和12,159。为了有效验证我们提出模型，我们采用ConvLSTM，SmaAt-UNet，LPT-QPN和Rainformer等四个模型与我们提出的模型进行量化和可视化两种对比方式。所有模型都使用深度学习平台框架Pytorch进行实现。实验过程中采用了AdamW优化器，该优化器是在Adam基础上增加了权值衰减正则化用于提升训练模型的收敛速度。学习率的初始值设置为0.001，我们同时选择余弦退火调整学习率，模型训练结束的迭代次数设置为最大60次，在训练过程中我们取验证集中的CSI值最大的模型用于测试集。为了量化模型的预报效果，我们采用命中率(probability of detection, POD)和临界成功子数(Critical Success Index, CSI)两个准则来量化预报结果，这两个指标的值越大表示预测的效果越好。我们使用HITS = (真实 = 1，预测 = 1)，FALSE = (真实 = 0，预测 = 1)，MISSES = (真实 = 1，预测 = 0)来表示真实情况和预测情况。其中POD和CSI的定义分别如公式5和6所示。

$POD=\frac{HITS}{HITS+FALSE}$ (5)

$CSI=\frac{HITS}{HITS+FALSE+MISSES}$ (6)

4.1. 实验结果量化对比

在实际的气象预报业务应用中对不同强度的回波值的关注度不一样，回波值越高表示发生的强对流天气的强度越大，因此为了更好的验证提出模型在预报强对流天气方面的作用及预报效果，我们也采用论文中将回波值划分为16，74，133，160，181，219等6个区间进行量化对比验证。我们对所给对比数据取得最好结果的记录用黑体标注，排在次好的记录用下划线表示。表1给出的是各个模型的平均POD值及6个像素值区间的POD值。我们从表1可以看出提出的模型在得到的平均POD值和多个阶段的POD值都远远高于排在第二位模型的数据。我们提出的模型得到的POD平均值比排在第二的LPT-QPN高接近2.7个百分点，同时提出模型得到的POD值在几个高强度值区间都高于其他模型，其中最多高的在POD-133中高近9个百分点，这表明提出模型在预报强对流天气时具有更好的优势。在POD值方面只是在219时比最好模型LPT-QPN的值低接近1.5。量化实验结果表明提出的模型具有非常好的预报能力，这表明提出模型的有效性。

Table 1. Quantitative comparison of the POD for various levels of all models

表1. 模型对比的各级POD值

	POD-M	POD-16	POD-74	POD-133	POD-160	POD-181	POD-219
ConvLSTM	0.4694	0.9006	0.7711	0.4579	0.3076	0.2570	0.1219
SmaAt-UNet	0.4437	0.9139	0.7778	0.4071	0.2677	0.2051	0.0905
LPT-QPN	0.5063	0.9138	0.7975	0.5033	0.3466	0.2973	0.1793
RainFormer	0.4982	0.9125	0.7831	0.4926	0.3432	0.2890	0.1687
RainbiFormer	0.5319	0.9179	0.8281	0.5925	0.3819	0.3051	0.1651

表2给出的是各个模型在测试集中取得的CSI平均值及各个区间的CSI值，从表中的数据可以看出，在量化指标CSI值方面没有哪个模型得到非常大的优势结果。SmarAt-UNet得到的效果最低，这可能是SmarAt-UNet模型只考虑了回波序列的空间特征而完全没有考虑时序特征的结果。ConvLSTM模型因为将回波序列的时空特征都进行了考虑并进行时空特征提取，因此得到了比SmarAt-UNet模型更好的结果，ConvLSTM模型得到的CSI平均值比SmarAt-UNet模型高2.2个百分点。三个基于Transformer的模型得到比ConvLSTM更好的CSI结果，这表明Transformer相对于CNN能够更加有效提取回波目标的全局特征信息，因此模型能够取得更好的外推结果。CSI平均值中LPT-QPN，RainFormer和我们提出的模型都非常接近，但我们提出的模型在高强度值133,160和181中都取得最好的结果，同时在最高强度值上取得的CSI值仅比最好的RainFormer仅仅相差0.23个百分点。这表明相对于其他模型，提出的模型对强对流预测结果具有更好和更综合的预测结果。

Table 2. CSI values for each level of model comparison

表2. 模型对比的各级CSI值

	CSI-M	CSI-16	CSI-74	CSI-133	CSI-160	CSI-181	CSI-219
ConvLSTM	0.4013	0.7546	0.6704	0.3772	0.2680	0.2258	0.1116
SmaAt-UNet	0.3784	0.7394	0.6654	0.3476	0.2418	0.1892	0.0867
LPT-QPN	0.4135	0.7408	0.6711	0.3890	0.2867	0.2478	0.1455
RainFormer	0.4175	0.7483	0.6758	0.3944	0.2917	0.2489	0.1457
RainbiFormer	0.4165	0.7387	0.6676	0.4015	0.2938	0.2540	0.1434

4.2. 实验结果视觉对比

为了量化图4中给出的预报结果，我们在表3和表4中分别给出预报的平均CSI和POD值，第10、20、30、40、50和60分钟的CSI及POD值。从表3中我们可以看出，提出模型在预报的CSI平均值及30分钟和其后的时间里的CSI值都是最高的，特别是在第60分钟进行的预报结果高于排在第二位模型接近4个百分点。而对于POD评价标准，提出模型在均值和所有时间段的结果都优于其他模型。通过这两个表给出的量化标准表明提出的模型RainbiFormer在较长预报时间时的结果优于其他模型，这表明我们提出模型能够更好的预测更长时间的结果，这充分证明了提出模型的有效性。

Table 3. Comparative results of CSI values at different time intervals for visual contrasts

表3. 视觉对比的不同时间间隔的CSI值对比结果

	CSI-M	CSI-10	CSI-20	CSI-30	CSI-40	CSI-50	CSI-60
ConvLSTM	0.5826	0.7073	0.6205	0.5943	0.5500	0.4839	0.4421
SmaAt-UNet	0.5737	0.7278	0.6025	0.5619	0.5296	0.4788	0.4426
LPT-QPN	0.5836	0.7216	0.6260	0.5808	0.5379	0.4898	0.4505
RainFormer	0.5870	0.7359	0.6445	0.5956	0.5451	0.4739	0.4365
RainbiFormer	0.6065	0.7310	0.6399	0.5971	0.5645	0.5310	0.4891

Table 4. Comparative results of POD values at different time intervals for visual comparison

表4. 视觉对比的不同时间间隔的POD值对比结果

	POD-M	POD-10	POD-20	POD-30	POD-40	POD-50	POD-60
ConvLSTM	0.6983	0.7963	0.7285	0.7114	0.6803	0.6098		0.5679
SmaAt-UNet	0.6540	0.8072	0.6834	0.6421	0.6126	0.5586		0.5185
LPT-QPN	0.7193	0.8342	0.7654	0.7375	0.6842	0.6269		0.5880
RainFormer	0.7094	0.8464	0.7608	0.7217	0.6804	0.6025		0.5680
RainbiFormer	0.7695	0.8602	0.7925	0.7638	0.7444	0.7097		0.6611

Figure 4. Visual comparison of all models

图4. 所有模型的视觉对比

5. 结论与展望

本研究针对临近预报应用中对强对流天气预报准确度不高的情况，提出了一种基于Transformer的时序空间预报模型，该模型采用Encoder-Forecaster架构，每个Encoder和Forecaster包含四个模块。Encoder中的模块采用时序空间处理方式对每个回波序列按时序分别进行单独和整体两种方式的处理，这样能够有效提取回波序列的时序和空间特征，从而实现对回波轨迹运动规律的高效建模。Forecaster端包含的四个模块的输入包括前一模块的输出和Encoder端对应模块的输出，这两个输出进行拼接后进行融合，这样能够有效组合回波特征。所有模块中使用了提取目标对象的局部特征的最大池化操作和均值池化操作及提取全局特征的双级路由注意力机制，这两种特征的结合实现了模型对高强度值回波的高准确度预报，实验结果验证了这一假设。

本研究提出的模型主要还是专注于从大量的气象数据中学习回波的运动轨迹，但没有考虑回波运动轨迹的气象物理规律及天气预报的不确定性规律。因此我们未来的研究将着重从深度学习模型提取回波运动的物理规律入手，同时考虑结合生成模型[22]建模天气预报的不确定性规律。通过这些方法研究的模型可能提高对高强度回波值的预测，从而能够提升对强对流天气预报结果的准确性。