通过潜在表示编码压缩高分辨率数据以降低大规模AI天气预报模型的尺度

Qian Liu2†, Bing Gong3†, Xiaoran Zhuang4†, Xiaohui Zhong1, Zhiming Kang5*, Hao Li1,2*

1复旦大学人工智能创新与孵化研究院，上海，200433，中国。2上海科学人工智能研究院，上海，200232，中国。3上海师范大学电子与计算机工程系，上海，200234，中国。4江苏省气象观测站，南京，江苏，210008，中国。5南京气象观测站，南京，江苏，210009，中国。*通讯作者。电子邮件(s): kangzm@cma.gov.cn; lihao lh@fudan.edu.cn; 贡献作者: liuqian@sais.com.cn; gongbing1112@gmail.com; zxrxz3212009@163.com; x7zhong@gmail.com; †这些作者对这项工作做出了同等贡献。

研究背景

随着人工智能（AI）在天气研究中的快速发展，其能够从大型高维数据集中学习的能力是推动这一进步的关键因素。然而，这一进步也带来了显著的挑战，尤其是在处理大量数据和计算资源限制方面的成本。本研究受到计算机视觉领域神经图像压缩（NIC）任务的启发，旨在通过压缩天气数据来应对这些挑战，并提高下游应用的效率。

研究意义

天气预报对社会和各行各业至关重要，它支持着农业、交通、灾害管理等领域的决策。传统的数值天气预报（NWP）模型虽然被广泛使用，但计算密集度高，需要大量的计算资源。深度学习的最新进展为NWP模型提供了有希望的替代方案，可能提供更快、同样准确的预报。然而，深度学习在天气和气候中的应用效果在很大程度上依赖于大规模数据集的可用性。获取、存储和管理这些数据的计算需求经常超出了只有有限设置的研究人员的能力，为那些缺乏高性能计算资源和数据存储的人创造了重大障碍。此外，为了确保准确的天气预报，全球众多超级计算和研究中心每天都要进行多次天气和气候模拟。例如，欧洲中期天气预报中心（ECMWF）管理着2300PB的数据，并每天处理约6亿个地球观测数据。由于预测模型的空间分辨率不断提高，预计在未来十年内，数据量将增加四倍。虽然数据量的增长为训练深度学习模型提供了更多机会，但它也带来了挑战，因为庞大的数据量可能会压垮现有的超级计算基础设施，并由于有限的网络带宽而复杂化天气产品的分发。因此，有效的数据压缩技术至关重要。

研究方法和数据

数据来源

为了评估所提出的压缩框架的性能，研究者使用HRCLDAS作为高分辨率真实数据进行降尺度作为概念验证应用。HRCLDAS是一个融合了站点观测、卫星数据和NWP数据的混合数据集，采用了先进的陆地表面和数据同化技术。在这项研究中，HRCLDAS数据包括2019年至2021年三年的数据，包括每小时的2米温度（T2M）、10米风的u分量（U10M）和10米风的v分量（V10M）。其原始尺寸为4384×6880×3，覆盖从15°到55°的纬度和从75°到135°的经度。研究者使用了2019年5月至10月和2020年7月至2021年8月的数据作为训练数据集，并使用2021年9月的数据作为数据压缩任务的测试集。对于降尺度任务，研究者使用了FuXi-2.0模型的预报作为低分辨率输入。FuXi-2.0模型是一个级联的机器学习天气预报模型，提供15天的全球预报，时间分辨率为1小时，空间分辨率为0.25°。在这项研究中，研究者选择了从1到24小时的预报，初始化时间为00 UTC和12 UTC，用于训练集。研究者从FuXi预报中选择了40个输入变量用于降尺度任务，这些变量在表1中有详细说明。数据集涵盖了从2019年5月至10月和2020年7月至2021年8月的数据，而测试集包括了2019年11月至2020年6月的数据，用于测试降尺度性能。在将数据输入降尺度模型之前，所有输入和输出变量都使用z分数归一化进行了标准化。

数据压缩框架概述

研究者提出了一个基于变分自编码器（VAE）的框架，用于压缩高分辨率天气数据。VAE的编码器通过变分推理过程生成量化的潜在高斯分布。通过比较训练策略，研究者发现预训练后进行微调可以获得最佳的重建性能，从而获得对原始数据的更优潜在表示。

实施细节

研究者首先训练了一个用于数据压缩的VAE模型。模型使用批量大小为8进行训练，并使用Adam优化器，学习率为1.6e-5进行优化。考虑到原始HRCLDAS数据的高分辨率，研究者将其分割成1000×1000的块，纬度方向有重叠，每个数据集得到35个块。为了降低计算成本并加速收敛，研究者进一步将这些1000×1000的块分割成更小的256×256的块。研究者在这些较小的块上对VAE模型进行了前10个周期的预训练，然后使用更大的块对其进行了额外5个周期的微调。与自然图像不同，天气变量之间的相互关系较弱。因此，研究者分别训练了T2m、U10M和V10M的单独VAE模型。对于降尺度模型的训练，研究者使用了从训练有素的VAE编码器生成的压缩数据Zt作为真实数据。研究者使用批量大小为16训练了50个周期的U-Net模型。模型使用Adam优化器，学习率为3.2e-5进行优化。所有模型都使用PyTorch框架实现，并在8个NVIDIA A100 GPU上进行训练。

研究结论

研究者使用所提出的VAE框架压缩了三年的高分辨率HRCLDAS数据，将数据大小从8.61TB减少到紧凑的204GB，这是一个令人印象深刻的42倍存储需求减少。结果表明，VAE框架通过微调策略优于其他基线方法，显著降低了重建误差，从而更准确地重建了2米温度、10米U分量风和10米V分量风。此外，基于VAE的模型在保留极值方面表现出色，有效地保持了原始HRCLDAS数据集中的高值和低值。为了进一步验证VAE框架生成的紧凑数据的有效性和可用性，研究者将其用于降尺度任务。性能指标显示，原始HRCLDAS数据和紧凑HRCLDAS数据之间在MSE和SSIM方面没有显著差异。尽管使用紧凑数据可能会由于U-Net模型使用的参数数量减少而导致性能略有下降，但根据功率谱的进一步评估表明，原始HRCLDAS数据和紧凑数据都显著优于基线双线性插值方法。虽然使用紧凑数据在更大尺度上会导致信息的轻微损失，但它在所有三个目标变量的更细尺度上有效地保留了关键信息，与原始HRCLDAS数据相似。这证明了紧凑数据用于下游降尺度任务的可用性和有效性。

不足与讨论

研究者指出，尽管VAE压缩方法在降尺度任务中表现出色，但在捕捉T2M的细尺度特征方面存在局限性。这一局限性可能是由于中国某些地区，如喜马拉雅山脉地区，地形对维持局部温度梯度至关重要。由于这项研究的主要重点是验证VAE压缩方法，因此当前模型中没有纳入地形数据。然而，这将在未来的工作中考虑。

未来工作方向

研究者提出，这种方法有潜力扩展到其他应用领域，如大规模天气和气候预测，在这些领域中，处理大量数据集的需求更为关键。此外，研究者还计划在未来的研究中考虑地形数据，以提高模型在捕捉细尺度特征方面的性能。

通过这项研究，我们可以看到，通过使用深度学习和机器学习技术，我们可以有效地处理和压缩大量的天气数据，这对于提高天气预报的准确性和效率具有重要意义。随着计算能力的提高和算法的改进，我们有望在未来实现更快速、更准确的天气预报。

文章中提到的VAE框架具体是如何工作的？

文章中提到的VAE（变分自编码器）框架是为了压缩高分辨率的天气数据集而设计的。VAE框架通过学习数据的潜在表示来压缩数据，同时尽可能保留关键信息。具体来说，VAE框架包括以下几个关键步骤：

1. 编码器（Encoder）：编码器是VAE框架的第一部分，它负责将高维输入数据压缩成一个低维的潜在表示。在这篇文章中，编码器通过一系列卷积层和残差块（ResNet blocks）来提取数据的特征，并逐步降低数据的空间维度。编码器的输出是一个潜在空间的分布参数，通常是一个均值（µ）和一个标准差（σ）。

2. 潜在空间（Latent Space）：编码器的输出是潜在空间中的一个分布，通常假设为高斯分布。这个分布的参数（均值和标准差）定义了输入数据在潜在空间中的表示。

3. 采样过程（Sampling）：由于潜在空间的表示是一个概率分布，因此在重建数据之前需要从这个分布中采样一个点。这是通过使用重参数化技巧（reparameterization trick）来实现的，它允许在训练过程中通过反向传播梯度。

4. 解码器（Decoder）：解码器是VAE框架的第二部分，它负责将潜在空间的表示重建回原始数据。解码器的结构通常与编码器对称，但操作是反向的。它使用一系列上采样层和卷积层将潜在表示逐步恢复到原始数据的空间维度。

5. 重建损失（Reconstruction Loss）：重建损失是衡量解码器输出与原始输入数据之间差异的指标。在这篇文章中，作者使用了Charbonnier损失函数，这是一种对异常值具有鲁棒性的损失函数。

6. Kullback-Leibler散度（KL Divergence）：KL散度是一种正则化项，用于确保潜在空间的分布接近先验分布，通常是标准正态分布。这有助于防止过拟合，并确保潜在空间的表示具有一般性。

7. 优化（Optimization）：VAE的训练涉及到最小化重建损失和KL散度的总和。通过优化这个过程，VAE学习如何有效地压缩和重建数据，同时保留关键信息。

作者提到了几种不同的训练策略，包括预训练后微调，以及分别对单一变量进行训练。这些策略有助于提高VAE框架在压缩天气数据时的性能。通过这种压缩方法，作者成功地将3年的HRCLDAS数据从8.61 TB压缩到204 GB，同时在尺度降低任务中保持了与原始数据相当的准确性。这表明VAE框架在压缩高分辨率天气数据方面的有效性和潜力。

研究中使用的数据压缩技术有哪些优势和局限性？

研究中使用的数据压缩技术，即变分自编码器（VAE）框架，具有以下优势和局限性：

优势

1. 高效的数据压缩：VAE框架能够显著减少数据的存储需求，本研究中将8.61 TB的数据压缩到204 GB，实现了约42倍的数据压缩率，这对于处理大规模天气数据集尤为重要。

2. 保留关键信息：尽管数据被大幅度压缩，VAE框架能够保留原始数据集中的关键信息，如极端值，这对于天气预报的准确性至关重要。

3. 提高计算效率：通过减少数据的体积，VAE框架能够降低下游应用的计算成本，使得在资源有限的环境下也能进行高效的天气预测。

4. 适用性：压缩后的数据在下游任务，如尺度降低中表现良好，表明该技术不仅适用于数据存储，也适用于提高数据处理的效率。

5. 灵活性：VAE框架可以针对特定的数据集进行定制和优化，例如，通过调整网络结构或训练策略来适应不同的数据特性。

6. 潜力：该技术展示了将深度学习技术应用于天气数据压缩的潜力，为未来天气研究和气候模型的开发提供了新的方向。

局限性

1. 计算资源需求：虽然VAE框架减少了数据存储的需求，但训练VAE模型本身需要大量的计算资源，这可能限制了其在资源受限环境中的应用。

2. 模型复杂性：VAE模型相对复杂，需要精心设计和调整以确保有效的数据压缩和信息保留，这可能增加了模型训练和调优的难度。

3. 对极端地形的敏感性：研究中提到，模型在处理具有极端地形变化的地区的数据时存在挑战，如喜马拉雅山脉地区的温度梯度，这表明模型可能需要进一步的改进以更好地处理复杂的地形。

4. 信息损失：尽管VAE框架在尺度降低任务中表现良好，但使用压缩数据可能会导致一定程度的信息损失，尤其是在较大尺度的特征上。

5. 泛化能力：虽然VAE框架在本研究中表现出色，但其在其他类型的天气数据集和不同地区的泛化能力还需要进一步验证。

6. 模型解释性：深度学习模型通常被认为是“黑箱”模型，其内部工作机制不易解释。这可能限制了模型在需要高度可解释性的应用中的使用。

VAE框架在数据压缩方面展示了显著的优势，尤其是在处理大规模天气数据集时。然而，它也有一些局限性，需要进一步的研究和改进以提高其在实际应用中的性能和可靠性。

声明：本公众号分享的资源均来自于互联网的公开资源，仅供个人学习和研究使用，不代表本人对其版权、真实性、合法性、准确性等方面做出任何保证。任何使用本号分享的资源的个人或组织，需自行承担相应的法律责任和后果；本号不承担任何法律责任和连带责任。使用本号分享的资源即代表您已完全了解并同意本免责声明的所有内容。