在前沿平台上对巨型空间基础模型进行预训练
人工智能
2024-04-19 v1
摘要
随着AI工作负载范围的不断扩大，小型任务特定模型的泛化能力变得具有挑战性，而且对大量标记训练样本的需求也随之增加。相反，基础模型（Foundation Models）通过自监督学习在互联网规模的无标签数据上进行预训练，已显示出能够以最小的微调适应各种任务。虽然大型基础模型在自然语言处理和计算机视觉领域已显示出显著影响，但针对空间应用的基础模型研究仍受限于较小规模的模型，因为预训练更大的模型需要配备最新硬件加速器的巨大计算资源。当前的卫星星座每天收集超过100TB的数据，导致图像包含数十亿像素且具有多模态特性。这种空间数据提出了独特的挑战，为开发基础模型开辟了新的机遇。我们通过在公开可用数据上进行预训练，研究针对空间应用的巨型基础模型及其HPC训练配置。我们从端到端地研究了通过扩大模型规模所带来的性能和影响。我们更大的30亿参数规模模型在与10亿参数模型进行比较时，实现了最高30%的top1场景分类准确率提升。此外，我们详细介绍了在美国首座超算系统Froniter上进行的性能实验，其中我们研究了使用PyTorch全分片数据并行库所需的不同模型和数据并行方法。具体而言，我们研究了Vision Transformer架构（ViT）的不同变体，进行了规模可达150亿参数的ViT模型的性能分析。通过在不同的并行配置下讨论吞吐量和性能瓶颈，提供了在开发用于空间图像应用的大型模型时，如何利用这类领导级HPC资源的见解。
引用
@article{arxiv.2404.11706,
title = {Pretraining Billion-scale Geospatial Foundational Models on Frontier},
author = {Aristeidis Tsaris and Philipe Ambrozio Dias and Abhishek Potnis and Junqi Yin and Feiyi Wang and Dalton Lunga},
journal= {arXiv preprint arXiv:2404.11706},
year = {2024}
}