Reproduce “Adapting Pre-trained Vision Transformers from 2D to 3D through Weight Inflation Improves Medical Image Segmentation” on OpenResearch

About this paper

Yuhui Zhang, Shih-Cheng Huang, Zhengping Zhou, Matthew P. Lungren, Serena Yeung

Researchers from Stanford University developed a weight inflation method to adapt 2D pre-trained Vision Transformers for 3D medical image segmentation. This approach successfully bridges the 2D-3D dimensionality gap, achieving an 87.13% Dice Similarity Coefficient on the BCV dataset and outperforming existing methods by combining transfer learning benefits with crucial depth-wise information.

Key insights

• Transfer learning from large-scale pre-trained models is crucial for performance, leading to an 11.18% improvement in Dice Similarity Coefficient (DSC) over random initialization on the BCV dataset.
• Incorporating depth information is valuable, with a randomly initialized 3D model outperforming a 2D slice-based model by 6.07% DSC.
• Centering inflation for weight adaptation and predicting only the center slice within a small window (e.g., 5 slices) proved to be the most effective strategy for 3D medical image segmentation.

Summary by alphaXiv.