Pretraining Billion-scale Geospatial Foundational Models on Frontier

随着AI工作负载范围的扩大，小型任务特定模型的泛化能力变得具有挑战性，它们对大量标记训练样本的需求也增加了。相反，基础模型（FMs）通过自监督学习使用互联网规模的未标记数据进行训练，并已证明可以在最少的微调下适应各种任务。尽管大型FMs在自然语言处理和计算机视觉方面已经展现了显著的影响，但是针对地理空间应用的FMs的努力仅限于较小的模型，因为预训练更大的模型需要配备最先进的硬件加速器的非常大的计算资源。当前的卫星星座每天收集超过100TB的数据，导致图像具有数十亿像素和多模式特性。这种地理空间数据提出了独特的挑战，为开发FMs开辟了新的机会。我们通过在公开可用的数据上进行预训练，研究了十亿级别的FMs和高性能计算（HPC）训练配置在地理空间应用中的性能和影响。我们通过扩展模型大小，从端到端地研究了解决方案的性能和影响。我们的更大的3B参数模型在与100M参数模型比较时，实现了高达30％的顶部1场景分类准确性的提高。此外，我们详细介绍了在美国第一台亿级计算机“Frontier”上进行的性能实验，研究了使用PyTorch的Fully Sharded Data Parallel库的不同模型和数据并行方法的性能。具体而言，我们研究了Vision Transformer架构（ViT）的变体，对大小高达15B参数的ViT模型进行了性能分析。通过讨论不同并行配置下的吞吐量和性能瓶颈，我们提供了有关如何在开发用于地理空间图像应用的大型模型时利用这种领先级HPC资源的见解。

本论文尝试解决在地理空间应用中使用Foundation Models（FM）的问题，尤其是对于大规模数据和更大的模型，需要更多的计算资源和硬件加速器。

论文通过在公开数据集上预训练大型FM，使用PyTorch的Fully Sharded Data Parallel库，研究了不同模型和数据并行方法，在Frontier超级计算机上进行了性能实验。论文的关键思路是使用大规模的FM进行自监督学习，以适应各种任务，从而实现地理空间数据的高精度分类和识别。

论文使用了公开数据集进行预训练，并在Frontier超级计算机上进行了性能实验。实验结果表明，相比于100M参数模型，3B参数模型在场景分类准确率上有30%的提升。此外，论文还探讨了不同的模型和数据并行方法，并提供了如何利用领先的HPC资源开发大型地理空间图像应用的见解。

最近的相关研究集中在自监督学习和大型模型的应用上，如GPT、BERT等。