Hi-End-MAE: Hierarchical encoder-driven masked autoencoders are stronger vision learners for medical image segmentation

📄 arXiv: 2502.08347v1 📥 PDF

作者: Fenghe Tang, Qingsong Yao, Wenxin Ma, Chenxu Wu, Zihang Jiang, S. Kevin Zhou

分类: cs.CV

发布日期: 2025-02-12

备注: 19 pages, Code: https://github.com/FengheTan9/Hi-End-MAE

🔗 代码/项目: GITHUB

💡 一句话要点

Hi-End-MAE：分层编码器驱动的掩码自编码器提升医学图像分割性能

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 医学图像分割 Vision Transformer 掩码图像建模 预训练 迁移学习

📋 核心要点

  1. 医学图像分割面临标签数据不足的挑战，现有ViT-MIM方法未能充分利用各层特征。
  2. Hi-End-MAE提出编码器驱动重建和分层密集解码，提升特征学习和信息利用。
  3. 在10K CT数据集上预训练，并在七个基准测试中表现出卓越的迁移学习能力。

📝 摘要（中文）

由于标签数据的稀缺性，医学图像分割仍然是一个巨大的挑战。通过在大型无标签医学数据集上使用掩码图像建模(MIM)预训练Vision Transformer (ViT)提供了一个有希望的解决方案，为各种下游任务提供计算效率和模型泛化能力。然而，目前基于ViT的MIM预训练框架主要强调输出层中的局部聚合表示，而未能利用不同ViT层中丰富的表示，这些表示可以更好地捕捉更精确的医学下游任务所需的细粒度语义信息。为了填补上述空白，我们提出了分层编码器驱动的MAE (Hi-End-MAE)，这是一个简单而有效的基于ViT的预训练解决方案，它以两个关键创新为中心：(1) 编码器驱动的重建，鼓励编码器学习更多信息性特征来指导掩码块的重建；(2) 分层密集解码，它实现了一种分层解码结构来捕获不同层中丰富的表示。我们在一个包含1万张CT扫描的大规模数据集上预训练了Hi-End-MAE，并在七个公共医学图像分割基准上评估了它的性能。大量的实验表明，Hi-End-MAE在各种下游任务中实现了卓越的迁移学习能力，揭示了ViT在医学成像应用中的潜力。

🔬 方法详解

问题定义：医学图像分割任务中，现有基于ViT的掩码图像建模预训练方法主要关注输出层的局部聚合特征，忽略了ViT编码器不同层级所蕴含的丰富语义信息，导致模型无法有效捕捉细粒度的特征，从而影响分割精度。现有方法未能充分利用编码器各层的特征，导致下游任务性能受限。

核心思路：Hi-End-MAE的核心思路是充分利用ViT编码器各层的特征信息，通过编码器驱动的重建任务和分层密集解码结构，促使模型学习更具信息量的特征表示，从而提升医学图像分割的性能。该方法旨在弥补现有方法对编码器多层级特征利用不足的缺陷。

技术框架：Hi-End-MAE的整体框架包括ViT编码器、掩码模块、编码器驱动的重建模块和分层密集解码器。首先，输入图像被分割成patch并随机掩码。然后，ViT编码器提取图像特征。编码器驱动的重建模块利用编码器的特征指导掩码区域的重建。最后，分层密集解码器融合来自编码器不同层的特征，生成最终的分割结果。

关键创新：Hi-End-MAE的关键创新在于两个方面：一是编码器驱动的重建，它改变了以往仅依赖解码器进行重建的方式，而是利用编码器的特征来指导重建过程，从而促使编码器学习更具信息量的特征。二是分层密集解码，它通过融合编码器不同层的特征，充分利用了ViT编码器各层所蕴含的语义信息，从而提升分割精度。

关键设计：在编码器驱动的重建中，使用了编码器输出的特征图来指导解码器的重建过程，具体实现方式未知。分层密集解码器采用了一种自上而下的方式，将编码器不同层的特征图逐层融合，并使用卷积操作进行特征提取和融合。损失函数方面，可能使用了像素级别的交叉熵损失或Dice损失来优化分割结果。具体的参数设置和网络结构细节在论文中可能有所描述，但此处未知。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Hi-End-MAE在七个公共医学图像分割基准上进行了评估，实验结果表明，该方法在各种下游任务中实现了卓越的迁移学习能力，显著优于现有的基于ViT的MIM预训练方法。具体的性能提升数据未知，但论文强调了其在医学成像应用中的潜力。

🎯 应用场景

Hi-End-MAE具有广泛的应用前景，可应用于多种医学图像分割任务，例如CT、MRI等影像的器官分割、病灶检测等。该研究成果有助于提高医学图像分析的自动化程度和准确性，辅助医生进行诊断和治疗，具有重要的临床应用价值。未来，该方法有望推广到其他医学影像分析任务中，例如图像配准、图像重建等。

📄 摘要（原文）

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