通过分层编码器引导重建,提升医学图像分割的预训练效果。
Hi-End-MAE: Hierarchical encoder-driven masked autoencoders are stronger vision learners for medical image segmentation
- 用编码器引导重建,增强特征表达能力。
- 在10,000张CT扫描上预训练,7个基准测试表现领先。
- 适合需要高精度分割的医学视觉任务研究者。
医学图像分割因标注数据稀缺而面临巨大挑战。在大规模未标注医学数据上通过掩码图像建模(MIM)预训练视觉变换器(ViT),可有效提升模型在下游任务中的计算效率与泛化能力。然而,现有基于ViT的MIM框架多聚焦输出层的局部聚合表征,未能充分利用不同层级间的丰富语义信息,限制了对精细医学任务的建模能力。为此,本文提出分层编码器驱动的掩码自编码器(Hi-End-MAE),包含两项核心创新:(1) 编码器驱动重建,促使编码器学习更具信息量的特征以指导被遮蔽块的重构;(2) 分层密集解码,构建跨层级的解码结构以捕捉多层次表示。我们在10,000张CT扫描数据集上预训练该模型,并在七个公开医学图像分割基准上评估其性能。大量实验表明,Hi-End-MAE在多种下游任务中展现出卓越的迁移学习能力,揭示了ViT在医学影像应用中的巨大潜力。代码已开源。
原文摘要 · Abstract (English)
Medical image segmentation remains a formidable challenge due to the label scarcity. Pre-training Vision Transformer (ViT) through masked image modeling (MIM) on large-scale unlabeled medical datasets presents a promising solution, providing both computational efficiency and model generalization for various downstream tasks. However, current ViT-based MIM pre-training frameworks predominantly emphasize local aggregation representations in output layers and fail to exploit the rich representations across different ViT layers that better capture fine-grained semantic information needed for more precise medical downstream tasks. To fill the above gap, we hereby present Hierarchical Encoder-driven MAE (Hi-End-MAE), a simple yet effective ViT-based pre-training solution, which centers on two key innovations: (1) Encoder-driven reconstruction, which encourages the encoder to learn more informative features to guide the reconstruction of masked patches; and (2) Hierarchical dense decoding, which implements a hierarchical decoding structure to capture rich representations across different layers. We pre-train Hi-End-MAE on a large-scale dataset of 10K CT scans and evaluated its performance across seven public medical image segmentation benchmarks. Extensive experiments demonstrate that Hi-End-MAE achieves superior transfer learning capabilities across various downstream tasks, revealing the potential of ViT in medical imaging applications. The code is available at: https://github.com/FengheTan9/Hi-End-MAE
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。