[论文解读] Foundation Models for Generalist Geospatial Artificial Intelligence

Johannes Jakubik, Sujit Roy+31 more|arXiv (Cornell University)|Oct 28, 2023

Flood Risk Assessment and Management被引用 11

一句话总结

介绍 Prithvi，这是一个100M参数的地理空间基础模型，使用1TB的 Harmonized Landsat-Sentinel-2 数据进行预训练，在云层 gap 填充、洪水/野火/作物分割，以及数据高效学习方面展示微调成功，并在 HuggingFace 开源发布。

ABSTRACT

Significant progress in the development of highly adaptable and reusable Artificial Intelligence (AI) models is expected to have a significant impact on Earth science and remote sensing. Foundation models are pre-trained on large unlabeled datasets through self-supervision, and then fine-tuned for various downstream tasks with small labeled datasets. This paper introduces a first-of-a-kind framework for the efficient pre-training and fine-tuning of foundational models on extensive geospatial data. We have utilized this framework to create Prithvi, a transformer-based geospatial foundational model pre-trained on more than 1TB of multispectral satellite imagery from the Harmonized Landsat-Sentinel 2 (HLS) dataset. Our study demonstrates the efficacy of our framework in successfully fine-tuning Prithvi to a range of Earth observation tasks that have not been tackled by previous work on foundation models involving multi-temporal cloud gap imputation, flood mapping, wildfire scar segmentation, and multi-temporal crop segmentation. Our experiments show that the pre-trained model accelerates the fine-tuning process compared to leveraging randomly initialized weights. In addition, pre-trained Prithvi compares well against the state-of-the-art, e.g., outperforming a conditional GAN model in multi-temporal cloud imputation by up to 5pp (or 5.7%) in the structural similarity index. Finally, due to the limited availability of labeled data in the field of Earth observation, we gradually reduce the quantity of available labeled data for refining the model to evaluate data efficiency and demonstrate that data can be decreased significantly without affecting the model's accuracy. The pre-trained 100 million parameter model and corresponding fine-tuning workflows have been released publicly as open source contributions to the global Earth sciences community through Hugging Face.

研究动机与目标

• 在地球科学中推动使用基础模型，以处理大规模未标注的多传感器遥感数据。
• 开发一个可扩展的框架，在大型多光谱时序数据上直接进行地理空间基础模型的预训练和微调。
• 展示 Prithvi 在有限有标注数据下对多样化下游任务的适应能力并评估数据效率。
• 提供模型权重、架构和推理工具的开源访问，以促进地球科学社区。

提出的方法

• 提出一个分布式、可扩展的框架，将数据发现、预处理、预训练和推理链接起来，适用于地理空间数据。
• 使用带 ViT 主干的掩码自编码器（MAE），在六个 HLS 波段上对 Prithvi（100M 参数）进行预训练。
• 用 3D 时空嵌入（3D 位置嵌入和 3D 补丁嵌入）扩展 MAE，以处理多时相、多光谱的输入。
• 在预训练期间，使用 Song 基于的数据加载与 Zarr，实现高效流式加载，降低 I/O 瓶颈。
• 对预训练好的编码器进行微调（解码头部为任务特定），用于下游任务，借助 mmsegmentation 及面向任务定制的头部和损失函数。
• 评估数据加载效率并比较微调策略（全模型、仅解码器、无预训练的基线）。

实验结果

研究问题

• RQ1 RQ1: 地球科学中设计和评估基础模型的关键因素是什么？
• RQ2 RQ2: 在重复、嘈杂的遥感数据上，如何高效地预训练基础模型，同时消除噪声和冗余？
• RQ3 RQ3: 基础模型能否利用多样的训练特征，在显著少于标注数据的情况下实现跨地球科学领域的泛化？

主要发现

• 预训练的 Prithvi 相较于随机初始化权重，显著加速了微调。
• Prithvi 在多时相云层插补方面的表现优于条件 GAN，提升幅度可达 5 个百分点（SSI 中为 5.7%）。
• 在地理空间统计量上执行分层采样能获得多样化的预训练数据，减少对过度代表景观的偏见。
• 基于 Zarr 的数据加载显著提升了 epoch 时间，优于 GeoTiff 加载，并能够扩展到大型 GPU 集群。
• 在 3D 时空补丁上的 MAE 预训练能够有效处理多时相、多光谱卫星影像。
• 在 HuggingFace 上开源发布 Prithvi-100M 权重和框架，支持可复现性和社区协作。