[论文解读] Adapting Pre-trained Vision Transformers from 2D to 3D through Weight Inflation Improves Medical Image Segmentation

Yuhui Zhang, Shih-Cheng Huang+3 more|arXiv (Cornell University)|Feb 8, 2023

Advanced Neural Network Applications被引用 4

一句话总结

本文提出了一种简单而有效的权重膨胀策略，用于将预训练的2D视觉Transformer适配到3D医学图像分割任务中，实现在保留深度信息的同时进行迁移学习。通过使用中心初始化方法将2D权重膨胀为3D，并采用小尺寸3D感受野结合中心切片预测，该方法在仅增加0.75%计算成本的情况下，在12个不同的3D医学影像数据集上实现了最先进性能。

ABSTRACT

Given the prevalence of 3D medical imaging technologies such as MRI and CT that are widely used in diagnosing and treating diverse diseases, 3D segmentation is one of the fundamental tasks of medical image analysis. Recently, Transformer-based models have started to achieve state-of-the-art performances across many vision tasks, through pre-training on large-scale natural image benchmark datasets. While works on medical image analysis have also begun to explore Transformer-based models, there is currently no optimal strategy to effectively leverage pre-trained Transformers, primarily due to the difference in dimensionality between 2D natural images and 3D medical images. Existing solutions either split 3D images into 2D slices and predict each slice independently, thereby losing crucial depth-wise information, or modify the Transformer architecture to support 3D inputs without leveraging pre-trained weights. In this work, we use a simple yet effective weight inflation strategy to adapt pre-trained Transformers from 2D to 3D, retaining the benefit of both transfer learning and depth information. We further investigate the effectiveness of transfer from different pre-training sources and objectives. Our approach achieves state-of-the-art performances across a broad range of 3D medical image datasets, and can become a standard strategy easily utilized by all work on Transformer-based models for 3D medical images, to maximize performance.

研究动机与目标

• 为解决将预训练的2D视觉Transformer适配到3D医学图像分割任务中的挑战，同时保留迁移学习的优势与深度信息。
• 探究权重膨胀是否能有效弥合2D自然图像与3D医学扫描之间的域差距。
• 建立一种可泛化、可标准化的策略，用于在3D医学视觉任务中利用预训练的2D模型。
• 评估不同预训练来源与目标对3D分割性能的影响。
• 在多样化的3D医学影像数据集上展示一致的性能提升，且仅需极少的超参数调优。

提出的方法

• 通过中心化策略对预训练2D视觉Transformer的权重进行膨胀，即将中心切片的权重复制，其余切片初始化为零，从而实现从2D到3D的权重扩展。
• 将3D输入体积分割为包含少数邻近切片的小型3D感受野，以保留空间上下文信息。
• 仅在每个感受野的中心切片上进行分割预测，模拟残差学习机制，实现深度信息的逐步融合。
• 采用标准的3D U-Net架构，将膨胀后的2D Transformer作为编码器，实现端到端的分割。
• 通过消融实验比较不同的膨胀策略，包括预测目标、权重初始化方式、步长以及输入切片数量。
• 在11个额外的3D医学影像数据集上评估泛化性能，涵盖不同解剖区域与成像模态（CT/MRI）。

Figure 1: Approach overview. Large-scale pre-trained Transformers are used as the encoder in the segmentation model for transfer learning, in which weights are adapted using the inflation strategy to support 3D inputs. Each 3D image is split into windows, which contain a small number of neighbor sli

实验结果

研究问题

• RQ1 在不修改网络架构的前提下，权重膨胀是否能有效将预训练的2D视觉Transformer适配到3D医学图像分割任务？
• RQ2 预训练来源的选择（自监督 vs. 有监督）如何影响3D医学分割中的迁移性能？
• RQ3 在使用权重膨胀进行3D适配时，最优的感受野尺寸与预测策略（如仅预测中心切片）是什么？
• RQ4 所提出的方法是否能在具有不同解剖结构与成像协议的多样化3D医学影像数据集上实现良好泛化？
• RQ5 与标准的2D或3D随机初始化相比，权重膨胀在性能提升方面有多大优势？

主要发现

• 所提出的权重膨胀策略在仅增加0.75%计算成本的情况下，相较于强基线模型，将平均Dice相似系数（DSC）提升了1.95%。
• 使用预训练权重初始化的模型在多器官分割数据集上的表现比随机初始化的3D模型高出11.18%，凸显了迁移学习的重要性。
• 中心化膨胀策略——即将2D权重复制到中心切片，其余切片初始化为零——取得了最佳性能，表明渐进式深度信息整合是有效的。
• 仅包含少数邻近切片（如3–5个）已足够实现良好性能，因为相邻切片提供的上下文信息比远距离切片更相关。
• 该方法在12个数据集上均达到最先进性能，包括主评估之外的11个额外数据集，展现出强大的泛化能力。
• 有监督预训练来源的迁移性能优于自监督预训练来源，表明标签丰富的预训练对3D医学分割更具优势。

Figure 2: Transfer learning reduces noise in local predictions, and depth information further improves prediction consistency across slices. We show five consecutive input images, ground-truth labels, predictions from Ours , Ours w/o D , Ours w/o T&D .