[论文解读] Swin MAE: Masked Autoencoders for Small Datasets

Zi'an Xu, Yin Dai+6 more|arXiv (Cornell University)|Dec 28, 2022

AI in cancer detection被引用 11

一句话总结

Swin MAE 在小规模医学图像数据集上提出了一种基于 Swin Transformer 主干网络的掩码自编码器，用于无监督表示学习。在未使用预训练权重的情况下，其迁移学习性能与在 ImageNet 上微调的监督式 Swin Transformer 模型相当或略优，证明了仅通过数千张图像的掩码图像重建即可学习到强大的特征表示。

ABSTRACT

The development of deep learning models in medical image analysis is majorly limited by the lack of large-sized and well-annotated datasets. Unsupervised learning does not require labels and is more suitable for solving medical image analysis problems. However, most of the current unsupervised learning methods need to be applied to large datasets. To make unsupervised learning applicable to small datasets, we proposed Swin MAE, which is a masked autoencoder with Swin Transformer as its backbone. Even on a dataset of only a few thousand medical images and without using any pre-trained models, Swin MAE is still able to learn useful semantic features purely from images. It can equal or even slightly outperform the supervised model obtained by Swin Transformer trained on ImageNet in terms of the transfer learning results of downstream tasks. The code is publicly available at https://github.com/Zian-Xu/Swin-MAE.

研究动机与目标

• 解决深度学习中受限且标注良好的医学图像数据集的挑战。
• 通过引入 Swin Transformer 的归纳偏置，克服视觉 Transformer 对大规模预训练的依赖。
• 在无需数据增强的前提下，实现在小规模数据集上的有效无监督表示学习，而数据增强在医学图像中往往不可行。
• 仅通过在小规模医学图像上进行自监督预训练，实现在下游任务中的竞争力迁移学习性能。
• 证明掩码自编码与 Swin MAE 可在无任何预训练权重的情况下，达到或超越在 ImageNet 上微调的监督模型性能。

提出的方法

• 采用掩码自编码器（MAE）框架，在预训练过程中随机掩码 75% 的图像块。
• 使用 Swin Transformer 作为编码器主干网络，通过移位窗口机制引入局部归纳偏置，提升在小数据集上的训练稳定性。
• 采用基于块的图像标记化方法，使用 4×4 的非重叠块，与 MAE 的 16×16 块不同，以更好地适配医学图像特征。
• 实施基于窗口的掩码策略，将整个 4×4 块作为单位进行掩码，以保持重建任务中的空间连贯性。
• 使用原始图像与重建图像之间的均方误差（MSE）损失进行训练，重点在于预测被掩码区域。
• 采用轻量级解码器头，从掩码的潜在表示中重建完整图像，且 MAE 与 Swin MAE 采用相同的架构设计以确保公平比较。

实验结果

研究问题

• RQ1 在未使用任何预训练权重的情况下，基于 Swin Transformer 主干的掩码自编码器能否从小规模医学图像数据集中学习到有用的语义特征？
• RQ2 在小数据集上，Swin MAE 与标准 MAE 相比，在收敛速度和训练稳定性方面表现如何？
• RQ3 与在 ImageNet 上微调的监督模型相比，Swin MAE 的无监督预训练在多大程度上能提升下游任务的迁移学习性能？
• RQ4 与标准视觉 Transformer 相比，Swin Transformer 的归纳偏置是否能提升在小规模、低资源医学图像数据集上的自监督学习能力？
• RQ5 在不依赖数据增强的前提下，掩码自编码能否有效应用于医学图像，而数据增强在单通道、完整医学扫描中往往不切实际？

主要发现

• Swin MAE 在下游任务中的迁移学习性能与在 ImageNet 上预训练的监督式 Swin Transformer 模型相当或略优，尽管其仅在数千张医学图像上进行训练。
• 与标准 MAE 相比，Swin MAE 在相同小数据集和相同超参数设置下收敛更快，训练损失更低，且训练曲线更平滑。
• Swin MAE 在小数据集上优于 MAE，归因于 Swin Transformer 的归纳偏置，该偏置在无预训练的情况下也稳定了训练并提升了特征学习能力。
• 重建结果表明，Swin MAE 能够有效恢复缺失的图像区域，包括边界轮廓和颜色一致性，即使在 75% 的掩码率下也表现良好。
• 缺乏预训练权重并未影响性能——Swin MAE 仅通过在小规模数据上的自监督重建即可学习到鲁棒的表示。
• 上游训练损失曲线表明，Swin MAE 在小数据集上的优化稳定性更好，收敛速度更快，表明其具备更强的无监督学习能力。