基于Transformer的不确定量化地磁指数预测框架

A Transformer-based Framework for Predicting Geomagnetic Indices with Uncertainty Quantification

A transformer-based framework for predicting geomagnetic indices with uncertainty quantification | Journal of Intelligent Information Systems (springer.com)

1. Abstract

地磁活动对地球有着至关重要的影响，能够波及航天器和电网系统。地球空间科学家使用一种名为Kp指数的地磁指数来描述地磁活动的总体水平。该指数是地球磁场扰动的重要指标，被美国空间天气预报中心用作可能受扰动影响的用户的警报和预警服务。另一个常用的指数，称为Ap指数，是由Kp指数转换而来的。对Kp和Ap指数的早期和准确预测对于防灾准备和风险管理至关重要。

本文提出了一种名为GNet的深度学习框架，用于对Kp和Ap指数进行短期预测。具体来说，GNet以NASA空间科学数据协调档案馆提供的太阳风参数值的时间序列作为输入，并针对给定的时间点t预测在未来t+w小时（其中w的范围为1到9）的Kp和Ap指数。GNet结合了Transformer编码器块与贝叶斯推理，能够量化预测中的随机不确定性（数据不确定性）和认知不确定性（模型不确定性）。实验结果表明，在均方根误差和R平方分数方面，GNet优于密切相关的机器学习方法。此外，GNet还能提供数据和模型不确定性的量化结果，这是现有方法无法提供的。这可能是首次将贝叶斯Transformer用于地磁活动预测。

2 Data

2.1 Data collection

本研究中使用的Kp和Ap值来自德国地球科学研究中心GFZ，可通过“Data - Kp index”访问。GFZ网站上的Kp和Ap值均以3小时为间隔进行收集。本研究中使用的太阳风参数来自NASA空间科学数据协调档案馆，可通过NASA空间科学数据协调档案馆的公共页面访问。本文选择太阳风参数的每小时平均值作为时间分辨率。本文收集了2002年1月1日至2022年12月31日期间的所有数据。

遵循Tan等人（2018）和Zhelavskaya等人（2019）的研究，本文考虑了八个太阳风参数，即行星际磁场（IMF）幅度平均值、磁场Bx、By和Bz分量、等离子体温度、质子密度、等离子体速度和流动压力。由于Kp和Ap使用3小时的时间间隔，而太阳风参数使用1小时的时间间隔（如King和Papitashvili (2005)以及Lethy等人(2018)所述），本文通过将Kp和Ap值与太阳风参数在时间上进行匹配来处理数据，以创建最终数据集。

2.2 Data labeling

在每个实验中，都有两组数据：训练集和测试集。本文基于GNet框架准备了两个模型。一个模型标记为GNet_k，用于Kp预测；另一个模型标记为GNet_a，用于Ap预测。数据记录按以下方式标记。设t为给定的时间点，w为t之前的时间窗口，其中w的范围为1到9小时。本文进行w小时前的预测。对于GNet_k（或GNet_a），在时间点t的记录标签定义为在时间点t+w的Kp值（或Ap值）。训练集中每个时间点t的训练记录有九个值，包括在时间点t的八个太阳风参数值和训练记录的标签。测试集中每个测试记录仅包含八个太阳风参数值；测试集中测试记录的标签将由训练好的GNet_k（或GNet_a）进行预测。

3 Methodology

3.1 The GNet architecture

图1展示了GNet的架构，该架构是使用TensorFlow的Keras框架创建的，可通过https://www.tensorflow.org 访问。为了增强GNet的学习能力和性能，本文在网络中添加了多层。GNet的输入由非重叠的记录序列x_p+1, x_p+2, …, x_p+n组成，其中在本文的研究中n被设置为512。这些序列被传递到一个具有64个内核的一维卷积（Conv1D）层，其中每个内核的大小为1。Conv1D已被证明非常适合于序列数据，并且之前也用于地磁指数预测（Siciliano等人，2021年）；它从输入数据序列中学习内部模式，并将它们传递给配置有300个神经元的双向长短期记忆（biLSTM）层。如我们后续消融研究所示，在处理时间序列时，结合Conv1D和biLSTM层可以显著提高性能。

biLSTM层将学习到的模式传递到变换器网络，该网络由b个变换器编码器块（TEB）组成。（遵循Vaswani等人（2017）的研究，在本研究中将b设置为8。）每个TEB由一个多头注意力层（Vaswani等人，2017）、一个批量归一化层和一个前馈网络组成。通常，用于自然语言处理（NLP）的变换器使用层归一化，这可以带来比批量归一化更显著的性能提升（Vaswani等人，2017）。然而，我们在这里使用批量归一化，以避免时间序列中不存在于NLP词嵌入中的异常值的影响（Zerveas等人，2021）。多头注意力层对顺序输入值进行转换，以获得大小为h的不同度量指标。在这里，h是注意力头的数量，设置为4，每个注意力头的大小也设置为4；其他参数保留其默认值。前馈网络由一个具有4个内核的Conv1D层组成，内核数量等于注意力头的数量，且每个内核大小为1，后面跟着一个dropout层，然后是一个具有250个神经元的biLSTM层。dropout层有助于减少过拟合（Srivastava等人，2014；Lim，2021）。请注意，由于我们处理的是时间序列而不是语言处理，因此不涉及翻译，所以每个TEB都使用不带解码器的变换器编码器（Zerveas等人，2021）。

此外，还添加了一个具有10个神经元的密集变分层（DVL）（Tran等人，2019），该层使用变分推理（Blei等人，2016）来近似模型权重上的后验分布。DVL与常规的密集层相似，但需要两个输入函数来定义模型权重上的先验分布和后验分布。DVL允许GNet通过分布而不是估计点来表示权重。此外，GNet还包括几个密集层和dropout层。每个密集层都与其前一层紧密连接，其中密集层中的每个神经元都与前一层中的每个神经元连接。每个dropout层都提供了一种机制，用于随机丢弃一定比例的隐藏神经元，以避免过拟合（Srivastava等人，2014；Lim，2021）。表2总结了本研究中使用的参数及其值。

Fig. 1

(a) Architecture of GNet. (b) Configuration details of a transformer encoder block (TEB) used in GNet. Each TEB has two dropout layers, a multi-head attention layer, a batch normalization layer, followed by a feed forward network that contains a Conv1D layer, and a bidirectional LSTM (biLSTM) layer

3.2 Uncertainty quantification

由于dropout层的存在，每次丢弃神经元时，本文框架的内部结构都会略有不同。这对于依赖于随机采样的蒙特卡洛（MC）类算法来说是一种重要行为，并能提供有用的信息（Gal 和 Ghahramani，2016）。我们利用这一技术引入了预测值的分布区间，如第4节所示。具体来说，为了量化我们GNet框架的不确定性，我们使用了一个先验概率P(W)，该概率是关于网络权重W的。在训练过程中，八个太阳风参数值和Kp值（或Ap值，分别对应）统称为D，用于训练GNet_k（或GNet_a）。根据贝叶斯定理，

计算精确的后验概率 P(W|D) 是不可行的，但本文可以使用变分推理（Graves，2011）来学习网络权重上由 θ 参数化的变分分布 q_θ(W)，通过最小化 q_θ(W) 和 P(W|D) 之间的 Kullback-Leibler (KL) 散度。根据 Gal 和 Ghahramani（2016）的研究，带有 dropout 层的网络提供了变分近似。为了最小化 KL 散度，我们使用图 1 中所示的密集变分层（DVL），并将 KL 权重分配为 1/N，其中 N 是训练集的大小。我们使用均方误差（MSE）损失函数和自适应矩估计（Adam）优化器，学习率设置为 0.0001，来训练 GNet_k（或 GNet_a）。令 θ 表示通过训练 GNet_k（或 GNet_a）获得的优化变分参数；我们使用 q_θ(W) 来表示优化后的权重分布。

在测试/预测期间，GNet_k（或 GNet_a）利用 MC dropout 采样技术产生概率预测结果，量化偶然不确定性和认知不确定性。通过处理测试数据 K 次，dropout 用于检索 K 个 MC 样本（Gal 和 Ghahramani，2016）。（在本研究中，K 设置为 100。）对于 K 个 MC 样本中的每一个，都随机从 q_θ^(W) 中抽取一组权重。对于每个预测的 Kp（或 Ap）值，本文在 K 个样本上得到一个均值和一个方差。根据 Jiang 等人（2021）和 Kwon 等人（2020）的研究，本文将方差分解为偶然不确定性和认知不确定性。偶然不确定性捕获了预测结果的固有随机性，这来自于输入测试数据。另一方面，认知不确定性来自于 W 的可变性，这反映了网络参数（权重）的不确定性。