遥感大模型解读系列 | 02 SelectiveMAE:基于大规模遥感数据集的高效掩码自编码器

新开一个文章系列，将对现有的遥感大模型论文进行解读，主要以foundation model为主，将涉及图像解译、视觉-语言、图像生成、音频-图像及其他特定任务模型，预计将持续数十篇。

本篇论文——Scaling Efficient Masked Autoencoder Learning on Large Remote Sensing Dataset

• 发布时间：2024 年 6 月 17 日

• 来源：https://arxiv.org/abs/2406.11933

SelectiveMAE:基于大规模遥感数据集的高效掩码自编码器

摘要

掩码图像建模（Masked Image Modeling, MIM）已成为开发基础视觉模型的重要方法。然而，当前的遥感（Remote Sensing, RS）数据集在数量和多样性上都存在局限性，严重制约了MIM方法学习通用表示的能力。

在本研究中，作者提出了RS-4M，这是一个大规模数据集，旨在使MIM训练在RS图像上变得高效。RS-4M包含400万张光学图像，涵盖丰富且细粒度的RS视觉任务，包括目标检测和像素级分割。相比自然图像，RS图像通常包含大量冗余背景像素，限制了传统MIM模型的训练效率。

为了解决这个问题，作者提出了一种高效的MIM方法，称为SelectiveMAE，它基于语义丰富度动态编码和重建一部分补丁标记。SelectiveMAE基于渐进语义标记选择模块，从重建语义相似标记到编码互补语义依赖，逐步演变。

该方法将传统的MIM训练转变为渐进特征学习过程，使SelectiveMAE能够高效地学习RS图像的鲁棒表示。大量实验表明，SelectiveMAE显著提高了训练效率，并增强了基线MIM模型在分类、检测和分割任务中的性能。

1. 引言

过去十年中，遥感技术和数据采集的进步显著推动了生态系统监测、自然灾害管理等应用的发展。这些应用依赖于场景分类、目标检测、变化检测和语义分割等基本能力。然而，每个下游任务通常需要大量计算资源来学习特定任务的特征表示并开发专门的模型。

由于自监督学习方法的显著进步，如掩码图像建模（MIM）技术，视觉基础模型的预训练取得了显著改进。因此，遥感基础模型（RSFMs）最近出现，提供了通用特征表示，并在各种遥感下游任务中表现出色。

然而，RSFMs的发展面临两个挑战：

（i）与ImageNet-21k数据集相比，之前的RS数据集样本数量显著较少，阻碍了大型骨干网络的充分MIM训练。
（ii）这些数据集主要关注全局场景语义，缺乏下游任务中遇到的RS场景的多样性和细粒度信息，限制了学习表示的泛化能力。

为了解决这些挑战，作者提出了一个大规模RS数据集RS-4M。该数据集包含400万张光学图像，旨在充分利用MIM方法在RS应用中的表示学习能力。RS-4M显著超过了之前的RS数据集，至少大四倍。此外，RS-4M涵盖了下游任务中遇到的各种RS场景，如目标检测和像素级分割。

尽管在使用MIM方法训练RSFM方面做出了大量努力，如利用通用图像知识、扩展参数规模、集成时空信息和学习多尺度特征，但在大规模RS数据集上进行MIM训练时，计算负担和收敛速度缓慢的问题仍然存在。

特别是，使用ViT-B骨干网络在1百万RS样本上进行预训练需要107小时。为了解决这个问题，作者提出了一种高效的MIM方法，称为SelectiveMAE，它基于语义丰富度动态编码和重建一部分补丁标记。

2. 相关工作

2.1 遥感数据集

近年来，许多RS数据集被创建用于场景分类、目标检测和分割等任务。免费、未标记的卫星图像的可用性推动了大规模RS数据集的发展。例如，SEN12MS结合了多种传感器数据，而MillionAID提供了百万规模的RS光学图像分类数据集。

然而，这些数据集主要针对场景分类，通常忽略了细粒度目标信息，限制了它们在各种下游任务中的实用性。

2.2 遥感基础模型

自监督学习方法最近被用于从未标记的RS数据中提取表示。尽管有些对比自监督方法取得了一定进展，但最近的研究主要集中在生成自监督方法上。

例如，RVSA利用视觉变换器骨干网络，并引入计算高效的旋转变尺寸窗口注意机制替代原始注意机制。此外，许多研究致力于通过利用通用图像知识、扩大参数规模、集成时空数据和处理多传感器数据来改进生成自监督算法。

然而，这些方法尚未有效解决自监督预训练在RS中的巨大计算负担问题。

2.3 掩码图像建模

受掩码语言建模（MLM）在自然语言处理中的成功启发，MIM被开发用于视觉预训练。MIM通过重建被掩码的标记来学习图像表示，关注各种回归目标、掩码策略和重建方法。MAE表明预测像素值与使用复杂目标同样有效。

MIM的主要挑战是其高计算需求和长预训练时间。为缓解这一问题，一些研究使用不对称编码器-解码器策略，减少输入补丁或使用新颖的损失函数。然而，这些方法未能考虑RS图像的独特特征，如稀疏的前景信息和复杂的背景。

3. 数据集

3.1 数据收集与组织

最近的自监督预训练RSFMs进展受到现有RS数据集规模和多样性较小的限制。为了解决这个问题，作者提出了RS-4M数据集，这是一个持续更新的大规模RS数据集，遵循多样性、丰富性和可扩展性（DiRS）原则。

作者收集和组织了过去十年的数据集，重点关注主流RS图像解释任务，包括场景分类、目标检测、语义分割和变化检测。这些数据集包含约100个子数据集，总计1350万张图像。排除了多光谱和SAR数据，仅保留光学图像。

未来的更新将纳入多光谱和SAR数据，以支持多模态自监督预训练。与MillionAID等数据集不同，数据集旨在促进下游检测和分割任务。

为了处理分割数据集中的大规模图像，随机裁剪高分辨率图像为较小的切片。像素低于64或高于1024的图像被移除，以促进自监督预训练。结合剩余图像，并使用两阶段过程消除重复图像：粗略阶段使用感知哈希，精细阶段进行人工审查。

通过哈希值选择和人工审查，仅排除高度相似的图像。最终，RS-4M数据集包含约400万张高质量的RS光学图像，比之前的代表性RS数据集大四倍。

3.2 RS-4M数据集的初步评估

RS-4M提供了比之前数据集大且多样性更高的图像集。为了评估其效果，作者使用vanilla MAE方法预训练一个ViT基础模型。为了比较，使用MillionAID数据集，并在训练期间保持相同数量的数据点：对于MillionAID的100万张图像训练800个epoch，对于RS-4M数据集训练200个epoch。

此外，通过从RS-4M中采样不同数量的图像创建新数据集进行进一步比较。预训练后，在下游RESISC-45数据集上微调模型，并评估Top-1分类准确率。结果表明，RS-4M数据集在各种样本规模下均优于MillionAID，使用整个数据集时达到最高准确率92.38%。即使使用较少的图像，RS-4M也始终提供更优的性能，突显了其更高的多样性和有效性。

这些发现强调了使用更大和更多样化的数据集进行无监督预训练的优势。

4. 方法

4.1 掩码自编码器预备知识

1. 掩码。与标准ViT的监督训练类似，MAE将图像划分为规则的、不重叠的补丁。然后，它从这些补丁中采样一部分并掩码其余部分。典型的掩码比例为75%，意味着只有25%的补丁输入编码器。这个随机采样遵循掩码比例的均匀分布。

2. MAE编码器。编码器是一个标准的ViT，仅应用于可见的未掩码补丁。它线性投影补丁，添加位置嵌入，并通过一系列变换器块处理它们。通过处理较小的补丁子集，编码器使得大模型的训练具有较低的计算和内存需求。

3. MAE解码器。编码标记和掩码标记被输入解码器，解码器由具有自注意层的变换器块组成。掩码标记是共享的、可学习的张量，增强了位置嵌入。解码器仅在预训练期间使用，生成这些掩码标记的输出预测。

4. 重建目标。MAE预测每个掩码补丁的像素值，解码器输出中的每个元素表示一个补丁的像素值向量。损失函数计算重建目标与原始补丁之间的均方误差（MSE）。

RS光学图像通常包含大量冗余背景像素，这与MAE的掩码策略非常吻合。在MAE中，编码器仅处理25%的补丁，显著减少了计算负担。基于MAE，作者希望利用RS图像中的冗余来加速训练。

具体来说，要解决两个问题：

1）在RS图像中，是否有必要重建所有掩码补丁？
2）是否可以进一步压缩输入到MAE编码器的可见补丁，以增强加速效果？

4.2 部分重建

对于问题1，之前的研究表明，对于一般图像，当MAE重建75%的补丁以计算损失时，设计良好的解码器不需要完全重建所有剩余补丁。实际上，仅重建50%甚至25%的补丁即可实现类似的性能并加速训练。然而，对于RS图像，如果随机采样补丁并移除大部分进行重建，重建的补丁可能不是语义丰富的补丁。使用随机子集进行重建会降低性能。

为了解决这个问题，选择语义丰富的补丁进行重建，而不是随机选择。
具体来说，给定一个输入图像，它被重塑为非重叠的补丁。这些补丁然后线性映射到补丁嵌入。为了保留位置信息，位置嵌入被添加到补丁中。根据掩码比例选择一部分补丁输入编码器，剩余的补丁作为解码器的重建目标，引入一个新的重建比例，即要重建的像素比例，计算剩余补丁的HOG特征，并根据重建比例选择具有高HOG特征值的补丁，而不是使用所有补丁。

4.3 试验结果

作者在RS-4M数据集上进行了大量实验,验证了SelectiveMAE方法的有效性。主要结果包括:

1) 训练效率:与基线方法相比,SelectiveMAE将预训练速度提高了2.2-2.7倍。

2) 下游任务性能:在目标检测任务上,SelectiveMAE相比基线方法提升了5.6%的mAP;在语义分割任务上,提升了1.4%的mF1分数。

3) 稳定性:通过PSTS模块,SelectiveMAE能够有效避免高掩码率导致的梯度爆炸问题,保证训练的稳定性。

如图4所示,基线方法在使用高掩码率时容易出现梯度爆炸,而SelectiveMAE通过PSTS模块实现了稳定的训练过程。

5. 结论与展望

本文提出的RS-4M数据集和SelectiveMAE方法为遥感领域的自监督学习开辟了新的方向。RS-4M数据集的大规模和多样性为模型训练提供了充足的数据支持,而SelectiveMAE方法则有效解决了遥感图像中冗余信息带来的训练效率问题。

未来的研究方向可能包括:

1) 进一步扩展RS-4M数据集,纳入多光谱、SAR等多模态数据。

2) 探索更高效的语义令牌选择策略,进一步提升训练效率。

3) 将SelectiveMAE方法应用于更多遥感下游任务,验证其泛化能力。